src/compliment_forest/backends/text.py

from __future__ import annotations

import json
from typing import Protocol
from urllib.parse import urlparse

import httpx

from compliment_forest.prompts import author_messages, critic_messages


class TextBackend(Protocol):
    def author(
        self,
        name: str,
        situation: str,
        *,
        feedback: dict[str, str] | None = None,
        original: dict[str, object] | None = None,
    ) -> str: ...

    def critic(self, name: str, situation: str, forest: dict[str, object]) -> str: ...


class DemoTextBackend:
    """Deterministic development backend with the same model contract."""

    _CLEARINGS = (
        (
            "The Patient Fox",
            "patience under pressure",
            "a gentle russet fox sitting in a mossy clearing, soft kind eyes",
            "I am allowed to learn.",
        ),
        (
            "The Listening Owl",
            "careful perspective",
            "a round tawny owl resting on a low branch, attentive kind eyes",
            "I can listen before I leap.",
        ),
        (
            "The Brave Snail",
            "quiet courage",
            "a tiny snail crossing a fern frond, softly glowing spiral shell",
            "I make progress at my pace.",
        ),
        (
            "The Steady Deer",
            "steadiness through change",
            "a small deer standing in morning mist, calm gentle expression",
            "I can meet one moment.",
        ),
        (
            "The Clear-Voiced Wren",
            "honest self-expression",
            "a tiny wren singing beside dusty rose wildflowers",
            "I can speak gently and clearly.",
        ),
    )

    def author(
        self,
        name: str,
        situation: str,
        *,
        feedback: dict[str, str] | None = None,
        original: dict[str, object] | None = None,
    ) -> str:
        clearings = []
        for creature, strength, image_prompt, spell in self._CLEARINGS:
            clearings.append(
                {
                    "creature": creature,
                    "strength": strength,
                    "line": (
                        f"{situation.rstrip('.')} may feel uncertain. Your {strength} "
                        "lets you respond without pretending the hard part is easy."
                    ),
                    "reflection": (
                        "What becomes possible when you ask for one honest next step "
                        "instead of a perfect answer?"
                    ),
                    "spell": spell,
                    "image_prompt": image_prompt,
                }
            )
        payload = {
            "forest_title": f"{name}'s Path Through the New",
            "proposed_strengths": [item[1] for item in self._CLEARINGS],
            "clearings": clearings,
        }
        return json.dumps(payload)

    def critic(self, name: str, situation: str, forest: dict[str, object]) -> str:
        count = min(len(forest.get("clearings", [])), 5)
        return json.dumps(
            {
                "keep_indices": list(range(count)),
                "revise_indices": [],
                "reasons": {},
            }
        )


class LlamaCppTextBackend:
    """OpenAI-compatible client restricted to a local llama.cpp server."""

    def __init__(
        self,
        base_url: str = "http://127.0.0.1:8080",
        model: str = "compliment-forest-minicpm5-1b",
        timeout: float = 90,
    ) -> None:
        parsed = urlparse(base_url)
        if parsed.hostname not in {"127.0.0.1", "localhost", "::1"}:
            raise ValueError("llama.cpp base URL must resolve to the local machine")
        self.base_url = base_url.rstrip("/")
        self.model = model
        self.client = httpx.Client(timeout=timeout)

    def _complete(self, messages: list[dict[str, str]], max_tokens: int) -> str:
        response = self.client.post(
            f"{self.base_url}/v1/chat/completions",
            json={
                "model": self.model,
                "messages": messages,
                "temperature": 0.35,
                "top_p": 0.9,
                "max_tokens": max_tokens,
                "response_format": {"type": "json_object"},
                "chat_template_kwargs": {"enable_thinking": False},
            },
        )
        response.raise_for_status()
        payload = response.json()
        return payload["choices"][0]["message"]["content"]

    def author(
        self,
        name: str,
        situation: str,
        *,
        feedback: dict[str, str] | None = None,
        original: dict[str, object] | None = None,
    ) -> str:
        return self._complete(
            author_messages(name, situation, feedback=feedback, original=original),
            max_tokens=1700,
        )

    def critic(self, name: str, situation: str, forest: dict[str, object]) -> str:
        return self._complete(critic_messages(name, situation, forest), max_tokens=800)