src/compliment_forest/backends/text.py

from __future__ import annotations

import hashlib
import hmac
import json
from collections.abc import Callable
from typing import Any, Protocol
from urllib.parse import urlparse

import httpx
from huggingface_hub import InferenceClient

from compliment_forest.prompts import (
    author_messages,
    critic_messages,
    intake_messages,
    planner_messages,
)


def sign_modal_text_request(
    key: str,
    messages: list[dict[str, str]],
    *,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    seed: int,
    enable_thinking: bool,
) -> str:
    """HMAC the load-bearing request fields. Mirrors serve_minicpm.sign_text_request."""
    body = "\n".join(f"{m['role']}\x1f{m['content']}" for m in messages)
    payload = f"{body}\n{max_new_tokens}\n{temperature}\n{top_p}\n{seed}\n{int(enable_thinking)}"
    return hmac.new(key.encode(), payload.encode(), hashlib.sha256).hexdigest()


_DEMO_INTAKE_QUESTIONS: tuple[dict[str, object], ...] = (
    {
        "question": "Which part of this feels loudest right now?",
        "options": [
            "What might go wrong",
            "What other people will think",
            "The unknown right after this",
        ],
    },
    {
        "question": "When does this feel hardest?",
        "options": [
            "When I have time alone with it",
            "When I'm around the people involved",
            "Right before I have to act on it",
        ],
    },
    {
        "question": "Who else feels part of this for you?",
        "options": [
            "Just me",
            "One specific person",
            "A group I belong to",
        ],
    },
    {
        "question": "What feels most at stake if this does not go well?",
        "options": [
            "How I see myself",
            "How others see me",
            "Time or chances I cannot get back",
        ],
    },
    {
        "question": "What would a small win here look like?",
        "options": [
            "Just getting through it",
            "Doing one part well",
            "Knowing I was honest about how I felt",
        ],
    },
)


