syndarix/mcp-servers/knowledge-base/server.py

"""
Knowledge Base MCP Server.

Provides RAG capabilities with pgvector for semantic search,
intelligent chunking, and collection management.
"""

import inspect
import logging
from contextlib import asynccontextmanager
from typing import Any, get_type_hints

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from fastmcp import FastMCP
from pydantic import Field

from collection_manager import CollectionManager, get_collection_manager
from collections.abc import AsyncIterator
from config import get_settings
from database import DatabaseManager, get_database_manager
from embeddings import EmbeddingGenerator, get_embedding_generator
from exceptions import KnowledgeBaseError
from models import (
    ChunkType,
    DeleteRequest,
    FileType,
    IngestRequest,
    SearchRequest,
    SearchType,
)
from search import SearchEngine, get_search_engine

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)

# Global instances
_database: DatabaseManager | None = None
_embeddings: EmbeddingGenerator | None = None
_search: SearchEngine | None = None
_collections: CollectionManager | None = None


@asynccontextmanager
async def lifespan(_app: FastAPI) -> AsyncIterator[None]:
    """Application lifespan handler."""
    global _database, _embeddings, _search, _collections

    logger.info("Starting Knowledge Base MCP Server...")

    # Initialize database
    _database = get_database_manager()
    await _database.initialize()

    # Initialize embedding generator
    _embeddings = get_embedding_generator()
    await _embeddings.initialize()

    # Initialize search engine
    _search = get_search_engine()

    # Initialize collection manager
    _collections = get_collection_manager()

    logger.info("Knowledge Base MCP Server started successfully")

    yield

    # Cleanup
    logger.info("Shutting down Knowledge Base MCP Server...")

    if _embeddings:
        await _embeddings.close()

    if _database:
        await _database.close()

    logger.info("Knowledge Base MCP Server shut down")


# Create FastMCP server
mcp = FastMCP("syndarix-knowledge-base")

# Create FastAPI app with lifespan
app = FastAPI(
    title="Knowledge Base MCP Server",
    description="RAG with pgvector for semantic search",
    version="0.1.0",
    lifespan=lifespan,
)


@app.get("/health")
async def health_check() -> dict[str, Any]:
    """Health check endpoint."""
    status: dict[str, Any] = {
        "status": "healthy",
        "service": "knowledge-base",
        "version": "0.1.0",
    }

    # Check database connection
    try:
        if _database and _database._pool:
            async with _database.acquire() as conn:
                await conn.fetchval("SELECT 1")
            status["database"] = "connected"
        else:
            status["database"] = "not initialized"
    except Exception as e:
        status["database"] = f"error: {e}"
        status["status"] = "degraded"

    return status


# Tool registry for JSON-RPC
_tool_registry: dict[str, Any] = {}


def _python_type_to_json_schema(python_type: Any) -> dict[str, Any]:
    """Convert Python type annotation to JSON Schema."""
    type_name = getattr(python_type, "__name__", str(python_type))

    if python_type is str or type_name == "str":
        return {"type": "string"}
    elif python_type is int or type_name == "int":
        return {"type": "integer"}
    elif python_type is float or type_name == "float":
        return {"type": "number"}
    elif python_type is bool or type_name == "bool":
        return {"type": "boolean"}
    elif type_name == "NoneType":
        return {"type": "null"}
    elif hasattr(python_type, "__origin__"):
        origin = python_type.__origin__
        args = getattr(python_type, "__args__", ())

        if origin is list:
            item_type = args[0] if args else Any
            return {"type": "array", "items": _python_type_to_json_schema(item_type)}
        elif origin is dict:
            return {"type": "object"}
        elif origin is type(None) or str(origin) == "typing.Union":
            # Handle Optional types (Union with None)
            non_none_args = [a for a in args if a is not type(None)]
            if len(non_none_args) == 1:
                schema = _python_type_to_json_schema(non_none_args[0])
                schema["nullable"] = True
                return schema
            return {"type": "object"}
    return {"type": "object"}


