victoria-latynina commited on
Commit
07514fa
·
verified ·
1 Parent(s): 3fee918

Update wellbe_agent.py

Browse files
Files changed (1) hide show
  1. 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
- "build_mcp_servers",
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
- healthcare_srv, whoop_srv = build_mcp_servers(whoop_email, whoop_password)
 
 
62
 
63
- async with healthcare_srv as hserver: #, whoop_srv as wserver:
64
  agent = Agent(
65
  name="WellBe+ Assistant",
66
  instructions=PROMPT_TEMPLATE,
67
  model=select_model(),
68
  mcp_servers=[
69
- hserver
70
- #, wserver
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)