Update wellbe_agent.py
Browse files- wellbe_agent.py +26 -28
wellbe_agent.py
CHANGED
@@ -1,6 +1,5 @@
|
|
1 |
"""
|
2 |
Core logic for the WellBe+ multi‑context AI assistant.
|
3 |
-
|
4 |
This module handles MCP server instantiation, model selection, and single‑shot
|
5 |
question answering. It is intentionally agnostic of any UI layer so that it can
|
6 |
be re‑used from a command‑line script, a notebook, or the Gradio front‑end.
|
@@ -10,7 +9,7 @@ from __future__ import annotations
|
|
10 |
|
11 |
import asyncio
|
12 |
from copy import deepcopy
|
13 |
-
from typing import Tuple, Union
|
14 |
|
15 |
from agents import Agent, Runner
|
16 |
from agents.extensions.models.litellm_model import LitellmModel
|
@@ -19,33 +18,30 @@ from agents.mcp.server import MCPServerStdio, MCPServerStreamableHttp
|
|
19 |
from config import PROMPT_TEMPLATE, MCP_CONFIGS
|
20 |
|
21 |
__all__ = [
|
22 |
-
"
|
23 |
"select_model",
|
24 |
"answer_question",
|
|
|
25 |
]
|
26 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
27 |
|
28 |
-
def build_mcp_servers(
|
29 |
-
whoop_email: str, whoop_password: str
|
30 |
-
) -> Tuple[MCPServerStreamableHttp, MCPServerStdio]:
|
31 |
-
"""Return configured Healthcare and Whoop MCP servers.
|
32 |
-
|
33 |
-
Parameters
|
34 |
-
----------
|
35 |
-
whoop_email, whoop_password
|
36 |
-
User credentials for the private Whoop feed. The credentials are merged
|
37 |
-
into a *deep‑copied* variant of ``MCP_CONFIGS`` so that the global
|
38 |
-
settings remain untouched.
|
39 |
-
"""
|
40 |
-
cfg = deepcopy(MCP_CONFIGS)
|
41 |
-
cfg["whoop"].update({"username": whoop_email, "password": whoop_password})
|
42 |
-
|
43 |
-
healthcare_server = MCPServerStreamableHttp(
|
44 |
-
params=cfg["healthcare-mcp-public"], name="Healthcare MCP Server"
|
45 |
-
)
|
46 |
-
# whoop_server = MCPServerStdio(params=cfg["whoop"], name="Whoop MCP Server")
|
47 |
-
whoop_server = "Whoop Tool Not available now"
|
48 |
-
return healthcare_server, whoop_server
|
49 |
|
50 |
def select_model() -> Union[str, LitellmModel]:
|
51 |
return "o3-mini"
|
@@ -58,16 +54,18 @@ async def answer_question(
|
|
58 |
whoop_password: str,
|
59 |
) -> str:
|
60 |
"""Run the WellBe+ agent on a single question and return the assistant reply."""
|
61 |
-
|
|
|
|
|
62 |
|
63 |
-
async with
|
64 |
agent = Agent(
|
65 |
name="WellBe+ Assistant",
|
66 |
instructions=PROMPT_TEMPLATE,
|
67 |
model=select_model(),
|
68 |
mcp_servers=[
|
69 |
-
hserver
|
70 |
-
|
71 |
],
|
72 |
)
|
73 |
result = await Runner.run(agent, question)
|
|
|
1 |
"""
|
2 |
Core logic for the WellBe+ multi‑context AI assistant.
|
|
|
3 |
This module handles MCP server instantiation, model selection, and single‑shot
|
4 |
question answering. It is intentionally agnostic of any UI layer so that it can
|
5 |
be re‑used from a command‑line script, a notebook, or the Gradio front‑end.
|
|
|
9 |
|
10 |
import asyncio
|
11 |
from copy import deepcopy
|
12 |
+
from typing import Tuple, Union, Optional
|
13 |
|
14 |
from agents import Agent, Runner
|
15 |
from agents.extensions.models.litellm_model import LitellmModel
|
|
|
18 |
from config import PROMPT_TEMPLATE, MCP_CONFIGS
|
19 |
|
20 |
__all__ = [
|
21 |
+
"initialize_mcp_servers",
|
22 |
"select_model",
|
23 |
"answer_question",
|
24 |
+
"answer_sync",
|
25 |
]
|
26 |
|
27 |
+
# Global variables to store initialized servers
|
28 |
+
_healthcare_server: Optional[MCPServerStreamableHttp] = None
|
29 |
+
_whoop_server: Optional[Union[MCPServerStdio, str]] = None
|
30 |
+
|
31 |
+
|
32 |
+
def initialize_mcp_servers(whoop_email: str, whoop_password: str) -> None:
|
33 |
+
"""Initialize and cache the MCP servers (once only)."""
|
34 |
+
global _healthcare_server, _whoop_server
|
35 |
+
if _healthcare_server is None:
|
36 |
+
cfg = deepcopy(MCP_CONFIGS)
|
37 |
+
cfg["whoop"].update({"username": whoop_email, "password": whoop_password})
|
38 |
+
|
39 |
+
_healthcare_server = MCPServerStreamableHttp(
|
40 |
+
params=cfg["healthcare-mcp-public"], name="Healthcare MCP Server"
|
41 |
+
)
|
42 |
+
# _whoop_server = MCPServerStdio(params=cfg["whoop"], name="Whoop MCP Server")
|
43 |
+
_whoop_server = "Whoop Tool Not available now"
|
44 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
45 |
|
46 |
def select_model() -> Union[str, LitellmModel]:
|
47 |
return "o3-mini"
|
|
|
54 |
whoop_password: str,
|
55 |
) -> str:
|
56 |
"""Run the WellBe+ agent on a single question and return the assistant reply."""
|
57 |
+
initialize_mcp_servers(whoop_email, whoop_password)
|
58 |
+
if _healthcare_server is None:
|
59 |
+
return "**Error:** MCP servers not initialized."
|
60 |
|
61 |
+
async with _healthcare_server as hserver: # , _whoop_server as wserver:
|
62 |
agent = Agent(
|
63 |
name="WellBe+ Assistant",
|
64 |
instructions=PROMPT_TEMPLATE,
|
65 |
model=select_model(),
|
66 |
mcp_servers=[
|
67 |
+
hserver
|
68 |
+
# , wserver
|
69 |
],
|
70 |
)
|
71 |
result = await Runner.run(agent, question)
|