def _get_tool_schema(func: Any) -> dict[str, Any]:
    """Extract JSON Schema from a tool function."""
    sig = inspect.signature(func)
    hints = get_type_hints(func) if hasattr(func, "__annotations__") else {}

    properties: dict[str, Any] = {}
    required: list[str] = []

    for name, param in sig.parameters.items():
        if name in ("self", "cls"):
            continue

        prop: dict[str, Any] = {}

        # Get type from hints
        if name in hints:
            prop = _python_type_to_json_schema(hints[name])

        # Get description and constraints from Field default (FieldInfo object)
        default_val = param.default
        if hasattr(default_val, "description") and default_val.description:
            prop["description"] = default_val.description
        if hasattr(default_val, "ge") and default_val.ge is not None:
            prop["minimum"] = default_val.ge
        if hasattr(default_val, "le") and default_val.le is not None:
            prop["maximum"] = default_val.le
        # Handle Field default value (check for PydanticUndefined)
        if hasattr(default_val, "default"):
            field_default = default_val.default
            # Check if it's the "required" sentinel (...)
            if field_default is not ... and not (
                hasattr(field_default, "__class__")
                and "PydanticUndefined" in field_default.__class__.__name__
            ):
                prop["default"] = field_default

        # Determine if required
        if param.default is inspect.Parameter.empty:
            required.append(name)
        elif hasattr(default_val, "default"):
            field_default = default_val.default
            # Required if default is ellipsis or PydanticUndefined
            if field_default is ... or (
                hasattr(field_default, "__class__")
                and "PydanticUndefined" in field_default.__class__.__name__
            ):
                required.append(name)

        properties[name] = prop

    return {
        "type": "object",
        "properties": properties,
        "required": required,
    }


def _register_tool(name: str, tool_or_func: Any, description: str | None = None) -> None:
    """Register a tool in the registry.

    Handles both raw functions and FastMCP FunctionTool objects.
    """
    # Extract the underlying function from FastMCP FunctionTool if needed
    if hasattr(tool_or_func, "fn"):
        func = tool_or_func.fn
        # Use FunctionTool's description if available
        if not description and hasattr(tool_or_func, "description") and tool_or_func.description:
            description = tool_or_func.description
    else:
        func = tool_or_func

    _tool_registry[name] = {
        "func": func,
        "description": description or (func.__doc__ or "").strip(),
        "schema": _get_tool_schema(func),
    }


@app.get("/mcp/tools")
async def list_mcp_tools() -> dict[str, Any]:
    """
    Return list of available MCP tools with their schemas.

    This endpoint enables tool discovery for the backend MCP client.
    """
    tools = []
    for name, info in _tool_registry.items():
        tools.append({
            "name": name,
            "description": info["description"],
            "inputSchema": info["schema"],
        })

    return {"tools": tools}


@app.post("/mcp")
async def mcp_rpc(request: Request) -> JSONResponse:
    """
    JSON-RPC 2.0 endpoint for MCP tool execution.

    Request format:
    {
        "jsonrpc": "2.0",
        "method": "<tool_name>",
        "params": {...},
        "id": <request_id>
    }

    Response format:
    {
        "jsonrpc": "2.0",
        "result": {...},
        "id": <request_id>
    }
    """
    try:
        body = await request.json()
    except Exception as e:
        return JSONResponse(
            status_code=400,
            content={
                "jsonrpc": "2.0",
                "error": {"code": -32700, "message": f"Parse error: {e}"},
                "id": None,
            },
        )

    # Validate JSON-RPC structure
    jsonrpc = body.get("jsonrpc")
    method = body.get("method")
    params = body.get("params", {})
    request_id = body.get("id")

    if jsonrpc != "2.0":
        return JSONResponse(
            status_code=400,
            content={
                "jsonrpc": "2.0",
                "error": {"code": -32600, "message": "Invalid Request: jsonrpc must be '2.0'"},
                "id": request_id,
            },
        )

    if not method:
        return JSONResponse(
            status_code=400,
            content={
                "jsonrpc": "2.0",
                "error": {"code": -32600, "message": "Invalid Request: method is required"},
                "id": request_id,
            },
        )

    # Look up tool
    tool_info = _tool_registry.get(method)
    if not tool_info:
        return JSONResponse(
            status_code=404,
            content={
                "jsonrpc": "2.0",
                "error": {"code": -32601, "message": f"Method not found: {method}"},
                "id": request_id,
            },
        )

    # Execute tool
    try:
        func = tool_info["func"]

        # Resolve Field defaults for missing parameters
        sig = inspect.signature(func)
        resolved_params = dict(params)
        for name, param in sig.parameters.items():
            if name not in resolved_params:
                default_val = param.default
                # Check if it's a FieldInfo with a default value
                if hasattr(default_val, "default"):
                    field_default = default_val.default
                    # Only use if it has an actual default (not required)
                    if field_default is not ... and not (
                        hasattr(field_default, "__class__")
                        and "PydanticUndefined" in field_default.__class__.__name__
                    ):
                        resolved_params[name] = field_default

        result = await func(**resolved_params)
        return JSONResponse(
            content={
                "jsonrpc": "2.0",
                "result": result,
                "id": request_id,
            }
        )
    except TypeError as e:
        return JSONResponse(
            status_code=400,
            content={
                "jsonrpc": "2.0",
                "error": {"code": -32602, "message": f"Invalid params: {e}"},
                "id": request_id,
            },
        )
    except Exception as e:
        logger.error(f"Tool execution error: {e}")
        return JSONResponse(
            status_code=500,
            content={
                "jsonrpc": "2.0",
                "error": {"code": -32000, "message": f"Server error: {e}"},
                "id": request_id,
            },
        )


