itsobaa/grabba-mcp repository overview

pip install grabba-mcp

grabba-mcp [transport_protocol]
# or for source: python src/server.py [transport_protocol]

import asyncio
import os
from typing import List, Dict, Optional
from langchain_core.tools import BaseTool, Tool
from mcp.models.mcp_server_config import McpServerConfig, McpServer
from mcp.client.transports.streamable_http import StreamableHttpConnection
from mcp.client.transports.stdio import StdioConnection
from mcp.client.multi_server_client import MultiServerMCPClient
from grabba import Job, JobStats, PuppetRegion # Import necessary Grabba Pydantic models
from dotenv import load_dotenv # For loading API key from .env

async def connect_and_use_mcp_tools(mcp_server_configs: List[McpServerConfig], api_key: Optional[str] = None) -> List[Tool]:
    """
    Connects to the MCP server(s), discovers its tools, and wraps them as LangChain Tools.
    Handles API key injection for HTTP connections.
    """
    try:
        mcp_client_config = {}
        for config in mcp_server_configs:
            # Pydantic V2 model validation
            mcp_server_model = McpServer.model_validate(config.mcp_server.model_dump())
            
            connection_headers = {}
            if api_key:
                # Use standard header name for API keys
                connection_headers["API_KEY"] = api_key 

            if mcp_server_model.transport == "streamable_http":
                server_params: StreamableHttpConnection = {
                    "transport": "streamable_http",
                    "url": str(mcp_server_model.url),
                    "env": config.env_variables or {}, # For other env variables, if any
                    "headers": connection_headers # Pass headers for HTTP transports
                }
            elif mcp_server_model.transport == "stdio":
                server_params: StdioConnection = {
                    "transport": "stdio",
                    "command": mcp_server_model.command, 
                    "args": mcp_server_model.args, 
                    "env": config.env_variables # For stdio, env maps to subprocess env vars
                }
            else:
                raise ValueError(f"Unsupported transport: {mcp_server_model.transport}")

            print(f"Client connecting with params: {server_params}")
            mcp_client_config[mcp_server_model.name] = server_params
        
        mcp_client = MultiServerMCPClient(mcp_client_config)
        tools: List[BaseTool] = await mcp_client.get_tools()
        print(f"Successfully loaded {len(tools)} tools.")
        return tools
    except Exception as e:
        print(f"Error connecting to MCP server or loading tools: {e}")
        return []


