feat(cache): type prompt cache keys across chat APIs - #1247
Conversation
Expose prompt_cache_key through both MessagesParams and CompletionParams so schema-derived gateways can accept it without relying on provider kwargs. Validate provider capability centrally: OpenAI supports the field, Otari and custom OpenAI-compatible endpoints pass it through, and other providers reject it with UnsupportedParameterError.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
WalkthroughThe PR adds ChangesPrompt cache key support
Possibly related issues
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Codecov Report✅ All modified and coverable lines are covered by tests.
... and 32 files with indirect coverage changes 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unit/test_completion.py`:
- Around line 18-30: Extend test_completion_params_exposes_prompt_cache_key with
a synchronous mock-provider invocation of completion(), passing
prompt_cache_key="tenant-1". Assert the provider.completion() call receives and
forwards that value, while preserving the existing schema and signature
assertions.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 1d62df71-db69-480e-a043-42820f33f0c0
📒 Files selected for processing (14)
src/any_llm/any_llm.pysrc/any_llm/api.pysrc/any_llm/providers/openai/custom.pysrc/any_llm/providers/openai/openai.pysrc/any_llm/providers/otari/otari.pysrc/any_llm/types/completion.pysrc/any_llm/types/messages.pysrc/any_llm/utils/messages_compat.pytests/unit/providers/test_anthropic_messages.pytests/unit/providers/test_openai_base_provider.pytests/unit/providers/test_openai_compatible_provider.pytests/unit/providers/test_otari_provider.pytests/unit/test_completion.pytests/unit/test_messages.py
There was a problem hiding this comment.
Pull request overview
Adds first-class, typed prompt_cache_key support across the Chat Completions and Anthropic Messages-style APIs so schema-driven callers can safely provide it. The change introduces provider capability gating so unsupported providers fail fast with UnsupportedParameterError before any SDK or HTTP request is attempted.
Changes:
- Add
prompt_cache_key: str | NonetoCompletionParamsandMessagesParams, and thread it through the publiccompletion/acompletion/messages/amessagesAPIs. - Introduce
AnyLLM.PROMPT_CACHE_KEY_SUPPORTplus centralized validation to reject unsupported providers pre-dispatch. - Update OpenAI, Otari, and OpenAI-compatible providers’ capability flags and expand unit test coverage for forwarding and rejection behavior.
Reviewed changes
Copilot reviewed 14 out of 14 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unit/test_messages.py | Adds schema + forwarding assertions and a new unsupported-provider rejection test for Messages. |
| tests/unit/test_completion.py | Adds schema + signature assertions and a new unsupported-provider rejection test for Completions. |
| tests/unit/providers/test_otari_provider.py | Verifies Otari capability and asserts prompt_cache_key is forwarded on completion and messages paths. |
| tests/unit/providers/test_openai_compatible_provider.py | Asserts OpenAI-compatible default capability is passthrough. |
| tests/unit/providers/test_openai_base_provider.py | Adds capability-flag coverage and an HTTP transport test ensuring OpenAI sends prompt_cache_key. |
| tests/unit/providers/test_anthropic_messages.py | Adds tests asserting Anthropic rejects prompt_cache_key before any HTTP call. |
| src/any_llm/utils/messages_compat.py | Extends Messages→Completions bridge conversion to include prompt_cache_key when set. |
| src/any_llm/types/messages.py | Adds prompt_cache_key to the typed MessagesParams model and schema. |
| src/any_llm/types/completion.py | Adds prompt_cache_key to the typed CompletionParams model and schema. |
| src/any_llm/providers/otari/otari.py | Marks Otari as prompt-cache-key passthrough. |
| src/any_llm/providers/openai/openai.py | Marks OpenAI as prompt-cache-key supported. |
| src/any_llm/providers/openai/custom.py | Marks OpenAI-compatible provider as prompt-cache-key passthrough. |
| src/any_llm/api.py | Exposes prompt_cache_key as a typed argument on top-level API functions and forwards into parameter models. |
| src/any_llm/any_llm.py | Introduces capability flag + validation and threads prompt_cache_key through provider entrypoints. |
Suppressed comments (2)
tests/unit/test_messages.py:79
- This test instantiates BedrockProvider, which can raise ImportError in environments without boto3 due to AnyLLM._verify_no_missing_packages. Since the purpose here is only to verify prompt_cache_key is rejected before any client call, use a lightweight AnyLLM subclass that does not require optional dependencies and still has PROMPT_CACHE_KEY_SUPPORT left as "unsupported".
client = Mock()
provider = BedrockProvider(client=client)
with pytest.raises(UnsupportedParameterError, match="prompt_cache_key"):
await provider.amessages(
tests/unit/test_completion.py:39
- This test uses BedrockProvider to represent an unsupported provider, but BedrockProvider can raise ImportError if boto3 is not installed due to AnyLLM._verify_no_missing_packages. Since you only need to exercise AnyLLM's prompt_cache_key validation before any SDK call, use a tiny AnyLLM subclass with a mocked client instead of a provider with optional dependencies.
client = Mock()
provider = BedrockProvider(client=client)
with pytest.raises(UnsupportedParameterError, match="prompt_cache_key"):
await provider.acompletion(
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Cover synchronous completion forwarding and isolate optional Bedrock imports so general API test modules remain usable without boto3.
Description
prompt_cache_keywas available only through untyped provider kwargs on Chat Completions and Messages, so schema-derived gateways (Otari) could not expose it safely. This adds it to both typed parameter models and validates provider capability before dispatch.OpenAI accepts the field, Otari and custom OpenAI-compatible endpoints pass it to the service that resolves support, and known unsupported providers raise
UnsupportedParameterErrorbefore an SDK or HTTP call.PR Type
Relevant issues
Fixes #1243
Follow-up: mozilla-ai/otari#514
Checklist
AI Usage Information
When answering questions by the reviewer, please respond yourself, do not copy/paste the reviewer comments into an AI system and paste back its answer. We want to discuss with you, not your AI :)
Summary by CodeRabbit
New Features
Tests