atlas-mcp/app/claims.py

"""Claim extraction helpers for Atlas layered outputs."""

from __future__ import annotations

from typing import Any

from app.models import AtlasClaim, AtlasEntity, AtlasProvenance


def _prov_to_dict(p: AtlasProvenance | None) -> dict[str, Any] | None:
    if p is None:
        return None
    return {
        "source": p.source,
        "method": p.retrieval_method,
        "confidence": p.confidence,
        "retrieved_at": p.retrieved_at,
    }


def _id_type_resource(identifier_type: str) -> str:
    if identifier_type == "mid":
        return "atlas:Mid"
    if identifier_type == "qid":
        return "atlas:WikidataQID"
    return f"atlas:{identifier_type}"


def _claim_to_dict(claim: AtlasClaim) -> dict[str, Any]:
    obj: dict[str, Any] = {
        "kind": claim.object.kind,
        "value": claim.object.value,
    }
    if claim.object.id_type:
        obj["id_type"] = _id_type_resource(claim.object.id_type)
    return {
        "claim_id": claim.claim_id,
        "layer": claim.layer,
        "status": claim.status,
        "subject": claim.subject,
        "predicate": claim.predicate,
        "object": obj,
        "provenance": _prov_to_dict(claim.provenance),
    }


def build_claim_sets(entity: AtlasEntity) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    raw_claims = [_claim_to_dict(c) for c in entity.claims if c.layer == "raw"]
    derived_claims = [_claim_to_dict(c) for c in entity.claims if c.layer == "derived"]
    return raw_claims, derived_claims