tests/test_llm_client.py

"""Tests for ``cortex.llm_client.LLMClient`` (Session 7a).

Per the approved spec:
  - Cumulative per-caller token counter (sum across all calls).
  - Source: API ``response.usage.prompt_tokens + .completion_tokens``;
    if ``usage`` is missing, increment by 0 and warn (no local tokenizer).
  - Caller IDs: short colon-separated strings ("inference:t3", "b1:t3",
    "cortex:epi:planner:t3"). Passed explicitly per call, not thread-local.
  - ``reset_counters`` is harness-driven: never auto-reset by the client.

The OpenAI SDK is stubbed via dependency injection: ``LLMClient`` accepts
an optional ``client`` kwarg in ``__init__`` so tests can pass a fake
without monkey-patching the global OpenAI module.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import List, Optional

from cortex.llm_client import ChatMessage, ChatResponse, LLMClient

# ============================================================================
# Stub OpenAI client — minimal subset of the SDK surface used by LLMClient
# ============================================================================


@dataclass
class _StubUsage:
    prompt_tokens: int
    completion_tokens: int


@dataclass
class _StubMessage:
    content: str


@dataclass
class _StubChoice:
    message: _StubMessage
    finish_reason: str = "stop"


@dataclass
class _StubCompletion:
    choices: List[_StubChoice]
    usage: Optional[_StubUsage]


class _StubChatCompletions:
    def __init__(self, responses: List[_StubCompletion]) -> None:
        self._responses = list(responses)
        self.calls: List[dict] = []

    def create(self, **kwargs) -> _StubCompletion:
        self.calls.append(kwargs)
        if not self._responses:
            raise RuntimeError("stub exhausted - test wants more responses than configured")
        return self._responses.pop(0)


class _StubChat:
    def __init__(self, completions: _StubChatCompletions) -> None:
        self.completions = completions


class _StubOpenAI:
    """Quacks-like-OpenAI: ``client.chat.completions.create(...)``."""

    def __init__(self, responses: List[_StubCompletion]) -> None:
        self.chat = _StubChat(_StubChatCompletions(responses))


def _resp(
    content: str, prompt_tokens: int, completion_tokens: int, usage: bool = True
) -> _StubCompletion:
    return _StubCompletion(
        choices=[_StubChoice(message=_StubMessage(content=content))],
        usage=_StubUsage(prompt_tokens, completion_tokens) if usage else None,
    )


# ============================================================================
# Tests
# ============================================================================


def test_token_counter_increments_per_call() -> None:
    """Single caller, two calls: counter is the cumulative sum of usages."""
    stub = _StubOpenAI(
        [
            _resp("ok-1", prompt_tokens=50, completion_tokens=30),
            _resp("ok-2", prompt_tokens=70, completion_tokens=10),
        ]
    )
    client = LLMClient(model="stub-model", client=stub)

    client.chat(caller_id="b1:t1", messages=[ChatMessage(role="user", content="hi")])
    assert client.tokens_used_for("b1:t1") == 80, "after 1 call: 50+30 = 80"

    client.chat(caller_id="b1:t1", messages=[ChatMessage(role="user", content="hi")])
    assert client.tokens_used_for("b1:t1") == 160, "after 2 calls: 80 + 70+10 = 160"


def test_token_counter_isolates_callers() -> None:
    """Two distinct caller_ids accumulate independently."""
    stub = _StubOpenAI(
        [
            _resp("a", prompt_tokens=100, completion_tokens=20),
            _resp("b", prompt_tokens=5, completion_tokens=5),
        ]
    )
    client = LLMClient(model="stub-model", client=stub)

    client.chat(caller_id="inference:t1", messages=[ChatMessage(role="user", content="x")])
    client.chat(caller_id="b1:t1", messages=[ChatMessage(role="user", content="y")])

    assert client.tokens_used_for("inference:t1") == 120
    assert client.tokens_used_for("b1:t1") == 10
    assert client.tokens_used_for("never:called") == 0, "unknown caller_id reads as 0, not KeyError"


def test_reset_counters_zeroes_callers() -> None:
    """``reset_counters()`` with no args clears all; with a prefix clears
    only matching keys. No automatic reset happens — the client never
    zeroes counters on its own."""
    stub = _StubOpenAI(
        [
            _resp("a", 30, 10),
            _resp("b", 50, 10),
            _resp("c", 5, 5),
        ]
    )
    client = LLMClient(model="stub-model", client=stub)

    client.chat(caller_id="inference:t1", messages=[ChatMessage(role="user", content="x")])
    client.chat(caller_id="b1:t1", messages=[ChatMessage(role="user", content="y")])
    client.chat(caller_id="b1:t2", messages=[ChatMessage(role="user", content="z")])
    assert client.tokens_used_for("inference:t1") == 40
    assert client.tokens_used_for("b1:t1") == 60
    assert client.tokens_used_for("b1:t2") == 10

    # Prefix-scoped reset: only b1:* keys zero, inference:* survives.
    client.reset_counters(caller_id_prefix="b1:")
    assert client.tokens_used_for("b1:t1") == 0
    assert client.tokens_used_for("b1:t2") == 0
    assert client.tokens_used_for("inference:t1") == 40, (
        "prefix reset must not touch unrelated caller_ids"
    )

    # Full reset.
    client.reset_counters()
    assert client.tokens_used_for("inference:t1") == 0


def test_chat_returns_typed_response_and_handles_missing_usage() -> None:
    """``ChatResponse`` carries content + finish_reason + token fields.

    When a provider returns no ``usage`` block, the counter must NOT
    fall over — increment by 0 and surface zero token fields, so callers
    can still read the content. (No local tokenizer fallback in MVP.)
    """
    stub = _StubOpenAI(
        [
            _resp("hello world", prompt_tokens=12, completion_tokens=5),
            _resp("no usage", prompt_tokens=0, completion_tokens=0, usage=False),
        ]
    )
    client = LLMClient(model="stub-model", client=stub)

    r1 = client.chat(
        caller_id="inference:t1",
        messages=[ChatMessage(role="user", content="hello")],
    )
    assert isinstance(r1, ChatResponse)
    assert r1.content == "hello world"
    assert r1.finish_reason == "stop"
    assert r1.prompt_tokens == 12
    assert r1.completion_tokens == 5
    assert client.tokens_used_for("inference:t1") == 17

    # Provider with no usage block: don't crash.
    r2 = client.chat(
        caller_id="inference:t2",
        messages=[ChatMessage(role="user", content="hi")],
    )
    assert r2.content == "no usage"
    assert r2.prompt_tokens == 0
    assert r2.completion_tokens == 0
    assert client.tokens_used_for("inference:t2") == 0


def test_chat_passes_messages_and_temperature_to_sdk() -> None:
    """Messages and temperature must reach the SDK call verbatim."""
    stub = _StubOpenAI([_resp("ok", 1, 1)])
    client = LLMClient(model="stub-model", client=stub, temperature=0.7, max_tokens=128)
    client.chat(
        caller_id="b1:t1",
        messages=[
            ChatMessage(role="system", content="sys"),
            ChatMessage(role="user", content="usr"),
        ],
    )
    call = stub.chat.completions.calls[0]
    assert call["model"] == "stub-model"
    assert call["temperature"] == 0.7
    assert call["max_tokens"] == 128
    assert call["messages"] == [
        {"role": "system", "content": "sys"},
        {"role": "user", "content": "usr"},
    ]