Refactor roleplay agent implementation and update session handling for improved message processing

sessions.json

@@ -1,16 +0,0 @@
-[
-  {
-    "id": "82a6779d-ad13-4edd-a046-575e563a4348",
-    "name": "New Conversation",
-    "created_at": "2025-08-21T11:57:23.992279",
-    "last_message": "[Audio message]",
-    "message_count": 37
-  },
-  {
-    "id": "4fbf6c50-6054-4f3d-ac4e-d8281c306d72",
-    "name": "New Conversation",
-    "created_at": "2025-08-21T11:57:23.993885",
-    "last_message": null,
-    "message_count": 0
-  }
-]

src/agents/role_play/flow.py

@@ -1,24 +1,45 @@
 from langgraph.graph import StateGraph, START, END
-from .func import State
+from .func import State, trim_history, call_roleplay, call_guiding_agent
 from langgraph.graph.state import CompiledStateGraph
-from langgraph.store.memory import InMemoryStore
+from langgraph.checkpoint.memory import InMemorySaver
 
 
-class PrimaryChatBot:
+class RolePlayAgent:
     def __init__(self):
-        self.builder = StateGraph(State)
+        pass
 
     @staticmethod
-    def routing(state: State):
-        pass
+    def route_to_active_agent(state: State) -> str:
+        if state["active_agent"] == "Roleplay Agent":
+            return "Roleplay Agent"
+        elif state["active_agent"] == "Guiding Agent":
+            return "Guiding Agent"
 
-    def node(self):
-        pass
+    def node(self, graph: StateGraph):
+        graph.add_node("trim_history", trim_history)
+        graph.add_node("Roleplay Agent", call_roleplay, destinations=("Guiding Agent",))
+        graph.add_node(
+            "Guiding Agent", call_guiding_agent, destinations=("Roleplay Agent",)
+        )
+        return graph
+
+    def edge(self, graph: StateGraph):
+        graph.add_edge(START, "trim_history")
+        graph.add_conditional_edges(
+            "trim_history",
+            self.route_to_active_agent,
+            {
+                "Roleplay Agent": "Roleplay Agent",
+                "Guiding Agent": "Guiding Agent",
+            },
+        )
+        return graph
+
+    def __call__(self, checkpointer=InMemorySaver()) -> CompiledStateGraph:
+        graph = StateGraph(State)
+        graph: StateGraph = self.node(graph)
+        graph: StateGraph = self.edge(graph)
+        return graph.compile(checkpointer=checkpointer)
 
-    def edge(self):
-        pass
 
-    def __call__(self) -> CompiledStateGraph:
-        self.node()
-        self.edge()
-        return self.builder.compile(checkpointer=InMemoryStore())
+role_play_agent = RolePlayAgent()

src/agents/role_play/func.py

@@ -1,17 +1,40 @@
 from typing import TypedDict
 from src.config.llm import model
-from langgraph.checkpoint.memory import InMemorySaver
 from langgraph.prebuilt import create_react_agent
-from langgraph_swarm import create_handoff_tool, create_swarm
+from langgraph_swarm import create_handoff_tool
+from langchain_core.messages import RemoveMessage
 from .prompt import roleplay_prompt, guiding_prompt
+from typing_extensions import TypedDict, Annotated
+from langchain_core.messages import AnyMessage
+from langgraph.graph import add_messages
+from loguru import logger
 
 
 class State(TypedDict):
-    pass
+    active_agent: str | None
+    messages: Annotated[list[AnyMessage], add_messages]
+    scenario_title: str
+    scenario_description: str
+    scenario_context: str
+    your_role: str
+    key_vocabulary: str
 
 
-def create_agents(scenario, checkpointer=InMemorySaver()):
+def trim_history(state: State):
+    if not state.get("active_agent"):
+        state["active_agent"] = "Roleplay Agent"
+    history = state.get("messages", [])
+    if len(history) > 25:
+        num_to_remove = len(history) - 5
+        remove_messages = [
+            RemoveMessage(id=history[i].id) for i in range(num_to_remove)
+        ]
+        state["messages"] = remove_messages
+    return state
 
