MCP Agent

import asyncio
from ai_query.agents import Agent, MemoryStorage
from ai_query.providers import openai, connect_mcp_http, step_count_is

class MCPAgent(Agent):
    def __init__(self, mcp_tools: dict = None):
        super().__init__(
            "mcp-agent",
            model=openai("gpt-4o-mini"),
            system="You are a helpful assistant. Use tools when needed.",
            storage=MemoryStorage(),
            initial_state={"queries": 0},
            stop_when=step_count_is(10),
            tools=mcp_tools or {}
        )

    def on_step_finish(self, event):
        if event.step.tool_calls:
            for tc in event.step.tool_calls:
                print(f"  [Tool] {tc.name}")

    async def on_start(self):
        print(f"Agent started with {len(self.messages)} messages")


async def main():
    # Connect to MCP server
    print("Connecting to MCP server...")
    mcp_server = await connect_mcp_http("https://your-mcp-server.com/mcp")

    try:
        print(f"Tools available: {list(mcp_server.tools.keys())}")

        # Create agent with MCP tools
        agent = MCPAgent(mcp_tools=mcp_server.tools)

        async with agent:
            while True:
                question = input("\nYou: ")
                if question.lower() in ('quit', 'exit'):
                    break

                print("\nAgent: ", end="", flush=True)
                async for chunk in agent.stream(question):
                    print(chunk, end="", flush=True)
                print()
    finally:
        await mcp_server.close()


if __name__ == "__main__":
    asyncio.run(main())

from ai_query.providers.google import google
from ai_query.providers.anthropic import anthropic

class GoogleMCPAgent(Agent):
    def __init__(self, mcp_tools: dict = None):
        super().__init__(
            "google-mcp-agent",
            model=google("gemini-2.0-flash"),
            storage=MemoryStorage(),
            tools=mcp_tools or {}
        )

class ClaudeMCPAgent(Agent):
    def __init__(self, mcp_tools: dict = None):
        super().__init__(
            "claude-mcp-agent",
            model=anthropic("claude-sonnet-4-20250514"),
            storage=MemoryStorage(),
            tools=mcp_tools or {}
        )

from ai_query import connect_mcp_sse

async with connect_mcp_sse("https://server.com/sse") as mcp_server:
    # Same usage pattern
    agent = MCPAgent(mcp_tools=mcp_server.tools)
    async with agent:
        response = await agent.chat("Hello!")

from ai_query.agents import Agent, SQLiteStorage
from ai_query.providers import openai

class PersistentMCPAgent(Agent):
    def __init__(self, mcp_tools: dict = None):
        super().__init__(
            "persistent-mcp-agent",
            model=openai("gpt-4o"),
            storage=SQLiteStorage("./mcp_agent.db"),
            tools=mcp_tools or {}
        )