Update app.py

app.py

@@ -7,7 +7,6 @@ from typing import List
 from langchain.agents.middleware import dynamic_prompt, ModelRequest
 from langchain.agents import create_agent
 from langchain_core.documents import Document
-from langchain_core.runnables import chain
 
 import gradio as gr
 import os
@@ -79,7 +78,7 @@ def initialize_vector_store(documents: List[Document]):
     vector_store = Milvus(
         embedding_function=embeddings,
         connection_args={"uri": uri},
-        index_params={"index_type": "FLAT", "metric_type": "L2"},
+        index_params={"index_type": "FLAT", "metric_type": "COSINE"},  # COSINE for semantic similarity
         drop_old=True
     )
     
@@ -89,15 +88,8 @@ def initialize_vector_store(documents: List[Document]):
     return vector_store
 
 # -----------------------------
-# Retriever
+# Context Formatting
 # -----------------------------
-@chain
-def create_retriever(vector_store):
-    """Create a retriever function with the vector store."""
-    def retriever(query: str) -> List[Document]:
-        return vector_store.similarity_search(query, k=K_RETRIEVE)
-    return retriever
-
 def format_context(docs: List[Document]) -> str:
     """
     Format retrieved documents with citations.
@@ -149,7 +141,7 @@ def create_prompt_middleware(vector_store):
             last_message = request.state["messages"][-1]
             last_query = getattr(last_message, "text", None) or getattr(last_message, "content", "")
             
-            # Retrieve relevant documents
+            # Retrieve relevant documents directly from vector store
             retrieved_docs = vector_store.similarity_search(last_query, k=K_RETRIEVE)
             docs_content = format_context(retrieved_docs)
             
@@ -206,15 +198,11 @@ def create_chat_function(agent):
                 last_message = step["messages"][-1]
                 results.append(last_message)
             
-            # Extract response content
-            # Try the standard approach first
-            if len(results) > 1 and hasattr(results[1], 'content'):
-                return results[1].content
-            
-            # Fallback: search through results for content
+            # Extract the latest assistant response
+            # Search from the end for the most recent content
             for msg in reversed(results):
                 content = getattr(msg, "content", None)
-                if content:
+                if content and content.strip():  # Ensure non-empty content
                     return content
             
             return "I apologize, but I couldn't generate a response. Please try rephrasing your question."