Revert "Update generate_legal_chat_responses for Pydantic v2 compatibility on Hugging Face"

This reverts commit e1e3b4d7993fd3a80c40a59045f280358afa691c.

api/main.py

@@ -256,43 +256,30 @@ 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 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"
+            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"
 
     # Yield final state if captured
     if final_state and isinstance(final_state, MultiCountryLegalState):
-        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"
+        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"
 
     yield f"data: {json.dumps({'type': 'end'})}\n\n"
 
+
 @app.get("/chat")
 async def chat_stream(
     message: str = Query(..., description="User message"),