initial mcp server setup
This commit is contained in:
OwusuBlessing
@@ -0,0 +1,176 @@
|
||||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
import os
|
||||
from contextlib import AsyncExitStack
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import nest_asyncio
|
||||
from dotenv import load_dotenv
|
||||
from mcp import ClientSession, StdioServerParameters
|
||||
from mcp.client.stdio import stdio_client
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
# Add parent directory to path to import Config
|
||||
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
from Config import Config
|
||||
|
||||
# Apply nest_asyncio to allow nested event loops (needed for Jupyter/IPython)
|
||||
nest_asyncio.apply()
|
||||
|
||||
|
||||
|
||||
class MCPOpenAIClient:
|
||||
"""Client for interacting with OpenAI models using MCP tools."""
|
||||
|
||||
def __init__(self, model: str = "gpt-4o"):
|
||||
"""Initialize the OpenAI MCP client.
|
||||
|
||||
Args:
|
||||
model: The OpenAI model to use.
|
||||
"""
|
||||
# Initialize session and client objects
|
||||
self.session: Optional[ClientSession] = None
|
||||
self.exit_stack = AsyncExitStack()
|
||||
self.openai_client = AsyncOpenAI(api_key=Config.OPENAI_API_KEY)
|
||||
self.model = model
|
||||
self.stdio: Optional[Any] = None
|
||||
self.write: Optional[Any] = None
|
||||
|
||||
async def connect_to_server(self, server_script_path: str = "server.py"):
|
||||
"""Connect to an MCP server.
|
||||
|
||||
Args:
|
||||
server_script_path: Path to the server script.
|
||||
"""
|
||||
# Server configuration
|
||||
server_params = StdioServerParameters(
|
||||
command="python",
|
||||
args=[server_script_path],
|
||||
)
|
||||
|
||||
# Connect to the server
|
||||
stdio_transport = await self.exit_stack.enter_async_context(
|
||||
stdio_client(server_params)
|
||||
)
|
||||
self.stdio, self.write = stdio_transport
|
||||
self.session = await self.exit_stack.enter_async_context(
|
||||
ClientSession(self.stdio, self.write)
|
||||
)
|
||||
|
||||
# Initialize the connection
|
||||
await self.session.initialize()
|
||||
|
||||
# List available tools
|
||||
tools_result = await self.session.list_tools()
|
||||
print("\nConnected to server with tools:")
|
||||
for tool in tools_result.tools:
|
||||
print(f" - {tool.name}: {tool.description}")
|
||||
|
||||
async def get_mcp_tools(self) -> List[Dict[str, Any]]:
|
||||
"""Get available tools from the MCP server in OpenAI format.
|
||||
|
||||
Returns:
|
||||
A list of tools in OpenAI format.
|
||||
"""
|
||||
tools_result = await self.session.list_tools()
|
||||
return [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tool.name,
|
||||
"description": tool.description,
|
||||
"parameters": tool.inputSchema,
|
||||
},
|
||||
}
|
||||
for tool in tools_result.tools
|
||||
]
|
||||
|
||||
async def process_query(self, query: str) -> str:
|
||||
"""Process a query using OpenAI and available MCP tools.
|
||||
|
||||
Args:
|
||||
query: The user query.
|
||||
|
||||
Returns:
|
||||
The response from OpenAI.
|
||||
"""
|
||||
# Get available tools
|
||||
tools = await self.get_mcp_tools()
|
||||
|
||||
# Initial OpenAI API call
|
||||
response = await self.openai_client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[{"role": "user", "content": query}],
|
||||
tools=tools,
|
||||
tool_choice="auto",
|
||||
)
|
||||
|
||||
# Get assistant's response
|
||||
assistant_message = response.choices[0].message
|
||||
|
||||
# Initialize conversation with user query and assistant response
|
||||
messages = [
|
||||
{"role": "user", "content": query},
|
||||
assistant_message,
|
||||
]
|
||||
|
||||
# Handle tool calls if present
|
||||
if assistant_message.tool_calls:
|
||||
# Process each tool call
|
||||
for tool_call in assistant_message.tool_calls:
|
||||
# Execute tool call
|
||||
result = await self.session.call_tool(
|
||||
tool_call.function.name,
|
||||
arguments=json.loads(tool_call.function.arguments),
|
||||
)
|
||||
|
||||
# Add tool response to conversation
|
||||
messages.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": tool_call.id,
|
||||
"content": result.content[0].text,
|
||||
}
|
||||
)
|
||||
|
||||
# Get final response from OpenAI with tool results
|
||||
final_response = await self.openai_client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
tool_choice="none", # Don't allow more tool calls
|
||||
)
|
||||
|
||||
return final_response.choices[0].message.content
|
||||
|
||||
# No tool calls, just return the direct response
|
||||
return assistant_message.content
|
||||
|
||||
async def cleanup(self):
|
||||
"""Clean up resources."""
|
||||
try:
|
||||
await self.exit_stack.aclose()
|
||||
except Exception as e:
|
||||
# Ignore cleanup errors that don't affect functionality
|
||||
pass
|
||||
|
||||
|
||||
async def main():
|
||||
"""Main entry point for the client."""
|
||||
client = MCPOpenAIClient()
|
||||
try:
|
||||
await client.connect_to_server("server.py")
|
||||
|
||||
# Example: Ask about company vacation policy
|
||||
query = "How many employees do we have?"
|
||||
print(f"\nQuery: {query}")
|
||||
|
||||
response = await client.process_query(query)
|
||||
print(f"\nResponse: {response}")
|
||||
finally:
|
||||
await client.cleanup()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,53 @@
|
||||
import os
|
||||
import json
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
|
||||
# Create an MCP server
|
||||
mcp = FastMCP(
|
||||
name="Knowledge Base",
|
||||
host="0.0.0.0", # only used for SSE transport (localhost)
|
||||
port=8050, # only used for SSE transport (set this to any port)
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def get_knowledge_base() -> str:
|
||||
"""Retrieve the entire knowledge base as a formatted string.
|
||||
|
||||
Returns:
|
||||
A formatted string containing all Q&A pairs from the knowledge base.
|
||||
"""
|
||||
try:
|
||||
kb_path = os.path.join(os.path.dirname(__file__), "data", "kb.json")
|
||||
with open(kb_path, "r") as f:
|
||||
kb_data = json.load(f)
|
||||
|
||||
# Format the knowledge base as a string
|
||||
kb_text = "Here is the retrieved knowledge base:\n\n"
|
||||
|
||||
if isinstance(kb_data, list):
|
||||
for i, item in enumerate(kb_data, 1):
|
||||
if isinstance(item, dict):
|
||||
question = item.get("question", "Unknown question")
|
||||
answer = item.get("answer", "Unknown answer")
|
||||
else:
|
||||
question = f"Item {i}"
|
||||
answer = str(item)
|
||||
|
||||
kb_text += f"Q{i}: {question}\n"
|
||||
kb_text += f"A{i}: {answer}\n\n"
|
||||
else:
|
||||
kb_text += f"Knowledge base content: {json.dumps(kb_data, indent=2)}\n\n"
|
||||
|
||||
return kb_text
|
||||
except FileNotFoundError:
|
||||
return "Error: Knowledge base file not found"
|
||||
except json.JSONDecodeError:
|
||||
return "Error: Invalid JSON in knowledge base file"
|
||||
except Exception as e:
|
||||
return f"Error: {str(e)}"
|
||||
|
||||
|
||||
# Run the server
|
||||
if __name__ == "__main__":
|
||||
mcp.run(transport="stdio")
|
||||
Reference in New Issue
Block a user