Update generate_legal_chat_responses for Pydantic v2 compatibility on Hugging Face

api/main.py

@@ -256,30 +256,43 @@ async def generate_legal_chat_responses(message: str, session_id: Optional[str]
         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"
+            if state is not None:
+                try:
+                    if isinstance(state, MultiCountryLegalState):
+                        final_state = state
+                        # Use model_dump_json for Pydantic v2 compatibility
+                        state_dict = json.loads(state.model_dump_json(exclude_unset=True)) if hasattr(state, "model_dump_json") else state.dict(exclude_unset=True)
+                        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"
+                    else:
+                        logger.warning(f"Unexpected state type: {type(state)}")
+                except Exception as e:
+                    logger.error(f"Error converting state to dict: {str(e)}")
+                    state_dict = {"error": "Failed to serialize state"}
+                    yield f"data: {json.dumps({'type': 'state', 'content': state_dict})}\n\n"
             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"
+        try:
+            # Use model_dump_json for Pydantic v2 compatibility
+            final_state_dict = json.loads(final_state.model_dump_json(exclude_unset=True)) if hasattr(final_state, "model_dump_json") else final_state.dict(exclude_unset=True)
+            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 converting final state to dict: {str(e)}")
+            yield f"data: {json.dumps({'type': 'error', 'message': 'An internal error occurred'})}\n\n"
 
     yield f"data: {json.dumps({'type': 'end'})}\n\n"
 
-
 @app.get("/chat")
 async def chat_stream(
     message: str = Query(..., description="User message"),