atlas-mcp/app/storage_service.py

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

from __future__ import annotations

import json
import logging
import os
from typing import Any, Awaitable, Callable
from urllib.request import Request, urlopen

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

logger = logging.getLogger(__name__)

ATLAS_GRAPH_IRI = os.getenv("ATLAS_GRAPH_IRI", "http://world.eu.org/atlas_data#")
VIRTUOSO_RPC_URL = os.getenv("ATLAS_VIRTUOSO_RPC_URL", "http://192.168.0.249:8501/rpc")
VIRTUOSO_RPC_TIMEOUT = float(os.getenv("ATLAS_VIRTUOSO_RPC_TIMEOUT", "20"))

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)

        request = Request(
            VIRTUOSO_RPC_URL,
            data=json.dumps({"tool": tool_name, "input": payload}).encode("utf-8"),
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with urlopen(request, timeout=VIRTUOSO_RPC_TIMEOUT) as response:
            data = json.loads(response.read().decode("utf-8"))
        if isinstance(data, dict) and data.get("error"):
            raise RuntimeError(f"Tool {tool_name} failed: {data['error']}")
        if isinstance(data, dict) and "result" in data:
            return data["result"]
        return data

    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:
            logger.warning(
                "Atlas persistence failed for %s into %s: %s",
                entity.atlas_id,
                ATLAS_GRAPH_IRI,
                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,
            }