fix from claude ai

api/main.py

@@ -189,11 +189,14 @@ 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) -> str:
     if not session_id:
         session_id = f"api_{uuid4()}"
 
-    input_state = {
+    # CRITICAL FIX: Create input state as a dictionary first, then convert to Pydantic model
+    # This ensures proper serialization for PostgreSQL checkpointing
+    input_state_dict = {
         "messages": [{"role": "user", "content": message, "meta": {}}],
         "legal_context": {
             "jurisdiction": "Unknown",
@@ -207,16 +210,39 @@ async def generate_legal_chat_responses(message: str, session_id: Optional[str]
         "route_explanation": None,
         "last_search_query": None,
         "detected_articles": [],
+        "supplemental_message": "",
+        "country": None,
+        "assistance_requested": False,
+        "user_email": None,
+        "assistance_description": None,
+        "email_status": None,
+        "assistance_step": None,
+        "pending_assistance_data": {},
+        "repair_type": None,
+        "original_query": None,
+        "misunderstanding_count": 0,
+        "primary_intent": None,
+        "approval_status": None,
+        "approval_reason": None,
+        "approved_by": None,
+        "approval_timestamp": None,
+        "summary_generated": False,
+        "last_summary_timestamp": None,
+        "search_metadata": {}
     }
 
+    # Convert to Pydantic model (this will use our custom model_dump for serialization)
+    input_state = MultiCountryLegalState(**input_state_dict)
+
     config = {
         "configurable": {
             "thread_id": session_id
         }
     }
 
+    # Stream events from the graph
     events = graph.astream_events(
-        MultiCountryLegalState(**input_state),
+        input_state,  # Pass the Pydantic model directly
         version="v2",
         config=config
     )
@@ -225,56 +251,52 @@ async def generate_legal_chat_responses(message: str, session_id: Optional[str]
     current_node = ""
     final_state = None
 
-    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"
-            
-        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"
+    try:
+        async for event in events:
+            event_type = event["event"]
+            node_name = event.get("name", "")
             
-        elif event_type == "on_tool_end":
-            tool_name = event["name"]
-            yield f"data: {json.dumps({'type': 'tool_complete', 'tool': tool_name})}\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_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"
+            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("output")
+                if state and isinstance(state, MultiCountryLegalState):
+                    final_state = state
+                    # Use our custom model_dump method for proper serialization
+                    state_dict = state.model_dump()
+                    yield f"data: {json.dumps({'type': 'state', 'content': state_dict})}\n\n"
+                yield f"data: {json.dumps({'type': 'graph_end'})}\n\n"
+
+    except Exception as e:
+        logger.error(f"Error in generate_legal_chat_responses: {e}", exc_info=True)
+        yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\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"]
-            ]
+        final_state_dict = final_state.model_dump()
         yield f"data: {json.dumps({'type': 'final_state', 'content': final_state_dict})}\n\n"
 
     yield f"data: {json.dumps({'type': 'end'})}\n\n"

models/state_models.py

@@ -1,7 +1,10 @@
-# [file name]: models/state_models.py
+# models/state_models.py
 from typing import List, Dict, Any, Optional, Annotated, Literal, Union
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, ConfigDict
+from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
 import operator
+import json
+
 
 class MultiCountryLegalState(BaseModel):
     messages: Annotated[List[Dict[str, Any]], operator.add] = Field(default_factory=list)
@@ -55,6 +58,114 @@ class MultiCountryLegalState(BaseModel):
     # NEW: Search-related fields to prevent storing complex data in legal_context
     search_metadata: Dict[str, Any] = Field(default_factory=dict)
 
+    # ============================================================================
+    # CRITICAL FIX FOR JSON SERIALIZATION (Pydantic v2 Configuration)
+    # This fixes: TypeError: Object of type MultiCountryLegalState is not JSON serializable
+    # ============================================================================
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,  # Allow LangChain message types if used
+        validate_assignment=True,
+        # CRITICAL: Tell Pydantic how to serialize this model to JSON
+        json_encoders={
+            # Any custom types can be added here
+        }
+    )
+    
+    def model_dump(self, **kwargs) -> Dict[str, Any]:
+        """
+        Override model_dump to ensure proper serialization for PostgreSQL checkpointing.
+        This fixes: TypeError: Object of type MultiCountryLegalState is not JSON serializable
+        """
+        data = super().model_dump(**kwargs)
+        
+        # Ensure all nested objects are JSON-serializable
+        # Messages should already be dicts, but double-check
+        if "messages" in data and data["messages"]:
+            serialized_messages = []
+            for msg in data["messages"]:
+                if isinstance(msg, dict):
+                    serialized_messages.append(msg)
+                elif isinstance(msg, BaseMessage):
+                    # Convert LangChain message objects to dicts
+                    serialized_messages.append({
+                        "role": "assistant" if isinstance(msg, AIMessage) else "user",
+                        "content": msg.content,
+                        "meta": getattr(msg, "additional_kwargs", {}),
+                    })
+                else:
+                    # Fallback for any other type
+                    serialized_messages.append({
+                        "role": "unknown",
+                        "content": str(msg),
+                        "meta": {}
+                    })
+            data["messages"] = serialized_messages
+        
+        # Ensure nested dicts are serializable
+        for key in ["legal_context", "pending_assistance_data", "search_metadata"]:
+            if key in data and data[key]:
+                # Convert any non-serializable objects to strings
+                data[key] = self._make_json_serializable(data[key])
+        
+        return data
+    
+    def model_dump_json(self, **kwargs) -> str:
+        """
+        Override model_dump_json for explicit JSON string conversion.
+        """
+        data = self.model_dump(**kwargs)
+        return json.dumps(data, default=str)
+    
+    @staticmethod
+    def _make_json_serializable(obj: Any) -> Any:
+        """
+        Recursively convert objects to JSON-serializable format.
+        """
+        if isinstance(obj, dict):
+            return {k: MultiCountryLegalState._make_json_serializable(v) for k, v in obj.items()}
+        elif isinstance(obj, list):
+            return [MultiCountryLegalState._make_json_serializable(item) for item in obj]
+        elif isinstance(obj, (str, int, float, bool, type(None))):
+            return obj
+        elif isinstance(obj, BaseMessage):
+            return {
+                "role": "assistant" if isinstance(obj, AIMessage) else "user",
+                "content": obj.content,
+                "meta": getattr(obj, "additional_kwargs", {}),
+            }
+        else:
+            # Convert any other type to string
+            return str(obj)
+    
+    @classmethod
+    def model_validate(cls, obj: Any) -> "MultiCountryLegalState":
+        """
+        Override model_validate to properly handle deserialization from checkpoints.
+        """
+        if isinstance(obj, dict):
+            # Messages should already be dicts, but handle BaseMessage objects if present
+            if "messages" in obj and obj["messages"]:
+                reconstructed_messages = []
+                for msg in obj["messages"]:
+                    if isinstance(msg, dict):
+                        reconstructed_messages.append(msg)
+                    elif isinstance(msg, BaseMessage):
+                        reconstructed_messages.append({
+                            "role": "assistant" if isinstance(msg, AIMessage) else "user",
+                            "content": msg.content,
+                            "meta": getattr(msg, "additional_kwargs", {}),
+                        })
+                    else:
+                        reconstructed_messages.append({
+                            "role": "unknown",
+                            "content": str(msg),
+                            "meta": {}
+                        })
+                obj["messages"] = reconstructed_messages
+        
+        return super().model_validate(obj)
+    # ============================================================================
+
     @staticmethod
     def detect_country(text: str) -> str:
         """
@@ -104,6 +215,7 @@ class RoutingResult(BaseModel):
     method: str
     explanation: str
 
+
 class SearchResult(BaseModel):
     documents: List[Any]
     detected_articles: List[str]