Skip to content

Commit 5a5f456

Browse files
author
zach
authored
refactor: use MCP protocol to access mcp.run via SSE/stdio (#27)
* refactor: use mcp client
1 parent fb0e93b commit 5a5f456

File tree

6 files changed

+388
-519
lines changed

6 files changed

+388
-519
lines changed

‎mcpx_py/__main__.py‎

Lines changed: 16 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -50,10 +50,11 @@ async def tool_cmd(client, args):
5050

5151

5252
EXIT_COUNT = 0
53+
CACHED_TOOLS = []
5354

5455

5556
async def chat_loop(chat):
56-
global EXIT_COUNT
57+
global EXIT_COUNT, CACHED_TOOLS
5758
try:
5859
msg = input("> ").strip()
5960
EXIT_COUNT = 0
@@ -77,8 +78,15 @@ async def chat_loop(chat):
7778
print("Chat history cleared")
7879
return True
7980
elif msg == "!tools":
81+
if len(CACHED_TOOLS) > 0:
82+
tools = CACHED_TOOLS
83+
else:
84+
tools = []
85+
async for tool in chat.list_tools():
86+
tools.append(tool)
87+
CACHED_TOOLS = tools
8088
print("\nAvailable tools:")
81-
for tool in chat.agent._function_tools.values():
89+
for tool in tools:
8290
if tool is None:
8391
continue
8492
print(f"- {tool.name.strip()}")
@@ -93,27 +101,25 @@ async def chat_loop(chat):
93101
if msg == "":
94102
return True
95103
async for res in chat.iter_content(msg):
96-
if not isinstance(res, pydantic_ai.models.ModelResponse):
104+
if not hasattr(res, "parts"):
97105
continue
106+
98107
for part in res.parts:
99108
if isinstance(part, pydantic_ai.messages.TextPart):
100109
print(part.content)
101110
elif isinstance(part, pydantic_ai.messages.ToolCallPart):
102111
args = part.args
103112
if isinstance(args, str):
104113
args = json.loads(args)
105-
if part.tool_name == "final_result":
106-
print(args["response"])
107-
else:
108-
print(
109-
f">> Tool: {part.tool_name} ({part.tool_call_id}) input={args}"
110-
)
114+
print(
115+
f">> Tool: {part.tool_name} ({part.tool_call_id}) input={args}"
116+
)
111117
except Exception:
112118
print("\nERROR>>", traceback.format_exc())
113119
return True
114120

115121

116-
async def chat_cmd(client, args):
122+
async def chat_cmd(client: Client, args):
117123
m = args.model
118124
if args.provider:
119125
if args.provider == "ollama" or args.provider == "llama":

‎mcpx_py/builtin_tools.py‎

Lines changed: 0 additions & 47 deletions
This file was deleted.

‎mcpx_py/chat.py‎

Lines changed: 103 additions & 107 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,7 @@
11
from mcpx_pydantic_ai import Agent, pydantic_ai
22

33

4-
from typing import TypedDict
5-
6-
from . import builtin_tools
4+
import asyncio
75

86

97
SYSTEM_PROMPT = """
@@ -32,7 +30,6 @@ class Chat:
3230
def __init__(
3331
self,
3432
*args,
35-
ignore_builtin_tools: bool = False,
3633
**kw,
3734
):
3835
if "system_prompt" not in kw:
@@ -42,14 +39,8 @@ def __init__(
4239
*args,
4340
**kw,
4441
)
45-
if not ignore_builtin_tools:
46-
self._register_builtins()
4742
self.history = []
4843

49-
def _register_builtins(self):
50-
for tool in builtin_tools.TOOLS:
51-
self.agent.register_tool(tool, getattr(self, "_tool_" + tool.name))
52-
5344
@property
5445
def client(self):
5546
"""
@@ -63,121 +54,126 @@ def clear_history(self):
6354
"""
6455
self.history = []
6556

66-
async def send_message(self, msg: str, *args, **kw):
57+
async def send_message(self, msg: str, run_mcp_servers: bool = True, *args, **kw):
6758
"""
6859
Send a chat message to the LLM
6960
"""
70-
with pydantic_ai.capture_run_messages() as messages:
71-
res = await self.agent.run(
72-
msg,
73-
message_history=self.history,
74-
*args,
75-
**kw,
76-
)
61+
62+
async def inner():
63+
with pydantic_ai.capture_run_messages() as messages:
64+
return (
65+
await self.agent.run(
66+
msg,
67+
message_history=self.history,
68+
*args,
69+
**kw,
70+
),
71+
messages,
72+
)
73+
74+
if run_mcp_servers:
75+
async with self.agent.run_mcp_servers():
76+
res, messages = await inner()
77+
else:
78+
res, messages = await inner()
7779
self.history.extend(messages)
7880
return res
7981

80-
def send_message_sync(self, msg, *args, **kw):
81-
"""
82-
Send a chat message to the LLM
82+
def send_message_sync(self, *args, **kw):
8383
"""
84-
with pydantic_ai.capture_run_messages() as messages:
85-
res = self.agent.run_sync(
86-
msg,
87-
message_history=self.history,
88-
*args,
89-
**kw,
90-
)
91-
self.history.extend(messages)
92-
return res
84+
Send a chat message to the LLM synchronously
9385
94-
async def iter(self, msg, *args, **kw):
86+
This creates a new event loop to run the async send_message method.
87+
"""
88+
# Create a new event loop to avoid warnings about coroutines not being awaited
89+
loop = asyncio.new_event_loop()
90+
try:
91+
return loop.run_until_complete(self.send_message(*args, **kw))
92+
finally:
93+
loop.close()
94+
95+
async def iter(self, msg: str, run_mcp_servers: bool = True, *args, **kw):
9596
"""
9697
Send a chat message to the LLM
9798
"""
98-
with pydantic_ai.capture_run_messages() as messages:
99-
async with self.agent.iter(
100-
msg, message_history=self.history, *args, **kw
101-
) as run:
102-
async for node in run:
103-
yield node
104-
self.history.extend(messages)
10599

106-
async def iter_content(self, msg, *args, **kw):
100+
async def inner():
101+
with pydantic_ai.capture_run_messages() as messages:
102+
async with self.agent.iter(
103+
msg, message_history=self.history, *args, **kw
104+
) as run:
105+
async for node in run:
106+
yield node
107+
self.history.extend(messages)
108+
109+
if run_mcp_servers:
110+
async with self.agent.run_mcp_servers():
111+
async for msg in inner():
112+
yield msg
113+
else:
114+
async for msg in inner():
115+
yield msg
116+
117+
async def iter_content(self, msg: str, run_mcp_servers: bool = True, *args, **kw):
107118
"""
108119
Send a chat message to the LLM
109120
"""
110-
with pydantic_ai.capture_run_messages() as messages:
111-
async with self.agent.iter(
112-
msg, message_history=self.history, *args, **kw
113-
) as run:
114-
async for node in run:
115-
if hasattr(node, "response"):
116-
content = node.response
117-
elif hasattr(node, "model_response"):
118-
content = node.model_response
119-
elif hasattr(node, "request"):
120-
content = node.request
121-
elif hasattr(node, "model_request"):
122-
content = node.model_request
123-
elif hasattr(node, "data"):
124-
content = node.data
125-
else:
126-
continue
127-
yield content
128-
self.history.extend(messages)
129121

130-
async def inspect(self, msg, *args, **kw):
122+
async def inner():
123+
with pydantic_ai.capture_run_messages() as messages:
124+
async with self.agent.iter(
125+
msg, message_history=self.history, *args, **kw
126+
) as run:
127+
async for node in run:
128+
if hasattr(node, "response"):
129+
content = node.response
130+
elif hasattr(node, "model_response"):
131+
content = node.model_response
132+
elif hasattr(node, "request"):
133+
content = node.request
134+
elif hasattr(node, "model_request"):
135+
content = node.model_request
136+
elif hasattr(node, "data"):
137+
content = node.data
138+
elif hasattr(node, "output"):
139+
content = node
140+
else:
141+
continue
142+
yield content
143+
self.history.extend(messages)
144+
145+
f = inner()
146+
if run_mcp_servers:
147+
async with self.agent.run_mcp_servers():
148+
async for msg in f:
149+
yield msg
150+
else:
151+
async for msg in f:
152+
yield msg
153+
154+
async def inspect(self, msg: str, run_mcp_servers: bool = True, *args, **kw):
131155
"""
132156
Send a chat message to the LLM
133157
"""
134-
with pydantic_ai.capture_run_messages() as messages:
135-
res = await self.send_message(msg, *args, **kw)
136-
return res, messages
137158

138-
def _tool_mcp_run_search_servlets(
139-
self, input: TypedDict("SearchServlets", {"q": str})
140-
):
141-
q = input.get("q", "")
142-
if q == "":
143-
return "ERROR: provide a query when searching"
144-
x = []
145-
for r in self.agent.client.search(input["q"]):
146-
x.append(
147-
{
148-
"slug": r.slug,
149-
"schema": {
150-
"name": r.meta.get("name"),
151-
"description": r.meta.get("description"),
152-
},
153-
"installation_count": r.installation_count,
154-
}
155-
)
156-
return x
157-
158-
def _tool_mcp_run_get_profiles(self, input: TypedDict("GetProfile", {})):
159-
p = []
160-
for user, u in self.agent.client.profiles.items():
161-
if user == "~":
162-
continue
163-
for profile in u.values():
164-
p.append(
165-
{
166-
"name": f"{user}/{profile.slug}", # Assume slug is string
167-
"description": profile.description,
168-
}
159+
async def inner():
160+
with pydantic_ai.capture_run_messages() as messages:
161+
res = await self.agent.run(
162+
msg,
163+
message_history=self.history,
164+
*args,
165+
**kw,
169166
)
170-
return p
171-
172-
def _tool_mcp_run_set_profile(
173-
self, input: TypedDict("SetProfile", {"profile": str})
174-
):
175-
profile = input["profile"]
176-
if "/" not in profile:
177-
profile = "~/" + profile
178-
self.agent.set_profile(profile)
179-
return f"Active profile set to {profile}"
180-
181-
def _tool_mcp_run_current_profile(self, input: TypedDict("CurrentProfile", {})):
182-
"""Get current profile name"""
183-
return self.agent.client.config.profile
167+
return res, messages
168+
169+
if run_mcp_servers:
170+
async with self.agent.run_mcp_servers():
171+
return await inner()
172+
else:
173+
return await inner()
174+
175+
async def list_tools(self):
176+
async with self.client.mcp_sse(self.agent.client.profile).connect() as session:
177+
tools = await session.list_tools()
178+
for tool in tools.tools:
179+
yield tool

‎pyproject.toml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
11
[project]
22
name = "mcpx-py"
3-
version = "0.6.0"
3+
version = "0.7.0"
44
description = "An mcp.run client for Python"
55
readme = "README.md"
66
requires-python = ">=3.12"
77
dependencies = [
8-
"mcpx-pydantic-ai>=0.5.0",
8+
"mcpx-pydantic-ai>=0.7.0",
99
"psutil>=7.0.0",
1010
"python-dotenv>=1.0.1",
1111
]

0 commit comments

Comments
 (0)