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()