mcp_example.py

"""MCP server usage example for the Lumo SDK."""

from time import sleep
from lumo_sdk import (
    LumoClient,
    Message,
    MessageRole,
    MCPServer,
    StreamDoneEvent,
    StreamStepEvent,
    StreamTokenEvent,
)

# Initialize the client
client = LumoClient("sk-")

# Define MCP servers
mcp_servers = [
    MCPServer(
        command="npx",
        args=[
            "-y",
            "mcp-remote",
            "https://mcp.linear.app/sse",
            "--header",
            "Authorization: Bearer ${AUTH_TOKEN}",
        ],
        env={"AUTH_TOKEN": "lin_api"},
    ),
    MCPServer(
        command="npx",
        args=["-y", "mcp-remote", "https://remote.mcpservers.org/fetch/mcp"],
    ),
]

# Example 1: Simple task with MCP servers
print("Example 1: Simple task with MCP servers")
response = client.run_task(
    task="what are the latest issues on my linear board?",
    model="gpt-4.1-mini",
    base_url="https://api.openai.com/v1/chat/completions",
    max_steps=3,
    agent_type="mcp",
    mcp_servers=mcp_servers,
)

print(f"Final Answer: {response.final_answer}")
print(f"Number of steps: {len(response.steps)}")
for i, step in enumerate(response.steps, 1):
    print(f"  Step {i}:")
    if step.llm_output:
        print(f"    LLM Output: {step.llm_output[:100]}...")
    if step.tool_calls:
        print(f"    Tools used: {len(step.tool_calls)}")
        for tool_call in step.tool_calls:
            tool_name = tool_call.function.name
            arguments = tool_call.function.arguments
            print(
                f"      - {tool_name} (ID: {tool_call.id}, Type: {tool_call.type}, Arguments: {arguments})"
            )
    if step.token_usage:
        print(f"    Token Usage: {step.token_usage.total_tokens} tokens")
    else:
        print("    Tools used: None")
print(f"Total Token Usage: {response.token_usage.total_tokens} tokens")
print()

# Example 2: Streaming task with MCP servers
print("Example 2: Streaming task with MCP servers")
step_counter = 0
response_text = ""

for event in client.stream_task(
    task="what is at docs.starlight-search.com",
    model="gpt-4.1-nano",
    base_url="https://api.openai.com/v1/chat/completions",
    agent_type="mcp",
    mcp_servers=mcp_servers,
):
    # print(event)
    if isinstance(event, StreamTokenEvent):
        # Print tokens as they arrive
        token = event.content
        response_text += token
        sleep(0.01)
        print(token, end="", flush=True)

    elif isinstance(event, StreamStepEvent):
        # Handle step events with tool calls
        step_counter += 1
        step_data = event.step

        print("\n" + "=" * 60)
        print(f"Step {step_counter}")
        print("=" * 60)

        # Display tool calls if present
        tool_calls = step_data.tool_calls
        if tool_calls:
            print(f"\nTools used: {len(tool_calls)}")
            for tool_call in tool_calls:
                tool_name = tool_call.function.name
                tool_args = tool_call.function.arguments
                print(f"\n  Tool: {tool_name}")
                if tool_args:
                    print(f"  Arguments:")
                    if isinstance(tool_args, dict):
                        for key, value in tool_args.items():
                            print(f"    {key}: {value}")
                    else:
                        print(f"    {tool_args}")

        # Display step output if present
        output = step_data.llm_output
        if output:
            print(f"\nOutput: {output}")

        print("\n" + "=" * 60)
        print("\nContinuing response...\n")

    elif isinstance(event, StreamDoneEvent):
        print("=" * 60)
        print("Token Usage:")
        print("=" * 60)
        print(event.token_usage)
        print("\n" + "=" * 60)
        print("✅ Task completed!")
        print("=" * 60)
        break

print(f"\n\nFinal response length: {len(response_text)} characters")