mcp_server.py

#!/usr/bin/env python3
"""
CocoIndex MCP Server for Claude Code

Provides Claude Code with direct access to CocoIndex functionality:
- search: Semantic search across indexed documents
- index: Index/re-index documents
- add_document: Add a new document to the index
- list_documents: List all indexed documents
- get_metadata: Get LLM-extracted metadata for a document

This MCP server implements the Model Context Protocol to give Claude Code
semantic search capabilities over your local documents.

Usage:
    python mcp_server.py

Configuration in Claude Code (~/.claude.json):
    "cocoindex": {
        "command": "/path/to/cocoindex-claude-code/.venv/bin/python",
        "args": ["/path/to/cocoindex-claude-code/mcp_server.py"],
        "env": {
            "COCOINDEX_DATABASE_URL": "postgres://cocoindex:cocoindex@localhost/cocoindex"
        }
    }
"""

import json
import os
import sys
import subprocess
from pathlib import Path
from typing import Any

# Workspace paths - dynamically determined from script location
WORKSPACE = Path(__file__).parent.resolve()
DOCUMENTS_DIR = WORKSPACE / "data" / "documents"
VENV_PYTHON = WORKSPACE / ".venv" / "bin" / "python"

# Fallback for Windows
if not VENV_PYTHON.exists():
    VENV_PYTHON = WORKSPACE / ".venv" / "Scripts" / "python.exe"

def activate_and_run(code: str) -> str:
    """Run Python code in the CocoIndex virtual environment."""
    env = os.environ.copy()
    # Use environment variable if set, otherwise use default
    env["COCOINDEX_DATABASE_URL"] = os.environ.get(
        "COCOINDEX_DATABASE_URL",
        "postgres://cocoindex:cocoindex@localhost/cocoindex"
    )
    env["OPENAI_API_KEY"] = os.environ.get("OPENAI_API_KEY", "")
    env["ANTHROPIC_API_KEY"] = os.environ.get("ANTHROPIC_API_KEY", "")

    result = subprocess.run(
        [str(VENV_PYTHON), "-c", code],
        capture_output=True,
        text=True,
        cwd=str(WORKSPACE),
        env=env
    )
    if result.returncode != 0:
        return f"Error: {result.stderr}"
    return result.stdout


def search(query: str, limit: int = 5) -> dict:
    """Search indexed documents using semantic vector search."""
    workspace_path = str(WORKSPACE)
    code = f'''
import os
import sys
sys.path.insert(0, ".")
os.chdir("{workspace_path}")

from dotenv import load_dotenv
load_dotenv()

import cocoindex
cocoindex.init()

from flows.text_embedding import search

output = search("{query}", top_k={limit})
import json
results = []
for r in output.results:
    results.append({{
        "filename": r["filename"],
        "text": r["text"][:500],
        "score": round(r["score"], 3)
    }})
print(json.dumps(results, indent=2))
'''
    output = activate_and_run(code)
    try:
        return {"results": json.loads(output)}
    except:
        return {"results": [], "error": output}


def index_documents() -> dict:
    """Index or re-index all documents in the workspace."""
    result = subprocess.run(
        ["bash", "-c", f"cd {WORKSPACE} && source .venv/bin/activate && cocoindex update main.py"],
        capture_output=True,
        text=True
    )
    return {
        "success": result.returncode == 0,
        "output": result.stdout + result.stderr
    }


def list_documents() -> dict:
    """List all documents in the documents directory."""
    docs = []
    for f in DOCUMENTS_DIR.glob("**/*"):
        if f.is_file() and f.suffix in [".md", ".txt", ".py"]:
            docs.append({
                "path": str(f.relative_to(DOCUMENTS_DIR)),
                "size": f.stat().st_size,
                "type": f.suffix
            })
    return {"documents": docs, "count": len(docs)}


