mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-15 09:12:18 +00:00
2be96aed13
The request scope keeps only the native batches. RedisBatch, its ops, the cluster fallback, the Python post-call deadline and the DualCache post-call helpers had no production caller left, and the pre-call SET, DEL and miss tracking that UserApiKeyCache never reached go with them. The mixed-families pipeline metadata had no producer left
273 lines
12 KiB
Python
273 lines
12 KiB
Python
"""``with_service_target`` carries the purpose of a datastore call to code that cannot see it from its own
|
|
frames, and every Redis producer on the proxy request path declares a key family so no request-path span renders
|
|
as a bare ``redis.get``."""
|
|
|
|
import ast
|
|
import asyncio
|
|
import contextvars
|
|
import re
|
|
from collections.abc import Generator
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Final
|
|
|
|
from litellm._internal_context import (
|
|
current_request_parent_otel_span,
|
|
current_service_target,
|
|
request_parent_otel_span,
|
|
service_target,
|
|
set_request_parent_otel_span,
|
|
with_service_target,
|
|
)
|
|
|
|
_REPO: Final = Path(__file__).resolve().parents[2]
|
|
|
|
_REDIS_PRODUCER_ROOTS: Final = ("litellm", "enterprise")
|
|
# The cache implementations and facades: they emit the service events, their callers declare the family.
|
|
_CACHE_LAYER_DIRS: Final = ("litellm/caching",)
|
|
# Helpers that act on a cache handed in by the declaring caller, forward to the response-cache facade, or call a
|
|
# ``litellm.redis.stores`` store, which declares its own family.
|
|
_CACHE_PARAMETER_HELPERS: Final = frozenset(
|
|
{
|
|
"litellm/llms/anthropic/prompt_cache_prediction.py",
|
|
"litellm/proxy/common_utils/cache_coordinator.py",
|
|
"litellm/proxy/common_utils/user_api_key_cache.py",
|
|
"litellm/router_rust/routing_stores.py",
|
|
"litellm/router_utils/routing_prefetch.py",
|
|
"litellm/utils.py",
|
|
}
|
|
)
|
|
# Callers whose every cache call hits a process-local ``InMemoryCache`` (a ``DualCache`` built without
|
|
# ``redis_cache``, a ``local_only=True`` call, the client / logger / tool-name caches), so no Redis span exists.
|
|
_IN_MEMORY_ONLY_CALLERS: Final = frozenset(
|
|
{
|
|
"litellm/integrations/datadog/datadog_team_handler.py",
|
|
"litellm/integrations/humanloop.py",
|
|
"litellm/integrations/langfuse/langfuse_handler.py",
|
|
"litellm/integrations/langfuse/langfuse_prompt_management.py",
|
|
"litellm/integrations/newrelic/newrelic_team_handler.py",
|
|
"litellm/integrations/shadow_eval_logger.py",
|
|
"litellm/litellm_core_utils/litellm_logging.py",
|
|
"litellm/litellm_core_utils/oauth_token_exchange.py",
|
|
"litellm/litellm_core_utils/prompt_templates/factory.py",
|
|
"litellm/litellm_core_utils/prompt_templates/image_handling.py",
|
|
"litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py",
|
|
"litellm/litellm_core_utils/specialty_caches/service_trace_id_cache.py",
|
|
"litellm/llms/azure/common_utils.py",
|
|
"litellm/llms/bedrock/base_aws_llm.py",
|
|
"litellm/llms/custom_httpx/http_handler.py",
|
|
"litellm/llms/databricks/common_utils.py",
|
|
"litellm/llms/gigachat/authenticator.py",
|
|
"litellm/llms/github_copilot/per_user_auth.py",
|
|
"litellm/llms/github_copilot/authenticator.py",
|
|
"litellm/llms/litellm_proxy/skills/handler.py",
|
|
"litellm/llms/openai/common_utils.py",
|
|
"litellm/llms/openai_like/model_info.py",
|
|
"litellm/llms/vertex_ai/vertex_ai_non_gemini.py",
|
|
"litellm/llms/watsonx/common_utils.py",
|
|
"litellm/proxy/mcp_server/byok_credential_cache.py",
|
|
"litellm/proxy/mcp_server/catalog.py",
|
|
"litellm/proxy/mcp_server/discoverable_endpoints.py",
|
|
"litellm/proxy/mcp_server/oauth_identity_binding.py",
|
|
"litellm/proxy/mcp_server/operations.py",
|
|
"litellm/proxy/mcp_server/outbound_credentials/sso_assertion_store.py",
|
|
"litellm/proxy/mcp_server/outbound_credentials/token_endpoint.py",
|
|
"litellm/proxy/agent_endpoints/databricks_oauth.py",
|
|
"litellm/proxy/common_utils/registry_read_through.py",
|
|
"litellm/proxy/container_endpoints/ownership.py",
|
|
"litellm/proxy/discovery_endpoints/agent_skills_endpoints.py",
|
|
"litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py",
|
|
"litellm/proxy/hooks/model_budget_shared_windows.py",
|
|
"litellm/proxy/spend_tracking/key_metadata_recovery.py",
|
|
"litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py",
|
|
"litellm/responses/litellm_completion_transformation/transformation.py",
|
|
"litellm/router_rust/surface/model_info.py",
|
|
"litellm/router_utils/client_initalization_utils.py",
|
|
"litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py",
|
|
"litellm/router_utils/router_callbacks/track_deployment_metrics.py",
|
|
"litellm/secret_managers/cyberark_secret_manager.py",
|
|
"litellm/secret_managers/google_secret_manager.py",
|
|
"litellm/secret_managers/hashicorp_secret_manager.py",
|
|
"litellm/secret_managers/main.py",
|
|
}
|
|
)
|
|
|
|
_CACHE_CALL: Final = re.compile(
|
|
r"\.(?:async_)?(?:get_cache|set_cache|batch_get_cache|batch_get_cache_shared|increment_cache|increment_if_present|increment"
|
|
r"|set_cache_pipeline|set_cache_pipeline_with_ttls|set_cache_sadd|delete_cache|batch_set_cache|increment_pipeline"
|
|
r"|rpush|lpop|scan_iter|get_ttl|mget)\("
|
|
r"|\b_prepare_batch_get\("
|
|
)
|
|
_DECLARES_TARGET: Final = re.compile(r"\b(?:with_service_target|service_target|response_cache_phase)\(")
|
|
_BUILDS_A_REDIS_CACHE: Final = re.compile(r"\bRedisCache\(|\bredis_cache=(?!None\b)")
|
|
|
|
|
|
def _redis_producers() -> tuple[str, ...]:
|
|
files: Final = tuple(
|
|
path for root in _REDIS_PRODUCER_ROOTS for path in sorted((_REPO / root).rglob("*.py"))
|
|
) # comprehension-ok: flatten the producer roots
|
|
relative: Final = tuple(
|
|
path.relative_to(_REPO).as_posix() for path in files if _CACHE_CALL.search(path.read_text())
|
|
)
|
|
return tuple(name for name in relative if not name.startswith(_CACHE_LAYER_DIRS))
|
|
|
|
|
|
def test_every_redis_producer_declares_a_key_family() -> None:
|
|
"""A module that reads or writes a shared cache without a declared target renders as a
|
|
bare ``redis.get`` / ``redis.mget`` (flat under the request span, or an unnamed INTERNAL root
|
|
for a background job), which is exactly what the sensitive-data pin read, the rate-limiter
|
|
MGET and the budget-reset job did in production. Only process-local callers are exempt."""
|
|
exempt: Final = _CACHE_PARAMETER_HELPERS | _IN_MEMORY_ONLY_CALLERS
|
|
undeclared: Final = tuple(
|
|
name
|
|
for name in _redis_producers()
|
|
if name not in exempt and not _DECLARES_TARGET.search((_REPO / name).read_text())
|
|
)
|
|
assert undeclared == ()
|
|
|
|
|
|
def test_every_in_memory_exemption_still_only_touches_a_process_local_cache() -> None:
|
|
"""The exemption list is a claim about each file, so a file that is deleted or starts building
|
|
or receiving a ``RedisCache`` has to leave the list (and declare a family) rather than stay exempt."""
|
|
producers: Final = frozenset(_redis_producers())
|
|
stale: Final = tuple(sorted(_IN_MEMORY_ONLY_CALLERS - producers))
|
|
assert stale == ()
|
|
redis_backed: Final = tuple(
|
|
name for name in sorted(_IN_MEMORY_ONLY_CALLERS) if _BUILDS_A_REDIS_CACHE.search((_REPO / name).read_text())
|
|
)
|
|
assert redis_backed == ()
|
|
|
|
|
|
def test_with_service_target_sets_the_target_for_sync_and_async_calls_and_restores_it() -> None:
|
|
@with_service_target("rate_limits")
|
|
def read() -> str | None:
|
|
return current_service_target()
|
|
|
|
@with_service_target("rate_limits")
|
|
async def read_async() -> str | None:
|
|
await asyncio.sleep(0)
|
|
return current_service_target()
|
|
|
|
assert read() == "rate_limits"
|
|
assert asyncio.run(read_async()) == "rate_limits"
|
|
assert current_service_target() is None
|
|
with service_target("auth_objects"):
|
|
assert read() == "rate_limits"
|
|
assert current_service_target() == "auth_objects"
|
|
|
|
|
|
def test_with_service_target_keeps_the_wrapped_signature_and_coroutine_ness() -> None:
|
|
import inspect
|
|
|
|
@with_service_target("rate_limits")
|
|
async def hook(self: object, data: dict[str, str], call_type: str) -> None:
|
|
return None
|
|
|
|
assert inspect.iscoroutinefunction(hook)
|
|
assert tuple(inspect.signature(hook).parameters) == ("self", "data", "call_type")
|
|
assert hook.__name__ == "hook"
|
|
|
|
|
|
class _Suspend:
|
|
def __await__(self) -> Generator[None]:
|
|
yield
|
|
|
|
|
|
def test_an_sdk_callers_span_anchors_the_block_and_the_tasks_it_starts_and_nothing_after() -> None:
|
|
sdk_span: Final = object()
|
|
|
|
async def sdk_call() -> tuple[object, object, object]:
|
|
with request_parent_otel_span(sdk_span): # pyright: ignore[reportArgumentType] # any object stands in for a span
|
|
inside: Final = current_request_parent_otel_span()
|
|
background: Final = asyncio.create_task(asyncio.to_thread(current_request_parent_otel_span))
|
|
return inside, await background, current_request_parent_otel_span()
|
|
|
|
assert asyncio.run(sdk_call()) == (sdk_span, sdk_span, None)
|
|
|
|
|
|
def test_a_request_that_anchored_its_own_span_keeps_it_over_one_an_sdk_caller_passes() -> None:
|
|
proxy_span: Final = object()
|
|
|
|
def proxy_request() -> object:
|
|
set_request_parent_otel_span(proxy_span) # pyright: ignore[reportArgumentType] # any object stands in for a span
|
|
with request_parent_otel_span(object()): # pyright: ignore[reportArgumentType] # any object stands in for a span
|
|
return current_request_parent_otel_span()
|
|
|
|
assert contextvars.copy_context().run(proxy_request) is proxy_span
|
|
|
|
|
|
def test_a_targeted_coroutine_closed_from_another_context_does_not_raise() -> None:
|
|
@with_service_target("router_usage")
|
|
async def sync_forever() -> None:
|
|
await _Suspend()
|
|
|
|
suspended: Final = sync_forever()
|
|
contextvars.copy_context().run(suspended.send, None)
|
|
contextvars.copy_context().run(suspended.close)
|
|
assert current_service_target() is None
|
|
|
|
|
|
_DIRECT_REDIS_CALL: Final = re.compile(r"\b_?redis_cache\.(?!async_register_script\b)(?:async_)?\w+\(")
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class _FunctionScan:
|
|
name: str
|
|
reaches_redis_directly: bool
|
|
declares_a_family: bool
|
|
referenced_names: frozenset[str]
|
|
|
|
|
|
def _scan_function(source: str, fn: ast.FunctionDef | ast.AsyncFunctionDef) -> _FunctionScan:
|
|
body: Final = ast.get_source_segment(source, fn) or ""
|
|
decorators: Final = "\n".join(ast.get_source_segment(source, d) or "" for d in fn.decorator_list)
|
|
nodes: Final = tuple(ast.walk(fn))
|
|
names: Final = frozenset(n.id for n in nodes if isinstance(n, ast.Name))
|
|
attrs: Final = frozenset(n.attr for n in nodes if isinstance(n, ast.Attribute))
|
|
return _FunctionScan(
|
|
name=fn.name,
|
|
reaches_redis_directly=bool(_DIRECT_REDIS_CALL.search(body)),
|
|
declares_a_family=bool(_DECLARES_TARGET.search(body + "\n" + decorators)),
|
|
referenced_names=(names | attrs) - {fn.name},
|
|
)
|
|
|
|
|
|
def _covered_by_callers(scans: tuple[_FunctionScan, ...], covered: frozenset[str]) -> frozenset[str]:
|
|
"""Close ``covered`` over functions whose every in-file caller already declares a family."""
|
|
callers: Final = {
|
|
scan.name: frozenset(
|
|
other.name for other in scans if other.name != scan.name and scan.name in other.referenced_names
|
|
)
|
|
for scan in scans
|
|
}
|
|
grown: Final = covered | frozenset(
|
|
name for name, callers_of in callers.items() if callers_of and callers_of <= covered
|
|
)
|
|
return grown if grown == covered else _covered_by_callers(scans, grown)
|
|
|
|
|
|
def _direct_redis_callers_without_a_family(name: str) -> tuple[str, ...]:
|
|
source: Final = (_REPO / name).read_text()
|
|
scans: Final = tuple(
|
|
_scan_function(source, node)
|
|
for node in ast.walk(ast.parse(source))
|
|
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
|
|
)
|
|
declared: Final = frozenset(scan.name for scan in scans if scan.declares_a_family)
|
|
covered: Final = _covered_by_callers(scans, declared)
|
|
return tuple(f"{name}::{scan.name}" for scan in scans if scan.reaches_redis_directly and scan.name not in covered)
|
|
|
|
|
|
def test_every_function_that_reaches_redis_directly_declares_its_family() -> None:
|
|
"""A file-level declaration hides the producer that lacks one: the Claude Code session router
|
|
binding read sat in ``router.py`` beside dozens of declared families and still shipped as a bare
|
|
``redis.get``. A function that bypasses the cache facades and calls ``redis_cache`` itself must
|
|
carry the family on itself, its decorator, or every one of its in-file callers."""
|
|
exempt_files: Final = _CACHE_PARAMETER_HELPERS | _IN_MEMORY_ONLY_CALLERS
|
|
undeclared: Final = tuple(
|
|
function
|
|
for name in _redis_producers()
|
|
if name not in exempt_files
|
|
for function in _direct_redis_callers_without_a_family(name)
|
|
) # comprehension-ok: flatten per-file findings
|
|
assert undeclared == ()
|