from ai_query import mcp

async with mcp(command, *args, env=None) as server:
    # server.tools is ready to use
    ...

from ai_query import mcp_sse

async with mcp_sse(url, headers=None) as server:
    ...

from ai_query import mcp_http

async with mcp_http(url, headers=None) as server:
    ...

from ai_query import connect_mcp

server = await connect_mcp(command, *args, env=None)
try:
    # Use server.tools
    ...
finally:
    await server.close()

from ai_query import connect_mcp_sse

server = await connect_mcp_sse(url, headers=None)

from ai_query import connect_mcp_http

server = await connect_mcp_http(url, headers=None)

from ai_query import merge_tools

all_tools = merge_tools(
    {"local_tool": my_tool},  # ToolSet dict
    mcp_server1,               # MCPServer
    mcp_server2,               # MCPServer
)

@dataclass
class MCPServer:
    tools: ToolSet                    # Tools ready for generate_text/stream_text
    raw_tools: list[dict[str, Any]]   # Raw tool definitions from server
    transport: TransportType          # "stdio" | "sse" | "streamable_http"

    async def close(self) -> None:
        """Close the connection."""
        ...

class MCPClient:
    async def connect_stdio(self, command, args=None, env=None) -> MCPServer: ...
    async def connect_sse(self, url, headers=None) -> MCPServer: ...
    async def connect_http(self, url, headers=None) -> MCPServer: ...
    async def disconnect(self) -> None: ...

TransportType = Literal["stdio", "sse", "streamable_http"]

from ai_query import generate_text
from ai_query.providers import openai, mcp, merge_tools, tool, Field

@tool(description="Add two numbers")
def add(a: int = Field(description="First number"),
        b: int = Field(description="Second number")) -> int:
    return a + b

async def main():
    async with mcp("python", "weather_server.py") as weather:
        tools = merge_tools({"add": add}, weather)

        result = await generate_text(
            model=openai("gpt-4o"),
            prompt="What's 5 + 3? Also, what's the weather in Paris?",
            tools=tools,
        )
        print(result.text)