ThothII/harness/tht/session/models.py

import json
import os
import tempfile
import uuid
from datetime import datetime
from pathlib import Path
from typing import Literal, Self

import portalocker
import yaml
from pydantic import BaseModel, Field, ConfigDict

from tht.decisions import DecisionRecord


# Stub locale di _YamlModel. In the reference implementation questa base vive in
# mschema/models.py; qui la si replica perche' SessionManifest ha bisogno di
# populate_by_name=True (db_schema usa l'alias "schema"), che la base mschema non
# imposta. Identici to_yaml/from_yaml della controparte mschema.
class _YamlModel(BaseModel):
    model_config = ConfigDict(populate_by_name=True)

    def to_yaml(self, path: Path) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        data = self.model_dump(by_alias=True, mode="json", exclude_defaults=False)
        path.write_text(
            yaml.safe_dump(data, sort_keys=False, allow_unicode=True, width=120)
        )

    @classmethod
    def from_yaml(cls, path: Path) -> Self:
        raw = yaml.safe_load(path.read_text())
        return cls.model_validate(raw)


class PrincipalContext(BaseModel):
    """Trusted owner identity supplied by the runtime, never by a session document."""

    issuer: str
    subject: str
    display_name: str | None = None
    is_admin: bool = False

    model_config = ConfigDict(frozen=True)


def _private_directory(path: Path) -> Path:
    path.mkdir(parents=True, exist_ok=True)
    if os.name != "nt":
        path.chmod(0o700)
    return path


def _write_private_json(path: Path, data: dict[str, str]) -> None:
    fd, temporary_name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
    temporary = Path(temporary_name)
    try:
        with os.fdopen(fd, "w") as handle:
            json.dump(data, handle, sort_keys=True)
            handle.write("\n")
            handle.flush()
            os.fsync(handle.fileno())
        if os.name != "nt":
            temporary.chmod(0o600)
        os.replace(temporary, path)
        if os.name != "nt":
            path.chmod(0o600)
    finally:
        temporary.unlink(missing_ok=True)


def local_principal(home: Path | None = None) -> PrincipalContext:
    """Load or create the stable UUID identity for a local ThothII home."""
    if home is None:
        from tht.config import local_tht_home

        home = local_tht_home()
    identity_path = _private_directory(home) / "identity.json"
    lock_path = identity_path.with_suffix(".lock")
    with portalocker.Lock(lock_path, mode="a+", timeout=10):
        if identity_path.exists():
            try:
                identity = json.loads(identity_path.read_text())
                principal = PrincipalContext.model_validate(identity)
            except (json.JSONDecodeError, OSError, ValueError) as exc:
                raise ValueError(f"Invalid local identity: {identity_path}") from exc
            if principal.issuer != "local":
                raise ValueError(f"Invalid local identity issuer: {identity_path}")
            try:
                parsed = uuid.UUID(principal.subject, version=4)
            except ValueError as exc:
                raise ValueError(f"Invalid local identity subject: {identity_path}") from exc
            if str(parsed) != principal.subject or parsed.version != 4:
                raise ValueError(f"Invalid local identity subject: {identity_path}")
            return principal

        principal = PrincipalContext(issuer="local", subject=str(uuid.uuid4()))
        _write_private_json(
            identity_path, {"issuer": principal.issuer, "subject": principal.subject}
        )
        return principal


class SessionManifest(_YamlModel):
    id: str
    created_at: datetime
    status: Literal["open", "closed", "failed", "finalized"] = "open"
    question: str
    database: str
    db_schema: str = Field(alias="schema")
    # D12/D15: autore della sessione (auth) e versione del workflow usato.
    author: str | None = None
    summary: str | None = None
    updated_at: datetime | None = None
    updated_by: str | None = None
    schema_version: int | None = None
    provider: str | None = None
    model: str | None = None
    thinking: str | None = None
    name: str | None = None
    archived: bool = False
    group: str | None = None


class SessionSnapshot(BaseModel):
    """The current persisted session state, excluding the non-persistent chat stream."""

    principal: PrincipalContext | None = None
    manifest: SessionManifest
    artifacts: dict[str, str] = Field(default_factory=dict)
    decisions: list[DecisionRecord] = Field(default_factory=list)


class Candidate(BaseModel):
    kind: Literal["table", "column"]
    name: str
    signals: dict = {}
    evidence: list[str] = []
    decision: Literal["promoted", "excluded", "pending"] = "pending"
    decision_seq: int | None = None
    # D14a: valori citati nella domanda ancorati a questa colonna/tabella.
    grounded_values: list[dict] = []


class Join(BaseModel):
    from_: str = Field(alias="from")
    to: str
    source: str = ""
    decision: Literal["promoted", "excluded", "pending"] = "promoted"
    decision_seq: int | None = None

    model_config = {"populate_by_name": True}


class ExcludedItem(BaseModel):
    kind: Literal["table", "column"]
    name: str
    decision_seq: int | None = None


class SchemaLinking(BaseModel):
    question: str
    candidates: list[Candidate] = []
    joins: list[Join] = []
    excluded: list[ExcludedItem] = []
    open_questions: list[str] = []
    # D14b: formule di concetto approvate, parte dello schema-linking.
    concept_formulas: list[dict] = []

    model_config = {"extra": "forbid"}