+
+async def call_roleplay(state: State):
+    logger.info("Calling roleplay agent...")
     roleplay_agent = create_react_agent(
         model,
         [
@@ -21,15 +44,21 @@ def create_agents(scenario, checkpointer=InMemorySaver()):
             ),
         ],
         prompt=roleplay_prompt.format(
-            scenario_title=scenario["scenario_title"],
-            scenario_description=scenario["scenario_description"],
-            scenario_context=scenario["scenario_context"],
-            your_role=scenario["your_role"],
-            key_vocabulary=scenario["key_vocabulary"],
+            scenario_title=state["scenario_title"],
+            scenario_description=state["scenario_description"],
+            scenario_context=state["scenario_context"],
+            your_role=state["your_role"],
+            key_vocabulary=state["key_vocabulary"],
         ),
         name="Roleplay Agent",
     )
+    response = await roleplay_agent.ainvoke({"messages": state["messages"]})
+
+    return {"messages": response["messages"]}
 
+
+async def call_guiding_agent(state: State):
+    logger.info("Calling guiding agent...")
     guiding_agent = create_react_agent(
         model,
         [
@@ -39,17 +68,20 @@ def create_agents(scenario, checkpointer=InMemorySaver()):
             ),
         ],
         prompt=guiding_prompt.format(
-            scenario_title=scenario["scenario_title"],
-            scenario_description=scenario["scenario_description"],
-            scenario_context=scenario["scenario_context"],
-            your_role=scenario["your_role"],
-            key_vocabulary=scenario["key_vocabulary"],
+            scenario_title=state["scenario_title"],
+            scenario_description=state["scenario_description"],
+            scenario_context=state["scenario_context"],
+            your_role=state["your_role"],
+            key_vocabulary=state["key_vocabulary"],
         ),
         name="Guiding Agent",
     )
+    response = await guiding_agent.ainvoke({"messages": state["messages"]})
+    return {"messages": response["messages"]}
 
-    workflow = create_swarm(
-        [roleplay_agent, guiding_agent], default_active_agent="Roleplay Agent"
-    )
 
-    return workflow.compile(checkpointer)
+def route_to_active_agent(state: State) -> str:
+    if state["active_agent"] == "Roleplay Agent":
+        return "Roleplay Agent"
+    elif state["active_agent"] == "Guiding Agent":
+        return "Guiding Agent"

src/apis/routes/chat_route.py

@@ -10,7 +10,7 @@ from fastapi import (
 )
 from fastapi.responses import JSONResponse
 from src.utils.logger import logger
-from src.agents.role_play.func import create_agents
+from src.agents.role_play.flow import role_play_agent
 from pydantic import BaseModel, Field
 from typing import List, Dict, Any, Optional
 from src.agents.role_play.scenarios import get_scenarios, get_scenario_by_id
@@ -200,21 +200,24 @@ async def roleplay(
     message = {"role": "user", "content": message_content}
 
     try:
-        response = await create_agents(scenario_dict).ainvoke(
+        response = await role_play_agent().ainvoke(
             {
                 "messages": [message],
+                "scenario_title": scenario_dict["scenario_title"],
+                "scenario_description": scenario_dict["scenario_description"],
+                "scenario_context": scenario_dict["scenario_context"],
+                "your_role": scenario_dict["your_role"],
+                "key_vocabulary": scenario_dict["key_vocabulary"],
             },
             {"configurable": {"thread_id": session_id}},
         )
-
-        # Update session with last message (use text if available, otherwise indicate audio)
         last_message = text_message if text_message else "[Audio message]"
         update_session_last_message(session_id, last_message)
-        
+
         # Extract AI response content
         ai_response = response["messages"][-1].content
         logger.info(f"AI response: {ai_response}")
-        
+
         return JSONResponse(content={"response": ai_response})
 
     except Exception as e:
@@ -228,7 +231,7 @@ async def get_messages(request: SessionRequest):
     try:
 
         # Create agent instance
-        agent = create_agents()
+        agent = role_play_agent()
 
         # Get current state
         current_state = agent.get_state(