Add trim_history function and update lesson chat route to handle text and audio inputs

src/agents/lesson_practice/flow.py

@@ -1,5 +1,5 @@
 from langgraph.graph import StateGraph, START, END
-from .func import State, agent, tool_node
+from .func import State, trim_history, agent, tool_node
 from langgraph.graph.state import CompiledStateGraph
 from langgraph.checkpoint.memory import InMemorySaver
 
@@ -18,12 +18,14 @@ class LessonPracticeAgent:
             return "continue"
 
     def node(self, graph: StateGraph):
+        graph.add_node("trim_history", trim_history)
         graph.add_node("agent", agent)
         graph.add_node("tools", tool_node)
         return graph
 
     def edge(self, graph: StateGraph):
-        graph.add_edge(START, "agent")
+        graph.add_edge(START, "trim_history")
+        graph.add_edge("trim_history", "agent")
         graph.add_conditional_edges(
             "agent", self.should_continue, {"end": END, "continue": "tools"}
         )

src/agents/lesson_practice/func.py

@@ -3,7 +3,7 @@ from typing import (
     Sequence,
     TypedDict,
 )
-from langchain_core.messages import ToolMessage, AnyMessage
+from langchain_core.messages import ToolMessage, AnyMessage, RemoveMessage
 from langgraph.graph.message import add_messages
 import json
 from .prompt import conversation_prompt
@@ -26,6 +26,19 @@ tools = []
 tools_by_name = {tool.name: tool for tool in tools}
 
 
+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
+
+
 # Define our tool node
 def tool_node(state: State):
     outputs = []

src/apis/routes/lesson_route.py

@@ -1,4 +1,13 @@
-from fastapi import APIRouter, status, Depends, BackgroundTasks, HTTPException
+from fastapi import (
+    APIRouter,
+    status,
+    Depends,
+    BackgroundTasks,
+    HTTPException,
+    File,
+    UploadFile,
+    Form,
+)
 from fastapi.responses import JSONResponse
 from src.utils.logger import logger
 from pydantic import BaseModel, Field
@@ -9,6 +18,7 @@ import json
 import os
 import uuid
 from datetime import datetime
+import base64
 
 router = APIRouter(prefix="/lesson", tags=["AI"])
 
@@ -144,18 +154,101 @@ async def search_lessons_by_unit(unit_name: str):
 
 
 @router.post("/chat")
-async def chat(request: LessonPracticeRequest):
-    response = await lesson_practice_agent().ainvoke(
-        {
-            "unit": request.unit,
-            "vocabulary": request.vocabulary,
-            "key_structures": request.key_structures,
-            "practice_questions": request.practice_questions,
-            "student_level": request.student_level,
-            "messages": [request.query],
-        },
-        {"configurable": {"thread_id": request.session_id}},
-    )
-    return JSONResponse(
-        content=response["messages"][-1].content, status_code=status.HTTP_200_OK
-    )
+async def chat(
+    session_id: str = Form(
+        ..., description="Session ID for tracking user interactions"
+    ),
+    lesson_data: str = Form(
+        ..., description="The lesson data as JSON string"
+    ),
+    text_message: Optional[str] = Form(None, description="Text message from user"),
+    audio_file: Optional[UploadFile] = File(None, description="Audio file from user"),
+):
+    """Send a message (text or audio) to the lesson practice agent"""
+
+    # Validate that at least one input is provided
+    if not text_message and not audio_file:
+        raise HTTPException(
+            status_code=400, detail="Either text_message or audio_file must be provided"
+        )
+
+    # Parse lesson data from JSON string
+    try:
+        lesson_dict = json.loads(lesson_data)
+    except json.JSONDecodeError:
+        raise HTTPException(status_code=400, detail="Invalid lesson_data JSON format")
+
+    if not lesson_dict:
+        raise HTTPException(status_code=400, detail="Lesson data not provided")
+
+    # Prepare message content
+    message_content = []
+
+    # Handle text input
+    if text_message:
+        message_content.append({"type": "text", "text": text_message})
+
+    # Handle audio input
+    if audio_file:
+        try:
+            # Read audio file content
+            audio_data = await audio_file.read()
+
+            # Convert to base64
+            audio_base64 = base64.b64encode(audio_data).decode("utf-8")
+
+            # Determine mime type based on file extension
+            file_extension = (
+                audio_file.filename.split(".")[-1].lower()
+                if audio_file.filename
+                else "wav"
+            )
+            mime_type_map = {
+                "wav": "audio/wav",
+                "mp3": "audio/mpeg",
+                "ogg": "audio/ogg",
+                "webm": "audio/webm",
+                "m4a": "audio/mp4",
+            }
+            mime_type = mime_type_map.get(file_extension, "audio/wav")
+
+            message_content.append(
+                {
+                    "type": "audio",
+                    "source_type": "base64",
+                    "data": audio_base64,
+                    "mime_type": mime_type,
+                }
+            )
+
+        except Exception as e:
+            logger.error(f"Error processing audio file: {str(e)}")
+            raise HTTPException(
+                status_code=400, detail=f"Error processing audio file: {str(e)}"
+            )
+
+    # Create message in the required format
+    message = {"role": "user", "content": message_content}
+
+    try:
+        response = await lesson_practice_agent().ainvoke(
+            {
+                "messages": [message],
+                "unit": lesson_dict.get("unit", ""),
+                "vocabulary": lesson_dict.get("vocabulary", []),
+                "key_structures": lesson_dict.get("key_structures", []),
+                "practice_questions": lesson_dict.get("practice_questions", []),
+                "student_level": lesson_dict.get("student_level", "beginner"),
+            },
+            {"configurable": {"thread_id": session_id}},
+        )
+        
+        # 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:
+        logger.error(f"Error in lesson practice: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")