Skip to content

refactor(client): simplify context management and improve code structure #8

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

refactor(client): simplify context management and improve code structure #8

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 6 additions & 10 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -482,7 +482,6 @@ import json
import os
from typing import Optional
from contextlib import AsyncExitStack
import time
from mcp import ClientSession
from mcp.client.sse import sse_client

Expand All @@ -501,11 +500,8 @@ class MCPClient:
async def connect_to_sse_server(self, server_url: str):
"""Connect to an MCP server running with SSE transport"""
# Store the context managers so they stay alive
self._streams_context = sse_client(url=server_url)
streams = await self._streams_context.__aenter__()

self._session_context = ClientSession(*streams)
self.session: ClientSession = await self._session_context.__aenter__()
streams = await self.exit_stack.enter_async_context(sse_client(url=server_url))
self.session = await self.exit_stack.enter_async_context(ClientSession(*streams))

# Initialize
await self.session.initialize()
Expand All @@ -519,10 +515,7 @@ class MCPClient:

async def cleanup(self):
"""Properly clean up the session and streams"""
if self._session_context:
await self._session_context.__aexit__(None, None, None)
if self._streams_context:
await self._streams_context.__aexit__(None, None, None)
await self.exit_stack.aclose()

async def process_query(self, query: str) -> str:
"""Process a query using OpenAI API and available tools"""
Expand Down Expand Up @@ -600,6 +593,7 @@ class MCPClient:
final_text.append(assistant_message.content)

return "\n".join(final_text)


async def chat_loop(self):
"""Run an interactive chat loop"""
Expand All @@ -619,6 +613,7 @@ class MCPClient:
except Exception as e:
print(f"\nError: {str(e)}")


async def main():
if len(sys.argv) < 2:
print("Usage: uv run client.py <URL of SSE MCP server (i.e. http://localhost:8080/sse)>")
Expand All @@ -631,6 +626,7 @@ async def main():
finally:
await client.cleanup()


if __name__ == "__main__":
import sys
asyncio.run(main())
Expand Down
25 changes: 12 additions & 13 deletions README_EN.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -462,7 +462,6 @@ import json
import os
from typing import Optional
from contextlib import AsyncExitStack
import time
from mcp import ClientSession
from mcp.client.sse import sse_client

Expand All @@ -481,11 +480,8 @@ class MCPClient:
async def connect_to_sse_server(self, server_url: str):
"""Connect to an MCP server running with SSE transport"""
# Store the context managers so they stay alive
self._streams_context = sse_client(url=server_url)
streams = await self._streams_context.__aenter__()

self._session_context = ClientSession(*streams)
self.session: ClientSession = await self._session_context.__aenter__()
streams = await self.exit_stack.enter_async_context(sse_client(url=server_url))
self.session = await self.exit_stack.enter_async_context(ClientSession(*streams))

# Initialize
await self.session.initialize()
Expand All @@ -499,10 +495,7 @@ class MCPClient:

async def cleanup(self):
"""Properly clean up the session and streams"""
if self._session_context:
await self._session_context.__aexit__(None, None, None)
if self._streams_context:
await self._streams_context.__aexit__(None, None, None)
await self.exit_stack.aclose()

async def process_query(self, query: str) -> str:
"""Process a query using OpenAI API and available tools"""
Expand Down Expand Up @@ -542,11 +535,14 @@ class MCPClient:
tool_name = tool_call.function.name
tool_args = json.loads(tool_call.function.arguments)



# Execute tool call
result = await self.session.call_tool(tool_name, tool_args)
tool_results.append({"call": tool_name, "result": result})
final_text.append(f"[Calling tool {tool_name} with args {tool_args}]")



# Continue conversation with tool results
messages.extend([
{
Expand All @@ -561,8 +557,8 @@ class MCPClient:
}
])

