scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

fix(prompt-management): compose vector store retrieval with anthropic cache control hook

Browse source 
Motivation:
With litellm.enable_anthropic_prompt_caching on, Logging.get_custom_logger_for_prompt_management
picked exactly one prompt-management hook per request via first-match-wins branching. The
AnthropicCacheControlHook branch, seeded by main.py's maybe_seed_default_injection_points, matched
before the vector store branch, so a model configured with vector_store_ids never ran its
knowledge-base retrieval. vector_store_ids was never popped from the request params and was
serialized straight into the provider body, which rejects it (e.g. Bedrock's "vector_store_ids:
Extra inputs are not permitted"). The same starvation applies whenever a model carries explicit
cache_control_injection_points, independent of the flag.

Approach:
Logging.async_get_chat_completion_prompt now also runs the vector store hook, ahead of the selected
AnthropicCacheControlHook, whenever the registry has vector stores to run for the request. Vector
store retrieval runs first so it can inject knowledge-base context and pop vector_store_ids before
the cache-control hook stamps breakpoints onto the augmented messages. Extracted the singleton
lookup and callback registration for the vector store hook into a small
_get_vector_store_pre_call_hook helper shared by the existing selector and the new composition path.
The sync get_chat_completion_prompt path is untouched: VectorStorePreCallHook only implements the
async variant, so vector store retrieval was already async-only before this change.

Validation:
Added test_prompt_hooks_compose_vector_search_with_anthropic_cache_control to
tests/test_litellm/litellm_core_utils/test_litellm_logging.py. It seeds
cache_control_injection_points the same way main.py does, registers a vector store, and asserts the
search hook is awaited, vector_store_ids and cache_control_injection_points are both consumed, the
retrieved context is prepended to the messages, and the cache_control breakpoint still lands on the
trailing user message. Confirmed the test fails on the pre-fix code (search mock never awaited) and
passes after the fix.

Ran:
- uv run pytest tests/test_litellm/litellm_core_utils/test_litellm_logging.py -v (303 passed, 1
  pre-existing failure unrelated to this change: TestAnthropicPassthroughCustomPricing::
  test_completion_cost_receives_custom_pricing_args fails on a clean tree too, from a missing
  orjson dependency in this sandbox)
- uv run pytest tests/test_litellm/integrations/test_anthropic_cache_control_hook.py
  tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py -q
  (234 passed, 2 pre-existing failures unrelated to this change: both fail identically on a clean
  tree from a missing websockets dependency in this sandbox)
- make lint (ruff, ruff-strict-gate, type-discipline-gate, test-quality-gate, type-check-gate,
  basedpyright, circular-import check): all checks passed with no new budget regressions

Report: https://github.com/BerriAI/litellm/issues/40908
Signed-off-by: Pujitha Paladugu <10557236+pujitha24@users.noreply.github.com>
Assisted-by: claude-sonnet-5 (via Claude Code)
This commit is contained in:
Pujitha Paladugu 2026-09-18 07:13:08 -07:00
parent 615ed7900f
commit e9fc70e766
No known key found for this signature in database
2 changed files with 139 additions and 28 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

79
litellm/litellm_core_utils/litellm_logging.py
View file

@ -1132,25 +1132,41 @@ class Logging(LiteLLMLoggingBaseClass):
dynamic_callback_params=self.standard_callback_dynamic_params,
)
if custom_logger:
breakpoints_before: Final = AnthropicCacheControlHook.count_request_cache_breakpoints(messages)
(
model,
messages,
non_default_params,
) = await custom_logger.async_get_chat_completion_prompt(
model=model,
messages=messages,
non_default_params=non_default_params or {},
prompt_id=prompt_id,
prompt_spec=prompt_spec,
prompt_variables=prompt_variables,
dynamic_callback_params=self.standard_callback_dynamic_params,
litellm_logging_obj=self,
tools=tools,
prompt_label=prompt_label,
prompt_version=prompt_version,
# The Anthropic cache-control hook and the vector store hook both claim the single
# prompt-management slot; a request that seeds cache_control_injection_points (e.g. via
# `enable_anthropic_prompt_caching`) would otherwise starve knowledge-base retrieval, so
# run the vector store hook first when both apply to the same request.
vector_store_logger: Final = (
self._get_vector_store_pre_call_hook()
if isinstance(custom_logger, AnthropicCacheControlHook)
and litellm.vector_store_registry is not None
and litellm.vector_store_registry.get_vector_store_ids_to_run(
non_default_params=non_default_params, tools=tools
)
else None
)
prompt_loggers: Final = tuple(logger for logger in (vector_store_logger, custom_logger) if logger is not None)
if prompt_loggers:
breakpoints_before: Final = AnthropicCacheControlHook.count_request_cache_breakpoints(messages)
for prompt_logger in prompt_loggers:
(
model,
messages,
non_default_params,
) = await prompt_logger.async_get_chat_completion_prompt(
model=model,
messages=messages,
non_default_params=non_default_params or {},
prompt_id=prompt_id,
prompt_spec=prompt_spec,
prompt_variables=prompt_variables,
dynamic_callback_params=self.standard_callback_dynamic_params,
litellm_logging_obj=self,
tools=tools,
prompt_label=prompt_label,
prompt_version=prompt_version,
)
if request_kwargs is not None:
AnthropicCacheControlHook.record_gateway_injection(
request_kwargs,
@ -1286,19 +1302,26 @@ class Logging(LiteLLMLoggingBaseClass):
# Vector Store / Knowledge Base hooks
#########################################################
if litellm.vector_store_registry is not None:
vector_store_custom_logger: Final = _init_custom_logger_compatible_class(
logging_integration="vector_store_pre_call_hook",
internal_usage_cache=None,
llm_router=None,
)
self.model_call_details["prompt_integration"] = vector_store_custom_logger.__class__.__name__
# Add to global callbacks so post-call hooks are invoked
if vector_store_custom_logger and vector_store_custom_logger not in litellm.callbacks:
litellm.logging_callback_manager.add_litellm_callback(vector_store_custom_logger)
return vector_store_custom_logger
vector_store_custom_logger: Final = self._get_vector_store_pre_call_hook()
if vector_store_custom_logger is not None:
return vector_store_custom_logger
return None
def _get_vector_store_pre_call_hook(self) -> CustomLogger | None:
vector_store_custom_logger: Final = _init_custom_logger_compatible_class(
logging_integration="vector_store_pre_call_hook",
internal_usage_cache=None,
llm_router=None,
)
if vector_store_custom_logger is None:
return None
self.model_call_details["prompt_integration"] = vector_store_custom_logger.__class__.__name__
# Add to global callbacks so post-call hooks are invoked
if vector_store_custom_logger not in litellm.callbacks:
litellm.logging_callback_manager.add_litellm_callback(vector_store_custom_logger)
return vector_store_custom_logger
def get_custom_logger_for_anthropic_cache_control_hook(self, non_default_params: dict) -> CustomLogger | None:
if non_default_params.get("cache_control_injection_points", None):
custom_logger: Final = _init_custom_logger_compatible_class(

88
tests/unit/litellm_core_utils/test_litellm_logging.py
View file

@ -7255,6 +7255,94 @@ def test_prompt_hooks_skip_prompt_managers_when_no_prompt_id(logging_obj, tmp_pa
litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.callbacks, hook)
@pytest.mark.asyncio
async def test_prompt_hooks_compose_vector_search_with_anthropic_cache_control(logging_obj, monkeypatch):
"""
Regression for https://github.com/BerriAI/litellm/issues/40908: with
`enable_anthropic_prompt_caching` on, `get_custom_logger_for_prompt_management` selects
AnthropicCacheControlHook over VectorStorePreCallHook (first-match-wins), so a model
configured with `vector_store_ids` never retrieved from its knowledge base and
`vector_store_ids` leaked into the provider request. Both hooks must now run.
"""
from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import (
VectorStorePreCallHook,
)
from litellm.litellm_core_utils import litellm_logging as logging_module
from litellm.types.vector_stores import (
LiteLLM_ManagedVectorStore,
VectorStoreResultContent,
VectorStoreSearchResponse,
VectorStoreSearchResult,
)
from litellm.vector_stores.vector_store_registry import VectorStoreRegistry
model = "bedrock/anthropic.claude-3-5-haiku-20241022-v1:0"
messages = [{"role": "user", "content": "What does the handbook say about refunds?"}]
params = {"custom_llm_provider": "bedrock", "vector_store_ids": ["vs_123"]}
monkeypatch.setattr(
litellm,
"vector_store_registry",
VectorStoreRegistry(
vector_stores=[LiteLLM_ManagedVectorStore(vector_store_id="vs_123", custom_llm_provider="bedrock")]
),
)
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
# main.py:504 seeds this key at runtime before the prompt-management hook is selected.
AnthropicCacheControlHook.maybe_seed_default_injection_points(
non_default_params=params, messages=messages, model=model, custom_llm_provider="bedrock"
)
assert "cache_control_injection_points" in params
selected = logging_obj.get_custom_logger_for_prompt_management(
model=model, non_default_params=params, tools=None, prompt_id=None
)
assert isinstance(selected, AnthropicCacheControlHook)
search: Final = AsyncMock(
return_value=VectorStoreSearchResponse(
object="vector_store.search_results.page",
search_query="What does the handbook say about refunds?",
data=[
VectorStoreSearchResult(
score=1.0,
content=[VectorStoreResultContent(text="Refunds take seven days", type="text")],
)
],
)
)
runtime = MagicMock()
runtime.llm_router.return_value = None
runtime.prisma_client.return_value = None
vector_store_hook = VectorStorePreCallHook(proxy_runtime=runtime)
previous_loggers = tuple(logging_module._in_memory_loggers)
logging_module._in_memory_loggers.clear()
logging_module._in_memory_loggers.append(vector_store_hook)
try:
with patch("litellm.vector_stores.asearch", search):
_, result_messages, remaining_params = await logging_obj.async_get_chat_completion_prompt(
model=model,
messages=messages,
non_default_params=params,
prompt_variables=None,
)
search.assert_awaited_once()
assert result_messages[0]["content"] == "Context:\n\nRefunds take seven days\n\n"
assert result_messages[1]["content"] == "What does the handbook say about refunds?"
assert result_messages[1]["cache_control"] == {"type": "ephemeral"}
assert "vector_store_ids" not in remaining_params
assert "cache_control_injection_points" not in remaining_params
finally:
litellm.logging_callback_manager.remove_callback_from_list_by_object(
litellm.callbacks, vector_store_hook, require_self=False
)
logging_module._in_memory_loggers.clear()
logging_module._in_memory_loggers.extend(previous_loggers)
def test_newrelic_dispatch_prefers_otel_v2_when_flag_on(monkeypatch):
"""With LITELLM_OTEL_V2 on and operator credentials present, the "newrelic"
callback builds the OTel v2 logger (per-team credential routing); with the