obsidian-sorter/sorter/jev.py

"""
TypeSafe JEV client — stdlib only.

Deliberately dependency-free (urllib, not the official SDK) so the sorter
deploys onto hosts without pip.

JEV answers three primitive question types:
    choice  → picks one label from your criteria, with calibrated confidence
    noul    → probability that a statement is true (0..1)
    score   → expected value on an ordered rubric

API: POST https://api.typesafe.ai/v1/systemone
"""

from __future__ import annotations

import glob
import json
import os
import re
import urllib.error
import urllib.request
from dataclasses import dataclass, field
from typing import Any

DEFAULT_BASE_URL = "https://api.typesafe.ai/v1"
DEFAULT_MODEL = "jev-latest"

# JEV pricing: $0.042 per 1M input tokens; output tokens are free.
USD_PER_INPUT_TOKEN = 0.042 / 1_000_000

# Fallback locations for the API key. On NixOS the key usually lives in
# environment.d, which is loaded for systemd user sessions but NOT for
# non-interactive SSH shells — so we read it directly.
SECRETS_GLOBS = (
    "~/.config/environment.d/*.conf",
    "~/.config/obsidian-sorter/secrets.conf",
)


def _key_from_secrets() -> str | None:
    """Find TYPESAFE_API_KEY in the environment.d secrets files."""
    for pattern in SECRETS_GLOBS:
        for path in sorted(glob.glob(os.path.expanduser(pattern))):
            try:
                text = open(path, encoding="utf-8").read()
            except OSError:
                continue
            m = re.search(
                r"^\s*(?:export\s+)?TYPESAFE_API_KEY\s*=\s*(.+?)\s*$",
                text, re.MULTILINE,
            )
            if m:
                value = m.group(1).strip().strip('"').strip("'")
                if value:
                    return value
    return None


class JevError(RuntimeError):
    """Raised when a JEV request fails."""


# ------------------------------------------------------------ question builders

def choice(instructions: str, criteria: dict[str, str | None]) -> dict[str, Any]:
    """A question that selects one label from `criteria`."""
    return {"type": "choice", "instructions": instructions, "criteria": criteria}


def noul(instructions: str, true: str | None = None, false: str | None = None) -> dict[str, Any]:
    """A yes/no question."""
    q: dict[str, Any] = {"type": "noul", "instructions": instructions}
    if true is not None or false is not None:
        q["criteria"] = {"true": true, "false": false}
    return q


def score(instructions: str, criteria: list[str]) -> dict[str, Any]:
    """An ordered rating; `criteria` is one description per level from 0."""
    return {"type": "score", "instructions": instructions, "criteria": criteria}


# --------------------------------------------------------------------- results

@dataclass
class Answer:
    """One answer returned by JEV."""

    name: str
    kind: str                       # "choice" | "noul" | "score"
    raw: dict[str, Any] = field(default_factory=dict)

    @property
    def label(self) -> str | None:
        return self.raw.get("choice")

    @property
    def confidence(self) -> float:
        """Confidence for choice/score; probability for noul."""
        if self.kind == "noul":
            return float(self.raw.get("noul", 0.5))
        return float(self.raw.get("confidence", 0.0))

    @property
    def probabilities(self) -> dict[str, float]:
        return self.raw.get("probabilities", {}) or {}

    @property
    def value(self) -> float | None:
        if self.kind == "score":
            return float(self.raw.get("score", 0.0))
        return None


@dataclass
class JevResult:
    """A complete JEV response."""

    answers: dict[str, Answer]
    input_tokens: int = 0
    output_tokens: int = 0
    model: str = ""
    error: str | None = None

    @property
    def ok(self) -> bool:
        return self.error is None

    @property
    def cost_usd(self) -> float:
        return self.input_tokens * USD_PER_INPUT_TOKEN

    def answer(self, name: str) -> Answer | None:
        return self.answers.get(name)


# ---------------------------------------------------------------------- client

class JevClient:
    """Thin JEV HTTP client."""

    def __init__(
        self,
        api_key: str | None = None,
        model: str = DEFAULT_MODEL,
        base_url: str = DEFAULT_BASE_URL,
        timeout: float = 90.0,
    ):
        key = (
            api_key
            or os.environ.get("TYPESAFE_API_KEY", "")
            or _key_from_secrets()
            or ""
        )
        if not key:
            raise JevError(
                "No API key. Set TYPESAFE_API_KEY in the environment, add it to "
                "~/.config/environment.d/*.conf, or pass api_key=."
            )
        self.api_key = key
        self.model = model
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout

    def ask(
        self,
        state: Any,
        questions: dict[str, dict[str, Any]],
    ) -> JevResult:
        """Send one system-one request. Never raises for API errors — returns a JevResult with .error set."""
        if not questions:
            raise JevError("At least one question is required.")

        payload = {"state": state, "model": self.model, "questions": questions}
        body = json.dumps(payload).encode("utf-8")

        req = urllib.request.Request(
            f"{self.base_url}/systemone",
            data=body,
            method="POST",
            headers={
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
                "Accept": "application/json",
                "User-Agent": "obsidian-sorter/0.1",
            },
        )

        try:
            with urllib.request.urlopen(req, timeout=self.timeout) as resp:
                data = json.loads(resp.read().decode("utf-8"))
        except urllib.error.HTTPError as e:
            detail = ""
            try:
                detail = e.read().decode("utf-8")[:500]
            except Exception:
                pass
            return JevResult(answers={}, error=f"HTTP {e.code}: {detail or e.reason}")
        except urllib.error.URLError as e:
            return JevResult(answers={}, error=f"network error: {e.reason}")
        except (json.JSONDecodeError, TimeoutError) as e:
            return JevResult(answers={}, error=f"{type(e).__name__}: {e}")

        answers: dict[str, Answer] = {}
        for name, raw in (data.get("answers") or {}).items():
            kind = raw.get("type", "choice")
            answers[name] = Answer(name=name, kind=kind, raw=raw)

        usage = data.get("usage") or {}
        return JevResult(
            answers=answers,
            input_tokens=int(usage.get("input_tokens") or 0),
            output_tokens=int(usage.get("output_tokens") or 0),
            model=data.get("model", self.model),
        )


class CostTracker:
    """Accumulates token usage and enforces a daily call cap."""

    def __init__(self, max_calls_per_run: int = 500):
        self.calls = 0
        self.input_tokens = 0
        self.output_tokens = 0
        self.max_calls_per_run = max_calls_per_run
        self.errors = 0

    def record(self, result: JevResult) -> None:
        self.calls += 1
        self.input_tokens += result.input_tokens
        self.output_tokens += result.output_tokens
        if not result.ok:
            self.errors += 1

    @property
    def cost_usd(self) -> float:
        return self.input_tokens * USD_PER_INPUT_TOKEN

    def budget_exceeded(self) -> bool:
        return self.calls >= self.max_calls_per_run

    def summary(self) -> dict[str, Any]:
        return {
            "calls": self.calls,
            "input_tokens": self.input_tokens,
            "output_tokens": self.output_tokens,
            "cost_usd": round(self.cost_usd, 6),
            "errors": self.errors,
        }