atlas-mcp/app/models.py

"""Atlas internal data models."""

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional


@dataclass
class AtlasIdentifier:
    value: str
    source: str
    identifier_type: str


@dataclass
class AtlasAlias:
    label: str
    language: str = "und"
    source: str = "query"


@dataclass
class AtlasProvenance:
    source: str
    retrieval_method: str
    confidence: float = 0.0
    retrieved_at: Optional[str] = None
    evidence_property: str | None = None
    provider: str | None = None
    model: str | None = None


@dataclass
class AtlasClaimObject:
    kind: str
    value: str
    id_type: str | None = None


@dataclass
class AtlasClaim:
    claim_id: str
    subject: str
    predicate: str
    object: AtlasClaimObject
    layer: str
    status: str = "active"
    provenance: AtlasProvenance | None = None
    created_at: str | None = None


@dataclass
class AtlasEntity:
    atlas_id: str
    canonical_label: str
    canonical_description: str | None = None
    entity_type: str = "unknown"
    aliases: List[AtlasAlias] = field(default_factory=list)
    claims: List[AtlasClaim] = field(default_factory=list)
    raw_payload: Dict[str, Any] = field(default_factory=dict)
    needs_curation: bool = False

    def active_identifier(self, identifier_type: str) -> str | None:
        for claim in self.claims:
            if (
                claim.status == "active"
                and claim.predicate == "atlas:hasIdentifier"
                and claim.object.kind == "identifier"
                and claim.object.id_type == identifier_type
            ):
                return claim.object.value
        return None


@dataclass
class AtlasEnrichmentDataset:
    seed_entity: AtlasEntity
    related_entities: List[AtlasEntity] = field(default_factory=list)
    query_context: Dict[str, Any] = field(default_factory=dict)
    depth: int = 1