# MCP Tools


@mcp.tool()
async def search_knowledge(
    project_id: str = Field(..., description="Project ID for scoping"),
    agent_id: str = Field(..., description="Agent ID making the request"),
    query: str = Field(..., description="Search query"),
    search_type: str = Field(
        default="hybrid",
        description="Search type: semantic, keyword, or hybrid",
    ),
    collection: str | None = Field(
        default=None,
        description="Collection to search (None = all)",
    ),
    limit: int = Field(
        default=10,
        ge=1,
        le=100,
        description="Maximum number of results",
    ),
    threshold: float = Field(
        default=0.7,
        ge=0.0,
        le=1.0,
        description="Minimum similarity score",
    ),
    file_types: list[str] | None = Field(
        default=None,
        description="Filter by file types (python, javascript, etc.)",
    ),
) -> dict[str, Any]:
    """
    Search the knowledge base for relevant content.

    Supports semantic (vector), keyword (full-text), and hybrid search.
    Returns chunks ranked by relevance to the query.
    """
    try:
        # Parse search type
        try:
            search_type_enum = SearchType(search_type.lower())
        except ValueError:
            valid_types = [t.value for t in SearchType]
            return {
                "success": False,
                "error": f"Invalid search type: {search_type}. Valid types: {valid_types}",
            }

        # Parse file types
        file_type_enums = None
        if file_types:
            try:
                file_type_enums = [FileType(ft.lower()) for ft in file_types]
            except ValueError as e:
                return {
                    "success": False,
                    "error": f"Invalid file type: {e}",
                }

        request = SearchRequest(
            project_id=project_id,
            agent_id=agent_id,
            query=query,
            search_type=search_type_enum,
            collection=collection,
            limit=limit,
            threshold=threshold,
            file_types=file_type_enums,
        )

        response = await _search.search(request)  # type: ignore[union-attr]

        return {
            "success": True,
            "query": response.query,
            "search_type": response.search_type,
            "results": [
                {
                    "id": r.id,
                    "content": r.content,
                    "score": r.score,
                    "source_path": r.source_path,
                    "start_line": r.start_line,
                    "end_line": r.end_line,
                    "chunk_type": r.chunk_type,
                    "file_type": r.file_type,
                    "collection": r.collection,
                    "metadata": r.metadata,
                }
                for r in response.results
            ],
            "total_results": response.total_results,
            "search_time_ms": response.search_time_ms,
        }

    except KnowledgeBaseError as e:
        logger.error(f"Search error: {e}")
        return {
            "success": False,
            "error": e.message,
            "code": e.code.value,
        }
    except Exception as e:
        logger.error(f"Unexpected search error: {e}")
        return {
            "success": False,
            "error": str(e),
        }


