atlas-mcp/app/cache.py

"""Local caches for Atlas."""

from __future__ import annotations

from collections import OrderedDict
from typing import Iterable, Optional

from app.models import AtlasEntity


class EntityCache:
    def __init__(self, max_entries: int = 512):
        self.max_entries = max_entries
        self._data: OrderedDict[str, AtlasEntity] = OrderedDict()

    def _normalize_token(self, token: str | None) -> str:
        return str(token or "").strip().lower()

    def get(self, token: str | None) -> Optional[AtlasEntity]:
        key = self._normalize_token(token)
        if not key:
            return None
        entity = self._data.get(key)
        if entity is not None:
            self._data.move_to_end(key)
        return entity

    def store(self, entity: AtlasEntity, extra_tokens: Iterable[str] | None = None) -> None:
        tokens = set(extra_tokens or [])
        tokens.add(entity.canonical_label)
        tokens.add(entity.atlas_id)
        for alias in entity.aliases:
            tokens.add(alias.label)
        mid = entity.active_identifier("mid")
        if mid:
            tokens.add(mid)
        for token in tokens:
            key = self._normalize_token(token)
            if not key:
                continue
            self._data[key] = entity
            self._data.move_to_end(key)
        while len(self._data) > self.max_entries:
            self._data.popitem(last=False)