Spaces:
Sleeping
Sleeping
File size: 3,459 Bytes
2197ab7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
import gradio as gr
import pyperclip
from src.fabrics_processor.database import initialize_qdrant_database
from src.search_qdrant.database_query import query_qdrant_database
from src.fabrics_processor.logger import setup_logger
import logging
import atexit
from src.fabrics_processor.config import config
import time
import os
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
# Configure logging
logger = setup_logger()
# Initialize the database client
client = None
def init_client():
global client
if client is None:
client = initialize_qdrant_database(api_key=os.environ.get("QDRANT_API_KEY"))
# Register cleanup function
atexit.register(lambda: client.close() if hasattr(client, '_transport') else None)
return client
def search_prompts(query):
"""Search for prompts based on the query."""
try:
client = init_client()
results = query_qdrant_database(
query=query,
client=client,
num_results=5,
collection_name=config.embedding.collection_name
)
if not results:
return gr.Radio(choices=[]), None
# Format results for radio buttons - just filenames
filenames = [r.metadata['filename'] for r in results]
# Store the full results for later use
global current_results
current_results = results
return gr.Radio(choices=filenames), None
except Exception as e:
logger.error(f"Error during search: {str(e)}")
return gr.Radio(choices=[]), None
def show_selected_prompt(selected_filename):
"""Display the content of the selected prompt."""
if not selected_filename or not current_results:
return ""
# Find the selected result
selected_prompt = next(
(r for r in current_results if r.metadata['filename'] == selected_filename),
None
)
if selected_prompt:
return selected_prompt.metadata['content']
return ""
# Store current results globally
global current_results
current_results = []
with gr.Blocks() as demo:
gr.Markdown("# Prompt finding and comparing")
with gr.Column():
query_input = gr.Textbox(
label="What are you trying to accomplish? I will then search for good prompts to give you a good start.",
lines=3,
autofocus=True, # This will focus the textbox when the page loads
interactive=True # This enables keyboard events
)
search_button = gr.Button("Search")
# Radio buttons for selecting prompts
results_radio = gr.Radio(
choices=[],
label="Select a prompt",
interactive=True
)
# Display area for selected prompt using Markdown
selected_prompt_display = gr.Markdown(label="Selected Prompt", show_copy_button=True)
# Set up event handlers
query_input.submit(
fn=search_prompts,
inputs=[query_input],
outputs=[results_radio, selected_prompt_display]
)
search_button.click(
fn=search_prompts,
inputs=[query_input],
outputs=[results_radio, selected_prompt_display]
)
results_radio.change(
fn=show_selected_prompt,
inputs=[results_radio],
outputs=[selected_prompt_display]
)
if __name__ == "__main__":
demo.launch(share=True)
|