def get_metadata(filename: str) -> dict:
    """Get LLM-extracted metadata for a specific document."""
    code = f'''
import os
import psycopg
import json

with psycopg.connect("postgres://cocoindex:cocoindex@localhost/cocoindex") as conn:
    with conn.cursor() as cur:
        cur.execute("SELECT metadata, stats FROM document_metadata WHERE filename = %s", ("{filename}",))
        row = cur.fetchone()
        if row:
            print(json.dumps({{"metadata": row[0], "stats": row[1]}}))
        else:
            print(json.dumps({{"error": "Document not found"}}))
'''
    output = activate_and_run(code)
    try:
        return json.loads(output)
    except:
        return {"error": output}


def add_document(content: str, filename: str) -> dict:
    """Add a new document to the index."""
    filepath = DOCUMENTS_DIR / filename
    filepath.parent.mkdir(parents=True, exist_ok=True)
    filepath.write_text(content)

    # Re-index
    index_result = index_documents()
    return {
        "success": True,
        "path": str(filepath),
        "indexed": index_result["success"]
    }


# MCP Server Protocol Implementation
def handle_request(request: dict) -> dict:
    """Handle MCP protocol requests."""
    method = request.get("method", "")
    params = request.get("params", {})

    if method == "initialize":
        return {
            "protocolVersion": "2024-11-05",
            "serverInfo": {"name": "cocoindex-mcp", "version": "1.0.0"},
            "capabilities": {"tools": {}}
        }

    elif method == "tools/list":
        return {
            "tools": [
                {
                    "name": "cocoindex_search",
                    "description": "Search indexed documents using semantic vector search. Use this to find relevant information in the knowledge base.",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "query": {"type": "string", "description": "The search query"},
                            "limit": {"type": "integer", "description": "Max results (default 5)", "default": 5}
                        },
                        "required": ["query"]
                    }
                },
                {
                    "name": "cocoindex_index",
                    "description": "Index or re-index all documents in the CocoIndex workspace.",
                    "inputSchema": {"type": "object", "properties": {}}
                },
                {
                    "name": "cocoindex_list",
                    "description": "List all documents currently in the index.",
                    "inputSchema": {"type": "object", "properties": {}}
                },
                {
                    "name": "cocoindex_metadata",
                    "description": "Get LLM-extracted metadata (title, summary, key points, topics) for a document.",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "filename": {"type": "string", "description": "The filename to get metadata for"}
                        },
                        "required": ["filename"]
                    }
                },
                {
                    "name": "cocoindex_add",
                    "description": "Add a new document to the index.",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "filename": {"type": "string", "description": "Filename for the new document"},
                            "content": {"type": "string", "description": "Content of the document"}
                        },
                        "required": ["filename", "content"]
                    }
                }
            ]
        }

    elif method == "tools/call":
        tool_name = params.get("name", "")
        args = params.get("arguments", {})

        if tool_name == "cocoindex_search":
            result = search(args.get("query", ""), args.get("limit", 5))
        elif tool_name == "cocoindex_index":
            result = index_documents()
        elif tool_name == "cocoindex_list":
            result = list_documents()
        elif tool_name == "cocoindex_metadata":
            result = get_metadata(args.get("filename", ""))
        elif tool_name == "cocoindex_add":
            result = add_document(args.get("content", ""), args.get("filename", ""))
        else:
            result = {"error": f"Unknown tool: {tool_name}"}

        return {"content": [{"type": "text", "text": json.dumps(result, indent=2)}]}

    return {"error": "Unknown method"}


def main():
    """Run the MCP server using stdio transport."""
    # Load environment from workspace
    from dotenv import load_dotenv
    load_dotenv(WORKSPACE / ".env")

    for line in sys.stdin:
        try:
            request = json.loads(line)
            response = handle_request(request)
            response["jsonrpc"] = "2.0"
            response["id"] = request.get("id")
            print(json.dumps(response), flush=True)
        except Exception as e:
            error_response = {
                "jsonrpc": "2.0",
                "id": request.get("id") if 'request' in dir() else None,
                "error": {"code": -32603, "message": str(e)}
            }
            print(json.dumps(error_response), flush=True)


if __name__ == "__main__":
    main()