Using MCP Servers with Yi-AI Router

MCP (Model Context Protocol) servers provide tools and resources for LLMs. With Yi-AI Router, you can connect MCP servers to any model that supports tool/function calling through our OpenAI-compatible endpoint.

How It Works

The MCP protocol defines tools that LLMs can call. The flow is:

1. Your MCP client connects to a local MCP server and discovers its available tools
2. The client converts MCP tool definitions to OpenAI function-calling format
3. Tools are sent alongside your request to Yi-AI Router
4. The model decides which tool to call and returns tool_calls in the response
5. Your client executes the tool via the MCP server and sends results back
6. The model produces its final response

Quick Start

Prerequisites

• Python 3.10+
• A Yi-AI account with an API Key (sk-yi-...)
• An MCP server to connect to (e.g., filesystem, fetch, database)

Example: Filesystem MCP Server

Create a .env file with your Yi-AI API key:

OPENAI_API_KEY=sk-yi-your-key

Create mcp-client.py:

import asyncio
from typing import Optional
from contextlib import AsyncExitStack
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from openai import OpenAI
from dotenv import load_dotenv
import json

load_dotenv()

MODEL = "claude-sonnet-4-20250514"

SERVER_CONFIG = {
    "command": "npx",
    "args": [
        "-y",
        "@modelcontextprotocol/server-filesystem",
        "/Applications",
    ],
    "env": None,
}


def convert_tool_format(tool):
    converted_tool = {
        "type": "function",
        "function": {
            "name": tool.name,
            "description": tool.description,
            "parameters": {
                "type": "object",
                "properties": tool.inputSchema["properties"],
                "required": tool.inputSchema["required"],
            },
        },
    }
    return converted_tool


class MCPClient:
    def __init__(self):
        self.session: Optional[ClientSession] = None
        self.exit_stack = AsyncExitStack()
        self.openai = OpenAI(
            base_url="https://www.yiairouter.com/api/v1",
        )

    async def connect_to_server(self, server_config):
        server_params = StdioServerParameters(**server_config)
        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)
        )
        await self.session.initialize()
        response = await self.session.list_tools()
        print(
            "\nConnected to server, tools:",
            [tool.name for tool in response.tools],
        )
        self.messages = []

    async def process_query(self, query: str) -> str:
        self.messages.append({"role": "user", "content": query})
        response = await self.session.list_tools()
        available_tools = [convert_tool_format(tool) for tool in response.tools]
        response = self.openai.chat.completions.create(
            model=MODEL,
            tools=available_tools,
            messages=self.messages,
        )
        self.messages.append(response.choices[0].message.model_dump())
        final_text = []
        content = response.choices[0].message
        if content.tool_calls is not None:
            tool_name = content.tool_calls[0].function.name
            tool_args = content.tool_calls[0].function.arguments
            tool_args = json.loads(tool_args) if tool_args else {}
            try:
                result = await self.session.call_tool(tool_name, tool_args)
                final_text.append(
                    f"[Called tool {tool_name} with args {tool_args}]"
                )
            except Exception as e:
                print(f"Error calling tool {tool_name}: {e}")
                result = None
            self.messages.append(
                {
                    "role": "tool",
                    "tool_call_id": content.tool_calls[0].id,
                    "name": tool_name,
                    "content": result.content,
                }
            )
            response = self.openai.chat.completions.create(
                model=MODEL,
                max_tokens=1000,
                messages=self.messages,
            )
            final_text.append(response.choices[0].message.content)
        else:
            final_text.append(content.content)
        return "\n".join(final_text)

    async def chat_loop(self):
        print("\nMCP Client started!")
        print("Enter your queries or type 'quit' to exit.")
        while True:
            try:
                query = input("\nQuery: ").strip()
                result = await self.process_query(query)
                print("Result:")
                print(result)
            except Exception as e:
                print(f"Error: {str(e)}")

    async def cleanup(self):
        await self.exit_stack.aclose()


async def main():
    client = MCPClient()
    try:
        await client.connect_to_server(SERVER_CONFIG)
        await client.chat_loop()
    finally:
        await client.cleanup()


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

Run It

pip install mcp openai python-dotenv
python mcp-client.py

The client will connect to the filesystem MCP server and start an interactive session where you can query files and directories through Yi-AI Router.

Supported Models

Any model on Yi-AI Router that supports function/tool calling works with MCP. This includes:

• Claude Sonnet 4 / Haiku 4
• GPT-4o / GPT-4o-mini
• Gemini 2.5 Flash / Pro
• DeepSeek V3 / R1
• And all other models with tool support

Other MCP Servers

The same pattern works with any MCP server. Replace the SERVER_CONFIG with your server's configuration:

# Database MCP
SERVER_CONFIG = {
    "command": "npx",
    "args": ["-y", "@anthropic/mcp-server-db", "postgresql://..."],
    "env": None,
}

# Fetch MCP
SERVER_CONFIG = {
    "command": "npx",
    "args": ["-y", "@anthropic/mcp-server-fetch"],
    "env": None,
}

# GitHub MCP
SERVER_CONFIG = {
    "command": "npx",
    "args": ["-y", "@modelcontextprotocol/server-github"],
    "env": {"GITHUB_TOKEN": "..."},
}

How Yi-AI Router Handles It

Yi-AI Router's /api/v1/chat/completions endpoint transparently forwards all OpenAI-format fields — including tools, tool_choice, and tool_calls. No special configuration is needed. The platform handles model mapping, routing, and billing automatically.