@mcp.tool()
async def ingest_content(
    project_id: str = Field(..., description="Project ID for scoping"),
    agent_id: str = Field(..., description="Agent ID making the request"),
    content: str = Field(..., description="Content to ingest"),
    source_path: str | None = Field(
        default=None,
        description="Source file path for reference",
    ),
    collection: str = Field(
        default="default",
        description="Collection to store in",
    ),
    chunk_type: str = Field(
        default="text",
        description="Content type: code, markdown, or text",
    ),
    file_type: str | None = Field(
        default=None,
        description="File type for code chunking (python, javascript, etc.)",
    ),
    metadata: dict[str, Any] | None = Field(
        default=None,
        description="Additional metadata to store",
    ),
) -> dict[str, Any]:
    """
    Ingest content into the knowledge base.

    Content is automatically chunked based on type, embedded using
    the LLM Gateway, and stored in pgvector for search.
    """
    try:
        # Validate content size to prevent DoS
        settings = get_settings()
        content_size = len(content.encode("utf-8"))
        if content_size > settings.max_document_size:
            return {
                "success": False,
                "error": f"Content size ({content_size} bytes) exceeds maximum allowed ({settings.max_document_size} bytes)",
            }

        # Parse chunk type
        try:
            chunk_type_enum = ChunkType(chunk_type.lower())
        except ValueError:
            valid_types = [t.value for t in ChunkType]
            return {
                "success": False,
                "error": f"Invalid chunk type: {chunk_type}. Valid types: {valid_types}",
            }

        # Parse file type
        file_type_enum = None
        if file_type:
            try:
                file_type_enum = FileType(file_type.lower())
            except ValueError:
                valid_types = [t.value for t in FileType]
                return {
                    "success": False,
                    "error": f"Invalid file type: {file_type}. Valid types: {valid_types}",
                }

        request = IngestRequest(
            project_id=project_id,
            agent_id=agent_id,
            content=content,
            source_path=source_path,
            collection=collection,
            chunk_type=chunk_type_enum,
            file_type=file_type_enum,
            metadata=metadata or {},
        )

        result = await _collections.ingest(request)  # type: ignore[union-attr]

        return {
            "success": result.success,
            "chunks_created": result.chunks_created,
            "embeddings_generated": result.embeddings_generated,
            "source_path": result.source_path,
            "collection": result.collection,
            "chunk_ids": result.chunk_ids,
            "error": result.error,
        }

    except KnowledgeBaseError as e:
        logger.error(f"Ingest error: {e}")
        return {
            "success": False,
            "error": e.message,
            "code": e.code.value,
        }
    except Exception as e:
        logger.error(f"Unexpected ingest error: {e}")
        return {
            "success": False,
            "error": str(e),
        }


@mcp.tool()
async def delete_content(
    project_id: str = Field(..., description="Project ID for scoping"),
    agent_id: str = Field(..., description="Agent ID making the request"),
    source_path: str | None = Field(
        default=None,
        description="Delete by source file path",
    ),
    collection: str | None = Field(
        default=None,
        description="Delete entire collection",
    ),
    chunk_ids: list[str] | None = Field(
        default=None,
        description="Delete specific chunk IDs",
    ),
) -> dict[str, Any]:
    """
    Delete content from the knowledge base.

    Specify either source_path, collection, or chunk_ids to delete.
    """
    try:
        request = DeleteRequest(
            project_id=project_id,
            agent_id=agent_id,
            source_path=source_path,
            collection=collection,
            chunk_ids=chunk_ids,
        )

        result = await _collections.delete(request)  # type: ignore[union-attr]

        return {
            "success": result.success,
            "chunks_deleted": result.chunks_deleted,
            "error": result.error,
        }

    except KnowledgeBaseError as e:
        logger.error(f"Delete error: {e}")
        return {
            "success": False,
            "error": e.message,
            "code": e.code.value,
        }
    except Exception as e:
        logger.error(f"Unexpected delete error: {e}")
        return {
            "success": False,
            "error": str(e),
        }


@mcp.tool()
async def list_collections(
    project_id: str = Field(..., description="Project ID for scoping"),
    agent_id: str = Field(..., description="Agent ID making the request"),  # noqa: ARG001
) -> dict[str, Any]:
    """
    List all collections in a project's knowledge base.

    Returns collection names with chunk counts and file types.
    """
    try:
        result = await _collections.list_collections(project_id)  # type: ignore[union-attr]

        return {
            "success": True,
            "project_id": result.project_id,
            "collections": [
                {
                    "name": c.name,
                    "chunk_count": c.chunk_count,
                    "total_tokens": c.total_tokens,
                    "file_types": c.file_types,
                    "created_at": c.created_at.isoformat(),
                    "updated_at": c.updated_at.isoformat(),
                }
                for c in result.collections
            ],
            "total_collections": result.total_collections,
        }

    except KnowledgeBaseError as e:
        logger.error(f"List collections error: {e}")
        return {
            "success": False,
            "error": e.message,
            "code": e.code.value,
        }
    except Exception as e:
        logger.error(f"Unexpected list collections error: {e}")
        return {
            "success": False,
            "error": str(e),
        }


