atlas-mcp/app/storage_service.py

"""Atlas persistence/read service via virtuoso-mcp (MCP transport)."""

from __future__ import annotations

import os
from typing import Any, Awaitable, Callable

from mcp import ClientSession
from mcp.client.sse import sse_client

from app.models import AtlasEntity
from app.triple_export import entity_to_turtle

ATLAS_GRAPH_IRI = os.getenv("ATLAS_GRAPH_IRI", "http://world.eu.org/atlas_data#")
VIRTUOSO_MCP_SSE_URL = os.getenv("ATLAS_VIRTUOSO_MCP_SSE_URL", "http://192.168.0.249:8501/mcp/sse")
VIRTUOSO_MCP_TIMEOUT = float(os.getenv("ATLAS_VIRTUOSO_MCP_TIMEOUT", "10"))
VIRTUOSO_MCP_SSE_READ_TIMEOUT = float(os.getenv("ATLAS_VIRTUOSO_MCP_SSE_READ_TIMEOUT", str(60 * 5)))

CallToolFn = Callable[[str, dict[str, Any]], Awaitable[dict[str, Any]]]


def _safe_fragment(value: str) -> str:
    value = (value or "").strip().lower()
    out = []
    for ch in value:
        if ch.isalnum() or ch in ["_", "-"]:
            out.append(ch)
        else:
            out.append("_")
    frag = "".join(out).strip("_")
    return frag or "entity"


def entity_iri(entity_id: str) -> str:
    return f"http://world.eu.org/atlas_data#entity_{_safe_fragment(entity_id)}"


class AtlasStorageService:
    def __init__(self, call_tool: CallToolFn | None = None):
        self._call_tool_override = call_tool

    async def _call_tool(self, tool_name: str, payload: dict[str, Any]) -> dict[str, Any]:
        if self._call_tool_override:
            return await self._call_tool_override(tool_name, payload)

        async with sse_client(
            VIRTUOSO_MCP_SSE_URL,
            timeout=VIRTUOSO_MCP_TIMEOUT,
            sse_read_timeout=VIRTUOSO_MCP_SSE_READ_TIMEOUT,
        ) as (read_stream, write_stream):
            async with ClientSession(read_stream, write_stream) as session:
                await session.initialize()
                result = await session.call_tool(tool_name, payload)
                if result.isError:
                    raise RuntimeError(f"Tool {tool_name} failed: {result.content}")
                if isinstance(result.structuredContent, dict):
                    return result.structuredContent
                return {"content": result.content}

    async def write_entity(self, entity: AtlasEntity) -> dict[str, Any]:
        ttl = entity_to_turtle(entity)
        try:
            result = await self._call_tool(
                "batch_insert",
                {
                    "ttl": ttl,
                    "graph": ATLAS_GRAPH_IRI,
                },
            )
            return {
                "status": "ok",
                "graph": ATLAS_GRAPH_IRI,
                "entity_id": entity.atlas_id,
                "result": result,
            }
        except Exception as exc:
            return {
                "status": "unfinished",
                "message": "Persistence path not fully available yet",
                "error": str(exc),
                "entity_id": entity.atlas_id,
            }

    async def read_entity_claims(self, entity_id: str, include_superseded: bool = False) -> dict[str, Any]:
        iri = entity_iri(entity_id)
        status_filter = "" if include_superseded else 'FILTER(?status = "active")'
        query = f"""
PREFIX atlas: <http://world.eu.org/atlas_ontology#>
SELECT ?entity ?label ?claim ?pred ?objIri ?objLit ?layer ?status ?prov ?src ?method ?conf ?ts
WHERE {{
  VALUES ?entity {{ <{iri}> }}

  ?entity a atlas:Entity ;
          atlas:canonicalLabel ?label ;
          atlas:hasClaim ?claim .

  ?claim atlas:claimSubjectIri ?entity ;
         atlas:claimPredicate ?pred ;
         atlas:claimLayer ?layer ;
         atlas:claimStatus ?status .

  OPTIONAL {{ ?claim atlas:claimObjectIri ?objIri . }}
  OPTIONAL {{ ?claim atlas:claimObjectLiteral ?objLit . }}

  OPTIONAL {{
    ?claim atlas:hasProvenance ?prov .
    ?prov atlas:provenanceSource ?src .
    OPTIONAL {{ ?prov atlas:retrievalMethod ?method . }}
    OPTIONAL {{ ?prov atlas:confidence ?conf . }}
    OPTIONAL {{ ?prov atlas:retrievedAt ?ts . }}
  }}

  {status_filter}
}}
ORDER BY ?claim
"""
        try:
            result = await self._call_tool("sparql_query", {"query": query})
            return {
                "status": "ok",
                "entity_id": entity_id,
                "query": query,
                "result": result,
            }
        except Exception as exc:
            return {
                "status": "unfinished",
                "message": "Read path not fully available yet",
                "error": str(exc),
                "entity_id": entity_id,
                "query": query,
            }