gogo-telefon/gogo/agent/llm.py

"""LLMProvider interface (§4) — Anthropic implementation + scripted fake for tests.

The shape mirrors the Anthropic Messages API tool-use loop: the provider returns
text and/or tool_use blocks; the conversation loop executes tools and feeds
tool_result blocks back.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Protocol

from gogo.config import get_settings


@dataclass
class ToolUse:
    id: str
    name: str
    input: dict[str, Any]


@dataclass
class LLMResponse:
    text: str  # concatenated text blocks ("" if pure tool call)
    tool_uses: list[ToolUse] = field(default_factory=list)
    stop_reason: str = "end_turn"


class LLMProvider(Protocol):
    async def complete(
        self,
        *,
        system: str,
        messages: list[dict],
        tools: list[dict],
        model: str | None = None,
        max_tokens: int = 1024,
    ) -> LLMResponse: ...


class AnthropicLLM:
    def __init__(self, api_key: str | None = None):
        import anthropic

        self._client = anthropic.AsyncAnthropic(
            api_key=api_key or get_settings().anthropic_api_key or None
        )

    async def complete(
        self,
        *,
        system: str,
        messages: list[dict],
        tools: list[dict],
        model: str | None = None,
        max_tokens: int = 1024,
    ) -> LLMResponse:
        resp = await self._client.messages.create(
            model=model or get_settings().llm_model,
            system=system,
            messages=messages,
            tools=tools,
            max_tokens=max_tokens,
        )
        text_parts: list[str] = []
        tool_uses: list[ToolUse] = []
        for block in resp.content:
            if block.type == "text":
                text_parts.append(block.text)
            elif block.type == "tool_use":
                tool_uses.append(ToolUse(id=block.id, name=block.name, input=block.input))
        return LLMResponse(
            text="\n".join(text_parts).strip(),
            tool_uses=tool_uses,
            stop_reason=resp.stop_reason or "end_turn",
        )


class ScriptedLLM:
    """Deterministic fake: replays a fixed sequence of LLMResponses.

    Used to test the conversation loop machinery (tool dispatch, transcripts,
    guards) without a real model.
    """

    def __init__(self, responses: list[LLMResponse]):
        self._responses = list(responses)
        self.calls: list[dict] = []  # recorded kwargs for assertions

    async def complete(self, **kwargs) -> LLMResponse:
        import copy

        self.calls.append(copy.deepcopy(kwargs))  # snapshot: the history list mutates
        if not self._responses:
            return LLMResponse(text="Doviđenja!", stop_reason="end_turn")
        return self._responses.pop(0)


_llm: LLMProvider | None = None


def get_llm() -> LLMProvider:
    global _llm
    if _llm is None:
        _llm = AnthropicLLM()
    return _llm


def set_llm(provider: LLMProvider | None) -> None:
    """Test hook."""
    global _llm
    _llm = provider