mcpserver/agent.py at master · AdLibML/mcpserver · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
import os
import logging
from langchain_mcp_adapters.client import MultiServerMCPClient
from langgraph.prebuilt import create_react_agent
from langchain_mcp_tools import convert_mcp_to_langchain_tools
from langchain_ollama import ChatOllama
import asyncio
from load_dotenv import load_dotenv
from src.utils.setup_logger import get_logger  # new import

load_dotenv()

logger = get_logger("agent")  # use centralized logger

mode = os.environ.get("MODE", "local")
logger.info(f"Running in {mode} mode")

if mode == "local":
    transport = "stdio"
else:
    transport = "sse"

model = ChatOllama(model="qwen2.5:7b")
logger.info(f"Using model: {model.model}")

if mode == "local":
    server_params = {
        "math": {
            "command": "python",
            "args": ["/Users/garyj/Code/mcpserver/src/servers/math_server.py"],
            "transport": transport,
        },
        "weather": {
            "command": "python",
            "args": ["/Users/garyj/Code/mcpserver/src/servers/weather_server.py"],
            "transport": transport,
        },
        "brave": {
            "command": "python",
            "args": ["/Users/garyj/Code/mcpserver/src/servers/brave_server.py"],
            "transport": transport,
        },
        "postgres": {
            "command": "python",
            "args": ["/Users/garyj/Code/mcpserver/src/servers/postgres_server.py"],
            "transport": transport,
        }
    }
else:
    server_params = {
        "math": {
            "command": "python",
            "url": os.environ.get("MATH_URL", "http://localhost:5001/mcp/sse"),
            "transport": transport,
        },
        "weather": {
            "command": "python",
            "url": os.environ.get("WEATHER_URL", "http://localhost:5000/mcp/sse"),
            "transport": transport,
        },
        "brave": {
            "command": "python",
            "url": os.environ.get("BRAVE_URL", "http://localhost:5002/mcp/sse"),
            "transport": transport,
        },
        "postgres": {
            "command": "python",
            "url": os.environ.get("POSTGRES_URL", "http://localhost:5003/mcp/sse"),
            "transport": transport,
        }
    }

async def main(query: str):
    logger.info("Agent starting with query: %s", query)
    async with MultiServerMCPClient(server_params) as client:
        agent = create_react_agent(model, client.get_tools())
        response = await agent.ainvoke({"messages": query})
        logger.info("Agent received response: %s", response)
    return response

if __name__ == "__main__":
    query = "What is (3+5) * 4 - 13"
    # query = "What is the weather in San Francisco?"
    query = "Who is Alex Karp?"
    response = asyncio.run(main(query))
    print(response)