basic model

app.py

@@ -5,26 +5,19 @@ from langchain.memory import ConversationBufferMemory
 from langchain_openai import ChatOpenAI, OpenAIEmbeddings
 from langchain.prompts import PromptTemplate
 from langchain_community.vectorstores import Chroma
-from queue import Queue
-from threading import Thread
 
-class StreamHandler:
-    def __init__(self, queue):
-        self.queue = queue
-        
-    def on_llm_new_token(self, token):
-        self.queue.put(token)
-
-def create_qa_chain(streaming_handler=None):
+def create_qa_chain():
     """
-    Create the QA chain with streaming capability
+    Create the QA chain with the loaded vectorstore
     """
+    # Initialize embeddings and load vectorstore
     embeddings = OpenAIEmbeddings()
     vectorstore = Chroma(
         persist_directory="./vectorstore",
         embedding_function=embeddings
     )
     
+    # Set up retriever
     retriever = vectorstore.as_retriever(
         search_type="mmr",
         search_kwargs={
@@ -34,12 +27,14 @@ def create_qa_chain(streaming_handler=None):
         }
     )
     
+    # Set up memory
     memory = ConversationBufferMemory(
         memory_key="chat_history",
         return_messages=True,
         output_key='answer'
     )
     
+    # Create prompt template
     qa_prompt = PromptTemplate.from_template("""You are an expert technical writer specializing in API documentation. 
 When describing API endpoints, structure your response in this exact format:
 
@@ -66,12 +61,11 @@ Question: {question}
 
 Technical answer (following the exact structure above):""")
     
+    # Create the chain
     qa_chain = ConversationalRetrievalChain.from_llm(
         llm=ChatOpenAI(
             temperature=0.1,
-            model_name="gpt-4-turbo-preview",
-            streaming=True,
-            callbacks=[streaming_handler] if streaming_handler else None
+            model_name="gpt-4-turbo-preview"
         ),
         retriever=retriever,
         memory=memory,
@@ -82,75 +76,54 @@ Technical answer (following the exact structure above):""")
     
     return qa_chain
 
-def predict(message, history):
+def chat(message, history):
     """
-    Process each message with streaming and handle quit commands
+    Process chat messages and return responses
     """
-    # Check for quit command
-    if message.lower() in ['quit', 'exit', 'bye']:
-        # Clear the QA chain if it exists
-        if hasattr(predict, 'qa_chain'):
-            predict.qa_chain.memory.clear()
-            delattr(predict, 'qa_chain')
-        return "Chat history cleared. Goodbye!"
-
-    token_queue = Queue()
-    stream_handler = StreamHandler(token_queue)
+    # Get or create QA chain
+    if not hasattr(chat, 'qa_chain'):
+        chat.qa_chain = create_qa_chain()
     
-    # Create or get QA chain
-    if not hasattr(predict, 'qa_chain'):
-        predict.qa_chain = create_qa_chain(stream_handler)
+    # Get response
+    result = chat.qa_chain({"question": message})
     
-    # Function to process the message and add to queue
-    def get_response():
-        result = predict.qa_chain({"question": message})
-        # Add sources to queue
-        sources = "\n\nSources:\n"
-        seen_components = set()
-        shown_sources = 0
-        
-        for doc in result["source_documents"]:
-            component = doc.metadata.get('component', '')
-            title = doc.metadata.get('title', '')
-            combo = (component, title)
-            
-            if combo not in seen_components and shown_sources < 3:
-                seen_components.add(combo)
-                shown_sources += 1
-                sources += f"\nSource {shown_sources}:\n"
-                sources += f"Title: {title}\n"
-                sources += f"Component: {component}\n"
-                sources += f"Content: {doc.page_content[:300]}...\n"
+    # Format sources
+    sources = "\n\nSources:\n"
+    seen_components = set()
+    shown_sources = 0
+    
+    for doc in result["source_documents"]:
+        component = doc.metadata.get('component', '')
+        title = doc.metadata.get('title', '')
+        combo = (component, title)
         
-        for char in sources:
-            token_queue.put(char)
-        token_queue.put(None)  # Signal end of response
+        if combo not in seen_components and shown_sources < 3:
+            seen_components.add(combo)
+            shown_sources += 1
+            sources += f"\nSource {shown_sources}:\n"
+            sources += f"Title: {title}\n"
+            sources += f"Component: {component}\n"
+            sources += f"Content: {doc.page_content[:300]}...\n"
     
-    # Start processing in a separate thread
-    thread = Thread(target=get_response)
-    thread.start()
+    # Combine response with sources
+    full_response = result["answer"] + sources
     
-    # Stream the response
-    response = ""
-    while True:
-        token = token_queue.get()
-        if token is None:
-            break
-        response += token
-        yield response
+    return full_response
 
-# Create the Gradio interface
-with gr.Blocks() as demo:
-    chatbot = gr.ChatInterface(
-        predict,
-        title="Apple Music API Documentation Assistant",
-        description="Ask questions about the Apple Music API documentation. Type 'quit', 'exit', or 'bye' to clear chat history.",
-        examples=[
-            "How to search for songs on Apple Music API?",
-            "What are the required parameters for searching songs?",
-            "Show me an example request with all parameters"
-        ]
-    )
+# Create and launch the Gradio interface
+demo = gr.ChatInterface(
+    chat,
+    title="Apple Music API Documentation Assistant",
+    description="Ask questions about the Apple Music API documentation.",
+    examples=[
+        "How to search for songs on Apple Music API?",
+        "What are the required parameters for searching songs?",
+        "Show me an example request with all parameters"
+    ],
+    retry_btn=None,
+    undo_btn=None,
+    clear_btn="Clear Chat"
+)
 
 if __name__ == "__main__":
-    demo.queue().launch()
+    demo.launch()