src/agent/llm.py

"""LLM client abstraction.

The agent loop talks to *an* LLM client; it doesn't care which. That keeps the loop
testable without weights and lets us swap runtimes:

* ``LlamaCppClient``  — the real, local-first runtime (🔌 + 🦙). Loads the fine-tuned
  MiniCPM GGUF and uses llama.cpp's function-calling. Import is guarded so the rest
  of the app works before weights are present.
* ``ScriptedClient``  — a deterministic stand-in that replays a fixed list of turns.
  Used by tests and the demo to exercise the full loop today.

A "turn" is either a set of tool calls or a final text answer.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Protocol


@dataclass
class ToolCall:
    name: str
    arguments: Dict[str, Any]
    id: str = ""


@dataclass
class AssistantTurn:
    text: Optional[str] = None
    tool_calls: List[ToolCall] = field(default_factory=list)

    @property
    def is_final(self) -> bool:
        return not self.tool_calls


class LLMClient(Protocol):
    def chat(self, messages: List[dict], tools: List[dict]) -> AssistantTurn:
        ...


class ScriptedClient:
    """Replays a fixed list of AssistantTurns — deterministic, no model required."""

    def __init__(self, turns: List[AssistantTurn]):
        self._turns = list(turns)
        self._i = 0
        self.seen_messages: List[List[dict]] = []

    def chat(self, messages: List[dict], tools: List[dict]) -> AssistantTurn:
        self.seen_messages.append(list(messages))
        if self._i >= len(self._turns):
            # Safety net: if the script runs dry, end the conversation.
            return AssistantTurn(text="(end of script)")
        turn = self._turns[self._i]
        self._i += 1
        return turn


class LlamaCppClient:
    """Local-first runtime via llama-cpp-python (activated once weights land).

    Loads the fine-tuned MiniCPM GGUF and exposes the same ``chat`` interface. The
    import and model load are lazy so importing this module never requires the
    dependency or the weights.
    """

    def __init__(self, model_path: str, n_ctx: int = 8192, **kwargs):
        try:
            from llama_cpp import Llama
        except ImportError as e:  # pragma: no cover - exercised only with the dep
            raise RuntimeError(
                "llama-cpp-python is not installed. `pip install llama-cpp-python` "
                "and point model_path at the fine-tuned MiniCPM GGUF."
            ) from e
        self._llm = Llama(model_path=model_path, n_ctx=n_ctx, **kwargs)

    def chat(self, messages: List[dict], tools: List[dict]) -> AssistantTurn:  # pragma: no cover
        resp = self._llm.create_chat_completion(
            messages=messages,
            tools=tools,
            tool_choice="auto",
        )
        choice = resp["choices"][0]["message"]
        calls = []
        for tc in choice.get("tool_calls") or []:
            import json
            fn = tc["function"]
            args = fn.get("arguments") or "{}"
            if isinstance(args, str):
                args = json.loads(args)
            calls.append(ToolCall(name=fn["name"], arguments=args, id=tc.get("id", "")))
        return AssistantTurn(text=choice.get("content"), tool_calls=calls)