feat: support function factory (#4)

backend/src/agents/__init__.py

@@ -1,4 +1,4 @@
-from .lead_agent import lead_agent
+from .lead_agent import make_lead_agent
 from .thread_state import SandboxState, ThreadState
 
-__all__ = ["lead_agent", "SandboxState", "ThreadState"]
+__all__ = ["make_lead_agent", "SandboxState", "ThreadState"]

backend/src/agents/lead_agent/__init__.py

@@ -1,3 +1,3 @@
-from .agent import lead_agent
+from .agent import make_lead_agent
 
-__all__ = ["lead_agent"]
+__all__ = ["make_lead_agent"]

backend/src/agents/lead_agent/agent.py

@@ -1,4 +1,5 @@
 from langchain.agents import create_agent
+from langchain_core.runnables import RunnableConfig
 
 from src.agents.lead_agent.prompt import apply_prompt_template
 from src.agents.middlewares.thread_data_middleware import ThreadDataMiddleware
@@ -11,10 +12,15 @@ from src.tools import get_available_tools
 # ThreadDataMiddleware must be before SandboxMiddleware to ensure thread_id is available
 middlewares = [ThreadDataMiddleware(), SandboxMiddleware(), TitleMiddleware()]
 
-lead_agent = create_agent(
-    model=create_chat_model(thinking_enabled=True),
-    tools=get_available_tools(),
-    middleware=middlewares,
-    system_prompt=apply_prompt_template(),
-    state_schema=ThreadState,
-)
+
+def make_lead_agent(config: RunnableConfig):
+    thinking_enabled = config.get("configurable", {}).get("thinking_enabled", True)
+    model_name = config.get("configurable", {}).get("model_name")
+    print(f"thinking_enabled: {thinking_enabled}, model_name: {model_name}")
+    return create_agent(
+        model=create_chat_model(name=model_name, thinking_enabled=thinking_enabled),
+        tools=get_available_tools(),
+        middleware=middlewares,
+        system_prompt=apply_prompt_template(),
+        state_schema=ThreadState,
+    )

backend/src/agents/lead_agent/prompt.py

@@ -1,4 +1,3 @@
-import os
 from datetime import datetime
 
 SYSTEM_PROMPT = f"""

backend/src/agents/middlewares/thread_data_middleware.py

@@ -71,8 +71,9 @@ class ThreadDataMiddleware(AgentMiddleware[ThreadDataMiddlewareState]):
     @override
     def before_agent(self, state: ThreadDataMiddlewareState, runtime: Runtime) -> dict | None:
         # Generate new thread ID and create directories
-        print(runtime.context)
-        thread_id = runtime.context["thread_id"]
+        thread_id = runtime.context.get("thread_id")
+        if thread_id is None:
+            raise ValueError("Thread ID is required in the context")
         paths = self._create_thread_directories(thread_id)
         print(f"Created thread data directories for thread {thread_id}")