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())