async def main():
    load_dotenv() # Load API key from a client-side .env file
    API_KEY = os.getenv("API_KEY", "YOUR_API_KEY_HERE_IF_NOT_ENV")

    # --- Configuration for Streamable HTTP Transport (Local or Public Instance) ---
    # For local: url="http://localhost:8283"
    # For public: url="https://mcp.grabba.dev/"
    http_mcp_config = McpServerConfig(
        mcp_server=McpServer(
            name="grabba-agent-http",
            transport="streamable_http",
            url="http://localhost:8283" # Or "https://mcp.grabba.dev/" for public
        )
    )

    print("\n--- Connecting via Streamable HTTP ---")
    http_tools = await connect_and_use_mcp_tools(
        mcp_server_configs=[http_mcp_config],
        api_key=API_KEY
    )

    if http_tools:
        print("\nAvailable HTTP Tools:")
        for tool in http_tools:
            print(f"- {tool.name}: {tool.description.split('.')[0]}.")
        
        # Example: Using the extract_data tool (adjust as per your Job Pydantic model)
        extract_tool = next((t for t in http_tools if t.name == "extract_data"), None)
        if extract_tool:
            print("\n--- Testing extract_data tool via HTTP ---")
            sample_job = Job(
                url="https://example.com/some-page",
                type="markdown", # or "pdf", "html" etc.
                parser="text-content",
                strategy="auto"
                # ... other required fields for Job
            )
            try:
                result_msg, result_data = await extract_tool.ainvoke({"extraction_data": sample_job})
                print(f"Extraction Result (HTTP): {result_msg}")
                if result_data:
                    print(f"Extraction Data (HTTP): {result_data.get('extracted_text', 'No text extracted')[:100]}...") # Print first 100 chars
            except Exception as e:
                print(f"Error calling extract_data via HTTP: {e}")
        else:
            print("extract_data tool not found in HTTP tools.")

        # Example: Using fetch_all_jobs tool
        fetch_jobs_tool = next((t for t in http_tools if t.name == "fetch_all_jobs"), None)
        if fetch_jobs_tool:
            print("\n--- Testing fetch_all_jobs tool via HTTP ---")
            try:
                result_msg, jobs_list = await fetch_jobs_tool.ainvoke({})
                print(f"Fetch Jobs Result (HTTP): {result_msg}")
                if jobs_list:
                    print(f"Fetched {len(jobs_list)} jobs.")
                    for job in jobs_list[:2]: # Print first 2 jobs
                        print(f"  - Job ID: {job.job_id}, URL: {job.url}")
            except Exception as e:
                print(f"Error calling fetch_all_jobs via HTTP: {e}")
        
        # Example: Using fetch_stats_data tool
        fetch_stats_tool = next((t for t in http_tools if t.name == "fetch_stats_data"), None)
        if fetch_stats_tool:
            print("\n--- Testing fetch_stats_data tool via HTTP ---")
            try:
                result_msg, stats_data = await fetch_stats_tool.ainvoke({})
                print(f"Fetch Stats Result (HTTP): {result_msg}")
                if stats_data:
                    print(f"Token Balance (HTTP): {stats_data.token_balance}")
                    print(f"Jobs Run (HTTP): {stats_data.jobs_run_count}")
            except Exception as e:
                print(f"Error calling fetch_stats_data via HTTP: {e}")

    # --- Configuration for Stdio Transport (e.g., to a Docker container running the server) ---
    # This assumes you have the 'itsobaa/grabba-mcp:latest' Docker image available.
    # The client launches a temporary Docker container for each tool call.
    stdio_mcp_config = McpServerConfig(
        mcp_server=McpServer(
            name="grabba-agent-stdio",
            transport="stdio",
            command="docker",
            args=[
                "run",
                "-i",          # Keep STDIN open for interactive communication
                "--rm",        # Remove container after exit
                "itsobaa/grabba-mcp:latest", # The Docker Hub image for Grabba MCP server
                "grabba-mcp", "stdio" # Command to run the server in stdio mode inside container
            ],
            env_variables={"API_KEY": API_KEY} # Pass API key as env var for stdio
        )
    )

    print("\n--- Connecting via Stdio (to Docker container as a subprocess) ---")
    stdio_tools = await connect_and_use_mcp_tools(
        mcp_server_configs=[stdio_mcp_config],
        api_key=API_KEY # Client might still pass for internal consistency, though env_variables is primary for stdio
    )

    if stdio_tools:
        print("\nAvailable Stdio Tools:")
        for tool in stdio_tools:
            print(f"- {tool.name}: {tool.description.split('.')[0]}.")
        
        # Example: Using the fetch_available_regions tool via Stdio
        fetch_regions_tool = next((t for t in stdio_tools if t.name == "fetch_available_regions"), None)
        if fetch_regions_tool:
            print("\n--- Testing fetch_available_regions tool via Stdio ---")
            try:
                result_msg, regions_list = await fetch_regions_tool.ainvoke({})
                print(f"Fetch Regions Result (Stdio): {result_msg}")
                if regions_list:
                    print(f"Fetched {len(regions_list)} regions.")
                    for region in regions_list[:3]: # Print first 3 regions
                        print(f"  - {region.display_name} ({region.code})")
            except Exception as e:
                print(f"Error calling fetch_available_regions via Stdio: {e}")
        else:
            print("fetch_available_regions tool not found in Stdio tools.")

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

your_project_root/
├── src/
│   └── server.py               # Main FastMCP server application
├── .env                        # Environment variables for local development
├── pyproject.toml              # Poetry project configuration
└── poetry.lock                 # Poetry dependency lock file
├── Dockerfile                  # Docker build instructions for the server
├── docker-compose.yml          # Docker Compose configuration for local development/deployment
├── .dockerignore               # Files to ignore during Docker build
├── .env                        # Example .env for docker-compose (for API_KEY)
├── README.md                   # This documentation file
├── pyproject.toml              # Root pyproject.toml (if using monorepo structure)
├── poetry.lock                 # Root poetry.lock (if using monorepo structure)
├── src/                        # Source code (often for the root project if it's a monorepo)
├── tests/                      # Project tests
└── ... (other project files like dist, docs, tox.ini, project.json etc.)

poetry run pytest