src/lib/tools/models.py

"""Typed models for the tools registry."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Literal

Status = Literal["active", "stubbed", "planned"]
Category = Literal[
    "llm", "embedding", "reranker", "vector_db", "ocr", "pdf", "ui", "misc"
]


@dataclass(frozen=True)
class Pricing:
    input_per_million_usd: float | None
    output_per_million_usd: float | None
    image_input_tokens_per_image: int | None
    per_request_usd: float | None     # tools billed flat per call (e.g. Mistral OCR ~$0.001/page)
    last_updated: str | None
    verified: bool

    @property
    def is_free(self) -> bool:
        return (
            (self.input_per_million_usd or 0) == 0
            and (self.output_per_million_usd or 0) == 0
            and (self.per_request_usd or 0) == 0
        )


@dataclass(frozen=True)
class Tool:
    name: str
    category: Category
    status: Status
    used_by: tuple[str, ...]            # what this tool ACTUALLY runs in right now
    compatible_with: tuple[str, ...]    # what this tool COULD run in (incl. used_by)
    source_file: str | None
    provider: str | None
    pricing: Pricing | None
    notes: str

    @property
    def is_active(self) -> bool:
        return self.status == "active"

    def can_serve(self, stage: str) -> bool:
        return stage in self.compatible_with