Refactor chat endpoint: extract agent logic to run_agent function and update response handling

agent_api/main.py

@@ -1,32 +1,44 @@
 from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
-
 from .core.llm import client
 
 app = FastAPI()
 
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
-# ---- request model ----
 class ChatRequest(BaseModel):
     message: str
+    session_id: str | None = None
+
+
+def run_agent(message: str, session_id: str | None = None):
+    response = client.chat.completions.create(
+        model="deepseek-chat",
+        messages=[
+            {"role": "system", "content": "You are a helpful agent."},
+            {"role": "user", "content": message}
+        ]
+    )
+    return response.choices[0].message.content
 
 
-# ---- basic health check ----
 @app.get("/")
 def root():
     return {"status": "ok"}
 
 
-# ---- test LLM endpoint ----
 @app.post("/chat")
 def chat(req: ChatRequest):
-    response = client.chat.completions.create(
-        model="deepseek-chat",
-        messages=[
-            {"role": "user", "content": req.message}
-        ]
-    )
+    response = run_agent(req.message, req.session_id)
 
     return {
-        "response": response.choices[0].message.content
+        "response": response,
+        "session_id": req.session_id
     }