class TextBackend(Protocol):
    def next_intake_question(
        self,
        name: str,
        situation: str,
        history: list[dict[str, str]],
        *,
        rejected_questions: list[str] | None = None,
        seed: int = 3407,
    ) -> str: ...

    def plan(self, name: str, situation: str, *, seed: int = 3407) -> str: ...

    def author(
        self,
        name: str,
        situation: str,
        *,
        plan: dict[str, object],
        feedback: dict[str, str] | None = None,
        original: dict[str, object] | None = None,
        seed: int = 3407,
    ) -> str: ...

    def critic(
        self,
        name: str,
        situation: str,
        forest: dict[str, object],
        *,
        plan: dict[str, object],
        seed: int = 3407,
    ) -> str: ...


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

    # Each narration uses {situation} for at least one direct reference so the
    # demo backend can satisfy the orchestrator's grounding and source-phrase
    # checks without relying on the previous deterministic fallback path.
    _CHAPTERS = (
        (
            "arrive",
            "The Threshold in Mist",
            "The path opens quietly, asking nothing of you yet.",
            "honesty about uncertainty",
            (
                "{situation} sounds painful because it touches something that matters"
                " to you.\n\nYou do not need to argue with the feeling before looking"
                " at the problem."
            ),
            "Which part of this hurts most right now?",
            "I can name the part that hurts.",
            "a quiet path reaching a misty threshold with a soft opening ahead",
        ),
        (
            "steady",
            "The Lantern Beside the Map",
            "A small light is already keeping pace beside the path.",
            "careful attention",
            (
                "A worried thought can feel like a fact when it repeats often."
                "\n\nIt can help to ask what happened, then separate it from what"
                " fear predicts will happen next."
            ),
            "What is a fact here, and what is still a prediction?",
            "I can separate facts from predictions.",
            "a small lantern beside an unfolded map on a quiet wooden table",
        ),
        (
            "widen",
            "The Window and the Horizon",
            "The path lifts, and the room around the worry begins to feel taller.",
            "room for more than one outcome",
            (
                "One hard moment can have several explanations. It may show a skill"
                " to practice, a question to ask, or a standard that was too harsh."
                "\n\nYou can compare those options before choosing the worst one."
            ),
            "Which explanation fits the facts best?",
            "I can consider more than one explanation.",
            "an open window looking toward a broad horizon after light rain",
        ),
        (
            "step",
            "The First Stepping Stone",
            "The horizon settles back into the path beneath the feet.",
            "permission to move without certainty",
            (
                "You could write down the smallest question this situation raises."
                "\n\nThen choose one safe way to get more information before making"
                " a larger decision."
            ),
            "What question would give you the most useful information?",
            "I can take one useful next step.",
            (
                "an adult seen from behind pausing before the first stepping stone "
                "across a shallow stream"
            ),
        ),
        (
            "carry",
            "The Quiet Companion",
            (
                "The walk begins to turn back toward the ordinary day, "
                "but it does not turn back alone."
            ),
            "a simple plan",
            (
                "A simple rule can make the next hard moment easier to handle:"
                " name the fact, name the fear, then choose one useful action."
                "\n\nYou do not need to solve the whole problem at once."
            ),
            "Which part of that rule would help you first?",
            "I can name one fact and one next step.",
            "a small fox-shaped silhouette walking quietly beside an adult at sunset",
        ),
    )

    def next_intake_question(
        self,
        name: str,
        situation: str,
        history: list[dict[str, str]],
        *,
        rejected_questions: list[str] | None = None,
        seed: int = 3407,
    ) -> str:
        index = min(len(history), len(_DEMO_INTAKE_QUESTIONS) - 1)
        payload = dict(_DEMO_INTAKE_QUESTIONS[index])
        payload["rationale"] = "deterministic demo intake question"
        return json.dumps(payload)

    def plan(self, name: str, situation: str, *, seed: int = 3407) -> str:
        return json.dumps(
            {
                "faithful_summary": f"{name} is carrying this situation: {situation}",
                "fact_anchors": [{"source_phrase": situation, "meaning": situation}],
                "central_uncertainty": "What will happen next",
                "desired_direction": "Meet the next part with care and agency",
            }
        )

    def author(
        self,
        name: str,
        situation: str,
        *,
        plan: dict[str, object],
        feedback: dict[str, str] | None = None,
        original: dict[str, object] | None = None,
        seed: int = 3407,
    ) -> str:
        anchors = plan.get("fact_anchors") or []
        source_phrase = situation
        if isinstance(anchors, list) and anchors:
            first = anchors[0]
            if isinstance(first, dict):
                source_phrase = str(first.get("source_phrase") or situation)
        situation_phrase = f"“{source_phrase}”" if len(source_phrase) <= 80 else source_phrase
        clearings = [
            {
                "arc_role": arc_role,
                "source_phrase": source_phrase,
                "scene_title": scene_title,
                "scene_intro": scene_intro,
                "strength": strength,
                "narration": narration.format(situation=situation_phrase),
                "reflection": reflection,
                "spell": spell,
                "image_prompt": image_prompt,
            }
            for (
                arc_role,
                scene_title,
                scene_intro,
                strength,
                narration,
                reflection,
                spell,
                image_prompt,
            ) in self._CHAPTERS
        ]
        return json.dumps(
            {
                "forest_title": f"{name}'s Path Through the New",
                "proposed_strengths": [chapter[3] for chapter in self._CHAPTERS[:5]],
                "clearings": clearings,
            }
        )

    def critic(
        self,
        name: str,
        situation: str,
        forest: dict[str, object],
        *,
        plan: dict[str, object],
        seed: int = 3407,
    ) -> 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, follow_redirects=True)

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

    def next_intake_question(
        self,
        name: str,
        situation: str,
        history: list[dict[str, str]],
        *,
        rejected_questions: list[str] | None = None,
        seed: int = 3407,
    ) -> str:
        return self._complete(
            intake_messages(
                name,
                situation,
                history=history,
                rejected_questions=rejected_questions,
                seed=seed,
            ),
            max_tokens=8192,
            seed=seed,
            temperature=0.4,
        )

    def plan(
        self,
        name: str,
        situation: str,
        *,
        seed: int = 3407,
    ) -> str:
        return self._complete(
            planner_messages(name, situation, seed=seed),
            max_tokens=8192,
            seed=seed,
            temperature=0.2,
        )

    def author(
        self,
        name: str,
        situation: str,
        *,
        plan: dict[str, object],
        feedback: dict[str, str] | None = None,
        original: dict[str, object] | None = None,
        seed: int = 3407,
    ) -> str:
        return self._complete(
            author_messages(
                name,
                situation,
                plan=plan,
                feedback=feedback,
                original=original,
                seed=seed,
            ),
            max_tokens=8192,
            seed=seed,
            temperature=0.65,
        )

    def critic(
        self,
        name: str,
        situation: str,
        forest: dict[str, object],
        *,
        plan: dict[str, object],
        seed: int = 3407,
    ) -> str:
        return self._complete(
            critic_messages(name, situation, forest, plan=plan, seed=seed),
            max_tokens=8192,
            seed=seed,
            temperature=0.15,
        )


