diff --git a/fastvideo/entrypoints/streaming/__init__.py b/fastvideo/entrypoints/streaming/__init__.py index afb3a1c1b..4c538d1e5 100644 --- a/fastvideo/entrypoints/streaming/__init__.py +++ b/fastvideo/entrypoints/streaming/__init__.py @@ -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", diff --git a/fastvideo/entrypoints/streaming/prompt/__init__.py b/fastvideo/entrypoints/streaming/prompt/__init__.py new file mode 100644 index 000000000..dc83fa70e --- /dev/null +++ b/fastvideo/entrypoints/streaming/prompt/__init__.py @@ -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", +] diff --git a/fastvideo/entrypoints/streaming/prompt/enhancer.py b/fastvideo/entrypoints/streaming/prompt/enhancer.py new file mode 100644 index 000000000..365671110 --- /dev/null +++ b/fastvideo/entrypoints/streaming/prompt/enhancer.py @@ -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"] diff --git a/fastvideo/entrypoints/streaming/prompt/providers/__init__.py b/fastvideo/entrypoints/streaming/prompt/providers/__init__.py new file mode 100644 index 000000000..27e5ee66a --- /dev/null +++ b/fastvideo/entrypoints/streaming/prompt/providers/__init__.py @@ -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", +] diff --git a/fastvideo/entrypoints/streaming/prompt/providers/_openai_compat.py b/fastvideo/entrypoints/streaming/prompt/providers/_openai_compat.py new file mode 100644 index 000000000..450bdd619 --- /dev/null +++ b/fastvideo/entrypoints/streaming/prompt/providers/_openai_compat.py @@ -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"] diff --git a/fastvideo/entrypoints/streaming/prompt/providers/base.py b/fastvideo/entrypoints/streaming/prompt/providers/base.py new file mode 100644 index 000000000..d803da838 --- /dev/null +++ b/fastvideo/entrypoints/streaming/prompt/providers/base.py @@ -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", +] diff --git a/fastvideo/entrypoints/streaming/prompt/providers/cerebras.py b/fastvideo/entrypoints/streaming/prompt/providers/cerebras.py new file mode 100644 index 000000000..5ff4fc062 --- /dev/null +++ b/fastvideo/entrypoints/streaming/prompt/providers/cerebras.py @@ -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"] diff --git a/fastvideo/entrypoints/streaming/prompt/providers/groq.py b/fastvideo/entrypoints/streaming/prompt/providers/groq.py new file mode 100644 index 000000000..37f486a82 --- /dev/null +++ b/fastvideo/entrypoints/streaming/prompt/providers/groq.py @@ -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"] diff --git a/fastvideo/tests/entrypoints/streaming/test_prompt_enhancer.py b/fastvideo/tests/entrypoints/streaming/test_prompt_enhancer.py new file mode 100644 index 000000000..e9bb5bc7c --- /dev/null +++ b/fastvideo/tests/entrypoints/streaming/test_prompt_enhancer.py @@ -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" diff --git a/fastvideo/tests/entrypoints/streaming/test_prompt_providers.py b/fastvideo/tests/entrypoints/streaming/test_prompt_providers.py new file mode 100644 index 000000000..a490e231e --- /dev/null +++ b/fastvideo/tests/entrypoints/streaming/test_prompt_providers.py @@ -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 = "oops" + + 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"