diff --git a/Architecture.md b/Architecture.md index 00617d1e..6e1be239 100644 --- a/Architecture.md +++ b/Architecture.md @@ -185,24 +185,27 @@ Thin wrapper over a ChromaDB Collection: - **`delete_by_doc_id()`** — removes all chunks for a document via metadata WHERE clause - **`get_stats()`** — returns chunk count, backend name, collection name -### `src/llm_handler.py` — LLM Provider Routing +### `src/llm_handler/` — LLM Provider Routing -Auto-detects provider from model name prefix: +`LLMHandler` (in `src/llm_handler/__init__.py`) auto-detects the provider from the model-name prefix and selects **one adapter** at construction: -| Prefix | Provider | API Key Env Var | -|--------|----------|-----------------| -| `gpt*`, `o1*`, `o3*` | OpenAI | `OPENAI_API_KEY` | -| `claude*` | Anthropic | `ANTHROPIC_API_KEY` | -| `glm*` | Zhipu AI (OpenAI-compatible) | `GLM_API_KEY` | -| everything else | Ollama (localhost:11434) | none | +| Prefix | Provider | Adapter | API Key Env Var | +|--------|----------|---------|-----------------| +| `gpt*`, `o1*`, `o3*` | OpenAI | `OpenAICompatibleAdapter` | `OPENAI_API_KEY` | +| `claude*` | Anthropic | `AnthropicAdapter` | `ANTHROPIC_API_KEY` | +| `glm*` | Zhipu AI (OpenAI-compatible) | `OpenAICompatibleAdapter` | `GLM_API_KEY` | +| everything else | Ollama (localhost:11434) | `OllamaAdapter` | none | -All providers are optional imports with graceful fallback to a dummy response generator. +Each provider lives behind a `ProviderAdapter` (`src/llm_handler/adapters/`) whose SDK client is **injected** via a zero-arg `client_factory`, so every provider path is unit-testable with a fake (`tests/test_llm_adapters.py`). Adapters return `GenerationResult(text, usage)`; streaming yields text chunks then a terminal `Usage`. Usage is provider-reported where the SDK supplies it, adapter-counted otherwise. See [ADR 0002](docs/adr/0002-provider-adapters.md). -Public API surfaces: +`LLMHandler` owns provider selection, the single-prompt → messages translation, and the fallback: only `ProviderUnavailableError` (missing SDK, missing GLM key, Ollama connection refused) routes to the `DummyAdapter`; real API errors propagate. + +Public API surfaces (unchanged for callers): - `generate()` / `stream_response()` — single prompt string - `generate_messages()` / `stream_messages()` — OpenAI-style messages list (for multi-turn chat) +- `generate_with_usage()` — returns `(text, prompt_tokens, completion_tokens)` from provider-reported usage -GPT-5 family and o-series models use `max_completion_tokens` instead of `max_tokens` and omit the `temperature` parameter (constrained to default). +GPT-5 family and o-series models use `max_completion_tokens` instead of `max_tokens` and omit the `temperature` parameter (constrained to default) — handled inside `OpenAICompatibleAdapter`. ### `src/database.py` — SQLite/SQLModel @@ -452,7 +455,7 @@ The `src/eval/` package provides a reproducible evaluation system over labeled g | Module | Responsibility | |--------|----------------| | `src/eval/schemas.py` | Pydantic contracts: `EvalQuestion`, `EvalResult`, `AggregatedMetric`, `RunMetadata`, `MetricDelta`, `CompareResult`. | -| `src/eval/pricing.py` | Hard-coded model price table + `cost_usd()` helper. | +| `src/telemetry/pricing.py`, `src/telemetry/tokens.py` | **Core** (not eval): model price table + `cost_usd()`, and `count_tokens()`. Imported by both production telemetry and the eval harness (see [ADR 0003](docs/adr/0003-telemetry-ownership.md)). | | `src/eval/statistics.py` | `bootstrap_ci()` and `paired_permutation_test()` for run-level confidence intervals and two-run significance testing. | | `src/eval/metrics/retrieval.py` | Recall@k, MRR@k, nDCG@k over `(gold_chunk_ids, retrieved_chunk_ids)`. | | `src/eval/metrics/operational.py` | Per-stage latency p50/p95/p99, cost, token aggregation. | @@ -462,7 +465,7 @@ The `src/eval/` package provides a reproducible evaluation system over labeled g | `src/eval/datasets/ml_papers.py` | Hand-labeled dev set loader + manifest SHA-256 verification. | | `src/eval/config.py` | YAML-loaded `EvalConfig`. | | `src/eval/storage.py` | Run-directory CRUD over `eval_runs//`. | -| `src/eval/pipeline_factory.py` + `src/eval/_telemetry.py` | Builds an isolated RAG pipeline per (config, dataset) using ephemeral Chroma. | +| `src/eval/pipeline_factory.py` | Builds an isolated RAG pipeline per (config, dataset) using ephemeral Chroma. Token counting and pricing come from `src/telemetry/` (core), not the eval package. | | `src/eval/aggregator.py` | Per-dataset + combined `AggregatedMetric` rows from per-question results. | | `src/eval/runner.py` | Orchestrates `git_sha`, ingest, query+score loop, aggregation, persistence. | | `src/eval/compare.py` | Two-run diff with paired permutation tests + per-question regressions/wins. | @@ -520,7 +523,7 @@ The system exports per-stage spans for every chat query via OpenTelemetry to [Ar ### Telemetry Payload -The same numbers are returned to the client as a `StageTelemetry` Pydantic model (`src/api/schemas/telemetry.py`): +The same numbers are returned to the client as a `StageTelemetry` Pydantic model (`src/api/schemas/telemetry.py`). Token counts are the **provider-reported usage** for the answer pass (from the LLM adapter's `GenerationResult` / streaming terminal `Usage`), with the adapter's local count as fallback — never a reconstructed prompt. Telemetry covers the answer pass only, not the chain-of-thought reasoning pass. - REST `POST /api/query` — includes a `telemetry` field in the response JSON. - WebSocket `/api/chat` — emits a final `{"type": "telemetry", "content": {...}}` event after the existing `done` event. diff --git a/CONTEXT.md b/CONTEXT.md new file mode 100644 index 00000000..7d0a11c0 --- /dev/null +++ b/CONTEXT.md @@ -0,0 +1,33 @@ +# CONTEXT — domain glossary + +The vocabulary of this codebase, so names stay consistent across modules, tests, +and ADRs. New module names enter here as the architecture-deepening spec +([issue #16](https://github.com/elkaix/rag-document-qa/issues/16)) lands them. + +Design terms use the `/codebase-design` vocabulary: **module** (a unit with an +interface hiding behaviour), **interface** (the public surface), **depth** (much +behaviour behind a small interface), **seam** (a boundary you can substitute at), +**adapter** (a module presenting one interface over another), **leverage**, +**locality**. + +## Modules & types + +- **ProviderAdapter** — the interface (Protocol) one LLM provider hides behind: + `generate(messages) -> GenerationResult` and `stream(messages)` yielding text + chunks then a terminal `Usage`. Implementations: `OpenAICompatibleAdapter` + (OpenAI + GLM), `AnthropicAdapter`, `OllamaAdapter`, `DummyAdapter`. Each is + constructed with an injected `client_factory` so it is testable with a fake. + See [ADR 0002](docs/adr/0002-provider-adapters.md). +- **Usage** — a value object of `(prompt_tokens, completion_tokens)` for one + generation call. Provider-reported where the SDK returns it, adapter-counted + otherwise. +- **GenerationResult** — the `(text, usage)` pair returned by a non-streaming + generation, so callers never rebuild a prompt to estimate what was billed. +- **LLMHandler** — owns provider *selection* (from the model-name prefix), the + single-prompt → messages translation, and the unconfigured-provider fallback + to `DummyAdapter`. Delegates all SDK work to one selected `ProviderAdapter`. +- **telemetry** (`src/telemetry/`) — core token counting (`count_tokens`) and + cost pricing (`cost_usd`, `MODEL_PRICES`). Owned by the core so both production + telemetry assembly and the eval harness use one source of truth; the eval + package imports from here, never the reverse. See + [ADR 0003](docs/adr/0003-telemetry-ownership.md). diff --git a/docs/adr/0002-provider-adapters.md b/docs/adr/0002-provider-adapters.md new file mode 100644 index 00000000..f3bbe830 --- /dev/null +++ b/docs/adr/0002-provider-adapters.md @@ -0,0 +1,63 @@ +# ADR 0002 — LLM provider adapters with injected clients + +- **Status:** Accepted +- **Sequencing:** Step 2 of the RAG architecture deepening spec ([issue #16](https://github.com/elkaix/rag-document-qa/issues/16)); resolves [#10](https://github.com/elkaix/rag-document-qa/issues/10). +- **Date:** 2026-07-12 + +> ADR numbering follows issue #16's sequencing steps. Step 1 (deletions + +> vector-store lookup, PR #17) landed without a written ADR; this is the first. + +## Context + +`src/llm_handler.py` was an 812-line class that repeated the same three-branch +provider dispatch **four times** (`generate`, `stream_response`, +`generate_messages`, `stream_messages`) and built every provider client from +module globals. Two consequences: + +1. **Untestable providers.** The only way to fake a provider was to monkeypatch a + module global (`openai.OpenAI`). That worked for OpenAI in CI but left + **Anthropic and Ollama with zero coverage** — there was no seam to inject a + fake at. +2. **No usage at the source.** Generation returned only `str`, so callers that + needed token counts (cost telemetry) rebuilt the prompt and re-tokenized it — + counting a *reconstruction*, not what the provider actually billed. + +## Decision + +Introduce a **`ProviderAdapter` Protocol** (`generate` -> `GenerationResult`, +`stream` -> text chunks then a terminal `Usage`) with one adapter per provider: + +- **`OpenAICompatibleAdapter`** serves both OpenAI and GLM. The token-parameter + and temperature quirks live here; GLM differs only by the injected client's + base URL, so it is not a second code path. +- **`AnthropicAdapter`** owns the system-message split. +- **`OllamaAdapter`** owns the `/api/chat` payload shape. +- **`DummyAdapter`** is the always-available fallback. + +**Clients are injected** via a zero-arg `client_factory` per adapter, so every +provider is unit-testable with a fake (see `tests/test_llm_adapters.py`) — no +module-global monkeypatching. `LLMHandler` selects one adapter at construction; +the four dispatch tables collapse to that single selection. Its public interface +is unchanged; prompt variants translate to the messages form the adapters speak. + +**Usage is provider-reported where available, adapter-counted as fallback** +(`base.counted_usage`). Streaming reports usage in its terminal `Usage` event. + +## Consequences + +- Anthropic and Ollama gain their first real unit coverage; OpenAI/GLM quirks are + asserted directly against a fake client. +- **Preserved behavior (deliberately):** only `ProviderUnavailableError` (missing + SDK / missing GLM key / Ollama connection refused) triggers the dummy fallback; + real API errors (auth, quota, malformed) still propagate. A missing *OpenAI* + key is left to the SDK to reject, exactly as before — the asymmetry with GLM's + explicit pre-check is intentional and preserved. +- **Consolidation:** the single-prompt Ollama `/api/generate` endpoint is gone — + every call is translated to a messages list first, so `/api/chat` covers both. + Semantically equivalent; the wire format differs. +- In this step, `LLMHandler`'s public `stream_*` methods still yield text only + (they drop the adapter's terminal `Usage`), keeping the backend untouched. + Step 3 surfaces that usage to the telemetry layer. +- The injected-client type is intentionally loose (`object`) at the SDK seam — + the adapters call methods on external SDK objects we do not own; no `Any` and + no `# type: ignore` are used. diff --git a/docs/adr/0003-telemetry-ownership.md b/docs/adr/0003-telemetry-ownership.md new file mode 100644 index 00000000..61ec37ec --- /dev/null +++ b/docs/adr/0003-telemetry-ownership.md @@ -0,0 +1,57 @@ +# ADR 0003 — Telemetry ownership and usage-driven cost + +- **Status:** Accepted +- **Sequencing:** Step 3 of the RAG architecture deepening spec ([issue #16](https://github.com/elkaix/rag-document-qa/issues/16)); resolves [#12](https://github.com/elkaix/rag-document-qa/issues/12). +- **Date:** 2026-07-12 + +## Context + +Two coupled problems in the cost-telemetry path: + +1. **Inverted dependency.** Token counting (`count_tokens`) and pricing + (`cost_usd`, `MODEL_PRICES`) lived in `src/eval/` (`_telemetry.py`, + `pricing.py`), yet **production** imported them — `src/backend.py` and + `src/llm_handler` reached into the eval package for core utilities. +2. **Cost computed from a reconstruction.** `RAGBackend` did not have the tokens + the provider billed, so it **rebuilt the prompt string** (a third copy of the + system + user prompt) purely to re-tokenize it with `count_tokens`. The number + shown to users was an estimate of a *copy*, in four duplicated assembly sites + (sync query, streaming conversation branch, streaming single-turn branch, and + `LLMHandler.generate_with_usage`). + +## Decision + +**Move token/pricing utilities to a core `src/telemetry/` package** (`tokens.py`, +`pricing.py`). The eval package now imports from core (`from src.telemetry ...`), +never the reverse; `src/eval/__init__.py` re-exports the pricing symbols so its +public API is unchanged. The old `src/eval/_telemetry.py` and `src/eval/pricing.py` +are deleted (moved with `git mv`, not shimmed). + +**Consume provider-reported usage** (from ADR 0002's adapters): + +- The **synchronous** path calls `LLMHandler.generate_with_usage(...)`, which + returns the answer plus real `(prompt_tokens, completion_tokens)`. The + reconstructed prompt and its `count_tokens` calls are gone; the now-orphaned + `generate_with_context` is deleted. +- The **streaming** path reads the `Usage` carried on the stream's terminal event + (`LLMHandler.stream_response` / `stream_messages` now yield `str | Usage`). The + backend discriminates with `isinstance(item, Usage)`. + +`cost_usd` is still computed in the backend from those counts. + +## Consequences + +- Production no longer imports from the eval package — a clean dependency + direction (verified: `grep` finds no `src.eval` import in backend/llm_handler). +- The cost footer reflects what the provider actually billed (or the adapter's + local count when a provider omits usage), not a re-tokenized copy of a + reconstructed prompt. +- **Behavior preserved (deliberately):** telemetry still covers the **answer pass + only** — the streaming reasoning pass's terminal `Usage` is discarded, keeping + the per-query cost display to one model's spend, exactly as before. The + no-documents early-return still emits zero telemetry. `StageTelemetry`'s shape + is unchanged, so the API and frontend are untouched. `test_backend_telemetry.py` + passes unchanged as the behavior-preservation proof. +- The four duplicated token-assembly sites collapse to two usage reads (one per + path); further consolidation of the answer prompt itself is step 4's + QueryEngine work, not this step. diff --git a/src/backend.py b/src/backend.py index 228e3adf..12a46223 100644 --- a/src/backend.py +++ b/src/backend.py @@ -51,14 +51,13 @@ TOP_K_RESULTS, ) from .document_loader import DocumentLoader, TextChunker -from .eval._telemetry import count_tokens -from .eval.pricing import cost_usd from .evaluation import ( evaluate_answer_relevancy, evaluate_context_precision, evaluate_faithfulness, ) -from .llm_handler import LLMHandler +from .llm_handler import LLMHandler, Usage +from .telemetry.pricing import cost_usd from .models.conversation import Conversation from .models.document import DocumentRecord from .models.evaluation import MessageEvaluation @@ -387,9 +386,9 @@ def query_with_telemetry( if model and model != self.llm.model: handler = LLMHandler(model=model) - # WHY reconstruct prompt strings here: generate_with_context() builds - # these internally before calling self.generate(). We reproduce them to - # count exactly the tokens that were billed, not a proxy string. + # Build the answer prompt (system + user). Passing it through + # generate_with_usage means telemetry counts the exact tokens the + # provider billed — no reconstruction of a proxy string. answer_system_prompt = ( "You are a helpful assistant. Answer the user's question based solely on the " "provided context. If the context does not contain enough information, say so." @@ -400,7 +399,9 @@ def query_with_telemetry( t_generate_start = time.perf_counter() with tracer.start_as_current_span("rag.generate") as generate_span: generate_span.set_attribute("model", handler.model) - answer = handler.generate_with_context(question, context) + answer, prompt_tokens, completion_tokens = handler.generate_with_usage( + answer_user_prompt, system_prompt=answer_system_prompt + ) generate_span.set_attribute("answer_len", len(answer)) generate_ms = (time.perf_counter() - t_generate_start) * 1000 @@ -421,11 +422,8 @@ def query_with_telemetry( confidence = max(0.0, min(1.0, sum(top_scores) / len(top_scores))) # ---- Telemetry assembly --------------------------------------------- - # WHY count system_prompt + user_prompt together: together they form - # the full input that the provider billed as prompt tokens. - prompt_text = answer_system_prompt + "\n" + answer_user_prompt - prompt_tokens = count_tokens(prompt_text, handler.model) - completion_tokens = count_tokens(answer, handler.model) + # Token counts come from the provider-reported usage above (or the + # adapter's local count fallback) — never a reconstructed prompt. total_cost = cost_usd(handler.model, prompt_tokens, completion_tokens) telemetry = StageTelemetry( @@ -582,10 +580,14 @@ def stream_query( ) try: - for token in self.reasoning_llm.stream_response( + for item in self.reasoning_llm.stream_response( reasoning_user, system_prompt=reasoning_system ): - yield ("reasoning", token) + # The reasoning pass reports its own usage, but telemetry covers + # the answer pass only — so drop the reasoning terminal Usage. + if isinstance(item, Usage): + continue + yield ("reasoning", item) except Exception as exc: # PATTERN: Reasoning is best-effort — a failure here must not block # the final answer. Log, emit a terse status, continue to the answer. @@ -612,8 +614,10 @@ def stream_query( ) user_prompt = f"Context:\n{context}\n\nQuestion: {question}\n\nAnswer:" - # Accumulate full response for persistence and token counting + # Accumulate full response for persistence; capture the answer pass's + # reported usage (yielded as the stream's terminal event) for telemetry. full_response: list[str] = [] + answer_usage: Usage | None = None t_generate_start = time.perf_counter() @@ -635,9 +639,12 @@ def stream_query( with tracer.start_as_current_span("rag.generate") as generate_span: generate_span.set_attribute("model", handler.model) generate_span.set_attribute("has_conversation", True) - for token in handler.stream_messages(messages): - full_response.append(token) - yield ("token", token) + for item in handler.stream_messages(messages): + if isinstance(item, Usage): + answer_usage = item + continue + full_response.append(item) + yield ("token", item) generate_span.set_attribute("answer_len", sum(len(t) for t in full_response)) generate_ms = (time.perf_counter() - t_generate_start) * 1000 @@ -661,41 +668,38 @@ def stream_query( "conversation_id": conversation_id, }) - # WHY count full messages list for prompt: the multi-turn path sends - # the sliding window + system prompt + user turn as one request. - # We join all message content to estimate the billed input tokens. - prompt_text = "\n".join(m["content"] for m in messages) - else: # No conversation — simple single-turn streaming with tracer.start_as_current_span("rag.generate") as generate_span: generate_span.set_attribute("model", handler.model) generate_span.set_attribute("has_conversation", False) - for token in handler.stream_response(user_prompt, system_prompt=system_prompt): - full_response.append(token) - yield ("token", token) + for item in handler.stream_response(user_prompt, system_prompt=system_prompt): + if isinstance(item, Usage): + answer_usage = item + continue + full_response.append(item) + yield ("token", item) generate_span.set_attribute("answer_len", sum(len(t) for t in full_response)) generate_ms = (time.perf_counter() - t_generate_start) * 1000 yield ("done", {"sources": sources}) - prompt_text = system_prompt + "\n" + user_prompt - # ---- Telemetry assembly (after done, additive) ----------------------- # WHY after done: the done event is what the client waits for to show - # sources. Telemetry is a secondary signal — emit it last so done's - # latency is not affected by token-counting arithmetic. - answer_text = "".join(full_response) - prompt_tokens = count_tokens(prompt_text, handler.model) - completion_tokens = count_tokens(answer_text, handler.model) - total_cost = cost_usd(handler.model, prompt_tokens, completion_tokens) + # sources. Telemetry is a secondary signal — emit it last. + # + # Usage is the answer pass's reported figure (provider-reported where the + # provider supplies it, adapter-counted otherwise), carried on the + # stream's terminal event — no prompt reconstruction. + usage = answer_usage or Usage(prompt_tokens=0, completion_tokens=0) + total_cost = cost_usd(handler.model, usage.prompt_tokens, usage.completion_tokens) yield ("telemetry", StageTelemetry( retrieve_ms=round(retrieve_ms, 2), generate_ms=round(generate_ms, 2), - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, + prompt_tokens=usage.prompt_tokens, + completion_tokens=usage.completion_tokens, cost_usd=total_cost, ).model_dump()) diff --git a/src/eval/__init__.py b/src/eval/__init__.py index ba197058..9b4773e6 100644 --- a/src/eval/__init__.py +++ b/src/eval/__init__.py @@ -10,7 +10,6 @@ from src.eval.compare import compare_runs from src.eval.config import EvalConfig, load_config -from src.eval.pricing import MODEL_PRICES, ModelPrice, cost_usd from src.eval.runner import EvalRunner from src.eval.schemas import ( AggregatedMetric, @@ -22,6 +21,9 @@ ) from src.eval.statistics import bootstrap_ci, paired_permutation_test from src.eval.storage import list_runs, load_run, save_run +# Pricing moved to the core telemetry package; re-exported here so the eval +# package's public API (`from src.eval import cost_usd`) is preserved. +from src.telemetry.pricing import MODEL_PRICES, ModelPrice, cost_usd __all__ = [ # 1A — schemas, pricing, statistics diff --git a/src/eval/pipeline_factory.py b/src/eval/pipeline_factory.py index 1ab889a4..2bed4747 100644 --- a/src/eval/pipeline_factory.py +++ b/src/eval/pipeline_factory.py @@ -38,7 +38,7 @@ import chromadb from src.document_loader import TextChunker -from src.eval._telemetry import count_tokens +from src.telemetry.tokens import count_tokens from src.eval.config import EvalConfig from src.eval.schemas import EvalQuestion from src.llm_handler import LLMHandler @@ -222,7 +222,7 @@ def query(self, question: str) -> tuple[list[SearchResult], str, dict]: cost_usd: float (generator side) rewriter_cost_usd: float (rewriter side, 0.0 when disabled) """ - from src.eval import pricing + from src.telemetry import pricing timings: dict[str, float] = {} rewriter_cost = 0.0 diff --git a/src/eval/transforms/query_rewriter.py b/src/eval/transforms/query_rewriter.py index cc5eb2d9..8040ab98 100644 --- a/src/eval/transforms/query_rewriter.py +++ b/src/eval/transforms/query_rewriter.py @@ -17,7 +17,7 @@ import re from typing import Protocol -from src.eval import pricing +from src.telemetry import pricing logger = logging.getLogger(__name__) diff --git a/src/llm_handler.py b/src/llm_handler.py deleted file mode 100644 index 52864342..00000000 --- a/src/llm_handler.py +++ /dev/null @@ -1,812 +0,0 @@ -""" -LLM handler module for RAG pipeline. - -Supports OpenAI, Anthropic, and Ollama providers with graceful fallback -to a dummy response generator when providers are unavailable. -""" - -from __future__ import annotations - -import logging -import os -from pathlib import Path -from typing import Any, Generator, List, Optional - -# Load .env from project root -try: - from dotenv import load_dotenv - load_dotenv(Path(__file__).resolve().parent.parent / ".env") -except ImportError: - pass - -logger = logging.getLogger(__name__) - - -# BUG FIX: `generate()` and `stream_response()` used to swallow every exception -# and return a dummy response. That turned auth failures, quota exhaustion, -# and network outages into HTTP 200 "Your prompt was N characters" answers — -# dangerous in production. The fallback now fires *only* when the provider is -# genuinely unconfigured (module missing, api key missing). Real API errors -# propagate to the caller so the API returns 5xx. -class _ProviderUnavailableError(RuntimeError): - """Raised when a provider module is not installed or has no API key.""" - -# --------------------------------------------------------------------------- # -# Provider detection at import time (no hard dependency) # -# --------------------------------------------------------------------------- # - -try: - import openai as _openai_module - - _OPENAI_AVAILABLE = True -except ImportError: - _openai_module = None # type: ignore - _OPENAI_AVAILABLE = False - -try: - import anthropic as _anthropic_module - - _ANTHROPIC_AVAILABLE = True -except ImportError: - _anthropic_module = None # type: ignore - _ANTHROPIC_AVAILABLE = False - -try: - import requests as _requests_module - - _REQUESTS_AVAILABLE = True -except ImportError: - _requests_module = None # type: ignore - _REQUESTS_AVAILABLE = False - - -# --------------------------------------------------------------------------- # -# Provider constants # -# --------------------------------------------------------------------------- # - -OPENAI_MODELS = [ - "gpt-4o", - "gpt-4o-mini", - "gpt-4-turbo", - "gpt-4", - "gpt-3.5-turbo", -] - -ANTHROPIC_MODELS = [ - "claude-opus-4-5", - "claude-sonnet-4-5", - "claude-3-5-haiku-20241022", - "claude-3-opus-20240229", -] - -OLLAMA_DEFAULT_MODELS = [ - "llama3", - "mistral", - "gemma", - "phi3", -] - - -class LLMHandler: - """Unified LLM interface supporting OpenAI, Anthropic, and Ollama. - - Provider selection is automatic based on the model prefix: - - Models starting with 'gpt' / 'o1' / 'o3' → OpenAI - - Models starting with 'claude' → Anthropic - - Others → Ollama (localhost) - - Falls back to a dummy response if the selected provider is unavailable. - """ - - def __init__( - self, - model: str = "gpt-4", - temperature: float = 0.7, - # WHY: 1024 tokens is too low for detailed Q&A responses — the LLM - # hits the limit mid-sentence and truncates the answer. - # 4096 gives enough room for thorough, formatted answers while - # still bounding cost and latency. - max_tokens: int = 4096, - api_key: Optional[str] = None, - ollama_base_url: str = "http://localhost:11434", - ) -> None: - """ - Args: - model: Model name (determines provider). - temperature: Sampling temperature. - max_tokens: Maximum output tokens. - api_key: API key (falls back to env vars OPENAI_API_KEY / ANTHROPIC_API_KEY). - ollama_base_url: Base URL for Ollama server. - """ - self.model = model - self.temperature = temperature - self.max_tokens = max_tokens - self.api_key = api_key - self.ollama_base_url = ollama_base_url.rstrip("/") - - self._provider = self._detect_provider(model) - logger.info("LLMHandler initialised: model=%s provider=%s", model, self._provider) - - # ------------------------------------------------------------------ # - # Public API # - # ------------------------------------------------------------------ # - - def generate( - self, - prompt: str, - system_prompt: Optional[str] = None, - ) -> str: - """Generate a response for a prompt. - - Args: - prompt: User message. - system_prompt: Optional system instructions. - - Returns: - Model response string. - """ - try: - if self._provider in ("openai", "glm"): - return self._openai_generate(prompt, system_prompt) - if self._provider == "anthropic": - return self._anthropic_generate(prompt, system_prompt) - if self._provider == "ollama": - return self._ollama_generate(prompt, system_prompt) - except _ProviderUnavailableError as exc: - # Genuinely unconfigured — silent fallback keeps the demo usable - # without any API keys (the TF-IDF path still returns sources). - logger.warning("Provider unavailable, using dummy: %s", exc) - return self._dummy_response(prompt) - # Real errors (auth, quota, network) re-raise so the API returns 5xx - # instead of a fake success. - return self._dummy_response(prompt) - - def generate_with_context(self, query: str, context: str) -> str: - """Generate a RAG response given a user query and retrieved context. - - Args: - query: User question. - context: Retrieved document context. - - Returns: - Answer string. - """ - system_prompt = ( - "You are a helpful assistant. Answer the user's question based solely on the " - "provided context. If the context does not contain enough information, say so." - ) - user_prompt = f"Context:\n{context}\n\nQuestion: {query}\n\nAnswer:" - return self.generate(user_prompt, system_prompt=system_prompt) - - def generate_with_usage( - self, - prompt: str, - system_prompt: str | None = None, - ) -> tuple[str, int, int]: - """Generate a response and return text plus prompt/completion token counts. - - WHY a separate method: the existing `generate()` returns only `str` and is - called in many places that don't need usage. Phase 2's cost ledger needs - token counts on every LLM call; rather than break callers, we add a parallel - method that uses the existing tokenizer to estimate counts client-side. - - Args: - prompt: User message. - system_prompt: Optional system instructions. - - Returns: - (response_text, prompt_tokens, completion_tokens). - """ - from src.eval._telemetry import count_tokens - - text = self.generate(prompt, system_prompt=system_prompt) - full_prompt = (system_prompt + "\n" + prompt) if system_prompt else prompt - prompt_tokens = count_tokens(full_prompt, self.model) - completion_tokens = count_tokens(text, self.model) - return text, prompt_tokens, completion_tokens - - def stream_response( - self, - prompt: str, - system_prompt: Optional[str] = None, - ) -> Generator[str, None, None]: - """Stream model output token by token. - - Args: - prompt: User message. - system_prompt: Optional system instructions. - - Yields: - Text chunks as they arrive. - """ - try: - if self._provider in ("openai", "glm"): - yield from self._openai_stream(prompt, system_prompt) - return - if self._provider == "anthropic": - yield from self._anthropic_stream(prompt, system_prompt) - return - if self._provider == "ollama": - yield from self._ollama_stream(prompt, system_prompt) - return - except _ProviderUnavailableError as exc: - logger.warning("Provider unavailable, using dummy stream: %s", exc) - for word in self._dummy_response(prompt).split(): - yield word + " " - return - # Any other provider error (auth/quota/network) re-raises so the - # WebSocket handler can surface it as a structured error event. - - def generate_messages(self, messages: List[dict]) -> str: - """Generate a response given a full conversation messages list. - - This is the messages-list counterpart to generate(). It accepts the - standard OpenAI-style messages format (list of role/content dicts), - which carries the full sliding-window chat history. - - RAG Pipeline Position: - Chat History → [generate_messages] → Provider → Answer - ^^^ - Whereas generate() wraps a single prompt string, this method passes - the conversation history straight through to the provider — enabling - multi-turn awareness at the LLM level. - - Args: - messages: Conversation history as OpenAI-style dicts, e.g.: - [ - {"role": "system", "content": "..."}, - {"role": "user", "content": "..."}, - {"role": "assistant", "content": "..."}, - {"role": "user", "content": "..."}, ← latest turn - ] - - Returns: - Model response string, or dummy fallback if provider fails. - """ - try: - if self._provider in ("openai", "glm"): - return self._openai_generate_messages(messages) - if self._provider == "anthropic": - return self._anthropic_generate_messages(messages) - if self._provider == "ollama": - return self._ollama_generate_messages(messages) - except _ProviderUnavailableError as exc: - logger.warning("Provider unavailable, using dummy: %s", exc) - return self._dummy_messages_response(messages) - # Real provider errors propagate so the API can return a proper 5xx. - return self._dummy_messages_response(messages) - - def stream_messages( - self, messages: List[dict] - ) -> Generator[str, None, None]: - """Stream a response given a full conversation messages list. - - Streaming counterpart to generate_messages(). Yields text chunks as - they arrive so the frontend can display tokens incrementally without - waiting for the full response. - - Args: - messages: Conversation history as OpenAI-style dicts (same shape - as generate_messages()). - - Yields: - Text chunks from the model, or dummy fallback tokens if the - provider fails. - """ - try: - if self._provider in ("openai", "glm"): - yield from self._openai_stream_messages(messages) - return - if self._provider == "anthropic": - yield from self._anthropic_stream_messages(messages) - return - if self._provider == "ollama": - yield from self._ollama_stream_messages(messages) - return - except _ProviderUnavailableError as exc: - logger.warning("Provider unavailable, using dummy stream: %s", exc) - for word in self._dummy_messages_response(messages).split(): - yield word + " " - return - # Real provider errors propagate — the WebSocket handler turns them - # into a structured {"type": "error"} event for the client. - - def list_models(self) -> List[str]: - """Return a list of available models for the current provider. - - Returns: - List of model name strings. - """ - if self._provider == "openai": - return self._openai_list_models() - if self._provider == "anthropic": - return list(ANTHROPIC_MODELS) - if self._provider == "ollama": - return self._ollama_list_models() - return ["dummy-model"] - - # ------------------------------------------------------------------ # - # OpenAI # - # ------------------------------------------------------------------ # - - def _is_constrained_openai_model(self) -> bool: - """Return True if the current model rejects `temperature` overrides. - - WHY: GPT-5 family and the o-series reasoning models only accept the - default temperature (1). Sending any other value raises a 400. - Older families (gpt-4, gpt-4o, gpt-4.1) still accept the full range. - """ - m = self.model.lower() - return ( - m.startswith("gpt-5") - or m.startswith("o1") - or m.startswith("o3") - or m.startswith("o4") - ) - - def _openai_kwargs(self, **extra: Any) -> dict: - """Build kwargs for an OpenAI chat completions call. - - PATTERN: Centralised so every call site picks up the same constraints: - - Always use `max_completion_tokens` (gpt-5*/o* require it; modern - gpt-4* accept it). - - Omit `temperature` for constrained models rather than sending the - default, so a future SDK default change doesn't surprise us. - """ - kwargs: dict[str, Any] = { - "model": self.model, - "max_completion_tokens": self.max_tokens, - **extra, - } - if not self._is_constrained_openai_model(): - kwargs["temperature"] = self.temperature - return kwargs - - def _openai_generate(self, prompt: str, system_prompt: Optional[str]) -> str: - client = self._openai_client_for_provider() - messages = [] - if system_prompt: - messages.append({"role": "system", "content": system_prompt}) - messages.append({"role": "user", "content": prompt}) - - response = client.chat.completions.create( - **self._openai_kwargs(messages=messages) - ) - return response.choices[0].message.content or "" - - def _openai_stream( - self, prompt: str, system_prompt: Optional[str] - ) -> Generator[str, None, None]: - client = self._openai_client_for_provider() - messages = [] - if system_prompt: - messages.append({"role": "system", "content": system_prompt}) - messages.append({"role": "user", "content": prompt}) - - stream = client.chat.completions.create( - **self._openai_kwargs(messages=messages, stream=True) - ) - for chunk in stream: - delta = chunk.choices[0].delta - if delta and delta.content: - yield delta.content - - def _openai_list_models(self) -> List[str]: - if not _OPENAI_AVAILABLE: - return list(OPENAI_MODELS) - try: - client = _openai_module.OpenAI( # type: ignore[union-attr] - api_key=self.api_key or os.getenv("OPENAI_API_KEY") - ) - models = client.models.list() - return [m.id for m in models.data if "gpt" in m.id] - except Exception as exc: - logger.warning("Could not list OpenAI models: %s", exc) - return list(OPENAI_MODELS) - - # ------------------------------------------------------------------ # - # Anthropic # - # ------------------------------------------------------------------ # - - def _anthropic_generate(self, prompt: str, system_prompt: Optional[str]) -> str: - if not _ANTHROPIC_AVAILABLE: - raise _ProviderUnavailableError("anthropic package not installed") - - client = _anthropic_module.Anthropic( # type: ignore[union-attr] - api_key=self.api_key or os.getenv("ANTHROPIC_API_KEY") - ) - kwargs: dict = dict( - model=self.model, - max_tokens=self.max_tokens, - messages=[{"role": "user", "content": prompt}], - ) - if system_prompt: - kwargs["system"] = system_prompt - - response = client.messages.create(**kwargs) - block = response.content[0] - return block.text if hasattr(block, "text") else str(block) - - def _anthropic_stream( - self, prompt: str, system_prompt: Optional[str] - ) -> Generator[str, None, None]: - if not _ANTHROPIC_AVAILABLE: - raise _ProviderUnavailableError("anthropic package not installed") - - client = _anthropic_module.Anthropic( # type: ignore[union-attr] - api_key=self.api_key or os.getenv("ANTHROPIC_API_KEY") - ) - kwargs: dict = dict( - model=self.model, - max_tokens=self.max_tokens, - messages=[{"role": "user", "content": prompt}], - ) - if system_prompt: - kwargs["system"] = system_prompt - - with client.messages.stream(**kwargs) as stream: - for text in stream.text_stream: - yield text - - # ------------------------------------------------------------------ # - # Ollama # - # ------------------------------------------------------------------ # - - @staticmethod - def _wrap_ollama_errors(exc: Exception) -> _ProviderUnavailableError: - """Convert a requests error from a local Ollama into unavailable. - - WHY: A local Ollama that's offline, returns 404 for an unknown model, - or refuses connections is functionally "not configured" — same bucket - as a missing package or missing API key. Wrapping these as - _ProviderUnavailableError lets the top-level dispatch fall back to - the dummy response without suppressing real provider auth/quota - errors from OpenAI, Anthropic, or GLM. - """ - return _ProviderUnavailableError(f"Ollama unavailable: {exc}") - - def _ollama_generate(self, prompt: str, system_prompt: Optional[str]) -> str: - if not _REQUESTS_AVAILABLE: - raise _ProviderUnavailableError("requests package not installed") - - payload: dict = { - "model": self.model, - "prompt": prompt, - "stream": False, - "options": { - "temperature": self.temperature, - "num_predict": self.max_tokens, - }, - } - if system_prompt: - payload["system"] = system_prompt - - try: - resp = _requests_module.post( # type: ignore[union-attr] - f"{self.ollama_base_url}/api/generate", - json=payload, - timeout=120, - ) - resp.raise_for_status() - except _requests_module.exceptions.RequestException as exc: # type: ignore[union-attr] - raise self._wrap_ollama_errors(exc) from exc - return resp.json().get("response", "") - - def _ollama_stream( - self, prompt: str, system_prompt: Optional[str] - ) -> Generator[str, None, None]: - if not _REQUESTS_AVAILABLE: - raise _ProviderUnavailableError("requests package not installed") - - import json - - payload: dict = { - "model": self.model, - "prompt": prompt, - "stream": True, - "options": { - "temperature": self.temperature, - "num_predict": self.max_tokens, - }, - } - if system_prompt: - payload["system"] = system_prompt - - try: - with _requests_module.post( # type: ignore[union-attr] - f"{self.ollama_base_url}/api/generate", - json=payload, - stream=True, - timeout=120, - ) as resp: - resp.raise_for_status() - for line in resp.iter_lines(): - if line: - data = json.loads(line) - token = data.get("response", "") - if token: - yield token - if data.get("done"): - break - except _requests_module.exceptions.RequestException as exc: # type: ignore[union-attr] - raise self._wrap_ollama_errors(exc) from exc - - # ------------------------------------------------------------------ # - # OpenAI — messages-list API # - # ------------------------------------------------------------------ # - - def _openai_generate_messages(self, messages: List[dict]) -> str: - """Pass a messages list directly to the OpenAI chat completions API. - - WHY: Unlike _openai_generate() which builds the messages list itself - from a single prompt, this method receives the already-assembled list - so the caller controls the full conversation history. - """ - client = self._openai_client_for_provider() - response = client.chat.completions.create( - **self._openai_kwargs(messages=messages) # type: ignore[arg-type] - ) - return response.choices[0].message.content or "" - - def _openai_stream_messages( - self, messages: List[dict] - ) -> Generator[str, None, None]: - """Stream a response from OpenAI given a messages list.""" - client = self._openai_client_for_provider() - stream = client.chat.completions.create( - **self._openai_kwargs(messages=messages, stream=True) # type: ignore[arg-type] - ) - for chunk in stream: - delta = chunk.choices[0].delta - if delta and delta.content: - yield delta.content - - # ------------------------------------------------------------------ # - # Anthropic — messages-list API # - # ------------------------------------------------------------------ # - - def _anthropic_generate_messages(self, messages: List[dict]) -> str: - """Generate a response from Anthropic given a messages list. - - WHY: Anthropic's API does NOT accept "system" role messages inside the - messages list — it expects system instructions as a separate top-level - `system=` kwarg. We therefore filter out system messages first and join - their content, then pass only user/assistant turns in `messages`. - - TRADE-OFF: Multiple system messages are joined with newlines. In practice, - well-formed chat histories have at most one system message (the first one), - but defensive handling is safer. - """ - if not _ANTHROPIC_AVAILABLE: - raise _ProviderUnavailableError("anthropic package not installed") - - client = _anthropic_module.Anthropic( # type: ignore[union-attr] - api_key=self.api_key or os.getenv("ANTHROPIC_API_KEY") - ) - - # PATTERN: Separate system messages from conversation turns. - # Anthropic requires this split at the API level. - system_parts = [m["content"] for m in messages if m["role"] == "system"] - non_system = [m for m in messages if m["role"] != "system"] - - kwargs: dict = dict( - model=self.model, - max_tokens=self.max_tokens, - messages=non_system, - ) - if system_parts: - kwargs["system"] = "\n".join(system_parts) - - response = client.messages.create(**kwargs) - block = response.content[0] - return block.text if hasattr(block, "text") else str(block) - - def _anthropic_stream_messages( - self, messages: List[dict] - ) -> Generator[str, None, None]: - """Stream a response from Anthropic given a messages list. - - See _anthropic_generate_messages() for the system-message separation - rationale — the same split applies here. - """ - if not _ANTHROPIC_AVAILABLE: - raise _ProviderUnavailableError("anthropic package not installed") - - client = _anthropic_module.Anthropic( # type: ignore[union-attr] - api_key=self.api_key or os.getenv("ANTHROPIC_API_KEY") - ) - - system_parts = [m["content"] for m in messages if m["role"] == "system"] - non_system = [m for m in messages if m["role"] != "system"] - - kwargs: dict = dict( - model=self.model, - max_tokens=self.max_tokens, - messages=non_system, - ) - if system_parts: - kwargs["system"] = "\n".join(system_parts) - - with client.messages.stream(**kwargs) as stream: - for text in stream.text_stream: - yield text - - # ------------------------------------------------------------------ # - # Ollama — messages-list API (/api/chat) # - # ------------------------------------------------------------------ # - - def _ollama_generate_messages(self, messages: List[dict]) -> str: - """Generate a response from Ollama using the /api/chat endpoint. - - WHY: The existing _ollama_generate() uses /api/generate, which accepts - a single prompt string. /api/chat is the correct Ollama endpoint for - multi-turn conversations — it accepts an OpenAI-style messages list - and understands role context (system/user/assistant). - - Response format differs from /api/generate: - /api/generate → {"response": "token", "done": false} - /api/chat → {"message": {"role": "assistant", "content": "token"}, "done": false} - """ - if not _REQUESTS_AVAILABLE: - raise _ProviderUnavailableError("requests package not installed") - - payload: dict = { - "model": self.model, - "messages": messages, - "stream": False, - "options": { - "temperature": self.temperature, - "num_predict": self.max_tokens, - }, - } - - try: - resp = _requests_module.post( # type: ignore[union-attr] - f"{self.ollama_base_url}/api/chat", - json=payload, - timeout=120, - ) - resp.raise_for_status() - except _requests_module.exceptions.RequestException as exc: # type: ignore[union-attr] - raise self._wrap_ollama_errors(exc) from exc - # WHY: /api/chat wraps the response inside a "message" object, unlike - # /api/generate which puts it at the top-level "response" key. - return resp.json().get("message", {}).get("content", "") - - def _ollama_stream_messages( - self, messages: List[dict] - ) -> Generator[str, None, None]: - """Stream a response from Ollama /api/chat endpoint. - - See _ollama_generate_messages() for endpoint and response-format details. - """ - if not _REQUESTS_AVAILABLE: - raise _ProviderUnavailableError("requests package not installed") - - import json - - payload: dict = { - "model": self.model, - "messages": messages, - "stream": True, - "options": { - "temperature": self.temperature, - "num_predict": self.max_tokens, - }, - } - - try: - with _requests_module.post( # type: ignore[union-attr] - f"{self.ollama_base_url}/api/chat", - json=payload, - stream=True, - timeout=120, - ) as resp: - resp.raise_for_status() - for line in resp.iter_lines(): - if line: - data = json.loads(line) - # WHY: /api/chat nests the token inside data["message"]["content"], - # not at data["response"] like /api/generate does. - token = data.get("message", {}).get("content", "") - if token: - yield token - if data.get("done"): - break - except _requests_module.exceptions.RequestException as exc: # type: ignore[union-attr] - raise self._wrap_ollama_errors(exc) from exc - - def _ollama_list_models(self) -> List[str]: - if not _REQUESTS_AVAILABLE: - return list(OLLAMA_DEFAULT_MODELS) - try: - resp = _requests_module.get( # type: ignore[union-attr] - f"{self.ollama_base_url}/api/tags", timeout=5 - ) - resp.raise_for_status() - return [m["name"] for m in resp.json().get("models", [])] - except Exception as exc: - logger.warning("Could not list Ollama models: %s", exc) - return list(OLLAMA_DEFAULT_MODELS) - - # ------------------------------------------------------------------ # - # Helpers # - # ------------------------------------------------------------------ # - - @staticmethod - def _detect_provider(model: str) -> str: - """Infer provider from model name. - - BUG FIX: `glm*` used to fall through to Ollama, which silently bricked - `GLM_API_KEY` configuration in docker-compose.prod.yml. GLM (Zhipu AI) - exposes an OpenAI-compatible API, so we route it through the OpenAI - SDK with a custom base_url (see `_openai_client_for_provider`). - """ - lower = model.lower() - if lower.startswith("gpt") or lower.startswith("o1") or lower.startswith("o3"): - return "openai" - if lower.startswith("claude"): - return "anthropic" - if lower.startswith("glm"): - return "glm" - return "ollama" - - # GLM / Zhipu AI — OpenAI-compatible endpoint. Configurable via env so - # the same class works for Zhipu's public API or any self-hosted proxy. - GLM_DEFAULT_BASE_URL = "https://open.bigmodel.cn/api/paas/v4" - - def _openai_client_for_provider(self): - """Build an OpenAI-SDK client pointed at the right endpoint. - - WHY: Zhipu's GLM API is wire-compatible with OpenAI's chat completions - endpoint, so we reuse the same SDK and just redirect base_url+api_key - when the provider is "glm". Keeping this in one place means every - chat method (generate, stream, messages) stays a one-liner. - """ - if not _OPENAI_AVAILABLE: - raise _ProviderUnavailableError("openai package not installed") - if self._provider == "glm": - api_key = self.api_key or os.getenv("GLM_API_KEY") - if not api_key: - raise _ProviderUnavailableError("GLM_API_KEY not set") - base_url = os.getenv("GLM_BASE_URL", self.GLM_DEFAULT_BASE_URL) - return _openai_module.OpenAI( # type: ignore[union-attr] - api_key=api_key, base_url=base_url, - ) - return _openai_module.OpenAI( # type: ignore[union-attr] - api_key=self.api_key or os.getenv("OPENAI_API_KEY") - ) - - @staticmethod - def _dummy_response(prompt: str) -> str: - """Return a safe fallback response when no provider is available.""" - return ( - "[LLM unavailable] This is a placeholder response. " - f"Your prompt was {len(prompt)} characters. " - "Please configure a valid LLM provider (OpenAI, Anthropic, or Ollama)." - ) - - @staticmethod - def _dummy_messages_response(messages: List[dict]) -> str: - """Return a safe fallback response for the messages-list API. - - WHY: generate_messages() receives a List[dict] instead of a plain string, - so we can't call _dummy_response(prompt) directly. We derive a meaningful - length metric from the last user message (same sentinel, different context). - - PATTERN: Extract last user turn for the character count so the fallback - message is as informative as the prompt-based one. - """ - # Find the last user message to report its length in the fallback string - last_user_content = "" - for msg in reversed(messages): - if msg.get("role") == "user": - last_user_content = msg.get("content", "") - break - - return ( - "[LLM unavailable] This is a placeholder response. " - f"Your last message was {len(last_user_content)} characters " - f"(conversation has {len(messages)} turn(s)). " - "Please configure a valid LLM provider (OpenAI, Anthropic, or Ollama)." - ) diff --git a/src/llm_handler/__init__.py b/src/llm_handler/__init__.py new file mode 100644 index 00000000..41575d97 --- /dev/null +++ b/src/llm_handler/__init__.py @@ -0,0 +1,180 @@ +"""LLM handler — provider selection over a set of swappable adapters. + +RAG Pipeline Position: + Query + Context -> Prompt -> [LLM HANDLER] -> Answer (+ Usage) + ^^^ + LLMHandler owns the single-prompt -> messages translation and the + unconfigured-provider fallback. Provider selection and SDK client + construction live in ``providers`` (via ``build_adapter``); the actual SDK + calls live in one adapter per provider (see ``adapters/``). + +Design Decision: + Before this refactor a single class repeated the same provider dispatch four + times and built its clients from module globals, so two of three real + providers had no test coverage. Now one adapter is selected at construction + and every provider is unit-testable with a fake client (see + ``tests/test_llm_adapters.py``). + + The public interface (generate / stream / *_messages / generate_with_usage / + list_models) is unchanged for callers. Prompt variants are thin translations + to the messages form the adapters speak. +""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Iterator + +# Load .env from the project root (two levels up from this package). +try: + from dotenv import load_dotenv + + load_dotenv(Path(__file__).resolve().parent.parent.parent / ".env") +except ImportError: + pass + +from .adapters.base import ( + GenerationResult, + ProviderAdapter, + ProviderUnavailableError, + Usage, +) +from .adapters.dummy import DummyAdapter +from .providers import build_adapter, detect_provider, list_models + +logger = logging.getLogger(__name__) + +__all__ = [ + "LLMHandler", + "GenerationResult", + "ProviderAdapter", + "ProviderUnavailableError", + "Usage", +] + + +def _to_messages(prompt: str, system_prompt: str | None) -> list[dict]: + """Translate a single prompt (+ optional system) into a messages list.""" + messages: list[dict] = [] + if system_prompt: + messages.append({"role": "system", "content": system_prompt}) + messages.append({"role": "user", "content": prompt}) + return messages + + +class LLMHandler: + """Unified LLM interface backed by one provider adapter per instance. + + Provider selection is automatic from the model prefix (gpt/o1/o3 -> OpenAI, + claude -> Anthropic, glm -> GLM, else Ollama). If the selected provider is + unconfigured, calls fall back to a dummy response so the demo stays usable. + """ + + def __init__( + self, + model: str = "gpt-4", + temperature: float = 0.7, + # WHY 4096: 1024 truncates detailed Q&A answers mid-sentence; 4096 gives + # room for thorough, formatted answers while bounding cost. + max_tokens: int = 4096, + api_key: str | None = None, + ollama_base_url: str = "http://localhost:11434", + ) -> None: + """Configure the handler and select its provider adapter. + + Args: + model: Model name (determines provider). + temperature: Sampling temperature. + max_tokens: Maximum output tokens. + api_key: API key (falls back to provider-specific env vars). + ollama_base_url: Base URL for a local Ollama server. + """ + self.model = model + self.temperature = temperature + self.max_tokens = max_tokens + self.api_key = api_key + self.ollama_base_url = ollama_base_url.rstrip("/") + + self._provider = detect_provider(model) + self._adapter = build_adapter( + model, temperature, max_tokens, api_key, self.ollama_base_url + ) + # The fallback is always ready — no client, no configuration. + self._dummy = DummyAdapter(model) + logger.info("LLMHandler initialised: model=%s provider=%s", model, self._provider) + + # ------------------------------------------------------------------ # + # Public API — non-streaming # + # ------------------------------------------------------------------ # + + def generate(self, prompt: str, system_prompt: str | None = None) -> str: + """Generate a response for a single prompt.""" + return self._generate(_to_messages(prompt, system_prompt)).text + + def generate_messages(self, messages: list[dict]) -> str: + """Generate a response for a full OpenAI-style messages list.""" + return self._generate(messages).text + + def generate_with_usage( + self, prompt: str, system_prompt: str | None = None + ) -> tuple[str, int, int]: + """Generate a response and return (text, prompt_tokens, completion_tokens). + + Usage is provider-reported where available, adapter-counted otherwise — + callers no longer rebuild the prompt just to estimate token counts. + """ + result = self._generate(_to_messages(prompt, system_prompt)) + return result.text, result.usage.prompt_tokens, result.usage.completion_tokens + + def _generate(self, messages: list[dict]) -> GenerationResult: + """Call the selected adapter, falling back to dummy if unconfigured.""" + try: + return self._adapter.generate(messages) + except ProviderUnavailableError as exc: + logger.warning("Provider unavailable, using dummy: %s", exc) + return self._dummy.generate(messages) + + # ------------------------------------------------------------------ # + # Public API — streaming # + # ------------------------------------------------------------------ # + + def stream_response( + self, prompt: str, system_prompt: str | None = None + ) -> Iterator[str | Usage]: + """Stream a response for a single prompt. + + Yields text chunks, then a terminal ``Usage``. Callers discriminate the + terminal event with ``isinstance(item, Usage)``. + """ + yield from self._stream(_to_messages(prompt, system_prompt)) + + def stream_messages(self, messages: list[dict]) -> Iterator[str | Usage]: + """Stream a response for a messages list. + + Yields text chunks, then a terminal ``Usage`` (see ``stream_response``). + """ + yield from self._stream(messages) + + def _stream(self, messages: list[dict]) -> Iterator[str | Usage]: + """Stream text chunks then a terminal Usage from the adapter. + + Falls back to the dummy adapter (which reports counted usage) when the + provider is unconfigured. The terminal ``Usage`` flows through to the + backend's telemetry assembly. + """ + try: + yield from self._adapter.stream(messages) + except ProviderUnavailableError as exc: + logger.warning("Provider unavailable, using dummy stream: %s", exc) + yield from self._dummy.stream(messages) + + # ------------------------------------------------------------------ # + # Model listing (out of scope for the adapter refactor — preserved) # + # ------------------------------------------------------------------ # + + def list_models(self) -> list[str]: + """Return available model names for the current provider.""" + return list_models( + self._provider, self.model, self.api_key, self.ollama_base_url + ) diff --git a/src/llm_handler/adapters/__init__.py b/src/llm_handler/adapters/__init__.py new file mode 100644 index 00000000..19f21df0 --- /dev/null +++ b/src/llm_handler/adapters/__init__.py @@ -0,0 +1,6 @@ +"""Provider adapters for the LLM layer. + +Each module here implements the ``ProviderAdapter`` protocol from ``base`` for +one provider family, hiding its SDK's request shape, response parsing, and +usage extraction behind the shared messages-first interface. +""" diff --git a/src/llm_handler/adapters/anthropic.py b/src/llm_handler/adapters/anthropic.py new file mode 100644 index 00000000..f4a6bea6 --- /dev/null +++ b/src/llm_handler/adapters/anthropic.py @@ -0,0 +1,99 @@ +"""Anthropic provider adapter. + +RAG Pipeline Position: + Messages -> [ANTHROPIC ADAPTER] -> Answer + Usage + +Design Decision: + Anthropic's API does not accept ``system`` role entries inside the messages + list — the system prompt is a separate top-level ``system=`` argument. This + adapter owns that split so no caller has to know the quirk. Multiple system + messages are joined with newlines (well-formed histories have at most one). + + Usage maps Anthropic's ``input_tokens`` / ``output_tokens`` onto the shared + ``Usage`` prompt/completion fields, with a local count as fallback. +""" + +from __future__ import annotations + +from typing import Callable, Iterator + +from .base import ( + GenerationResult, + Usage, + counted_usage, + join_message_text, +) + + +class AnthropicAdapter: + """Messages adapter for Anthropic's Claude models.""" + + def __init__( + self, + model: str, + max_tokens: int, + client_factory: Callable[[], object], + ) -> None: + """Store generation settings and the injected client factory. + + Args: + model: Claude model name. + max_tokens: Output token cap. + client_factory: Zero-arg builder for an Anthropic-SDK-shaped client, + called once per request. + + Note: + No ``temperature`` — the previous implementation never sent one to + Anthropic, and this adapter preserves that exactly. + """ + self.model = model + self.max_tokens = max_tokens + self._client_factory = client_factory + + def _request_kwargs(self, messages: list[dict]) -> dict[str, object]: + """Split system messages out and build the create()/stream() kwargs.""" + system_parts = [m["content"] for m in messages if m["role"] == "system"] + non_system = [m for m in messages if m["role"] != "system"] + kwargs: dict[str, object] = { + "model": self.model, + "max_tokens": self.max_tokens, + "messages": non_system, + } + if system_parts: + kwargs["system"] = "\n".join(system_parts) + return kwargs + + def generate(self, messages: list[dict], **kwargs: object) -> GenerationResult: + """Call messages.create and return text plus usage.""" + client = self._client_factory() + response = client.messages.create(**self._request_kwargs(messages)) + block = response.content[0] + text = block.text if hasattr(block, "text") else str(block) + usage = _usage_from(response) + if usage is None: + usage = counted_usage(join_message_text(messages), text, self.model) + return GenerationResult(text=text, usage=usage) + + def stream(self, messages: list[dict], **kwargs: object) -> Iterator[str | Usage]: + """Stream tokens, then a terminal Usage (reported or counted).""" + client = self._client_factory() + collected: list[str] = [] + with client.messages.stream(**self._request_kwargs(messages)) as stream: + for text in stream.text_stream: + collected.append(text) + yield text + usage = _usage_from(stream.get_final_message()) + yield usage or counted_usage( + join_message_text(messages), "".join(collected), self.model + ) + + +def _usage_from(message: object) -> Usage | None: + """Map Anthropic's input/output token usage, or None when absent.""" + usage = getattr(message, "usage", None) + if usage is None: + return None + return Usage( + prompt_tokens=usage.input_tokens, + completion_tokens=usage.output_tokens, + ) diff --git a/src/llm_handler/adapters/base.py b/src/llm_handler/adapters/base.py new file mode 100644 index 00000000..61d6cff4 --- /dev/null +++ b/src/llm_handler/adapters/base.py @@ -0,0 +1,140 @@ +"""Provider-adapter contract for the LLM layer. + +RAG Pipeline Position: + Query + Context -> Prompt -> [LLM PROVIDER ADAPTER] -> Answer + Usage + ^^^ + This is the seam between LLMHandler (which owns provider *selection* and the + prompt/messages translation) and the concrete SDK calls. Each provider lives + behind one adapter that speaks a single messages-first interface. + +Design Decision: + A ``Protocol`` (not an ABC) defines the contract. WHY: the project standard + prefers structural typing for narrow, swappable interfaces — an adapter is + "a ProviderAdapter" if it has ``generate`` and ``stream`` with the right + shapes, with no inheritance ceremony. ``@runtime_checkable`` lets tests assert + conformance with ``isinstance``. + + Client construction is *injected*: every adapter receives a zero-arg + ``client_factory`` so a unit test can pass a fake instead of monkeypatching a + module global. Before this refactor, two of three real providers had zero + coverage precisely because the only way to fake them was to reach into module + internals. + + Usage is reported alongside text: provider-supplied counts where the SDK + returns them, adapter-internal token counting as the fallback (see each + adapter). Callers never rebuild a prompt just to estimate what was billed. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Iterator, Protocol, runtime_checkable + +from src.telemetry.tokens import count_tokens + + +class ProviderUnavailableError(RuntimeError): + """Raised when a provider is not usable for configuration reasons. + + "Configuration reasons" means the SDK package is not installed, no API key is + set, or a local Ollama server refuses the connection — i.e. the provider is + *unconfigured*, not *failing*. LLMHandler catches only this error and falls + back to the dummy adapter, so the demo stays usable with no keys. + + Genuine runtime errors (auth rejected, quota exhausted, malformed request) + are NOT this type: they propagate so the API surfaces a real 5xx instead of a + fake "success" dummy answer. + """ + + +@dataclass(frozen=True) +class Usage: + """Token counts for one generation call. + + Attributes: + prompt_tokens: Tokens billed for the input (system + messages). + completion_tokens: Tokens billed for the generated output. + """ + + prompt_tokens: int + completion_tokens: int + + +@dataclass(frozen=True) +class GenerationResult: + """The text and usage returned by a non-streaming generation. + + Attributes: + text: The model's answer. + usage: Provider-reported usage where available, adapter-counted otherwise. + """ + + text: str + usage: Usage + + +def join_message_text(messages: list[dict]) -> str: + """Concatenate message contents for fallback token counting. + + Joins with newlines so the count matches what LLMHandler used to reconstruct + for the same messages — keeping fallback numbers stable across the refactor. + """ + return "\n".join(str(m.get("content", "")) for m in messages) + + +def counted_usage(input_text: str, output_text: str, model: str) -> Usage: + """Estimate usage with the local tokenizer when a provider reports none. + + This is the "adapter-internal counting" fallback: the dummy adapter always + uses it, and the real adapters fall back to it only when a provider omits + usage (e.g. a stream that carries no terminal usage chunk). + """ + return Usage( + prompt_tokens=count_tokens(input_text, model), + completion_tokens=count_tokens(output_text, model), + ) + + +@runtime_checkable +class ProviderAdapter(Protocol): + """One provider behind a messages-first interface. + + Implementations own their SDK's request shape, response parsing, usage + extraction, and the provider-specific quirks (base URLs, token-parameter + names, system-message handling). LLMHandler owns provider *selection* and the + single-prompt -> messages translation, so adapters only ever see messages. + """ + + def generate(self, messages: list[dict], **kwargs: object) -> GenerationResult: + """Generate a full response for an OpenAI-style messages list. + + Args: + messages: Conversation as ``[{"role": ..., "content": ...}, ...]``. + **kwargs: Reserved for future per-call overrides. + + Returns: + The answer text plus token usage. + + Raises: + ProviderUnavailableError: If the provider is unconfigured. + """ + ... + + def stream(self, messages: list[dict], **kwargs: object) -> Iterator[str | Usage]: + """Stream a response, yielding text chunks then a terminal ``Usage``. + + Consumers discriminate the terminal event with ``isinstance(item, Usage)``: + every ``str`` is an answer chunk; exactly one final ``Usage`` reports the + call's token counts. + + Args: + messages: Conversation as ``[{"role": ..., "content": ...}, ...]``. + **kwargs: Reserved for future per-call overrides. + + Yields: + ``str`` answer chunks, then one ``Usage``. + + Raises: + ProviderUnavailableError: If the provider is unconfigured. + """ + ... diff --git a/src/llm_handler/adapters/dummy.py b/src/llm_handler/adapters/dummy.py new file mode 100644 index 00000000..780e4a5f --- /dev/null +++ b/src/llm_handler/adapters/dummy.py @@ -0,0 +1,58 @@ +"""Dummy fallback adapter — the answer when no provider is configured. + +RAG Pipeline Position: + ... -> [DUMMY ADAPTER] -> placeholder answer + LLMHandler routes here when the selected provider raises + ``ProviderUnavailableError`` (missing SDK, missing key, local server down), + so the demo stays usable — the retrieval half still returns real sources. + +Design Decision: + Usage is always the local-count fallback: there is no provider to report it. +""" + +from __future__ import annotations + +from typing import Iterator + +from .base import ( + GenerationResult, + Usage, + counted_usage, + join_message_text, +) + + +class DummyAdapter: + """Return a safe placeholder response for any messages list.""" + + def __init__(self, model: str) -> None: + """Store the model name (used only to pick a tokenizer for usage counts).""" + self.model = model + + def generate(self, messages: list[dict], **kwargs: object) -> GenerationResult: + """Return the placeholder answer with locally-counted usage.""" + text = self._placeholder(messages) + usage = counted_usage(join_message_text(messages), text, self.model) + return GenerationResult(text=text, usage=usage) + + def stream(self, messages: list[dict], **kwargs: object) -> Iterator[str | Usage]: + """Stream the placeholder word by word, then a terminal Usage.""" + text = self._placeholder(messages) + for word in text.split(): + yield word + " " + yield counted_usage(join_message_text(messages), text, self.model) + + @staticmethod + def _placeholder(messages: list[dict]) -> str: + """Build the sentinel string, sized from the last user turn.""" + last_user_content = "" + for message in reversed(messages): + if message.get("role") == "user": + last_user_content = message.get("content", "") + break + return ( + "[LLM unavailable] This is a placeholder response. " + f"Your last message was {len(last_user_content)} characters " + f"(conversation has {len(messages)} turn(s)). " + "Please configure a valid LLM provider (OpenAI, Anthropic, or Ollama)." + ) diff --git a/src/llm_handler/adapters/ollama.py b/src/llm_handler/adapters/ollama.py new file mode 100644 index 00000000..536dd969 --- /dev/null +++ b/src/llm_handler/adapters/ollama.py @@ -0,0 +1,150 @@ +"""Ollama provider adapter — local models over HTTP. + +RAG Pipeline Position: + Messages -> [OLLAMA ADAPTER] -> Answer + Usage + +Design Decision: + Ollama is a local server, so a refused connection or an unknown model is + treated as "unconfigured" (``ProviderUnavailableError``) — the same bucket as + a missing SDK — letting LLMHandler fall back to the dummy rather than 500. + + This adapter speaks only ``/api/chat`` (the messages-native endpoint). The + older single-prompt ``/api/generate`` path is gone: LLMHandler now translates + every single-prompt call into a messages list before it reaches an adapter, + so one endpoint covers both. + + Usage comes from Ollama's ``prompt_eval_count`` / ``eval_count`` fields (on + the response, or the terminal streamed line), with a local count as fallback. +""" + +from __future__ import annotations + +import json +from typing import Callable, Iterator + +from .base import ( + GenerationResult, + ProviderUnavailableError, + Usage, + counted_usage, + join_message_text, +) + +# Requests errors (connection refused, timeout, HTTP error) all subclass +# RequestException. Catching it narrowly wraps a down/unknown local server +# without swallowing genuine bugs. +_REQUEST_EXCEPTION: type[Exception] | tuple[type[Exception], ...] +try: + import requests + + _REQUEST_EXCEPTION = requests.exceptions.RequestException +except ImportError: # pragma: no cover - requests is a hard dependency + # If requests is missing, LLMHandler's factory raises ProviderUnavailableError + # before this adapter ever posts, so this except clause never fires. + _REQUEST_EXCEPTION = () + + +class OllamaAdapter: + """Messages adapter for a local Ollama server.""" + + def __init__( + self, + model: str, + temperature: float, + max_tokens: int, + base_url: str, + client_factory: Callable[[], object], + ) -> None: + """Store generation settings, the server URL, and the injected factory. + + Args: + model: Ollama model tag (e.g. ``llama3``). + temperature: Sampling temperature, sent under ``options``. + max_tokens: Output cap, sent as ``options.num_predict``. + base_url: Ollama server base URL (no trailing slash). + client_factory: Zero-arg builder for an HTTP client exposing + ``post`` (the ``requests`` module by default), called per request. + """ + self.model = model + self.temperature = temperature + self.max_tokens = max_tokens + self.base_url = base_url + self._client_factory = client_factory + + def _payload(self, messages: list[dict], *, stream: bool) -> dict[str, object]: + """Build the /api/chat request body.""" + return { + "model": self.model, + "messages": messages, + "stream": stream, + "options": { + "temperature": self.temperature, + "num_predict": self.max_tokens, + }, + } + + def generate(self, messages: list[dict], **kwargs: object) -> GenerationResult: + """POST to /api/chat and return content plus usage.""" + http = self._client_factory() + try: + resp = http.post( + f"{self.base_url}/api/chat", + json=self._payload(messages, stream=False), + timeout=120, + ) + resp.raise_for_status() + except _REQUEST_EXCEPTION as exc: + raise _unavailable(exc) from exc + + data = resp.json() + text = data.get("message", {}).get("content", "") + usage = _usage_from(data) + if usage is None: + usage = counted_usage(join_message_text(messages), text, self.model) + return GenerationResult(text=text, usage=usage) + + def stream(self, messages: list[dict], **kwargs: object) -> Iterator[str | Usage]: + """Stream tokens from /api/chat, then a terminal Usage.""" + http = self._client_factory() + collected: list[str] = [] + reported: Usage | None = None + try: + with http.post( + f"{self.base_url}/api/chat", + json=self._payload(messages, stream=True), + stream=True, + timeout=120, + ) as resp: + resp.raise_for_status() + for line in resp.iter_lines(): + if not line: + continue + data = json.loads(line) + token = data.get("message", {}).get("content", "") + if token: + collected.append(token) + yield token + if data.get("done"): + reported = _usage_from(data) + break + except _REQUEST_EXCEPTION as exc: + raise _unavailable(exc) from exc + + yield reported or counted_usage( + join_message_text(messages), "".join(collected), self.model + ) + + +def _usage_from(data: dict) -> Usage | None: + """Read Ollama's token counts, or None when the line omits them.""" + if "prompt_eval_count" not in data and "eval_count" not in data: + return None + return Usage( + prompt_tokens=data.get("prompt_eval_count", 0), + completion_tokens=data.get("eval_count", 0), + ) + + +def _unavailable(exc: Exception) -> ProviderUnavailableError: + """Wrap a requests failure as a provider-unavailable signal.""" + return ProviderUnavailableError(f"Ollama unavailable: {exc}") diff --git a/src/llm_handler/adapters/openai_compatible.py b/src/llm_handler/adapters/openai_compatible.py new file mode 100644 index 00000000..f957c5f8 --- /dev/null +++ b/src/llm_handler/adapters/openai_compatible.py @@ -0,0 +1,133 @@ +"""OpenAI-compatible provider adapter (serves both OpenAI and GLM). + +RAG Pipeline Position: + Messages -> [OPENAI-COMPATIBLE ADAPTER] -> Answer + Usage + +Design Decision: + OpenAI and Zhipu's GLM speak the same chat-completions wire protocol, so one + adapter serves both. The only difference is the client's base URL and key, + which LLMHandler bakes into the injected ``client_factory`` — so GLM support + is a factory concern, not a second code path here. + + The token-parameter and temperature quirks DO live here: + - Always send ``max_completion_tokens`` (the gpt-5 / o-series families + require it; modern gpt-4* accept it). + - Omit ``temperature`` for constrained models (gpt-5*, o1*, o3*, o4*) — + they reject any non-default value with a 400. + + Usage is provider-reported (``response.usage`` / the stream's terminal usage + chunk requested via ``stream_options``) with a local count as fallback. +""" + +from __future__ import annotations + +from typing import Callable, Iterator + +from .base import ( + GenerationResult, + Usage, + counted_usage, + join_message_text, +) + + +def _is_constrained(model: str) -> bool: + """Return True if the model rejects a ``temperature`` override. + + WHY these prefixes: the GPT-5 family and the o-series reasoning models only + accept the default temperature; older gpt-4* families accept the full range. + """ + lower = model.lower() + return ( + lower.startswith("gpt-5") + or lower.startswith("o1") + or lower.startswith("o3") + or lower.startswith("o4") + ) + + +class OpenAICompatibleAdapter: + """Chat-completions adapter for OpenAI and any OpenAI-compatible endpoint.""" + + def __init__( + self, + model: str, + temperature: float, + max_tokens: int, + client_factory: Callable[[], object], + ) -> None: + """Store generation settings and the injected client factory. + + Args: + model: Model name (also decides the constrained-model quirks). + temperature: Sampling temperature (dropped for constrained models). + max_tokens: Output cap, sent as ``max_completion_tokens``. + client_factory: Zero-arg builder for an OpenAI-SDK-shaped client. + Called once per request (no caching) so a rotated key or a test + fake is always picked up. + """ + self.model = model + self.temperature = temperature + self.max_tokens = max_tokens + self._client_factory = client_factory + + def _request_kwargs(self, **extra: object) -> dict[str, object]: + """Assemble the create() kwargs shared by generate() and stream().""" + kwargs: dict[str, object] = { + "model": self.model, + "max_completion_tokens": self.max_tokens, + **extra, + } + if not _is_constrained(self.model): + kwargs["temperature"] = self.temperature + return kwargs + + def generate(self, messages: list[dict], **kwargs: object) -> GenerationResult: + """Call chat.completions.create and return text plus usage.""" + client = self._client_factory() + response = client.chat.completions.create( + **self._request_kwargs(messages=messages) + ) + text = response.choices[0].message.content or "" + usage = _usage_from_response(response) + if usage is None: + usage = counted_usage(join_message_text(messages), text, self.model) + return GenerationResult(text=text, usage=usage) + + def stream(self, messages: list[dict], **kwargs: object) -> Iterator[str | Usage]: + """Stream tokens, then a terminal Usage (reported or counted).""" + client = self._client_factory() + stream = client.chat.completions.create( + **self._request_kwargs( + messages=messages, + stream=True, + # WHY: without include_usage the stream carries no token counts, + # forcing the local-count fallback. Requesting it lets the + # provider's own numbers flow through. + stream_options={"include_usage": True}, + ) + ) + collected: list[str] = [] + reported: Usage | None = None + for chunk in stream: + reported = _usage_from_response(chunk) or reported + if not chunk.choices: + continue + delta = chunk.choices[0].delta + if delta and delta.content: + collected.append(delta.content) + yield delta.content + yield reported or counted_usage( + join_message_text(messages), "".join(collected), self.model + ) + + +def _usage_from_response(response: object) -> Usage | None: + """Extract OpenAI-style usage, or None when the provider omits it.""" + usage = getattr(response, "usage", None) + if usage is None: + return None + return Usage( + prompt_tokens=usage.prompt_tokens, + completion_tokens=usage.completion_tokens, + ) diff --git a/src/llm_handler/providers.py b/src/llm_handler/providers.py new file mode 100644 index 00000000..7d631410 --- /dev/null +++ b/src/llm_handler/providers.py @@ -0,0 +1,199 @@ +"""Provider selection and client construction for LLMHandler. + +RAG Pipeline Position: + model name -> [PROVIDER WIRING] -> a configured ProviderAdapter + +Owns: which provider a model routes to, how to build that provider's SDK client +(the injected factory), the model-listing calls, and the graceful optional-import +guards. Split out of the LLMHandler facade so each module has one responsibility +and stays under the line ceiling. +""" + +from __future__ import annotations + +import logging +import os +from types import ModuleType +from typing import Callable + +from .adapters.anthropic import AnthropicAdapter +from .adapters.base import ProviderAdapter, ProviderUnavailableError +from .adapters.ollama import OllamaAdapter +from .adapters.openai_compatible import OpenAICompatibleAdapter + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- # +# Optional SDK availability (checked once at import, no hard dependency) # +# --------------------------------------------------------------------------- # + +_openai_module: ModuleType | None +try: + import openai as _openai_module +except ImportError: + _openai_module = None + +_anthropic_module: ModuleType | None +try: + import anthropic as _anthropic_module +except ImportError: + _anthropic_module = None + +_requests_module: ModuleType | None +try: + import requests as _requests_module +except ImportError: + _requests_module = None + + +# --------------------------------------------------------------------------- # +# Provider constants # +# --------------------------------------------------------------------------- # + +OPENAI_MODELS = ["gpt-4o", "gpt-4o-mini", "gpt-4-turbo", "gpt-4", "gpt-3.5-turbo"] +ANTHROPIC_MODELS = [ + "claude-opus-4-5", + "claude-sonnet-4-5", + "claude-3-5-haiku-20241022", + "claude-3-opus-20240229", +] +OLLAMA_DEFAULT_MODELS = ["llama3", "mistral", "gemma", "phi3"] + +# GLM / Zhipu AI — OpenAI-compatible endpoint, overridable via env. +GLM_DEFAULT_BASE_URL = "https://open.bigmodel.cn/api/paas/v4" + + +def detect_provider(model: str) -> str: + """Infer the provider from the model-name prefix. + + GLM (Zhipu AI) exposes an OpenAI-compatible API, so it routes to the + OpenAI-compatible adapter with a custom base URL (see the client factory). + + Args: + model: Model name. + + Returns: + One of ``"openai"``, ``"anthropic"``, ``"glm"``, ``"ollama"``. + """ + lower = model.lower() + if lower.startswith("gpt") or lower.startswith("o1") or lower.startswith("o3"): + return "openai" + if lower.startswith("claude"): + return "anthropic" + if lower.startswith("glm"): + return "glm" + return "ollama" + + +def build_adapter( + model: str, + temperature: float, + max_tokens: int, + api_key: str | None, + ollama_base_url: str, +) -> ProviderAdapter: + """Select and construct the one adapter for a model's provider. + + Each adapter is given an injected zero-arg client factory (built here) so the + SDK client is constructed per request and can be faked in tests. + + Args: + model: Model name (also decides the provider and constrained-model quirks). + temperature: Sampling temperature. + max_tokens: Output token cap. + api_key: Explicit key, or None to read the provider's env var. + ollama_base_url: Base URL for a local Ollama server. + + Returns: + A ready-to-use ``ProviderAdapter``. + """ + provider = detect_provider(model) + if provider in ("openai", "glm"): + return OpenAICompatibleAdapter( + model, temperature, max_tokens, _openai_client_factory(provider, api_key) + ) + if provider == "anthropic": + return AnthropicAdapter(model, max_tokens, _anthropic_client_factory(api_key)) + return OllamaAdapter( + model, temperature, max_tokens, ollama_base_url, _ollama_client_factory() + ) + + +def _openai_client_factory(provider: str, api_key: str | None) -> Callable[[], object]: + """Build a factory for an OpenAI-SDK client, redirected to GLM when needed. + + Missing SDK -> ProviderUnavailableError (dummy fallback). Missing GLM key -> + ProviderUnavailableError. A missing OpenAI key is left to the SDK, which + raises its own error that propagates (unchanged behaviour). + """ + def factory() -> object: + if _openai_module is None: + raise ProviderUnavailableError("openai package not installed") + if provider == "glm": + key = api_key or os.getenv("GLM_API_KEY") + if not key: + raise ProviderUnavailableError("GLM_API_KEY not set") + base_url = os.getenv("GLM_BASE_URL", GLM_DEFAULT_BASE_URL) + return _openai_module.OpenAI(api_key=key, base_url=base_url) + return _openai_module.OpenAI(api_key=api_key or os.getenv("OPENAI_API_KEY")) + + return factory + + +def _anthropic_client_factory(api_key: str | None) -> Callable[[], object]: + """Build a factory for an Anthropic-SDK client (missing SDK -> unavailable).""" + def factory() -> object: + if _anthropic_module is None: + raise ProviderUnavailableError("anthropic package not installed") + return _anthropic_module.Anthropic(api_key=api_key or os.getenv("ANTHROPIC_API_KEY")) + + return factory + + +def _ollama_client_factory() -> Callable[[], object]: + """Build a factory returning the HTTP client for Ollama (the requests module).""" + def factory() -> object: + if _requests_module is None: + raise ProviderUnavailableError("requests package not installed") + return _requests_module + + return factory + + +def list_models( + provider: str, model: str, api_key: str | None, ollama_base_url: str +) -> list[str]: + """Return available model names for a provider (out of scope; preserved).""" + if provider == "openai": + return _openai_list_models(api_key) + if provider == "anthropic": + return list(ANTHROPIC_MODELS) + if provider == "ollama": + return _ollama_list_models(ollama_base_url) + return ["dummy-model"] + + +def _openai_list_models(api_key: str | None) -> list[str]: + """Query OpenAI for GPT models, falling back to the static list.""" + if _openai_module is None: + return list(OPENAI_MODELS) + try: + client = _openai_module.OpenAI(api_key=api_key or os.getenv("OPENAI_API_KEY")) + models = client.models.list() + return [m.id for m in models.data if "gpt" in m.id] + except Exception as exc: + logger.warning("Could not list OpenAI models: %s", exc) + return list(OPENAI_MODELS) + + +def _ollama_list_models(ollama_base_url: str) -> list[str]: + """Query the local Ollama server for tags, falling back to defaults.""" + if _requests_module is None: + return list(OLLAMA_DEFAULT_MODELS) + try: + resp = _requests_module.get(f"{ollama_base_url}/api/tags", timeout=5) + resp.raise_for_status() + return [m["name"] for m in resp.json().get("models", [])] + except Exception as exc: + logger.warning("Could not list Ollama models: %s", exc) + return list(OLLAMA_DEFAULT_MODELS) diff --git a/src/telemetry/__init__.py b/src/telemetry/__init__.py new file mode 100644 index 00000000..4668a0a9 --- /dev/null +++ b/src/telemetry/__init__.py @@ -0,0 +1,18 @@ +"""Core telemetry — token counting and cost pricing. + +RAG Pipeline Position: + Generation (reported usage) -> [TELEMETRY] -> StageTelemetry / EvalResult + +Owned by the core so both the production telemetry assembly and the eval harness +depend on one source of truth. Production no longer imports these from the eval +package (the dependency direction the architecture-deepening spec inverts). + +See [ADR 0003](../../docs/adr/0003-telemetry-ownership.md). +""" + +from __future__ import annotations + +from src.telemetry.pricing import MODEL_PRICES, ModelPrice, cost_usd +from src.telemetry.tokens import count_tokens + +__all__ = ["MODEL_PRICES", "ModelPrice", "cost_usd", "count_tokens"] diff --git a/src/eval/pricing.py b/src/telemetry/pricing.py similarity index 86% rename from src/eval/pricing.py rename to src/telemetry/pricing.py index 1327a362..9be6490d 100644 --- a/src/eval/pricing.py +++ b/src/telemetry/pricing.py @@ -1,11 +1,13 @@ """ -Model pricing table and cost arithmetic. - -Eval Harness Position: - Pipeline → tokens → [PRICING] → cost_usd → EvalResult.cost_usd - ^^^^^^^^ - Pure data + a 4-line function. Hard-coded prices are fine for Phase 1; - if prices change frequently we move to a JSON file in a later phase. +Model pricing table and cost arithmetic — core telemetry utility. + +RAG Pipeline Position: + tokens (reported usage) → [PRICING] → cost_usd → StageTelemetry / EvalResult + ^^^^^^^^ + Owned by the core (``src/telemetry/``) so both the production telemetry + assembly and the eval harness compute cost from one price table — production + no longer imports it from the eval package. Pure data + a small function; + hard-coded prices are fine here (see below). Design decisions: - Hard-coded table, not env-driven, so price changes are visible in diff --git a/src/eval/_telemetry.py b/src/telemetry/tokens.py similarity index 71% rename from src/eval/_telemetry.py rename to src/telemetry/tokens.py index cde5f6f3..d14286b6 100644 --- a/src/eval/_telemetry.py +++ b/src/telemetry/tokens.py @@ -1,14 +1,14 @@ -"""Internal telemetry helpers for the eval harness. +"""Token counting — core telemetry utility. -This module provides token counting and timing utilities used by -EvalPipeline.query() to track cost and latency. It's kept separate from -pipeline_factory.py to keep that file under the project's 250-line ceiling. +Owned by the core (``src/telemetry/``) and used by both production +(``LLMHandler`` adapters, as the usage fallback) and the eval harness. It used +to live under ``src/eval/``, which forced production code to import from the eval +package — the dependency this module's move inverts (core never imports eval). Token Counting: - Counting tokens is essential for eval cost tracking. This module tries - tiktoken first (exact model-aware tokenization) and falls back to a - word-count heuristic if tiktoken doesn't know the model. The fallback - ensures eval doesn't hard-fail on new model releases. + This module tries tiktoken first (exact model-aware tokenization) and falls + back to a word-count heuristic if tiktoken doesn't know the model. The fallback + ensures nothing hard-fails on a new model release before tiktoken is updated. """ from __future__ import annotations diff --git a/tests/test_backend_telemetry.py b/tests/test_backend_telemetry.py index 31a90102..4ea86d7b 100644 --- a/tests/test_backend_telemetry.py +++ b/tests/test_backend_telemetry.py @@ -210,6 +210,37 @@ def test_stream_query_emits_telemetry_on_empty_store(self, backend: RAGBackend): assert last_data["generate_ms"] == 0.0 assert last_data["prompt_tokens"] == 0 + def test_stream_query_with_conversation_emits_telemetry_and_persists( + self, ingested_backend: RAGBackend + ): + """The conversation branch captures answer usage and still persists. + + WHY: stream_query's multi-turn branch is the one that reads the answer + stream's terminal Usage AND saves the assistant message. This drives it + with a real conversation_id so both behaviours are exercised together — + the non-conversation telemetry tests never take this path. + """ + conv_id = ingested_backend.create_conversation()["id"] + + events = list( + ingested_backend.stream_query("What is RAG?", conversation_id=conv_id) + ) + + # Telemetry still last, with non-negative usage from the captured Usage. + last_type, last_data = events[-1] + assert last_type == "telemetry" + assert last_data["prompt_tokens"] >= 0 + assert last_data["completion_tokens"] >= 0 + + # The done event carries the persisted message + conversation ids. + done_data = next(d for t, d in events if t == "done") + assert done_data["conversation_id"] == conv_id + assert done_data["message_id"] + + # The assistant message was actually saved to the conversation. + detail = ingested_backend.get_conversation(conv_id) + assert any(m["role"] == "assistant" for m in detail["messages"]) + def test_stream_query_existing_events_order_preserved( self, ingested_backend: RAGBackend ): diff --git a/tests/test_llm_adapters.py b/tests/test_llm_adapters.py new file mode 100644 index 00000000..99b39262 --- /dev/null +++ b/tests/test_llm_adapters.py @@ -0,0 +1,453 @@ +"""Contract tests for the LLM provider adapters. + +What concept this teaches: + Each provider lives behind one adapter that speaks the same messages-first + interface (``generate`` / ``stream``). Because the SDK client is *injected* + via a zero-arg factory, every provider — including Anthropic and Ollama, + which previously had no coverage — is unit-testable with a fake client. No + monkeypatching of module globals. + +Seams under test: + Each adapter's public ``generate`` (-> GenerationResult) and ``stream`` + (-> str chunks then a terminal Usage), plus the ProviderUnavailableError + triggers that make LLMHandler fall back to the dummy. +""" + +from __future__ import annotations + +import sys +from pathlib import Path +from types import SimpleNamespace + +PROJECT_ROOT = Path(__file__).parent.parent +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +import pytest + +from src.llm_handler.adapters.base import ( + GenerationResult, + ProviderAdapter, + ProviderUnavailableError, + Usage, +) +import requests + +from src.llm_handler import LLMHandler +from src.llm_handler.adapters.anthropic import AnthropicAdapter +from src.llm_handler.adapters.dummy import DummyAdapter +from src.llm_handler.adapters.ollama import OllamaAdapter +from src.llm_handler.adapters.openai_compatible import OpenAICompatibleAdapter +from src.llm_handler.providers import build_adapter, detect_provider + + +# --------------------------------------------------------------------------- # +# DummyAdapter — the always-available fallback # +# --------------------------------------------------------------------------- # + +class TestDummyAdapter: + """The dummy adapter needs no client and always returns a placeholder.""" + + def test_conforms_to_protocol(self) -> None: + assert isinstance(DummyAdapter("any-model"), ProviderAdapter) + + def test_generate_returns_marker_and_counted_usage(self) -> None: + adapter = DummyAdapter("dummy-model") + + result = adapter.generate([{"role": "user", "content": "What is RAG?"}]) + + assert isinstance(result, GenerationResult) + assert "[LLM unavailable]" in result.text + # Fallback usage is counted locally, so a non-empty prompt is > 0 tokens. + assert result.usage.prompt_tokens > 0 + assert result.usage.completion_tokens > 0 + + def test_stream_yields_tokens_then_terminal_usage(self) -> None: + adapter = DummyAdapter("dummy-model") + + items = list(adapter.stream([{"role": "user", "content": "Hello"}])) + + # Exactly one terminal Usage, and it is last. + assert isinstance(items[-1], Usage) + assert sum(isinstance(i, Usage) for i in items) == 1 + text = "".join(i for i in items if isinstance(i, str)) + assert "[LLM unavailable]" in text + + +# --------------------------------------------------------------------------- # +# OpenAI-compatible adapter (serves both OpenAI and GLM) # +# --------------------------------------------------------------------------- # + +class FakeOpenAIClient: + """Mimics the subset of the OpenAI SDK the adapter calls. + + Records every ``chat.completions.create`` kwargs list in ``calls`` so tests + can assert the token-parameter/temperature quirks. ``usage`` / ``stream_usage`` + of ``None`` simulate a provider that omits usage (forcing the count fallback). + """ + + def __init__( + self, + *, + usage: tuple[int, int] | None = (10, 8), + stream_usage: tuple[int, int] | None = (11, 9), + calls: list[dict] | None = None, + ) -> None: + self.calls = calls if calls is not None else [] + self._usage = usage + self._stream_usage = stream_usage + self.chat = SimpleNamespace( + completions=SimpleNamespace(create=self._create) + ) + + def _create(self, **kwargs: object): + self.calls.append(kwargs) + if kwargs.get("stream"): + return self._stream_chunks() + return self._response() + + def _response(self): + choice = SimpleNamespace( + message=SimpleNamespace(content="OpenAI answer", role="assistant"), + finish_reason="stop", + index=0, + ) + usage = ( + SimpleNamespace(prompt_tokens=self._usage[0], completion_tokens=self._usage[1]) + if self._usage + else None + ) + return SimpleNamespace(choices=[choice], usage=usage) + + def _stream_chunks(self): + for piece in ("OpenAI ", "answer"): + delta = SimpleNamespace(content=piece, role="assistant") + yield SimpleNamespace( + choices=[SimpleNamespace(delta=delta, finish_reason=None, index=0)], + usage=None, + ) + if self._stream_usage: + # Matches the real SDK with stream_options include_usage: a final + # chunk with empty choices carries the usage totals. + yield SimpleNamespace( + choices=[], + usage=SimpleNamespace( + prompt_tokens=self._stream_usage[0], + completion_tokens=self._stream_usage[1], + ), + ) + + +class TestOpenAICompatibleAdapter: + """OpenAI and GLM share this adapter — GLM differs only by injected client.""" + + def _adapter(self, model: str, fake: FakeOpenAIClient) -> OpenAICompatibleAdapter: + return OpenAICompatibleAdapter( + model=model, temperature=0.5, max_tokens=256, client_factory=lambda: fake + ) + + def test_conforms_to_protocol(self) -> None: + assert isinstance(self._adapter("gpt-4", FakeOpenAIClient()), ProviderAdapter) + + def test_generate_returns_reported_usage(self) -> None: + fake = FakeOpenAIClient(usage=(10, 8)) + result = self._adapter("gpt-4", fake).generate( + [{"role": "user", "content": "Q"}] + ) + assert result.text == "OpenAI answer" + assert result.usage == Usage(prompt_tokens=10, completion_tokens=8) + + def test_generate_falls_back_to_counted_usage_when_provider_omits_it(self) -> None: + fake = FakeOpenAIClient(usage=None) + result = self._adapter("gpt-4", fake).generate( + [{"role": "user", "content": "Q"}] + ) + assert result.usage.prompt_tokens > 0 # counted locally, not reported + + def test_constrained_model_omits_temperature_and_uses_completion_tokens(self) -> None: + fake = FakeOpenAIClient() + self._adapter("gpt-5-mini", fake).generate([{"role": "user", "content": "Q"}]) + kwargs = fake.calls[-1] + assert "temperature" not in kwargs + assert kwargs["max_completion_tokens"] == 256 + + def test_unconstrained_model_includes_temperature(self) -> None: + fake = FakeOpenAIClient() + self._adapter("gpt-4", fake).generate([{"role": "user", "content": "Q"}]) + assert fake.calls[-1]["temperature"] == 0.5 + + def test_stream_yields_tokens_then_reported_usage(self) -> None: + fake = FakeOpenAIClient(stream_usage=(11, 9)) + items = list(self._adapter("gpt-4", fake).stream([{"role": "user", "content": "Q"}])) + assert isinstance(items[-1], Usage) + assert items[-1] == Usage(prompt_tokens=11, completion_tokens=9) + assert "".join(i for i in items if isinstance(i, str)) == "OpenAI answer" + # The adapter must request usage on the stream. + assert fake.calls[-1]["stream_options"] == {"include_usage": True} + + def test_stream_falls_back_to_counted_usage_without_usage_chunk(self) -> None: + fake = FakeOpenAIClient(stream_usage=None) + items = list(self._adapter("gpt-4", fake).stream([{"role": "user", "content": "Q"}])) + assert isinstance(items[-1], Usage) + assert items[-1].completion_tokens > 0 # counted from streamed text + + +# --------------------------------------------------------------------------- # +# Anthropic adapter — owns the system-message split # +# --------------------------------------------------------------------------- # + +class _FakeAnthropicStream: + """Context-manager stream mirroring anthropic's messages.stream().""" + + def __init__(self, usage: tuple[int, int] | None) -> None: + self._usage = usage + + def __enter__(self) -> "_FakeAnthropicStream": + return self + + def __exit__(self, *exc: object) -> bool: + return False + + @property + def text_stream(self): + yield "Claude " + yield "answer" + + def get_final_message(self): + usage = ( + SimpleNamespace(input_tokens=self._usage[0], output_tokens=self._usage[1]) + if self._usage + else None + ) + return SimpleNamespace(usage=usage) + + +class FakeAnthropicClient: + """Mimics the subset of the Anthropic SDK the adapter calls.""" + + def __init__( + self, + *, + usage: tuple[int, int] | None = (12, 7), + stream_usage: tuple[int, int] | None = (13, 6), + calls: list[dict] | None = None, + ) -> None: + self.calls = calls if calls is not None else [] + self._usage = usage + self._stream_usage = stream_usage + self.messages = SimpleNamespace(create=self._create, stream=self._stream) + + def _create(self, **kwargs: object): + self.calls.append(kwargs) + usage = ( + SimpleNamespace(input_tokens=self._usage[0], output_tokens=self._usage[1]) + if self._usage + else None + ) + return SimpleNamespace(content=[SimpleNamespace(text="Claude answer")], usage=usage) + + def _stream(self, **kwargs: object) -> _FakeAnthropicStream: + self.calls.append(kwargs) + return _FakeAnthropicStream(self._stream_usage) + + +class TestAnthropicAdapter: + """Anthropic needs system messages hoisted out of the messages list.""" + + def _adapter(self, fake: FakeAnthropicClient) -> AnthropicAdapter: + return AnthropicAdapter( + model="claude-sonnet-4-5", max_tokens=256, client_factory=lambda: fake + ) + + def test_conforms_to_protocol(self) -> None: + assert isinstance(self._adapter(FakeAnthropicClient()), ProviderAdapter) + + def test_generate_maps_reported_usage(self) -> None: + result = self._adapter(FakeAnthropicClient(usage=(12, 7))).generate( + [{"role": "user", "content": "Q"}] + ) + assert result.text == "Claude answer" + # Anthropic reports input/output tokens -> prompt/completion. + assert result.usage == Usage(prompt_tokens=12, completion_tokens=7) + + def test_generate_hoists_system_message(self) -> None: + fake = FakeAnthropicClient() + self._adapter(fake).generate( + [ + {"role": "system", "content": "Be terse."}, + {"role": "user", "content": "Q"}, + ] + ) + kwargs = fake.calls[-1] + assert kwargs["system"] == "Be terse." + # The system turn must NOT remain in the messages list. + assert all(m["role"] != "system" for m in kwargs["messages"]) + + def test_stream_yields_tokens_then_reported_usage(self) -> None: + items = list( + self._adapter(FakeAnthropicClient(stream_usage=(13, 6))).stream( + [{"role": "user", "content": "Q"}] + ) + ) + assert isinstance(items[-1], Usage) + assert items[-1] == Usage(prompt_tokens=13, completion_tokens=6) + assert "".join(i for i in items if isinstance(i, str)) == "Claude answer" + + def test_stream_falls_back_to_counted_usage(self) -> None: + items = list( + self._adapter(FakeAnthropicClient(stream_usage=None)).stream( + [{"role": "user", "content": "Q"}] + ) + ) + assert isinstance(items[-1], Usage) + assert items[-1].completion_tokens > 0 + + +# --------------------------------------------------------------------------- # +# Ollama adapter — local /api/chat with prompt_eval_count / eval_count usage # +# --------------------------------------------------------------------------- # + +class _FakeOllamaResponse: + def __init__(self, payload: dict) -> None: + self._payload = payload + + def raise_for_status(self) -> None: + pass + + def json(self) -> dict: + return self._payload + + +class _FakeOllamaStreamResponse: + def __init__(self, lines: list[bytes]) -> None: + self._lines = lines + + def __enter__(self) -> "_FakeOllamaStreamResponse": + return self + + def __exit__(self, *exc: object) -> bool: + return False + + def raise_for_status(self) -> None: + pass + + def iter_lines(self): + yield from self._lines + + +class FakeOllamaHttp: + """Mimics the requests.post interface the adapter calls.""" + + def __init__( + self, + *, + content: str = "Ollama answer", + usage: tuple[int, int] | None = (20, 15), + fail: bool = False, + stream_lines: list[bytes] | None = None, + calls: list[dict] | None = None, + ) -> None: + self.calls = calls if calls is not None else [] + self._content = content + self._usage = usage + self._fail = fail + self._stream_lines = stream_lines + + def post(self, url: str, json: dict | None = None, timeout: int | None = None, stream: bool = False): + self.calls.append({"url": url, "json": json, "stream": stream}) + if self._fail: + raise requests.exceptions.ConnectionError("connection refused") + if stream: + return _FakeOllamaStreamResponse(self._stream_lines or []) + payload: dict = {"message": {"content": self._content}} + if self._usage: + payload["prompt_eval_count"] = self._usage[0] + payload["eval_count"] = self._usage[1] + return _FakeOllamaResponse(payload) + + +class TestOllamaAdapter: + """Ollama runs locally; connection refusal means 'unconfigured', not error.""" + + def _adapter(self, http: FakeOllamaHttp) -> OllamaAdapter: + return OllamaAdapter( + model="llama3", + temperature=0.5, + max_tokens=256, + base_url="http://localhost:11434", + client_factory=lambda: http, + ) + + def test_conforms_to_protocol(self) -> None: + assert isinstance(self._adapter(FakeOllamaHttp()), ProviderAdapter) + + def test_generate_posts_to_chat_and_returns_usage(self) -> None: + http = FakeOllamaHttp(usage=(20, 15)) + result = self._adapter(http).generate([{"role": "user", "content": "Q"}]) + assert result.text == "Ollama answer" + assert result.usage == Usage(prompt_tokens=20, completion_tokens=15) + assert http.calls[-1]["url"].endswith("/api/chat") + + def test_connection_failure_raises_provider_unavailable(self) -> None: + adapter = self._adapter(FakeOllamaHttp(fail=True)) + with pytest.raises(ProviderUnavailableError): + adapter.generate([{"role": "user", "content": "Q"}]) + + def test_stream_yields_tokens_then_reported_usage(self) -> None: + lines = [ + b'{"message": {"content": "Ollama "}, "done": false}', + b'{"message": {"content": "answer"}, "done": false}', + b'{"done": true, "prompt_eval_count": 21, "eval_count": 14}', + ] + items = list( + self._adapter(FakeOllamaHttp(stream_lines=lines)).stream( + [{"role": "user", "content": "Q"}] + ) + ) + assert isinstance(items[-1], Usage) + assert items[-1] == Usage(prompt_tokens=21, completion_tokens=14) + assert "".join(i for i in items if isinstance(i, str)) == "Ollama answer" + + def test_stream_falls_back_to_counted_usage_without_counts(self) -> None: + lines = [ + b'{"message": {"content": "hi"}, "done": false}', + b'{"done": true}', + ] + items = list( + self._adapter(FakeOllamaHttp(stream_lines=lines)).stream( + [{"role": "user", "content": "Q"}] + ) + ) + assert isinstance(items[-1], Usage) + assert items[-1].completion_tokens > 0 + + +# --------------------------------------------------------------------------- # +# Provider selection — model prefix -> adapter, incl. the GLM branch # +# --------------------------------------------------------------------------- # + +class TestProviderSelection: + """LLMHandler picks one adapter per model; GLM shares the OpenAI adapter.""" + + def test_detect_provider_routes_by_prefix(self) -> None: + assert detect_provider("gpt-4") == "openai" + assert detect_provider("claude-sonnet-4-5") == "anthropic" + assert detect_provider("glm-5.1") == "glm" + assert detect_provider("llama3") == "ollama" + + def test_glm_model_selects_openai_compatible_adapter(self) -> None: + adapter = build_adapter( + "glm-5.1", 0.7, 256, api_key="k", ollama_base_url="http://localhost:11434" + ) + assert isinstance(adapter, OpenAICompatibleAdapter) + + def test_glm_without_key_falls_back_to_dummy(self, monkeypatch) -> None: + """A GLM model with no GLM_API_KEY is unconfigured -> dummy fallback. + + WHY this branch matters: GLM's key/base-url swap is the one part of the + selection logic the OpenAI adapter's own tests don't cover. + """ + monkeypatch.delenv("GLM_API_KEY", raising=False) + handler = LLMHandler("glm-5.1") # api_key=None, so no GLM key available + + assert "[LLM unavailable]" in handler.generate("hi") diff --git a/tests/test_llm_handler.py b/tests/test_llm_handler.py index c6c8c8f0..f0fcc424 100644 --- a/tests/test_llm_handler.py +++ b/tests/test_llm_handler.py @@ -25,7 +25,7 @@ import pytest -from src.llm_handler import LLMHandler +from src.llm_handler import LLMHandler, Usage # --------------------------------------------------------------------------- # @@ -107,10 +107,10 @@ def test_stream_messages_falls_back_to_dummy(self) -> None: handler = LLMHandler(model=_DUMMY_MODEL) messages = [{"role": "user", "content": "Hello, world!"}] - tokens = list(handler.stream_messages(messages)) - full_response = "".join(tokens) + items = list(handler.stream_messages(messages)) + full_response = "".join(i for i in items if isinstance(i, str)) - assert len(tokens) > 0, "stream_messages must yield at least one token" + assert full_response, "stream_messages must yield answer text" assert "[LLM unavailable]" in full_response, ( "Streamed fallback must contain '[LLM unavailable]' marker" ) @@ -125,8 +125,22 @@ def test_stream_messages_accepts_sliding_window(self) -> None: """ handler = LLMHandler(model=_DUMMY_MODEL) - tokens = list(handler.stream_messages(_SLIDING_WINDOW_MESSAGES)) - full_response = "".join(tokens) + items = list(handler.stream_messages(_SLIDING_WINDOW_MESSAGES)) + full_response = "".join(i for i in items if isinstance(i, str)) - assert len(tokens) > 0, "Must yield tokens for multi-turn input" assert len(full_response) > 0, "Joined tokens must form a non-empty response" + + def test_stream_messages_ends_with_terminal_usage(self) -> None: + """The stream yields text chunks, then exactly one terminal Usage. + + WHY: cost telemetry needs the token counts the provider (or the local + fallback) reports for the streamed call, delivered as the final event so + text chunks stay pure strings for the UI. + """ + handler = LLMHandler(model=_DUMMY_MODEL) + + items = list(handler.stream_messages([{"role": "user", "content": "Hi"}])) + + assert isinstance(items[-1], Usage), "last streamed item must be Usage" + assert sum(isinstance(i, Usage) for i in items) == 1 + assert all(isinstance(i, str) for i in items[:-1]), "chunks before it are text" diff --git a/tests/test_eval_pricing.py b/tests/test_telemetry_pricing.py similarity index 93% rename from tests/test_eval_pricing.py rename to tests/test_telemetry_pricing.py index 51c0b670..6a1c0dc7 100644 --- a/tests/test_eval_pricing.py +++ b/tests/test_telemetry_pricing.py @@ -1,4 +1,4 @@ -"""Tests for src.eval.pricing.""" +"""Tests for src.telemetry.pricing.""" from __future__ import annotations @@ -6,7 +6,7 @@ import pytest -from src.eval.pricing import MODEL_PRICES, ModelPrice, cost_usd +from src.telemetry.pricing import MODEL_PRICES, ModelPrice, cost_usd class TestModelPriceTable: