Update app.py

app.py

@@ -194,32 +194,39 @@ class QASystem:
             logger.error(f"System initialization error: {str(e)}")
             return False
 
-    def process_query(self, query: str) -> List[Dict[str, str]]:
+    def process_query(self, query: str) -> Dict[str, str]:
+        """Process a query and return a single final response"""
         try:
-            responses = []
+            # Generate a unique thread ID for production use
+            # For simplicity, using a fixed ID here
+            thread_id = "abc123"
             
-            # 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(
+            # Use invoke instead of stream to get only the final result
+            final_state = self.graph.invoke(
                 {"messages": [HumanMessage(content=query)]},
-                stream_mode="values",
                 config={"configurable": {"thread_id": thread_id}}
-            ):
-                if step["messages"]:
-                    # 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
-                        })
+            )
+            
+            # Extract only the last AI message from the final state
+            ai_messages = [m for m in final_state["messages"] if m.type == "ai"]
+            
+            if ai_messages:
+                # Return only the last AI message
+                return {
+                    'content': ai_messages[-1].content,
+                    'type': ai_messages[-1].type
+                }
+            return {
+                'content': "No response generated",
+                'type': 'error'
+            }
             
-            return responses
         except Exception as e:
             logger.error(f"Query processing error: {str(e)}")
-            return [{'content': f"Query processing error: {str(e)}", 'type': 'error'}]
+            return {
+                'content': f"Query processing error: {str(e)}",
+                'type': 'error'
+            }
 
 qa_system = QASystem()
 if qa_system.initialize_system():
@@ -229,5 +236,6 @@ else:
 
 @app.post("/query")
 async def query_api(query: str):
-    responses = qa_system.process_query(query)
-    return {"responses": responses}
+    """API endpoint that returns a single response for a query"""
+    response = qa_system.process_query(query)
+    return {"response": response}