class HfInferenceTextBackend:
    """Hosted chat completion backend for the development Space."""

    def __init__(
        self,
        model: str = "openbmb/MiniCPM4.1-8B",
        *,
        client: Any | None = None,
        provider: str = "auto",
        timeout: float = 120,
    ) -> None:
        self.model = model
        self.client = client or InferenceClient(provider=provider, timeout=timeout)

    @staticmethod
    def _content(response: Any) -> str:
        content = response.choices[0].message.content
        if not isinstance(content, str):
            raise ValueError("hosted model response did not contain text")
        return content

    def next_intake_question(
        self,
        name: str,
        situation: str,
        history: list[dict[str, str]],
        *,
        rejected_questions: list[str] | None = None,
        seed: int = 3407,
    ) -> str:
        response = self.client.chat_completion(
            model=self.model,
            messages=intake_messages(
                name,
                situation,
                history=history,
                rejected_questions=rejected_questions,
                seed=seed,
            ),
            max_tokens=8192,
            temperature=0.4,
            top_p=0.9,
            seed=seed,
            response_format={"type": "json_object"},
            chat_template_kwargs={"enable_thinking": True},
        )
        return self._content(response)

    def plan(
        self,
        name: str,
        situation: str,
        *,
        seed: int = 3407,
    ) -> str:
        response = self.client.chat_completion(
            model=self.model,
            messages=planner_messages(name, situation, seed=seed),
            max_tokens=8192,
            temperature=0.2,
            top_p=0.9,
            seed=seed,
            response_format={"type": "json_object"},
            chat_template_kwargs={"enable_thinking": True},
        )
        return self._content(response)

    def author(
        self,
        name: str,
        situation: str,
        *,
        plan: dict[str, object],
        feedback: dict[str, str] | None = None,
        original: dict[str, object] | None = None,
        seed: int = 3407,
    ) -> str:
        response = self.client.chat_completion(
            model=self.model,
            messages=author_messages(
                name,
                situation,
                plan=plan,
                feedback=feedback,
                original=original,
                seed=seed,
            ),
            max_tokens=8192,
            temperature=0.82,
            top_p=0.92,
            seed=seed,
            response_format={"type": "json_object"},
            chat_template_kwargs={"enable_thinking": True},
        )
        return self._content(response)

    def critic(
        self,
        name: str,
        situation: str,
        forest: dict[str, object],
        *,
        plan: dict[str, object],
        seed: int = 3407,
    ) -> str:
        response = self.client.chat_completion(
            model=self.model,
            messages=critic_messages(name, situation, forest, plan=plan, seed=seed),
            max_tokens=8192,
            temperature=0.15,
            top_p=0.9,
            seed=seed,
            response_format={"type": "json_object"},
            chat_template_kwargs={"enable_thinking": True},
        )
        return self._content(response)


class TransformersTextBackend:
    """Routes prompts through an externally provided GPU generator closure."""

    def __init__(
        self,
        model: str = "openbmb/MiniCPM4.1-8B",
        *,
        generator: Callable[[list[dict[str, str]], dict[str, object]], str],
    ) -> None:
        self.model = model
        self._generate = generator

    def _call(
        self,
        messages: list[dict[str, str]],
        *,
        max_new_tokens: int,
        temperature: float,
        top_p: float,
        seed: int,
    ) -> str:
        params: dict[str, object] = {
            "max_new_tokens": max_new_tokens,
            "temperature": temperature,
            "top_p": top_p,
            "seed": seed,
            "chat_template_kwargs": {"enable_thinking": True},
        }
        return self._generate(messages, params)

    def next_intake_question(
        self,
        name: str,
        situation: str,
        history: list[dict[str, str]],
        *,
        rejected_questions: list[str] | None = None,
        seed: int = 3407,
    ) -> str:
        return self._call(
            intake_messages(
                name,
                situation,
                history=history,
                rejected_questions=rejected_questions,
                seed=seed,
            ),
            max_new_tokens=8192,
            temperature=0.4,
            top_p=0.9,
            seed=seed,
        )

    def plan(self, name: str, situation: str, *, seed: int = 3407) -> str:
        return self._call(
            planner_messages(name, situation, seed=seed),
            max_new_tokens=8192,
            temperature=0.2,
            top_p=0.9,
            seed=seed,
        )

    def author(
        self,
        name: str,
        situation: str,
        *,
        plan: dict[str, object],
        feedback: dict[str, str] | None = None,
        original: dict[str, object] | None = None,
        seed: int = 3407,
    ) -> str:
        return self._call(
            author_messages(
                name,
                situation,
                plan=plan,
                feedback=feedback,
                original=original,
                seed=seed,
            ),
            max_new_tokens=8192,
            temperature=0.82,
            top_p=0.92,
            seed=seed,
        )

    def critic(
        self,
        name: str,
        situation: str,
        forest: dict[str, object],
        *,
        plan: dict[str, object],
        seed: int = 3407,
    ) -> str:
        return self._call(
            critic_messages(name, situation, forest, plan=plan, seed=seed),
            max_new_tokens=8192,
            temperature=0.15,
            top_p=0.9,
            seed=seed,
        )


