atlas-mcp/app/mcp_server.py

"""FastMCP transport for Atlas tools."""

from __future__ import annotations

from pathlib import Path

from mcp.server.fastmcp import FastMCP
from mcp.server.transport_security import TransportSecuritySettings

from .atlas import enrich_entity, resolve_entity
from .claims import build_claim_sets
from .triple_export import entity_to_turtle

mcp = FastMCP(
    "atlas",
    transport_security=TransportSecuritySettings(
        enable_dns_rebinding_protection=False
    ),
)


@mcp.tool(name="resolve_entity", description="Resolve a subject string to a canonical Atlas entity.")
async def resolve_entity_tool(subject: str, context: str | None = None, debug: bool = False, debug_path: str | None = None):
    entity = await resolve_entity(subject, context)
    result = {
        "atlas_id": entity.atlas_id,
        "canonical_label": entity.canonical_label,
        "canonical_description": entity.canonical_description,
        "entity_type": entity.entity_type,
        "needs_curation": entity.needs_curation,
        "aliases": [alias.label for alias in entity.aliases],
        "identifiers": [
            {
                "value": identifier.value,
                "source": identifier.source,
                "identifier_type": identifier.identifier_type,
            }
            for identifier in entity.identifiers
        ],
        "provenance": [
            {
                "source": provenance.source,
                "retrieval_method": provenance.retrieval_method,
                "confidence": provenance.confidence,
                "retrieved_at": provenance.retrieved_at,
            }
            for provenance in entity.provenance
        ],
        "g_trends_payload": {k: v for k, v in entity.raw_payload.items() if k != "wikidata"},
        "wikidata_payload": entity.raw_payload.get("wikidata"),
    }
    if debug:
        raw_claims, derived_claims = build_claim_sets(entity)
        turtle = entity_to_turtle(entity)
        result["raw_claims"] = raw_claims
        result["derived_claims"] = derived_claims
        result["turtle"] = turtle
        if debug_path:
            path = Path(debug_path)
            path.parent.mkdir(parents=True, exist_ok=True)
            path.write_text(turtle, encoding="utf-8")
            result["turtle_path"] = str(path)
    return result


@mcp.tool(name="enrich_entity", description="Enrich a canonical Atlas entity.")
async def enrich_entity_tool(subject: str, depth: int = 1, context: str | None = None):
    entity = await resolve_entity(subject, context)
    result = enrich_entity(entity, depth=depth)
    return {
        "seed": {
            "atlas_id": result.seed_entity.atlas_id,
            "canonical_label": result.seed_entity.canonical_label,
        },
        "related_entities": [
            {
                "atlas_id": item.atlas_id,
                "canonical_label": item.canonical_label,
                "entity_type": item.entity_type,
            }
            for item in result.related_entities
        ],
        "query_context": result.query_context,
        "depth": result.depth,
    }