aethon_open_structure/model.py

from __future__ import annotations

import sys
from dataclasses import dataclass
from pathlib import Path

from huggingface_hub import snapshot_download


@dataclass(frozen=True)
class AethonOpenStructureResponse:
    answer: str
    text: str
    explanation: str
    proof: tuple[str, ...]
    reasoning: tuple[str, ...]
    mode: str


class AethonOpenStructureModel:
    """Portable model-facing wrapper for the public Aethon Open Structure release."""

    def __init__(self, release_dir: str | Path) -> None:
        self.release_dir = Path(release_dir)
        runtime_root = self.release_dir / "runtime"
        if str(runtime_root) not in sys.path:
            sys.path.insert(0, str(runtime_root))

        from aethon.rfi_bundle import NativeBundleManager  # type: ignore

        self._runtime = NativeBundleManager.load(self.release_dir / "bundle")
        self.metadata = getattr(self._runtime, "metadata", None)

    @classmethod
    def from_hub(
        cls,
        repo_id: str,
        *,
        local_dir: str | Path = "aethon_open_structure_release",
    ) -> "AethonOpenStructureModel":
        release_dir = snapshot_download(
            repo_id=repo_id,
            local_dir=str(local_dir),
            local_dir_use_symlinks=False,
        )
        return cls(release_dir)

    def ask(self, question: str) -> AethonOpenStructureResponse:
        response = self._runtime.ask(question)
        return AethonOpenStructureResponse(
            answer=response.answer,
            text=response.text,
            explanation=response.explanation,
            proof=tuple(response.proof),
            reasoning=tuple(response.reasoning),
            mode=response.mode,
        )

    def ask_messages(self, messages: list[dict[str, str]]) -> AethonOpenStructureResponse:
        response = self._runtime.ask_messages(messages)
        return AethonOpenStructureResponse(
            answer=response.answer,
            text=response.text,
            explanation=response.explanation,
            proof=tuple(response.proof),
            reasoning=tuple(response.reasoning),
            mode=response.mode,
        )

    def learn(self, text: str) -> dict[str, object]:
        return self._runtime.learn(text)

    def close(self) -> None:
        self._runtime.close()