print(f"Tool {tool_name} returned: {result.content[0].text}")
print("messages", messages)
# print(f"Tool {tool_name} returned: {result.content[0].text}")
# print("messages", messages)
# Get next response from OpenAI
completion = await self.openai.chat.completions.create(
model=os.getenv("OPENAI_MODEL"),
Expand All @@ -580,6 +576,7 @@ class MCPClient:
final_text.append(assistant_message.content)

return "\n".join(final_text)


async def chat_loop(self):
"""Run an interactive chat loop"""
Expand All @@ -599,6 +596,7 @@ class MCPClient:
except Exception as e:
print(f"\nError: {str(e)}")


async def main():
if len(sys.argv) < 2:
print("Usage: uv run client.py <URL of SSE MCP server (i.e. http://localhost:8080/sse)>")
Expand All @@ -611,6 +609,7 @@ async def main():
finally:
await client.cleanup()


if __name__ == "__main__":
import sys
asyncio.run(main())
Expand Down
17 changes: 5 additions & 12 deletions client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
import os
from typing import Optional
from contextlib import AsyncExitStack
import time
from mcp import ClientSession
from mcp.client.sse import sse_client

Expand All @@ -22,11 +21,8 @@ def __init__(self):
async def connect_to_sse_server(self, server_url: str):
"""Connect to an MCP server running with SSE transport"""
# Store the context managers so they stay alive
self._streams_context = sse_client(url=server_url)
streams = await self._streams_context.__aenter__()

self._session_context = ClientSession(*streams)
self.session: ClientSession = await self._session_context.__aenter__()
streams = await self.exit_stack.enter_async_context(sse_client(url=server_url))
self.session = await self.exit_stack.enter_async_context(ClientSession(*streams))

# Initialize
await self.session.initialize()
Expand All @@ -40,10 +36,7 @@ async def connect_to_sse_server(self, server_url: str):

async def cleanup(self):
"""Properly clean up the session and streams"""
if self._session_context:
await self._session_context.__aexit__(None, None, None)
if self._streams_context:
await self._streams_context.__aexit__(None, None, None)
await self.exit_stack.aclose()

async def process_query(self, query: str) -> str:
"""Process a query using OpenAI API and available tools"""
Expand Down Expand Up @@ -105,8 +98,8 @@ async def process_query(self, query: str) -> str:
}
])

# print(f"Tool {tool_name} returned: {result.content[0].text}")
# print("messages", messages)
print(f"Tool {tool_name} returned: {result.content[0].text}")
print("messages", messages)
# Get next response from OpenAI
completion = await self.openai.chat.completions.create(
model=os.getenv("OPENAI_MODEL"),
Expand Down
21 changes: 9 additions & 12 deletions main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,8 @@
# main.py
from mcp.server.fastmcp import FastMCP
from dotenv import load_dotenv
import httpx
import json
import os
from bs4 import BeautifulSoup
from typing import Any
import httpx
from mcp.server.fastmcp import FastMCP
from starlette.applications import Starlette
Expand Down Expand Up @@ -92,13 +89,13 @@ async def get_docs(query: str, library: str):
# Stdio协议
if __name__ == "__main__":
mcp.run(transport="stdio")


# # SSE协议
# def create_starlette_app(mcp_server: Server, *, debug: bool = False) -> Starlette:
# """Create a Starlette application that can server the provied mcp server with SSE."""
# sse = SseServerTransport("/messages/")

#
# async def handle_sse(request: Request) -> None:
# async with sse.connect_sse(
# request.scope,
Expand All @@ -110,26 +107,26 @@ async def get_docs(query: str, library: str):
# write_stream,
# mcp_server.create_initialization_options(),
# )

#
# return Starlette(
# debug=debug,
# routes=[
# Route("/sse", endpoint=handle_sse),
# Mount("/messages/", app=sse.handle_post_message),
# ],
# )

#
# if __name__ == "__main__":
# mcp_server = mcp._mcp_server

# mcp_server = mcp._mcp_server
#
# import argparse
#
# parser = argparse.ArgumentParser(description='Run MCP SSE-based server')
# parser.add_argument('--host', default='0.0.0.0', help='Host to bind to')
# parser.add_argument('--port', type=int, default=8020, help='Port to listen on')
# args = parser.parse_args()

#
# # Bind SSE request handling to MCP server
# starlette_app = create_starlette_app(mcp_server, debug=True)

#
# uvicorn.run(starlette_app, host=args.host, port=args.port)