schema.py

"""Schema: entity types + relation types + LLM hints + recommended extractors.

A Schema is what the user (or an LLM) iterates on during a Session. It is
the **only** thing that decides what kind of graph comes out of a corpus.

Phase 0 = data class + preset loader skeleton. Real preset library lands
in Phase 1 (#5).
"""
from __future__ import annotations

import json
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any


@dataclass
class EntityType:
    name: str
    description: str = ""
    attrs: list[str] = field(default_factory=list)


@dataclass
class RelationType:
    name: str
    description: str = ""
    domain: list[str] = field(default_factory=list)   # source entity types
    range: list[str] = field(default_factory=list)    # target entity types


@dataclass
class Schema:
    """Live schema for a Session. Mutable through Session.refine()."""

    name: str = "ad_hoc"
    description: str = ""
    entities: list[EntityType] = field(default_factory=list)
    relations: list[RelationType] = field(default_factory=list)
    extractor_hint: str | None = None         # which extractor to prefer
    prompt_hints: dict[str, str] = field(default_factory=dict)
    version: int = 1

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "Schema":
        ents = [EntityType(**e) for e in data.get("entities", [])]
        rels = [RelationType(**r) for r in data.get("relations", [])]
        return cls(
            name=data.get("name", "ad_hoc"),
            description=data.get("description", ""),
            entities=ents,
            relations=rels,
            extractor_hint=data.get("extractor_hint"),
            prompt_hints=data.get("prompt_hints", {}),
            version=data.get("version", 1),
        )

    def add_entity(self, name: str, *, description: str = "", attrs: list[str] | None = None) -> None:
        if any(e.name == name for e in self.entities):
            return
        self.entities.append(EntityType(name=name, description=description, attrs=attrs or []))
        self.version += 1

    def remove_entity(self, name: str) -> bool:
        before = len(self.entities)
        self.entities = [e for e in self.entities if e.name != name]
        if len(self.entities) != before:
            self.version += 1
            return True
        return False

    def add_relation(
        self,
        name: str,
        *,
        description: str = "",
        domain: list[str] | None = None,
        range: list[str] | None = None,
    ) -> None:
        if any(r.name == name for r in self.relations):
            return
        self.relations.append(
            RelationType(name=name, description=description, domain=domain or [], range=range or [])
        )
        self.version += 1

    def remove_relation(self, name: str) -> bool:
        before = len(self.relations)
        self.relations = [r for r in self.relations if r.name != name]
        if len(self.relations) != before:
            self.version += 1
            return True
        return False

    def is_empty(self) -> bool:
        return not self.entities and not self.relations


# ---------------------------------------------------------------------------
# Preset library
# ---------------------------------------------------------------------------

PRESETS_DIR = Path(__file__).parent / "schemas"


def list_presets() -> list[dict[str, str]]:
    """Return built-in presets with name + one-line description.

    Sorted alphabetically; lazy-loaded yaml so missing pyyaml just yields
    descriptions=''. Used by CLI ``graphanything presets`` and Skill
    ``graphanything_list_presets``.
    """
    if not PRESETS_DIR.exists():
        return []

    out: list[dict[str, str]] = []
    paths = sorted(PRESETS_DIR.glob("*.yaml")) + sorted(PRESETS_DIR.glob("*.json"))
    for p in paths:
        desc = ""
        try:
            if p.suffix == ".yaml":
                import yaml  # type: ignore
                data = yaml.safe_load(p.read_text()) or {}
            else:
                data = json.loads(p.read_text())
            desc = (data or {}).get("description", "") or ""
        except Exception:  # pragma: no cover — listing must never crash
            desc = ""
        out.append({"name": p.stem, "description": desc})
    return out


def load_preset(name: str) -> Schema:
    """Load a preset by name. Tries .yaml first, then .json."""
    yaml_path = PRESETS_DIR / f"{name}.yaml"
    json_path = PRESETS_DIR / f"{name}.json"

    if yaml_path.exists():
        try:
            import yaml  # type: ignore
        except ImportError as e:
            raise RuntimeError(
                f"Preset '{name}' is YAML; install pyyaml or use a .json preset."
            ) from e
        data = yaml.safe_load(yaml_path.read_text())
        return Schema.from_dict(data)

    if json_path.exists():
        data = json.loads(json_path.read_text())
        return Schema.from_dict(data)

    available = [p["name"] for p in list_presets()]
    raise FileNotFoundError(
        f"No preset named '{name}' under {PRESETS_DIR}. "
        f"Available: {available}"
    )