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


@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)
    identifiers: List[AtlasIdentifier] = field(default_factory=list)
    provenance: List[AtlasProvenance] = field(default_factory=list)
    raw_payload: Dict[str, Any] = field(default_factory=dict)
    needs_curation: bool = False


@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