class ModalTextBackend:
    """Call the private-token Modal service that hosts MiniCPM4.1-8B."""

    def __init__(
        self,
        endpoint: str,
        signing_key: str,
        *,
        client: Any | None = None,
        timeout: float = 600,
    ) -> None:
        if urlparse(endpoint).scheme != "https":
            raise ValueError("modal text endpoint must use HTTPS")
        self.endpoint = endpoint
        self.signing_key = signing_key
        self.client = client or httpx.Client(timeout=timeout, follow_redirects=True)

    def _call(
        self,
        messages: list[dict[str, str]],
        *,
        max_new_tokens: int,
        temperature: float,
        top_p: float,
        seed: int,
        enable_thinking: bool = True,
    ) -> str:
        signature = sign_modal_text_request(
            self.signing_key,
            messages,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
            seed=seed,
            enable_thinking=enable_thinking,
        )
        try:
            response = self.client.post(
                self.endpoint,
                json={
                    "messages": messages,
                    "max_new_tokens": max_new_tokens,
                    "temperature": temperature,
                    "top_p": top_p,
                    "seed": seed,
                    "enable_thinking": enable_thinking,
                    "signature": signature,
                },
            )
            response.raise_for_status()
        except httpx.HTTPStatusError as error:
            raise ValueError(
                f"Modal text request failed with HTTP {error.response.status_code}"
            ) from error
        except httpx.RequestError as error:
            raise ValueError(f"Modal text request failed: {error}") from error
        payload = response.json()
        content = payload.get("content")
        if not isinstance(content, str):
            raise ValueError("Modal text response did not contain a string content")
        return content

    def next_intake_question(
        self,
        name: str,
        situation: str,
        history: list[dict[str, str]],
        *,
        rejected_questions: list[str] | None = None,
        seed: int = 3407,
    ) -> str:
        return self._call(
            intake_messages(
                name,
                situation,
                history=history,
                rejected_questions=rejected_questions,
                seed=seed,
            ),
            max_new_tokens=2048,
            temperature=0.4,
            top_p=0.9,
            seed=seed,
            enable_thinking=False,
        )

    def plan(self, name: str, situation: str, *, seed: int = 3407) -> str:
        return self._call(
            planner_messages(name, situation, seed=seed),
            max_new_tokens=8192,
            temperature=0.2,
            top_p=0.9,
            seed=seed,
            enable_thinking=False,
        )

    def author(
        self,
        name: str,
        situation: str,
        *,
        plan: dict[str, object],
        feedback: dict[str, str] | None = None,
        original: dict[str, object] | None = None,
        seed: int = 3407,
    ) -> str:
        return self._call(
            author_messages(
                name,
                situation,
                plan=plan,
                feedback=feedback,
                original=original,
                seed=seed,
            ),
            max_new_tokens=8192,
            temperature=0.82,
            top_p=0.92,
            seed=seed,
            enable_thinking=False,
        )

    def critic(
        self,
        name: str,
        situation: str,
        forest: dict[str, object],
        *,
        plan: dict[str, object],
        seed: int = 3407,
    ) -> str:
        return self._call(
            critic_messages(name, situation, forest, plan=plan, seed=seed),
            max_new_tokens=8192,
            temperature=0.15,
            top_p=0.9,
            seed=seed,
            enable_thinking=False,
        )