Update app.py

app.py

@@ -15,7 +15,7 @@ from qdrant_client.http.models import Distance, VectorParams
 from qdrant_client.models import PointIdsList
 
 from langgraph.graph import MessagesState, StateGraph
-from langchain_core.messages import SystemMessage, HumanMessage
+from langchain_core.messages import SystemMessage, HumanMessage, ToolMessage
 from langgraph.prebuilt import ToolNode
 from langgraph.graph import END
 from langgraph.prebuilt import tools_condition
@@ -114,52 +114,76 @@ class QASystem:
             
             graph_builder = StateGraph(MessagesState)
 
-            def query_or_respond(state: MessagesState):
-                retrieved_docs = [m for m in state["messages"] if m.type == "tool"]
-    
-                if retrieved_docs:
-                    context = ' '.join(m.content for m in retrieved_docs)
-                else:
-                    context = "mountain bicycle documentation knowledge"
-
-                system_prompt = (
-                    "You are an AI assistant embedded within the Interactive Electronic Technical Manual (IETM) for Mountain Cycles.. "
-                    "Always provide accurate responses with references to provided data. "
-                    "If the user query is not technical-specific, still respond from a IETM perspective."
-                    f"\n\nContext:\n{context}"
-                )
-    
-                messages = [SystemMessage(content=system_prompt)] + state["messages"]
-    
-                logger.info(f"Sending to LLM: {[m.content for m in messages]}")  # Debugging log
-
-                response = llm.invoke(messages)
-                return {"messages": [response]}
+            # Define a retrieval node that fetches relevant docs
+            def retrieve_docs(state: MessagesState):
+                # Get the most recent human message
+                human_messages = [m for m in state["messages"] if m.type == "human"]
+                if not human_messages:
+                    return {"messages": state["messages"]}
+                
+                user_query = human_messages[-1].content
+                logger.info(f"Retrieving documents for query: {user_query}")
+                
+                # Query the vector store
+                try:
+                    retrieved_docs = self.vector_store.similarity_search(user_query, k=3)
+                    
+                    # Create tool messages for each retrieved document
+                    tool_messages = []
+                    for i, doc in enumerate(retrieved_docs):
+                        tool_messages.append(
+                            ToolMessage(
+                                content=f"Document {i+1}: {doc.page_content}",
+                                tool_call_id=f"retrieval_{i}"
+                            )
+                        )
+                    
+                    logger.info(f"Retrieved {len(tool_messages)} relevant documents")
+                    return {"messages": state["messages"] + tool_messages}
+                
+                except Exception as e:
+                    logger.error(f"Error retrieving documents: {str(e)}")
+                    return {"messages": state["messages"]}
 
+            # Updated generate function that uses retrieved documents
             def generate(state: MessagesState):
-                retrieved_docs = [m for m in reversed(state["messages"]) if m.type == "tool"][::-1]
-    
-                context = ' '.join(m.content for m in retrieved_docs) if retrieved_docs else "mountain bicycle documentation knowledge"
+                # Extract retrieved documents (tool messages)
+                tool_messages = [m for m in state["messages"] if m.type == "tool"]
+                
+                # Collect context from retrieved documents
+                if tool_messages:
+                    context = "\n".join([m.content for m in tool_messages])
+                    logger.info(f"Using context from {len(tool_messages)} retrieved documents")
+                else:
+                    context = "No specific mountain bicycle documentation available."
+                    logger.info("No relevant documents retrieved, using default context")
 
                 system_prompt = (
                     "You are an AI assistant embedded within the Interactive Electronic Technical Manual (IETM) for Mountain Cycles. "
-                    "Your responses MUST be accurate, concise (5 sentences max)."
-                    f"\n\nContext:\n{context}"
+                    "Always provide accurate responses with references to provided data. "
+                    "If the user query is not technical-specific, still respond from a IETM perspective."
+                    f"\n\nContext from mountain bicycle documentation:\n{context}"
                 )
 
-                messages = [SystemMessage(content=system_prompt)] + state["messages"]
-    
-                logger.info(f"Sending to LLM: {[m.content for m in messages]}")  # Debugging log
-
+                # Get all messages excluding tool messages to avoid redundancy
+                human_and_ai_messages = [m for m in state["messages"] if m.type != "tool"]
+                
+                # Create the full message history for the LLM
+                messages = [SystemMessage(content=system_prompt)] + human_and_ai_messages
+                
+                logger.info(f"Sending query to LLM with {len(messages)} messages")
+                
+                # Generate the response
                 response = llm.invoke(messages)
-                return {"messages": [response]}
-
+                return {"messages": state["messages"] + [response]}
 
-            graph_builder.add_node("query_or_respond", query_or_respond)
+            # Add nodes to the graph
+            graph_builder.add_node("retrieve_docs", retrieve_docs)
             graph_builder.add_node("generate", generate)
             
-            graph_builder.set_entry_point("query_or_respond")
-            graph_builder.add_edge("query_or_respond", "generate")
+            # Set the flow of the graph
+            graph_builder.set_entry_point("retrieve_docs")
+            graph_builder.add_edge("retrieve_docs", "generate")
             graph_builder.add_edge("generate", END)
             
             self.memory = MemorySaver()
@@ -173,16 +197,25 @@ class QASystem:
     def process_query(self, query: str) -> List[Dict[str, str]]:
         try:
             responses = []
+            
+            # Use a unique thread_id for each conversation
+            thread_id = "abc123"  # In production, generate a unique ID for each conversation
+            
+            # Stream the responses
             for step in self.graph.stream(
                 {"messages": [HumanMessage(content=query)]},
                 stream_mode="values",
-                config={"configurable": {"thread_id": "abc123"}}
+                config={"configurable": {"thread_id": thread_id}}
             ):
                 if step["messages"]:
-                    responses.append({
-                        'content': step["messages"][-1].content,
-                        'type': step["messages"][-1].type
-                    })
+                    # Only include AI messages in the response
+                    ai_messages = [m for m in step["messages"] if m.type == "ai"]
+                    if ai_messages:
+                        responses.append({
+                            'content': ai_messages[-1].content,
+                            'type': ai_messages[-1].type
+                        })
+            
             return responses
         except Exception as e:
             logger.error(f"Query processing error: {str(e)}")
@@ -197,4 +230,4 @@ else:
 @app.post("/query")
 async def query_api(query: str):
     responses = qa_system.process_query(query)
-    return {"responses": responses}
+    return {"responses": responses}