Replacing graph.aget(config) with state capture from on_graph_end

api/main.py

@@ -10,6 +10,7 @@ from fastapi import FastAPI, Query, HTTPException
 from fastapi.responses import StreamingResponse, HTMLResponse
 from fastapi.middleware.cors import CORSMiddleware
 from langchain_core.messages import AIMessageChunk
+from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage
 import json
 from uuid import uuid4
 import logging
@@ -20,6 +21,9 @@ import asyncio
 from core.system_initializer import setup_system
 from models.state_models import MultiCountryLegalState
 
+# utils functions
+from utils.helpers import message_obj_to_dict, dict_to_message_obj
+
 # Setup logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -185,92 +189,97 @@ def serialize_ai_message_chunk(chunk):
         raise TypeError(
             f"Object of type {type(chunk).__name__} is not correctly formatted for serialisation"
         )
-
-async def generate_legal_chat_responses(message: str, session_id: Optional[str] = None):
-    """Generate streaming responses for legal chat"""
-    if not system_initialized:
-        yield f"data: {json.dumps({'type': 'error', 'message': 'System is still starting up. Please try again in a moment.'})}\n\n"
-        yield f"data: {json.dumps({'type': 'end'})}\n\n"
-        return
-    
-    is_new_conversation = session_id is None
-    
-    if is_new_conversation:
+async def generate_legal_chat_responses(message: str, session_id: Optional[str] = None) -> str:
+    if not session_id:
         session_id = f"api_{uuid4()}"
-        logger.info(f"🆕 New conversation session: {session_id}")
-        yield f"data: {json.dumps({'type': 'session', 'session_id': session_id})}\n\n"
-    else:
-        logger.info(f"🔄 Continuing session: {session_id}")
 
-    try:
-        input_state = {
-            "messages": [{"role": "user", "content": message, "meta": {}}],
-            "legal_context": {
-                "jurisdiction": "Unknown", 
-                "user_type": "general", 
-                "document_type": "legal",
-                "detected_country": "unknown"
-            },
-            "session_id": session_id,
-            "router_decision": None,
-            "search_results": None,
-            "route_explanation": None,
-            "last_search_query": None,
-            "detected_articles": [],
-        }
+    input_state = {
+        "messages": [{"role": "user", "content": message, "meta": {}}],
+        "legal_context": {
+            "jurisdiction": "Unknown",
+            "user_type": "general",
+            "document_type": "legal",
+            "detected_country": "unknown"
+        },
+        "session_id": session_id,
+        "router_decision": None,
+        "search_results": None,
+        "route_explanation": None,
+        "last_search_query": None,
+        "detected_articles": [],
+    }
 
-        config = {
-            "configurable": {
-                "thread_id": session_id
-            }
+    config = {
+        "configurable": {
+            "thread_id": session_id
         }
+    }
 
-        events = graph.astream_events(
-            MultiCountryLegalState(**input_state),
-            version="v2",
-            config=config
-        )
+    events = graph.astream_events(
+        MultiCountryLegalState(**input_state),
+        version="v2",
+        config=config
+    )
 
-        current_content = ""
-        current_node = ""
+    current_content = ""
+    current_node = ""
+    final_state = None
 
-        async for event in events:
-            event_type = event["event"]
-            node_name = event.get("name", "")
+    async for event in events:
+        event_type = event["event"]
+        node_name = event.get("name", "")
+        
+        if node_name != current_node:
+            current_node = node_name
+            yield f"data: {json.dumps({'type': 'node_transition', 'node': node_name})}\n\n"
+
+        if event_type == "on_chat_model_stream":
+            chunk_content = serialize_ai_message_chunk(event["data"]["chunk"])
+            current_content += chunk_content
+            yield f"data: {json.dumps({'type': 'content', 'content': chunk_content})}\n\n"
+            
+        elif event_type == "on_chat_model_end":
+            yield f"data: {json.dumps({'type': 'content_end'})}\n\n"
+            
+        elif event_type == "on_chain_start" and "retrieval" in node_name:
+            country = node_name.replace("_retrieval", "")
+            yield f"data: {json.dumps({'type': 'search_start', 'country': country})}\n\n"
             
-            if node_name != current_node:
-                current_node = node_name
-                yield f"data: {json.dumps({'type': 'node_transition', 'node': node_name})}\n\n"
+        elif event_type == "on_chain_end" and "retrieval" in node_name:
+            country = node_name.replace("_retrieval", "")
+            yield f"data: {json.dumps({'type': 'search_end', 'country': country})}\n\n"
+            
+        elif event_type == "on_tool_end":
+            tool_name = event["name"]
+            yield f"data: {json.dumps({'type': 'tool_complete', 'tool': tool_name})}\n\n"
 
-            if event_type == "on_chat_model_stream":
-                chunk_content = serialize_ai_message_chunk(event["data"]["chunk"])
-                current_content += chunk_content
-                yield f"data: {json.dumps({'type': 'content', 'content': chunk_content})}\n\n"
-                
-            elif event_type == "on_chat_model_end":
-                yield f"data: {json.dumps({'type': 'content_end'})}\n\n"
-                
-            elif event_type == "on_chain_start" and "retrieval" in node_name:
-                country = node_name.replace("_retrieval", "")
-                yield f"data: {json.dumps({'type': 'search_start', 'country': country})}\n\n"
-                
-            elif event_type == "on_chain_end" and "retrieval" in node_name:
-                country = node_name.replace("_retrieval", "")
-                yield f"data: {json.dumps({'type': 'search_end', 'country': country})}\n\n"
-                
-            elif event_type == "on_tool_end":
-                tool_name = event["name"]
-                yield f"data: {json.dumps({'type': 'tool_complete', 'tool': tool_name})}\n\n"
+        elif event_type == "on_graph_end":
+            # Capture and convert the final state
+            state = event.get("data", {}).get("state")
+            if state and isinstance(state, MultiCountryLegalState):
+                final_state = state
+                state_dict = state.model_dump() if hasattr(state, "model_dump") else state.dict()
+                if "messages" in state_dict and isinstance(state_dict["messages"], list):
+                    state_dict["messages"] = [
+                        msg if isinstance(msg, dict) else {"role": "unknown", "content": str(msg), "meta": {}}
+                        for msg in state_dict["messages"]
+                    ]
+                yield f"data: {json.dumps({'type': 'state', 'content': state_dict})}\n\n"
+            yield f"data: {json.dumps({'type': 'graph_end'})}\n\n"
 
-            elif event_type == "on_graph_end":
-                yield f"data: {json.dumps({'type': 'graph_end'})}\n\n"
+    # Yield final state if captured
+    if final_state and isinstance(final_state, MultiCountryLegalState):
+        final_state_dict = final_state.model_dump() if hasattr(final_state, "model_dump") else final_state.dict()
+        if "messages" in final_state_dict and isinstance(final_state_dict["messages"], list):
+            final_state_dict["messages"] = [
+                msg if isinstance(msg, dict) else {"role": "unknown", "content": str(msg), "meta": {}}
+                for msg in final_state_dict["messages"]
+            ]
+        yield f"data: {json.dumps({'type': 'final_state', 'content': final_state_dict})}\n\n"
 
-    except Exception as e:
-        logger.error(f"Error in streaming: {e}")
-        yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"
-    
     yield f"data: {json.dumps({'type': 'end'})}\n\n"
 
+
 @app.get("/chat")
 async def chat_stream(
     message: str = Query(..., description="User message"),