import asyncio
import sys
from openai import AsyncOpenAI
from openai.types.chat import ChatCompletionUserMessageParam
from mcp import ClientSession
from mcp.client.sse import sse_client


async def main():
    print("测试LiteLLM的MCP功能...")
    
    try:
        # Initialize OpenAI client
        print("初始化OpenAI客户端...")
        client = AsyncOpenAI(api_key="sk-1234", base_url="http://localhost:4000")
        print("OpenAI客户端初始化完成")

        # Test basic LiteLLM functionality first
        print("测试基本的LiteLLM功能...")
        response = await client.chat.completions.create(
            model="gemini-flash",
            messages=[
                {"role": "user", "content": "Hello, this is a test message."}
            ]
        )
        print(f"LiteLLM响应: {response.choices[0].message.content}")
        
        # Now test MCP endpoint
        print("\n测试MCP端点...")
        
        # 添加超时处理
        try:
            async with asyncio.timeout(10):  # 10秒超时
                print("正在建立SSE连接到 /mcp/ 端点...")
                async with sse_client("http://localhost:4000/mcp/") as (read, write):
                    print("SSE连接建立成功，初始化会话...")
                    async with ClientSession(read, write) as session:
                        print("正在初始化MCP会话...")
                        await session.initialize()
                        print("MCP会话初始化成功！")

                        # List available tools
                        print("获取可用工具列表...")
                        tools = await session.list_tools()
                        print(f"找到 {len(tools.tools)} 个工具:")
                        for tool in tools.tools:
                            print(f"  - {tool.name}: {tool.description}")
                            
        except asyncio.TimeoutError:
            print("MCP连接超时！")
            print("这可能意味着:")
            print("1. LiteLLM版本不支持MCP功能")
            print("2. MCP功能需要额外配置")
            print("3. /mcp/ 端点不存在")
            return
            
    except Exception as e:
        print(f"发生错误: {type(e).__name__}: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    print("启动LiteLLM MCP测试...")
    asyncio.run(main())
    print("测试完成")