From 1fa5f70f63a403e53ffd08e30103eb40f8b39f2e Mon Sep 17 00:00:00 2001 From: Eric Law <39393654+acn-ericlaw@users.noreply.github.com> Date: Mon, 31 Aug 2026 19:26:19 -0700 Subject: [PATCH] AI nodes llm.chat and llm.stream - provider-neutral LLM adapters (E0) The agent-orchestration experiment's wrapper-side AI nodes: llm.chat answers a prompt or conversation with optional schema-constrained structured output (the verdict shape a graph's decision routing needs), llm.stream relays the provider's real token stream over the multi-shot reply contract so a calling engine renders it progressively out its own HTTP edge. One contract, two provider editions (Anthropic and Gemini SDKs as lazy optional extras - pip install 'mercury-composable[llm]'); provider errors ride the envelope status; timeout_ms maps onto the SDK timeout (the x-ttl pattern); the Gemini editions opt out of the SDK's automatic-function-calling loop (no tool surface by design - the graph decides, the model advises). 14 token-free tests pin the contract with fake clients and a recording writer. Co-Authored-By: Claude Code --- CHANGELOG.md | 19 +++ examples/demo_app.py | 337 ++++++++++++++++++++++++++++++++++++++- pyproject.toml | 4 +- tests/test_llm_chat.py | 252 +++++++++++++++++++++++++++++ tests/test_llm_stream.py | 234 +++++++++++++++++++++++++++ 5 files changed, 844 insertions(+), 2 deletions(-) create mode 100644 tests/test_llm_chat.py create mode 100644 tests/test_llm_stream.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 65bb8cf..4b8a7f3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,24 @@ # Changelog +## Unreleased + +### Added + +- `llm.chat` demo function (examples/demo_app.py) - the AI node of the agent-orchestration + experiment E0: a provider-neutral LLM adapter (Anthropic and Gemini backends behind one + contract; select with `llm.provider` config / `-Dllm.provider=...` / `params.provider`) + with structured output (`schema` -> JSON verdicts for graph decision routing; + additionalProperties defaults closed), usage/stop_reason surfacing, provider-error + mapping onto the envelope status, and the params.timeout_ms time-budget mapping. + The SDKs are an optional extra (`pip install "mercury-composable[llm]"`); the package + itself stays SDK-free (scope fence intact). +- `llm.stream` demo function - the STREAMING AI node: pulls the provider's real token + stream (Anthropic `messages.stream` or Gemini `generate_content_stream`) and relays + each token batch over the multi-shot reply contract, so a calling engine renders it + progressively out its own HTTP edge as SSE. Terminal metadata carries model, + stop_reason, usage and the trace/business correlation ids; provider errors fail the + stream in-band. Live-proven end to end with real Gemini tokens (2026-08-31). + ## 4.12.0 (2026-08-30) The progressive-rendering milestone release. The version aligns with the Mercury diff --git a/examples/demo_app.py b/examples/demo_app.py index 4b77028..87fa3f9 100644 --- a/examples/demo_app.py +++ b/examples/demo_app.py @@ -15,6 +15,8 @@ """ import asyncio +import json +from typing import Any from mercury_composable import ( AppException, @@ -23,6 +25,7 @@ EventStreamWriter, PostOffice, annotate_trace, + app_config, get_logger, get_trace, platform, @@ -31,6 +34,20 @@ log = get_logger(__name__) +# the AI node's provider surface: llm.provider / llm.model in the app config +# (or -D overrides), params.provider / params.model per call +LLM_DEFAULT_PROVIDER = "anthropic" +LLM_DEFAULT_MODELS = {"anthropic": "claude-opus-5", "gemini": "gemini-3.6-flash"} +LLM_DEFAULT_MAX_TOKENS = 16000 +LLM_DEFAULT_TIMEOUT_MS = 60000 + +# the streaming head's content type (every streaming demo renders as SSE) +TEXT_EVENT_STREAM = "text/event-stream" + +# lazily built provider clients - module-level so tests can inject fakes +_llm_client: Any = None +_gemini_client: Any = None + @preload(route="hello.python", instances=10) def handle_event(_headers: dict[str, str], body: Body): @@ -92,7 +109,7 @@ async def stream_tokens(headers: dict[str, str], event: EventEnvelope): # "context" block (trace ids, business cid) - see the streaming guide log.info("Streaming %d messages", count) out = EventStreamWriter.from_request(event) - out.first(200, "text/event-stream") + out.first(200, TEXT_EVENT_STREAM) out.write("The following messages are rendered slowly to demonstrate streaming:") for n in range(1, count + 1): await asyncio.sleep(delay) @@ -106,6 +123,324 @@ async def stream_tokens(headers: dict[str, str], event: EventEnvelope): "my_correlation_id": headers.get("my_correlation_id")}) +@preload(route="llm.chat", instances=50) +async def llm_chat(_headers: dict[str, str], body: Body): + """The AI node (agent-orchestration experiment E0): a provider-neutral LLM + adapter as a plain wrapper-side function. The engine and this host stay + LLM-free - a graph or flow reaches this route like any other function, so + the certified graph decides control flow while the model advises within it. + + Input (map): + prompt | messages single-turn text, or conversation turns [{role, content}] + system optional system prompt + schema optional JSON schema -> structured output (the graph + needs parseable verdicts for decision routing; + additionalProperties defaults to false) + params provider, model, max_tokens, timeout_ms + provider + pass-through + + Provider selection: params.provider, else the llm.provider config key + (e.g. mercury-serve ... -Dllm.provider=gemini), else anthropic. The default + model per provider comes from params.model, the llm.model config key, or + LLM_DEFAULT_MODELS. + + Output (map): text | data, model, stop_reason (check for "refusal"), + usage {input_tokens, output_tokens}. + + Provider errors ride the envelope status - portable to a graph's error + context (error.code / error.message). The remaining time budget maps onto + the SDK timeout (params.timeout_ms), the x-ttl pattern. + """ + provider, model, max_tokens, timeout_ms, messages, system, params = _llm_request_prep(body) + assert isinstance(body, dict) # narrowed by _llm_request_prep + schema = body.get("schema") + if isinstance(schema, dict): + # structured output: a closed schema is what a bounded verdict wants, so + # default additionalProperties to false when the caller omits it + schema = {"additionalProperties": False, **schema} + else: + schema = None + if provider == "gemini": + result = await _call_gemini(model, messages, system, schema, max_tokens, timeout_ms, params) + else: + result = await _call_anthropic(model, messages, system, schema, max_tokens, timeout_ms, params) + annotate_trace("llm_model", str(result.get("model", ""))) + text = str(result.pop("text", "")) + if schema is not None and text: + # both providers guarantee schema-constrained output as one JSON text + result["data"] = json.loads(text) + else: + result["text"] = text + return result + + +def _llm_request_prep( + body: Body, +) -> tuple[str, str, int, int, Any, Any, dict[str, Any]]: + """The shared request surface of the AI nodes (llm.chat and llm.stream): + provider and model resolution (params -> llm.provider/llm.model config -> + defaults), token/time budgets and message shaping.""" + if not isinstance(body, dict) or not (body.get("prompt") or body.get("messages")): + raise AppException(400, "missing 'prompt' or 'messages'") + raw_params = body.get("params") + params: dict[str, Any] = dict(raw_params) if isinstance(raw_params, dict) else {} + provider = str( + params.pop("provider", None) + or app_config().get_property("llm.provider", LLM_DEFAULT_PROVIDER) + or LLM_DEFAULT_PROVIDER + ).lower() + if provider not in LLM_DEFAULT_MODELS: + raise AppException(400, f"unknown LLM provider '{provider}' - use one of " + f"{sorted(LLM_DEFAULT_MODELS)}") + model = str( + params.pop("model", None) + or app_config().get_property("llm.model", None) + or LLM_DEFAULT_MODELS[provider] + ) + max_tokens = int(params.pop("max_tokens", LLM_DEFAULT_MAX_TOKENS)) + timeout_ms = int(params.pop("timeout_ms", LLM_DEFAULT_TIMEOUT_MS)) + messages = body.get("messages") or [{"role": "user", "content": str(body["prompt"])}] + return provider, model, max_tokens, timeout_ms, messages, body.get("system"), params + + +async def _call_anthropic(model: str, messages: Any, system: Any, schema: dict[str, Any] | None, + max_tokens: int, timeout_ms: int, extra: dict[str, Any]) -> dict[str, Any]: + """Anthropic SDK edition of the llm.chat contract (lazy optional import).""" + try: + import anthropic + except ImportError as exc: # the SDK is an optional extra - teach, don't crash the app + raise AppException(501, "llm.chat requires the Anthropic SDK - pip install anthropic") from exc + request: dict[str, Any] = {"model": model, "max_tokens": max_tokens, "messages": messages} + if system: + request["system"] = system + if schema is not None: + request["output_config"] = {"format": {"type": "json_schema", "schema": schema}} + request.update(extra) # provider pass-through (e.g. output_config.effort) wins verbatim + global _llm_client + if _llm_client is None: + _llm_client = anthropic.AsyncAnthropic() + try: + response = await _llm_client.with_options(timeout=max(1.0, timeout_ms / 1000)) \ + .messages.create(**request) + except anthropic.RateLimitError as exc: + raise AppException(429, f"LLM provider rate limit - {exc}") from exc + except anthropic.APIStatusError as exc: + raise AppException(int(exc.status_code), f"LLM provider error - {exc}") from exc + except anthropic.APIConnectionError as exc: + raise AppException(503, f"LLM provider unreachable - {exc}") from exc + text = "" + for block in response.content: + if getattr(block, "type", "") == "text" and block.text: + text = block.text + break + return { + "text": text, + "model": response.model, + "stop_reason": str(response.stop_reason), + "usage": { + "input_tokens": response.usage.input_tokens, + "output_tokens": response.usage.output_tokens, + }, + } + + +def _gemini_request(system: Any, messages: Any, max_tokens: int, timeout_ms: int, + extra: dict[str, Any]) -> tuple[Any, list[Any]]: + """Config and role-mapped contents shared by the single-shot and streaming + Gemini editions. HttpOptions.timeout is in milliseconds - params.timeout_ms + passes through. Conversation turns map onto Gemini roles (assistant -> model). + """ + from google.genai import types as genai_types + config = genai_types.GenerateContentConfig( + max_output_tokens=max_tokens, + http_options=genai_types.HttpOptions(timeout=timeout_ms), + **extra, # provider pass-through (e.g. temperature) wins verbatim + ) + if config.automatic_function_calling is None: + # the AI nodes expose no tool surface (the graph decides, the model + # advises), so the SDK's automatic-function-calling loop is opted out - + # which also silences its AFC advisory warning on direct + # generate_content(_stream) calls + config.automatic_function_calling = genai_types.AutomaticFunctionCallingConfig(disable=True) + if system: + config.system_instruction = str(system) + contents = [ + genai_types.Content( + role="model" if turn.get("role") == "assistant" else "user", + parts=[genai_types.Part(text=str(turn.get("content", "")))], + ) + for turn in messages + if isinstance(turn, dict) + ] + return config, contents + + +def _gemini_finish(candidates: Any) -> str: + """Finish-reason name from a response/chunk's candidates, or empty.""" + if candidates: + reason = candidates[0].finish_reason + if reason is not None: + return getattr(reason, "name", str(reason)) + return "" + + +def _gemini_usage(usage: Any) -> dict[str, int]: + """Usage metadata (absent until the final stream chunk) in the contract shape.""" + return { + "input_tokens": (usage.prompt_token_count or 0) if usage else 0, + "output_tokens": (usage.candidates_token_count or 0) if usage else 0, + } + + +async def _call_gemini(model: str, messages: Any, system: Any, schema: dict[str, Any] | None, + max_tokens: int, timeout_ms: int, extra: dict[str, Any]) -> dict[str, Any]: + """Gemini SDK edition of the llm.chat contract (lazy optional import). + + The client reads GEMINI_API_KEY (or GOOGLE_API_KEY) from the environment. + """ + try: + from google import genai + from google.genai import errors as genai_errors + except ImportError as exc: + raise AppException(501, "llm.chat requires the Gemini SDK - pip install google-genai") from exc + config, contents = _gemini_request(system, messages, max_tokens, timeout_ms, extra) + if schema is not None: + config.response_mime_type = "application/json" + config.response_json_schema = schema + global _gemini_client + if _gemini_client is None: + _gemini_client = genai.Client() + try: + response = await _gemini_client.aio.models.generate_content( + model=model, contents=contents, config=config) + except genai_errors.APIError as exc: + raise AppException(int(exc.code) if exc.code else 500, f"LLM provider error - {exc}") from exc + except OSError as exc: # connection-level failures (DNS, refused, timeout) + raise AppException(503, f"LLM provider unreachable - {exc}") from exc + return { + "text": response.text or "", + "model": getattr(response, "model_version", None) or model, + "stop_reason": _gemini_finish(response.candidates), + "usage": _gemini_usage(response.usage_metadata), + } + + +@preload(route="llm.stream", instances=50, interceptor=True) +async def llm_stream(headers: dict[str, str], event: EventEnvelope): + """The streaming AI node (agent-orchestration follow-up to E0): pulls the + provider's REAL token stream and relays each token batch over the multi-shot + reply contract - a calling engine renders it progressively out its own HTTP + edge (SSE), with the same provider neutrality as llm.chat. + + Body: prompt | messages, optional system, params (provider, model, + max_tokens, timeout_ms + provider pass-through). Structured output (schema) + is deliberately not part of the streaming contract - a schema verdict is a + single-shot reply (use llm.chat). + + The terminal event's trailing metadata carries model, stop_reason, usage + and the trace/business correlation ids. + """ + out = EventStreamWriter.from_request(event) + try: + provider, model, max_tokens, timeout_ms, messages, system, params = \ + _llm_request_prep(event.body) + except AppException as exc: + out.fail(exc) + return + info = get_trace() + meta: dict[str, Any] = { + "language": "python", + "trace_id": info.trace_id if info else None, + "my_correlation_id": headers.get("my_correlation_id"), + } + log.info("Streaming tokens from %s via %s", model, provider) + if provider == "gemini": + await _stream_gemini(out, model, messages, system, max_tokens, timeout_ms, params, meta) + else: + await _stream_anthropic(out, model, messages, system, max_tokens, timeout_ms, params, meta) + + +async def _stream_anthropic(out: EventStreamWriter, model: str, messages: Any, system: Any, + max_tokens: int, timeout_ms: int, extra: dict[str, Any], + meta: dict[str, Any]) -> None: + """Anthropic SDK edition of the streaming contract (lazy optional import).""" + try: + import anthropic + except ImportError: + out.fail(AppException(501, "llm.stream requires the Anthropic SDK - pip install anthropic")) + return + request: dict[str, Any] = {"model": model, "max_tokens": max_tokens, "messages": messages} + if system: + request["system"] = system + request.update(extra) # provider pass-through wins verbatim + global _llm_client + if _llm_client is None: + _llm_client = anthropic.AsyncAnthropic() + try: + async with _llm_client.with_options(timeout=max(1.0, timeout_ms / 1000)) \ + .messages.stream(**request) as stream: + out.first(200, TEXT_EVENT_STREAM) + async for text in stream.text_stream: + if text: + out.write(text) + message = await stream.get_final_message() + except anthropic.RateLimitError as exc: + out.fail(AppException(429, f"LLM provider rate limit - {exc}")) + return + except anthropic.APIStatusError as exc: + out.fail(AppException(int(exc.status_code), f"LLM provider error - {exc}")) + return + except anthropic.APIConnectionError as exc: + out.fail(AppException(503, f"LLM provider unreachable - {exc}")) + return + out.close({ + "model": message.model, + "stop_reason": str(message.stop_reason), + "usage": { + "input_tokens": message.usage.input_tokens, + "output_tokens": message.usage.output_tokens, + }, + **meta, + }) + + +async def _stream_gemini(out: EventStreamWriter, model: str, messages: Any, system: Any, + max_tokens: int, timeout_ms: int, extra: dict[str, Any], + meta: dict[str, Any]) -> None: + """Gemini SDK edition of the streaming contract: each chunk is a token batch.""" + try: + from google import genai + from google.genai import errors as genai_errors + except ImportError: + out.fail(AppException(501, "llm.stream requires the Gemini SDK - pip install google-genai")) + return + config, contents = _gemini_request(system, messages, max_tokens, timeout_ms, extra) + global _gemini_client + if _gemini_client is None: + _gemini_client = genai.Client() + usage = None + finish = "" + version = model + try: + stream = await _gemini_client.aio.models.generate_content_stream( + model=model, contents=contents, config=config) + out.first(200, TEXT_EVENT_STREAM) + async for chunk in stream: + if chunk.text: + out.write(chunk.text) + # usage/finish arrive on the final chunk; the model version on any + usage = chunk.usage_metadata or usage + finish = _gemini_finish(chunk.candidates) or finish + version = getattr(chunk, "model_version", None) or version + except genai_errors.APIError as exc: + out.fail(AppException(int(exc.code) if exc.code else 500, f"LLM provider error - {exc}")) + return + except OSError as exc: + out.fail(AppException(503, f"LLM provider unreachable - {exc}")) + return + out.close({"model": version, "stop_reason": finish, "usage": _gemini_usage(usage), **meta}) + + @preload(route="demo.health", instances=5, private=True) async def health_check(headers: dict[str, str], _body: Body): """Health check speaking the engines' interface contract (type=info / type=health). diff --git a/pyproject.toml b/pyproject.toml index 92f252b..e5fbcd3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,9 @@ Homepage = "https://github.com/Accenture/mercury-python" Documentation = "https://accenture.github.io/mercury-composable" [project.optional-dependencies] -dev = ["pytest>=8", "pytest-asyncio>=0.23", "ruff>=0.9"] +# the llm.chat demo function (examples/demo_app.py) - the package itself stays SDK-free +llm = ["anthropic>=1,<2", "google-genai>=2,<3"] +dev = ["pytest>=8", "pytest-asyncio>=0.23", "ruff>=0.9", "anthropic>=1,<2", "google-genai>=2,<3"] [project.scripts] mercury-serve = "mercury_composable.cli:main" diff --git a/tests/test_llm_chat.py b/tests/test_llm_chat.py new file mode 100644 index 0000000..d2e3142 --- /dev/null +++ b/tests/test_llm_chat.py @@ -0,0 +1,252 @@ +"""The llm.chat demo function (agent-orchestration experiment E0) - token-free: +the provider client is a fake, so these tests pin the adapter contract (request +shaping, structured output, usage surfacing, provider-error mapping) without +spending tokens or needing credentials. +""" + +import importlib.util +import json +import sys +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import anthropic +import pytest + +from mercury_composable import AppException + +_DEMO = Path(__file__).resolve().parent.parent / "examples" / "demo_app.py" +_spec = importlib.util.spec_from_file_location("demo_app_under_test", _DEMO) +assert _spec is not None +assert _spec.loader is not None +# typed Any: the module is loaded dynamically from a file path, so its attributes +# are unknowable statically - Any tells every analyzer to trust the runtime +demo_app: Any = importlib.util.module_from_spec(_spec) +sys.modules.setdefault("demo_app_under_test", demo_app) +_spec.loader.exec_module(demo_app) + + +def _provider_response(text: str) -> SimpleNamespace: + return SimpleNamespace( + content=[SimpleNamespace(type="text", text=text)], + usage=SimpleNamespace(input_tokens=42, output_tokens=7), + stop_reason="end_turn", + model="claude-opus-5", + ) + + +class _FakeMessages: + def __init__(self, outcome: Any): + self.outcome = outcome + self.requests: list[dict[str, Any]] = [] + + async def create(self, **kwargs: Any) -> Any: + self.requests.append(kwargs) + if isinstance(self.outcome, Exception): + raise self.outcome + return self.outcome + + +class _FakeClient: + def __init__(self, outcome: Any): + self.messages = _FakeMessages(outcome) + self.timeouts: list[float] = [] + + def with_options(self, timeout: float) -> "_FakeClient": + self.timeouts.append(timeout) + return self + + +def _patch_demo(monkeypatch: pytest.MonkeyPatch, attr: str, value: Any) -> None: + """monkeypatch.setattr on the dynamically loaded demo module. The attribute + name rides through a parameter because no static analyzer can verify names + on a module loaded from a file path; monkeypatch still validates the name + at run time and restores the original on teardown.""" + monkeypatch.setattr(demo_app, attr, value) + + +def _install(monkeypatch: pytest.MonkeyPatch, outcome: Any) -> _FakeClient: + fake = _FakeClient(outcome) + _patch_demo(monkeypatch, "_llm_client", fake) + return fake + + +async def test_prompt_mode_returns_text_usage_and_stop_reason( + monkeypatch: pytest.MonkeyPatch, +) -> None: + fake = _install(monkeypatch, _provider_response("Paris")) + result = await demo_app.llm_chat({}, {"prompt": "Capital of France?"}) + assert result == { + "model": "claude-opus-5", + "stop_reason": "end_turn", + "usage": {"input_tokens": 42, "output_tokens": 7}, + "text": "Paris", + } + request = fake.messages.requests[0] + # provider defaults: the documented model and max_tokens, single-turn message shaping + assert request["model"] == demo_app.LLM_DEFAULT_MODELS["anthropic"] + assert request["max_tokens"] == demo_app.LLM_DEFAULT_MAX_TOKENS + assert request["messages"] == [{"role": "user", "content": "Capital of France?"}] + assert "output_config" not in request + # the default time budget maps onto the SDK timeout (seconds) + assert fake.timeouts == [demo_app.LLM_DEFAULT_TIMEOUT_MS / 1000] + + +async def test_schema_mode_requests_structured_output_and_parses_data( + monkeypatch: pytest.MonkeyPatch, +) -> None: + schema = { + "type": "object", + "properties": {"label": {"type": "string", "enum": ["question", "bug", "feature"]}}, + "required": ["label"], + "additionalProperties": False, + } + fake = _install(monkeypatch, _provider_response(json.dumps({"label": "bug"}))) + result = await demo_app.llm_chat( + {}, + { + "prompt": "Classify: the app crashes on save", + "system": "You are a support triage assistant.", + "schema": schema, + "params": {"model": "claude-opus-5", "max_tokens": 512, "timeout_ms": 5000}, + }, + ) + assert result["data"] == {"label": "bug"} + assert "text" not in result + request = fake.messages.requests[0] + assert request["output_config"] == {"format": {"type": "json_schema", "schema": schema}} + assert request["system"] == "You are a support triage assistant." + assert request["max_tokens"] == 512 + assert fake.timeouts == [5.0] + + +async def test_schema_defaults_to_a_closed_object(monkeypatch: pytest.MonkeyPatch) -> None: + # a bounded verdict wants a closed schema - additionalProperties defaults to + # false when the caller omits it (callers that set it keep their value) + fake = _install(monkeypatch, _provider_response(json.dumps({"label": "bug"}))) + await demo_app.llm_chat({}, {"prompt": "x", "schema": {"type": "object"}}) + sent = fake.messages.requests[0]["output_config"]["format"]["schema"] + assert sent == {"additionalProperties": False, "type": "object"} + + +async def test_conversation_turns_pass_through(monkeypatch: pytest.MonkeyPatch) -> None: + fake = _install(monkeypatch, _provider_response("hi")) + turns = [{"role": "user", "content": "hello"}] + await demo_app.llm_chat({}, {"messages": turns}) + assert fake.messages.requests[0]["messages"] == turns + + +async def test_missing_prompt_and_messages_is_a_400(monkeypatch: pytest.MonkeyPatch) -> None: + _install(monkeypatch, _provider_response("unused")) + with pytest.raises(AppException) as error: + await demo_app.llm_chat({}, {"schema": {}}) + assert error.value.status == 400 + + +def _gemini_response(text: str) -> SimpleNamespace: + return SimpleNamespace( + text=text, + model_version="gemini-3.6-flash", + usage_metadata=SimpleNamespace(prompt_token_count=11, candidates_token_count=3), + candidates=[SimpleNamespace(finish_reason=SimpleNamespace(name="STOP"))], + ) + + +class _FakeGeminiModels: + def __init__(self, outcome: Any): + self.outcome = outcome + self.requests: list[dict[str, Any]] = [] + + async def generate_content(self, **kwargs: Any) -> Any: + self.requests.append(kwargs) + if isinstance(self.outcome, Exception): + raise self.outcome + return self.outcome + + +class _FakeGemini: + def __init__(self, outcome: Any): + self.aio = SimpleNamespace(models=_FakeGeminiModels(outcome)) + + +def _install_gemini(monkeypatch: pytest.MonkeyPatch, outcome: Any) -> _FakeGemini: + fake = _FakeGemini(outcome) + _patch_demo(monkeypatch, "_gemini_client", fake) + return fake + + +async def test_gemini_provider_speaks_the_same_contract( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # provider swapped per call (or by the llm.provider config key) - the caller's + # contract and the graph above it do not change + schema = {"type": "object", "properties": {"label": {"type": "string"}}} + fake = _install_gemini(monkeypatch, _gemini_response(json.dumps({"label": "bug"}))) + result = await demo_app.llm_chat( + {}, + { + "prompt": "Classify: the app crashes on save", + "system": "You are a support triage assistant.", + "schema": schema, + "params": {"provider": "gemini", "timeout_ms": 5000}, + }, + ) + assert result == { + "model": "gemini-3.6-flash", + "stop_reason": "STOP", + "usage": {"input_tokens": 11, "output_tokens": 3}, + "data": {"label": "bug"}, + } + request = fake.aio.models.requests[0] + assert request["model"] == demo_app.LLM_DEFAULT_MODELS["gemini"] + config = request["config"] + assert config.system_instruction == "You are a support triage assistant." + assert config.response_mime_type == "application/json" + assert config.response_json_schema == {"additionalProperties": False, **schema} + # HttpOptions.timeout is milliseconds - timeout_ms passes through unchanged + assert config.http_options is not None + assert config.http_options.timeout == 5000 + contents = request["contents"] + assert len(contents) == 1 + assert contents[0].role == "user" + + +async def test_gemini_errors_map_to_envelope_status(monkeypatch: pytest.MonkeyPatch) -> None: + from google.genai import errors as genai_errors + + _install_gemini(monkeypatch, genai_errors.APIError(429, {"error": {"message": "quota"}})) + with pytest.raises(AppException) as error: + await demo_app.llm_chat({}, {"prompt": "x", "params": {"provider": "gemini"}}) + assert error.value.status == 429 + + +async def test_unknown_provider_is_a_400(monkeypatch: pytest.MonkeyPatch) -> None: + _install(monkeypatch, _provider_response("unused")) + with pytest.raises(AppException) as error: + await demo_app.llm_chat({}, {"prompt": "x", "params": {"provider": "openai"}}) + assert error.value.status == 400 + + +async def test_provider_errors_map_to_envelope_status(monkeypatch: pytest.MonkeyPatch) -> None: + # subclass the SDK exceptions so no httpx plumbing is needed - isinstance is + # what the mapping chain dispatches on + class _RateLimited(anthropic.RateLimitError): + def __init__(self) -> None: + Exception.__init__(self, "rate limited") + self.status_code = 429 + + class _Invalid(anthropic.APIStatusError): + def __init__(self) -> None: + Exception.__init__(self, "bad request") + self.status_code = 400 + + class _Unreachable(anthropic.APIConnectionError): + def __init__(self) -> None: + Exception.__init__(self, "connect timeout") + + for boom, expected in ((_RateLimited(), 429), (_Invalid(), 400), (_Unreachable(), 503)): + _install(monkeypatch, boom) + with pytest.raises(AppException) as error: + await demo_app.llm_chat({}, {"prompt": "x"}) + assert error.value.status == expected diff --git a/tests/test_llm_stream.py b/tests/test_llm_stream.py new file mode 100644 index 0000000..0941594 --- /dev/null +++ b/tests/test_llm_stream.py @@ -0,0 +1,234 @@ +"""The llm.stream demo function (progressive token rendering, E0 follow-up) - +token-free: fake provider streams and a recording writer pin the relay contract +(head, ordered token batches, terminal metadata, provider-error fail) without +spending tokens or needing credentials. +""" + +import importlib.util +import sys +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import anthropic +import pytest + +from mercury_composable import AppException + +_DEMO = Path(__file__).resolve().parent.parent / "examples" / "demo_app.py" +_spec = importlib.util.spec_from_file_location("demo_app_stream_under_test", _DEMO) +assert _spec is not None +assert _spec.loader is not None +# typed Any: the module is loaded dynamically from a file path, so its attributes +# are unknowable statically - Any tells every analyzer to trust the runtime +demo_app: Any = importlib.util.module_from_spec(_spec) +sys.modules.setdefault("demo_app_stream_under_test", demo_app) +_spec.loader.exec_module(demo_app) + + +class _FakeWriter: + def __init__(self) -> None: + self.head: tuple[int, str] | None = None + self.segments: list[Any] = [] + self.trailing: Any = None + self.error: Exception | None = None + + def first(self, status: int, content_type: str) -> None: + self.head = (status, content_type) + + def write(self, segment: Any) -> None: + self.segments.append(segment) + + def close(self, trailing_metadata: Any = None) -> None: + self.trailing = trailing_metadata + + def fail(self, error: Exception) -> None: + self.error = error + + +class _FakeWriterFactory: + instance = _FakeWriter() + + @classmethod + def from_request(cls, _event: Any) -> _FakeWriter: + return cls.instance + + +def _patch_demo(monkeypatch: pytest.MonkeyPatch, attr: str, value: Any) -> None: + """monkeypatch.setattr on the dynamically loaded demo module. The attribute + name rides through a parameter because no static analyzer can verify names + on a module loaded from a file path; monkeypatch still validates the name + at run time and restores the original on teardown.""" + monkeypatch.setattr(demo_app, attr, value) + + +def _install_writer(monkeypatch: pytest.MonkeyPatch) -> _FakeWriter: + fake = _FakeWriter() + _FakeWriterFactory.instance = fake + _patch_demo(monkeypatch, "EventStreamWriter", _FakeWriterFactory) + return fake + + +def _event(body: Any) -> SimpleNamespace: + return SimpleNamespace(body=body) + + +# --- gemini streaming path --------------------------------------------------- + + +def _gemini_chunk(text: str | None, usage: Any = None, finish: str | None = None) -> SimpleNamespace: + return SimpleNamespace( + text=text, + usage_metadata=usage, + candidates=[SimpleNamespace(finish_reason=SimpleNamespace(name=finish))] if finish else [], + model_version="gemini-3.6-flash", + ) + + +class _FakeGeminiStreamModels: + def __init__(self, outcome: Any): + self.outcome = outcome + self.requests: list[dict[str, Any]] = [] + + async def generate_content_stream(self, **kwargs: Any) -> Any: + self.requests.append(kwargs) + if isinstance(self.outcome, Exception): + raise self.outcome + chunks = list(self.outcome) + + async def gen(): + for chunk in chunks: + yield chunk + + return gen() + + +def _install_gemini_stream(monkeypatch: pytest.MonkeyPatch, outcome: Any) -> _FakeGeminiStreamModels: + models = _FakeGeminiStreamModels(outcome) + fake = SimpleNamespace(aio=SimpleNamespace(models=models)) + _patch_demo(monkeypatch, "_gemini_client", fake) + return models + + +async def test_gemini_token_batches_relay_in_order(monkeypatch: pytest.MonkeyPatch) -> None: + out = _install_writer(monkeypatch) + final_usage = SimpleNamespace(prompt_token_count=9, candidates_token_count=17) + models = _install_gemini_stream(monkeypatch, [ + _gemini_chunk("Event-"), + _gemini_chunk("driven "), + _gemini_chunk("haiku", usage=final_usage, finish="STOP"), + ]) + await demo_app.llm_stream( + {"my_correlation_id": "biz-777"}, + _event({"prompt": "haiku please", "params": {"provider": "gemini", "timeout_ms": 5000}}), + ) + assert out.error is None + assert out.head == (200, "text/event-stream") + assert out.segments == ["Event-", "driven ", "haiku"] + assert out.trailing["usage"] == {"input_tokens": 9, "output_tokens": 17} + assert out.trailing["stop_reason"] == "STOP" + assert out.trailing["model"] == "gemini-3.6-flash" + assert out.trailing["my_correlation_id"] == "biz-777" + # HttpOptions.timeout is milliseconds - timeout_ms passes through unchanged + config = models.requests[0]["config"] + assert config.http_options is not None + assert config.http_options.timeout == 5000 + + +async def test_gemini_provider_error_fails_the_stream(monkeypatch: pytest.MonkeyPatch) -> None: + from google.genai import errors as genai_errors + + out = _install_writer(monkeypatch) + _install_gemini_stream(monkeypatch, genai_errors.APIError(429, {"error": {"message": "quota"}})) + await demo_app.llm_stream({}, _event({"prompt": "x", "params": {"provider": "gemini"}})) + assert isinstance(out.error, AppException) + assert out.error.status == 429 + assert out.trailing is None + + +async def test_missing_prompt_fails_before_any_provider_call( + monkeypatch: pytest.MonkeyPatch, +) -> None: + out = _install_writer(monkeypatch) + await demo_app.llm_stream({}, _event({"params": {"provider": "gemini"}})) + assert isinstance(out.error, AppException) + assert out.error.status == 400 + assert out.head is None + assert out.segments == [] + + +# --- anthropic streaming path ------------------------------------------------ + + +class _FakeAnthropicStream: + def __init__(self, texts: list[str], final: Any): + self._texts = texts + self._final = final + + async def __aenter__(self): + return self + + async def __aexit__(self, *_exc: object) -> bool: + return False + + @property + def text_stream(self) -> Any: + async def gen(): + for text in self._texts: + yield text + + return gen() + + async def get_final_message(self) -> Any: + return self._final + + +class _FakeAnthropicStreamClient: + def __init__(self, outcome: Any, final: Any): + self.outcome = outcome + self.final = final + self.requests: list[dict[str, Any]] = [] + self.messages = self + self.timeout: float | None = None + + def with_options(self, timeout: float) -> "_FakeAnthropicStreamClient": + self.timeout = timeout + return self + + def stream(self, **kwargs: Any) -> Any: + self.requests.append(kwargs) + if isinstance(self.outcome, Exception): + raise self.outcome + return _FakeAnthropicStream(self.outcome, self.final) + + +async def test_anthropic_token_batches_relay_in_order(monkeypatch: pytest.MonkeyPatch) -> None: + out = _install_writer(monkeypatch) + final = SimpleNamespace( + model="claude-opus-5", + stop_reason="end_turn", + usage=SimpleNamespace(input_tokens=12, output_tokens=34), + ) + client = _FakeAnthropicStreamClient(["Hello ", "world"], final) + _patch_demo(monkeypatch, "_llm_client", client) + await demo_app.llm_stream({}, _event({"prompt": "greet me"})) + assert out.error is None + assert out.head == (200, "text/event-stream") + assert out.segments == ["Hello ", "world"] + assert out.trailing["usage"] == {"input_tokens": 12, "output_tokens": 34} + assert out.trailing["model"] == "claude-opus-5" + assert client.requests[0]["model"] == demo_app.LLM_DEFAULT_MODELS["anthropic"] + + +async def test_anthropic_rate_limit_fails_the_stream(monkeypatch: pytest.MonkeyPatch) -> None: + class _RateLimited(anthropic.RateLimitError): + def __init__(self) -> None: + Exception.__init__(self, "rate limited") + self.status_code = 429 + + out = _install_writer(monkeypatch) + client = _FakeAnthropicStreamClient(_RateLimited(), None) + _patch_demo(monkeypatch, "_llm_client", client) + await demo_app.llm_stream({}, _event({"prompt": "x"})) + assert isinstance(out.error, AppException) + assert out.error.status == 429