nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[feat] Improve API: streaming prompt enhancer with LLMProvider abstraction (#1258)

Browse Source 
Co-authored-by: Junda (David) Su <90978028+Davids048@users.noreply.github.com>
Co-authored-by: Matthew Noto <99706358+RandNMR73@users.noreply.github.com>
Co-authored-by: XOR-op <17672363+XOR-op@users.noreply.github.com>
Co-authored-by: Zhang Peiyuan <42993249+jzhang38@users.noreply.github.com>
This commit is contained in:
William Lin
2026-05-04 13:44:40 -07:00
committed by GitHub
co-authored by Junda Su Matthew Noto XOR-op Zhang Peiyuan
parent eb0a41528a
commit f673423b51
10 changed files with 1044 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/entrypoints/streaming/__init__.py
+6
View File
@@ -17,6 +17,10 @@ from fastvideo.entrypoints.streaming.gpu_pool import (
PoolAcquireTimeout,
SubprocessGpuPool,
)
from fastvideo.entrypoints.streaming.prompt import (
LLMProvider,
PromptEnhancer,
)
from fastvideo.entrypoints.streaming.stream import (
FragmentedMP4Chunk,
FragmentedMP4Encoder,
@@ -30,7 +34,9 @@ __all__ = [
"InMemoryBlobStore",
"InMemorySessionStore",
"InProcessGpuPool",
"LLMProvider",
"PoolAcquireTimeout",
"PromptEnhancer",
"Session",
"SessionManager",
"SessionState",
fastvideo/entrypoints/streaming/prompt/__init__.py
+36
View File
@@ -0,0 +1,36 @@
# SPDX-License-Identifier: Apache-2.0
"""Prompt pipeline for the streaming server.
* :mod:`providers` — LLM backend abstraction + built-in adapters
* :mod:`enhancer` — provider-agnostic enhance / auto-extend / rewrite
operations on top of the provider layer
All of this is optional; the streaming server runs fine without it
(PR 7.5's skeleton never invokes the enhancer). When the operator
enables ``ServeConfig.streaming.prompt.enabled``, the server routes
each ``session_init_v2`` curated prompt through ``enhance`` before the
first segment.
"""
from fastvideo.entrypoints.streaming.prompt.enhancer import (
PromptEnhancer,
PromptOperation,
)
from fastvideo.entrypoints.streaming.prompt.providers.base import (
LLMMessage,
LLMProvider,
LLMProviderError,
LLMRequest,
LLMResponse,
LLMTimeoutError,
)
__all__ = [
"LLMMessage",
"LLMProvider",
"LLMProviderError",
"LLMRequest",
"LLMResponse",
"LLMTimeoutError",
"PromptEnhancer",
"PromptOperation",
]
fastvideo/entrypoints/streaming/prompt/enhancer.py
+197
View File
@@ -0,0 +1,197 @@
# SPDX-License-Identifier: Apache-2.0
"""Provider-agnostic prompt orchestration for the streaming server.
Three operations the streaming server needs:
* ``enhance`` — polish a user prompt (add cinematic detail, fix syntax)
* ``auto_extend`` — generate a follow-on prompt for loop generation
* ``rewrite`` — rewrite a seed prompt for a user-directed rewrite flow
All three share the same orchestration: pick a provider in priority
order, submit an ``LLMRequest``, fall back to the next provider on
retryable errors, and surface a structured :class:`LLMResponse` back
to the caller.
System prompts are loaded from ``system_prompt_dir`` on construction
and can be hot-reloaded via :meth:`PromptEnhancer.reload_system_prompts`.
The streaming server's management endpoint calls that method in
response to a ``rewrite_seed_prompts_started`` frame.
"""
from __future__ import annotations
import enum
import os
from collections.abc import Sequence
from dataclasses import dataclass, replace
from fastvideo.entrypoints.streaming.prompt.providers.base import (
LLMMessage,
LLMProvider,
LLMProviderError,
LLMRequest,
LLMResponse,
)
from fastvideo.logger import init_logger
logger = init_logger(__name__)
class PromptOperation(enum.Enum):
ENHANCE = "enhance"
AUTO_EXTEND = "auto_extend"
REWRITE = "rewrite"
@dataclass
class _SystemPrompts:
enhance: str
auto_extend: str
rewrite: str
_DEFAULT_SYSTEM_PROMPTS = _SystemPrompts(
enhance=("You are a prompt enhancer for cinematic video generation. Given "
"a user prompt, produce an enhanced prompt that is more vivid, "
"specific, and concrete. Keep the subject intact; add lighting, "
"camera, and motion detail. Reply with just the enhanced prompt."),
auto_extend=("You are a video continuation assistant. Given the current "
"sequence of prompts, produce one new prompt that naturally "
"continues the sequence. Reply with just the next prompt."),
rewrite=("You are a creative prompt rewriter. Given a seed prompt, produce "
"a set of alternative prompts that explore different angles, "
"styles, and moods. Reply with one prompt per line."),
)
class PromptEnhancer:
"""Orchestrates prompt operations across a priority-ordered provider
list with structured fallback + hot-reloadable system prompts.
Usage::
enhancer = PromptEnhancer(
providers=[CerebrasProvider(), GroqProvider()],
model="gpt-oss-120b",
system_prompt_dir="/etc/fastvideo/prompts",
)
response = await enhancer.enhance("a fox running through snow")
"""
def __init__(
self,
*,
providers: Sequence[LLMProvider],
model: str,
timeout_ms: int = 20000,
temperature: float = 0.7,
max_tokens: int | None = 256,
system_prompt_dir: str | None = None,
) -> None:
if not providers:
raise ValueError("PromptEnhancer requires at least one LLMProvider")
self._providers = list(providers)
self._model = model
self._timeout_ms = timeout_ms
self._temperature = temperature
self._max_tokens = max_tokens
self._system_prompt_dir = system_prompt_dir
self._system_prompts = self._load_system_prompts()
@property
def providers(self) -> list[LLMProvider]:
return list(self._providers)
def register_provider(self, provider: LLMProvider, *, priority: int = -1) -> None:
"""Insert an additional provider. ``priority=0`` makes it primary;
``priority=-1`` (default) appends as a fallback."""
if priority < 0:
self._providers.append(provider)
else:
self._providers.insert(priority, provider)
def reload_system_prompts(self) -> None:
"""Re-read the system prompt files from ``system_prompt_dir``.
The streaming server exposes this via a management endpoint so
operators can iterate on prompt templates without restarting
workers.
"""
self._system_prompts = self._load_system_prompts()
logger.info("prompt enhancer: reloaded system prompts from %s", self._system_prompt_dir or "defaults")
async def enhance(self, prompt: str) -> LLMResponse:
return await self._run(
PromptOperation.ENHANCE,
system=self._system_prompts.enhance,
user=prompt,
)
async def auto_extend(self, prior_prompts: Sequence[str]) -> LLMResponse:
user = "\n".join(prior_prompts)
return await self._run(
PromptOperation.AUTO_EXTEND,
system=self._system_prompts.auto_extend,
user=user,
)
async def rewrite(self, seed_prompt: str) -> LLMResponse:
return await self._run(
PromptOperation.REWRITE,
system=self._system_prompts.rewrite,
user=seed_prompt,
)
async def _run(
self,
operation: PromptOperation,
*,
system: str,
user: str,
) -> LLMResponse:
request = LLMRequest(
messages=[
LLMMessage(role="system", content=system),
LLMMessage(role="user", content=user),
],
model=self._model,
max_tokens=self._max_tokens,
temperature=self._temperature,
timeout_ms=self._timeout_ms,
)
last_error: LLMProviderError | None = None
for idx, provider in enumerate(self._providers):
try:
response = await provider.complete(request)
if idx > 0:
# Mark the fallback flag without losing any other
# response fields the provider populated.
response = replace(response, fallback_used=True)
return response
except LLMProviderError as exc:
logger.warning("prompt %s: provider %s failed: %s; trying next", operation.value, provider.name, exc)
last_error = exc
if not exc.retryable:
break
assert last_error is not None
raise last_error
def _load_system_prompts(self) -> _SystemPrompts:
if not self._system_prompt_dir:
return _DEFAULT_SYSTEM_PROMPTS
return _SystemPrompts(
enhance=_read_prompt(self._system_prompt_dir, "enhance.txt", _DEFAULT_SYSTEM_PROMPTS.enhance),
auto_extend=_read_prompt(self._system_prompt_dir, "auto_extend.txt", _DEFAULT_SYSTEM_PROMPTS.auto_extend),
rewrite=_read_prompt(self._system_prompt_dir, "rewrite.txt", _DEFAULT_SYSTEM_PROMPTS.rewrite),
)
def _read_prompt(dirname: str, filename: str, default: str) -> str:
path = os.path.join(dirname, filename)
if not os.path.exists(path):
return default
with open(path, encoding="utf-8") as f:
content = f.read().strip()
return content or default
__all__ = ["PromptEnhancer", "PromptOperation"]
fastvideo/entrypoints/streaming/prompt/providers/__init__.py
+24
View File
@@ -0,0 +1,24 @@
# SPDX-License-Identifier: Apache-2.0
"""LLM provider implementations used by the prompt enhancer."""
from fastvideo.entrypoints.streaming.prompt.providers.base import (
LLMMessage,
LLMProvider,
LLMProviderError,
LLMRequest,
LLMResponse,
LLMTimeoutError,
)
from fastvideo.entrypoints.streaming.prompt.providers.cerebras import (
CerebrasProvider, )
from fastvideo.entrypoints.streaming.prompt.providers.groq import GroqProvider
__all__ = [
"CerebrasProvider",
"GroqProvider",
"LLMMessage",
"LLMProvider",
"LLMProviderError",
"LLMRequest",
"LLMResponse",
"LLMTimeoutError",
]
fastvideo/entrypoints/streaming/prompt/providers/_openai_compat.py
+101
View File
@@ -0,0 +1,101 @@
# SPDX-License-Identifier: Apache-2.0
"""Shared HTTP path for OpenAI-compatible ``/chat/completions`` providers.
Cerebras and Groq both expose the OpenAI chat-completions schema, so
the request shape, error mapping, and response decoding are identical
between them. This module centralizes that logic; the per-provider
modules stay thin (just defaults + env var wiring).
"""
from __future__ import annotations
import time
from fastvideo.entrypoints.streaming.prompt.providers.base import (
LLMProviderError,
LLMRequest,
LLMResponse,
LLMTimeoutError,
)
async def complete_openai_compatible(
*,
api_key: str | None,
api_key_hint: str,
base_url: str,
provider_name: str,
request: LLMRequest,
) -> LLMResponse:
"""Issue a chat-completions call and decode the OpenAI response."""
if not api_key:
raise LLMProviderError(
f"{provider_name} provider requires {api_key_hint} "
"(or explicit api_key=...)",
retryable=False,
)
try:
import httpx
except ImportError as exc: # pragma: no cover - optional dep
raise LLMProviderError(
f"{provider_name} provider requires httpx; install httpx",
retryable=False,
) from exc
timeout_s = (request.timeout_ms or 20000) / 1000.0
t0 = time.perf_counter()
try:
async with httpx.AsyncClient(timeout=timeout_s) as client:
response = await client.post(
f"{base_url}/chat/completions",
headers={
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
},
json={
"model": request.model,
"messages": [{
"role": m.role,
"content": m.content
} for m in request.messages],
"max_tokens": request.max_tokens,
"temperature": request.temperature,
},
)
except httpx.TimeoutException as exc:
raise LLMTimeoutError(f"{provider_name} timed out after {timeout_s}s") from exc
except httpx.HTTPError as exc:
raise LLMProviderError(f"{provider_name} HTTP error: {exc}") from exc
if response.status_code >= 400:
# 5xx and 429 (rate-limit) are retryable: another provider may
# succeed. 4xx (auth, bad-request, etc.) are client errors —
# the enhancer should stop fallback traversal.
retryable = (response.status_code >= 500 or response.status_code == 429)
raise LLMProviderError(
f"{provider_name} returned {response.status_code}: "
f"{response.text[:200]}",
retryable=retryable,
)
try:
data = response.json()
except Exception as exc:
# Non-JSON body usually means a proxy / load-balancer error
# page; leave it retryable so a fallback provider can try.
raise LLMProviderError(f"{provider_name} returned non-JSON body: {exc}") from exc
choices = data.get("choices") or []
if not choices:
raise LLMProviderError(f"{provider_name} returned no choices")
content = choices[0].get("message", {}).get("content") or ""
latency_ms = (time.perf_counter() - t0) * 1000.0
return LLMResponse(
content=content.strip(),
provider=provider_name,
model=request.model,
latency_ms=latency_ms,
)
__all__ = ["complete_openai_compatible"]
fastvideo/entrypoints/streaming/prompt/providers/base.py
+85
View File
@@ -0,0 +1,85 @@
# SPDX-License-Identifier: Apache-2.0
"""LLM provider protocol + DTOs used by the prompt enhancer.
Third-party users add a new provider by implementing
:class:`LLMProvider` and registering it with a prompt enhancer
instance. The shipped providers live in sibling modules
(``cerebras.py``, ``groq.py``) and each is ~100-200 LOC — the
provider layer is intentionally thin so the enhancer stays
provider-agnostic.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Literal, Protocol, runtime_checkable
@dataclass
class LLMMessage:
role: Literal["system", "user", "assistant"]
content: str
@dataclass
class LLMRequest:
messages: list[LLMMessage]
model: str
max_tokens: int | None = None
temperature: float | None = None
timeout_ms: int | None = None
@dataclass
class LLMResponse:
content: str
provider: str
model: str
latency_ms: float
fallback_used: bool = False
class LLMProviderError(RuntimeError):
"""Raised when an LLM provider fails a request.
``retryable`` controls whether the enhancer falls back to the next
provider. It is settable per-instance so the same exception type
can describe retryable transport errors (5xx, 429) and
non-retryable client errors (4xx auth/bad-request) without forcing
a separate subclass for every status family.
"""
def __init__(self, message: str, *, retryable: bool = True) -> None:
super().__init__(message)
self.retryable = retryable
class LLMTimeoutError(LLMProviderError):
"""Raised when an LLM provider times out — always retryable."""
def __init__(self, message: str) -> None:
super().__init__(message, retryable=True)
@runtime_checkable
class LLMProvider(Protocol):
"""Provider interface every LLM adapter implements.
Providers are async-first because every built-in implementation
talks to an HTTP API. Synchronous providers can wrap their call in
``asyncio.to_thread`` internally.
"""
name: str
async def complete(self, request: LLMRequest) -> LLMResponse:
...
__all__ = [
"LLMMessage",
"LLMProvider",
"LLMProviderError",
"LLMRequest",
"LLMResponse",
"LLMTimeoutError",
]
fastvideo/entrypoints/streaming/prompt/providers/cerebras.py
+44
View File
@@ -0,0 +1,44 @@
# SPDX-License-Identifier: Apache-2.0
"""Cerebras LLM provider (OpenAI-compatible chat endpoint)."""
from __future__ import annotations
import os
from dataclasses import dataclass
from fastvideo.entrypoints.streaming.prompt.providers._openai_compat import (
complete_openai_compatible, )
from fastvideo.entrypoints.streaming.prompt.providers.base import (
LLMRequest,
LLMResponse,
)
_DEFAULT_BASE_URL = "https://api.cerebras.ai/v1"
_API_KEY_ENV = "CEREBRAS_API_KEY"
@dataclass
class CerebrasProvider:
"""Cerebras inference adapter.
``api_key`` falls back to ``CEREBRAS_API_KEY`` when unset.
"""
api_key: str | None = None
base_url: str = _DEFAULT_BASE_URL
name: str = "cerebras"
def __post_init__(self) -> None:
if self.api_key is None:
self.api_key = os.environ.get(_API_KEY_ENV)
async def complete(self, request: LLMRequest) -> LLMResponse:
return await complete_openai_compatible(
api_key=self.api_key,
api_key_hint=_API_KEY_ENV,
base_url=self.base_url,
provider_name=self.name,
request=request,
)
__all__ = ["CerebrasProvider"]
fastvideo/entrypoints/streaming/prompt/providers/groq.py
+46
View File
@@ -0,0 +1,46 @@
# SPDX-License-Identifier: Apache-2.0
"""Groq LLM provider (OpenAI-compatible chat endpoint)."""
from __future__ import annotations
import os
from dataclasses import dataclass
from fastvideo.entrypoints.streaming.prompt.providers._openai_compat import (
complete_openai_compatible, )
from fastvideo.entrypoints.streaming.prompt.providers.base import (
LLMRequest,
LLMResponse,
)
_DEFAULT_BASE_URL = "https://api.groq.com/openai/v1"
_API_KEY_ENV = "GROQ_API_KEY"
@dataclass
class GroqProvider:
"""Groq inference adapter.
Identical wire format to :class:`CerebrasProvider`; both go through
:func:`complete_openai_compatible`. The two providers differ only
in base URL, env var, and model id conventions.
"""
api_key: str | None = None
base_url: str = _DEFAULT_BASE_URL
name: str = "groq"
def __post_init__(self) -> None:
if self.api_key is None:
self.api_key = os.environ.get(_API_KEY_ENV)
async def complete(self, request: LLMRequest) -> LLMResponse:
return await complete_openai_compatible(
api_key=self.api_key,
api_key_hint=_API_KEY_ENV,
base_url=self.base_url,
provider_name=self.name,
request=request,
)
__all__ = ["GroqProvider"]
fastvideo/tests/entrypoints/streaming/test_prompt_enhancer.py
+207
View File
@@ -0,0 +1,207 @@
# SPDX-License-Identifier: Apache-2.0
"""Tests for the provider-agnostic prompt enhancer."""
from __future__ import annotations
import asyncio
from dataclasses import dataclass
import pytest
from fastvideo.entrypoints.streaming.prompt import (
LLMProvider,
LLMProviderError,
LLMRequest,
LLMResponse,
PromptEnhancer,
)
@dataclass
class _StaticProvider:
name: str
content: str
async def complete(self, request: LLMRequest) -> LLMResponse:
return LLMResponse(
content=self.content,
provider=self.name,
model=request.model,
latency_ms=1.0,
)
@dataclass
class _FailingProvider:
name: str
message: str = "boom"
retryable: bool = True
async def complete(self, request: LLMRequest) -> LLMResponse:
raise LLMProviderError(self.message, retryable=self.retryable)
class TestConstruction:
def test_requires_at_least_one_provider(self):
with pytest.raises(ValueError):
PromptEnhancer(providers=[], model="m")
def test_registers_system_prompts_from_defaults(self, tmp_path):
enh = PromptEnhancer(
providers=[_StaticProvider("p", "ok")],
model="m",
)
# Defaults live inside _DEFAULT_SYSTEM_PROMPTS; no hot-reload file
# means the enhancer falls back to the shipped values.
assert "prompt enhancer" in enh._system_prompts.enhance.lower()
def test_reads_override_files_when_present(self, tmp_path):
(tmp_path / "enhance.txt").write_text("customized enhance")
(tmp_path / "auto_extend.txt").write_text("") # empty -> default
enh = PromptEnhancer(
providers=[_StaticProvider("p", "ok")],
model="m",
system_prompt_dir=str(tmp_path),
)
assert enh._system_prompts.enhance == "customized enhance"
# Empty file fell back to the default.
assert "continuation assistant" in enh._system_prompts.auto_extend
class TestEnhance:
def test_enhance_returns_primary_provider_response(self):
enh = PromptEnhancer(
providers=[_StaticProvider("primary", "enhanced!")],
model="m",
)
response = asyncio.run(enh.enhance("a fox"))
assert response.content == "enhanced!"
assert response.provider == "primary"
assert response.fallback_used is False
def test_enhance_falls_back_on_provider_error(self):
enh = PromptEnhancer(
providers=[
_FailingProvider("a"),
_StaticProvider("b", "fallback!"),
],
model="m",
)
response = asyncio.run(enh.enhance("x"))
assert response.provider == "b"
assert response.content == "fallback!"
assert response.fallback_used is True
def test_enhance_stops_on_non_retryable_error(self):
enh = PromptEnhancer(
providers=[
_FailingProvider("a", message="hard-fail", retryable=False),
_StaticProvider("b", "should-not-be-reached"),
],
model="m",
)
with pytest.raises(LLMProviderError, match="hard-fail"):
asyncio.run(enh.enhance("x"))
def test_enhance_raises_when_all_providers_fail(self):
enh = PromptEnhancer(
providers=[
_FailingProvider("a"),
_FailingProvider("b"),
],
model="m",
)
with pytest.raises(LLMProviderError):
asyncio.run(enh.enhance("x"))
class TestAutoExtendAndRewrite:
def test_auto_extend_joins_prior_prompts(self):
captured: list[str] = []
class _Capturer:
name = "cap"
async def complete(self, request: LLMRequest) -> LLMResponse:
user = next(m.content for m in request.messages
if m.role == "user")
captured.append(user)
return LLMResponse(
content="next",
provider=self.name,
model=request.model,
latency_ms=1.0,
)
enh = PromptEnhancer(providers=[_Capturer()], model="m")
asyncio.run(enh.auto_extend(["first", "second"]))
assert captured[0] == "first\nsecond"
def test_rewrite_passes_seed_through(self):
captured: list[str] = []
class _Capturer:
name = "cap"
async def complete(self, request: LLMRequest) -> LLMResponse:
user = next(m.content for m in request.messages
if m.role == "user")
captured.append(user)
return LLMResponse(
content="one\ntwo\nthree",
provider=self.name,
model=request.model,
latency_ms=1.0,
)
enh = PromptEnhancer(providers=[_Capturer()], model="m")
response = asyncio.run(enh.rewrite("seed prompt"))
assert captured == ["seed prompt"]
assert response.content.splitlines() == ["one", "two", "three"]
class TestRegisterProvider:
def test_register_appends_by_default(self):
enh = PromptEnhancer(
providers=[_StaticProvider("primary", "a")],
model="m",
)
enh.register_provider(_StaticProvider("extra", "b"))
assert [p.name for p in enh.providers] == ["primary", "extra"]
def test_register_priority_zero_makes_primary(self):
enh = PromptEnhancer(
providers=[_StaticProvider("old", "a")],
model="m",
)
enh.register_provider(
_StaticProvider("new-primary", "b"), priority=0)
assert enh.providers[0].name == "new-primary"
def test_registered_provider_is_used_in_enhance(self):
enh = PromptEnhancer(
providers=[_FailingProvider("broken")],
model="m",
)
enh.register_provider(_StaticProvider("fallback", "ok"))
response = asyncio.run(enh.enhance("x"))
assert response.content == "ok"
assert response.fallback_used is True
class TestHotReload:
def test_reload_picks_up_new_file(self, tmp_path):
(tmp_path / "enhance.txt").write_text("first version")
enh = PromptEnhancer(
providers=[_StaticProvider("p", "ok")],
model="m",
system_prompt_dir=str(tmp_path),
)
assert enh._system_prompts.enhance == "first version"
(tmp_path / "enhance.txt").write_text("second version")
enh.reload_system_prompts()
assert enh._system_prompts.enhance == "second version"
fastvideo/tests/entrypoints/streaming/test_prompt_providers.py
+298
View File
@@ -0,0 +1,298 @@
# SPDX-License-Identifier: Apache-2.0
"""Tests for the LLM provider protocol + built-in adapters.
Real Cerebras/Groq HTTP calls are stubbed with a fake ``httpx`` module
inserted into ``sys.modules`` so the unit tests don't depend on
external API availability or paid keys.
"""
from __future__ import annotations
import asyncio
import sys
import types
from dataclasses import dataclass
from typing import Any
import pytest
from fastvideo.entrypoints.streaming.prompt.providers import (
CerebrasProvider,
GroqProvider,
LLMMessage,
LLMProvider,
LLMProviderError,
LLMRequest,
LLMTimeoutError,
)
@dataclass
class _FakeResponse:
status_code: int
payload: dict[str, Any]
text: str = ""
def json(self) -> dict[str, Any]:
return self.payload
class _FakeAsyncClient:
def __init__(self, response_or_exc, *, captured: list) -> None:
self._response_or_exc = response_or_exc
self._captured = captured
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, tb):
return None
async def post(self, url, **kwargs):
self._captured.append({"url": url, **kwargs})
if isinstance(self._response_or_exc, Exception):
raise self._response_or_exc
return self._response_or_exc
def _install_fake_httpx(
monkeypatch,
*,
response: _FakeResponse | None = None,
exception: Exception | None = None,
) -> list[dict]:
"""Replace ``sys.modules['httpx']`` with a stub exposing the bits
providers touch: ``AsyncClient``, ``HTTPError``, ``TimeoutException``.
Returns a list the test can inspect to see what requests went out.
"""
captured: list[dict] = []
class _HTTPError(Exception):
pass
class _TimeoutException(_HTTPError):
pass
payload = response if exception is None else exception
def _client_factory(*_args, **_kwargs):
return _FakeAsyncClient(payload, captured=captured)
stub = types.SimpleNamespace(
AsyncClient=_client_factory,
HTTPError=_HTTPError,
TimeoutException=_TimeoutException,
)
monkeypatch.setitem(sys.modules, "httpx", stub)
return captured
# ----------------------------------------------------------------------
# Cerebras
# ----------------------------------------------------------------------
class TestCerebrasProvider:
def test_is_llm_provider(self):
assert isinstance(CerebrasProvider(api_key="x"), LLMProvider)
def test_requires_api_key(self, monkeypatch):
monkeypatch.delenv("CEREBRAS_API_KEY", raising=False)
provider = CerebrasProvider()
with pytest.raises(LLMProviderError, match="CEREBRAS_API_KEY"):
asyncio.run(provider.complete(
LLMRequest(messages=[], model="m")))
def test_api_key_from_env(self, monkeypatch):
monkeypatch.setenv("CEREBRAS_API_KEY", "from-env")
provider = CerebrasProvider()
assert provider.api_key == "from-env"
def test_explicit_api_key_wins(self, monkeypatch):
monkeypatch.setenv("CEREBRAS_API_KEY", "from-env")
provider = CerebrasProvider(api_key="explicit")
assert provider.api_key == "explicit"
def test_success_path(self, monkeypatch):
captured = _install_fake_httpx(
monkeypatch,
response=_FakeResponse(
status_code=200,
payload={
"choices": [{"message": {"content": "enhanced"}}],
},
),
)
provider = CerebrasProvider(api_key="secret")
result = asyncio.run(provider.complete(
LLMRequest(
messages=[LLMMessage(role="user", content="a fox")],
model="gpt-oss-120b",
max_tokens=64,
temperature=0.7,
)))
assert result.content == "enhanced"
assert result.provider == "cerebras"
assert result.model == "gpt-oss-120b"
# Verify the HTTP request was shaped correctly.
body = captured[0]["json"]
assert body["model"] == "gpt-oss-120b"
assert body["messages"] == [{"role": "user", "content": "a fox"}]
assert captured[0]["headers"]["Authorization"] == "Bearer secret"
def test_http_error_wrapped(self, monkeypatch):
# Stage the stub first, then raise that stub's own HTTPError so
# the provider's ``except httpx.HTTPError`` catches it.
stub = types.SimpleNamespace()
class _HTTPError(Exception):
pass
class _TimeoutException(_HTTPError):
pass
stub.HTTPError = _HTTPError
stub.TimeoutException = _TimeoutException
stub.AsyncClient = lambda *a, **k: _FakeAsyncClient(
_HTTPError("boom"), captured=[])
monkeypatch.setitem(sys.modules, "httpx", stub)
provider = CerebrasProvider(api_key="k")
with pytest.raises(LLMProviderError, match="HTTP"):
asyncio.run(provider.complete(
LLMRequest(messages=[], model="m")))
def test_timeout_wrapped(self, monkeypatch):
_install_fake_httpx(monkeypatch)
# Install timeout exception after stubbing.
stub = sys.modules["httpx"]
def raising_factory(*_a, **_kw):
raise stub.TimeoutException("timed out") # type: ignore[attr-defined]
stub.AsyncClient = raising_factory # type: ignore[attr-defined]
provider = CerebrasProvider(api_key="k")
with pytest.raises(LLMTimeoutError):
asyncio.run(provider.complete(
LLMRequest(messages=[], model="m")))
def test_4xx_raises_non_retryable_provider_error(self, monkeypatch):
_install_fake_httpx(
monkeypatch,
response=_FakeResponse(
status_code=401,
payload={},
text="unauthorized",
),
)
provider = CerebrasProvider(api_key="k")
with pytest.raises(LLMProviderError, match="401") as excinfo:
asyncio.run(provider.complete(
LLMRequest(messages=[], model="m")))
assert excinfo.value.retryable is False
def test_429_raises_retryable_provider_error(self, monkeypatch):
_install_fake_httpx(
monkeypatch,
response=_FakeResponse(
status_code=429,
payload={},
text="rate limited",
),
)
provider = CerebrasProvider(api_key="k")
with pytest.raises(LLMProviderError, match="429") as excinfo:
asyncio.run(provider.complete(
LLMRequest(messages=[], model="m")))
assert excinfo.value.retryable is True
def test_5xx_raises_retryable_provider_error(self, monkeypatch):
_install_fake_httpx(
monkeypatch,
response=_FakeResponse(
status_code=503,
payload={},
text="service unavailable",
),
)
provider = CerebrasProvider(api_key="k")
with pytest.raises(LLMProviderError, match="503") as excinfo:
asyncio.run(provider.complete(
LLMRequest(messages=[], model="m")))
assert excinfo.value.retryable is True
def test_non_json_body_raises_provider_error(self, monkeypatch):
# Simulate a proxy/load-balancer HTML error page slipping in
# with a 200 status: response.json() raises, and the provider
# must wrap it in an LLMProviderError instead of bubbling.
class _BadJsonResponse:
status_code = 200
text = "<html>oops</html>"
def json(self):
raise ValueError("Expecting value: line 1 column 1 (char 0)")
_install_fake_httpx(
monkeypatch,
response=_BadJsonResponse(), # type: ignore[arg-type]
)
provider = CerebrasProvider(api_key="k")
with pytest.raises(LLMProviderError, match="non-JSON"):
asyncio.run(provider.complete(
LLMRequest(messages=[], model="m")))
def test_empty_choices_raises(self, monkeypatch):
_install_fake_httpx(
monkeypatch,
response=_FakeResponse(
status_code=200, payload={"choices": []}),
)
provider = CerebrasProvider(api_key="k")
with pytest.raises(LLMProviderError, match="no choices"):
asyncio.run(provider.complete(
LLMRequest(messages=[], model="m")))
# ----------------------------------------------------------------------
# Groq
# ----------------------------------------------------------------------
class TestGroqProvider:
def test_is_llm_provider(self):
assert isinstance(GroqProvider(api_key="x"), LLMProvider)
def test_requires_api_key(self, monkeypatch):
monkeypatch.delenv("GROQ_API_KEY", raising=False)
provider = GroqProvider()
with pytest.raises(LLMProviderError, match="GROQ_API_KEY"):
asyncio.run(provider.complete(
LLMRequest(messages=[], model="m")))
def test_api_key_from_env(self, monkeypatch):
monkeypatch.setenv("GROQ_API_KEY", "from-env")
provider = GroqProvider()
assert provider.api_key == "from-env"
def test_success_path(self, monkeypatch):
captured = _install_fake_httpx(
monkeypatch,
response=_FakeResponse(
status_code=200,
payload={
"choices": [{"message": {"content": "groq out"}}],
},
),
)
provider = GroqProvider(api_key="secret")
result = asyncio.run(provider.complete(
LLMRequest(
messages=[LLMMessage(role="user", content="a deer")],
model="llama-3.1-70b",
)))
assert result.content == "groq out"
assert result.provider == "groq"
assert captured[0]["headers"]["Authorization"] == "Bearer secret"