@mcp.tool()
async def get_collection_stats(
    project_id: str = Field(..., description="Project ID for scoping"),
    agent_id: str = Field(..., description="Agent ID making the request"),  # noqa: ARG001
    collection: str = Field(..., description="Collection name"),
) -> dict[str, Any]:
    """
    Get detailed statistics for a collection.

    Returns chunk counts, token totals, and type breakdowns.
    """
    try:
        stats = await _collections.get_collection_stats(project_id, collection)  # type: ignore[union-attr]

        return {
            "success": True,
            "collection": stats.collection,
            "project_id": stats.project_id,
            "chunk_count": stats.chunk_count,
            "unique_sources": stats.unique_sources,
            "total_tokens": stats.total_tokens,
            "avg_chunk_size": stats.avg_chunk_size,
            "chunk_types": stats.chunk_types,
            "file_types": stats.file_types,
            "oldest_chunk": stats.oldest_chunk.isoformat() if stats.oldest_chunk else None,
            "newest_chunk": stats.newest_chunk.isoformat() if stats.newest_chunk else None,
        }

    except KnowledgeBaseError as e:
        logger.error(f"Get collection stats error: {e}")
        return {
            "success": False,
            "error": e.message,
            "code": e.code.value,
        }
    except Exception as e:
        logger.error(f"Unexpected get collection stats error: {e}")
        return {
            "success": False,
            "error": str(e),
        }


@mcp.tool()
async def update_document(
    project_id: str = Field(..., description="Project ID for scoping"),
    agent_id: str = Field(..., description="Agent ID making the request"),
    source_path: str = Field(..., description="Source file path"),
    content: str = Field(..., description="New content"),
    collection: str = Field(
        default="default",
        description="Collection name",
    ),
    chunk_type: str = Field(
        default="text",
        description="Content type: code, markdown, or text",
    ),
    file_type: str | None = Field(
        default=None,
        description="File type for code chunking",
    ),
    metadata: dict[str, Any] | None = Field(
        default=None,
        description="Additional metadata",
    ),
) -> dict[str, Any]:
    """
    Update a document in the knowledge base.

    Replaces all existing chunks for the source path with new content.
    """
    try:
        # Validate content size to prevent DoS
        settings = get_settings()
        content_size = len(content.encode("utf-8"))
        if content_size > settings.max_document_size:
            return {
                "success": False,
                "error": f"Content size ({content_size} bytes) exceeds maximum allowed ({settings.max_document_size} bytes)",
            }

        # Parse chunk type
        try:
            chunk_type_enum = ChunkType(chunk_type.lower())
        except ValueError:
            valid_types = [t.value for t in ChunkType]
            return {
                "success": False,
                "error": f"Invalid chunk type: {chunk_type}. Valid types: {valid_types}",
            }

        # Parse file type
        file_type_enum = None
        if file_type:
            try:
                file_type_enum = FileType(file_type.lower())
            except ValueError:
                valid_types = [t.value for t in FileType]
                return {
                    "success": False,
                    "error": f"Invalid file type: {file_type}. Valid types: {valid_types}",
                }

        result = await _collections.update_document(  # type: ignore[union-attr]
            project_id=project_id,
            agent_id=agent_id,
            source_path=source_path,
            content=content,
            collection=collection,
            chunk_type=chunk_type_enum,
            file_type=file_type_enum,
            metadata=metadata,
        )

        return {
            "success": result.success,
            "chunks_created": result.chunks_created,
            "embeddings_generated": result.embeddings_generated,
            "source_path": result.source_path,
            "collection": result.collection,
            "chunk_ids": result.chunk_ids,
            "error": result.error,
        }

    except KnowledgeBaseError as e:
        logger.error(f"Update document error: {e}")
        return {
            "success": False,
            "error": e.message,
            "code": e.code.value,
        }
    except Exception as e:
        logger.error(f"Unexpected update document error: {e}")
        return {
            "success": False,
            "error": str(e),
        }


# Register tools in the JSON-RPC registry
# This must happen after tool functions are defined
_register_tool("search_knowledge", search_knowledge)
_register_tool("ingest_content", ingest_content)
_register_tool("delete_content", delete_content)
_register_tool("list_collections", list_collections)
_register_tool("get_collection_stats", get_collection_stats)
_register_tool("update_document", update_document)


def main() -> None:
    """Run the server."""
    import uvicorn

    settings = get_settings()

    uvicorn.run(
        "server:app",
        host=settings.host,
        port=settings.port,
        reload=settings.debug,
        log_level="info",
    )


if __name__ == "__main__":
    main()