Update app.py

app.py

@@ -1,89 +1,64 @@
-import os
-from openai import OpenAI
-import requests
 import gradio as gr
+import requests
 
-# Initialize OpenAI client
-client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
+API_URL = "http://154.12.226.68:8000"
 
-# Vector database search function
-def search_document(query, k=5):
-    url = "http://154.12.226.68:8000/search"
-    
-    payload = {
-        "text": query,
-        "k": k
-    }
+def search_document(index_name, query, k):
+    url = f"{API_URL}/search/{index_name}"
+    payload = {"text": query, "k": k}
+    headers = {"Content-Type": "application/json"}
+    response = requests.post(url, json=payload, headers=headers)
+    results = response.json()
     
-    headers = {
-        "Content-Type": "application/json"
-    }
+    formatted_results = []
+    for result in results.get('results', []):
+        metadata = result.get('metadata', {})
+        formatted_result = f"Source: {metadata.get('source', 'Unknown')}\n"
+        formatted_result += f"Page: {metadata.get('page', 'Unknown')}\n"
+        formatted_result += f"Content: {metadata.get('content', 'No content available')}\n"
+        formatted_result += f"Distance: {result.get('distance', 'Unknown')}\n"
+        formatted_results.append(formatted_result)
     
-    try:
-        response = requests.post(url, json=payload, headers=headers)
-        response.raise_for_status()
-        return response.json()
-    except requests.exceptions.RequestException as e:
-        return f"An error occurred: {e}"
+    return "\n\n".join(formatted_results)
 
-# Function to query OpenAI
-def query_openai(prompt):
-    try:
-        response = client.chat.completions.create(
-            model="gpt-4o-mini",
-            messages=[
-                {"role": "system", "content": "You are a helpful assistant. Answer the question based on the provided context."},
-                {"role": "user", "content": prompt}
-            ]
-        )
-        return response.choices[0].message.content
-    except Exception as e:
-        return f"An error occurred while querying OpenAI: {e}"
-
-# Function to perform vector search and format results
-def vector_search(query):
-    results = search_document(query)
-    if isinstance(results, str):  # Error occurred
-        return results
-    if not isinstance(results, dict) or 'results' not in results:
-        return "Unexpected format in vector database response."
+def qa_document(index_name, question, k):
+    url = f"{API_URL}/qa/{index_name}"
+    payload = {"text": question, "k": k}
+    headers = {"Content-Type": "application/json"}
+    response = requests.post(url, json=payload, headers=headers)
+    result = response.json()
     
-    formatted_results = ""
-    for i, result in enumerate(results['results'], 1):
-        content = result['metadata']['content']
-        source = f"Source {i}: {result['metadata'].get('source', 'Unknown source')}, page {result['metadata'].get('page', 'Unknown page')}"
-        metadata = ", ".join([f"{k}: {v}" for k, v in result['metadata'].items() if k != 'content'])
-        formatted_results += f"{source}\nMetadata: {metadata}\nContent: {content}\n\n"
+    answer = result.get('answer', 'No answer available')
+    sources = result.get('sources', [])
     
-    return formatted_results
-
-# Combined function for search and query
-def search_and_query(question):
-    # First, perform the vector search
-    search_results = vector_search(question)
+    formatted_sources = []
+    for source in sources:
+        formatted_source = f"Source: {source.get('source', 'Unknown')}\n"
+        formatted_source += f"Relevance Score: {source.get('relevance_score', 'Unknown')}"
+        formatted_sources.append(formatted_source)
     
-    # Then, use these results to query OpenAI
-    prompt = f"""Given the following context and question, provide a comprehensive and accurate answer. Use ONLY the information provided in the context to answer. If the context doesn't contain relevant information to answer the question, state that clearly.
+    formatted_result = f"Answer: {answer}\n\nSources:\n" + "\n\n".join(formatted_sources)
+    return formatted_result
 
-Context:
-{search_results}
-
-Question: {question}
-
-Answer:"""
+with gr.Blocks() as demo:
+    gr.Markdown("# Document Search and Question Answering System")
     
-    openai_response = query_openai(prompt)
+    index_name = gr.Textbox(label="Index Name", value="default")
     
-    # Return both the search results and the OpenAI response
-    return search_results, openai_response
-
-# Gradio interface
-with gr.Blocks() as demo:
-    question_input = gr.Textbox(label="Enter your question")
-    search_output = gr.Textbox(label="Vector Search Results")
-    answer_output = gr.Textbox(label="OpenAI Answer")
-    query_button = gr.Button("Get Answer")
+    with gr.Tab("Search"):
+        search_input = gr.Textbox(label="Search Query")
+        search_k = gr.Slider(1, 10, 5, step=1, label="Number of Results")
+        search_button = gr.Button("Search")
+        search_output = gr.Textbox(label="Search Results", lines=10)
+        
+        search_button.click(search_document, inputs=[index_name, search_input, search_k], outputs=search_output)
     
-    query_button.click(search_and_query, inputs=question_input, outputs=[search_output, answer_output])
+    with gr.Tab("Question Answering"):
+        qa_input = gr.Textbox(label="Question")
+        qa_k = gr.Slider(1, 10, 5, step=1, label="Number of Contexts to Consider")
+        qa_button = gr.Button("Ask Question")
+        qa_output = gr.Textbox(label="Answer and Sources", lines=10)
+        
+        qa_button.click(qa_document, inputs=[index_name, qa_input, qa_k], outputs=qa_output)
 
 demo.launch()