add mcp server

src/api/v1/router.py

@@ -20,7 +20,8 @@ async def insurance_chat(request: models.InsuranceChatRequest):
             session_id=result["session_id"],
             answer=result["answer"],
             sources=result["sources"],
-            history=result["history"]
+            history=result["history"],
+            hooks=result["hooks"],
         )
         
     except Exception as e:

src/mcp/server.py

@@ -0,0 +1,37 @@
+from fastmcp import FastMCP
+
+from src.cache.redis_cache import fetch_plans, get_plan_by_id
+
+mcp = FastMCP(
+    name="LollyInsurancePlans",
+    instructions="""
+    You are a helpful assistant that can fetch insurance plans from the Lolly API.
+    You can fetch the plan id list by using get_plan_list() tool.
+    You can fetch the particular plan information by using get_plan_by_id() tool.
+    """
+)
+
+
+@mcp.tool()
+async def get_plan_by_id(plan_id: int) -> dict:
+    plan = await get_plan_by_id(plan_id)
+    return plan
+
+@mcp.tool()
+async def get_plan_list() -> list[dict]:
+    plans = await fetch_plans()
+    return [
+        {
+            "id": p["id"],
+            "name": p["name"],
+            "category": p["category"],
+        } for p in plans
+    ]
+
+if __name__ == "__main__":
+    mcp.run(
+        "http",
+        host="0.0.0.0",
+        port=8000,
+        path="/mcp",
+    )

src/models.py

@@ -4,16 +4,16 @@ from datetime import date
 
 class Applicant(BaseModel):
     applicant: int
-    firstName: str
-    lastName: str
-    midName: Optional[str] = Field("", description="Middle name")
-    phone: Optional[str] = Field("", description="Phone number")
-    gender: str
-    dob: date
-    nicotine: bool
-    weight: float
-    heightFt: int
-    heightIn: int
+    firstName: str | None = None
+    lastName: str | None = None
+    midName: str | None
+    phone: str | None
+    gender: str | None = None
+    dob: date | None = None
+    nicotine: bool | None = None
+    weight: float | None = None
+    heightFt: int | None = None
+    heightIn: int | None = None
 
 class Plan(BaseModel):
     id: int
@@ -99,11 +99,26 @@ class HistoryItem(BaseModel):
     role: str
     message: str
 
+class PlansParam(BaseModel):
+    plans: list[int]
+
+class ApplicantParam(BaseModel):
+    applicants: list[Applicant]
+
+class ChatHook(BaseModel):
+    tool: str
+    params: PlansParam | ApplicantParam
+
+class AIChatResponse(BaseModel):
+    answer: str
+    hooks: List[ChatHook] = []
+
 class InsuranceChatResponse(BaseModel):
     session_id: str
     answer: str
     sources: List[Source] = []
     history: List[HistoryItem] = []
+    hooks: List[ChatHook] = []
 
 class SessionCreateResponse(BaseModel):
     session_id: str

src/services/chat_service.py

@@ -1,3 +1,4 @@
+import json
 import httpx
 from typing import Dict, Any, List, Optional
 from ..config import settings
@@ -111,14 +112,17 @@ class ChatService:
             
             chat_response = await self.send_message(session_id, message)
             history = await self.get_chat_history(session_id)
-            
-            sources = self._extract_sources_from_response(chat_response.get("message", ""))
+
+            ai_response = json.loads(chat_response.get("message", {}))
+            ai_message = ai_response.get("message", "")
+            hooks = ai_response.get("hooks", [])
             
             return {
                 "session_id": session_id,
-                "answer": chat_response.get("message", "No response received"),
-                "sources": sources,
-                "history": history
+                "answer": ai_message,
+                "sources": [],
+                "history": history,
+                "hooks": hooks,
             }
             
         except Exception as e: