mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_lit_7174_stream_tool_call_rewrites
This commit is contained in:
commit
b456caa05e
140 changed files with 9316 additions and 3415 deletions
2
.github/workflows/_test-unit-base.yml
vendored
2
.github/workflows/_test-unit-base.yml
vendored
|
|
@ -113,7 +113,7 @@ jobs:
|
|||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 8
|
||||
run: |
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra mongodb
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml
|
||||
uv run --no-sync python -c 'import os, sys; print(sys.version); assert f"{sys.version_info.major}.{sys.version_info.minor}" == os.environ["UV_PYTHON"]'
|
||||
|
||||
- name: Cache Prisma binaries
|
||||
|
|
|
|||
4
.github/workflows/test-litellm-ui-unit.yml
vendored
4
.github/workflows/test-litellm-ui-unit.yml
vendored
|
|
@ -66,7 +66,7 @@ jobs:
|
|||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
|
||||
run: |
|
||||
full_suite() { npm run test -- --run --pool forks --poolOptions.forks.maxForks=14; }
|
||||
full_suite() { npm run test -- --run --pool forks --maxWorkers=14; }
|
||||
|
||||
if [ -z "$BASE_SHA" ]; then
|
||||
echo "Push to $GITHUB_REF_NAME: running the full suite"
|
||||
|
|
@ -95,4 +95,4 @@ jobs:
|
|||
|
||||
echo "Pull request: running tests related to ${#changed_files[@]} changed UI files"
|
||||
npm run test -- related "${changed_files[@]}" --run --passWithNoTests \
|
||||
--pool forks --poolOptions.forks.maxForks=14
|
||||
--pool forks --maxWorkers=14
|
||||
|
|
|
|||
|
|
@ -67,7 +67,6 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr
|
|||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--extra bedrock-realtime \
|
||||
--extra mongodb \
|
||||
--python python3.13
|
||||
|
||||
# Copy full source tree
|
||||
|
|
@ -90,7 +89,6 @@ RUN uv sync --frozen --no-default-groups --no-editable \
|
|||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--extra bedrock-realtime \
|
||||
--extra mongodb \
|
||||
--python python3.13
|
||||
|
||||
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
|
||||
|
|
|
|||
|
|
@ -65,7 +65,6 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr
|
|||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--extra bedrock-realtime \
|
||||
--extra mongodb \
|
||||
--python python3.13
|
||||
|
||||
# Copy full source tree
|
||||
|
|
@ -88,7 +87,6 @@ RUN uv sync --frozen --no-default-groups --no-editable \
|
|||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--extra bedrock-realtime \
|
||||
--extra mongodb \
|
||||
--python python3.13
|
||||
|
||||
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
|
||||
|
|
|
|||
|
|
@ -71,7 +71,6 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
|
|||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--extra bedrock-realtime \
|
||||
--extra mongodb \
|
||||
--python python3.13
|
||||
|
||||
# Copy full source tree
|
||||
|
|
@ -100,7 +99,6 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
|
|||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--extra bedrock-realtime \
|
||||
--extra mongodb \
|
||||
--python python3.13 \
|
||||
--no-sources-package litellm-proxy-extras; \
|
||||
else \
|
||||
|
|
@ -111,7 +109,6 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
|
|||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--extra bedrock-realtime \
|
||||
--extra mongodb \
|
||||
--python python3.13; \
|
||||
fi
|
||||
|
||||
|
|
|
|||
|
|
@ -47,7 +47,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
|||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra bedrock-realtime \
|
||||
--extra mongodb \
|
||||
--python python3.13
|
||||
|
||||
# Stage 2 — copy source and install the project + workspace members.
|
||||
|
|
@ -60,7 +59,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
|||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra bedrock-realtime \
|
||||
--extra mongodb \
|
||||
--python python3.13
|
||||
|
||||
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
|
||||
|
|
|
|||
|
|
@ -326,6 +326,9 @@ ssl_certificate: Optional[str] = None
|
|||
user_url_validation: bool = True
|
||||
user_url_allowed_hosts: List[str] = []
|
||||
provider_url_destination_allowed_hosts: List[str] = []
|
||||
#: "override" (default) or "additive": whether a key or team destination replaces
|
||||
#: the operator's exporter for that backend or exports alongside it.
|
||||
otel_tenant_destination_mode: str | None = None
|
||||
ssl_ecdh_curve: Optional[str] = None # Set to 'X25519' to disable PQC and improve performance
|
||||
disable_streaming_logging: bool = False
|
||||
disable_token_counter: bool = False
|
||||
|
|
|
|||
|
|
@ -1374,6 +1374,7 @@ bedrock_embedding_models: Final[set] = set(
|
|||
"cohere.embed-multilingual-v3",
|
||||
"cohere.embed-v4:0",
|
||||
"twelvelabs.marengo-embed-2-7-v1:0",
|
||||
"twelvelabs.marengo-embed-3-0-v1:0",
|
||||
]
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -367,6 +367,13 @@
|
|||
"ui_name": "Headers",
|
||||
"description": "Headers for OTEL exporter (e.g., x-honeycomb-team=YOUR_API_KEY)",
|
||||
"required": false
|
||||
},
|
||||
"otel_exporter_otlp_protocol": {
|
||||
"type": "select",
|
||||
"ui_name": "Export Protocol",
|
||||
"description": "OTLP wire format for trace exports. Use http/json for collectors that cannot decode protobuf",
|
||||
"options": ["http/protobuf", "http/json"],
|
||||
"required": false
|
||||
}
|
||||
},
|
||||
"description": "OpenTelemetry Logging Integration"
|
||||
|
|
|
|||
|
|
@ -133,17 +133,17 @@ class MlflowLogger(CustomLogger):
|
|||
if final_response:
|
||||
end_time_ns: Final = int(end_time.timestamp() * 1e9)
|
||||
|
||||
self._extract_and_set_chat_attributes(span, kwargs, final_response)
|
||||
self._end_span_or_trace(
|
||||
span=span,
|
||||
outputs=final_response,
|
||||
status=SpanStatusCode.OK,
|
||||
end_time_ns=end_time_ns,
|
||||
)
|
||||
|
||||
# Remove the stream_id from the map
|
||||
with self._lock:
|
||||
self._stream_id_to_span.pop(litellm_call_id)
|
||||
try:
|
||||
self._extract_and_set_chat_attributes(span, kwargs, final_response)
|
||||
self._end_span_or_trace(
|
||||
span=span,
|
||||
outputs=final_response,
|
||||
status=SpanStatusCode.OK,
|
||||
end_time_ns=end_time_ns,
|
||||
)
|
||||
finally:
|
||||
with self._lock:
|
||||
self._stream_id_to_span.pop(litellm_call_id, None)
|
||||
|
||||
def _add_chunk_events(self, span, response_obj):
|
||||
from mlflow.entities import SpanEvent
|
||||
|
|
@ -282,15 +282,15 @@ class MlflowLogger(CustomLogger):
|
|||
"""End an MLflow span or a trace."""
|
||||
if span.parent_id is None:
|
||||
self._client.end_trace(
|
||||
trace_id=span.request_id,
|
||||
span.request_id,
|
||||
outputs=outputs,
|
||||
status=status,
|
||||
end_time_ns=end_time_ns,
|
||||
)
|
||||
else:
|
||||
self._client.end_span(
|
||||
trace_id=span.request_id,
|
||||
span_id=span.span_id,
|
||||
span.request_id,
|
||||
span.span_id,
|
||||
outputs=outputs,
|
||||
status=status,
|
||||
end_time_ns=end_time_ns,
|
||||
|
|
|
|||
|
|
@ -16,9 +16,11 @@ from opentelemetry.trace import (
|
|||
Span,
|
||||
Tracer,
|
||||
get_current_span,
|
||||
get_tracer_provider,
|
||||
set_span_in_context,
|
||||
use_span,
|
||||
)
|
||||
from opentelemetry.trace import TracerProvider as ApiTracerProvider
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -63,6 +65,7 @@ from litellm.integrations.otel.plumbing.metrics import (
|
|||
create_genai_metrics,
|
||||
)
|
||||
from litellm.integrations.otel.plumbing.providers import (
|
||||
attach_tenant_fan_out,
|
||||
build_tracer_provider,
|
||||
get_event_logger,
|
||||
get_meter,
|
||||
|
|
@ -85,6 +88,7 @@ if TYPE_CHECKING:
|
|||
)
|
||||
|
||||
LITELLM_TRACER_NAME: Final = "litellm"
|
||||
_published_v2_provider: ApiTracerProvider | None = None
|
||||
|
||||
|
||||
def _span_error_from_exception(
|
||||
|
|
@ -180,7 +184,9 @@ class OpenTelemetryV2(CustomLogger):
|
|||
self.config: OpenTelemetryV2Config = config or OpenTelemetryV2Config(**kwargs)
|
||||
self.callback_name = callback_name
|
||||
self._tracer_provider: TracerProvider = (
|
||||
tracer_provider if tracer_provider is not None else build_tracer_provider(self.config)
|
||||
tracer_provider
|
||||
if tracer_provider is not None
|
||||
else build_tracer_provider(self.config, tenant_overrides=True)
|
||||
)
|
||||
self.tracer: Tracer = get_tracer(self._tracer_provider, LITELLM_TRACER_NAME)
|
||||
self._metrics_recorder = self._init_metrics(meter_provider)
|
||||
|
|
@ -195,6 +201,11 @@ class OpenTelemetryV2(CustomLogger):
|
|||
self._open_llm_calls: OrderedDict[str, _LLMCallSpan] = OrderedDict()
|
||||
self._init_otel_logger_on_litellm_proxy()
|
||||
|
||||
@property
|
||||
def tracer_provider(self) -> TracerProvider:
|
||||
"""The provider this logger emits through, read-only to its callers."""
|
||||
return self._tracer_provider
|
||||
|
||||
def _init_metrics(self, meter_provider: "MeterProvider | None") -> "GenAIMetricRecorder | None":
|
||||
"""Create the six GenAI histograms when metrics are enabled, else ``None``.
|
||||
|
||||
|
|
@ -863,12 +874,33 @@ def publish_global_otel_v2_provider(
|
|||
``opentelemetry.trace.set_tracer_provider``) are injected so the publish step is
|
||||
unit-testable without reading or mutating real global OTel state. Returns the
|
||||
logger whose provider was published.
|
||||
|
||||
The published provider is also the one that fans spans out to key/team
|
||||
destinations, because it is the only provider the whole request tree passes
|
||||
through; see :func:`attach_tenant_fan_out`. It is remembered for
|
||||
:func:`fan_out_provider` because neither the OTel global (``set_tracer_provider``
|
||||
keeps the first provider it was ever handed) nor
|
||||
``proxy_server.open_telemetry_logger`` (a legacy v1 logger can hold that slot)
|
||||
reliably leads back to it.
|
||||
"""
|
||||
global _published_v2_provider
|
||||
logger: Final = select_global_otel_v2_logger(in_memory_loggers, registered=registered)
|
||||
set_global_provider(logger._tracer_provider)
|
||||
attach_tenant_fan_out(logger.tracer_provider, *_v2_configs(in_memory_loggers, logger))
|
||||
set_global_provider(logger.tracer_provider)
|
||||
_published_v2_provider = logger.tracer_provider # rebind-ok: startup records the one provider carrying the fan-out
|
||||
return logger
|
||||
|
||||
|
||||
def _v2_configs(in_memory_loggers: Sequence[object], logger: "OpenTelemetryV2") -> tuple[OpenTelemetryV2Config, ...]:
|
||||
"""Every v2 logger's config, the published logger's first.
|
||||
|
||||
Each preset keeps its own provider and exporters, so the accounts the operator
|
||||
writes to are spread over all of them, not held by the published logger alone.
|
||||
"""
|
||||
others: Final = tuple(cb.config for cb in in_memory_loggers if isinstance(cb, OpenTelemetryV2) and cb is not logger)
|
||||
return (logger.config, *others)
|
||||
|
||||
|
||||
def _registered_v2_logger() -> "OpenTelemetryV2 | None":
|
||||
try:
|
||||
from litellm.proxy import proxy_server
|
||||
|
|
@ -904,6 +936,25 @@ def seed_request_identity(user_api_key_dict: object, model: str | None = None) -
|
|||
logger.seed_request_identity(user_api_key_dict, model=model)
|
||||
|
||||
|
||||
def fan_out_provider() -> ApiTracerProvider:
|
||||
"""The provider :func:`publish_global_otel_v2_provider` gave the tenant fan-out.
|
||||
|
||||
Read off the publish itself, not the OTel global and not the registered logger:
|
||||
the global keeps whichever provider claimed it first (auto-instrumentation, a
|
||||
legacy logger), and the registered slot can hold a v1 logger while the publish
|
||||
picked a v2 one from ``_in_memory_loggers``. Either detour lands on a provider
|
||||
with no fan-out and drops every destination at auth.
|
||||
"""
|
||||
published: Final = _published_v2_provider
|
||||
if published is not None:
|
||||
return published
|
||||
logger: Final = _registered_v2_logger()
|
||||
if logger is not None:
|
||||
attach_tenant_fan_out(logger.tracer_provider, logger.config)
|
||||
return logger.tracer_provider
|
||||
return get_tracer_provider()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def phase_span(name: str) -> "Iterator[Span | None]":
|
||||
logger: Final = _registered_v2_logger()
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ from litellm.integrations.otel.model.payloads import (
|
|||
ServiceSpanData,
|
||||
ToolDefinition,
|
||||
)
|
||||
from litellm.integrations.otel.model.semconv import Error
|
||||
|
||||
# Attribute keys in the semconv-ai / Traceloop vocabulary.
|
||||
_LEGACY_SYSTEM: Final = "gen_ai.system"
|
||||
|
|
@ -36,7 +37,7 @@ _LEGACY_PRESENCE_PENALTY: Final = "llm.presence_penalty"
|
|||
_LEGACY_STOP_SEQUENCES: Final = "llm.chat.stop_sequences"
|
||||
_LEGACY_SERVICE: Final = "service"
|
||||
_LEGACY_CALL_TYPE: Final = "call_type"
|
||||
_LEGACY_ERROR: Final = "error"
|
||||
_LEGACY_ERROR: Final = Error.MESSAGE_LEGACY
|
||||
|
||||
|
||||
class LegacyMapper:
|
||||
|
|
|
|||
|
|
@ -69,7 +69,7 @@ class ExporterSpec(BaseModel):
|
|||
|
||||
kind: str = Field(
|
||||
default="console",
|
||||
description="console | in_memory | otlp_http | otlp_grpc | <factory kind>",
|
||||
description="console | in_memory | otlp_http | http/json | otlp_grpc | <factory kind>",
|
||||
)
|
||||
endpoint: str | None = None
|
||||
traces_endpoint: str | None = Field(
|
||||
|
|
@ -269,7 +269,9 @@ class OpenTelemetryV2Config(BaseSettings):
|
|||
if (self.endpoint or self.traces_endpoint) and self.exporter == "console":
|
||||
self.exporter = "otlp_http"
|
||||
# When no explicit destinations are given, fold the single-destination
|
||||
# shorthand into one spec so the provider always has a destination.
|
||||
# shorthand into one spec so the provider always has a destination. A spec
|
||||
# with no fields set is how the presets tell "nothing configured" from an
|
||||
# operator who asked for the console by name.
|
||||
if not self.exporters:
|
||||
self.exporters = [
|
||||
ExporterSpec(
|
||||
|
|
@ -278,6 +280,8 @@ class OpenTelemetryV2Config(BaseSettings):
|
|||
traces_endpoint=self.traces_endpoint,
|
||||
headers=self.headers,
|
||||
)
|
||||
if not self.model_fields_set.isdisjoint(("exporter", "endpoint", "headers"))
|
||||
else ExporterSpec()
|
||||
]
|
||||
# Ensure ``genai`` is always present and first.
|
||||
names = list(self.mapper_names)
|
||||
|
|
|
|||
49
litellm/integrations/otel/model/destination.py
Normal file
49
litellm/integrations/otel/model/destination.py
Normal file
|
|
@ -0,0 +1,49 @@
|
|||
"""The resolved OTLP destination a request's traces export to.
|
||||
|
||||
Backend-agnostic on purpose: every OTEL backend reduces to an endpoint plus auth
|
||||
headers. The per-backend field mapping lives in ``presets.destinations``.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
from urllib.parse import quote
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
|
||||
class OtelDestination(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
endpoint: str
|
||||
headers: Mapping[str, str] = Field(default_factory=dict)
|
||||
resource_attributes: Mapping[str, str] = Field(default_factory=dict)
|
||||
callback_name: str | None = None
|
||||
protocol: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"OTLP transport, defaulting to the backend's own. Not derivable from the "
|
||||
"scheme: Arize's ``https://otlp.arize.com/v1`` is gRPC."
|
||||
),
|
||||
)
|
||||
|
||||
def header_string(self) -> str:
|
||||
"""Render headers as the ``k=v,k2=v2`` form an ``ExporterSpec`` expects.
|
||||
|
||||
Values are percent-encoded because ``providers.parse_headers`` decodes them
|
||||
with the SDK's W3C-Baggage parser: a value carrying a ``,`` or ``=`` (a
|
||||
Langfuse project name, a base64 Authorization payload ending in ``==``)
|
||||
would otherwise be split into bogus pairs on the way back out.
|
||||
"""
|
||||
return ",".join(f"{key}={quote(value, safe='')}" for key, value in self.headers.items())
|
||||
|
||||
def cache_key(self) -> tuple[str, tuple[tuple[str, str], ...], tuple[tuple[str, str], ...], str | None]:
|
||||
"""Identity for processor reuse, so one destination means one exporter."""
|
||||
return (
|
||||
self.endpoint,
|
||||
tuple(sorted(self.headers.items())),
|
||||
tuple(sorted(self.resource_attributes.items())),
|
||||
self.protocol,
|
||||
)
|
||||
|
||||
|
||||
NO_DESTINATIONS: Final[tuple[OtelDestination, ...]] = ()
|
||||
|
|
@ -204,6 +204,9 @@ class Error:
|
|||
|
||||
TYPE: Final = "error.type"
|
||||
MESSAGE: Final = "error.message"
|
||||
# The same text under the bare key the semconv-ai / Traceloop vocabulary uses
|
||||
# (see ``LegacyMapper``), so anything reading or redacting error text covers both.
|
||||
MESSAGE_LEGACY: Final = "error"
|
||||
|
||||
|
||||
class LiteLLMError:
|
||||
|
|
|
|||
|
|
@ -1,8 +1,9 @@
|
|||
"""Trace-context + Baggage helpers."""
|
||||
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from contextvars import ContextVar, Token
|
||||
from typing import Final
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from opentelemetry import baggage
|
||||
from opentelemetry.context import Context, get_current
|
||||
|
|
@ -21,6 +22,9 @@ from opentelemetry.trace.propagation.tracecontext import (
|
|||
|
||||
from litellm.integrations.otel.model.semconv import HTTP
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.integrations.otel.model.destination import OtelDestination
|
||||
|
||||
_PROPAGATOR: Final = TraceContextTextMapPropagator()
|
||||
|
||||
# The request's root span — the FastAPI-owned SERVER span — captured ONCE when the
|
||||
|
|
@ -304,3 +308,65 @@ def extract_traceparent(headers: Mapping[str, str]) -> Context | None:
|
|||
return None
|
||||
carrier: Final = {str(key).lower(): value for key, value in headers.items()}
|
||||
return _PROPAGATOR.extract(carrier)
|
||||
|
||||
|
||||
# The OTLP destinations this request's key or team pointed its traces at, resolved
|
||||
# once during auth. A ``ContextVar`` for the same reason the root span above is one:
|
||||
# it rides the request task's context into the ``asyncio.create_task`` children that
|
||||
# close the LLM span, and it is visible to every ``SpanProcessor.on_end`` that fires
|
||||
# on the request task. Stateful MCP handlers set and reset it per message; the
|
||||
# request-task value otherwise dies with that task.
|
||||
_request_destinations: Final['ContextVar[tuple["OtelDestination", ...]]'] = ContextVar(
|
||||
"litellm_otel_request_destinations", default=()
|
||||
)
|
||||
|
||||
|
||||
def set_request_destinations(destinations: 'tuple["OtelDestination", ...]') -> "Token[tuple[OtelDestination, ...]]":
|
||||
"""Anchor the destinations this request exports to and return a reset token."""
|
||||
return _request_destinations.set(destinations)
|
||||
|
||||
|
||||
def reset_request_destinations(token: "Token[tuple[OtelDestination, ...]]") -> None:
|
||||
_request_destinations.reset(token)
|
||||
|
||||
|
||||
def request_destinations() -> 'tuple["OtelDestination", ...]':
|
||||
"""The destinations resolved for this request, empty outside a proxy request."""
|
||||
return _request_destinations.get()
|
||||
|
||||
|
||||
#: ``litellm_settings: otel_tenant_destination_mode`` and its env equivalent.
|
||||
ADDITIVE_DESTINATION_MODE: Final = "additive"
|
||||
OTEL_TENANT_DESTINATION_MODE_ENV: Final = "LITELLM_OTEL_TENANT_DESTINATION_MODE"
|
||||
|
||||
|
||||
def tenant_destinations_are_additive() -> bool:
|
||||
"""Whether a tenant destination exports alongside the operator's own exporter.
|
||||
|
||||
Override is the default: the tenant's traffic reaches the tenant's account and
|
||||
nowhere else. Operators running one org-wide backend across every team set this
|
||||
to ``additive`` so the same trace lands in both places.
|
||||
"""
|
||||
import litellm
|
||||
|
||||
configured: Final = litellm.otel_tenant_destination_mode or os.environ.get(OTEL_TENANT_DESTINATION_MODE_ENV)
|
||||
return isinstance(configured, str) and configured.strip().lower() == ADDITIVE_DESTINATION_MODE
|
||||
|
||||
|
||||
def destination_backends() -> frozenset[str]:
|
||||
"""Backends this request resolved a tenant destination for.
|
||||
|
||||
The fan-out already carries the whole trace to those destinations, so the
|
||||
per-request tracer route must never send a second copy, in either mode.
|
||||
"""
|
||||
return frozenset(d.callback_name for d in _request_destinations.get() if d.callback_name)
|
||||
|
||||
|
||||
def suppressed_backends() -> frozenset[str]:
|
||||
"""Backends whose operator-level exporters this request must NOT reach.
|
||||
|
||||
Empty under ``additive``, where the operator keeps its copy of every span.
|
||||
"""
|
||||
if tenant_destinations_are_additive():
|
||||
return frozenset()
|
||||
return destination_backends()
|
||||
|
|
|
|||
70
litellm/integrations/otel/plumbing/otlp_json.py
Normal file
70
litellm/integrations/otel/plumbing/otlp_json.py
Normal file
|
|
@ -0,0 +1,70 @@
|
|||
"""OTLP/HTTP span exporter that sends the OTLP/JSON encoding instead of protobuf.
|
||||
|
||||
The SDK only ships a protobuf OTLP/HTTP exporter; this reuses its transport and
|
||||
retry loop and swaps the payload for OTLP/JSON (enums as integers, ids as hex).
|
||||
"""
|
||||
|
||||
import base64
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Final, TypeAlias
|
||||
|
||||
from google.protobuf.json_format import MessageToDict
|
||||
from opentelemetry.exporter.otlp.proto.common.trace_encoder import encode_spans
|
||||
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
|
||||
from opentelemetry.sdk.trace import ReadableSpan
|
||||
|
||||
JSON_CONTENT_TYPE: Final = "application/json"
|
||||
_HEX_ID_KEYS: Final = frozenset({"traceId", "spanId", "parentSpanId"})
|
||||
|
||||
_JsonValue: TypeAlias = "Mapping[str, _JsonValue] | Sequence[_JsonValue] | str | int | float | bool | None"
|
||||
_JsonObject: TypeAlias = Mapping[str, "_JsonValue"]
|
||||
|
||||
|
||||
def _objects(node: _JsonObject, key: str) -> tuple[_JsonObject, ...]:
|
||||
items: Final = node.get(key)
|
||||
if isinstance(items, str) or not isinstance(items, Sequence):
|
||||
return ()
|
||||
return tuple(item for item in items if isinstance(item, Mapping))
|
||||
|
||||
|
||||
def _hex_ids(node: _JsonObject) -> _JsonObject:
|
||||
return MappingProxyType(
|
||||
{
|
||||
key: base64.b64decode(item).hex() if key in _HEX_ID_KEYS and isinstance(item, str) else item
|
||||
for key, item in node.items()
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _hex_span(span: _JsonObject) -> _JsonObject:
|
||||
links: Final = _objects(span, "links")
|
||||
if not links:
|
||||
return _hex_ids(span)
|
||||
return MappingProxyType({**_hex_ids(span), "links": tuple(_hex_ids(link) for link in links)})
|
||||
|
||||
|
||||
def _hex_scope_spans(scope: _JsonObject) -> _JsonObject:
|
||||
return MappingProxyType({**scope, "spans": tuple(_hex_span(span) for span in _objects(scope, "spans"))})
|
||||
|
||||
|
||||
def _hex_resource_spans(resource: _JsonObject) -> _JsonObject:
|
||||
scope_spans: Final = tuple(_hex_scope_spans(scope) for scope in _objects(resource, "scopeSpans"))
|
||||
return MappingProxyType({**resource, "scopeSpans": scope_spans})
|
||||
|
||||
|
||||
def encode_spans_json(spans: Sequence[ReadableSpan]) -> bytes:
|
||||
payload: Final[_JsonObject] = MessageToDict(encode_spans(spans), use_integers_for_enums=True)
|
||||
resource_spans: Final = tuple(_hex_resource_spans(resource) for resource in _objects(payload, "resourceSpans"))
|
||||
hexed: Final[_JsonObject] = MappingProxyType({**payload, "resourceSpans": resource_spans})
|
||||
return json.dumps(hexed, default=dict, separators=(",", ":")).encode()
|
||||
|
||||
|
||||
class OTLPJsonSpanExporter(OTLPSpanExporter):
|
||||
def __init__(self, endpoint: str | None, headers: dict[str, str]) -> None: # mutable-ok: SDK __init__ takes Dict
|
||||
super().__init__(endpoint=endpoint, headers=headers)
|
||||
self._session.headers["Content-Type"] = JSON_CONTENT_TYPE
|
||||
|
||||
def _serialize_spans(self, spans: Sequence[ReadableSpan]) -> bytes:
|
||||
return encode_spans_json(spans)
|
||||
|
|
@ -1,9 +1,14 @@
|
|||
"""Provider / exporter factory + the Baggage span processor."""
|
||||
|
||||
from collections.abc import Callable, Iterable
|
||||
import queue
|
||||
import threading
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Callable, Iterable, Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal
|
||||
|
||||
from opentelemetry import _logs, baggage, metrics
|
||||
from opentelemetry import _logs, baggage, metrics, trace
|
||||
from opentelemetry._events import EventLogger
|
||||
from opentelemetry._logs import LoggerProvider, NoOpLoggerProvider
|
||||
from opentelemetry.context import Context
|
||||
|
|
@ -19,7 +24,8 @@ from opentelemetry.sdk._logs.export import (
|
|||
)
|
||||
from opentelemetry.sdk.metrics import MeterProvider as SDKMeterProvider
|
||||
from opentelemetry.sdk.resources import Resource
|
||||
from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor, TracerProvider
|
||||
from opentelemetry.sdk.trace import Event, ReadableSpan, SpanProcessor, TracerProvider
|
||||
from opentelemetry.sdk.trace import Span as SDKSpan
|
||||
from opentelemetry.sdk.trace.export import (
|
||||
BatchSpanProcessor,
|
||||
ConsoleSpanExporter,
|
||||
|
|
@ -29,18 +35,35 @@ from opentelemetry.sdk.trace.export import (
|
|||
from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
|
||||
InMemorySpanExporter,
|
||||
)
|
||||
from opentelemetry.trace import Span, SpanKind, Tracer
|
||||
from opentelemetry.trace import Span, SpanKind, Status, Tracer
|
||||
from opentelemetry.util.re import parse_env_headers
|
||||
from opentelemetry.util.types import Attributes, AttributeValue
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm._version import version as litellm_version
|
||||
from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
|
||||
from litellm.integrations.otel.model.semconv import LiteLLM
|
||||
from litellm.integrations.otel.model.semconv import (
|
||||
DB,
|
||||
MCP,
|
||||
Error,
|
||||
ExceptionEvent,
|
||||
GenAI,
|
||||
LiteLLM,
|
||||
LiteLLMError,
|
||||
Server,
|
||||
)
|
||||
from litellm.integrations.otel.model.spans import LiteLLMSpanKind
|
||||
from litellm.integrations.otel.plumbing.context import (
|
||||
request_destinations,
|
||||
suppressed_backends,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from opentelemetry.metrics import Meter
|
||||
from opentelemetry.sdk.metrics.export import MetricReader
|
||||
|
||||
from litellm.integrations.otel.model.destination import OtelDestination
|
||||
|
||||
_SPAN_KIND_BY_ROLE_KIND: Final[dict[LiteLLMSpanKind, SpanKind]] = {
|
||||
LiteLLMSpanKind.SERVER: SpanKind.SERVER,
|
||||
LiteLLMSpanKind.CLIENT: SpanKind.CLIENT,
|
||||
|
|
@ -136,7 +159,8 @@ def parse_headers(raw: str | None) -> dict[str, str]:
|
|||
|
||||
|
||||
_IN_MEMORY_KINDS: Final = ("in_memory", "inmemory", "memory")
|
||||
_OTLP_HTTP_KINDS: Final = ("otlp_http", "http", "http/protobuf", "http/json")
|
||||
_OTLP_HTTP_JSON_KINDS: Final = ("http/json",)
|
||||
_OTLP_HTTP_KINDS: Final = ("otlp_http", "http", "http/protobuf", *_OTLP_HTTP_JSON_KINDS)
|
||||
_OTLP_GRPC_KINDS: Final = ("otlp_grpc", "grpc")
|
||||
|
||||
|
||||
|
|
@ -164,6 +188,13 @@ def _exporter_from_spec(spec: ExporterSpec) -> SpanExporter:
|
|||
return factory(spec)
|
||||
if kind in _IN_MEMORY_KINDS:
|
||||
return InMemorySpanExporter()
|
||||
if kind in _OTLP_HTTP_JSON_KINDS:
|
||||
from litellm.integrations.otel.plumbing.otlp_json import OTLPJsonSpanExporter
|
||||
|
||||
return OTLPJsonSpanExporter(
|
||||
endpoint=spec.traces_endpoint or _otlp_traces_endpoint(spec.endpoint),
|
||||
headers=parse_headers(spec.headers),
|
||||
)
|
||||
if kind in _OTLP_HTTP_KINDS:
|
||||
from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
|
||||
OTLPSpanExporter as HTTPExporter,
|
||||
|
|
@ -194,6 +225,555 @@ def _processor_for(exporter: SpanExporter, use_simple: bool | None) -> SpanProce
|
|||
return SimpleSpanProcessor(exporter) if use_simple else BatchSpanProcessor(exporter)
|
||||
|
||||
|
||||
#: Distinct tenant destinations whose exporters stay alive. Each holds a connection
|
||||
#: pool and a batch thread, so the cache is bounded and evicts least-recently-used.
|
||||
_MAX_CACHED_DESTINATION_PROCESSORS: Final = 32
|
||||
|
||||
#: Workers closing shed destination processors, bounding the threads a tenant can
|
||||
#: create by cycling its destination config.
|
||||
_DRAIN_WORKERS: Final = 2
|
||||
|
||||
#: Shed processors waiting to be closed before the fan-out stops building new ones.
|
||||
#: Each still owns a batch thread until its close returns, and a collector that never
|
||||
#: answers makes every close take the exporter's full timeout, so past this many the
|
||||
#: operator's exporter keeps the span instead (see ``deliverable``).
|
||||
_MAX_PENDING_DRAINS: Final = 64
|
||||
|
||||
#: How long ``shutdown`` waits for spans already being forwarded, so teardown closes
|
||||
#: no processor under one. Bounded: an exporter that never returns must not hold the
|
||||
#: proxy open.
|
||||
_SHUTDOWN_DRAIN_SECONDS: Final = 5.0
|
||||
|
||||
#: An exporter's account: its normalized endpoint and the credentials it presents.
|
||||
_SinkKey = tuple[str, tuple[tuple[str, str], ...]]
|
||||
|
||||
#: Header names that spell one credential two ways. Arize's operator exporter sends
|
||||
#: ``space_id`` where a tenant destination sends ``arize-space-id``.
|
||||
_CREDENTIAL_ALIASES: Final = MappingProxyType({"arize_space_id": "space_id"})
|
||||
|
||||
|
||||
class _DrainPool:
|
||||
"""Closes shed destination processors off the span-export path.
|
||||
|
||||
``shutdown`` flushes over the network and is reached from ``on_end``, so closing
|
||||
one inline would let a single unreachable tenant collector stall every other
|
||||
tenant's spans behind it. A fixed set of workers rather than a thread per
|
||||
processor means a tenant cycling its destination config cannot spawn threads as
|
||||
fast as it can send requests; slow shutdowns queue behind each other.
|
||||
|
||||
The workers are daemons and belong to the fan-out that sheds the processors, so
|
||||
neither an unreachable collector nor a lazily built process-wide singleton can
|
||||
hold the proxy open on the way down.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
workers: int = _DRAIN_WORKERS,
|
||||
pending: "queue.Queue[SpanProcessor | None] | None" = None,
|
||||
capacity: int = _MAX_PENDING_DRAINS,
|
||||
) -> None:
|
||||
self._workers: Final = workers
|
||||
self._capacity: Final = capacity
|
||||
self._lock: Final = threading.Lock()
|
||||
self._closed = False
|
||||
self._backlog = 0 # guarded by ``_lock``: submitted processors whose close has not returned
|
||||
self._pending: Final[queue.Queue[SpanProcessor | None]] = pending if pending is not None else queue.Queue()
|
||||
self._threads: Final = tuple(
|
||||
threading.Thread(target=self._drain_until_closed, daemon=True, name="litellm-otel-destination-drain")
|
||||
for _ in range(workers)
|
||||
)
|
||||
for worker in self._threads:
|
||||
worker.start()
|
||||
|
||||
def submit(self, processor: SpanProcessor) -> None:
|
||||
"""Queue ``processor`` for closing, or hand it off once the pool is retired.
|
||||
|
||||
The check and the put share one lock. Reading a closed flag on its own leaves
|
||||
room for :meth:`close` to run in between, and the processor would land behind
|
||||
the sentinels every worker has already exited on.
|
||||
|
||||
Past close there is no worker left to take it, and the caller is whichever
|
||||
thread just ended a span, so closing it inline would park that thread on a
|
||||
network flush the shutdown deadline has already stopped waiting for. The extra
|
||||
thread is bounded by the same close: the fan-out stops handing processors out
|
||||
at that point, so only the ones already exporting when it happened arrive here.
|
||||
"""
|
||||
with self._lock:
|
||||
if not self._closed:
|
||||
self._backlog += 1
|
||||
self._pending.put(processor)
|
||||
return
|
||||
threading.Thread(
|
||||
target=_shutdown_quietly,
|
||||
args=(processor,),
|
||||
daemon=True,
|
||||
name="litellm-otel-destination-drain-straggler",
|
||||
).start()
|
||||
|
||||
def saturated(self) -> bool:
|
||||
"""Whether enough closes are outstanding that building another processor must wait.
|
||||
|
||||
The workers close in order and each close blocks for as long as its exporter
|
||||
does, so a collector that stopped answering would otherwise turn every new
|
||||
destination into one more batch thread parked behind them, for as long as the
|
||||
tenants keep rotating. Holding the count here rather than reading the queue
|
||||
keeps the two processors a worker is mid-close on in the total.
|
||||
"""
|
||||
with self._lock:
|
||||
return self._backlog >= self._capacity
|
||||
|
||||
def close(self, timeout: float | None = None) -> None:
|
||||
"""Retire the workers once they have closed everything already queued.
|
||||
|
||||
A proxy that rebuilds its telemetry builds another fan-out, so workers that
|
||||
outlive the one that started them are two more threads per reload, forever.
|
||||
|
||||
``timeout`` bounds how long the caller waits for that draining to finish. The
|
||||
workers are daemons, so whatever is still flushing when it expires is dropped
|
||||
by the interpreter rather than holding it open.
|
||||
"""
|
||||
with self._lock:
|
||||
if self._closed:
|
||||
return
|
||||
self._closed = True
|
||||
for _ in range(self._workers):
|
||||
self._pending.put(None)
|
||||
if timeout is None:
|
||||
return
|
||||
deadline: Final = time.monotonic() + timeout
|
||||
for worker in self._threads:
|
||||
worker.join(timeout=max(0.0, deadline - time.monotonic()))
|
||||
|
||||
def _drain_until_closed(self) -> None:
|
||||
while True:
|
||||
processor: SpanProcessor | None = self._pending.get() # rebind-ok: loop variable
|
||||
if processor is None:
|
||||
return
|
||||
_shutdown_quietly(processor)
|
||||
with self._lock:
|
||||
self._backlog -= 1
|
||||
|
||||
|
||||
_NO_ATTRIBUTES: Final[Mapping[str, AttributeValue]] = MappingProxyType({})
|
||||
_DB_SYSTEM_KEYS: Final = frozenset({DB.SYSTEM_NAME, DB.SYSTEM_LEGACY})
|
||||
# Keys on a database span that describe the proxy's own datastore: its host, its
|
||||
# port, and its schema.
|
||||
_DATASTORE_ENDPOINT_KEYS: Final = frozenset({Server.ADDRESS, Server.PORT, DB.NAMESPACE})
|
||||
# A span carrying one of these describes the tenant's own call (the model call, the
|
||||
# MCP call, the guardrail), so its error text is theirs to see. Every other span is
|
||||
# the proxy's own work, whose error text names the operator's infrastructure.
|
||||
_TENANT_OWNED_KEYS: Final = frozenset({GenAI.OPERATION_NAME, MCP.METHOD_NAME, LiteLLM.GUARDRAIL_NAME})
|
||||
_PROXY_ERROR_TEXT_KEYS: Final = frozenset({Error.MESSAGE, Error.MESSAGE_LEGACY})
|
||||
# A guardrail that never answered carries the exception it raised as its response,
|
||||
# which names the operator's guardrail endpoint. The second spelling is the legacy
|
||||
# status the request-level logger still maps.
|
||||
_GUARDRAIL_UNREACHABLE_STATUSES: Final = frozenset({"guardrail_failed_to_respond", "failure"})
|
||||
# Attribute prefixes the FastAPI instrumentor uses for headers the operator opted to
|
||||
# capture (``OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_*``). The request
|
||||
# side carries the caller's bearer token verbatim.
|
||||
_CAPTURED_HEADER_PREFIXES: Final = ("http.request.header.", "http.response.header.")
|
||||
# The instrumentor stamps the request URL on the server span with its query string,
|
||||
# under the old convention and the new one, and litellm accepts a virtual key as a
|
||||
# ``?key=`` query parameter.
|
||||
_URL_KEYS: Final = frozenset({"http.url", "http.target", "url.full"})
|
||||
_URL_QUERY_KEY: Final = "url.query"
|
||||
|
||||
|
||||
class _TenantSpanView(ReadableSpan):
|
||||
"""A ``ReadableSpan`` view for one destination, leaving the operator's own span alone."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
inner: ReadableSpan,
|
||||
resource: Resource,
|
||||
attributes: Attributes,
|
||||
events: Sequence[Event],
|
||||
status: Status,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
name=inner.name,
|
||||
context=inner.context,
|
||||
parent=inner.parent,
|
||||
resource=resource,
|
||||
attributes=attributes,
|
||||
events=events,
|
||||
links=inner.links,
|
||||
kind=inner.kind,
|
||||
status=status,
|
||||
start_time=inner.start_time,
|
||||
end_time=inner.end_time,
|
||||
instrumentation_scope=inner.instrumentation_scope,
|
||||
)
|
||||
|
||||
|
||||
def _is_database_span(attributes: Mapping[str, AttributeValue]) -> bool:
|
||||
return any(key in attributes for key in _DB_SYSTEM_KEYS)
|
||||
|
||||
|
||||
def _is_tenant_owned_span(attributes: Mapping[str, AttributeValue]) -> bool:
|
||||
return any(key in attributes for key in _TENANT_OWNED_KEYS)
|
||||
|
||||
|
||||
def _guardrail_unreachable(attributes: Mapping[str, AttributeValue]) -> bool:
|
||||
return attributes.get(LiteLLM.GUARDRAIL_STATUS) in _GUARDRAIL_UNREACHABLE_STATUSES
|
||||
|
||||
|
||||
def _tenant_visible(key: str, database: bool, owned: bool, unreachable_guardrail: bool) -> bool:
|
||||
if key.startswith(_CAPTURED_HEADER_PREFIXES) or key in (LiteLLMError.STACK_TRACE, _URL_QUERY_KEY):
|
||||
return False
|
||||
if database and key in _DATASTORE_ENDPOINT_KEYS:
|
||||
return False
|
||||
if unreachable_guardrail and key == LiteLLM.GUARDRAIL_RESPONSE:
|
||||
return False
|
||||
return owned or key not in _PROXY_ERROR_TEXT_KEYS
|
||||
|
||||
|
||||
def _without_query(key: str, value: AttributeValue) -> AttributeValue:
|
||||
if key not in _URL_KEYS or not isinstance(value, str):
|
||||
return value
|
||||
return value.partition("?")[0]
|
||||
|
||||
|
||||
def _same_attributes(kept: Mapping[str, AttributeValue], attributes: Mapping[str, AttributeValue]) -> bool:
|
||||
return len(kept) == len(attributes) and all(kept[key] is value for key, value in attributes.items())
|
||||
|
||||
|
||||
def _without_stack_trace(event: Event) -> Event:
|
||||
attributes: Final = event.attributes or _NO_ATTRIBUTES
|
||||
if ExceptionEvent.STACKTRACE not in attributes:
|
||||
return event
|
||||
return Event(
|
||||
name=event.name,
|
||||
attributes=MappingProxyType(
|
||||
{key: value for key, value in attributes.items() if key != ExceptionEvent.STACKTRACE}
|
||||
),
|
||||
timestamp=event.timestamp,
|
||||
)
|
||||
|
||||
|
||||
def _for_destination(span: ReadableSpan, destination: "OtelDestination") -> ReadableSpan:
|
||||
"""The view of ``span`` a tenant destination receives.
|
||||
|
||||
A span the tenant's own call produced keeps its error text. Every other span is
|
||||
the proxy's own work (the request root, auth, the database), and its error text,
|
||||
its events and its status description come off, since a Prisma failure there
|
||||
spells out the operator's Postgres endpoint. A database span loses that endpoint
|
||||
too, and a guardrail that failed to respond loses its response text, which is the
|
||||
exception it raised and names the operator's guardrail endpoint. Stack traces walk
|
||||
the operator's install and come off every span, as do the headers the operator
|
||||
captures on the server span, whose request side holds the caller's bearer token,
|
||||
and the query string of the request URL, which can hold the same key. The span
|
||||
itself stays, so the tenant still gets the whole trace tree.
|
||||
"""
|
||||
extra: Final = destination.resource_attributes
|
||||
attributes: Final = span.attributes or _NO_ATTRIBUTES
|
||||
database: Final = _is_database_span(attributes)
|
||||
owned: Final = _is_tenant_owned_span(attributes)
|
||||
unreachable: Final = _guardrail_unreachable(attributes)
|
||||
kept: Final = MappingProxyType(
|
||||
{
|
||||
key: _without_query(key, value)
|
||||
for key, value in attributes.items()
|
||||
if _tenant_visible(key, database, owned, unreachable)
|
||||
}
|
||||
)
|
||||
recorded: Final = span.events
|
||||
events: Final = tuple(_without_stack_trace(event) for event in recorded) if owned else ()
|
||||
unchanged: Final = owned and _same_attributes(kept, attributes) and all(a is b for a, b in zip(events, recorded))
|
||||
if not extra and unchanged:
|
||||
return span
|
||||
resource: Final = span.resource.merge(Resource(extra)) if extra else span.resource
|
||||
status: Final = span.status if owned else Status(span.status.status_code)
|
||||
return _TenantSpanView(span, resource, kept, events, status)
|
||||
|
||||
|
||||
class TenantFanOutSpanProcessor(SpanProcessor):
|
||||
"""Export every finished span to each destination this request resolved.
|
||||
|
||||
Destinations ride a request-scoped ``ContextVar`` set during auth, so concurrent
|
||||
requests stay isolated. The forwarded view keeps the original trace and parent
|
||||
ids, so the tenant gets the same tree the operator would have received.
|
||||
|
||||
Exactly one provider carries this processor, the one published as the OTel global
|
||||
(see :func:`attach_tenant_fan_out`). That provider is the only one every span
|
||||
passes through: the FastAPI server span, the auth span and the post-call database
|
||||
spans are emitted on the global, while a second v2 logger's provider sees only
|
||||
that logger's own gen-AI span. Attaching the fan-out per logger would hand a
|
||||
tenant a one-span trace whenever its backend is not the global one, and two
|
||||
copies of the model call whenever it is.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
processor_factory: 'Callable[["OtelDestination"], SpanProcessor | None] | None' = None,
|
||||
shutdown_drain_seconds: float = _SHUTDOWN_DRAIN_SECONDS,
|
||||
operator_sinks: frozenset[_SinkKey] = frozenset(),
|
||||
pending_drains: int = _MAX_PENDING_DRAINS,
|
||||
drain_pool: _DrainPool | None = None,
|
||||
) -> None:
|
||||
self._operator_sinks: Final = operator_sinks
|
||||
self._drain_seconds: Final = shutdown_drain_seconds
|
||||
self._lock: Final = threading.Condition()
|
||||
self._closed = False # guarded by ``_lock``: an unlocked read races the teardown it gates
|
||||
self._build: Final = processor_factory if processor_factory is not None else _destination_processor
|
||||
self._processors: OrderedDict[object, SpanProcessor] = OrderedDict() # mutable-ok: bounded LRU
|
||||
self._retired: OrderedDict[int, SpanProcessor] = OrderedDict() # mutable-ok: drains as exports finish
|
||||
self._exporting: dict[int, int] = {} # mutable-ok: per-processor in-flight export count
|
||||
self._drain: Final = drain_pool if drain_pool is not None else _DrainPool(capacity=pending_drains)
|
||||
|
||||
def on_start(self, span: SDKSpan, parent_context: Context | None = None) -> None:
|
||||
return None
|
||||
|
||||
def on_end(self, span: ReadableSpan) -> None:
|
||||
suppressed: Final = suppressed_backends()
|
||||
for destination in request_destinations():
|
||||
if self._operator_already_writes(destination, suppressed):
|
||||
continue
|
||||
processor = self._acquire(destination) # rebind-ok: loop variable; pyright forbids Final in a loop
|
||||
if processor is None:
|
||||
continue
|
||||
try:
|
||||
processor.on_end(_for_destination(span, destination))
|
||||
except Exception as exc: # noqa: BLE001 # one destination's failure must not cost the others their span
|
||||
verbose_logger.debug("OTel V2 fan-out: forwarding to %s failed: %s", destination.endpoint, exc)
|
||||
finally:
|
||||
self._release(processor)
|
||||
|
||||
def _operator_already_writes(self, destination: "OtelDestination", suppressed: frozenset[str]) -> bool:
|
||||
"""Whether the operator's own exporter is sending this span to the same account.
|
||||
|
||||
Only reachable under ``additive``, where nothing is suppressed: a team that
|
||||
names the operator's own project would otherwise have every span written
|
||||
there twice, once by the operator's exporter and once by the fan-out.
|
||||
"""
|
||||
return (
|
||||
destination.callback_name not in suppressed
|
||||
and _sink_key(destination.endpoint, destination.headers) in self._operator_sinks
|
||||
)
|
||||
|
||||
def shutdown(self) -> None:
|
||||
"""Close every destination processor, once the spans in flight have landed.
|
||||
|
||||
``on_end`` runs on whichever thread ends a span and can reach this fan-out
|
||||
while the SDK is tearing the provider down, so closing blind would drop a
|
||||
trace mid-forward and would hand the next caller a fresh exporter nothing
|
||||
will ever close. Refusing new work and then waiting out the in-flight ones
|
||||
keeps both from happening. A straggler past the bound is retired instead of
|
||||
closed: the thread still exporting it closes it through the drain as soon as
|
||||
its export returns, so no span is dropped mid-forward.
|
||||
|
||||
Every close then goes to the drain rather than running here. Closing a
|
||||
destination processor flushes it over the network and the SDK joins its own
|
||||
worker with no timeout of its own, so one tenant collector that answers but
|
||||
never finishes a response would otherwise hold process teardown open for as
|
||||
long as it likes. The drain's workers are daemons, and the whole teardown
|
||||
shares one deadline.
|
||||
"""
|
||||
deadline: Final = time.monotonic() + self._drain_seconds
|
||||
with self._lock:
|
||||
self._closed = True
|
||||
self._lock.wait_for(lambda: not self._exporting, timeout=self._drain_seconds)
|
||||
live: Final = tuple((id(p), p) for p in (*self._processors.values(), *self._retired.values()))
|
||||
closing: Final = tuple(p for ident, p in live if ident not in self._exporting)
|
||||
self._processors.clear()
|
||||
self._retired = OrderedDict( # mutable-ok: the same bounded map, keeping only what is still exporting
|
||||
(ident, p) for ident, p in live if ident in self._exporting
|
||||
)
|
||||
for processor in closing:
|
||||
self._drain.submit(processor)
|
||||
self._drain.close(timeout=max(0.0, deadline - time.monotonic()))
|
||||
|
||||
def force_flush(self, timeout_millis: int = 30000) -> bool:
|
||||
results: Final = tuple(self._flush_one(processor, timeout_millis) for processor in self._snapshot())
|
||||
return all(results)
|
||||
|
||||
def _snapshot(self) -> tuple[SpanProcessor, ...]:
|
||||
with self._lock:
|
||||
return (*self._processors.values(), *self._retired.values())
|
||||
|
||||
@staticmethod
|
||||
def _flush_one(processor: SpanProcessor, timeout_millis: int) -> bool:
|
||||
try:
|
||||
return processor.force_flush(timeout_millis)
|
||||
except Exception: # noqa: BLE001 # one exporter's flush failure must not fail the whole flush
|
||||
return False
|
||||
|
||||
def deliverable(self, destinations: Iterable["OtelDestination"]) -> tuple["OtelDestination", ...]:
|
||||
"""The subset of ``destinations`` this fan-out can actually export to.
|
||||
|
||||
A destination whose exporter will not build (a protocol whose package is not
|
||||
installed, a malformed endpoint) has to be dropped before the request anchors
|
||||
it, not when its first span ends. By then the operator's own exporter has been
|
||||
told to hold that backend's spans back for this request, so dropping there
|
||||
loses the span outright instead of leaving it where it would have gone with no
|
||||
override at all.
|
||||
"""
|
||||
return tuple(destination for destination in destinations if self._buildable(destination))
|
||||
|
||||
def _buildable(self, destination: "OtelDestination") -> bool:
|
||||
"""Whether a processor for ``destination`` exists or can be built right now."""
|
||||
with self._lock:
|
||||
if self._closed:
|
||||
return False
|
||||
built: Final = self._cached_or_built_locked(destination, anchored=False)
|
||||
drained: Final = self._drainable_locked()
|
||||
for shed in drained:
|
||||
self._drain.submit(shed)
|
||||
return built is not None
|
||||
|
||||
def _acquire(self, destination: "OtelDestination") -> SpanProcessor | None:
|
||||
"""The processor for ``destination``, marked busy until ``_release``.
|
||||
|
||||
The build happens under the same lock that reads the cache, so a cold cache
|
||||
met by a burst of concurrent requests yields one exporter rather than one per
|
||||
thread with all but the winner shed. Building an exporter opens no connection,
|
||||
so the cost of holding the lock is a constructor, once per destination.
|
||||
"""
|
||||
with self._lock:
|
||||
if self._closed:
|
||||
return None
|
||||
processor: Final = self._cached_or_built_locked(destination, anchored=True)
|
||||
if processor is None:
|
||||
return None
|
||||
self._exporting[id(processor)] = self._exporting.get(id(processor), 0) + 1
|
||||
drained: Final = self._drainable_locked()
|
||||
for shed in drained:
|
||||
self._drain.submit(shed)
|
||||
return processor
|
||||
|
||||
def _cached_or_built_locked(self, destination: "OtelDestination", *, anchored: bool) -> SpanProcessor | None:
|
||||
"""The cached processor for ``destination``, or a new one if the drain can take it.
|
||||
|
||||
Every build past the cache cap sheds one processor into the drain, so while the
|
||||
shed ones are stuck closing against a collector that stopped answering, a
|
||||
destination that is not yet anchored is refused rather than parked behind them:
|
||||
``deliverable`` then leaves its spans with the operator's exporter until the
|
||||
drain catches up. One the request already anchored is rebuilt regardless. The
|
||||
operator's exporter has stood down for it, so refusing here would drop the span,
|
||||
and other tenants' auths can evict it in the meantime, with that eviction being
|
||||
what tips the drain over. Eviction holds while the drain is saturated, so such a
|
||||
rebuild costs the cache one entry rather than shedding another processor, and
|
||||
the total stays at one per destination in flight.
|
||||
"""
|
||||
key: Final = destination.cache_key()
|
||||
if (cached := self._processors.get(key)) is not None:
|
||||
self._processors.move_to_end(key)
|
||||
self._retire_overflow_locked()
|
||||
return cached
|
||||
if not anchored and self._drain.saturated():
|
||||
verbose_logger.debug("OTel V2 fan-out: drain saturated, not building for %s", destination.endpoint)
|
||||
return None
|
||||
return self._build_locked(destination, key)
|
||||
|
||||
def _build_locked(self, destination: "OtelDestination", key: object) -> SpanProcessor | None:
|
||||
built: Final = self._build(destination)
|
||||
if built is None:
|
||||
return None
|
||||
self._processors[key] = built
|
||||
self._retire_overflow_locked()
|
||||
return built
|
||||
|
||||
def _release(self, processor: SpanProcessor) -> None:
|
||||
with self._lock:
|
||||
remaining: Final = self._exporting.get(id(processor), 1) - 1
|
||||
if remaining > 0:
|
||||
self._exporting[id(processor)] = remaining
|
||||
else:
|
||||
self._exporting.pop(id(processor), None)
|
||||
if not self._exporting:
|
||||
self._lock.notify_all()
|
||||
drained: Final = self._drainable_locked()
|
||||
for retired in drained:
|
||||
self._drain.submit(retired)
|
||||
|
||||
def _retire_overflow_locked(self) -> None:
|
||||
"""Move the LRU processor out of the cache once it is past the cap, drain permitting.
|
||||
|
||||
Eviction is what feeds the drain, and a destination a request already anchored
|
||||
is rebuilt on its next span, which would shed another one. While the shed ones
|
||||
are stuck closing against a collector that stopped answering, evicting would
|
||||
churn the cache at one more processor, and one more batch thread, per span.
|
||||
Holding above the cap instead keeps the total at one processor per destination
|
||||
in flight, since ``deliverable`` anchors no new destination while the drain is
|
||||
saturated. Once it has room again, every hit and build trims one entry.
|
||||
"""
|
||||
if len(self._processors) <= _MAX_CACHED_DESTINATION_PROCESSORS or self._drain.saturated():
|
||||
return
|
||||
_, evicted = self._processors.popitem(last=False)
|
||||
self._retired[id(evicted)] = evicted
|
||||
|
||||
def _drainable_locked(self) -> tuple[SpanProcessor, ...]:
|
||||
"""Retired processors no thread is exporting through, removed from the list.
|
||||
|
||||
``on_end`` holds a processor across an export, so closing an evicted one there
|
||||
drops the span it is holding. A retiree is out of the cache and can never be
|
||||
handed out again, so once its export count reaches zero it stays there.
|
||||
"""
|
||||
idle: Final = tuple(key for key in self._retired if self._exporting.get(key, 0) == 0)
|
||||
return tuple(self._retired.pop(key) for key in idle)
|
||||
|
||||
|
||||
def _destination_processor(destination: "OtelDestination") -> SpanProcessor | None:
|
||||
"""A batching OTLP processor aimed at ``destination``, or ``None`` if unbuildable.
|
||||
|
||||
A protocol that resolves to a headerless exporter is unbuildable too: the
|
||||
console fallback would swallow the tenant's credentials and print its spans to
|
||||
the proxy's stdout while the operator's exporter stands down for them.
|
||||
"""
|
||||
kind: Final = destination.protocol or "otlp_http"
|
||||
if exporter_transport(kind) == "headerless":
|
||||
verbose_logger.debug("OTel V2 fan-out: no OTLP transport for protocol %r at %s", kind, destination.endpoint)
|
||||
return None
|
||||
try:
|
||||
spec: Final = ExporterSpec(
|
||||
kind=kind,
|
||||
endpoint=destination.endpoint,
|
||||
headers=destination.header_string(),
|
||||
owner=None,
|
||||
)
|
||||
return _processor_for(_exporter_from_spec(spec), use_simple=False)
|
||||
except Exception as exc: # noqa: BLE001 # a malformed destination must not break the request or the other destinations
|
||||
verbose_logger.debug("OTel V2 fan-out: no processor for %s: %s", destination.endpoint, exc)
|
||||
return None
|
||||
|
||||
|
||||
def _shutdown_quietly(processor: SpanProcessor) -> None:
|
||||
try:
|
||||
processor.shutdown()
|
||||
except Exception as exc: # noqa: BLE001 # defensive: shedding a spare processor must not raise
|
||||
verbose_logger.debug("OTel V2 fan-out: discarding processor failed: %s", exc)
|
||||
|
||||
|
||||
class _OverriddenBackendFilter(SpanProcessor):
|
||||
"""Hold a span back from ``owner``'s operator-level exporter when the request
|
||||
pointed ``owner`` at a tenant's own account.
|
||||
|
||||
Wrapping is the only place this works: ``SynchronousMultiSpanProcessor.on_end``
|
||||
ignores return values, so a sibling processor can never veto the export.
|
||||
|
||||
Under ``additive`` mode nothing is suppressed, so the wrapper passes every span
|
||||
straight through and the operator keeps its copy.
|
||||
"""
|
||||
|
||||
def __init__(self, inner: SpanProcessor, owner: str) -> None:
|
||||
self._inner: Final = inner
|
||||
self._owner: Final = owner
|
||||
|
||||
def on_start(self, span: SDKSpan, parent_context: Context | None = None) -> None:
|
||||
self._inner.on_start(span, parent_context)
|
||||
|
||||
def on_end(self, span: ReadableSpan) -> None:
|
||||
if self._owner in suppressed_backends():
|
||||
return
|
||||
self._inner.on_end(span)
|
||||
|
||||
def shutdown(self) -> None:
|
||||
self._inner.shutdown()
|
||||
|
||||
def force_flush(self, timeout_millis: int = 30000) -> bool:
|
||||
return self._inner.force_flush(timeout_millis)
|
||||
|
||||
|
||||
def build_span_exporter(config: OpenTelemetryV2Config) -> SpanExporter:
|
||||
"""Build a single exporter from the top-level config fields.
|
||||
|
||||
|
|
@ -444,6 +1024,7 @@ def build_tracer_provider(
|
|||
exporter: SpanExporter | None = None,
|
||||
baggage_processor: SpanProcessor | None = None,
|
||||
use_simple_processor: bool | None = None,
|
||||
tenant_overrides: bool = False,
|
||||
) -> TracerProvider:
|
||||
"""Build the shared :class:`TracerProvider`.
|
||||
|
||||
|
|
@ -452,6 +1033,13 @@ def build_tracer_provider(
|
|||
``config.exporters`` entry — this is what fans spans out to multiple
|
||||
backends. ``exporter`` and ``use_simple_processor`` are explicit overrides:
|
||||
pass a single exporter to attach exactly that one (used by tests).
|
||||
|
||||
``tenant_overrides`` wraps each owned exporter so a request that pointed that
|
||||
backend at a key's or team's own account skips it. Every v2 logger's provider
|
||||
wants it, since any of them may own the overridden backend; delivering to the
|
||||
tenant is a separate job, done once by :func:`attach_tenant_fan_out`. The
|
||||
per-tenant providers this same function builds must leave it off, or they would
|
||||
filter out the very spans they exist to carry.
|
||||
"""
|
||||
provider: Final = TracerProvider(resource=build_resource(config))
|
||||
if baggage_processor is None:
|
||||
|
|
@ -468,15 +1056,107 @@ def build_tracer_provider(
|
|||
if spec.requires_headers and not spec.headers:
|
||||
continue
|
||||
exp = _exporter_from_spec(spec)
|
||||
processor = _processor_for(
|
||||
exp,
|
||||
(spec.use_simple_processor if spec.use_simple_processor is not None else use_simple_processor),
|
||||
)
|
||||
owner = spec.owner.value if spec.owner is not None else None
|
||||
provider.add_span_processor(
|
||||
_processor_for(
|
||||
exp,
|
||||
(spec.use_simple_processor if spec.use_simple_processor is not None else use_simple_processor),
|
||||
)
|
||||
_OverriddenBackendFilter(processor, owner) if tenant_overrides and owner is not None else processor
|
||||
)
|
||||
return provider
|
||||
|
||||
|
||||
_FAN_OUT_ATTACH_LOCK: Final = threading.Lock()
|
||||
|
||||
|
||||
def attach_tenant_fan_out(provider: TracerProvider, *configs: OpenTelemetryV2Config) -> None:
|
||||
"""Give ``provider`` the fan-out that delivers spans to key/team destinations.
|
||||
|
||||
Called on the one provider published as the OTel global, and idempotent so a
|
||||
second publish (a test, a re-initialized proxy) cannot double-export. Concurrent
|
||||
first calls (requests racing to anchor before any publish) serialize on one lock
|
||||
so exactly one fan-out lands. ``configs`` name the operator's own exporters, one
|
||||
config per v2 logger since each keeps its own provider and still writes its
|
||||
account, so an additive destination pointing at any of them is delivered once
|
||||
rather than twice.
|
||||
"""
|
||||
with _FAN_OUT_ATTACH_LOCK:
|
||||
if any(isinstance(processor, TenantFanOutSpanProcessor) for processor in _attached_processors(provider)):
|
||||
return
|
||||
provider.add_span_processor(TenantFanOutSpanProcessor(operator_sinks=operator_sink_keys(*configs)))
|
||||
|
||||
|
||||
def deliverable_destinations(
|
||||
destinations: Iterable["OtelDestination"],
|
||||
provider: trace.TracerProvider | None = None,
|
||||
) -> tuple["OtelDestination", ...]:
|
||||
"""The destinations a request can anchor, given what is published to carry them.
|
||||
|
||||
Anchoring a destination is what tells the operator's own exporter to stand down
|
||||
for that backend, so one nothing can deliver has to be dropped here: with no
|
||||
fan-out attached, or with an exporter that will not build, the request keeps
|
||||
exactly the routing it would have had without any override.
|
||||
"""
|
||||
fan_out: Final = next(
|
||||
(
|
||||
processor
|
||||
for processor in _attached_processors(provider if provider is not None else trace.get_tracer_provider())
|
||||
if isinstance(processor, TenantFanOutSpanProcessor)
|
||||
),
|
||||
None,
|
||||
)
|
||||
return fan_out.deliverable(destinations) if fan_out is not None else ()
|
||||
|
||||
|
||||
def operator_sink_keys(*configs: OpenTelemetryV2Config) -> frozenset[_SinkKey]:
|
||||
"""The accounts the operator's own exporters write to, in destination terms.
|
||||
|
||||
Every v2 logger's config counts, since each logger exports through its own
|
||||
provider. An exporter with no endpoint of its own resolves one from the
|
||||
environment at export time, so it has no comparable identity and is left out,
|
||||
and so is one that never reaches the wire: a console kind ignores the endpoint,
|
||||
and a header-gated spec with no credentials is skipped when the provider is built.
|
||||
"""
|
||||
return frozenset(
|
||||
key
|
||||
for config in configs
|
||||
for spec in config.exporters
|
||||
if _exports_to_the_wire(spec) and (key := _sink_key(spec.endpoint, parse_headers(spec.headers))) is not None
|
||||
)
|
||||
|
||||
|
||||
def _exports_to_the_wire(spec: ExporterSpec) -> bool:
|
||||
"""Whether ``build_tracer_provider`` gives ``spec`` an exporter that sends OTLP."""
|
||||
return exporter_transport(spec.kind) != "headerless" and not (spec.requires_headers and not spec.headers)
|
||||
|
||||
|
||||
def _sink_key(endpoint: str | None, headers: Mapping[str, str]) -> "_SinkKey | None":
|
||||
"""The account an exporter writes to, or ``None`` when it has no fixed one.
|
||||
|
||||
Normalized on the three counts that make one account look like two: the operator's
|
||||
spec carries the signal path a tenant destination leaves for the exporter to
|
||||
append, header names survive one round trip lowercased and the other not, and one
|
||||
credential answers to more than one name (see :data:`_CREDENTIAL_ALIASES`).
|
||||
"""
|
||||
normalized: Final = _otlp_traces_endpoint(endpoint)
|
||||
if normalized is None:
|
||||
return None
|
||||
return (normalized, tuple(sorted((_credential_name(name), value) for name, value in headers.items())))
|
||||
|
||||
|
||||
def _credential_name(header: str) -> str:
|
||||
"""The credential a header carries, under whichever name the backend spells it."""
|
||||
normalized: Final = header.strip().lower().replace("-", "_")
|
||||
return _CREDENTIAL_ALIASES.get(normalized, normalized)
|
||||
|
||||
|
||||
def _attached_processors(provider: trace.TracerProvider) -> "tuple[SpanProcessor, ...]":
|
||||
"""The processors already on ``provider``, or empty when the SDK hides them."""
|
||||
multi: Final = getattr(provider, "_active_span_processor", None)
|
||||
return tuple(getattr(multi, "_span_processors", ()))
|
||||
|
||||
|
||||
def get_tracer(provider: TracerProvider, name: str = "litellm") -> Tracer:
|
||||
# Stamp the instrumentation scope with the LiteLLM package version so every
|
||||
# emitted span carries a deterministic ``scope.version`` (the standard OTel
|
||||
|
|
|
|||
|
|
@ -25,6 +25,7 @@ from opentelemetry.trace import Tracer
|
|||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import OTEL_SERVICE_NAME_METADATA_KEYS
|
||||
from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
|
||||
from litellm.integrations.otel.plumbing.context import destination_backends
|
||||
from litellm.integrations.otel.plumbing.providers import (
|
||||
build_tracer_provider,
|
||||
exporter_transport,
|
||||
|
|
@ -231,10 +232,21 @@ class TenantTracerCache:
|
|||
concurrent overflow eviction can't shut it down between selection and
|
||||
the caller's span start. The caller must ``release`` it exactly once.
|
||||
"""
|
||||
# A backend with a destination is delivered by the fan-out processor, which
|
||||
# carries the whole trace and already carries this tenant's credentials and
|
||||
# service name. Routing here too would detach this span onto a second provider,
|
||||
# so the tenant would get the request tree plus a stray one-span trace.
|
||||
if self._callback_name is not None and self._callback_name in destination_backends():
|
||||
return TenantRoute(tracer=default, detached=False)
|
||||
credential_headers: Final = self._credential_headers(dynamic_params)
|
||||
project_headers: Final = self._project_headers(auth_metadata)
|
||||
service_name: Final = tenant_service_name(auth_metadata)
|
||||
if not credential_headers and not project_headers and service_name is None:
|
||||
tenant_account: Final = bool(credential_headers) or bool(project_headers)
|
||||
# A service name on its own only relabels the operator's own backend, so moving
|
||||
# the span to a second provider for it while some other backend has a
|
||||
# destination would drop the model call out of the trace the fan-out delivers.
|
||||
# The destination stamps the same service name itself.
|
||||
if not tenant_account and (service_name is None or destination_backends()):
|
||||
return TenantRoute(tracer=default, detached=False)
|
||||
# A fixed per-integration region endpoint (New Relic us/eu), never a
|
||||
# caller-supplied host; ``None`` keeps the preset's own endpoint.
|
||||
|
|
@ -255,7 +267,7 @@ class TenantTracerCache:
|
|||
_shutdown_provider(evicted)
|
||||
return TenantRoute(
|
||||
tracer=get_tracer(provider, self._tracer_name),
|
||||
detached=bool(project_headers) or bool(credential_headers),
|
||||
detached=tenant_account,
|
||||
provider=provider,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -39,6 +39,7 @@ class _AgentOpsSettings(BaseSettings):
|
|||
def agentops_preset(
|
||||
*,
|
||||
config_overrides: OpenTelemetryV2Config | None = None,
|
||||
allow_missing_credentials: bool = False,
|
||||
) -> OpenTelemetryV2Config:
|
||||
"""Build the AgentOps config without any network I/O.
|
||||
|
||||
|
|
|
|||
|
|
@ -26,10 +26,12 @@ class _ArizeSettings(BaseSettings):
|
|||
def arize_preset(
|
||||
*,
|
||||
config_overrides: OpenTelemetryV2Config | None = None,
|
||||
allow_missing_credentials: bool = False,
|
||||
) -> OpenTelemetryV2Config:
|
||||
base: Final = config_overrides or OpenTelemetryV2Config()
|
||||
mappers: Final = ensure_mappers(base.mapper_names, "openinference")
|
||||
arize_cfg: Final = _V1ArizeLogger.get_arize_config()
|
||||
headers: Final = _arize_headers(arize_cfg)
|
||||
base: Final = config_overrides or OpenTelemetryV2Config()
|
||||
return base.model_copy(
|
||||
update={
|
||||
"exporters": [
|
||||
|
|
@ -41,7 +43,7 @@ def arize_preset(
|
|||
owner=ExporterOwner.ARIZE_AX,
|
||||
),
|
||||
],
|
||||
"mapper_names": ensure_mappers(base.mapper_names, "openinference"),
|
||||
"mapper_names": mappers,
|
||||
"resource_attributes": {
|
||||
**base.resource_attributes,
|
||||
**({"model_id": arize_cfg.project_name} if arize_cfg.project_name else {}),
|
||||
|
|
|
|||
|
|
@ -18,6 +18,18 @@ class Preset(Protocol):
|
|||
|
||||
``config_overrides`` lets one preset layer onto another's config (or onto
|
||||
test-supplied defaults); the factory calls presets with no arguments.
|
||||
|
||||
``allow_missing_credentials`` lets a credential-mandatory backend (langfuse and
|
||||
weave) degrade to an exporter-less, mapper-only config instead of raising when the
|
||||
operator set no env credentials of their own. That is a real
|
||||
deployment: every team brings its own account and the operator keeps none, and
|
||||
without it the whole V2 path silently falls back to the legacy integration, so
|
||||
no team destination is ever reached. Credential-optional backends ignore it.
|
||||
"""
|
||||
|
||||
def __call__(self, *, config_overrides: OpenTelemetryV2Config | None = None) -> OpenTelemetryV2Config: ...
|
||||
def __call__(
|
||||
self,
|
||||
*,
|
||||
config_overrides: OpenTelemetryV2Config | None = None,
|
||||
allow_missing_credentials: bool = False,
|
||||
) -> OpenTelemetryV2Config: ...
|
||||
|
|
|
|||
152
litellm/integrations/otel/presets/destinations.py
Normal file
152
litellm/integrations/otel/presets/destinations.py
Normal file
|
|
@ -0,0 +1,152 @@
|
|||
"""Map a key's or team's callback vars to the OTLP destination its traces export to.
|
||||
|
||||
Header building is delegated to each preset's existing ``*_dynamic_headers`` builder,
|
||||
so a destination authenticates exactly the way the per-request tracer route already
|
||||
did; only the endpoint and transport need a per-backend rule.
|
||||
"""
|
||||
|
||||
import os
|
||||
from collections.abc import Callable, Mapping
|
||||
from functools import lru_cache
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.otel.model.destination import OtelDestination
|
||||
from litellm.litellm_core_utils.url_utils import is_url_destination_allowed_by_host
|
||||
from litellm.types.utils import StandardCallbackDynamicParams
|
||||
|
||||
#: An endpoint plus the OTLP transport to reach it with, or ``None`` when the backend
|
||||
#: names no destination. The transport is ``None`` where the backend has only one.
|
||||
_Destination = tuple[str, str | None]
|
||||
|
||||
|
||||
@lru_cache(maxsize=128)
|
||||
def _warn_host_not_allowlisted(host: str) -> None:
|
||||
"""Cached so one misconfigured team logs once rather than once per request."""
|
||||
verbose_logger.warning(
|
||||
"OTel V2: not exporting to key/team Langfuse host '%s'. Add it to "
|
||||
"litellm_settings.provider_url_destination_allowed_hosts to permit it",
|
||||
host,
|
||||
)
|
||||
|
||||
|
||||
def _langfuse_destination(params: StandardCallbackDynamicParams) -> "_Destination | None":
|
||||
"""The tenant's own Langfuse host, else the operator's, else Langfuse US cloud.
|
||||
|
||||
A host the tenant named has to be allowlisted by the operator, the same way a
|
||||
URL-valued ``model`` is: anyone who can mint a key can write it, and it becomes an
|
||||
endpoint the proxy posts the request's whole trace to, carrying the tenant's own
|
||||
credentials. The operator's own ``LANGFUSE_HOST`` is not checked, since an internal
|
||||
collector there is a deployment choice.
|
||||
"""
|
||||
from litellm.integrations.langfuse.langfuse_otel import (
|
||||
LANGFUSE_CLOUD_US_ENDPOINT,
|
||||
LangfuseOtelLogger,
|
||||
)
|
||||
|
||||
tenant_host: Final = params.get("langfuse_host") or None
|
||||
host: Final = tenant_host or LangfuseOtelLogger._get_langfuse_otel_host() # pyright: ignore[reportPrivateUsage] # reuse the backend's own env host resolver rather than duplicating it
|
||||
if not host:
|
||||
return (LANGFUSE_CLOUD_US_ENDPOINT, None)
|
||||
normalized: Final = host if host.startswith("http") else f"https://{host}"
|
||||
endpoint: Final = f"{normalized.rstrip('/')}/api/public/otel"
|
||||
if tenant_host is None:
|
||||
return (endpoint, None)
|
||||
if not is_url_destination_allowed_by_host(endpoint, litellm.provider_url_destination_allowed_hosts):
|
||||
_warn_host_not_allowlisted(host)
|
||||
return None
|
||||
return (endpoint, None)
|
||||
|
||||
|
||||
def _arize_destination(params: StandardCallbackDynamicParams) -> "_Destination | None":
|
||||
from litellm.integrations.arize.arize import ArizeLogger
|
||||
|
||||
config: Final = ArizeLogger.get_arize_config()
|
||||
return (config.endpoint, config.protocol)
|
||||
|
||||
|
||||
def _weave_destination(params: StandardCallbackDynamicParams) -> "_Destination | None":
|
||||
from litellm.integrations.weave.weave_otel import weave_otel_endpoint
|
||||
|
||||
return (weave_otel_endpoint(os.environ.get("WANDB_HOST")), None)
|
||||
|
||||
|
||||
def _newrelic_destination(params: StandardCallbackDynamicParams) -> "_Destination | None":
|
||||
from litellm.integrations.otel.presets.newrelic import newrelic_dynamic_endpoint
|
||||
|
||||
endpoint: Final = newrelic_dynamic_endpoint(params)
|
||||
return (endpoint, None) if endpoint else None
|
||||
|
||||
|
||||
#: Callback name -> destination resolver. A backend is destination-capable exactly
|
||||
#: when it appears here AND in ``DYNAMIC_HEADERS_BY_CALLBACK``: without a header
|
||||
#: builder the destination would carry no tenant credentials, and the exporter
|
||||
#: would post the tenant's traffic to the operator's account.
|
||||
_DESTINATION_BY_CALLBACK: Final[Mapping[str, Callable[[StandardCallbackDynamicParams], "_Destination | None"]]] = (
|
||||
MappingProxyType(
|
||||
{
|
||||
"langfuse_otel": _langfuse_destination,
|
||||
"arize": _arize_destination,
|
||||
"weave_otel": _weave_destination,
|
||||
"newrelic": _newrelic_destination,
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
#: Headers a destination must carry to authenticate. Several dynamic-header builders
|
||||
#: gate each credential independently, so a half-configured backend yields a non-empty
|
||||
#: but unusable header set; accepting it would suppress the operator's own exporter and
|
||||
#: send the request's whole trace where it cannot be stored.
|
||||
_REQUIRED_HEADERS_BY_CALLBACK: Final[Mapping[str, frozenset[str]]] = MappingProxyType(
|
||||
{
|
||||
"langfuse_otel": frozenset({"Authorization"}),
|
||||
"arize": frozenset({"arize-space-id", "api_key"}),
|
||||
"weave_otel": frozenset({"Authorization", "project_id"}),
|
||||
"newrelic": frozenset({"api-key"}),
|
||||
}
|
||||
)
|
||||
|
||||
_NO_ATTRS: Final[Mapping[str, str]] = MappingProxyType({})
|
||||
|
||||
|
||||
def destination_capable_backends() -> frozenset[str]:
|
||||
"""Backends a key or team can point at its own account."""
|
||||
from litellm.integrations.otel.presets import DYNAMIC_HEADERS_BY_CALLBACK
|
||||
|
||||
return frozenset(_DESTINATION_BY_CALLBACK) & frozenset(DYNAMIC_HEADERS_BY_CALLBACK)
|
||||
|
||||
|
||||
def destination_for(
|
||||
callback_name: str,
|
||||
params: StandardCallbackDynamicParams,
|
||||
service_name: str | None = None,
|
||||
) -> OtelDestination | None:
|
||||
"""The destination ``params`` names for ``callback_name``, or ``None``.
|
||||
|
||||
``None`` means the caller configured nothing usable for this backend, so the
|
||||
request keeps the operator's global exporters. ``service_name`` is the key's or
|
||||
team's ``otel_service_name``, which the per-request tracer route applies when the
|
||||
backend is not overridden and the destination has to apply once it is.
|
||||
"""
|
||||
from litellm.integrations.otel.presets import DYNAMIC_HEADERS_BY_CALLBACK
|
||||
|
||||
header_builder: Final = DYNAMIC_HEADERS_BY_CALLBACK.get(callback_name)
|
||||
destination_builder: Final = _DESTINATION_BY_CALLBACK.get(callback_name)
|
||||
if header_builder is None or destination_builder is None:
|
||||
return None
|
||||
headers: Final = header_builder(params)
|
||||
if not headers or not _REQUIRED_HEADERS_BY_CALLBACK[callback_name] <= frozenset(headers):
|
||||
return None
|
||||
resolved: Final = destination_builder(params)
|
||||
if resolved is None:
|
||||
return None
|
||||
endpoint, protocol = resolved
|
||||
return OtelDestination(
|
||||
endpoint=endpoint,
|
||||
headers=MappingProxyType(dict(headers)), # mutable-ok: MappingProxyType needs a concrete mapping to wrap
|
||||
resource_attributes=MappingProxyType({"service.name": service_name}) if service_name else _NO_ATTRS,
|
||||
callback_name=callback_name,
|
||||
protocol=protocol,
|
||||
)
|
||||
|
|
@ -10,17 +10,32 @@ from litellm.integrations.otel.model.config import (
|
|||
ExporterSpec,
|
||||
OpenTelemetryV2Config,
|
||||
)
|
||||
from litellm.integrations.otel.presets.utils import ensure_mappers
|
||||
from litellm.integrations.otel.presets.utils import (
|
||||
credential_gated_exporters,
|
||||
ensure_mappers,
|
||||
)
|
||||
from litellm.types.utils import StandardCallbackDynamicParams
|
||||
|
||||
|
||||
def langfuse_preset(
|
||||
*,
|
||||
config_overrides: OpenTelemetryV2Config | None = None,
|
||||
allow_missing_credentials: bool = False,
|
||||
) -> OpenTelemetryV2Config:
|
||||
cfg: Final = _V1Langfuse.get_langfuse_otel_config()
|
||||
kind: Final = cfg.exporter if isinstance(cfg.exporter, str) else "otlp_http"
|
||||
base: Final = config_overrides or OpenTelemetryV2Config()
|
||||
mappers: Final = ensure_mappers(base.mapper_names, "langfuse")
|
||||
try:
|
||||
cfg: Final = _V1Langfuse.get_langfuse_otel_config()
|
||||
except Exception:
|
||||
if not allow_missing_credentials:
|
||||
raise
|
||||
return base.model_copy(
|
||||
update={ # mutable-ok: pydantic model_copy takes a plain update mapping
|
||||
"exporters": credential_gated_exporters(base.exporters, ExporterOwner.LANGFUSE_OTEL),
|
||||
"mapper_names": mappers,
|
||||
}
|
||||
)
|
||||
kind: Final = cfg.exporter if isinstance(cfg.exporter, str) else "otlp_http"
|
||||
return base.model_copy(
|
||||
update={
|
||||
"exporters": [
|
||||
|
|
@ -32,7 +47,7 @@ def langfuse_preset(
|
|||
owner=ExporterOwner.LANGFUSE_OTEL,
|
||||
),
|
||||
],
|
||||
"mapper_names": ensure_mappers(base.mapper_names, "langfuse"),
|
||||
"mapper_names": mappers,
|
||||
}
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ from litellm.integrations.otel.presets.utils import ensure_mappers
|
|||
def langtrace_preset(
|
||||
*,
|
||||
config_overrides: OpenTelemetryV2Config | None = None,
|
||||
allow_missing_credentials: bool = False,
|
||||
) -> OpenTelemetryV2Config:
|
||||
"""Compose the Langtrace mapper on top of the customer's OTLP destination.
|
||||
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ from litellm.integrations.otel.model.config import (
|
|||
def levo_preset(
|
||||
*,
|
||||
config_overrides: OpenTelemetryV2Config | None = None,
|
||||
allow_missing_credentials: bool = False,
|
||||
) -> OpenTelemetryV2Config:
|
||||
cfg: Final = _V1Levo.get_levo_config()
|
||||
base: Final = config_overrides or OpenTelemetryV2Config()
|
||||
|
|
|
|||
|
|
@ -44,6 +44,7 @@ class _NewRelicSettings(BaseSettings):
|
|||
def newrelic_preset(
|
||||
*,
|
||||
config_overrides: OpenTelemetryV2Config | None = None,
|
||||
allow_missing_credentials: bool = False,
|
||||
) -> OpenTelemetryV2Config:
|
||||
settings: Final = _NewRelicSettings()
|
||||
base: Final = config_overrides or OpenTelemetryV2Config()
|
||||
|
|
|
|||
|
|
@ -60,6 +60,7 @@ def phoenix_project_headers(auth_metadata: Mapping[str, str] | None) -> Mapping[
|
|||
def phoenix_preset(
|
||||
*,
|
||||
config_overrides: OpenTelemetryV2Config | None = None,
|
||||
allow_missing_credentials: bool = False,
|
||||
) -> OpenTelemetryV2Config:
|
||||
cfg: Final = _V1Phoenix.get_arize_phoenix_config()
|
||||
headers: Final = cfg.otlp_auth_headers if hasattr(cfg, "otlp_auth_headers") else None
|
||||
|
|
|
|||
|
|
@ -3,6 +3,8 @@
|
|||
from collections.abc import Iterable
|
||||
from typing import Final
|
||||
|
||||
from litellm.integrations.otel.model.config import ExporterOwner, ExporterSpec
|
||||
|
||||
|
||||
def ensure_mappers(mapper_names: Iterable[str], *names: str) -> list[str]:
|
||||
"""Return ``mapper_names`` with each of ``names`` appended if not already present.
|
||||
|
|
@ -15,3 +17,32 @@ def ensure_mappers(mapper_names: Iterable[str], *names: str) -> list[str]:
|
|||
if name not in result:
|
||||
result.append(name)
|
||||
return result
|
||||
|
||||
|
||||
def credential_gated_exporters(
|
||||
exporters: "Iterable[ExporterSpec]", owner: "ExporterOwner"
|
||||
) -> "tuple[ExporterSpec, ...]":
|
||||
"""``exporters`` with the operator's destination replaced by a header-gated one.
|
||||
|
||||
Used when a credential-mandatory backend is asked to build without the operator's
|
||||
own credentials, so only key/team destinations receive spans. Two things have to
|
||||
happen for that to mean "export nowhere": the placeholder console spec that
|
||||
``OpenTelemetryV2Config`` folds in for an empty exporter list is dropped, or every
|
||||
span would be printed to stdout, and the gated spec keeps the owner so the
|
||||
override filter still recognises which backend this provider speaks for.
|
||||
"""
|
||||
return (
|
||||
*(spec for spec in exporters if not is_unconfigured_placeholder(spec)),
|
||||
ExporterSpec(owner=owner, requires_headers=True),
|
||||
)
|
||||
|
||||
|
||||
def is_unconfigured_placeholder(spec: "ExporterSpec") -> bool:
|
||||
"""Whether ``spec`` is the one ``_normalize`` folds in when nothing was configured.
|
||||
|
||||
No field set is what says the operator asked for nothing: an exporter they did
|
||||
configure survives, even ``OTEL_EXPORTER=console`` whose value matches the default,
|
||||
and so does the gated spec this module appends, which would otherwise eat itself
|
||||
when one preset layers onto another.
|
||||
"""
|
||||
return not spec.model_fields_set
|
||||
|
|
|
|||
|
|
@ -7,7 +7,10 @@ from litellm.integrations.otel.model.config import (
|
|||
ExporterSpec,
|
||||
OpenTelemetryV2Config,
|
||||
)
|
||||
from litellm.integrations.otel.presets.utils import ensure_mappers
|
||||
from litellm.integrations.otel.presets.utils import (
|
||||
credential_gated_exporters,
|
||||
ensure_mappers,
|
||||
)
|
||||
from litellm.integrations.weave.weave_otel import (
|
||||
_get_weave_authorization_header,
|
||||
get_weave_otel_config,
|
||||
|
|
@ -18,9 +21,21 @@ from litellm.types.utils import StandardCallbackDynamicParams
|
|||
def weave_preset(
|
||||
*,
|
||||
config_overrides: OpenTelemetryV2Config | None = None,
|
||||
allow_missing_credentials: bool = False,
|
||||
) -> OpenTelemetryV2Config:
|
||||
weave_cfg: Final = get_weave_otel_config()
|
||||
base: Final = config_overrides or OpenTelemetryV2Config()
|
||||
mappers: Final = ensure_mappers(base.mapper_names, "openinference", "weave")
|
||||
try:
|
||||
weave_cfg: Final = get_weave_otel_config()
|
||||
except Exception:
|
||||
if not allow_missing_credentials:
|
||||
raise
|
||||
return base.model_copy(
|
||||
update={ # mutable-ok: pydantic model_copy takes a plain update mapping
|
||||
"exporters": credential_gated_exporters(base.exporters, ExporterOwner.WEAVE_OTEL),
|
||||
"mapper_names": mappers,
|
||||
}
|
||||
)
|
||||
return base.model_copy(
|
||||
update={
|
||||
"exporters": [
|
||||
|
|
@ -33,7 +48,7 @@ def weave_preset(
|
|||
),
|
||||
],
|
||||
# Weave consumes OpenInference + a small Weave-specific overlay.
|
||||
"mapper_names": ensure_mappers(base.mapper_names, "openinference", "weave"),
|
||||
"mapper_names": mappers,
|
||||
}
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -117,6 +117,14 @@ def _get_weave_authorization_header(api_key: str) -> str:
|
|||
return f"Basic {auth_header}"
|
||||
|
||||
|
||||
def weave_otel_endpoint(host: str | None) -> str:
|
||||
"""The OTLP traces endpoint for a self-managed ``host``, else Weave cloud."""
|
||||
if not host:
|
||||
return WEAVE_BASE_URL + WEAVE_OTEL_ENDPOINT
|
||||
normalized: Final = host if host.startswith("http") else f"https://{host}"
|
||||
return normalized.rstrip("/") + WEAVE_OTEL_ENDPOINT
|
||||
|
||||
|
||||
def get_weave_otel_config() -> WeaveOtelConfig:
|
||||
"""
|
||||
Retrieves the Weave OpenTelemetry configuration based on environment variables.
|
||||
|
|
@ -134,7 +142,6 @@ def get_weave_otel_config() -> WeaveOtelConfig:
|
|||
"""
|
||||
api_key: Final = os.getenv("WANDB_API_KEY")
|
||||
project_id: Final = os.getenv("WANDB_PROJECT_ID")
|
||||
host = os.getenv("WANDB_HOST")
|
||||
|
||||
if not api_key:
|
||||
raise ValueError("WANDB_API_KEY must be set for Weave OpenTelemetry integration.")
|
||||
|
|
@ -144,15 +151,8 @@ def get_weave_otel_config() -> WeaveOtelConfig:
|
|||
"WANDB_PROJECT_ID must be set for Weave OpenTelemetry integration. Format: <entity>/<project_name>"
|
||||
)
|
||||
|
||||
if host:
|
||||
if not host.startswith("http"):
|
||||
host = "https://" + host
|
||||
# Self-managed instances use a different path
|
||||
endpoint = host.rstrip("/") + WEAVE_OTEL_ENDPOINT
|
||||
verbose_logger.debug("Using Weave OTEL endpoint from host: %s", endpoint)
|
||||
else:
|
||||
endpoint = WEAVE_BASE_URL + WEAVE_OTEL_ENDPOINT
|
||||
verbose_logger.debug("Using Weave cloud endpoint: %s", endpoint)
|
||||
endpoint: Final = weave_otel_endpoint(os.getenv("WANDB_HOST"))
|
||||
verbose_logger.debug("Using Weave OTEL endpoint: %s", endpoint)
|
||||
|
||||
# Weave uses Basic auth with format: api:<WANDB_API_KEY>
|
||||
auth_header: Final = _get_weave_authorization_header(api_key=api_key)
|
||||
|
|
|
|||
|
|
@ -202,6 +202,7 @@ if TYPE_CHECKING:
|
|||
from mcp.types import EmbeddedResource, ImageContent, TextContent
|
||||
|
||||
from litellm.integrations.otel.logger import OpenTelemetryV2
|
||||
from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
|
||||
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig
|
||||
try:
|
||||
from litellm_enterprise.enterprise_callbacks.callback_controls import (
|
||||
|
|
@ -4850,31 +4851,83 @@ def _maybe_construct_otel_v2(callback_name: str, _in_memory_loggers: list[Custom
|
|||
|
||||
Returns ``None`` when V2 is off OR when there's no preset registered for
|
||||
``callback_name`` — callers should then fall through to the legacy path.
|
||||
|
||||
A preset that needs operator credentials it cannot find is allowed to build
|
||||
only when this request has a key/team destination for that backend and another
|
||||
V2 logger is already registered to carry the fan-out. The resulting logger keeps
|
||||
only its credential-gated exporter, while the registered logger owns operator
|
||||
delivery. Without that carrier, a preset that raises or that ends up with nothing
|
||||
but its gated exporter and the default console placeholder returns ``None``, so the
|
||||
caller falls through to the legacy path exactly as before V2 landed.
|
||||
"""
|
||||
from litellm.integrations.otel.model.config import is_otel_v2_enabled
|
||||
|
||||
if not is_otel_v2_enabled():
|
||||
return None
|
||||
from litellm.integrations.otel.logger import OpenTelemetryV2, build_otel_v2_logger
|
||||
from litellm.integrations.otel.plumbing.context import destination_backends
|
||||
from litellm.integrations.otel.presets import PRESET_BY_CALLBACK
|
||||
|
||||
preset_fn: Final = PRESET_BY_CALLBACK.get(callback_name)
|
||||
if preset_fn is None:
|
||||
return None
|
||||
serves_a_destination: Final = callback_name in destination_backends()
|
||||
has_v2_logger: Final = any(isinstance(callback, OpenTelemetryV2) for callback in _in_memory_loggers)
|
||||
carried: Final = serves_a_destination and has_v2_logger
|
||||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, OpenTelemetryV2) and getattr(callback, "callback_name", None) == callback_name:
|
||||
if (
|
||||
isinstance(callback, OpenTelemetryV2)
|
||||
and getattr(callback, "callback_name", None) == callback_name
|
||||
and (serves_a_destination or not _exports_nowhere(callback.config))
|
||||
):
|
||||
return callback
|
||||
try:
|
||||
config: Final = preset_fn()
|
||||
built: Final = preset_fn(allow_missing_credentials=carried)
|
||||
except Exception:
|
||||
# If env vars are missing or the preset raises, defer to the legacy path
|
||||
# so customers get the same error story they had before V2 landed.
|
||||
return None
|
||||
gated: Final = _is_credential_gated(built)
|
||||
if gated and not carried and not _has_operator_exporter(built):
|
||||
return None
|
||||
config: Final = _only_the_gated_exporter(built) if gated and carried else built
|
||||
if _exports_nowhere(config):
|
||||
verbose_logger.warning(
|
||||
"OTel V2: no operator credentials for '%s'; only key/team destinations will receive its traces",
|
||||
callback_name,
|
||||
)
|
||||
v2_logger: Final = build_otel_v2_logger(config=config, callback_name=callback_name)
|
||||
_in_memory_loggers.append(v2_logger)
|
||||
return v2_logger
|
||||
|
||||
|
||||
def _exports_nowhere(config: "OpenTelemetryV2Config") -> bool:
|
||||
"""Whether every exporter in ``config`` is waiting on credentials it never got."""
|
||||
return all(_is_gated(spec) for spec in config.exporters)
|
||||
|
||||
|
||||
def _is_credential_gated(config: "OpenTelemetryV2Config") -> bool:
|
||||
"""Whether the preset built without the operator's own credentials for its backend."""
|
||||
return any(_is_gated(spec) for spec in config.exporters)
|
||||
|
||||
|
||||
def _has_operator_exporter(config: "OpenTelemetryV2Config") -> bool:
|
||||
"""Whether the operator configured somewhere real to export, beyond the default console placeholder."""
|
||||
from litellm.integrations.otel.presets.utils import is_unconfigured_placeholder
|
||||
|
||||
return any(not _is_gated(spec) and not is_unconfigured_placeholder(spec) for spec in config.exporters)
|
||||
|
||||
|
||||
def _only_the_gated_exporter(config: "OpenTelemetryV2Config") -> "OpenTelemetryV2Config":
|
||||
return config.model_copy(
|
||||
update={"exporters": [spec for spec in config.exporters if _is_gated(spec)]} # mutable-ok: model_copy update
|
||||
)
|
||||
|
||||
|
||||
def _is_gated(spec: "ExporterSpec") -> bool:
|
||||
return spec.requires_headers and not spec.headers
|
||||
|
||||
|
||||
def _maybe_auto_initialize_arize_phoenix(_in_memory_loggers: list[CustomLogger]) -> None:
|
||||
"""
|
||||
Auto-initialize ArizePhoenixLogger when Phoenix env vars are detected.
|
||||
|
|
|
|||
|
|
@ -780,6 +780,7 @@ class PromptTokensDetailsResult(TypedDict):
|
|||
image_count: int
|
||||
video_length_seconds: float
|
||||
audio_length_seconds: float
|
||||
query_count: int
|
||||
|
||||
|
||||
def parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult:
|
||||
|
|
@ -828,6 +829,7 @@ def parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult:
|
|||
)
|
||||
or 0.0
|
||||
)
|
||||
query_count: Final = _coerce_token_count(getattr(usage.prompt_tokens_details, "query_count", 0))
|
||||
|
||||
return PromptTokensDetailsResult(
|
||||
cache_hit_tokens=cache_hit_tokens,
|
||||
|
|
@ -841,6 +843,7 @@ def parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult:
|
|||
image_count=image_count,
|
||||
video_length_seconds=float(video_length_seconds),
|
||||
audio_length_seconds=float(audio_length_seconds),
|
||||
query_count=query_count,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -978,6 +981,11 @@ def _calculate_input_cost(
|
|||
prompt_tokens_details["audio_length_seconds"],
|
||||
)
|
||||
|
||||
if prompt_tokens_details["query_count"]:
|
||||
prompt_cost += calculate_cost_component(
|
||||
model_info, "input_cost_per_query", prompt_tokens_details["query_count"]
|
||||
)
|
||||
|
||||
return prompt_cost
|
||||
|
||||
|
||||
|
|
@ -1149,6 +1157,7 @@ def generic_cost_per_token(
|
|||
image_count=0,
|
||||
video_length_seconds=0.0,
|
||||
audio_length_seconds=0.0,
|
||||
query_count=0,
|
||||
)
|
||||
if usage.prompt_tokens_details:
|
||||
prompt_tokens_details = parse_prompt_tokens_details(usage)
|
||||
|
|
|
|||
|
|
@ -183,6 +183,7 @@ class _RemoteSource:
|
|||
class RemoteMedia:
|
||||
url: str
|
||||
fields: Mapping[str, object]
|
||||
part_type: str
|
||||
|
||||
|
||||
_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
|
@ -192,6 +193,10 @@ def inline_every_remote_url(_media: RemoteMedia) -> bool:
|
|||
return True
|
||||
|
||||
|
||||
def inline_remote_image_urls(media: RemoteMedia) -> bool:
|
||||
return media.part_type == "image_url"
|
||||
|
||||
|
||||
def _parse_remote_image(fields: Mapping[str, object]) -> _RemoteImage | None:
|
||||
if fields.get("type") != "image_url":
|
||||
return None
|
||||
|
|
@ -223,11 +228,11 @@ def _parse_remote_part(part: object) -> _RemoteImage | _RemoteFile | _RemoteSour
|
|||
def _remote_media(remote: _RemoteImage | _RemoteFile | _RemoteSource) -> RemoteMedia:
|
||||
match remote:
|
||||
case _RemoteImage(_, image_url, url):
|
||||
return RemoteMedia(url, image_url if image_url is not None else _NO_FIELDS)
|
||||
return RemoteMedia(url, image_url if image_url is not None else _NO_FIELDS, "image_url")
|
||||
case _RemoteFile(_, file, url):
|
||||
return RemoteMedia(url, file)
|
||||
case _RemoteSource(_, source, url):
|
||||
return RemoteMedia(url, source)
|
||||
return RemoteMedia(url, file, "file")
|
||||
case _RemoteSource(part, source, url):
|
||||
return RemoteMedia(url, source, str(part.get("type")))
|
||||
|
||||
|
||||
_PDF_FORMAT: Final = MappingProxyType({"format": "application/pdf"})
|
||||
|
|
|
|||
|
|
@ -368,27 +368,92 @@ class AnthropicChatCompletion(BaseLLM):
|
|||
if config is None:
|
||||
raise ValueError(f"Provider config not found for model: {model} and provider: {custom_llm_provider}")
|
||||
|
||||
def build_request() -> tuple[dict, dict]: # mutable-ok: rewritten in place downstream
|
||||
"""Translate the request the Python way, returning `(headers, data)`.
|
||||
transform_params: Final = {**optional_params, "is_vertex_request": is_vertex_request}
|
||||
|
||||
def finish_request(request_data: dict) -> tuple[dict, dict]: # mutable-ok: rewritten in place downstream
|
||||
"""Filter beta headers and emit pre_call, returning `(headers, data)`.
|
||||
|
||||
The pair stays mutable because the streaming path rewrites it in
|
||||
place (`data["stream"] = True`) before sending.
|
||||
|
||||
Shared by the normal path and by the Rust path's fallback, which
|
||||
builds it only when the Rust call did not serve the request.
|
||||
place (`data["stream"] = True`) before sending. A Rust attempt that
|
||||
declined already emitted pre_call for this request, so skip it there.
|
||||
"""
|
||||
request_data: Final = config.transform_request(
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params={**optional_params, "is_vertex_request": is_vertex_request},
|
||||
litellm_params=litellm_params,
|
||||
headers=headers,
|
||||
)
|
||||
return update_request_with_filtered_beta(
|
||||
request_headers, data = update_request_with_filtered_beta(
|
||||
headers=headers,
|
||||
request_data=request_data,
|
||||
provider=custom_llm_provider,
|
||||
)
|
||||
if not serves_via_rust:
|
||||
logging_obj.pre_call(
|
||||
input=messages,
|
||||
api_key=api_key,
|
||||
additional_args={
|
||||
"complete_input_dict": data,
|
||||
"api_base": api_base,
|
||||
"headers": request_headers,
|
||||
},
|
||||
)
|
||||
print_verbose(f"_is_function_call: {_is_function_call}")
|
||||
return request_headers, data
|
||||
|
||||
async def acompletion_dispatch() -> "ModelResponse | CustomStreamWrapper":
|
||||
"""Translate then send, so the provider config can inline remote media off the event loop."""
|
||||
request_headers, data = finish_request(
|
||||
await config.async_transform_request(
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params=transform_params,
|
||||
litellm_params=litellm_params,
|
||||
headers=headers,
|
||||
)
|
||||
)
|
||||
if (
|
||||
stream is True
|
||||
): # if function call - fake the streaming (need complete blocks for output parsing in openai format)
|
||||
print_verbose("makes async anthropic streaming POST request")
|
||||
data["stream"] = stream
|
||||
return await self.acompletion_stream_function(
|
||||
model=model,
|
||||
messages=messages,
|
||||
data=data,
|
||||
api_base=api_base,
|
||||
custom_prompt_dict=custom_prompt_dict,
|
||||
model_response=model_response,
|
||||
print_verbose=print_verbose,
|
||||
encoding=encoding,
|
||||
api_key=api_key,
|
||||
logging_obj=logging_obj,
|
||||
optional_params=optional_params,
|
||||
stream=stream,
|
||||
_is_function_call=_is_function_call,
|
||||
json_mode=json_mode,
|
||||
litellm_params=litellm_params,
|
||||
logger_fn=logger_fn,
|
||||
headers=request_headers,
|
||||
timeout=timeout,
|
||||
client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None),
|
||||
)
|
||||
return await self.acompletion_function(
|
||||
model=model,
|
||||
messages=messages,
|
||||
data=data,
|
||||
api_base=api_base,
|
||||
custom_prompt_dict=custom_prompt_dict,
|
||||
model_response=model_response,
|
||||
print_verbose=print_verbose,
|
||||
encoding=encoding,
|
||||
api_key=api_key,
|
||||
provider_config=config,
|
||||
logging_obj=logging_obj,
|
||||
optional_params=optional_params,
|
||||
stream=stream,
|
||||
_is_function_call=_is_function_call,
|
||||
litellm_params=litellm_params,
|
||||
logger_fn=logger_fn,
|
||||
headers=request_headers,
|
||||
client=client,
|
||||
json_mode=json_mode,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
# The Rust core owns the whole call for the subset it accepts, so ask
|
||||
# before transforming: whichever path runs emits pre_call exactly once.
|
||||
|
|
@ -424,35 +489,6 @@ class AnthropicChatCompletion(BaseLLM):
|
|||
additional_args=rust_logging_args,
|
||||
)
|
||||
if acompletion is True:
|
||||
|
||||
async def python_fallback() -> "ModelResponse | CustomStreamWrapper":
|
||||
# pre_call already fired for this request above. The Rust
|
||||
# path only declines before the provider is called, so this
|
||||
# is the same attempt continuing, not a second one.
|
||||
fallback_headers, fallback_data = build_request()
|
||||
return await self.acompletion_function(
|
||||
model=model,
|
||||
messages=messages,
|
||||
data=fallback_data,
|
||||
api_base=api_base,
|
||||
custom_prompt_dict=custom_prompt_dict,
|
||||
model_response=model_response,
|
||||
print_verbose=print_verbose,
|
||||
encoding=encoding,
|
||||
api_key=api_key,
|
||||
provider_config=config,
|
||||
logging_obj=logging_obj,
|
||||
optional_params=optional_params,
|
||||
stream=stream,
|
||||
_is_function_call=_is_function_call,
|
||||
litellm_params=litellm_params,
|
||||
logger_fn=logger_fn,
|
||||
headers=fallback_headers,
|
||||
client=client,
|
||||
json_mode=json_mode,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
return rust_chat_completions_bridge.achat_completions_or_fallback(
|
||||
model=model,
|
||||
messages=messages,
|
||||
|
|
@ -464,7 +500,7 @@ class AnthropicChatCompletion(BaseLLM):
|
|||
extra_headers=headers,
|
||||
timeout=timeout,
|
||||
on_response=log_rust_post_call,
|
||||
python_fallback=python_fallback,
|
||||
python_fallback=acompletion_dispatch,
|
||||
)
|
||||
rust_response: Final = rust_chat_completions_bridge.chat_completions(
|
||||
model=model,
|
||||
|
|
@ -481,74 +517,18 @@ class AnthropicChatCompletion(BaseLLM):
|
|||
if rust_response is not None:
|
||||
return rust_response
|
||||
|
||||
headers, data = build_request()
|
||||
|
||||
## LOGGING
|
||||
# Reaching here with `serves_via_rust` set means the Rust attempt
|
||||
# declined at call time, before the provider was called, and already
|
||||
# logged this request. That is the same attempt continuing.
|
||||
if not serves_via_rust:
|
||||
logging_obj.pre_call(
|
||||
input=messages,
|
||||
api_key=api_key,
|
||||
additional_args={
|
||||
"complete_input_dict": data,
|
||||
"api_base": api_base,
|
||||
"headers": headers,
|
||||
},
|
||||
)
|
||||
print_verbose(f"_is_function_call: {_is_function_call}")
|
||||
if acompletion is True:
|
||||
if (
|
||||
stream is True
|
||||
): # if function call - fake the streaming (need complete blocks for output parsing in openai format)
|
||||
print_verbose("makes async anthropic streaming POST request")
|
||||
data["stream"] = stream
|
||||
return self.acompletion_stream_function(
|
||||
model=model,
|
||||
messages=messages,
|
||||
data=data,
|
||||
api_base=api_base,
|
||||
custom_prompt_dict=custom_prompt_dict,
|
||||
model_response=model_response,
|
||||
print_verbose=print_verbose,
|
||||
encoding=encoding,
|
||||
api_key=api_key,
|
||||
logging_obj=logging_obj,
|
||||
optional_params=optional_params,
|
||||
stream=stream,
|
||||
_is_function_call=_is_function_call,
|
||||
json_mode=json_mode,
|
||||
litellm_params=litellm_params,
|
||||
logger_fn=logger_fn,
|
||||
headers=headers,
|
||||
timeout=timeout,
|
||||
client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None),
|
||||
)
|
||||
else:
|
||||
return self.acompletion_function(
|
||||
model=model,
|
||||
messages=messages,
|
||||
data=data,
|
||||
api_base=api_base,
|
||||
custom_prompt_dict=custom_prompt_dict,
|
||||
model_response=model_response,
|
||||
print_verbose=print_verbose,
|
||||
encoding=encoding,
|
||||
api_key=api_key,
|
||||
provider_config=config,
|
||||
logging_obj=logging_obj,
|
||||
optional_params=optional_params,
|
||||
stream=stream,
|
||||
_is_function_call=_is_function_call,
|
||||
litellm_params=litellm_params,
|
||||
logger_fn=logger_fn,
|
||||
headers=headers,
|
||||
client=client,
|
||||
json_mode=json_mode,
|
||||
timeout=timeout,
|
||||
)
|
||||
return acompletion_dispatch()
|
||||
else:
|
||||
headers, data = finish_request(
|
||||
config.transform_request(
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params=transform_params,
|
||||
litellm_params=litellm_params,
|
||||
headers=headers,
|
||||
)
|
||||
)
|
||||
## COMPLETION CALL
|
||||
if (
|
||||
stream is True
|
||||
|
|
|
|||
|
|
@ -26,6 +26,11 @@ from litellm.litellm_core_utils.core_helpers import map_finish_reason
|
|||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
sanitize_input_schema_for_anthropic,
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.image_handling import (
|
||||
RemoteMedia,
|
||||
async_inline_remote_media,
|
||||
inline_remote_image_urls,
|
||||
)
|
||||
from litellm.llms.base_llm.base_utils import type_to_response_format_param
|
||||
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
|
||||
from litellm.types.llms.anthropic import (
|
||||
|
|
@ -1840,6 +1845,25 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
break
|
||||
return headers
|
||||
|
||||
def inlines_remote_media(self, media: RemoteMedia) -> bool:
|
||||
return inline_remote_image_urls(media) and media.url.startswith("http://")
|
||||
|
||||
async def async_transform_request(
|
||||
self,
|
||||
model: str,
|
||||
messages: list[AllMessageValues], # mutable-ok: BaseConfig signature
|
||||
optional_params: dict[str, object], # mutable-ok: BaseConfig signature
|
||||
litellm_params: dict[str, object], # mutable-ok: BaseConfig signature
|
||||
headers: dict[str, object], # mutable-ok: BaseConfig signature
|
||||
) -> dict[str, object]: # mutable-ok: BaseConfig signature
|
||||
return self.transform_request(
|
||||
model=model,
|
||||
messages=await async_inline_remote_media(messages, should_inline=self.inlines_remote_media),
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
def transform_request(
|
||||
self,
|
||||
model: str,
|
||||
|
|
|
|||
|
|
@ -121,6 +121,9 @@ class RouterVectorStoreEmbeddingExecutor:
|
|||
|
||||
|
||||
class BaseVectorStoreConfig:
|
||||
def validate_create_vector_store(self) -> None:
|
||||
return None
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list[VECTOR_STORE_OPENAI_PARAMS]:
|
||||
return []
|
||||
|
||||
|
|
|
|||
|
|
@ -35,7 +35,7 @@ from .amazon_titan_multimodal_transformation import (
|
|||
)
|
||||
from .amazon_titan_v2_transformation import AmazonTitanV2Config
|
||||
from .cohere_transformation import BedrockCohereEmbeddingConfig
|
||||
from .twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig
|
||||
from .twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig, drop_params_enabled
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
|
@ -239,7 +239,7 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
returned_response = AmazonTitanG1Config()._transform_response(response_list=response_list, model=model)
|
||||
elif provider == "twelvelabs":
|
||||
returned_response = TwelveLabsMarengoEmbeddingConfig()._transform_response(
|
||||
response_list=response_list, model=model
|
||||
response_list=response_list, model=model, batch_data=batch_data
|
||||
)
|
||||
elif provider == "nova":
|
||||
returned_response = AmazonNovaEmbeddingConfig()._transform_response(
|
||||
|
|
@ -484,12 +484,13 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
elif provider == "twelvelabs":
|
||||
batch_data = []
|
||||
for i in input:
|
||||
twelvelabs_request = TwelveLabsMarengoEmbeddingConfig()._transform_request(
|
||||
twelvelabs_request = TwelveLabsMarengoEmbeddingConfig(model=model)._transform_request(
|
||||
input=i,
|
||||
inference_params=inference_params,
|
||||
async_invoke_route=has_async_invoke,
|
||||
model_id=modelId,
|
||||
output_s3_uri=inference_params.get("output_s3_uri"),
|
||||
drop_params=drop_params_enabled(litellm_params),
|
||||
)
|
||||
batch_data.append(twelvelabs_request)
|
||||
elif provider == "nova":
|
||||
|
|
|
|||
|
|
@ -0,0 +1,239 @@
|
|||
"""
|
||||
Request builder for Bedrock TwelveLabs Marengo Embed 3.0, whose payload nests the input under a key named after
|
||||
``inputType`` instead of the flat 2.7 layout.
|
||||
|
||||
Docs - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-marengo-3.html
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm.llms.bedrock.common_utils import BedrockError
|
||||
from litellm.types.llms.bedrock import (
|
||||
TWELVELABS_MARENGO_3_EMBEDDING_OPTIONS,
|
||||
TWELVELABS_MARENGO_3_EMBEDDING_SCOPES,
|
||||
TWELVELABS_MARENGO_3_EMBEDDING_TYPES,
|
||||
TWELVELABS_MARENGO_3_INPUT_TYPES,
|
||||
TwelveLabsMarengo3AudioRequest,
|
||||
TwelveLabsMarengo3EmbeddingRequest,
|
||||
TwelveLabsMarengo3ImageRequest,
|
||||
TwelveLabsMarengo3MultiInputRequest,
|
||||
TwelveLabsMarengo3NamedMediaSource,
|
||||
TwelveLabsMarengo3RequestBase,
|
||||
TwelveLabsMarengo3Segmentation,
|
||||
TwelveLabsMarengo3TextImageRequest,
|
||||
TwelveLabsMarengo3TextRequest,
|
||||
TwelveLabsMarengo3TimedMediaInput,
|
||||
TwelveLabsMarengo3TimedMediaOptions,
|
||||
TwelveLabsMarengo3VideoRequest,
|
||||
TwelveLabsMediaSource,
|
||||
TwelveLabsS3Location,
|
||||
)
|
||||
from litellm.utils import get_base64_str
|
||||
|
||||
MARENGO_3_MODEL_MARKER: Final = "marengo-embed-3-"
|
||||
S3_URI_PREFIX: Final = "s3://"
|
||||
TIMED_MEDIA_OPTION_FIELDS: Final = MappingProxyType(
|
||||
{
|
||||
"startSec": True,
|
||||
"endSec": True,
|
||||
"segmentation": True,
|
||||
"embeddingOption": True,
|
||||
"embeddingType": True,
|
||||
"embeddingScope": True,
|
||||
}
|
||||
)
|
||||
TIMED_MEDIA_OPTIONS: Final = TypeAdapter(TwelveLabsMarengo3TimedMediaOptions)
|
||||
TIMED_INPUT_TYPES: Final = frozenset({"video", "audio"})
|
||||
MARENGO_2_7_ONLY_PARAMS: Final = ("textTruncate", "lengthSec", "useFixedLengthSec", "minClipSec")
|
||||
MARENGO_2_7_ONLY_FIELDS: Final = MappingProxyType({name: True for name in MARENGO_2_7_ONLY_PARAMS})
|
||||
|
||||
|
||||
def is_marengo_3_model(model: str | None) -> bool:
|
||||
return MARENGO_3_MODEL_MARKER in (model or "")
|
||||
|
||||
|
||||
class Marengo3Params(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore", frozen=True)
|
||||
|
||||
inputType: TWELVELABS_MARENGO_3_INPUT_TYPES | None = None
|
||||
input_type: TWELVELABS_MARENGO_3_INPUT_TYPES | None = None
|
||||
media_source: str | None = None
|
||||
media_sources: Mapping[str, str] | None = None
|
||||
bucketOwner: str | None = None
|
||||
startSec: float | None = None
|
||||
endSec: float | None = None
|
||||
segmentation: TwelveLabsMarengo3Segmentation | None = None
|
||||
embeddingOption: tuple[TWELVELABS_MARENGO_3_EMBEDDING_OPTIONS, ...] | None = None
|
||||
embeddingType: tuple[TWELVELABS_MARENGO_3_EMBEDDING_TYPES, ...] | None = None
|
||||
embeddingScope: tuple[TWELVELABS_MARENGO_3_EMBEDDING_SCOPES, ...] | None = None
|
||||
inferenceId: str | None = None
|
||||
textTruncate: object = None
|
||||
lengthSec: object = None
|
||||
useFixedLengthSec: object = None
|
||||
minClipSec: object = None
|
||||
|
||||
@property
|
||||
def resolved_input_type(self) -> TWELVELABS_MARENGO_3_INPUT_TYPES:
|
||||
return self.inputType or self.input_type or "text"
|
||||
|
||||
def timed_media_options(self) -> TwelveLabsMarengo3TimedMediaOptions:
|
||||
return TIMED_MEDIA_OPTIONS.validate_python(self.given_timed_media_options())
|
||||
|
||||
def given_timed_media_options(self) -> dict[str, object]:
|
||||
return self.model_dump(include=TIMED_MEDIA_OPTION_FIELDS, exclude_none=True)
|
||||
|
||||
def given_2_7_only_params(self) -> dict[str, object]:
|
||||
return self.model_dump(include=MARENGO_2_7_ONLY_FIELDS, exclude_none=True)
|
||||
|
||||
|
||||
def _require_bucket_owner(bucket_owner: str | None) -> str:
|
||||
if bucket_owner is None:
|
||||
raise BedrockError(
|
||||
status_code=400,
|
||||
message="s3:// media requires the 'bucketOwner' parameter, the account id that owns the bucket",
|
||||
)
|
||||
return bucket_owner
|
||||
|
||||
|
||||
def _media_source(media: str, bucket_owner: str | None) -> TwelveLabsMediaSource:
|
||||
if not media.startswith(S3_URI_PREFIX):
|
||||
inline: Final[TwelveLabsMediaSource] = {"base64String": get_base64_str(media)}
|
||||
return inline
|
||||
s3_location: Final[TwelveLabsS3Location] = {"uri": media, "bucketOwner": _require_bucket_owner(bucket_owner)}
|
||||
remote: Final[TwelveLabsMediaSource] = {"s3Location": s3_location}
|
||||
return remote
|
||||
|
||||
|
||||
def _named_media_source(name: str, media: str, bucket_owner: str | None) -> TwelveLabsMarengo3NamedMediaSource:
|
||||
named: Final[TwelveLabsMarengo3NamedMediaSource] = {
|
||||
"name": name,
|
||||
"mediaType": "image",
|
||||
**_media_source(media, bucket_owner),
|
||||
}
|
||||
return named
|
||||
|
||||
|
||||
def _timed_media_input(media: str, params: Marengo3Params) -> TwelveLabsMarengo3TimedMediaInput:
|
||||
timed: Final[TwelveLabsMarengo3TimedMediaInput] = {
|
||||
"mediaSource": _media_source(media, params.bucketOwner),
|
||||
**params.timed_media_options(),
|
||||
}
|
||||
return timed
|
||||
|
||||
|
||||
def _describe(error: ValidationError) -> str:
|
||||
return "; ".join(
|
||||
f"{'.'.join(str(part) for part in problem['loc'])}: {problem['msg']}" for problem in error.errors()
|
||||
)
|
||||
|
||||
|
||||
def _validated_params(inference_params: Mapping[str, object]) -> Marengo3Params:
|
||||
try:
|
||||
return Marengo3Params.model_validate(inference_params)
|
||||
except ValidationError as error:
|
||||
raise BedrockError(status_code=400, message=f"Invalid Marengo 3.0 parameters: {_describe(error)}") from error
|
||||
|
||||
|
||||
def _reject_unless_dropped(given: Mapping[str, object], drop_params: bool, reason: str) -> None:
|
||||
if not given or drop_params:
|
||||
return
|
||||
raise BedrockError(status_code=400, message=f"{reason} {', '.join(given)}; set drop_params to drop them")
|
||||
|
||||
|
||||
def _require(value: str | None, input_type: str, param_name: str) -> str:
|
||||
if value is None:
|
||||
raise BedrockError(status_code=400, message=f"Input type '{input_type}' requires the '{param_name}' parameter")
|
||||
return value
|
||||
|
||||
|
||||
def _require_media_sources(value: Mapping[str, str] | None) -> Mapping[str, str]:
|
||||
if not value:
|
||||
raise BedrockError(
|
||||
status_code=400,
|
||||
message="Input type 'multi_input' requires a non-empty 'media_sources' mapping of name to media",
|
||||
)
|
||||
return value
|
||||
|
||||
|
||||
def _request_base(inference_id: str | None) -> TwelveLabsMarengo3RequestBase:
|
||||
if inference_id is None:
|
||||
anonymous: Final[TwelveLabsMarengo3RequestBase] = {}
|
||||
return anonymous
|
||||
identified: Final[TwelveLabsMarengo3RequestBase] = {"inferenceId": inference_id}
|
||||
return identified
|
||||
|
||||
|
||||
def build_marengo_3_request(
|
||||
input: str, inference_params: Mapping[str, object], drop_params: bool = False
|
||||
) -> TwelveLabsMarengo3EmbeddingRequest:
|
||||
params: Final = _validated_params(inference_params)
|
||||
base: Final = _request_base(params.inferenceId)
|
||||
input_type: Final = params.resolved_input_type
|
||||
_reject_unless_dropped(
|
||||
params.given_2_7_only_params(), drop_params, "Marengo 3.0 does not accept the Marengo 2.7 parameters"
|
||||
)
|
||||
if input_type not in TIMED_INPUT_TYPES:
|
||||
_reject_unless_dropped(
|
||||
params.given_timed_media_options(), drop_params, f"Input type '{input_type}' does not accept"
|
||||
)
|
||||
match input_type:
|
||||
case "text":
|
||||
text_request: Final[TwelveLabsMarengo3TextRequest] = {
|
||||
**base,
|
||||
"inputType": "text",
|
||||
"text": {"inputText": input},
|
||||
}
|
||||
return text_request
|
||||
case "image":
|
||||
image_request: Final[TwelveLabsMarengo3ImageRequest] = {
|
||||
**base,
|
||||
"inputType": "image",
|
||||
"image": {"mediaSource": _media_source(input, params.bucketOwner)},
|
||||
}
|
||||
return image_request
|
||||
case "video":
|
||||
video_request: Final[TwelveLabsMarengo3VideoRequest] = {
|
||||
**base,
|
||||
"inputType": "video",
|
||||
"video": _timed_media_input(input, params),
|
||||
}
|
||||
return video_request
|
||||
case "audio":
|
||||
audio_request: Final[TwelveLabsMarengo3AudioRequest] = {
|
||||
**base,
|
||||
"inputType": "audio",
|
||||
"audio": _timed_media_input(input, params),
|
||||
}
|
||||
return audio_request
|
||||
case "text_image":
|
||||
text_image_request: Final[TwelveLabsMarengo3TextImageRequest] = {
|
||||
**base,
|
||||
"inputType": "text_image",
|
||||
"text_image": {
|
||||
"inputText": input,
|
||||
"mediaSource": _media_source(
|
||||
_require(params.media_source, input_type, "media_source"), params.bucketOwner
|
||||
),
|
||||
},
|
||||
}
|
||||
return text_image_request
|
||||
case "multi_input":
|
||||
media_sources: Final = tuple(
|
||||
_named_media_source(name, media, params.bucketOwner)
|
||||
for name, media in _require_media_sources(params.media_sources).items()
|
||||
)
|
||||
multi_input_request: Final[TwelveLabsMarengo3MultiInputRequest] = {
|
||||
**base,
|
||||
"inputType": "multi_input",
|
||||
"multi_input": {"inputText": input, "mediaSources": media_sources}
|
||||
if input
|
||||
else {"mediaSources": media_sources},
|
||||
}
|
||||
return multi_input_request
|
||||
case _:
|
||||
assert_never(input_type)
|
||||
|
|
@ -4,19 +4,120 @@ Transformation logic from OpenAI /v1/embeddings format to Bedrock TwelveLabs Mar
|
|||
Why separate file? Make it easy to see how transformation works
|
||||
|
||||
Docs - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-marengo.html
|
||||
Marengo 3.0 docs - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-marengo-3.html
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Final, cast
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter
|
||||
from typing_extensions import assert_never
|
||||
|
||||
import litellm
|
||||
from litellm.llms.bedrock.embed.twelvelabs_marengo_3_transformation import (
|
||||
MARENGO_2_7_ONLY_PARAMS,
|
||||
build_marengo_3_request,
|
||||
is_marengo_3_model,
|
||||
)
|
||||
from litellm.types.llms.bedrock import (
|
||||
TWELVELABS_EMBEDDING_INPUT_TYPES,
|
||||
TWELVELABS_MARENGO_3_INPUT_TYPES,
|
||||
TwelveLabsAsyncInvokeRequest,
|
||||
TwelveLabsMarengo3EmbeddingRequest,
|
||||
TwelveLabsMarengoEmbeddingRequest,
|
||||
TwelveLabsOutputDataConfig,
|
||||
TwelveLabsS3Location,
|
||||
TwelveLabsS3OutputDataConfig,
|
||||
)
|
||||
from litellm.types.utils import Embedding, EmbeddingResponse, Usage
|
||||
from litellm.types.utils import Embedding, EmbeddingResponse, PromptTokensDetailsWrapper, Usage
|
||||
|
||||
|
||||
class MarengoEmbeddingItem(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore", frozen=True)
|
||||
|
||||
embedding: tuple[float, ...] | None = None
|
||||
|
||||
|
||||
class MarengoInvokeResponse(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore", frozen=True)
|
||||
|
||||
data: tuple[MarengoEmbeddingItem, ...] = ()
|
||||
embedding: tuple[float, ...] | None = None
|
||||
embeddings: tuple[MarengoEmbeddingItem, ...] = ()
|
||||
|
||||
def vectors(self) -> tuple[tuple[float, ...], ...]:
|
||||
if self.data:
|
||||
return tuple(item.embedding for item in self.data if item.embedding is not None)
|
||||
if self.embedding is not None:
|
||||
return (self.embedding,)
|
||||
return tuple(item.embedding for item in self.embeddings if item.embedding is not None)
|
||||
|
||||
|
||||
class MarengoBilledMultiInput(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore", frozen=True)
|
||||
|
||||
inputText: str | None = None
|
||||
mediaSources: tuple[Mapping[str, object], ...] = ()
|
||||
|
||||
|
||||
class MarengoBilledRequest(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore", frozen=True)
|
||||
|
||||
inputType: TWELVELABS_MARENGO_3_INPUT_TYPES | None = None
|
||||
multi_input: MarengoBilledMultiInput | None = None
|
||||
|
||||
|
||||
INVOKE_RESPONSES: Final = TypeAdapter(tuple[MarengoInvokeResponse, ...])
|
||||
BILLED_REQUESTS: Final = TypeAdapter(tuple[MarengoBilledRequest, ...])
|
||||
|
||||
|
||||
def _billed_units(request: MarengoBilledRequest) -> tuple[int, int]:
|
||||
input_type: Final = request.inputType
|
||||
match input_type:
|
||||
case "text":
|
||||
return (1, 0)
|
||||
case "image":
|
||||
return (0, 1)
|
||||
case "text_image":
|
||||
return (1, 1)
|
||||
case "multi_input":
|
||||
multi_input: Final = request.multi_input or MarengoBilledMultiInput()
|
||||
return (1 if multi_input.inputText else 0, len(multi_input.mediaSources))
|
||||
case "video" | "audio" | None:
|
||||
return (0, 0)
|
||||
case _:
|
||||
assert_never(input_type)
|
||||
|
||||
|
||||
def _billed_usage(batch_data: list[dict] | None) -> Usage:
|
||||
units: Final = tuple(_billed_units(request) for request in BILLED_REQUESTS.validate_python(batch_data or ()))
|
||||
query_count: Final = sum(text_requests for text_requests, _ in units)
|
||||
image_count: Final = sum(images for _, images in units)
|
||||
details: Final = (
|
||||
PromptTokensDetailsWrapper(query_count=query_count or None, image_count=image_count or None)
|
||||
if query_count or image_count
|
||||
else None
|
||||
)
|
||||
return Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0, prompt_tokens_details=details)
|
||||
|
||||
|
||||
MARENGO_SHARED_PARAMS: Final = (
|
||||
"encoding_format",
|
||||
"embeddingOption",
|
||||
"startSec",
|
||||
"input_type",
|
||||
"endSec",
|
||||
"segmentation",
|
||||
"embeddingType",
|
||||
"embeddingScope",
|
||||
"inferenceId",
|
||||
"media_source",
|
||||
"media_sources",
|
||||
)
|
||||
|
||||
|
||||
def drop_params_enabled(litellm_params: Mapping[str, object]) -> bool:
|
||||
return litellm.drop_params is True or litellm_params.get("drop_params") is True
|
||||
|
||||
|
||||
class TwelveLabsMarengoEmbeddingConfig:
|
||||
|
|
@ -26,28 +127,24 @@ class TwelveLabsMarengoEmbeddingConfig:
|
|||
Supports text, image, video, and audio inputs.
|
||||
- InvokeModel: text and image inputs
|
||||
- StartAsyncInvoke: video, audio, image, and text inputs
|
||||
|
||||
Marengo 3.0 (model ids containing "marengo-embed-3") nests the input under a key named after inputType and
|
||||
adds the text_image and multi_input input types; that payload is built by build_marengo_3_request.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
def __init__(self, model: str | None = None) -> None:
|
||||
self.is_marengo_3: Final = is_marengo_3_model(model)
|
||||
|
||||
def get_supported_openai_params(self) -> list[str]:
|
||||
return [
|
||||
"encoding_format",
|
||||
"textTruncate",
|
||||
"embeddingOption",
|
||||
"startSec",
|
||||
"lengthSec",
|
||||
"useFixedLengthSec",
|
||||
"minClipSec",
|
||||
"input_type",
|
||||
]
|
||||
if self.is_marengo_3:
|
||||
return list(MARENGO_SHARED_PARAMS)
|
||||
return [*MARENGO_SHARED_PARAMS, *MARENGO_2_7_ONLY_PARAMS]
|
||||
|
||||
def map_openai_params(self, non_default_params: dict, optional_params: dict) -> dict:
|
||||
for k, v in non_default_params.items():
|
||||
if k == "encoding_format":
|
||||
# TwelveLabs doesn't have encoding_format, but we can map it to embeddingOption
|
||||
if v == "float":
|
||||
if v == "float" and not self.is_marengo_3:
|
||||
optional_params["embeddingOption"] = ["visual-text", "visual-image"]
|
||||
elif k == "textTruncate":
|
||||
optional_params["textTruncate"] = v
|
||||
|
|
@ -56,7 +153,19 @@ class TwelveLabsMarengoEmbeddingConfig:
|
|||
elif k == "input_type":
|
||||
# Map input_type to inputType for Bedrock
|
||||
optional_params["inputType"] = v
|
||||
elif k in ["startSec", "lengthSec", "useFixedLengthSec", "minClipSec"]:
|
||||
elif k in (
|
||||
"startSec",
|
||||
"lengthSec",
|
||||
"useFixedLengthSec",
|
||||
"minClipSec",
|
||||
"endSec",
|
||||
"segmentation",
|
||||
"embeddingType",
|
||||
"embeddingScope",
|
||||
"inferenceId",
|
||||
"media_source",
|
||||
"media_sources",
|
||||
):
|
||||
optional_params[k] = v
|
||||
return optional_params
|
||||
|
||||
|
|
@ -77,7 +186,8 @@ class TwelveLabsMarengoEmbeddingConfig:
|
|||
async_invoke_route: bool = False,
|
||||
model_id: str | None = None,
|
||||
output_s3_uri: str | None = None,
|
||||
) -> TwelveLabsMarengoEmbeddingRequest | TwelveLabsAsyncInvokeRequest:
|
||||
drop_params: bool = False,
|
||||
) -> TwelveLabsMarengoEmbeddingRequest | TwelveLabsMarengo3EmbeddingRequest | TwelveLabsAsyncInvokeRequest:
|
||||
"""
|
||||
Transform OpenAI-style input to TwelveLabs Marengo format/async-invoke format.
|
||||
|
||||
|
|
@ -87,20 +197,29 @@ class TwelveLabsMarengoEmbeddingConfig:
|
|||
- Video inputs (async-invoke only)
|
||||
- Audio inputs (async-invoke only)
|
||||
- S3 URLs for all media types (async-invoke only)
|
||||
- Marengo 3.0 only: text_image and multi_input inputs (nested payload)
|
||||
"""
|
||||
# Get input_type or default to "text"
|
||||
input_type: Final = cast(
|
||||
TWELVELABS_EMBEDDING_INPUT_TYPES,
|
||||
inference_params.get("inputType") or inference_params.get("input_type") or "text",
|
||||
)
|
||||
|
||||
# Validate that async-invoke is used for video/audio
|
||||
if input_type in ["video", "audio"] and not async_invoke_route:
|
||||
raise ValueError(
|
||||
f"Input type '{input_type}' requires async_invoke route. "
|
||||
f"Use model format: 'bedrock/async_invoke/model_id'"
|
||||
)
|
||||
|
||||
if self.is_marengo_3:
|
||||
marengo_3_request: Final = build_marengo_3_request(
|
||||
input=input, inference_params=inference_params, drop_params=drop_params
|
||||
)
|
||||
if async_invoke_route and model_id:
|
||||
return self._wrap_async_invoke_request(
|
||||
model_input=marengo_3_request, model_id=model_id, output_s3_uri=output_s3_uri
|
||||
)
|
||||
return marengo_3_request
|
||||
|
||||
transformed_request: Final[TwelveLabsMarengoEmbeddingRequest] = {"inputType": input_type}
|
||||
|
||||
if input_type == "text":
|
||||
|
|
@ -154,7 +273,7 @@ class TwelveLabsMarengoEmbeddingConfig:
|
|||
|
||||
def _wrap_async_invoke_request(
|
||||
self,
|
||||
model_input: TwelveLabsMarengoEmbeddingRequest,
|
||||
model_input: TwelveLabsMarengoEmbeddingRequest | TwelveLabsMarengo3EmbeddingRequest,
|
||||
model_id: str,
|
||||
output_s3_uri: str | None = None,
|
||||
) -> TwelveLabsAsyncInvokeRequest:
|
||||
|
|
@ -188,62 +307,16 @@ class TwelveLabsMarengoEmbeddingConfig:
|
|||
),
|
||||
)
|
||||
|
||||
def _transform_response(self, response_list: list[dict], model: str) -> EmbeddingResponse:
|
||||
"""
|
||||
Transform TwelveLabs response to OpenAI format.
|
||||
Handles the actual TwelveLabs response format: {"data": [{"embedding": [...]}]}
|
||||
"""
|
||||
embeddings: Final[list[Embedding]] = []
|
||||
total_tokens = 0
|
||||
|
||||
for response in response_list:
|
||||
# TwelveLabs response format has a "data" field containing the embeddings
|
||||
if "data" in response and isinstance(response["data"], list):
|
||||
for item in response["data"]:
|
||||
if "embedding" in item:
|
||||
# Single embedding response
|
||||
embedding = Embedding(
|
||||
embedding=item["embedding"],
|
||||
index=len(embeddings),
|
||||
object="embedding",
|
||||
)
|
||||
embeddings.append(embedding)
|
||||
|
||||
# Estimate token count (rough approximation)
|
||||
if "inputTextTokenCount" in item:
|
||||
total_tokens += item["inputTextTokenCount"]
|
||||
else:
|
||||
# Rough estimate: 1 token per 4 characters for text, or use embedding size
|
||||
total_tokens += len(item["embedding"]) // 4
|
||||
elif "embedding" in response:
|
||||
# Direct embedding response (fallback for other formats)
|
||||
embedding = Embedding(
|
||||
embedding=response["embedding"],
|
||||
index=len(embeddings),
|
||||
object="embedding",
|
||||
)
|
||||
embeddings.append(embedding)
|
||||
|
||||
# Estimate token count (rough approximation)
|
||||
if "inputTextTokenCount" in response:
|
||||
total_tokens += response["inputTextTokenCount"]
|
||||
else:
|
||||
# Rough estimate: 1 token per 4 characters for text
|
||||
total_tokens += len(response.get("inputText", "")) // 4
|
||||
elif "embeddings" in response:
|
||||
# Multiple embeddings response (from video/audio)
|
||||
for i, emb in enumerate(response["embeddings"]):
|
||||
embedding = Embedding(
|
||||
embedding=emb["embedding"],
|
||||
index=len(embeddings),
|
||||
object="embedding",
|
||||
)
|
||||
embeddings.append(embedding)
|
||||
total_tokens += len(emb["embedding"]) // 4 # Rough estimate
|
||||
|
||||
usage: Final = Usage(prompt_tokens=total_tokens, total_tokens=total_tokens)
|
||||
|
||||
return EmbeddingResponse(data=embeddings, model=model, usage=usage)
|
||||
def _transform_response(
|
||||
self, response_list: list[dict], model: str, batch_data: list[dict] | None = None
|
||||
) -> EmbeddingResponse:
|
||||
vectors: Final = tuple(
|
||||
vector for response in INVOKE_RESPONSES.validate_python(response_list) for vector in response.vectors()
|
||||
)
|
||||
embeddings: Final = [
|
||||
Embedding(embedding=list(vector), index=index, object="embedding") for index, vector in enumerate(vectors)
|
||||
]
|
||||
return EmbeddingResponse(data=embeddings, model=model, usage=_billed_usage(batch_data))
|
||||
|
||||
def _transform_async_invoke_response(self, response: dict, model: str) -> EmbeddingResponse:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -9826,7 +9826,7 @@ class BaseLLMHTTPHandler:
|
|||
vector_store_search_optional_params=vector_store_search_optional_params,
|
||||
api_base=api_base,
|
||||
litellm_logging_obj=logging_obj,
|
||||
litellm_params=dict(litellm_params),
|
||||
litellm_params=MappingProxyType(dict(litellm_params, timeout=timeout)),
|
||||
extra_body=extra_body,
|
||||
embedding_executor=embedding_executor,
|
||||
)
|
||||
|
|
@ -9871,6 +9871,12 @@ class BaseLLMHTTPHandler:
|
|||
data=request_data,
|
||||
timeout=timeout,
|
||||
)
|
||||
except httpx.TimeoutException:
|
||||
raise vector_store_provider_config.get_error_class(
|
||||
error_message="Vector store search exceeded the caller timeout.",
|
||||
status_code=408,
|
||||
headers=httpx.Headers(),
|
||||
) from None
|
||||
except Exception as e:
|
||||
raise self._handle_error(e=e, provider_config=vector_store_provider_config)
|
||||
|
||||
|
|
@ -9955,7 +9961,7 @@ class BaseLLMHTTPHandler:
|
|||
vector_store_search_optional_params=vector_store_search_optional_params,
|
||||
api_base=api_base,
|
||||
litellm_logging_obj=logging_obj,
|
||||
litellm_params=dict(litellm_params),
|
||||
litellm_params=MappingProxyType(dict(litellm_params, timeout=timeout)),
|
||||
extra_body=extra_body,
|
||||
embedding_executor=embedding_executor,
|
||||
)
|
||||
|
|
@ -10000,7 +10006,14 @@ class BaseLLMHTTPHandler:
|
|||
url=url,
|
||||
headers=headers,
|
||||
data=request_data,
|
||||
timeout=timeout,
|
||||
)
|
||||
except httpx.TimeoutException:
|
||||
raise vector_store_provider_config.get_error_class(
|
||||
error_message="Vector store search exceeded the caller timeout.",
|
||||
status_code=408,
|
||||
headers=httpx.Headers(),
|
||||
) from None
|
||||
except Exception as e:
|
||||
raise self._handle_error(e=e, provider_config=vector_store_provider_config)
|
||||
|
||||
|
|
@ -10030,6 +10043,8 @@ class BaseLLMHTTPHandler:
|
|||
else:
|
||||
async_httpx_client = client
|
||||
|
||||
vector_store_provider_config.validate_create_vector_store()
|
||||
|
||||
headers: Final = vector_store_provider_config.validate_environment(
|
||||
headers=extra_headers or {}, litellm_params=litellm_params
|
||||
)
|
||||
|
|
@ -10100,6 +10115,8 @@ class BaseLLMHTTPHandler:
|
|||
else:
|
||||
sync_httpx_client = client
|
||||
|
||||
vector_store_provider_config.validate_create_vector_store()
|
||||
|
||||
headers: Final = vector_store_provider_config.validate_environment(
|
||||
headers=extra_headers or {}, litellm_params=litellm_params
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,303 +0,0 @@
|
|||
"""Shared helpers for the MongoDB integrations. pymongo lives in the optional ``mongodb`` extra,
|
||||
so every import of it is deferred to call time."""
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
import weakref
|
||||
from asyncio import AbstractEventLoop
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, TypeAlias, TypeVar
|
||||
|
||||
from litellm.exceptions import BadRequestError, ServiceUnavailableError, Timeout
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pymongo import AsyncMongoClient, MongoClient
|
||||
|
||||
PYMONGO_INSTALL_HINT: Final = (
|
||||
"The MongoDB vector store requires the 'pymongo' package. "
|
||||
"Run 'pip install litellm[mongodb]' (or 'pip install pymongo') to install it."
|
||||
)
|
||||
|
||||
MONGODB_PROVIDER: Final = "mongodb"
|
||||
|
||||
|
||||
def config_error(message: str) -> BadRequestError:
|
||||
"""400 rather than the 500 a bare ValueError becomes once litellm.exception_type wraps it."""
|
||||
return BadRequestError(message=message, model=None, llm_provider=MONGODB_PROVIDER)
|
||||
|
||||
|
||||
def timeout_error(message: str) -> Timeout:
|
||||
return Timeout(message=message, model=None, llm_provider=MONGODB_PROVIDER)
|
||||
|
||||
|
||||
def unavailable_error(message: str) -> ServiceUnavailableError:
|
||||
"""litellm only retries 408, 409, 429 and 5xx, so a 400 here would make a failover permanent."""
|
||||
return ServiceUnavailableError(message=message, model=None, llm_provider=MONGODB_PROVIDER)
|
||||
|
||||
|
||||
DEFAULT_CONNECT_TIMEOUT_MS: Final = 10_000
|
||||
DEFAULT_SOCKET_TIMEOUT_MS: Final = 30_000
|
||||
DEFAULT_SERVER_SELECTION_TIMEOUT_MS: Final = 10_000
|
||||
|
||||
_MAX_CACHED_CLIENTS: Final = 32
|
||||
|
||||
_APP_NAME: Final = "litellm"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class MongoClientKey:
|
||||
connection_string: str
|
||||
connect_timeout_ms: int
|
||||
socket_timeout_ms: int
|
||||
server_selection_timeout_ms: int
|
||||
|
||||
|
||||
SyncClientFactory: TypeAlias = Callable[..., "MongoClient"]
|
||||
AsyncClientFactory: TypeAlias = Callable[..., "AsyncMongoClient"]
|
||||
|
||||
_K = TypeVar("_K")
|
||||
_V = TypeVar("_V")
|
||||
|
||||
_AsyncClientCacheKey: TypeAlias = tuple[MongoClientKey, int]
|
||||
# CPython recycles id() aggressively, so the id alone would hand a new loop a closed loop's client
|
||||
_AsyncClientEntry: TypeAlias = tuple["weakref.ref[AbstractEventLoop]", "AsyncMongoClient"]
|
||||
|
||||
_SyncClientCache: TypeAlias = "OrderedDict[MongoClientKey, MongoClient]"
|
||||
_AsyncClientCache: TypeAlias = "OrderedDict[_AsyncClientCacheKey, _AsyncClientEntry]"
|
||||
|
||||
_sync_clients: Final[_SyncClientCache] = OrderedDict() # mutable-ok: process-level client cache
|
||||
_async_clients: Final[_AsyncClientCache] = OrderedDict() # mutable-ok: same cache, per loop
|
||||
# async searches reach the sync client through executor threads, so both caches are shared state
|
||||
_cache_lock: Final = threading.Lock()
|
||||
|
||||
|
||||
def _store_bounded(cache: "OrderedDict[_K, _V]", cache_key: "_K", value: "_V") -> None:
|
||||
"""Eviction only drops this cache's reference; an in-flight search keeps its client alive."""
|
||||
with _cache_lock:
|
||||
cache[cache_key] = value # mutable-ok: an LRU cache is mutable state by definition
|
||||
cache.move_to_end(cache_key)
|
||||
while len(cache) > _MAX_CACHED_CLIENTS:
|
||||
cache.popitem(last=False)
|
||||
|
||||
|
||||
def _mark_used(cache: "OrderedDict[_K, _V]", cache_key: "_K") -> None:
|
||||
with _cache_lock:
|
||||
if cache_key in cache:
|
||||
cache.move_to_end(cache_key)
|
||||
|
||||
|
||||
def import_sync_mongo_client() -> "type[MongoClient]":
|
||||
try:
|
||||
from pymongo import MongoClient as SyncMongoClient
|
||||
except ImportError as e:
|
||||
raise config_error(PYMONGO_INSTALL_HINT) from e
|
||||
return SyncMongoClient
|
||||
|
||||
|
||||
def import_async_mongo_client() -> "type[AsyncMongoClient]":
|
||||
try:
|
||||
from pymongo import AsyncMongoClient as AsyncMongoClientClass
|
||||
except ImportError as e:
|
||||
raise config_error(PYMONGO_INSTALL_HINT) from e
|
||||
return AsyncMongoClientClass
|
||||
|
||||
|
||||
def _client_kwargs(key: MongoClientKey) -> Mapping[str, object]:
|
||||
return MappingProxyType(
|
||||
{
|
||||
"connectTimeoutMS": key.connect_timeout_ms,
|
||||
"socketTimeoutMS": key.socket_timeout_ms,
|
||||
"serverSelectionTimeoutMS": key.server_selection_timeout_ms,
|
||||
"appname": _APP_NAME,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def get_sync_client(key: MongoClientKey, client_class: SyncClientFactory | None = None) -> "MongoClient":
|
||||
cached: Final = _sync_clients.get(key)
|
||||
if cached is not None:
|
||||
_mark_used(_sync_clients, key)
|
||||
return cached
|
||||
build: Final = client_class if client_class is not None else import_sync_mongo_client()
|
||||
client: Final = build(key.connection_string, **_client_kwargs(key))
|
||||
_store_bounded(_sync_clients, key, client)
|
||||
return client
|
||||
|
||||
|
||||
def _purge_dead_loops() -> None:
|
||||
"""A cached client holds its loop alive, so a closed loop's entry would pin that client and its
|
||||
sockets for the life of the process."""
|
||||
with _cache_lock:
|
||||
for stale in tuple(
|
||||
cache_key
|
||||
for cache_key, (loop_ref, _) in _async_clients.items()
|
||||
if (cached_loop := loop_ref()) is None or cached_loop.is_closed()
|
||||
):
|
||||
del _async_clients[stale]
|
||||
|
||||
|
||||
def get_async_client(key: MongoClientKey, client_class: AsyncClientFactory | None = None) -> "AsyncMongoClient":
|
||||
"""Async clients bind to the loop that created them, so the cache is keyed per loop."""
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
loop_key: Final = (key, id(loop))
|
||||
cached: Final = _async_clients.get(loop_key)
|
||||
if cached is not None and cached[0]() is loop:
|
||||
_mark_used(_async_clients, loop_key)
|
||||
return cached[1]
|
||||
_purge_dead_loops()
|
||||
build: Final = client_class if client_class is not None else import_async_mongo_client()
|
||||
client: Final = build(key.connection_string, **_client_kwargs(key))
|
||||
_store_bounded(_async_clients, loop_key, (weakref.ref(loop), client))
|
||||
return client
|
||||
|
||||
|
||||
def reset_client_cache() -> None:
|
||||
with _cache_lock:
|
||||
_sync_clients.clear()
|
||||
_async_clients.clear()
|
||||
|
||||
|
||||
_AUTHENTICATION_FAILED_CODE: Final = 18
|
||||
_UNAUTHORIZED_CODE: Final = 13
|
||||
# Atlas reports a rejected user as code 8000 "AtlasError" where a self-managed mongod reports 18
|
||||
_AUTHENTICATION_MESSAGE_MARKERS: Final = ("bad auth", "authentication failed", "not authorized")
|
||||
_RESOLUTION_TIMEOUT_MARKERS: Final = ("resolution lifetime expired", "dns operation timed out")
|
||||
_UNKNOWN_HOSTNAME_MARKERS: Final = ("dns query name does not exist", "name or service not known")
|
||||
_CREDENTIAL_ESCAPING_MARKERS: Final = ("must be escaped according to rfc 3986", "bad database name")
|
||||
|
||||
|
||||
def _index_hint(index_name: str, database: str, collection: str) -> str:
|
||||
return (
|
||||
f"No queryable MongoDB Vector Search index named '{index_name}' was found on "
|
||||
f"'{database}.{collection}'. Confirm the index exists on that exact collection, that its "
|
||||
"status is READY rather than still building, and that the vector store id matches the index name."
|
||||
)
|
||||
|
||||
|
||||
def missing_index_error(index_name: str, database: str, collection: str) -> BadRequestError:
|
||||
"""$vectorSearch against a missing index, database or collection returns zero documents rather
|
||||
than failing, so an empty result set is checked against the catalogue and reported as this."""
|
||||
return config_error(
|
||||
f"{_index_hint(index_name, database, collection)} A vector search against a database, "
|
||||
"collection or index that does not exist returns no results rather than an error, so this "
|
||||
"was reported as an empty result set by MongoDB."
|
||||
)
|
||||
|
||||
|
||||
def index_not_ready_error(index_name: str, database: str, collection: str, status: str) -> BadRequestError:
|
||||
return config_error(
|
||||
f"The MongoDB Vector Search index '{index_name}' on '{database}.{collection}' is not queryable "
|
||||
f"yet; its status is {status}. Searches against it return no results until the build finishes."
|
||||
)
|
||||
|
||||
|
||||
def translate_mongo_error(error: Exception, index_name: str, database: str, collection: str) -> Exception:
|
||||
"""Returns the exception to raise, so callers keep the driver error as ``__cause__``."""
|
||||
try:
|
||||
from pymongo.errors import (
|
||||
ConfigurationError,
|
||||
ConnectionFailure,
|
||||
ExecutionTimeout,
|
||||
InvalidOperation,
|
||||
NetworkTimeout,
|
||||
OperationFailure,
|
||||
ServerSelectionTimeoutError,
|
||||
)
|
||||
except ImportError:
|
||||
return error
|
||||
|
||||
if isinstance(error, ServerSelectionTimeoutError):
|
||||
return timeout_error(
|
||||
"Could not reach the MongoDB deployment before the timeout. On Atlas this is usually the "
|
||||
"project's IP access list not containing this host, or a paused cluster. On a self-managed "
|
||||
"deployment it is usually the host or port in the URI, or a firewall between this process "
|
||||
f"and mongod. Either way it can also be an unresolvable hostname. Driver detail: {error}"
|
||||
)
|
||||
# ExecutionTimeout subclasses OperationFailure, so it has to be matched before it
|
||||
if isinstance(error, (NetworkTimeout, ExecutionTimeout)):
|
||||
return timeout_error(
|
||||
f"The MongoDB vector search against '{database}.{collection}' timed out before returning. "
|
||||
f"Driver detail: {error}"
|
||||
)
|
||||
# ServerSelectionTimeoutError and NetworkTimeout also subclass ConnectionFailure, so this only
|
||||
# sees what those branches left
|
||||
if isinstance(error, ConnectionFailure):
|
||||
return unavailable_error(
|
||||
f"The connection to '{database}.{collection}' was dropped or refused. That is usually a "
|
||||
"replica set failover or a restarted node, so the search is worth retrying. If it keeps "
|
||||
"happening: on Atlas the usual cause is a connection string with no username and password, "
|
||||
"or a TLS failure, so confirm the URI is the one Atlas shows under Connect, Drivers; on a "
|
||||
"self-managed deployment, check that mongod is listening on the host and port in the URI. "
|
||||
f"Driver detail: {error}"
|
||||
)
|
||||
if isinstance(error, OperationFailure):
|
||||
code: Final = error.code
|
||||
detail: Final = str(error).lower()
|
||||
if code in (_AUTHENTICATION_FAILED_CODE, _UNAUTHORIZED_CODE) or any(
|
||||
marker in detail for marker in _AUTHENTICATION_MESSAGE_MARKERS
|
||||
):
|
||||
return config_error(
|
||||
"MongoDB rejected the credentials in mongodb_connection_string, or the database user "
|
||||
f"lacks read access to '{database}.{collection}'. Driver detail: {error.details}"
|
||||
)
|
||||
if "dimension" in detail:
|
||||
return config_error(
|
||||
"The query embedding does not match the vector dimensions the index was built for. "
|
||||
"litellm_embedding_model must be the same model that produced the stored vectors. "
|
||||
f"Driver detail: {error}"
|
||||
)
|
||||
if "is not indexed as vector" in detail:
|
||||
return config_error(
|
||||
"mongodb_embedding_field names a field the MongoDB Vector Search index does not cover. "
|
||||
f"It must match the 'path' the index '{index_name}' was created on. Driver detail: {error}"
|
||||
)
|
||||
if "index" in detail and ("not found" in detail or "does not exist" in detail or "unknown" in detail):
|
||||
return config_error(f"{_index_hint(index_name, database, collection)} Driver detail: {error}")
|
||||
return config_error(
|
||||
f"MongoDB rejected the vector search against '{database}.{collection}' using index "
|
||||
f"'{index_name}'. Driver detail: {error}"
|
||||
)
|
||||
if isinstance(error, ConfigurationError):
|
||||
configuration_detail: Final = str(error).lower()
|
||||
if any(marker in configuration_detail for marker in _RESOLUTION_TIMEOUT_MARKERS):
|
||||
return timeout_error(
|
||||
"The DNS lookup for the cluster in mongodb_connection_string did not finish in time. "
|
||||
"A mongodb+srv:// URI needs an SRV lookup before any connection is attempted, so this "
|
||||
f"is DNS or the configured timeout, not MongoDB. Driver detail: {error}"
|
||||
)
|
||||
if any(marker in configuration_detail for marker in _UNKNOWN_HOSTNAME_MARKERS):
|
||||
return config_error(
|
||||
"The hostname in mongodb_connection_string does not exist in DNS. On Atlas, check the "
|
||||
"cluster name against the URI shown under Connect, Drivers. On a self-managed deployment, "
|
||||
f"check that the hostname resolves from this process. Driver detail: {error}"
|
||||
)
|
||||
if any(marker in configuration_detail for marker in _CREDENTIAL_ESCAPING_MARKERS):
|
||||
return config_error(
|
||||
"mongodb_connection_string could not be parsed. A username or password containing "
|
||||
"'@', '/', ':' or '%' has to be percent-encoded per RFC 3986, so 'p@ss/word' becomes "
|
||||
"'p%40ss%2Fword'. If the credentials are already encoded, check the database name in "
|
||||
f"the URI path instead. Driver detail: {error}"
|
||||
)
|
||||
return config_error(
|
||||
f"mongodb_connection_string is not a usable MongoDB connection string. Driver detail: {error}"
|
||||
)
|
||||
if isinstance(error, InvalidOperation):
|
||||
return config_error(f"The MongoDB client was already closed or is unusable. Driver detail: {error}")
|
||||
# An unreadable tlsCAFile or tlsCertificateKeyFile raises OSError, not a PyMongoError
|
||||
if isinstance(error, OSError) and error.filename:
|
||||
return config_error(
|
||||
f"'{error.filename}', named by a TLS option in mongodb_connection_string, could not be read. "
|
||||
"Check that tlsCAFile and tlsCertificateKeyFile point at files this process can open; inside "
|
||||
f"a container that is the path in the container, not on the host. Driver detail: {error}"
|
||||
)
|
||||
# pymongo raises a plain ValueError, not a PyMongoError, for an unusable port
|
||||
if isinstance(error, ValueError):
|
||||
return config_error(
|
||||
"The host and port in mongodb_connection_string could not be parsed. If the port is a "
|
||||
"number between 0 and 65535, the cause is usually an unescaped ':' in the password, which "
|
||||
f"has to be percent-encoded per RFC 3986 as '%3A'. Driver detail: {error}"
|
||||
)
|
||||
return error
|
||||
|
|
@ -1,37 +1,29 @@
|
|||
"""MongoDB Vector Search has no HTTP query API, so this is a direct provider that runs the
|
||||
``$vectorSearch`` aggregation through pymongo. ``vector_store_id`` is the search index name."""
|
||||
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from collections.abc import Mapping, Sequence
|
||||
from ipaddress import ip_address
|
||||
from math import isfinite
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, NoReturn
|
||||
from typing import TYPE_CHECKING, Final, Literal, NoReturn
|
||||
from urllib.parse import quote, urlsplit
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
|
||||
from litellm.exceptions import AuthenticationError, BadRequestError, ServiceUnavailableError, Timeout
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.base_llm.vector_store.transformation import (
|
||||
BaseDirectVectorStoreConfig,
|
||||
BaseQueryEmbeddingVectorStoreConfig,
|
||||
LiteLLMVectorStoreEmbeddingExecutor,
|
||||
VectorStoreEmbeddingExecutor,
|
||||
)
|
||||
from litellm.llms.mongodb.common_utils import (
|
||||
DEFAULT_CONNECT_TIMEOUT_MS,
|
||||
DEFAULT_SERVER_SELECTION_TIMEOUT_MS,
|
||||
DEFAULT_SOCKET_TIMEOUT_MS,
|
||||
MongoClientKey,
|
||||
config_error,
|
||||
get_async_client,
|
||||
get_sync_client,
|
||||
index_not_ready_error,
|
||||
missing_index_error,
|
||||
translate_mongo_error,
|
||||
)
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import EmbeddingResponse
|
||||
from litellm.types.vector_stores import (
|
||||
BaseVectorStoreAuthCredentials,
|
||||
VectorStoreCreateOptionalRequestParams,
|
||||
VectorStoreResultContent,
|
||||
VectorStoreIndexEndpoints,
|
||||
VectorStoreSearchOptionalRequestParams,
|
||||
VectorStoreSearchResponse,
|
||||
VectorStoreSearchResult,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -39,26 +31,45 @@ if TYPE_CHECKING:
|
|||
|
||||
DEFAULT_EMBEDDING_FIELD_NAME: Final = "embedding"
|
||||
DEFAULT_TEXT_FIELD_NAME: Final = "text"
|
||||
SCORE_FIELD_NAME: Final = "score"
|
||||
|
||||
DEFAULT_MAX_NUM_RESULTS: Final = 10
|
||||
MIN_MAX_NUM_RESULTS: Final = 1
|
||||
MAX_MAX_NUM_RESULTS: Final = 50
|
||||
|
||||
NUM_CANDIDATES_MULTIPLIER: Final = 10
|
||||
MIN_NUM_CANDIDATES: Final = 100
|
||||
MAX_NUM_CANDIDATES: Final = 10_000
|
||||
|
||||
MAX_QUERY_CHARACTERS: Final = 32_000
|
||||
|
||||
_EMPTY_EMBEDDING_CONFIG: Final = MappingProxyType({})
|
||||
|
||||
_SEARCH_ONLY_MESSAGE: Final = (
|
||||
"MongoDB vector store is search-only. Create the collection and its MongoDB Vector Search "
|
||||
"index in MongoDB directly, then register it here by index name."
|
||||
)
|
||||
|
||||
|
||||
def config_error(message: str) -> BadRequestError:
|
||||
return BadRequestError(message=message, model=None, llm_provider="mongodb")
|
||||
|
||||
|
||||
class _Content(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, strict=True)
|
||||
type: Literal["text"]
|
||||
text: str
|
||||
|
||||
|
||||
class _Result(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, strict=True, allow_inf_nan=False)
|
||||
score: float | None
|
||||
content: Sequence[_Content]
|
||||
file_id: str | None
|
||||
filename: str | None
|
||||
|
||||
|
||||
class _SearchResponse(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, strict=True)
|
||||
object: Literal["vector_store.search_results.page"]
|
||||
search_query: str
|
||||
data: Sequence[_Result]
|
||||
|
||||
|
||||
class _MongoDBSearchParams(BaseModel):
|
||||
"""Typed view over the vector store's litellm_params; unrelated keys are ignored."""
|
||||
|
||||
|
|
@ -66,7 +77,6 @@ class _MongoDBSearchParams(BaseModel):
|
|||
|
||||
litellm_embedding_model: str | None = None
|
||||
litellm_embedding_config: Mapping[str, object] | None = None
|
||||
mongodb_connection_string: str | None = None
|
||||
mongodb_database: str | None = None
|
||||
mongodb_collection: str | None = None
|
||||
mongodb_text_field: str | None = None
|
||||
|
|
@ -91,21 +101,6 @@ class _MongoDBSearchParams(BaseModel):
|
|||
)
|
||||
return self.litellm_embedding_model
|
||||
|
||||
def require_connection_string(self) -> str:
|
||||
if not self.mongodb_connection_string:
|
||||
raise config_error(
|
||||
"mongodb_connection_string is required in litellm_params for the MongoDB vector store. "
|
||||
"Example: mongodb+srv://<user>:<password>@<cluster>.mongodb.net for Atlas, or "
|
||||
"mongodb://<user>:<password>@<host>:27017 for a self-managed deployment"
|
||||
)
|
||||
scheme: Final = self.mongodb_connection_string.split("://", 1)[0].lower()
|
||||
if scheme not in ("mongodb", "mongodb+srv"):
|
||||
raise config_error(
|
||||
"mongodb_connection_string must start with 'mongodb://' or 'mongodb+srv://', "
|
||||
f"got '{self.mongodb_connection_string.split('://', 1)[0]}://'"
|
||||
)
|
||||
return self.mongodb_connection_string
|
||||
|
||||
def require_database(self) -> str:
|
||||
if not self.mongodb_database:
|
||||
raise config_error(
|
||||
|
|
@ -127,30 +122,28 @@ _MONGODB_PARAM_PREFIX: Final = "mongodb_"
|
|||
_KNOWN_MONGODB_PARAMS: Final = frozenset(
|
||||
name for name in _MongoDBSearchParams.model_fields if name.startswith(_MONGODB_PARAM_PREFIX)
|
||||
)
|
||||
_RESPONSE_ADAPTER: Final = TypeAdapter(VectorStoreSearchResponse)
|
||||
|
||||
|
||||
class MongoDBVectorStoreConfig(BaseDirectVectorStoreConfig):
|
||||
def __init__(
|
||||
self,
|
||||
embedding_executor: VectorStoreEmbeddingExecutor | None = None,
|
||||
sync_client_factory: Callable[[MongoClientKey], object] | None = None,
|
||||
async_client_factory: Callable[[MongoClientKey], object] | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.embedding_executor: Final[VectorStoreEmbeddingExecutor] = (
|
||||
embedding_executor if embedding_executor is not None else LiteLLMVectorStoreEmbeddingExecutor()
|
||||
)
|
||||
self.sync_client_factory: Final[Callable[[MongoClientKey], object]] = (
|
||||
sync_client_factory if sync_client_factory is not None else get_sync_client
|
||||
)
|
||||
self.async_client_factory: Final[Callable[[MongoClientKey], object]] = (
|
||||
async_client_factory if async_client_factory is not None else get_async_client
|
||||
)
|
||||
class MongoDBVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig):
|
||||
def __init__(self, embedding_executor: VectorStoreEmbeddingExecutor | None = None) -> None:
|
||||
self.embedding_executor: Final = embedding_executor or LiteLLMVectorStoreEmbeddingExecutor()
|
||||
|
||||
def get_auth_credentials(self, litellm_params: Mapping[str, object]) -> BaseVectorStoreAuthCredentials:
|
||||
return BaseVectorStoreAuthCredentials()
|
||||
|
||||
def get_vector_store_endpoints_by_type(self) -> VectorStoreIndexEndpoints:
|
||||
return VectorStoreIndexEndpoints(read=[], write=[]) # mutable-ok: the TypedDict declares list fields
|
||||
|
||||
@staticmethod
|
||||
def _reject_unknown_params(litellm_params: Mapping[str, object]) -> None:
|
||||
"""Without this a mistyped mongodb_collection reads as 'mongodb_collection is required',
|
||||
naming a key the reader can see they have set."""
|
||||
if litellm_params.get("mongodb_connection_string") is not None:
|
||||
raise config_error(
|
||||
"MongoDB vector stores now use the BETA sidecar. Move mongodb_connection_string to "
|
||||
"MONGODB_CONNECTION_STRING in the sidecar, remove it from LiteLLM, and configure api_base and api_key."
|
||||
)
|
||||
unknown: Final = sorted(
|
||||
key for key in litellm_params if key.startswith(_MONGODB_PARAM_PREFIX) and key not in _KNOWN_MONGODB_PARAMS
|
||||
)
|
||||
|
|
@ -191,239 +184,203 @@ class MongoDBVectorStoreConfig(BaseDirectVectorStoreConfig):
|
|||
return configured
|
||||
return min(max(limit * NUM_CANDIDATES_MULTIPLIER, MIN_NUM_CANDIDATES), MAX_NUM_CANDIDATES)
|
||||
|
||||
@staticmethod
|
||||
def _timeout_ms(timeout: float | httpx.Timeout | None) -> tuple[int, int]:
|
||||
"""The connect and socket budgets pymongo is built with, in that order."""
|
||||
if isinstance(timeout, httpx.Timeout):
|
||||
return (
|
||||
int((timeout.connect or DEFAULT_CONNECT_TIMEOUT_MS / 1000) * 1000),
|
||||
int((timeout.read or DEFAULT_SOCKET_TIMEOUT_MS / 1000) * 1000),
|
||||
def validate_environment(
|
||||
self, headers: Mapping[str, object], litellm_params: GenericLiteLLMParams | None
|
||||
) -> dict[str, object]: # mutable-ok: the shared HTTP handler requires writable headers
|
||||
if litellm_params is None:
|
||||
raise config_error("Configure api_base and api_key for the MongoDB BETA sidecar.")
|
||||
self._reject_unknown_params(MappingProxyType(dict(litellm_params)))
|
||||
api_key: Final = litellm_params.api_key or get_secret_str("MONGODB_SIDECAR_API_KEY")
|
||||
if not api_key:
|
||||
raise config_error("MongoDB sidecar api_key is required. Set api_key or MONGODB_SIDECAR_API_KEY.")
|
||||
return {
|
||||
**headers,
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
} # mutable-ok: writable HTTP headers
|
||||
|
||||
def get_complete_url(self, api_base: str | None, litellm_params: Mapping[str, object]) -> str:
|
||||
if not api_base:
|
||||
raise config_error("MongoDB sidecar api_base is required, for example http://127.0.0.1:8080.")
|
||||
try:
|
||||
parsed: Final = urlsplit(api_base)
|
||||
valid: Final = parsed.scheme in ("http", "https") and bool(parsed.hostname) and parsed.port != 0
|
||||
except ValueError:
|
||||
raise config_error("MongoDB sidecar api_base must be a valid HTTP or HTTPS URL.") from None
|
||||
if not valid or parsed.username or parsed.password or parsed.query or parsed.fragment:
|
||||
raise config_error(
|
||||
"MongoDB sidecar api_base must be an HTTP or HTTPS URL without credentials, query, or fragment."
|
||||
)
|
||||
if timeout is None:
|
||||
return DEFAULT_CONNECT_TIMEOUT_MS, DEFAULT_SOCKET_TIMEOUT_MS
|
||||
return min(int(float(timeout) * 1000), DEFAULT_CONNECT_TIMEOUT_MS), int(float(timeout) * 1000)
|
||||
if parsed.scheme == "http":
|
||||
try:
|
||||
loopback: Final = ip_address(parsed.hostname or "").is_loopback
|
||||
except ValueError:
|
||||
raise config_error(
|
||||
"MongoDB sidecar requires HTTPS. HTTP is supported only for a loopback IP such as 127.0.0.1."
|
||||
) from None
|
||||
if not loopback:
|
||||
raise config_error(
|
||||
"MongoDB sidecar requires HTTPS. HTTP is supported only for a loopback IP such as 127.0.0.1."
|
||||
)
|
||||
return api_base.rstrip("/")
|
||||
|
||||
@staticmethod
|
||||
def _timeout_ms(value: object) -> int:
|
||||
seconds: Final = value.read if isinstance(value, httpx.Timeout) else value
|
||||
if seconds is None:
|
||||
return 30_000
|
||||
if not isinstance(seconds, (int, float)) or not isfinite(seconds) or seconds <= 0:
|
||||
raise config_error("MongoDB search timeout must be a positive finite number.")
|
||||
try:
|
||||
return max(1, int(seconds * 1000))
|
||||
except (ValueError, OverflowError):
|
||||
raise config_error("MongoDB search timeout must be a positive finite number.") from None
|
||||
|
||||
@classmethod
|
||||
def _client_key(cls, params: _MongoDBSearchParams, timeout: float | httpx.Timeout | None) -> MongoClientKey:
|
||||
connect_ms, socket_ms = cls._timeout_ms(timeout)
|
||||
return MongoClientKey(
|
||||
connection_string=params.require_connection_string(),
|
||||
connect_timeout_ms=connect_ms,
|
||||
socket_timeout_ms=socket_ms,
|
||||
server_selection_timeout_ms=min(connect_ms, DEFAULT_SERVER_SELECTION_TIMEOUT_MS),
|
||||
)
|
||||
def _params(
|
||||
cls,
|
||||
litellm_params: Mapping[str, object],
|
||||
optional_params: VectorStoreSearchOptionalRequestParams,
|
||||
extra_body: Mapping[str, object] | None,
|
||||
) -> _MongoDBSearchParams:
|
||||
cls._reject_unknown_params(litellm_params)
|
||||
if extra_body:
|
||||
raise config_error("MongoDB vector store does not support extra_body overrides.")
|
||||
for unsupported in ("filters", "ranking_options", "rewrite_query"):
|
||||
if optional_params.get(unsupported) is not None:
|
||||
raise config_error(f"MongoDB vector store does not support the {unsupported} parameter.")
|
||||
try:
|
||||
params: Final = _MongoDBSearchParams.model_validate(litellm_params)
|
||||
except ValidationError:
|
||||
raise config_error(
|
||||
"Invalid MongoDB vector-store configuration. Check the database, collection, fields, and candidate count."
|
||||
) from None
|
||||
params.require_database()
|
||||
params.require_collection()
|
||||
params.require_embedding_model()
|
||||
cls._num_candidates(cls._limit(optional_params), params.mongodb_num_candidates)
|
||||
cls._timeout_ms(litellm_params.get("timeout"))
|
||||
return params
|
||||
|
||||
@classmethod
|
||||
def _pipeline(
|
||||
def _request(
|
||||
cls,
|
||||
vector_store_id: str,
|
||||
query_vector: Sequence[float],
|
||||
query_text: str,
|
||||
params: _MongoDBSearchParams,
|
||||
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
|
||||
) -> Sequence[Mapping[str, object]]:
|
||||
if vector_store_search_optional_params.get("filters") is not None:
|
||||
optional_params: VectorStoreSearchOptionalRequestParams,
|
||||
api_base: str,
|
||||
embedding_response: EmbeddingResponse,
|
||||
timeout: object,
|
||||
) -> tuple[str, dict[str, object]]: # mutable-ok: the provider contract returns a writable JSON request body
|
||||
if not embedding_response.data:
|
||||
raise config_error(
|
||||
"MongoDB vector store does not support the filters parameter yet. "
|
||||
"Restrict the collection or the MongoDB Vector Search index definition instead."
|
||||
"The embedding model returned no embedding for the search query. Check litellm_embedding_model."
|
||||
)
|
||||
if vector_store_search_optional_params.get("ranking_options") is not None:
|
||||
raise config_error(
|
||||
"MongoDB vector store does not support the ranking_options parameter yet. "
|
||||
"Every result already carries the vectorSearchScore, so filter or re-rank "
|
||||
"on that rather than having the threshold silently ignored."
|
||||
)
|
||||
if vector_store_search_optional_params.get("rewrite_query") is not None:
|
||||
raise config_error(
|
||||
"MongoDB vector store does not support the rewrite_query parameter. The query is "
|
||||
"embedded exactly as sent; rewrite it before calling if you need that."
|
||||
)
|
||||
limit: Final = cls._limit(vector_store_search_optional_params)
|
||||
search: Final = MappingProxyType(
|
||||
{
|
||||
"index": vector_store_id,
|
||||
"path": params.embedding_field,
|
||||
"queryVector": tuple(query_vector),
|
||||
"numCandidates": cls._num_candidates(limit, params.mongodb_num_candidates),
|
||||
"limit": limit,
|
||||
}
|
||||
)
|
||||
projection: Final = MappingProxyType(
|
||||
{params.text_field: 1, SCORE_FIELD_NAME: MappingProxyType({"$meta": "vectorSearchScore"})}
|
||||
)
|
||||
return [ # mutable-ok: pymongo rejects any non-list pipeline in common.validate_list
|
||||
MappingProxyType({"$vectorSearch": search}),
|
||||
MappingProxyType({"$project": projection}),
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def _field_value(cls, document: Mapping[str, object], dotted_path: str) -> str | None:
|
||||
"""None means absent, which is what separates a mistyped field from genuinely empty text."""
|
||||
head, _, rest = dotted_path.partition(".")
|
||||
if head not in document:
|
||||
return None
|
||||
value: Final = document[head]
|
||||
if not rest:
|
||||
return None if value is None else str(value)
|
||||
return cls._field_value(value, rest) if isinstance(value, Mapping) else None
|
||||
|
||||
@classmethod
|
||||
def _to_result(cls, document: Mapping[str, object], text_field: str) -> VectorStoreSearchResult:
|
||||
document_id: Final = document.get("_id")
|
||||
identifier: Final = None if document_id is None else str(document_id)
|
||||
content: Final = [ # mutable-ok: VectorStoreSearchResult declares a list of content parts
|
||||
VectorStoreResultContent(text=cls._field_value(document, text_field) or "", type="text")
|
||||
]
|
||||
raw_score: Final = document.get(SCORE_FIELD_NAME)
|
||||
return VectorStoreSearchResult(
|
||||
score=float(raw_score) if isinstance(raw_score, (int, float)) else None,
|
||||
content=content,
|
||||
file_id=identifier,
|
||||
filename=identifier,
|
||||
vector: Final = embedding_response.data[0]["embedding"]
|
||||
if not vector or any(not isinstance(value, (float, int)) or not isfinite(value) for value in vector):
|
||||
raise config_error("The embedding model must return a non-empty, finite query vector.")
|
||||
limit: Final = cls._limit(optional_params)
|
||||
return (
|
||||
f"{api_base}/v1/vector_stores/{quote(vector_store_id, safe='')}/search",
|
||||
{ # mutable-ok: JSON transport requires a dict
|
||||
"query": query_text,
|
||||
"query_vector": tuple(vector),
|
||||
"mongodb_database": params.require_database(),
|
||||
"mongodb_collection": params.require_collection(),
|
||||
"mongodb_embedding_field": params.embedding_field,
|
||||
"mongodb_text_field": params.text_field,
|
||||
"mongodb_num_candidates": cls._num_candidates(limit, params.mongodb_num_candidates),
|
||||
"max_num_results": limit,
|
||||
"timeout_ms": cls._timeout_ms(timeout),
|
||||
},
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _raise_for_missing_text_field(
|
||||
cls, documents: Sequence[Mapping[str, object]], text_field: str, database: str, collection: str
|
||||
) -> None:
|
||||
"""$vectorSearch matches documents carrying no text, so a mistyped mongodb_text_field
|
||||
returns well-scored results with empty content instead of failing."""
|
||||
if documents and all(cls._field_value(document, text_field) is None for document in documents):
|
||||
raise config_error(
|
||||
f"None of the {len(documents)} matched documents in '{database}.{collection}' has a "
|
||||
f"'{text_field}' field, so every result would carry empty text. Set mongodb_text_field "
|
||||
"to the field holding the readable text; it accepts a dotted path such as metadata.body."
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _to_response(
|
||||
cls, documents: Sequence[Mapping[str, object]], query_text: str, text_field: str
|
||||
) -> VectorStoreSearchResponse:
|
||||
return VectorStoreSearchResponse(
|
||||
object="vector_store.search_results.page",
|
||||
search_query=query_text,
|
||||
data=[ # mutable-ok: VectorStoreSearchResponse declares data as a list
|
||||
cls._to_result(document, text_field) for document in documents
|
||||
],
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _raise_for_unusable_index(
|
||||
catalogue: Sequence[Mapping[str, object]], index_name: str, database: str, collection: str
|
||||
) -> None:
|
||||
"""mongod returns zero documents both for a query that matched nothing and for a missing
|
||||
database, collection or index, so the catalogue decides which one happened."""
|
||||
if not catalogue:
|
||||
raise missing_index_error(index_name, database, collection)
|
||||
entry: Final = catalogue[0]
|
||||
if not entry.get("queryable"):
|
||||
raise index_not_ready_error(index_name, database, collection, str(entry.get("status") or "unknown"))
|
||||
|
||||
@staticmethod
|
||||
def _embedding_vector(embedding_response: EmbeddingResponse) -> Sequence[float]:
|
||||
data: Final = embedding_response.data
|
||||
if not data:
|
||||
raise config_error(
|
||||
"The embedding model returned no embedding for the search query, so there is nothing "
|
||||
"to search MongoDB with. Check the embedding deployment named by litellm_embedding_model."
|
||||
)
|
||||
return data[0]["embedding"]
|
||||
|
||||
def execute_search_vector_store_request(
|
||||
def transform_search_vector_store_request(
|
||||
self,
|
||||
vector_store_id: str,
|
||||
query: str | Sequence[str],
|
||||
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
|
||||
api_base: str,
|
||||
litellm_logging_obj: "LiteLLMLoggingObj",
|
||||
litellm_params: Mapping[str, object],
|
||||
extra_body: Mapping[str, object] | None = None,
|
||||
embedding_executor: VectorStoreEmbeddingExecutor | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
) -> VectorStoreSearchResponse:
|
||||
self._reject_unknown_params(litellm_params)
|
||||
params: Final = _MongoDBSearchParams.model_validate(litellm_params)
|
||||
) -> tuple[str, dict[str, object]]: # mutable-ok: the provider contract returns a writable JSON request body
|
||||
params: Final = self._params(litellm_params, vector_store_search_optional_params, extra_body)
|
||||
query_text: Final = self._query_text(query)
|
||||
key: Final = self._client_key(params, timeout)
|
||||
database: Final = params.require_database()
|
||||
collection: Final = params.require_collection()
|
||||
|
||||
embedding_response: Final = (embedding_executor or self.embedding_executor).embed(
|
||||
params.require_embedding_model(),
|
||||
response: Final = (embedding_executor or self.embedding_executor).embed(
|
||||
params.require_embedding_model(), query_text, params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG
|
||||
)
|
||||
return self._request(
|
||||
vector_store_id,
|
||||
query_text,
|
||||
params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG,
|
||||
)
|
||||
pipeline: Final = self._pipeline(
|
||||
vector_store_id, self._embedding_vector(embedding_response), params, vector_store_search_optional_params
|
||||
params,
|
||||
vector_store_search_optional_params,
|
||||
api_base,
|
||||
response,
|
||||
litellm_params.get("timeout"),
|
||||
)
|
||||
|
||||
try:
|
||||
client: Final = self.sync_client_factory(key)
|
||||
target: Final = client[database][collection] # pyright: ignore[reportIndexIssue] # factory is typed as returning object so injected doubles are accepted
|
||||
documents: Final = tuple(target.aggregate(pipeline))
|
||||
except Exception as e:
|
||||
raise translate_mongo_error(e, index_name=vector_store_id, database=database, collection=collection) from e
|
||||
if not documents:
|
||||
try:
|
||||
catalogue: Final = tuple(target.list_search_indexes(vector_store_id))
|
||||
except Exception as e:
|
||||
raise translate_mongo_error(
|
||||
e, index_name=vector_store_id, database=database, collection=collection
|
||||
) from e
|
||||
self._raise_for_unusable_index(catalogue, vector_store_id, database, collection)
|
||||
self._raise_for_missing_text_field(documents, params.text_field, database, collection)
|
||||
return self._to_response(documents, query_text, params.text_field)
|
||||
|
||||
async def aexecute_search_vector_store_request(
|
||||
async def atransform_search_vector_store_request(
|
||||
self,
|
||||
vector_store_id: str,
|
||||
query: str | Sequence[str],
|
||||
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
|
||||
api_base: str,
|
||||
litellm_logging_obj: "LiteLLMLoggingObj",
|
||||
litellm_params: Mapping[str, object],
|
||||
extra_body: Mapping[str, object] | None = None,
|
||||
embedding_executor: VectorStoreEmbeddingExecutor | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
) -> VectorStoreSearchResponse:
|
||||
self._reject_unknown_params(litellm_params)
|
||||
params: Final = _MongoDBSearchParams.model_validate(litellm_params)
|
||||
) -> tuple[str, dict[str, object]]: # mutable-ok: the provider contract returns a writable JSON request body
|
||||
params: Final = self._params(litellm_params, vector_store_search_optional_params, extra_body)
|
||||
query_text: Final = self._query_text(query)
|
||||
key: Final = self._client_key(params, timeout)
|
||||
database: Final = params.require_database()
|
||||
collection: Final = params.require_collection()
|
||||
|
||||
embedding_response: Final = await (embedding_executor or self.embedding_executor).aembed(
|
||||
params.require_embedding_model(),
|
||||
response: Final = await (embedding_executor or self.embedding_executor).aembed(
|
||||
params.require_embedding_model(), query_text, params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG
|
||||
)
|
||||
return self._request(
|
||||
vector_store_id,
|
||||
query_text,
|
||||
params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG,
|
||||
)
|
||||
pipeline: Final = self._pipeline(
|
||||
vector_store_id, self._embedding_vector(embedding_response), params, vector_store_search_optional_params
|
||||
params,
|
||||
vector_store_search_optional_params,
|
||||
api_base,
|
||||
response,
|
||||
litellm_params.get("timeout"),
|
||||
)
|
||||
|
||||
def transform_search_vector_store_response(
|
||||
self, response: httpx.Response, litellm_logging_obj: "LiteLLMLoggingObj"
|
||||
) -> VectorStoreSearchResponse:
|
||||
try:
|
||||
client: Final = self.async_client_factory(key)
|
||||
target: Final = client[database][collection] # pyright: ignore[reportIndexIssue] # factory is typed as returning object so injected doubles are accepted
|
||||
cursor: Final = await target.aggregate(pipeline)
|
||||
documents: Final = [ # mutable-ok: an async comprehension cannot build a tuple directly
|
||||
document async for document in cursor
|
||||
]
|
||||
except Exception as e:
|
||||
raise translate_mongo_error(e, index_name=vector_store_id, database=database, collection=collection) from e
|
||||
if not documents:
|
||||
try:
|
||||
index_cursor: Final = await target.list_search_indexes(vector_store_id)
|
||||
catalogue: Final = [ # mutable-ok: an async comprehension cannot build a tuple directly
|
||||
entry async for entry in index_cursor
|
||||
]
|
||||
except Exception as e:
|
||||
raise translate_mongo_error(
|
||||
e, index_name=vector_store_id, database=database, collection=collection
|
||||
) from e
|
||||
self._raise_for_unusable_index(catalogue, vector_store_id, database, collection)
|
||||
self._raise_for_missing_text_field(documents, params.text_field, database, collection)
|
||||
return self._to_response(documents, query_text, params.text_field)
|
||||
validated: Final = _SearchResponse.model_validate_json(response.content)
|
||||
return _RESPONSE_ADAPTER.validate_python(validated.model_dump())
|
||||
except ValidationError:
|
||||
raise ServiceUnavailableError(
|
||||
message="MongoDB sidecar returned an invalid search response. Check the sidecar version and deployment.",
|
||||
model=None,
|
||||
llm_provider="mongodb",
|
||||
) from None
|
||||
|
||||
def get_error_class(
|
||||
self, error_message: str, status_code: int, headers: Mapping[str, object] | httpx.Headers
|
||||
) -> BaseLLMException:
|
||||
if status_code == 400:
|
||||
raise config_error(error_message)
|
||||
if status_code == 401:
|
||||
raise AuthenticationError(message="MongoDB sidecar rejected api_key.", model=None, llm_provider="mongodb")
|
||||
if status_code == 408:
|
||||
raise Timeout(message=error_message, model=None, llm_provider="mongodb")
|
||||
raise ServiceUnavailableError(
|
||||
message="MongoDB sidecar is unavailable. Check its address, health, and logs.",
|
||||
model=None,
|
||||
llm_provider="mongodb",
|
||||
)
|
||||
|
||||
def validate_create_vector_store(self) -> NoReturn:
|
||||
raise config_error(_SEARCH_ONLY_MESSAGE)
|
||||
|
||||
def transform_create_vector_store_request(
|
||||
self,
|
||||
vector_store_create_optional_params: VectorStoreCreateOptionalRequestParams,
|
||||
api_base: str,
|
||||
self, vector_store_create_optional_params: VectorStoreCreateOptionalRequestParams, api_base: str
|
||||
) -> NoReturn:
|
||||
raise config_error(_SEARCH_ONLY_MESSAGE)
|
||||
|
||||
|
|
|
|||
|
|
@ -16,6 +16,10 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
|
|||
custom_prompt,
|
||||
ollama_pt,
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.image_handling import (
|
||||
async_inline_remote_media,
|
||||
inline_remote_image_urls,
|
||||
)
|
||||
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
|
||||
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionUsageBlock
|
||||
|
|
@ -344,6 +348,26 @@ class OllamaConfig(BaseConfig):
|
|||
)
|
||||
return model_response
|
||||
|
||||
@property
|
||||
def uses_async_transform_request(self) -> bool:
|
||||
return True
|
||||
|
||||
async def async_transform_request(
|
||||
self,
|
||||
model: str,
|
||||
messages: list[AllMessageValues], # mutable-ok: BaseConfig signature
|
||||
optional_params: dict[str, object], # mutable-ok: BaseConfig signature
|
||||
litellm_params: dict[str, object], # mutable-ok: BaseConfig signature
|
||||
headers: dict[str, object], # mutable-ok: BaseConfig signature
|
||||
) -> dict[str, object]: # mutable-ok: BaseConfig signature
|
||||
return self.transform_request(
|
||||
model=model,
|
||||
messages=await async_inline_remote_media(messages, should_inline=inline_remote_image_urls),
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
def transform_request(
|
||||
self,
|
||||
model: str,
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Final
|
|||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.prompt_templates.image_handling import RemoteMedia, inline_remote_image_urls
|
||||
from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
|
@ -51,6 +52,9 @@ class VertexAIAnthropicConfig(AnthropicConfig):
|
|||
def custom_llm_provider(self) -> str | None:
|
||||
return "vertex_ai"
|
||||
|
||||
def inlines_remote_media(self, media: RemoteMedia) -> bool:
|
||||
return inline_remote_image_urls(media)
|
||||
|
||||
def should_strip_billing_metadata(self) -> bool:
|
||||
return True
|
||||
|
||||
|
|
|
|||
|
|
@ -157,11 +157,9 @@ class IBMWatsonXChatConfig(IBMWatsonXMixin, OpenAIGPTConfig):
|
|||
@staticmethod
|
||||
async def aapply_prompt_template(model: str, messages: list[dict[str, str]]) -> str | None:
|
||||
"""Apply prompt template (async version)"""
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
ahf_chat_template,
|
||||
custom_prompt,
|
||||
hf_chat_template,
|
||||
ibm_granite_pt,
|
||||
mistral_instruct_pt,
|
||||
)
|
||||
|
|
@ -179,11 +177,7 @@ class IBMWatsonXChatConfig(IBMWatsonXMixin, OpenAIGPTConfig):
|
|||
else:
|
||||
hf_model = model
|
||||
try:
|
||||
# Use sync if cached, async if not
|
||||
if hf_model in litellm.known_tokenizer_config:
|
||||
result = hf_chat_template(model=hf_model, messages=messages)
|
||||
else:
|
||||
result = await ahf_chat_template(model=hf_model, messages=messages)
|
||||
result = await ahf_chat_template(model=hf_model, messages=messages)
|
||||
# Return result if it's truthy (not None and not empty string)
|
||||
# The caller (_aconvert_watsonx_messages_core) will handle None/empty by falling back to default
|
||||
if result:
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ from ..common_utils import (
|
|||
IBMWatsonXMixin,
|
||||
WatsonXAIError,
|
||||
_get_api_params,
|
||||
aconvert_watsonx_messages_to_prompt,
|
||||
convert_watsonx_messages_to_prompt,
|
||||
)
|
||||
|
||||
|
|
@ -236,7 +237,11 @@ class IBMWatsonXAIConfig(IBMWatsonXMixin, BaseConfig):
|
|||
**watsonx_auth_payload,
|
||||
}
|
||||
|
||||
async def atransform_request(
|
||||
@property
|
||||
def uses_async_transform_request(self) -> bool:
|
||||
return True
|
||||
|
||||
async def async_transform_request(
|
||||
self,
|
||||
model: str,
|
||||
messages: list[AllMessageValues],
|
||||
|
|
@ -244,11 +249,6 @@ class IBMWatsonXAIConfig(IBMWatsonXMixin, BaseConfig):
|
|||
litellm_params: dict,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
"""Async version of transform_request"""
|
||||
from litellm.llms.watsonx.common_utils import (
|
||||
aconvert_watsonx_messages_to_prompt,
|
||||
)
|
||||
|
||||
provider: Final = model.split("/")[0]
|
||||
prompt: Final = await aconvert_watsonx_messages_to_prompt(
|
||||
model=model, messages=messages, provider=provider, custom_prompt_dict={}
|
||||
|
|
|
|||
|
|
@ -6545,7 +6545,7 @@ def embedding(
|
|||
client=client,
|
||||
timeout=timeout,
|
||||
aembedding=aembedding,
|
||||
litellm_params={},
|
||||
litellm_params=litellm_params_dict,
|
||||
api_base=api_base,
|
||||
print_verbose=print_verbose,
|
||||
extra_headers=headers,
|
||||
|
|
|
|||
|
|
@ -650,7 +650,10 @@
|
|||
},
|
||||
"twelvelabs.marengo-embed-2-7-v1:0": {
|
||||
"deprecation_date": "2026-11-30",
|
||||
"input_cost_per_token": 7e-05,
|
||||
"input_cost_per_query": 7e-05,
|
||||
"input_cost_per_video_per_second": 0.0007,
|
||||
"input_cost_per_audio_per_second": 0.00014,
|
||||
"input_cost_per_image": 0.0001,
|
||||
"litellm_provider": "bedrock",
|
||||
"max_input_tokens": 77,
|
||||
"max_tokens": 77,
|
||||
|
|
@ -662,7 +665,7 @@
|
|||
},
|
||||
"us.twelvelabs.marengo-embed-2-7-v1:0": {
|
||||
"deprecation_date": "2026-11-30",
|
||||
"input_cost_per_token": 7e-05,
|
||||
"input_cost_per_query": 7e-05,
|
||||
"input_cost_per_video_per_second": 0.0007,
|
||||
"input_cost_per_audio_per_second": 0.00014,
|
||||
"input_cost_per_image": 0.0001,
|
||||
|
|
@ -677,7 +680,7 @@
|
|||
},
|
||||
"eu.twelvelabs.marengo-embed-2-7-v1:0": {
|
||||
"deprecation_date": "2026-11-30",
|
||||
"input_cost_per_token": 7e-05,
|
||||
"input_cost_per_query": 7e-05,
|
||||
"input_cost_per_video_per_second": 0.0007,
|
||||
"input_cost_per_audio_per_second": 0.00014,
|
||||
"input_cost_per_image": 0.0001,
|
||||
|
|
@ -690,6 +693,48 @@
|
|||
"supports_embedding_image_input": true,
|
||||
"supports_image_input": true
|
||||
},
|
||||
"twelvelabs.marengo-embed-3-0-v1:0": {
|
||||
"input_cost_per_query": 7e-05,
|
||||
"input_cost_per_video_per_second": 0.0007,
|
||||
"input_cost_per_audio_per_second": 0.00014,
|
||||
"input_cost_per_image": 0.0001,
|
||||
"litellm_provider": "bedrock",
|
||||
"max_input_tokens": 500,
|
||||
"max_tokens": 500,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 0.0,
|
||||
"output_vector_size": 512,
|
||||
"supports_embedding_image_input": true,
|
||||
"supports_image_input": true
|
||||
},
|
||||
"us.twelvelabs.marengo-embed-3-0-v1:0": {
|
||||
"input_cost_per_query": 7e-05,
|
||||
"input_cost_per_video_per_second": 0.0007,
|
||||
"input_cost_per_audio_per_second": 0.00014,
|
||||
"input_cost_per_image": 0.0001,
|
||||
"litellm_provider": "bedrock",
|
||||
"max_input_tokens": 500,
|
||||
"max_tokens": 500,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 0.0,
|
||||
"output_vector_size": 512,
|
||||
"supports_embedding_image_input": true,
|
||||
"supports_image_input": true
|
||||
},
|
||||
"eu.twelvelabs.marengo-embed-3-0-v1:0": {
|
||||
"input_cost_per_query": 7e-05,
|
||||
"input_cost_per_video_per_second": 0.0007,
|
||||
"input_cost_per_audio_per_second": 0.00014,
|
||||
"input_cost_per_image": 0.0001,
|
||||
"litellm_provider": "bedrock",
|
||||
"max_input_tokens": 500,
|
||||
"max_tokens": 500,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 0.0,
|
||||
"output_vector_size": 512,
|
||||
"supports_embedding_image_input": true,
|
||||
"supports_image_input": true
|
||||
},
|
||||
"twelvelabs.pegasus-1-2-v1:0": {
|
||||
"input_cost_per_video_per_second": 0.00049,
|
||||
"output_cost_per_token": 7.5e-06,
|
||||
|
|
|
|||
|
|
@ -19,7 +19,7 @@ from typing import TYPE_CHECKING, Any, Final, Protocol
|
|||
|
||||
import httpx
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from pydantic import AnyUrl, ConfigDict
|
||||
from pydantic import AnyUrl, ConfigDict, TypeAdapter, ValidationError
|
||||
from starlette.requests import Request as StarletteRequest
|
||||
from starlette.responses import JSONResponse
|
||||
from starlette.types import Message, Receive, Scope, Send
|
||||
|
|
@ -108,9 +108,9 @@ _MAX_STATEFUL_SESSIONS_PER_OWNER: Final = 100
|
|||
# prevents an authenticated client from forcing the proxy to buffer an
|
||||
# arbitrarily large body just to make a routing decision.
|
||||
_MCP_ROUTING_PEEK_MAX_BYTES: Final = 4096
|
||||
# ASGI scope key holding the tracing span of the request carrying an MCP
|
||||
# message, written on the request task and read back by the message handler.
|
||||
# ASGI scope keys carrying OTel request state into a stateful MCP message handler.
|
||||
_MCP_TRANSPORT_SPAN_SCOPE_KEY: Final = "litellm_otel_transport_span"
|
||||
_MCP_DESTINATIONS_SCOPE_KEY: Final = "litellm_otel_request_destinations"
|
||||
|
||||
|
||||
def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None:
|
||||
|
|
@ -328,18 +328,17 @@ def _otel_publish_transport_span_on_scope(scope: Scope) -> None:
|
|||
scope[_MCP_TRANSPORT_SPAN_SCOPE_KEY] = span
|
||||
|
||||
|
||||
def _otel_transport_span_from_message(req_ctx: object) -> object:
|
||||
"""The tracing span of the HTTP request that carried this MCP message.
|
||||
|
||||
Read off that request's ASGI scope, reached through the ``Request`` the
|
||||
streamable-HTTP transport attaches to each message, so it is this message's
|
||||
transport and not whichever request happens to have touched the session last.
|
||||
Returns whatever the scope holds; the otel plumbing validates it."""
|
||||
def _otel_value_from_message_scope(req_ctx: object, key: str) -> object:
|
||||
request: Final = getattr(req_ctx, "request", None)
|
||||
scope: Final = getattr(request, "scope", None)
|
||||
if not isinstance(scope, Mapping):
|
||||
return None
|
||||
return scope.get(_MCP_TRANSPORT_SPAN_SCOPE_KEY)
|
||||
return scope.get(key)
|
||||
|
||||
|
||||
def _otel_transport_span_from_message(req_ctx: object) -> object:
|
||||
"""The tracing span of the HTTP request that carried this MCP message."""
|
||||
return _otel_value_from_message_scope(req_ctx, _MCP_TRANSPORT_SPAN_SCOPE_KEY)
|
||||
|
||||
|
||||
def _otel_set_mcp_transport_span(span: object) -> object:
|
||||
|
|
@ -372,6 +371,44 @@ def _otel_reset_mcp_transport_span(token: object) -> None:
|
|||
return
|
||||
|
||||
|
||||
def _otel_publish_request_destinations_on_scope(scope: Scope) -> None:
|
||||
try:
|
||||
from litellm.integrations.otel.plumbing.context import request_destinations
|
||||
|
||||
scope[_MCP_DESTINATIONS_SCOPE_KEY] = request_destinations()
|
||||
except ImportError:
|
||||
return
|
||||
|
||||
|
||||
def _otel_set_mcp_request_destinations(req_ctx: object) -> object:
|
||||
destinations: Final = _otel_value_from_message_scope(req_ctx, _MCP_DESTINATIONS_SCOPE_KEY)
|
||||
if not isinstance(destinations, tuple):
|
||||
return None
|
||||
try:
|
||||
from litellm.integrations.otel.model.destination import OtelDestination
|
||||
from litellm.integrations.otel.plumbing.context import set_request_destinations
|
||||
|
||||
destination_adapter: Final[TypeAdapter[tuple[OtelDestination, ...]]] = TypeAdapter(
|
||||
tuple[OtelDestination, ...],
|
||||
config=ConfigDict(revalidate_instances="always"),
|
||||
)
|
||||
validated_destinations: Final = destination_adapter.validate_python(destinations, strict=True)
|
||||
return set_request_destinations(validated_destinations)
|
||||
except (ImportError, ValidationError):
|
||||
return None
|
||||
|
||||
|
||||
def _otel_reset_mcp_request_destinations(token: object) -> None:
|
||||
if token is None:
|
||||
return
|
||||
try:
|
||||
from litellm.integrations.otel.plumbing.context import reset_request_destinations
|
||||
|
||||
reset_request_destinations(token)
|
||||
except ImportError:
|
||||
return
|
||||
|
||||
|
||||
def _proxy_exception_to_http_exception(exc: ProxyException) -> HTTPException:
|
||||
"""Map a ``ProxyException`` to an ``HTTPException`` that preserves its real
|
||||
status code and headers.
|
||||
|
|
@ -763,10 +800,12 @@ if MCP_AVAILABLE:
|
|||
_session_reset_token = active_mcp_session_var.set(req_ctx.session)
|
||||
_trace_token = None
|
||||
_transport_token = None
|
||||
_destinations_token = None
|
||||
|
||||
try:
|
||||
_trace_token = _otel_set_mcp_trace_carrier(_mcp_meta_trace_carrier(req_ctx))
|
||||
_transport_token = _otel_set_mcp_transport_span(_otel_transport_span_from_message(req_ctx))
|
||||
_destinations_token = _otel_set_mcp_request_destinations(req_ctx)
|
||||
# Get user authentication from context variable
|
||||
(
|
||||
user_api_key_auth,
|
||||
|
|
@ -828,6 +867,7 @@ if MCP_AVAILABLE:
|
|||
# This prevents the HTTP stream from failing and allows the client to get a response
|
||||
return []
|
||||
finally:
|
||||
_otel_reset_mcp_request_destinations(_destinations_token)
|
||||
_otel_reset_mcp_transport_span(_transport_token)
|
||||
_otel_reset_mcp_trace_carrier(_trace_token)
|
||||
if _session_reset_token is not None:
|
||||
|
|
@ -1021,10 +1061,12 @@ if MCP_AVAILABLE:
|
|||
_session_reset_token = active_mcp_session_var.set(req_ctx.session)
|
||||
_trace_token = None
|
||||
_transport_token = None
|
||||
_destinations_token = None
|
||||
|
||||
try:
|
||||
_trace_token = _otel_set_mcp_trace_carrier(_mcp_meta_trace_carrier(req_ctx))
|
||||
_transport_token = _otel_set_mcp_transport_span(_otel_transport_span_from_message(req_ctx))
|
||||
_destinations_token = _otel_set_mcp_request_destinations(req_ctx)
|
||||
# Validate arguments
|
||||
(
|
||||
user_api_key_auth,
|
||||
|
|
@ -1163,6 +1205,7 @@ if MCP_AVAILABLE:
|
|||
|
||||
return response
|
||||
finally:
|
||||
_otel_reset_mcp_request_destinations(_destinations_token)
|
||||
_otel_reset_mcp_transport_span(_transport_token)
|
||||
_otel_reset_mcp_trace_carrier(_trace_token)
|
||||
if _session_reset_token is not None:
|
||||
|
|
@ -4493,6 +4536,7 @@ if MCP_AVAILABLE:
|
|||
|
||||
async def _dispatch() -> None:
|
||||
_otel_publish_transport_span_on_scope(scope)
|
||||
_otel_publish_request_destinations_on_scope(scope)
|
||||
auth_user: Final = _set_or_update_auth_context(
|
||||
user_api_key_auth=user_api_key_auth,
|
||||
mcp_auth_header=mcp_auth_header,
|
||||
|
|
|
|||
|
|
@ -835,6 +835,14 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/team/daily/activity/aggregated",
|
||||
"/team/spend/by_user",
|
||||
"/team/{team_id}/members/me",
|
||||
# POST/GET the team's logging callbacks, and DELETE one of them. Every
|
||||
# handler calls _verify_team_access, which admits only a proxy admin, an
|
||||
# org admin for the team, or an admin of this team.
|
||||
#
|
||||
# team_id is a free-form string, so it spells these with the same path
|
||||
# converter the router uses; the gate matches that converter.
|
||||
"/team/{team_id:path}/callback",
|
||||
"/team/{team_id:path}/callback/{callback_name}",
|
||||
"/model/new",
|
||||
"/model/update",
|
||||
"/model/delete",
|
||||
|
|
@ -3584,6 +3592,7 @@ class AllCallbacks(LiteLLMPydanticObjectBase):
|
|||
ui_callback_name="OpenTelemetry",
|
||||
litellm_callback_params=[
|
||||
"OTEL_EXPORTER",
|
||||
"OTEL_EXPORTER_OTLP_PROTOCOL",
|
||||
"OTEL_ENDPOINT",
|
||||
"OTEL_TRACES_ENDPOINT",
|
||||
"OTEL_HEADERS",
|
||||
|
|
|
|||
|
|
@ -144,31 +144,32 @@ def _validate_push_notification_url(url: str) -> None:
|
|||
raise HTTPException(status_code=400, detail=str(e)) from e
|
||||
|
||||
|
||||
def _caller_identity_headers(user_api_key_dict: UserAPIKeyAuth) -> dict[str, str]:
|
||||
headers: Final[dict[str, str]] = {}
|
||||
if user_api_key_dict.user_id:
|
||||
headers["X-LiteLLM-User-Id"] = user_api_key_dict.user_id
|
||||
if user_api_key_dict.team_id:
|
||||
headers["X-LiteLLM-Team-Id"] = user_api_key_dict.team_id
|
||||
return headers
|
||||
def _caller_identity_headers(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, str]:
|
||||
return MappingProxyType(
|
||||
{
|
||||
name: value
|
||||
for name, value in (
|
||||
("X-LiteLLM-User-Id", user_api_key_dict.user_id),
|
||||
("X-LiteLLM-Team-Id", user_api_key_dict.team_id),
|
||||
)
|
||||
if value
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _forwarding_headers(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
caller_identity: Mapping[str, str],
|
||||
request_data: Mapping[str, object],
|
||||
agent_extra_headers: Mapping[str, str] | None,
|
||||
) -> Mapping[str, str] | None:
|
||||
sanitized: Final = (
|
||||
{k: v for k, v in agent_extra_headers.items() if not k.lower().startswith("x-litellm-")}
|
||||
if agent_extra_headers
|
||||
else None
|
||||
) -> dict[str, str] | None:
|
||||
passthrough: Final = tuple(
|
||||
(name, value)
|
||||
for name, value in (agent_extra_headers.items() if agent_extra_headers else ())
|
||||
if not name.lower().startswith("x-litellm-")
|
||||
)
|
||||
merged: Final = merge_agent_headers(dynamic_headers=sanitized, static_headers=None) or {}
|
||||
identity: Final = _caller_identity_headers(user_api_key_dict)
|
||||
trace_id: Final = request_data.get("litellm_trace_id")
|
||||
if trace_id:
|
||||
identity["X-LiteLLM-Trace-Id"] = str(trace_id)
|
||||
merged.update(identity)
|
||||
trace: Final = (("X-LiteLLM-Trace-Id", str(trace_id)),) if trace_id else ()
|
||||
merged: Final = dict((*passthrough, *caller_identity.items(), *trace))
|
||||
return merged or None
|
||||
|
||||
|
||||
|
|
@ -755,6 +756,7 @@ async def invoke_agent_a2a(
|
|||
ProxyBaseLLMRequestProcessing,
|
||||
)
|
||||
|
||||
caller_identity: Final = _caller_identity_headers(user_api_key_dict)
|
||||
processor: Final = ProxyBaseLLMRequestProcessing(data=body)
|
||||
data, logging_obj = await processor.common_processing_pre_call_logic(
|
||||
request=request,
|
||||
|
|
@ -793,9 +795,13 @@ async def invoke_agent_a2a(
|
|||
if header_name:
|
||||
dynamic_headers[header_name] = val
|
||||
|
||||
agent_extra_headers = merge_agent_headers(
|
||||
dynamic_headers=dynamic_headers or None,
|
||||
static_headers=static_headers or None,
|
||||
agent_extra_headers = _forwarding_headers(
|
||||
caller_identity=caller_identity,
|
||||
request_data=data,
|
||||
agent_extra_headers=merge_agent_headers(
|
||||
dynamic_headers=dynamic_headers or None,
|
||||
static_headers=static_headers or None,
|
||||
),
|
||||
)
|
||||
|
||||
# Databricks App endpoints require a short-lived OAuth M2M token rather
|
||||
|
|
@ -942,12 +948,7 @@ async def invoke_agent_a2a(
|
|||
"method": method,
|
||||
"params": params,
|
||||
}
|
||||
caller_headers: Final = _forwarding_headers(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=data,
|
||||
agent_extra_headers=agent_extra_headers,
|
||||
)
|
||||
result = await _forward_jsonrpc(agent_url, forward_body, extra_headers=caller_headers)
|
||||
result = await _forward_jsonrpc(agent_url, forward_body, extra_headers=agent_extra_headers)
|
||||
if method == "agent/getAuthenticatedExtendedCard":
|
||||
card: Final = result.get("result")
|
||||
if isinstance(card, dict):
|
||||
|
|
@ -988,16 +989,11 @@ async def invoke_agent_a2a(
|
|||
"method": method,
|
||||
"params": params,
|
||||
}
|
||||
sse_caller_headers: Final = _forwarding_headers(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=data,
|
||||
agent_extra_headers=agent_extra_headers,
|
||||
)
|
||||
return await _forward_jsonrpc_sse(
|
||||
agent_url,
|
||||
forward_body,
|
||||
request_id=request_id,
|
||||
extra_headers=sse_caller_headers,
|
||||
extra_headers=agent_extra_headers,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=data,
|
||||
|
|
|
|||
|
|
@ -25,6 +25,11 @@ from litellm.proxy.common_request_processing import (
|
|||
proxy_exception_from_http_exception,
|
||||
)
|
||||
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
|
||||
from litellm.proxy.common_utils.openai_error_payload import (
|
||||
error_status_code,
|
||||
openai_error_param,
|
||||
openai_error_type,
|
||||
)
|
||||
from litellm.types.utils import TokenCountResponse
|
||||
|
||||
router: Final = APIRouter()
|
||||
|
|
@ -243,9 +248,9 @@ async def anthropic_response(
|
|||
return _anthropic_error_json_response(
|
||||
ProxyException(
|
||||
message=getattr(e, "message", error_msg),
|
||||
type=getattr(e, "type", "None"),
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", 500),
|
||||
type=openai_error_type(e, error_status_code(e, 500)),
|
||||
param=openai_error_param(e),
|
||||
code=error_status_code(e, 500),
|
||||
headers=headers,
|
||||
),
|
||||
request,
|
||||
|
|
|
|||
|
|
@ -497,10 +497,22 @@ class RouteChecks:
|
|||
|
||||
def _placeholder_to_regex(match: re.Match) -> str:
|
||||
placeholder: Final = match.group(0).strip("{}")
|
||||
if placeholder.endswith(":path"):
|
||||
# allow "/" in the placeholder value, but don't eat the route suffix after ":"
|
||||
return r"[^:]+"
|
||||
return r"[^/]+"
|
||||
if not placeholder.endswith(":path"):
|
||||
return r"[^/]+"
|
||||
# A ":path" placeholder takes whatever the router's own path
|
||||
# converter takes, slashes and colons alike, so an id spelled with
|
||||
# either (or both) still matches the template it was mounted under.
|
||||
#
|
||||
# Unless the template puts a ":" literal of its own after the
|
||||
# placeholder: the Google routes end in ":generateContent" and
|
||||
# friends, and there the value has to stop before that suffix
|
||||
# rather than swallow it and match a different verb.
|
||||
#
|
||||
# "[\s\S]" rather than ".", because "." stops at a newline and the
|
||||
# path converter does not: a %0A anywhere in the value would leave
|
||||
# the route unmatched here while still reaching the handler, which
|
||||
# turns this gate into a bypass for the lists built on it.
|
||||
return r"[^:]+" if ":" in match.string[match.end() :] else r"[\s\S]+"
|
||||
|
||||
pattern = re.sub(r"\{[^}]+\}", _placeholder_to_regex, pattern)
|
||||
# Anchor the pattern to match the entire string
|
||||
|
|
|
|||
|
|
@ -2837,6 +2837,43 @@ async def _authorize_authenticated_request(
|
|||
|
||||
|
||||
@tracer.wrap()
|
||||
def _seed_request_destinations(user_api_key_dict: UserAPIKeyAuth, request: Request | None = None) -> None:
|
||||
"""Anchor the OTLP destinations this key or team overrides its traces to.
|
||||
|
||||
Called inside the ``auth`` phase span so that span reaches the tenant's account
|
||||
as well, and on the request task so the ``ContextVar`` is inherited by the logging
|
||||
tasks that close the LLM span. Best-effort: trace routing must never fail auth.
|
||||
|
||||
``request`` carries the headers, so a backend this request disabled with
|
||||
``x-litellm-disable-callbacks`` resolves to no destination.
|
||||
|
||||
Only destinations the published fan-out can build are anchored. Anchoring one is
|
||||
what tells the operator's exporter to hold that backend's spans back under
|
||||
``override``, so an unbuildable one would leave the span with nowhere to go.
|
||||
|
||||
The ``postgres`` spans under ``auth`` close before this runs, because they are the
|
||||
reads that resolve the identity being read here. They never reach the tenant's
|
||||
account, and they are never withheld from the operator's backend, whichever mode
|
||||
is set.
|
||||
"""
|
||||
try:
|
||||
from litellm.integrations.otel.logger import fan_out_provider
|
||||
from litellm.integrations.otel.plumbing.context import set_request_destinations
|
||||
from litellm.integrations.otel.plumbing.providers import deliverable_destinations
|
||||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
resolve_tenant_otel_destinations,
|
||||
)
|
||||
|
||||
set_request_destinations(
|
||||
deliverable_destinations(
|
||||
resolve_tenant_otel_destinations(user_api_key_dict, _safe_get_request_headers(request)),
|
||||
fan_out_provider(),
|
||||
)
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 # telemetry routing is best-effort and must never break authentication
|
||||
verbose_proxy_logger.debug("OTel V2: tenant destination resolution failed: %s", exc)
|
||||
|
||||
|
||||
async def user_api_key_auth(
|
||||
request: Request,
|
||||
api_key: str = fastapi.Security(api_key_header),
|
||||
|
|
@ -2883,6 +2920,7 @@ async def user_api_key_auth(
|
|||
raise body_parse_exception
|
||||
raise
|
||||
user_api_key_auth_obj.budget_reservation = None
|
||||
_seed_request_destinations(user_api_key_auth_obj, request)
|
||||
|
||||
# A body that never parsed is authenticated (so the trace carries identity
|
||||
# and this ``auth`` span) but not authorized: there is no model to check it
|
||||
|
|
|
|||
|
|
@ -54,6 +54,12 @@ from litellm.proxy.common_utils.callback_utils import (
|
|||
get_logging_caching_headers,
|
||||
get_remaining_tokens_and_requests_from_request_data,
|
||||
)
|
||||
from litellm.proxy.common_utils.openai_error_payload import (
|
||||
attribute_of,
|
||||
error_status_code,
|
||||
openai_error_param,
|
||||
openai_error_type,
|
||||
)
|
||||
from litellm.proxy.common_utils.sse_keepalive import (
|
||||
SSE_COMMENT_PING_BYTES,
|
||||
coerce_keepalive_interval,
|
||||
|
|
@ -464,46 +470,6 @@ def _stream_usage_tracking_updates(
|
|||
}
|
||||
|
||||
|
||||
def _getattr_object(value: object, name: str, default: object = None) -> object:
|
||||
return getattr(value, name, default)
|
||||
|
||||
|
||||
_OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType(
|
||||
{
|
||||
status.HTTP_401_UNAUTHORIZED: "authentication_error",
|
||||
status.HTTP_403_FORBIDDEN: "permission_error",
|
||||
status.HTTP_429_TOO_MANY_REQUESTS: "rate_limit_error",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _error_status_code(exc: object, default: int) -> int:
|
||||
"""The HTTP status an exception carries, or ``default`` when it carries none."""
|
||||
carried: Final = _getattr_object(exc, "status_code")
|
||||
return carried if isinstance(carried, int) and not isinstance(carried, bool) else default
|
||||
|
||||
|
||||
def _openai_error_type(exc: object, status_code: int) -> str:
|
||||
"""OpenAI types ``error.type`` as a required string, so an exception carrying none
|
||||
falls back to the type its status code stands for."""
|
||||
carried: Final = _getattr_object(exc, "type")
|
||||
if isinstance(carried, str):
|
||||
return carried
|
||||
mapped: Final = _OPENAI_ERROR_TYPE_BY_STATUS.get(status_code)
|
||||
if mapped is not None:
|
||||
return mapped
|
||||
if status_code < status.HTTP_500_INTERNAL_SERVER_ERROR:
|
||||
return "invalid_request_error"
|
||||
return "internal_server_error"
|
||||
|
||||
|
||||
def _openai_error_param(exc: object) -> str | None:
|
||||
"""OpenAI types ``error.param`` as nullable, so an exception carrying none
|
||||
serializes as JSON ``null``."""
|
||||
carried: Final = _getattr_object(exc, "param")
|
||||
return carried if isinstance(carried, str) else None
|
||||
|
||||
|
||||
class _UpstreamHttpResponse(Protocol):
|
||||
@property
|
||||
def status_code(self) -> int: ...
|
||||
|
|
@ -573,15 +539,15 @@ def serialize_http_exception_detail(
|
|||
|
||||
|
||||
def proxy_exception_from_http_exception(exc: HTTPException, headers: dict[str, str]) -> ProxyException:
|
||||
raw_detail: Final = _getattr_object(exc, "detail", str(exc))
|
||||
raw_detail: Final = attribute_of(exc, "detail", str(exc))
|
||||
message, structured_fields = serialize_http_exception_detail(raw_detail)
|
||||
existing_fields: Final = getattr(exc, "provider_specific_fields", None) or {}
|
||||
merged_fields: Final = {**existing_fields, **structured_fields} if structured_fields else (existing_fields or None)
|
||||
error_status: Final = _error_status_code(exc, status.HTTP_400_BAD_REQUEST)
|
||||
error_status: Final = error_status_code(exc, status.HTTP_400_BAD_REQUEST)
|
||||
return ProxyException(
|
||||
message=message,
|
||||
type=_openai_error_type(exc, error_status),
|
||||
param=_openai_error_param(exc),
|
||||
type=openai_error_type(exc, error_status),
|
||||
param=openai_error_param(exc),
|
||||
code=error_status,
|
||||
provider_specific_fields=merged_fields,
|
||||
headers=headers,
|
||||
|
|
@ -865,8 +831,8 @@ def sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]:
|
|||
are byte-identical.
|
||||
"""
|
||||
# Preserve status code from HTTPException (e.g. guardrail blocks)
|
||||
error_status: Final = _error_status_code(exc, status.HTTP_500_INTERNAL_SERVER_ERROR)
|
||||
raw_detail: Final = _getattr_object(exc, "detail", "Error processing stream start")
|
||||
error_status: Final = error_status_code(exc, status.HTTP_500_INTERNAL_SERVER_ERROR)
|
||||
raw_detail: Final = attribute_of(exc, "detail", "Error processing stream start")
|
||||
message, structured_fields = serialize_http_exception_detail(raw_detail)
|
||||
|
||||
existing_fields: Final = getattr(exc, "provider_specific_fields", None) or {}
|
||||
|
|
@ -874,8 +840,8 @@ def sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]:
|
|||
|
||||
error_obj: Final = {
|
||||
"message": message,
|
||||
"type": _openai_error_type(exc, error_status),
|
||||
"param": _openai_error_param(exc),
|
||||
"type": openai_error_type(exc, error_status),
|
||||
"param": openai_error_param(exc),
|
||||
"code": str(error_status),
|
||||
}
|
||||
if not merged_fields:
|
||||
|
|
@ -2815,10 +2781,10 @@ class ProxyBaseLLMRequestProcessing:
|
|||
``ResponsesAPIResponse`` directly. Handle both shapes so the
|
||||
container-ownership recording path can walk ``.output`` either way.
|
||||
"""
|
||||
completed: Final = _getattr_object(stream_response, "completed_response")
|
||||
completed: Final = attribute_of(stream_response, "completed_response")
|
||||
if completed is None:
|
||||
return None
|
||||
response_obj: Final = _getattr_object(completed, "response")
|
||||
response_obj: Final = attribute_of(completed, "response")
|
||||
if response_obj is not None:
|
||||
return response_obj
|
||||
return completed
|
||||
|
|
@ -3468,7 +3434,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
headers = getattr(e, "headers", None) or {}
|
||||
if not headers:
|
||||
# Try to get headers from e.response.headers (httpx.Response)
|
||||
_response: Final = _getattr_object(e, "response")
|
||||
_response: Final = attribute_of(e, "response")
|
||||
if _response is not None:
|
||||
_response_headers: Final = getattr(_response, "headers", None)
|
||||
if _response_headers:
|
||||
|
|
@ -3543,8 +3509,8 @@ class ProxyBaseLLMRequestProcessing:
|
|||
_code = status.HTTP_500_INTERNAL_SERVER_ERROR
|
||||
raise ProxyException(
|
||||
message=redact_internal_details_from_client_message(getattr(e, "message", error_msg)),
|
||||
type=_openai_error_type(e, _code),
|
||||
param=_openai_error_param(e),
|
||||
type=openai_error_type(e, _code),
|
||||
param=openai_error_param(e),
|
||||
openai_code=getattr(e, "code", None),
|
||||
code=_code,
|
||||
provider_specific_fields=getattr(e, "provider_specific_fields", None),
|
||||
|
|
@ -3754,11 +3720,11 @@ class ProxyBaseLLMRequestProcessing:
|
|||
|
||||
if isinstance(e, HTTPException):
|
||||
raise e
|
||||
stream_error_status: Final = _error_status_code(e, status.HTTP_500_INTERNAL_SERVER_ERROR)
|
||||
stream_error_status: Final = error_status_code(e, status.HTTP_500_INTERNAL_SERVER_ERROR)
|
||||
proxy_exception: Final = ProxyException(
|
||||
message=redact_internal_details_from_client_message(getattr(e, "message", str(e))),
|
||||
type=_openai_error_type(e, stream_error_status),
|
||||
param=_openai_error_param(e),
|
||||
type=openai_error_type(e, stream_error_status),
|
||||
param=openai_error_param(e),
|
||||
code=stream_error_status,
|
||||
)
|
||||
stream_completed = True
|
||||
|
|
|
|||
|
|
@ -44,6 +44,91 @@ def _langfuse_environment_error(callback_vars: Mapping[str, str]) -> str | None:
|
|||
return None
|
||||
|
||||
|
||||
# Which credential family a dynamic variable belongs to. The families are the
|
||||
# integrations that share one account: every langfuse_* variable configures the
|
||||
# same Langfuse project whether it rides the classic callback or the OTel one,
|
||||
# and every dd_* variable configures the same Datadog account.
|
||||
_VAR_FAMILIES: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{
|
||||
"arize_": "Arize",
|
||||
"dd_": "Datadog",
|
||||
"gcs_": "GCS",
|
||||
"humanloop_": "Humanloop",
|
||||
"langfuse_": "Langfuse",
|
||||
"langsmith_": "LangSmith",
|
||||
"newrelic_": "New Relic",
|
||||
"posthog_": "PostHog",
|
||||
"wandb_": "Weights & Biases",
|
||||
"weave_": "Weights & Biases",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _family_of(var: str) -> str | None:
|
||||
"""The credential family ``var`` configures, or ``None`` if it configures none.
|
||||
|
||||
``turn_off_message_logging`` and friends belong to no backend, so they carry
|
||||
no credentials anyone could redirect.
|
||||
"""
|
||||
return next((family for prefix, family in _VAR_FAMILIES.items() if var.startswith(prefix)), None)
|
||||
|
||||
|
||||
def cross_entry_family_error(
|
||||
callback_vars: Mapping[str, str] | None,
|
||||
stored_vars_by_entry: Sequence[Mapping[str, str]],
|
||||
) -> str | None:
|
||||
"""Reject an entry that changes what a family another entry holds resolves to.
|
||||
|
||||
Every stored entry's variables are flattened into one dict before a request
|
||||
reads them, and the flattened dict is what the exporter authenticates and
|
||||
addresses with. So an entry naming only a destination is enough to redirect
|
||||
credentials that were written somewhere else: a host on a second entry pairs
|
||||
with the key from the first, and the request carries that key to the new
|
||||
host.
|
||||
|
||||
Two rules together keep the flattened dict out of the caller's hands. A
|
||||
variable the family already configures has to keep the value it has, so
|
||||
nothing already in use can be moved. A variable the family does not yet
|
||||
configure may only carry a value the family already holds, which is what lets
|
||||
the same credential go in under its other spelling (``langfuse_secret`` and
|
||||
``langfuse_secret_key`` are one key) without anything here having to list the
|
||||
spellings. Between them, no value the caller chose can enter the family, and
|
||||
repeating the family as it stands is still allowed -- that is how one
|
||||
integration gets registered for both the success and the failure event.
|
||||
|
||||
A team admin who does want to move a family deletes the entry holding it
|
||||
first, which reveals nothing.
|
||||
|
||||
Only the writers this endpoint newly admits are held to this, because a proxy
|
||||
admin already holds every credential the proxy has.
|
||||
|
||||
``stored_vars_by_entry`` has to arrive decrypted; the credential values are
|
||||
encrypted at rest and ciphertext never equals the plaintext coming in.
|
||||
"""
|
||||
if not callback_vars:
|
||||
return None
|
||||
stored_by_var: Final = {
|
||||
var: value for entry in stored_vars_by_entry for var, value in entry.items() if _family_of(var) is not None
|
||||
}
|
||||
family_values: Final = frozenset(
|
||||
(family, value)
|
||||
for entry in stored_vars_by_entry
|
||||
for var, value in entry.items()
|
||||
if (family := _family_of(var)) is not None
|
||||
)
|
||||
held_families: Final = frozenset(family for family, _ in family_values)
|
||||
return next(
|
||||
(
|
||||
f"{family} is already configured by another callback entry on this team. "
|
||||
f"Remove that entry before setting {var} here."
|
||||
for var, value, family in ((v, callback_vars[v], _family_of(v)) for v in callback_vars)
|
||||
if family in held_families
|
||||
and (stored_by_var[var] != value if var in stored_by_var else (family, value) not in family_values)
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def logging_metadata_config_error(metadata: Mapping[str, object] | None) -> str | None:
|
||||
"""Validate every ``logging`` entry of a team/key metadata payload."""
|
||||
if not metadata:
|
||||
|
|
|
|||
52
litellm/proxy/common_utils/openai_error_payload.py
Normal file
52
litellm/proxy/common_utils/openai_error_payload.py
Normal file
|
|
@ -0,0 +1,52 @@
|
|||
"""Shapes the ``error`` object the proxy answers with so it matches OpenAI's contract:
|
||||
``type`` is a required string and ``param`` is nullable, neither of which the literal
|
||||
string ``"None"`` satisfies."""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from fastapi import status
|
||||
|
||||
_OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType(
|
||||
{
|
||||
status.HTTP_401_UNAUTHORIZED: "authentication_error",
|
||||
status.HTTP_403_FORBIDDEN: "permission_error",
|
||||
status.HTTP_429_TOO_MANY_REQUESTS: "rate_limit_error",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def attribute_of(value: object, name: str, default: object = None) -> object:
|
||||
return getattr(value, name, default)
|
||||
|
||||
|
||||
def error_status_code(exc: object, default: int) -> int:
|
||||
"""The HTTP status an exception carries as ``status_code`` or, the way ``ProxyException``
|
||||
stores it, as a stringified ``code``; ``default`` when it carries neither."""
|
||||
carried: Final = attribute_of(exc, "status_code")
|
||||
if isinstance(carried, int) and not isinstance(carried, bool):
|
||||
return carried
|
||||
stringified: Final = attribute_of(exc, "code")
|
||||
return int(stringified) if isinstance(stringified, str) and stringified.isdecimal() else default
|
||||
|
||||
|
||||
def openai_error_type(exc: object, status_code: int) -> str:
|
||||
"""OpenAI types ``error.type`` as a required string, so an exception carrying none
|
||||
falls back to the type its status code stands for."""
|
||||
carried: Final = attribute_of(exc, "type")
|
||||
if isinstance(carried, str):
|
||||
return carried
|
||||
mapped: Final = _OPENAI_ERROR_TYPE_BY_STATUS.get(status_code)
|
||||
if mapped is not None:
|
||||
return mapped
|
||||
if status_code < status.HTTP_500_INTERNAL_SERVER_ERROR:
|
||||
return "invalid_request_error"
|
||||
return "internal_server_error"
|
||||
|
||||
|
||||
def openai_error_param(exc: object) -> str | None:
|
||||
"""OpenAI types ``error.param`` as nullable, so an exception carrying none
|
||||
serializes as JSON ``null``."""
|
||||
carried: Final = attribute_of(exc, "param")
|
||||
return carried if isinstance(carried, str) else None
|
||||
|
|
@ -1063,7 +1063,7 @@ class DBSpendUpdateWriter:
|
|||
|
||||
await enqueue_spend_logs(prisma_client, (payload,))
|
||||
if payload.get("call_type") in RESPONSES_SESSION_CALL_TYPES:
|
||||
request_spend_log_flush()
|
||||
request_spend_log_flush(prisma_client)
|
||||
else:
|
||||
verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.")
|
||||
|
||||
|
|
|
|||
|
|
@ -64,6 +64,7 @@ DISABLE_PREPARED_STATEMENTS_ENV_VAR: Final = "DATABASE_DISABLE_PREPARED_STATEMEN
|
|||
DisablePreparedStatementsFlag = Annotated[
|
||||
bool, BeforeValidator(partial(token_auth_flag_enabled, env_var=DISABLE_PREPARED_STATEMENTS_ENV_VAR))
|
||||
]
|
||||
MAX_IDLE_CONNECTION_LIFETIME_ENV_VAR: Final = "DATABASE_MAX_IDLE_CONNECTION_LIFETIME"
|
||||
|
||||
# schema.prisma pins `provider = "postgresql"`, so these are the only schemes
|
||||
# Prisma can actually connect with.
|
||||
|
|
@ -217,6 +218,9 @@ class DatabaseURLSettings(BaseSettings):
|
|||
disable_prepared_statements: DisablePreparedStatementsFlag = Field(
|
||||
default=False, validation_alias=DISABLE_PREPARED_STATEMENTS_ENV_VAR
|
||||
)
|
||||
max_idle_connection_lifetime: int | None = Field(
|
||||
default=None, validation_alias=MAX_IDLE_CONNECTION_LIFETIME_ENV_VAR
|
||||
)
|
||||
|
||||
# Writer
|
||||
database_url: str | None = Field(default=None, validation_alias="DATABASE_URL")
|
||||
|
|
@ -453,6 +457,12 @@ class DatabaseURLSettings(BaseSettings):
|
|||
if url:
|
||||
os.environ[env_var] = add_missing_query_params(url, MappingProxyType({"pgbouncer": "true"}))
|
||||
|
||||
lifetime_params: Final = idle_lifetime_params(self.max_idle_connection_lifetime)
|
||||
for env_var in ("DATABASE_URL", "DIRECT_URL"):
|
||||
url = os.environ.get(env_var)
|
||||
if url:
|
||||
os.environ[env_var] = add_missing_query_params(url, lifetime_params)
|
||||
|
||||
# The reader inherits the writer's connection params (pool size, timeouts,
|
||||
# pgbouncer mode). Without this the reader pool ignores the configured cap
|
||||
# and falls back to Prisma's `num_physical_cpus * 2 + 1` default.
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ from datetime import datetime, timedelta
|
|||
from typing import Any, Final, Protocol
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy.db.db_url_settings import add_missing_query_params, connection_params_from_url
|
||||
from litellm.proxy.db.token_auth import (
|
||||
DEFAULT_POSTGRES_PORT,
|
||||
DatabaseTokenAuth,
|
||||
|
|
@ -438,7 +439,10 @@ class PrismaWrapper:
|
|||
return None
|
||||
|
||||
endpoint: Final = self._iam_endpoint if self._iam_endpoint is not None else self._endpoint_from_env()
|
||||
db_url: Final = endpoint.build_url(mint_database_token(auth, endpoint))
|
||||
db_url: Final = add_missing_query_params(
|
||||
endpoint.build_url(mint_database_token(auth, endpoint)),
|
||||
connection_params_from_url(os.environ.get(self._db_url_env_var, "")),
|
||||
)
|
||||
os.environ[self._db_url_env_var] = db_url
|
||||
return db_url
|
||||
|
||||
|
|
|
|||
|
|
@ -20,6 +20,11 @@ from litellm.proxy.common_utils.http_parsing_utils import (
|
|||
coerce_numeric_form_fields,
|
||||
numeric_form_fields,
|
||||
)
|
||||
from litellm.proxy.common_utils.openai_error_payload import (
|
||||
error_status_code,
|
||||
openai_error_param,
|
||||
openai_error_type,
|
||||
)
|
||||
from litellm.proxy.route_llm_request import route_request
|
||||
from litellm.types.images.main import ImageEditRequestParams
|
||||
from litellm.types.llms.openai import ChatCompletionUserMessage
|
||||
|
|
@ -200,18 +205,18 @@ async def image_generation(
|
|||
if isinstance(e, HTTPException):
|
||||
raise ProxyException(
|
||||
message=getattr(e, "message", str(e)),
|
||||
type=getattr(e, "type", "None"),
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
|
||||
type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)),
|
||||
param=openai_error_param(e),
|
||||
code=error_status_code(e, status.HTTP_400_BAD_REQUEST),
|
||||
)
|
||||
else:
|
||||
error_msg: Final = f"{e}"
|
||||
raise ProxyException(
|
||||
message=getattr(e, "message", error_msg),
|
||||
type=getattr(e, "type", "None"),
|
||||
param=getattr(e, "param", "None"),
|
||||
type=openai_error_type(e, error_status_code(e, 500)),
|
||||
param=openai_error_param(e),
|
||||
openai_code=getattr(e, "code", None),
|
||||
code=getattr(e, "status_code", 500),
|
||||
code=error_status_code(e, 500),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ from types import MappingProxyType
|
|||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
|
||||
from fastapi import HTTPException, Request
|
||||
from pydantic import TypeAdapter
|
||||
from pydantic import ValidationError as PydanticValidationError
|
||||
from starlette.datastructures import Headers
|
||||
|
||||
|
|
@ -28,6 +29,7 @@ from litellm.constants import (
|
|||
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
|
||||
SESSION_ID_GENERATED_METADATA_KEY,
|
||||
SESSION_ID_OMITTED_METADATA_KEY,
|
||||
X_LITELLM_DISABLE_CALLBACKS,
|
||||
)
|
||||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
||||
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
|
||||
|
|
@ -159,6 +161,7 @@ from litellm.types.utils import (
|
|||
CustomPricingLiteLLMParams,
|
||||
LlmProviders,
|
||||
ProviderSpecificHeader,
|
||||
StandardCallbackDynamicParams,
|
||||
StandardLoggingUserAPIKeyMetadata,
|
||||
SupportedCacheControls,
|
||||
)
|
||||
|
|
@ -172,6 +175,7 @@ _ENABLE_TEAM_STALE_ALIAS_BYPASS: bool | None = None
|
|||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.integrations.otel.model.destination import OtelDestination
|
||||
from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig
|
||||
from litellm.types.proxy.policy_engine import Policy, PolicyMatchContext
|
||||
|
||||
|
|
@ -975,6 +979,142 @@ def _get_dynamic_logging_metadata(
|
|||
return callback_settings_obj
|
||||
|
||||
|
||||
_TENANT_OTEL_PARAMS: Final = TypeAdapter(StandardCallbackDynamicParams)
|
||||
|
||||
|
||||
def _tenant_otel_params(callback_vars: Mapping[str, str]) -> StandardCallbackDynamicParams:
|
||||
try:
|
||||
return _TENANT_OTEL_PARAMS.validate_python(callback_vars)
|
||||
except PydanticValidationError:
|
||||
return StandardCallbackDynamicParams()
|
||||
|
||||
|
||||
_NO_REQUEST_HEADERS: Final[Mapping[str, str]] = MappingProxyType({})
|
||||
|
||||
|
||||
def _dynamically_disabled_backends(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
request_headers: Mapping[str, str] | None,
|
||||
) -> frozenset[str]:
|
||||
"""The callbacks this request turned off, read the way dispatch reads them.
|
||||
|
||||
Same sources, precedence, and premium gate ``EnterpriseCallbackControls`` applies
|
||||
before it skips a callback: the ``x-litellm-disable-callbacks`` header wins over the
|
||||
key's stored list, team settings are not a source, and a non-premium proxy honours
|
||||
neither. A destination has to agree with that decision, or a backend the key turned
|
||||
off would still be exported to, now through the fan-out instead of the callback.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import premium_user
|
||||
|
||||
if litellm.allow_dynamic_callback_disabling is not True or not premium_user:
|
||||
return frozenset()
|
||||
header: Final = (request_headers if request_headers is not None else _NO_REQUEST_HEADERS).get(
|
||||
X_LITELLM_DISABLE_CALLBACKS
|
||||
)
|
||||
if header is not None:
|
||||
return frozenset(name.strip().lower() for name in header.split(","))
|
||||
metadata: Final = user_api_key_dict.metadata
|
||||
disabled: Final = metadata.get("litellm_disabled_callbacks") if metadata else None
|
||||
if not isinstance(disabled, list):
|
||||
return frozenset()
|
||||
return frozenset(name.lower() for name in disabled if isinstance(name, str))
|
||||
|
||||
|
||||
def resolve_tenant_otel_destinations(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
request_headers: Mapping[str, str] | None = None,
|
||||
) -> "tuple[OtelDestination, ...]":
|
||||
"""The OTLP destinations this request's key or team config overrides its traces to.
|
||||
|
||||
Key settings win over team settings outright, the same precedence
|
||||
``_get_dynamic_logging_metadata`` applies, so one caller never exports the same
|
||||
backend to two accounts. An empty key-level list counts as configured, since that
|
||||
is what disabling a key's callbacks writes. Returns empty when OTEL V2 is off, when
|
||||
neither level named a destination-capable backend, or when the config is
|
||||
incomplete, and the request then keeps the operator's own exporters.
|
||||
|
||||
Two entries naming the same backend merge their ``callback_vars`` last-wins, the
|
||||
way ``convert_key_logging_metadata_to_callback`` merges them, so the destination
|
||||
and the per-request tracer routing cannot read one config two ways.
|
||||
|
||||
A ``failure``-only entry is skipped: a destination is resolved during auth, before
|
||||
the request has an outcome, so honouring the filter would mean holding every span
|
||||
back until the call finishes. Those entries keep today's behaviour instead, where
|
||||
the tenant's credentials reach the backend through per-request tracer routing and
|
||||
the operator's exporter is left alone.
|
||||
|
||||
A backend the request disabled dynamically, through the key's
|
||||
``litellm_disabled_callbacks`` or the ``x-litellm-disable-callbacks`` header in
|
||||
``request_headers``, resolves to no destination, so the fan-out never carries the
|
||||
request tree to that account and the operator's exporter is never suppressed for
|
||||
it. That leaves the request exactly where it stood before destinations existed:
|
||||
the OTel V2 logger itself is not on the disable list's class registry, so its own
|
||||
span still routes to the tenant's credentials the way it did then.
|
||||
"""
|
||||
from litellm.integrations.otel.model.config import is_otel_v2_enabled
|
||||
from litellm.integrations.otel.presets.destinations import destination_for
|
||||
|
||||
if not is_otel_v2_enabled():
|
||||
return ()
|
||||
key_entries: Final = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict)
|
||||
entries: Final = (
|
||||
key_entries
|
||||
if key_entries is not None
|
||||
else KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict)
|
||||
)
|
||||
if not entries:
|
||||
return ()
|
||||
disabled: Final = _dynamically_disabled_backends(user_api_key_dict, request_headers)
|
||||
callbacks: Final = tuple(
|
||||
callback
|
||||
for item in entries
|
||||
if (callback := _get_validated_callback_metadata(item=item, source="otel-destination")) is not None
|
||||
if callback.callback_type != "failure"
|
||||
if callback.callback_name.lower() not in disabled
|
||||
)
|
||||
return tuple(
|
||||
destination
|
||||
for name in dict.fromkeys(callback.callback_name for callback in callbacks)
|
||||
if (
|
||||
destination := destination_for(
|
||||
name,
|
||||
_tenant_otel_params(
|
||||
MappingProxyType(
|
||||
{
|
||||
var: value
|
||||
for callback in callbacks
|
||||
if callback.callback_name == name
|
||||
for var, value in callback.callback_vars.items()
|
||||
}
|
||||
)
|
||||
),
|
||||
_tenant_service_name(user_api_key_dict),
|
||||
)
|
||||
)
|
||||
is not None
|
||||
)
|
||||
|
||||
|
||||
def _tenant_service_name(user_api_key_dict: UserAPIKeyAuth) -> str | None:
|
||||
"""The ``service.name`` this key or team configured, the key winning over its team.
|
||||
|
||||
Same fields and same precedence the request-metadata build applies, read straight
|
||||
off the auth object because destinations resolve during auth, before that metadata
|
||||
is assembled.
|
||||
"""
|
||||
sources: Final = (user_api_key_dict.metadata, user_api_key_dict.team_metadata)
|
||||
return next(
|
||||
(
|
||||
stripped
|
||||
for source in sources
|
||||
if source
|
||||
for field in OTEL_SERVICE_NAME_METADATA_KEYS
|
||||
if isinstance(value := source.get(field), str) and (stripped := value.strip())
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def clean_headers(
|
||||
headers: Headers,
|
||||
litellm_key_header_name: str | None = None,
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ from litellm.proxy._types import (
|
|||
LiteLLM_AuditLogs,
|
||||
LiteLLM_TeamTable,
|
||||
LitellmTableNames,
|
||||
LitellmUserRoles,
|
||||
ProxyErrorTypes,
|
||||
ProxyException,
|
||||
TeamCallbackDeleteResponse,
|
||||
|
|
@ -28,7 +29,10 @@ from litellm.proxy._types import (
|
|||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_utils.callback_config_validation import callback_config_error
|
||||
from litellm.proxy.common_utils.callback_config_validation import (
|
||||
callback_config_error,
|
||||
cross_entry_family_error,
|
||||
)
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
_CALLBACK_VAR_ENCRYPTED_PREFIX,
|
||||
decrypt_callback_vars,
|
||||
|
|
@ -230,6 +234,22 @@ def _callback_error(status_code: int, message: str) -> HTTPException:
|
|||
)
|
||||
|
||||
|
||||
def _unknown_team_error(team_id: str, user_api_key_dict: UserAPIKeyAuth, status_code: int) -> HTTPException:
|
||||
"""Report an unknown team without telling an unauthorized caller that it is unknown.
|
||||
|
||||
These routes are reachable by any authenticated caller so that a team admin can
|
||||
get as far as _verify_team_access. A distinct "does not exist" would therefore let
|
||||
any valid key probe which team ids exist, so a caller who could not have managed
|
||||
the team either way gets the same 403 body _verify_team_access raises.
|
||||
"""
|
||||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
|
||||
return _callback_error(status_code, f"Team id = {team_id} does not exist.")
|
||||
return HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail="You do not have access to this team",
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/team/{team_id:path}/callback",
|
||||
tags=["team management"],
|
||||
|
|
@ -304,10 +324,7 @@ async def add_team_callbacks(
|
|||
# Check if team_id exists already
|
||||
_existing_team = await prisma_client.get_data(team_id=team_id, table_name="team", query_type="find_unique")
|
||||
if _existing_team is None:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={"error": f"Team id = {team_id} does not exist. Please use a different team id."},
|
||||
)
|
||||
raise _unknown_team_error(team_id, user_api_key_dict, status.HTTP_400_BAD_REQUEST)
|
||||
|
||||
# IDOR guard: only proxy admins / org admins / team admins of THIS
|
||||
# team may write callback credentials. Without this, any
|
||||
|
|
@ -326,6 +343,28 @@ async def add_team_callbacks(
|
|||
if team_callback_settings is None or not isinstance(team_callback_settings, list):
|
||||
team_callback_settings = []
|
||||
|
||||
# One entry has to own a credential family end to end. The entries are
|
||||
# flattened into one dict before a request reads them, so an entry
|
||||
# naming only a destination would pair with a key written on another
|
||||
# entry and carry it to that destination -- a key a team admin can read
|
||||
# back nowhere. Repeating a value the owning entry already stores is
|
||||
# fine, which is how one integration covers both events. Proxy admins
|
||||
# are exempt: they already hold every credential the proxy has.
|
||||
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
|
||||
# Decrypted, because the check compares the incoming values against
|
||||
# the stored ones and the credentials are encrypted at rest.
|
||||
decrypted_logging: Final = decrypt_callback_vars(team_metadata).get("logging")
|
||||
stored_entries: Final = decrypted_logging if isinstance(decrypted_logging, list) else ()
|
||||
stored_entry_vars: Final = [ # mutable-ok: read-only input to the check, never stored
|
||||
entry.get("callback_vars") or {} for entry in stored_entries
|
||||
]
|
||||
family_error: Final = cross_entry_family_error(data.callback_vars, stored_entry_vars)
|
||||
if family_error is not None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=family_error,
|
||||
)
|
||||
|
||||
## check if it already exists, for the same callback event
|
||||
for callback in team_callback_settings:
|
||||
if (
|
||||
|
|
@ -452,7 +491,7 @@ async def delete_team_callback(
|
|||
team_id=team_id, table_name="team", query_type="find_unique"
|
||||
)
|
||||
if _existing_team is None:
|
||||
raise _callback_error(404, f"Team id = {team_id} does not exist.")
|
||||
raise _unknown_team_error(team_id, user_api_key_dict, status.HTTP_404_NOT_FOUND)
|
||||
|
||||
# IDOR guard: only proxy admins / org admins / team admins of THIS team may
|
||||
# deregister its callbacks, otherwise any authenticated key holder could
|
||||
|
|
@ -726,10 +765,7 @@ async def get_team_callbacks(
|
|||
# Check if team_id exists
|
||||
_existing_team = await prisma_client.get_data(team_id=team_id, table_name="team", query_type="find_unique")
|
||||
if _existing_team is None:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail={"error": f"Team id = {team_id} does not exist."},
|
||||
)
|
||||
raise _unknown_team_error(team_id, user_api_key_dict, status.HTTP_404_NOT_FOUND)
|
||||
|
||||
# IDOR guard: callback metadata holds third-party API credentials
|
||||
# (Langfuse / Langsmith / GCS). Only proxy admins / org admins /
|
||||
|
|
|
|||
|
|
@ -45,6 +45,11 @@ from litellm.proxy.common_utils.openai_endpoint_utils import (
|
|||
get_custom_llm_provider_from_request_headers,
|
||||
get_custom_llm_provider_from_request_query,
|
||||
)
|
||||
from litellm.proxy.common_utils.openai_error_payload import (
|
||||
error_status_code,
|
||||
openai_error_param,
|
||||
openai_error_type,
|
||||
)
|
||||
from litellm.proxy.openai_files_endpoints.batch_file_validation import (
|
||||
check_batch_file_upload,
|
||||
raise_batch_file_validation_failure,
|
||||
|
|
@ -296,22 +301,22 @@ async def route_create_file(
|
|||
if managed_files_obj is None:
|
||||
raise ProxyException(
|
||||
message="Managed files hook not found",
|
||||
type="None",
|
||||
param="None",
|
||||
type=ProxyErrorTypes.internal_server_error.value,
|
||||
param=None,
|
||||
code=500,
|
||||
)
|
||||
if llm_router is None:
|
||||
raise ProxyException(
|
||||
message="LLM Router not found",
|
||||
type="None",
|
||||
param="None",
|
||||
type=ProxyErrorTypes.internal_server_error.value,
|
||||
param=None,
|
||||
code=500,
|
||||
)
|
||||
if not isinstance(managed_files_obj, BaseFileEndpoints):
|
||||
raise ProxyException(
|
||||
message="Managed files hook is not a BaseFileEndpoints",
|
||||
type="None",
|
||||
param="None",
|
||||
type=ProxyErrorTypes.internal_server_error.value,
|
||||
param=None,
|
||||
code=500,
|
||||
)
|
||||
# Managed files internally calls llm_router.acreate_file() which includes loadbalancing
|
||||
|
|
@ -713,17 +718,17 @@ async def create_file(
|
|||
if isinstance(e, HTTPException):
|
||||
raise ProxyException(
|
||||
message=getattr(e, "message", str(e.detail)),
|
||||
type=getattr(e, "type", "None"),
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
|
||||
type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)),
|
||||
param=openai_error_param(e),
|
||||
code=error_status_code(e, status.HTTP_400_BAD_REQUEST),
|
||||
)
|
||||
else:
|
||||
error_msg: Final = f"{e}"
|
||||
raise ProxyException(
|
||||
message=getattr(e, "message", error_msg),
|
||||
type=getattr(e, "type", "None"),
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", 500),
|
||||
type=openai_error_type(e, error_status_code(e, 500)),
|
||||
param=openai_error_param(e),
|
||||
code=error_status_code(e, 500),
|
||||
)
|
||||
finally:
|
||||
for spool in spools:
|
||||
|
|
@ -812,22 +817,22 @@ async def get_file_content(
|
|||
if managed_files_obj is None:
|
||||
raise ProxyException(
|
||||
message="Managed files hook not found",
|
||||
type="None",
|
||||
param="None",
|
||||
type=ProxyErrorTypes.internal_server_error.value,
|
||||
param=None,
|
||||
code=500,
|
||||
)
|
||||
if llm_router is None:
|
||||
raise ProxyException(
|
||||
message="LLM Router not found",
|
||||
type="None",
|
||||
param="None",
|
||||
type=ProxyErrorTypes.internal_server_error.value,
|
||||
param=None,
|
||||
code=500,
|
||||
)
|
||||
if not isinstance(managed_files_obj, BaseFileEndpoints):
|
||||
raise ProxyException(
|
||||
message="Managed files hook is not a BaseFileEndpoints",
|
||||
type="None",
|
||||
param="None",
|
||||
type=ProxyErrorTypes.internal_server_error.value,
|
||||
param=None,
|
||||
code=500,
|
||||
)
|
||||
|
||||
|
|
@ -1021,17 +1026,17 @@ async def get_file_content(
|
|||
if isinstance(e, HTTPException):
|
||||
raise ProxyException(
|
||||
message=getattr(e, "message", str(e.detail)),
|
||||
type=getattr(e, "type", "None"),
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
|
||||
type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)),
|
||||
param=openai_error_param(e),
|
||||
code=error_status_code(e, status.HTTP_400_BAD_REQUEST),
|
||||
)
|
||||
else:
|
||||
error_msg: Final = f"{e}"
|
||||
raise ProxyException(
|
||||
message=getattr(e, "message", error_msg),
|
||||
type=getattr(e, "type", "None"),
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", 500),
|
||||
type=openai_error_type(e, error_status_code(e, 500)),
|
||||
param=openai_error_param(e),
|
||||
code=error_status_code(e, 500),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1151,15 +1156,15 @@ async def get_file(
|
|||
if managed_files_obj is None:
|
||||
raise ProxyException(
|
||||
message="Managed files hook not found",
|
||||
type="None",
|
||||
param="None",
|
||||
type=ProxyErrorTypes.internal_server_error.value,
|
||||
param=None,
|
||||
code=500,
|
||||
)
|
||||
if not isinstance(managed_files_obj, BaseFileEndpoints):
|
||||
raise ProxyException(
|
||||
message="Managed files hook is not a BaseFileEndpoints",
|
||||
type="None",
|
||||
param="None",
|
||||
type=ProxyErrorTypes.internal_server_error.value,
|
||||
param=None,
|
||||
code=500,
|
||||
)
|
||||
response = await managed_files_obj.afile_retrieve(
|
||||
|
|
@ -1215,17 +1220,17 @@ async def get_file(
|
|||
if isinstance(e, HTTPException):
|
||||
raise ProxyException(
|
||||
message=getattr(e, "message", str(e.detail)),
|
||||
type=getattr(e, "type", "None"),
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
|
||||
type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)),
|
||||
param=openai_error_param(e),
|
||||
code=error_status_code(e, status.HTTP_400_BAD_REQUEST),
|
||||
)
|
||||
else:
|
||||
error_msg: Final = f"{e}"
|
||||
raise ProxyException(
|
||||
message=getattr(e, "message", error_msg),
|
||||
type=getattr(e, "type", "None"),
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", 500),
|
||||
type=openai_error_type(e, error_status_code(e, 500)),
|
||||
param=openai_error_param(e),
|
||||
code=error_status_code(e, 500),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1355,22 +1360,22 @@ async def delete_file(
|
|||
if managed_files_obj is None:
|
||||
raise ProxyException(
|
||||
message="Managed files hook not found",
|
||||
type="None",
|
||||
param="None",
|
||||
type=ProxyErrorTypes.internal_server_error.value,
|
||||
param=None,
|
||||
code=500,
|
||||
)
|
||||
if llm_router is None:
|
||||
raise ProxyException(
|
||||
message="LLM Router not found",
|
||||
type="None",
|
||||
param="None",
|
||||
type=ProxyErrorTypes.internal_server_error.value,
|
||||
param=None,
|
||||
code=500,
|
||||
)
|
||||
if not isinstance(managed_files_obj, BaseFileEndpoints):
|
||||
raise ProxyException(
|
||||
message="Managed files hook is not a BaseFileEndpoints",
|
||||
type="None",
|
||||
param="None",
|
||||
type=ProxyErrorTypes.internal_server_error.value,
|
||||
param=None,
|
||||
code=500,
|
||||
)
|
||||
|
||||
|
|
@ -1427,17 +1432,17 @@ async def delete_file(
|
|||
if isinstance(e, HTTPException):
|
||||
raise ProxyException(
|
||||
message=getattr(e, "message", str(e.detail)),
|
||||
type=getattr(e, "type", "None"),
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
|
||||
type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)),
|
||||
param=openai_error_param(e),
|
||||
code=error_status_code(e, status.HTTP_400_BAD_REQUEST),
|
||||
)
|
||||
else:
|
||||
error_msg: Final = f"{e}"
|
||||
raise ProxyException(
|
||||
message=getattr(e, "message", error_msg),
|
||||
type=getattr(e, "type", "None"),
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", 500),
|
||||
type=openai_error_type(e, error_status_code(e, 500)),
|
||||
param=openai_error_param(e),
|
||||
code=error_status_code(e, 500),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1629,15 +1634,15 @@ async def list_files(
|
|||
if isinstance(e, HTTPException):
|
||||
raise ProxyException(
|
||||
message=getattr(e, "message", str(e.detail)),
|
||||
type=getattr(e, "type", "None"),
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
|
||||
type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)),
|
||||
param=openai_error_param(e),
|
||||
code=error_status_code(e, status.HTTP_400_BAD_REQUEST),
|
||||
)
|
||||
else:
|
||||
error_msg: Final = f"{e}"
|
||||
raise ProxyException(
|
||||
message=getattr(e, "message", error_msg),
|
||||
type=getattr(e, "type", "None"),
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", 500),
|
||||
type=openai_error_type(e, error_status_code(e, 500)),
|
||||
param=openai_error_param(e),
|
||||
code=error_status_code(e, 500),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -78,6 +78,11 @@ from litellm.proxy.common_utils.http_parsing_utils import (
|
|||
_read_request_body,
|
||||
_safe_get_request_headers,
|
||||
)
|
||||
from litellm.proxy.common_utils.openai_error_payload import (
|
||||
error_status_code,
|
||||
openai_error_param,
|
||||
openai_error_type,
|
||||
)
|
||||
from litellm.proxy.common_utils.sse_keepalive import (
|
||||
wrap_passthrough_sse_bytes_with_keepalive_pings,
|
||||
)
|
||||
|
|
@ -311,9 +316,9 @@ async def chat_completion_pass_through_endpoint(
|
|||
error_msg: Final = f"{e}"
|
||||
raise ProxyException(
|
||||
message=getattr(e, "message", error_msg),
|
||||
type=getattr(e, "type", "None"),
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", 500),
|
||||
type=openai_error_type(e, error_status_code(e, 500)),
|
||||
param=openai_error_param(e),
|
||||
code=error_status_code(e, 500),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1728,18 +1733,18 @@ async def pass_through_request(
|
|||
if isinstance(e, HTTPException):
|
||||
raise ProxyException(
|
||||
message=getattr(e, "message", str(getattr(e, "detail", str(e)))),
|
||||
type=getattr(e, "type", "None"),
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
|
||||
type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)),
|
||||
param=openai_error_param(e),
|
||||
code=error_status_code(e, status.HTTP_400_BAD_REQUEST),
|
||||
headers=custom_headers,
|
||||
)
|
||||
else:
|
||||
error_msg: Final = f"{e}"
|
||||
raise ProxyException(
|
||||
message=getattr(e, "message", error_msg),
|
||||
type=getattr(e, "type", "None"),
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", 500),
|
||||
type=openai_error_type(e, error_status_code(e, 500)),
|
||||
param=openai_error_param(e),
|
||||
code=error_status_code(e, 500),
|
||||
headers=custom_headers,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -17,6 +17,11 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
|||
encrypt_value_helper,
|
||||
)
|
||||
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
|
||||
from litellm.proxy.common_utils.openai_error_payload import (
|
||||
error_status_code,
|
||||
openai_error_param,
|
||||
openai_error_type,
|
||||
)
|
||||
from litellm.types.realtime import (
|
||||
RealtimeClientSecretRequest,
|
||||
RealtimeClientSecretResponse,
|
||||
|
|
@ -304,15 +309,15 @@ async def create_realtime_client_secret(
|
|||
if isinstance(e, HTTPException):
|
||||
raise ProxyException(
|
||||
message=getattr(e, "message", str(e)),
|
||||
type=getattr(e, "type", "None"),
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", http_status.HTTP_400_BAD_REQUEST),
|
||||
type=openai_error_type(e, error_status_code(e, http_status.HTTP_400_BAD_REQUEST)),
|
||||
param=openai_error_param(e),
|
||||
code=error_status_code(e, http_status.HTTP_400_BAD_REQUEST),
|
||||
)
|
||||
raise ProxyException(
|
||||
message=getattr(e, "message", str(e)),
|
||||
type=getattr(e, "type", "None"),
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", 500),
|
||||
type=openai_error_type(e, error_status_code(e, 500)),
|
||||
param=openai_error_param(e),
|
||||
code=error_status_code(e, 500),
|
||||
)
|
||||
|
||||
if upstream_resp.status_code != 200:
|
||||
|
|
@ -495,15 +500,15 @@ async def proxy_realtime_calls(
|
|||
if isinstance(e, HTTPException):
|
||||
raise ProxyException(
|
||||
message=getattr(e, "message", str(e)),
|
||||
type=getattr(e, "type", "None"),
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", http_status.HTTP_400_BAD_REQUEST),
|
||||
type=openai_error_type(e, error_status_code(e, http_status.HTTP_400_BAD_REQUEST)),
|
||||
param=openai_error_param(e),
|
||||
code=error_status_code(e, http_status.HTTP_400_BAD_REQUEST),
|
||||
)
|
||||
raise ProxyException(
|
||||
message=getattr(e, "message", str(e)),
|
||||
type=getattr(e, "type", "None"),
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", 500),
|
||||
type=openai_error_type(e, error_status_code(e, 500)),
|
||||
param=openai_error_param(e),
|
||||
code=error_status_code(e, 500),
|
||||
)
|
||||
|
||||
return Response(
|
||||
|
|
@ -608,15 +613,15 @@ async def create_realtime_transcription_session(
|
|||
if isinstance(e, HTTPException):
|
||||
raise ProxyException(
|
||||
message=getattr(e, "detail", getattr(e, "message", str(e))),
|
||||
type=getattr(e, "type", "None"),
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", http_status.HTTP_400_BAD_REQUEST),
|
||||
type=openai_error_type(e, error_status_code(e, http_status.HTTP_400_BAD_REQUEST)),
|
||||
param=openai_error_param(e),
|
||||
code=error_status_code(e, http_status.HTTP_400_BAD_REQUEST),
|
||||
)
|
||||
raise ProxyException(
|
||||
message=getattr(e, "message", str(e)),
|
||||
type=getattr(e, "type", "None"),
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", 500),
|
||||
type=openai_error_type(e, error_status_code(e, 500)),
|
||||
param=openai_error_param(e),
|
||||
code=error_status_code(e, 500),
|
||||
)
|
||||
|
||||
if upstream_resp.status_code != 200:
|
||||
|
|
|
|||
|
|
@ -11,6 +11,11 @@ from litellm._logging import verbose_proxy_logger
|
|||
from litellm.proxy._types import *
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
from litellm.proxy.common_utils.openai_error_payload import (
|
||||
error_status_code,
|
||||
openai_error_param,
|
||||
openai_error_type,
|
||||
)
|
||||
|
||||
router: Final = APIRouter()
|
||||
|
||||
|
|
@ -112,15 +117,15 @@ async def rerank(
|
|||
if isinstance(e, HTTPException):
|
||||
raise ProxyException(
|
||||
message=getattr(e, "message", str(e)),
|
||||
type=getattr(e, "type", "None"),
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
|
||||
type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)),
|
||||
param=openai_error_param(e),
|
||||
code=error_status_code(e, status.HTTP_400_BAD_REQUEST),
|
||||
)
|
||||
else:
|
||||
error_msg: Final = f"{e}"
|
||||
raise ProxyException(
|
||||
message=getattr(e, "message", error_msg),
|
||||
type=getattr(e, "type", "None"),
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", 500),
|
||||
type=openai_error_type(e, error_status_code(e, 500)),
|
||||
param=openai_error_param(e),
|
||||
code=error_status_code(e, 500),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -37,6 +37,7 @@ from litellm.proxy._types import (
|
|||
SpendLogsMetadata,
|
||||
SpendLogsPayload,
|
||||
)
|
||||
from litellm.proxy.common_utils.openai_error_payload import openai_error_param
|
||||
from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.proxy.model_listing import ModelInfoResponse
|
||||
|
|
@ -3806,7 +3807,7 @@ class _StaleReadEngine:
|
|||
class PrismaClient:
|
||||
spend_log_transactions: list = []
|
||||
_spend_log_transactions_lock = asyncio.Lock()
|
||||
spend_log_flush_requested: ClassVar[asyncio.Event] = asyncio.Event()
|
||||
spend_log_flush_requested: "asyncio.Event | None" = None
|
||||
spend_log_queue_bytes: ClassVar[int] = 0
|
||||
spend_logs_queue_monitor_task: "asyncio.Task[None] | None" = None
|
||||
tool_usage_transactions: list["ToolUsageTransaction"] = []
|
||||
|
|
@ -6532,23 +6533,27 @@ async def enqueue_spend_logs(
|
|||
)
|
||||
|
||||
|
||||
def request_spend_log_flush() -> None:
|
||||
"""Wake the queue monitor now rather than leaving the rows for its next poll.
|
||||
def request_spend_log_flush(prisma_client: PrismaClient) -> None:
|
||||
"""Wake this client's queue monitor now rather than leaving the rows for its next poll.
|
||||
|
||||
The Responses API hands the client an id it can chain from straight away, and that
|
||||
lookup reads the DB, so the row cannot sit in this worker's queue for a poll interval.
|
||||
Repeated requests coalesce into the monitor's next pass, so the batching holds.
|
||||
A request made before the monitor is running is dropped, and loses nothing: the
|
||||
monitor reads the queue on its first pass, before it ever waits on a request.
|
||||
"""
|
||||
PrismaClient.spend_log_flush_requested.set()
|
||||
flush_requested: Final = prisma_client.spend_log_flush_requested
|
||||
if flush_requested is not None:
|
||||
flush_requested.set()
|
||||
|
||||
|
||||
async def _wait_for_spend_log_flush_request(interval: float) -> bool:
|
||||
async def _wait_for_spend_log_flush_request(flush_requested: asyncio.Event, interval: float) -> bool:
|
||||
"""Wait out ``interval``, returning early and True when a flush was requested."""
|
||||
try:
|
||||
await asyncio.wait_for(PrismaClient.spend_log_flush_requested.wait(), timeout=interval)
|
||||
await asyncio.wait_for(flush_requested.wait(), timeout=interval)
|
||||
except asyncio.TimeoutError:
|
||||
return False
|
||||
PrismaClient.spend_log_flush_requested.clear()
|
||||
flush_requested.clear()
|
||||
return True
|
||||
|
||||
|
||||
|
|
@ -6975,6 +6980,8 @@ async def _monitor_spend_logs_queue(
|
|||
max_backoff: Final = 30.0 # Maximum backoff interval in seconds
|
||||
backoff_multiplier: Final = 1.5 # Exponential backoff multiplier
|
||||
current_interval = base_interval
|
||||
flush_requested: Final = asyncio.Event()
|
||||
prisma_client.spend_log_flush_requested = flush_requested # rebind-ok: the client owns its monitor's flush signal
|
||||
|
||||
verbose_proxy_logger.info(
|
||||
"Starting spend logs queue monitor (threshold: %s, poll_interval: %ss)", threshold, base_interval
|
||||
|
|
@ -7013,7 +7020,7 @@ async def _monitor_spend_logs_queue(
|
|||
# Exponential backoff when no logs to process
|
||||
current_interval = min(current_interval * backoff_multiplier, max_backoff)
|
||||
|
||||
if await _wait_for_spend_log_flush_request(current_interval):
|
||||
if await _wait_for_spend_log_flush_request(flush_requested, current_interval):
|
||||
current_interval = base_interval
|
||||
except Exception as e:
|
||||
spend_log_error("Error in spend logs queue monitor: %s", str(e), exc=e)
|
||||
|
|
@ -7398,7 +7405,7 @@ def handle_exception_on_proxy(e: Exception) -> ProxyException:
|
|||
return ProxyException(
|
||||
message=getattr(e, "detail", f"error({e})"),
|
||||
type=ProxyErrorTypes.internal_server_error,
|
||||
param=getattr(e, "param", "None"),
|
||||
param=openai_error_param(e),
|
||||
code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR),
|
||||
)
|
||||
elif isinstance(e, ProxyException):
|
||||
|
|
@ -7407,7 +7414,7 @@ def handle_exception_on_proxy(e: Exception) -> ProxyException:
|
|||
return ProxyException(
|
||||
message=str(e),
|
||||
type=ProxyErrorTypes.internal_server_error,
|
||||
param=getattr(e, "param", "None"),
|
||||
param=openai_error_param(e),
|
||||
code=_status_code,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -3656,6 +3656,13 @@ class Router:
|
|||
effective_model_info: Final = kwargs.get("model_info") or deployment.get("model_info") or MappingProxyType({})
|
||||
self._set_failed_deployment_id_on_exception(exception, MappingProxyType({"model_info": effective_model_info}))
|
||||
|
||||
@staticmethod
|
||||
def _stamp_retry_skip_deployment_id(exception: Exception, kwargs: Mapping[str, object]) -> None:
|
||||
effective_model_info: Final = kwargs.get("model_info")
|
||||
deployment_id: Final = effective_model_info.get("id") if isinstance(effective_model_info, Mapping) else None
|
||||
if isinstance(deployment_id, str) and deployment_id:
|
||||
exception.retry_skip_deployment_id = deployment_id # pyright: ignore[reportAttributeAccessIssue] # dynamic stamp, read by _deployment_ids_to_skip_on_retry
|
||||
|
||||
def _update_kwargs_with_default_litellm_params(
|
||||
self, kwargs: dict, metadata_variable_name: str | None = "metadata"
|
||||
) -> None:
|
||||
|
|
@ -4358,6 +4365,7 @@ class Router:
|
|||
model=model,
|
||||
messages=[{"role": "user", "content": "prompt"}],
|
||||
specific_deployment=kwargs.pop("specific_deployment", None),
|
||||
request_kwargs=kwargs,
|
||||
)
|
||||
self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
|
||||
data: Final = deployment["litellm_params"].copy()
|
||||
|
|
@ -4388,6 +4396,7 @@ class Router:
|
|||
verbose_router_logger.info("litellm.image_generation(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e)
|
||||
if model_name is not None:
|
||||
self.fail_calls[model_name] += 1
|
||||
self._stamp_retry_skip_deployment_id(e, kwargs)
|
||||
raise e
|
||||
|
||||
async def aimage_generation(self, prompt: str, model: str, **kwargs):
|
||||
|
|
@ -4472,6 +4481,7 @@ class Router:
|
|||
verbose_router_logger.info("litellm.aimage_generation(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e)
|
||||
if model_name is not None:
|
||||
self.fail_calls[model_name] += 1
|
||||
self._stamp_retry_skip_deployment_id(e, kwargs)
|
||||
raise e
|
||||
|
||||
async def atranscription(self, file: FileTypes, model: str, **kwargs):
|
||||
|
|
@ -4576,6 +4586,7 @@ class Router:
|
|||
verbose_router_logger.info("litellm.atranscription(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e)
|
||||
if model_name is not None:
|
||||
self.fail_calls[model_name] += 1
|
||||
self._stamp_retry_skip_deployment_id(e, kwargs)
|
||||
raise e
|
||||
|
||||
async def aspeech(self, model: str, input: str, voice: str | None = None, **kwargs):
|
||||
|
|
@ -4690,6 +4701,7 @@ class Router:
|
|||
verbose_router_logger.info("litellm.aspeech(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e)
|
||||
if model_name is not None:
|
||||
self.fail_calls[model_name] += 1
|
||||
self._stamp_retry_skip_deployment_id(e, kwargs)
|
||||
raise e
|
||||
|
||||
async def arerank(self, model: str, **kwargs):
|
||||
|
|
@ -4748,6 +4760,7 @@ class Router:
|
|||
verbose_router_logger.info("litellm.arerank(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e)
|
||||
if model_name is not None:
|
||||
self.fail_calls[model_name] += 1
|
||||
self._stamp_retry_skip_deployment_id(e, kwargs)
|
||||
raise e
|
||||
|
||||
def text_completion(
|
||||
|
|
@ -4882,6 +4895,7 @@ class Router:
|
|||
verbose_router_logger.info("litellm.atext_completion(model=%s)\x1b[31m Exception %s\x1b[0m", model, e)
|
||||
if model is not None:
|
||||
self.fail_calls[model] += 1
|
||||
self._stamp_retry_skip_deployment_id(e, kwargs)
|
||||
raise e
|
||||
|
||||
async def aadapter_completion(
|
||||
|
|
@ -4972,6 +4986,7 @@ class Router:
|
|||
verbose_router_logger.info("litellm.aadapter_completion(model=%s)\x1b[31m Exception %s\x1b[0m", model, e)
|
||||
if model is not None:
|
||||
self.fail_calls[model] += 1
|
||||
self._stamp_retry_skip_deployment_id(e, kwargs)
|
||||
raise e
|
||||
|
||||
async def _asearch_with_fallbacks(self, original_function: Callable, **kwargs):
|
||||
|
|
@ -5754,6 +5769,7 @@ class Router:
|
|||
model=model,
|
||||
input=input,
|
||||
specific_deployment=kwargs.pop("specific_deployment", None),
|
||||
request_kwargs=kwargs,
|
||||
)
|
||||
self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
|
||||
data: Final = deployment["litellm_params"].copy()
|
||||
|
|
@ -5792,6 +5808,7 @@ class Router:
|
|||
verbose_router_logger.info("litellm.embedding(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e)
|
||||
if model_name is not None:
|
||||
self.fail_calls[model_name] += 1
|
||||
self._stamp_retry_skip_deployment_id(e, kwargs)
|
||||
raise e
|
||||
|
||||
async def aembedding(
|
||||
|
|
@ -5879,6 +5896,7 @@ class Router:
|
|||
verbose_router_logger.info("litellm.aembedding(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e)
|
||||
if model_name is not None:
|
||||
self.fail_calls[model_name] += 1
|
||||
self._stamp_retry_skip_deployment_id(e, kwargs)
|
||||
raise e
|
||||
|
||||
#### FILES API ####
|
||||
|
|
@ -6252,6 +6270,7 @@ class Router:
|
|||
)
|
||||
if model is not None:
|
||||
self.fail_calls[model] += 1
|
||||
self._stamp_retry_skip_deployment_id(e, kwargs)
|
||||
raise e
|
||||
|
||||
async def aretrieve_batch(
|
||||
|
|
@ -6472,6 +6491,7 @@ class Router:
|
|||
)
|
||||
if model is not None:
|
||||
self.fail_calls[model] += 1
|
||||
self._stamp_retry_skip_deployment_id(e, kwargs)
|
||||
raise e
|
||||
|
||||
async def alist_batches(
|
||||
|
|
@ -7589,7 +7609,9 @@ class Router:
|
|||
|
||||
@staticmethod
|
||||
def _deployment_ids_to_skip_on_retry(exception: Exception, already_skipped: object) -> tuple[str, ...]:
|
||||
failed_deployment_id: Final[str | None] = getattr(exception, "failed_deployment_id", None)
|
||||
failed_deployment_id: Final[str | None] = getattr(exception, "retry_skip_deployment_id", None) or getattr(
|
||||
exception, "failed_deployment_id", None
|
||||
)
|
||||
status_code: Final = getattr(exception, "status_code", None)
|
||||
if not failed_deployment_id or not isinstance(status_code, int):
|
||||
return ()
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import json
|
||||
from collections.abc import Sequence
|
||||
from enum import Enum
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias
|
||||
|
||||
from typing_extensions import ReadOnly, Required, TypedDict, override
|
||||
|
||||
|
|
@ -557,7 +557,7 @@ class AmazonTitanMultimodalEmbeddingResponse(TypedDict):
|
|||
message: str # Specifies any errors that occur during generation.
|
||||
|
||||
|
||||
# TwelveLabs Marengo Embed 2.7 types
|
||||
# TwelveLabs Marengo Embed types
|
||||
TWELVELABS_EMBEDDING_INPUT_TYPES = Literal["text", "image", "video", "audio"]
|
||||
TWELVELABS_EMBEDDING_OPTIONS = Literal["visual-text", "visual-image", "audio"]
|
||||
|
||||
|
|
@ -591,6 +591,113 @@ class TwelveLabsMarengoEmbeddingResponse(TypedDict):
|
|||
endSec: float
|
||||
|
||||
|
||||
TWELVELABS_MARENGO_3_INPUT_TYPES: TypeAlias = Literal["text", "image", "video", "audio", "text_image", "multi_input"]
|
||||
TWELVELABS_MARENGO_3_EMBEDDING_OPTIONS: TypeAlias = Literal["visual", "audio", "transcription"]
|
||||
TWELVELABS_MARENGO_3_EMBEDDING_TYPES: TypeAlias = Literal["separate_embedding", "fused_embedding"]
|
||||
TWELVELABS_MARENGO_3_EMBEDDING_SCOPES: TypeAlias = Literal["clip", "asset"]
|
||||
|
||||
|
||||
class TwelveLabsMarengo3FixedSegmentationConfig(TypedDict):
|
||||
durationSec: ReadOnly[int]
|
||||
|
||||
|
||||
class TwelveLabsMarengo3FixedSegmentation(TypedDict):
|
||||
method: ReadOnly[Literal["fixed"]]
|
||||
fixed: ReadOnly[TwelveLabsMarengo3FixedSegmentationConfig]
|
||||
|
||||
|
||||
class TwelveLabsMarengo3DynamicSegmentationConfig(TypedDict):
|
||||
minDurationSec: ReadOnly[int]
|
||||
|
||||
|
||||
class TwelveLabsMarengo3DynamicSegmentation(TypedDict):
|
||||
method: ReadOnly[Literal["dynamic"]]
|
||||
dynamic: ReadOnly[TwelveLabsMarengo3DynamicSegmentationConfig]
|
||||
|
||||
|
||||
TwelveLabsMarengo3Segmentation: TypeAlias = TwelveLabsMarengo3FixedSegmentation | TwelveLabsMarengo3DynamicSegmentation
|
||||
|
||||
|
||||
class TwelveLabsMarengo3TextInput(TypedDict):
|
||||
inputText: ReadOnly[str]
|
||||
|
||||
|
||||
class TwelveLabsMarengo3ImageInput(TypedDict):
|
||||
mediaSource: ReadOnly[TwelveLabsMediaSource]
|
||||
|
||||
|
||||
class TwelveLabsMarengo3TimedMediaOptions(TypedDict, total=False):
|
||||
startSec: ReadOnly[float]
|
||||
endSec: ReadOnly[float]
|
||||
segmentation: ReadOnly[TwelveLabsMarengo3Segmentation]
|
||||
embeddingOption: ReadOnly[Sequence[TWELVELABS_MARENGO_3_EMBEDDING_OPTIONS]]
|
||||
embeddingType: ReadOnly[Sequence[TWELVELABS_MARENGO_3_EMBEDDING_TYPES]]
|
||||
embeddingScope: ReadOnly[Sequence[TWELVELABS_MARENGO_3_EMBEDDING_SCOPES]]
|
||||
|
||||
|
||||
class TwelveLabsMarengo3TimedMediaInput(TwelveLabsMarengo3TimedMediaOptions):
|
||||
mediaSource: Required[ReadOnly[TwelveLabsMediaSource]]
|
||||
|
||||
|
||||
class TwelveLabsMarengo3TextImageInput(TypedDict):
|
||||
inputText: ReadOnly[str]
|
||||
mediaSource: ReadOnly[TwelveLabsMediaSource]
|
||||
|
||||
|
||||
class TwelveLabsMarengo3NamedMediaSource(TwelveLabsMediaSource):
|
||||
name: Required[ReadOnly[str]]
|
||||
mediaType: Required[ReadOnly[Literal["image"]]]
|
||||
|
||||
|
||||
class TwelveLabsMarengo3MultiInput(TypedDict, total=False):
|
||||
inputText: ReadOnly[str]
|
||||
mediaSources: Required[ReadOnly[Sequence[TwelveLabsMarengo3NamedMediaSource]]]
|
||||
|
||||
|
||||
class TwelveLabsMarengo3RequestBase(TypedDict, total=False):
|
||||
inferenceId: ReadOnly[str]
|
||||
|
||||
|
||||
class TwelveLabsMarengo3TextRequest(TwelveLabsMarengo3RequestBase):
|
||||
inputType: ReadOnly[Literal["text"]]
|
||||
text: ReadOnly[TwelveLabsMarengo3TextInput]
|
||||
|
||||
|
||||
class TwelveLabsMarengo3ImageRequest(TwelveLabsMarengo3RequestBase):
|
||||
inputType: ReadOnly[Literal["image"]]
|
||||
image: ReadOnly[TwelveLabsMarengo3ImageInput]
|
||||
|
||||
|
||||
class TwelveLabsMarengo3VideoRequest(TwelveLabsMarengo3RequestBase):
|
||||
inputType: ReadOnly[Literal["video"]]
|
||||
video: ReadOnly[TwelveLabsMarengo3TimedMediaInput]
|
||||
|
||||
|
||||
class TwelveLabsMarengo3AudioRequest(TwelveLabsMarengo3RequestBase):
|
||||
inputType: ReadOnly[Literal["audio"]]
|
||||
audio: ReadOnly[TwelveLabsMarengo3TimedMediaInput]
|
||||
|
||||
|
||||
class TwelveLabsMarengo3TextImageRequest(TwelveLabsMarengo3RequestBase):
|
||||
inputType: ReadOnly[Literal["text_image"]]
|
||||
text_image: ReadOnly[TwelveLabsMarengo3TextImageInput]
|
||||
|
||||
|
||||
class TwelveLabsMarengo3MultiInputRequest(TwelveLabsMarengo3RequestBase):
|
||||
inputType: ReadOnly[Literal["multi_input"]]
|
||||
multi_input: ReadOnly[TwelveLabsMarengo3MultiInput]
|
||||
|
||||
|
||||
TwelveLabsMarengo3EmbeddingRequest: TypeAlias = (
|
||||
TwelveLabsMarengo3TextRequest
|
||||
| TwelveLabsMarengo3ImageRequest
|
||||
| TwelveLabsMarengo3VideoRequest
|
||||
| TwelveLabsMarengo3AudioRequest
|
||||
| TwelveLabsMarengo3TextImageRequest
|
||||
| TwelveLabsMarengo3MultiInputRequest
|
||||
)
|
||||
|
||||
|
||||
class TwelveLabsS3OutputDataConfig(TypedDict):
|
||||
s3Uri: str
|
||||
|
||||
|
|
@ -601,7 +708,7 @@ class TwelveLabsOutputDataConfig(TypedDict):
|
|||
|
||||
class TwelveLabsAsyncInvokeRequest(TypedDict):
|
||||
modelId: str
|
||||
modelInput: TwelveLabsMarengoEmbeddingRequest
|
||||
modelInput: ReadOnly[TwelveLabsMarengoEmbeddingRequest | TwelveLabsMarengo3EmbeddingRequest]
|
||||
outputDataConfig: TwelveLabsOutputDataConfig
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -272,7 +272,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
|
|||
input_cost_per_token_above_272k_tokens_flex: float | None
|
||||
input_cost_per_token_above_512k_tokens: float | None # MiniMax-M3: prompts >512K priced at 2x input
|
||||
input_cost_per_character_above_128k_tokens: float | None # only for vertex ai models
|
||||
input_cost_per_query: float | None # only for rerank models
|
||||
input_cost_per_query: float | None # per-request pricing: rerank, search, and Bedrock Marengo embeddings
|
||||
input_cost_per_image: float | None # only for vertex ai models
|
||||
input_cost_per_image_token: float | None # for gpt-image-1 and similar models
|
||||
input_cost_per_video_token: float | None # for gemini omni models with video input
|
||||
|
|
@ -1694,6 +1694,9 @@ class PromptTokensDetailsWrapper(
|
|||
audio_length_seconds: float | None = None
|
||||
"""Length of audio sent to the model. Used for multimodal embeddings priced per audio-second."""
|
||||
|
||||
query_count: int | None = None
|
||||
"""Number of billable requests sent to the model. Used for embeddings priced per request, such as Bedrock Marengo."""
|
||||
|
||||
cache_write_tokens: int | None = None
|
||||
"""Number of cache write (creation) tokens sent to the model. OpenAI naming (prompt_tokens_details.cache_write_tokens); this is the canonical field."""
|
||||
|
||||
|
|
@ -1735,6 +1738,8 @@ class PromptTokensDetailsWrapper(
|
|||
del self.video_length_seconds
|
||||
if self.audio_length_seconds is None:
|
||||
del self.audio_length_seconds
|
||||
if self.query_count is None:
|
||||
del self.query_count
|
||||
if self.web_search_requests is None:
|
||||
del self.web_search_requests
|
||||
if self.google_maps_grounding_requests is None:
|
||||
|
|
|
|||
|
|
@ -3632,7 +3632,7 @@ def get_optional_params_embeddings(
|
|||
elif "cohere.embed" in model:
|
||||
object = litellm.BedrockCohereEmbeddingConfig()
|
||||
elif "twelvelabs" in model or "marengo" in model:
|
||||
object = litellm.TwelveLabsMarengoEmbeddingConfig()
|
||||
object = litellm.TwelveLabsMarengoEmbeddingConfig(model=model)
|
||||
elif "nova" in model.lower():
|
||||
object = litellm.AmazonNovaEmbeddingConfig()
|
||||
else: # unmapped model
|
||||
|
|
@ -6043,7 +6043,7 @@ def get_model_info(
|
|||
input_cost_per_character_above_128k_tokens: Optional[
|
||||
float
|
||||
] # only for vertex ai models
|
||||
input_cost_per_query: Optional[float] # only for rerank models
|
||||
input_cost_per_query: Optional[float] # per-request pricing: rerank, search, and Bedrock Marengo embeddings
|
||||
input_cost_per_image: Optional[float] # only for vertex ai models
|
||||
input_cost_per_audio_token: Optional[float]
|
||||
input_cost_per_audio_per_second: Optional[float] # only for vertex ai models
|
||||
|
|
|
|||
|
|
@ -650,7 +650,10 @@
|
|||
},
|
||||
"twelvelabs.marengo-embed-2-7-v1:0": {
|
||||
"deprecation_date": "2026-11-30",
|
||||
"input_cost_per_token": 7e-05,
|
||||
"input_cost_per_query": 7e-05,
|
||||
"input_cost_per_video_per_second": 0.0007,
|
||||
"input_cost_per_audio_per_second": 0.00014,
|
||||
"input_cost_per_image": 0.0001,
|
||||
"litellm_provider": "bedrock",
|
||||
"max_input_tokens": 77,
|
||||
"max_tokens": 77,
|
||||
|
|
@ -662,7 +665,7 @@
|
|||
},
|
||||
"us.twelvelabs.marengo-embed-2-7-v1:0": {
|
||||
"deprecation_date": "2026-11-30",
|
||||
"input_cost_per_token": 7e-05,
|
||||
"input_cost_per_query": 7e-05,
|
||||
"input_cost_per_video_per_second": 0.0007,
|
||||
"input_cost_per_audio_per_second": 0.00014,
|
||||
"input_cost_per_image": 0.0001,
|
||||
|
|
@ -677,7 +680,7 @@
|
|||
},
|
||||
"eu.twelvelabs.marengo-embed-2-7-v1:0": {
|
||||
"deprecation_date": "2026-11-30",
|
||||
"input_cost_per_token": 7e-05,
|
||||
"input_cost_per_query": 7e-05,
|
||||
"input_cost_per_video_per_second": 0.0007,
|
||||
"input_cost_per_audio_per_second": 0.00014,
|
||||
"input_cost_per_image": 0.0001,
|
||||
|
|
@ -690,6 +693,48 @@
|
|||
"supports_embedding_image_input": true,
|
||||
"supports_image_input": true
|
||||
},
|
||||
"twelvelabs.marengo-embed-3-0-v1:0": {
|
||||
"input_cost_per_query": 7e-05,
|
||||
"input_cost_per_video_per_second": 0.0007,
|
||||
"input_cost_per_audio_per_second": 0.00014,
|
||||
"input_cost_per_image": 0.0001,
|
||||
"litellm_provider": "bedrock",
|
||||
"max_input_tokens": 500,
|
||||
"max_tokens": 500,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 0.0,
|
||||
"output_vector_size": 512,
|
||||
"supports_embedding_image_input": true,
|
||||
"supports_image_input": true
|
||||
},
|
||||
"us.twelvelabs.marengo-embed-3-0-v1:0": {
|
||||
"input_cost_per_query": 7e-05,
|
||||
"input_cost_per_video_per_second": 0.0007,
|
||||
"input_cost_per_audio_per_second": 0.00014,
|
||||
"input_cost_per_image": 0.0001,
|
||||
"litellm_provider": "bedrock",
|
||||
"max_input_tokens": 500,
|
||||
"max_tokens": 500,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 0.0,
|
||||
"output_vector_size": 512,
|
||||
"supports_embedding_image_input": true,
|
||||
"supports_image_input": true
|
||||
},
|
||||
"eu.twelvelabs.marengo-embed-3-0-v1:0": {
|
||||
"input_cost_per_query": 7e-05,
|
||||
"input_cost_per_video_per_second": 0.0007,
|
||||
"input_cost_per_audio_per_second": 0.00014,
|
||||
"input_cost_per_image": 0.0001,
|
||||
"litellm_provider": "bedrock",
|
||||
"max_input_tokens": 500,
|
||||
"max_tokens": 500,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 0.0,
|
||||
"output_vector_size": 512,
|
||||
"supports_embedding_image_input": true,
|
||||
"supports_image_input": true
|
||||
},
|
||||
"twelvelabs.pegasus-1-2-v1:0": {
|
||||
"input_cost_per_video_per_second": 0.00049,
|
||||
"output_cost_per_token": 7.5e-06,
|
||||
|
|
|
|||
|
|
@ -114,7 +114,6 @@ caching = ["diskcache>=5.6.3,<6.0"]
|
|||
mcp = ["mcp>=1.28.1,<2.0"]
|
||||
# Driver for the MongoDB Atlas vector store; Atlas Vector Search has no HTTP query API.
|
||||
# The floor is 4.9 because that is the release AsyncMongoClient landed in.
|
||||
mongodb = ["pymongo>=4.9,<5.0"]
|
||||
# SAML SSO for the admin UI. python3-saml pulls in xmlsec/lxml, whose wheels
|
||||
# bundle the native libxmlsec1/libxml2 libraries, so no system packages are
|
||||
# required. Kept out of the base `proxy` extra so it stays optional.
|
||||
|
|
|
|||
|
|
@ -6,14 +6,18 @@ import json
|
|||
import threading
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import replace
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from http.server import BaseHTTPRequestHandler, HTTPServer, ThreadingHTTPServer
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("opentelemetry")
|
||||
|
||||
from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( # noqa: E402
|
||||
ExportTraceServiceRequest,
|
||||
)
|
||||
from opentelemetry.sdk.metrics import MeterProvider # noqa: E402
|
||||
from opentelemetry.sdk.metrics.export import InMemoryMetricReader # noqa: E402
|
||||
from opentelemetry.sdk.trace import TracerProvider # noqa: E402
|
||||
from opentelemetry.sdk.trace.export import ( # noqa: E402
|
||||
BatchSpanProcessor,
|
||||
ConsoleSpanExporter,
|
||||
|
|
@ -541,6 +545,89 @@ def test_build_span_exporter_variants():
|
|||
assert "OTLPSpanExporter" in type(http_exporter).__name__
|
||||
|
||||
|
||||
def _export_one_trace_to_local_collector(exporter_kind: str) -> tuple[list[dict], tuple[int, int, int]]:
|
||||
"""Run a parent/child trace through the configured exporter against a
|
||||
throwaway HTTP collector. Returns the requests as the collector saw them
|
||||
(child first, since it ends first) and (trace_id, parent span_id, child span_id)."""
|
||||
received: list[dict] = []
|
||||
|
||||
class Collector(BaseHTTPRequestHandler):
|
||||
def do_POST(self):
|
||||
body = self.rfile.read(int(self.headers["Content-Length"]))
|
||||
received.append({"path": self.path, "headers": dict(self.headers), "body": body})
|
||||
self.send_response(200)
|
||||
self.end_headers()
|
||||
|
||||
def log_message(self, *_args):
|
||||
pass
|
||||
|
||||
server = HTTPServer(("127.0.0.1", 0), Collector)
|
||||
threading.Thread(target=server.serve_forever, daemon=True).start()
|
||||
try:
|
||||
config = OpenTelemetryV2Config(
|
||||
exporter=exporter_kind,
|
||||
endpoint=f"http://127.0.0.1:{server.server_port}",
|
||||
headers="x-collector-token=secret",
|
||||
)
|
||||
provider = TracerProvider()
|
||||
provider.add_span_processor(SimpleSpanProcessor(providers.build_span_exporter(config)))
|
||||
tracer = provider.get_tracer("test")
|
||||
with tracer.start_as_current_span("parent", kind=SpanKind.SERVER) as parent:
|
||||
with tracer.start_as_current_span("child") as child:
|
||||
ids = (
|
||||
parent.get_span_context().trace_id,
|
||||
parent.get_span_context().span_id,
|
||||
child.get_span_context().span_id,
|
||||
)
|
||||
provider.shutdown()
|
||||
finally:
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
assert len(received) == 2
|
||||
return received, ids
|
||||
|
||||
|
||||
def _only_span(request: dict) -> dict:
|
||||
scope_spans = json.loads(request["body"])["resourceSpans"][0]["scopeSpans"][0]["spans"]
|
||||
assert len(scope_spans) == 1
|
||||
return scope_spans[0]
|
||||
|
||||
|
||||
def test_http_json_exporter_posts_otlp_json_to_traces_endpoint():
|
||||
"""``http/json`` must put the OTLP/JSON mapping on the wire (camelCase
|
||||
fields, integer enums, hex ids) with a JSON content type, so collectors that
|
||||
cannot decode protobuf can ingest the trace. Headers still travel."""
|
||||
(child_request, parent_request), (trace_id, parent_id, child_id) = _export_one_trace_to_local_collector("http/json")
|
||||
|
||||
assert parent_request["path"] == "/v1/traces"
|
||||
assert parent_request["headers"]["Content-Type"] == "application/json"
|
||||
assert parent_request["headers"]["x-collector-token"] == "secret"
|
||||
parent = _only_span(parent_request)
|
||||
assert parent["name"] == "parent"
|
||||
assert parent["kind"] == 2
|
||||
assert parent["traceId"] == format(trace_id, "032x")
|
||||
assert parent["spanId"] == format(parent_id, "016x")
|
||||
assert "parentSpanId" not in parent
|
||||
child = _only_span(child_request)
|
||||
assert child["traceId"] == format(trace_id, "032x")
|
||||
assert child["spanId"] == format(child_id, "016x")
|
||||
assert child["parentSpanId"] == format(parent_id, "016x")
|
||||
|
||||
|
||||
def test_http_protobuf_exporter_still_posts_protobuf():
|
||||
(_child_request, parent_request), (trace_id, _parent_id, _child_id) = _export_one_trace_to_local_collector(
|
||||
"http/protobuf"
|
||||
)
|
||||
|
||||
assert parent_request["path"] == "/v1/traces"
|
||||
assert parent_request["headers"]["Content-Type"] == "application/x-protobuf"
|
||||
assert format(trace_id, "032x").encode() not in parent_request["body"]
|
||||
decoded = ExportTraceServiceRequest.FromString(parent_request["body"])
|
||||
span = decoded.resource_spans[0].scope_spans[0].spans[0]
|
||||
assert span.name == "parent"
|
||||
assert span.trace_id == trace_id.to_bytes(16, "big")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def otlp_collector() -> Iterator[tuple[str, list[str]]]:
|
||||
received_paths: list[str] = []
|
||||
|
|
@ -612,6 +699,21 @@ def test_traces_endpoint_per_exporter_coexists_with_default_normalization(otlp_c
|
|||
assert sorted(received_paths) == ["/services/collector/traces", "/v1/traces"]
|
||||
|
||||
|
||||
def test_http_json_exporter_honors_traces_endpoint(otlp_collector):
|
||||
base_url, received_paths = otlp_collector
|
||||
cfg = OpenTelemetryV2Config(
|
||||
exporters=[
|
||||
{
|
||||
"kind": "http/json",
|
||||
"endpoint": base_url,
|
||||
"traces_endpoint": f"{base_url}/services/collector/traces",
|
||||
}
|
||||
]
|
||||
)
|
||||
_export_one_span(cfg)
|
||||
assert received_paths == ["/services/collector/traces"]
|
||||
|
||||
|
||||
def test_otlp_metric_exporter_uses_cumulative_histogram_temporality():
|
||||
"""Histograms must export as cumulative, not delta.
|
||||
|
||||
|
|
|
|||
2719
tests/test_litellm/integrations/otel/test_otel_v2_destinations.py
Normal file
2719
tests/test_litellm/integrations/otel/test_otel_v2_destinations.py
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -2041,7 +2041,7 @@ def test_select_global_otel_v2_logger_builds_one_when_none_registered():
|
|||
assert isinstance(chosen, OpenTelemetryV2)
|
||||
|
||||
|
||||
def test_publish_global_otel_v2_provider_sets_selected_logger_provider():
|
||||
def test_publish_global_otel_v2_provider_sets_selected_logger_provider(monkeypatch):
|
||||
"""The startup publish must set the OTel global provider to the *selected*
|
||||
logger's provider (the preset logger that owns every exporter), so the FastAPI
|
||||
server span and the gen-ai spans share one provider and one trace.
|
||||
|
|
@ -2051,8 +2051,10 @@ def test_publish_global_otel_v2_provider_sets_selected_logger_provider():
|
|||
test would otherwise miss: that the published provider is the selected logger's,
|
||||
not some other.
|
||||
"""
|
||||
from litellm.integrations.otel import logger as otel_logger
|
||||
from litellm.integrations.otel.logger import publish_global_otel_v2_provider
|
||||
|
||||
monkeypatch.setattr(otel_logger, "_published_v2_provider", None)
|
||||
cfg = OpenTelemetryV2Config(exporter="in_memory")
|
||||
tp = providers.build_tracer_provider(cfg)
|
||||
preset_logger = OpenTelemetryV2(
|
||||
|
|
|
|||
|
|
@ -867,6 +867,45 @@ async def test_azure_sentinel_concurrent_threshold_sends_collapse_into_one_attem
|
|||
assert getattr(logger, queue_attr) == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES)
|
||||
async def test_azure_sentinel_batch_size_bounds_every_request_under_concurrent_events(
|
||||
queue_attr, send_method, build_payloads
|
||||
):
|
||||
"""Lowering batch_size is the documented way to stay under the ingestion cap, so no request may
|
||||
carry more than batch_size records even when events keep landing while a send is on the wire,
|
||||
and every one of those records still has to arrive exactly once."""
|
||||
logger = _build_logger(batch_size=5)
|
||||
records = build_payloads(40)
|
||||
|
||||
attempts = []
|
||||
first_send_started = asyncio.Event()
|
||||
release_first_send = asyncio.Event()
|
||||
|
||||
async def _on_ingest(data):
|
||||
attempts.append([record["id"] for record in json.loads(data.decode("utf-8"))])
|
||||
if len(attempts) == 1:
|
||||
first_send_started.set()
|
||||
await release_first_send.wait()
|
||||
return _accepted()
|
||||
|
||||
_install_ingestion(logger, _on_ingest)
|
||||
|
||||
sends = [asyncio.create_task(_log(logger, queue_attr, record)) for record in records]
|
||||
await asyncio.wait_for(first_send_started.wait(), timeout=10)
|
||||
|
||||
assert attempts == [[record["id"] for record in records[:5]]]
|
||||
assert getattr(logger, queue_attr) == records[5:]
|
||||
|
||||
release_first_send.set()
|
||||
await asyncio.wait_for(asyncio.gather(*sends), timeout=10)
|
||||
await logger.flush_queue()
|
||||
|
||||
assert max(len(attempt) for attempt in attempts) <= 5
|
||||
assert [record_id for attempt in attempts for record_id in attempt] == [record["id"] for record in records]
|
||||
assert getattr(logger, queue_attr) == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES)
|
||||
async def test_azure_sentinel_requeues_a_cancelled_send(
|
||||
|
|
|
|||
|
|
@ -195,3 +195,70 @@ def test_mlflow_stream_handler_uses_async_complete_response():
|
|||
is final_response
|
||||
)
|
||||
assert "abc123" not in mlflow_logger._stream_id_to_span
|
||||
|
||||
|
||||
def test_mlflow_stream_handler_pops_span_when_end_raises():
|
||||
modules = _mock_mlflow_modules()
|
||||
with patch.dict("sys.modules", modules):
|
||||
from litellm.integrations.mlflow import MlflowLogger
|
||||
|
||||
mlflow_logger = MlflowLogger()
|
||||
mlflow_logger._start_span_or_trace = MagicMock(return_value="mock_span")
|
||||
mlflow_logger._end_span_or_trace = MagicMock(
|
||||
side_effect=TypeError("unexpected keyword argument 'trace_id'")
|
||||
)
|
||||
mlflow_logger._extract_and_set_chat_attributes = MagicMock()
|
||||
|
||||
response_obj = MagicMock()
|
||||
response_obj.choices = []
|
||||
|
||||
kwargs = {
|
||||
"litellm_call_id": "leak123",
|
||||
"complete_streaming_response": MagicMock(),
|
||||
}
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
mlflow_logger._handle_stream_event(
|
||||
kwargs=kwargs,
|
||||
response_obj=response_obj,
|
||||
start_time=datetime.utcnow(),
|
||||
end_time=datetime.utcnow(),
|
||||
)
|
||||
|
||||
assert "leak123" not in mlflow_logger._stream_id_to_span
|
||||
|
||||
|
||||
class _Mlflow2StyleClient:
|
||||
"""Mimics the mlflow 2.x client signatures, which have no trace_id kwarg."""
|
||||
|
||||
def __init__(self):
|
||||
self.ended_traces = []
|
||||
self.ended_spans = []
|
||||
|
||||
def end_trace(self, request_id, outputs=None, attributes=None, status="OK", end_time_ns=None):
|
||||
self.ended_traces.append(request_id)
|
||||
|
||||
def end_span(self, request_id, span_id, outputs=None, attributes=None, status="OK", end_time_ns=None):
|
||||
self.ended_spans.append((request_id, span_id))
|
||||
|
||||
|
||||
def test_mlflow_end_span_or_trace_works_with_mlflow_2x_client():
|
||||
modules = _mock_mlflow_modules()
|
||||
with patch.dict("sys.modules", modules):
|
||||
from litellm.integrations.mlflow import MlflowLogger
|
||||
|
||||
mlflow_logger = MlflowLogger()
|
||||
client = _Mlflow2StyleClient()
|
||||
mlflow_logger._client = client
|
||||
|
||||
root_span = MagicMock(parent_id=None, request_id="req-1")
|
||||
mlflow_logger._end_span_or_trace(
|
||||
span=root_span, outputs="out", end_time_ns=1, status="OK"
|
||||
)
|
||||
assert client.ended_traces == ["req-1"]
|
||||
|
||||
child_span = MagicMock(parent_id="parent-1", request_id="req-2", span_id="span-2")
|
||||
mlflow_logger._end_span_or_trace(
|
||||
span=child_span, outputs="out", end_time_ns=1, status="OK"
|
||||
)
|
||||
assert client.ended_spans == [("req-2", "span-2")]
|
||||
|
|
|
|||
|
|
@ -2658,6 +2658,7 @@ def test_cache_writing_cost_with_zero_creation_tokens_and_ephemeral_details():
|
|||
"image_count": 0,
|
||||
"video_length_seconds": 0.0,
|
||||
"audio_length_seconds": 0.0,
|
||||
"query_count": 0,
|
||||
}
|
||||
|
||||
model_info: ModelInfo = {}
|
||||
|
|
@ -3239,6 +3240,37 @@ def test_image_count_prevents_text_tokens_fallback(_local_model_cost_map):
|
|||
assert completion_cost == 0.0
|
||||
|
||||
|
||||
def test_query_count_bills_input_cost_per_query(_local_model_cost_map):
|
||||
usage = Usage(
|
||||
prompt_tokens=0,
|
||||
completion_tokens=0,
|
||||
total_tokens=0,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(query_count=3, image_count=1),
|
||||
)
|
||||
|
||||
prompt_cost, completion_cost = generic_cost_per_token(
|
||||
model="us.twelvelabs.marengo-embed-3-0-v1:0",
|
||||
usage=usage,
|
||||
custom_llm_provider="bedrock",
|
||||
)
|
||||
|
||||
assert prompt_cost == pytest.approx(3 * 7e-05 + 1e-04)
|
||||
assert completion_cost == 0.0
|
||||
|
||||
|
||||
def test_query_count_is_free_without_a_per_query_price(_local_model_cost_map):
|
||||
usage = Usage(
|
||||
prompt_tokens=0,
|
||||
completion_tokens=0,
|
||||
total_tokens=0,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(query_count=1),
|
||||
)
|
||||
|
||||
prompt_cost, _ = generic_cost_per_token(model="text-embedding-3-small", usage=usage, custom_llm_provider="openai")
|
||||
|
||||
assert prompt_cost == 0.0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Data-residency (OpenAI regional processing) tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -6074,15 +6074,17 @@ def test_prompt_hooks_skip_prompt_managers_when_no_prompt_id(logging_obj, tmp_pa
|
|||
for hook in [cb for cb in litellm.callbacks if isinstance(cb, VectorStorePreCallHook)]:
|
||||
litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.callbacks, hook)
|
||||
def test_newrelic_dispatch_prefers_otel_v2_when_flag_on(monkeypatch):
|
||||
"""With LITELLM_OTEL_V2 on, the "newrelic" callback builds the OTel v2
|
||||
logger (per-team credential routing); with the flag off (default) it keeps
|
||||
the legacy agent-based logger, so existing deployments are untouched."""
|
||||
"""With LITELLM_OTEL_V2 on and operator credentials present, the "newrelic"
|
||||
callback builds the OTel v2 logger (per-team credential routing); with the
|
||||
flag off (default) it keeps the legacy agent-based logger, so existing
|
||||
deployments are untouched."""
|
||||
from litellm.integrations.otel.logger import OpenTelemetryV2
|
||||
from litellm.integrations.otel.model.config import is_otel_v2_enabled
|
||||
from litellm.litellm_core_utils import litellm_logging as logging_module
|
||||
|
||||
logging_module._in_memory_loggers.clear()
|
||||
monkeypatch.setenv("LITELLM_OTEL_V2", "true")
|
||||
monkeypatch.setenv("NEW_RELIC_LICENSE_KEY", "test-license-key")
|
||||
is_otel_v2_enabled.cache_clear()
|
||||
try:
|
||||
v2_logger = logging_module._init_custom_logger_compatible_class(
|
||||
|
|
@ -6137,6 +6139,7 @@ def test_get_custom_logger_compatible_class_finds_v2_newrelic(monkeypatch):
|
|||
|
||||
logging_module._in_memory_loggers.clear()
|
||||
monkeypatch.setenv("LITELLM_OTEL_V2", "true")
|
||||
monkeypatch.setenv("NEW_RELIC_LICENSE_KEY", "test-license-key")
|
||||
is_otel_v2_enabled.cache_clear()
|
||||
try:
|
||||
created = logging_module._init_custom_logger_compatible_class(
|
||||
|
|
|
|||
|
|
@ -9,7 +9,8 @@ import pytest
|
|||
import litellm
|
||||
from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
|
||||
from litellm.llms.anthropic.chat.handler import ModelResponseIterator, make_call
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm._uuid import uuid
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.types.llms.openai import (
|
||||
ChatCompletionToolCallChunk,
|
||||
ChatCompletionToolCallFunctionChunk,
|
||||
|
|
@ -85,6 +86,54 @@ def test_anthropic_completion_does_not_send_deployment_default_limits():
|
|||
assert "default_api_key_tpm_limit" not in request_body
|
||||
|
||||
|
||||
async def test_anthropic_async_completion_inlines_http_images_off_the_event_loop(async_only_image_fetch):
|
||||
http_image_url = f"http://img.example/{uuid.uuid4()}.png"
|
||||
https_image_url = f"https://img.example/{uuid.uuid4()}.png"
|
||||
captured = {}
|
||||
|
||||
def handle(request):
|
||||
captured["body"] = json.loads(request.content)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-4-6",
|
||||
"content": [{"type": "text", "text": "Green"}],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1},
|
||||
},
|
||||
)
|
||||
|
||||
client = AsyncHTTPHandler()
|
||||
client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
|
||||
|
||||
response = await litellm.acompletion(
|
||||
model="anthropic/claude-sonnet-4-6",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What colour is this?"},
|
||||
{"type": "image_url", "image_url": {"url": http_image_url}},
|
||||
{"type": "image_url", "image_url": {"url": https_image_url}},
|
||||
],
|
||||
}
|
||||
],
|
||||
api_key="test-key",
|
||||
client=client,
|
||||
)
|
||||
|
||||
assert response.choices[0].message.content == "Green"
|
||||
assert async_only_image_fetch.fetched == [http_image_url]
|
||||
sources = [part["source"] for part in captured["body"]["messages"][0]["content"] if part["type"] == "image"]
|
||||
assert sources == [
|
||||
{"type": "base64", "media_type": "image/png", "data": async_only_image_fetch.base64_png},
|
||||
{"type": "url", "url": https_image_url},
|
||||
]
|
||||
|
||||
|
||||
def test_redacted_thinking_content_block_delta():
|
||||
chunk = {
|
||||
"type": "content_block_start",
|
||||
|
|
|
|||
|
|
@ -184,6 +184,45 @@ class TestBedrockAsyncInvokeEmbedding:
|
|||
request_url = mock_post.call_args.kwargs.get("url", "")
|
||||
assert "/async-invoke" in request_url
|
||||
|
||||
def test_async_invoke_marengo_3_wraps_the_nested_payload_with_the_base_model_id(self):
|
||||
client = HTTPHandler()
|
||||
|
||||
with patch.object(client, "post") as mock_post:
|
||||
mock_response = Mock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.text = json.dumps(async_invoke_response)
|
||||
mock_response.json = lambda: json.loads(mock_response.text)
|
||||
mock_post.return_value = mock_response
|
||||
|
||||
response = litellm.embedding(
|
||||
model="bedrock/async_invoke/twelvelabs.marengo-embed-3-0-v1:0",
|
||||
input="s3://test-bucket/clip.mp4",
|
||||
client=client,
|
||||
aws_region_name="us-east-1",
|
||||
aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com",
|
||||
api_key="test-bearer-token-12345",
|
||||
input_type="video",
|
||||
embeddingOption=["visual", "audio"],
|
||||
segmentation={"method": "fixed", "fixed": {"durationSec": 6}},
|
||||
bucketOwner="123456789012",
|
||||
output_s3_uri="s3://test-bucket/async-invoke-output/",
|
||||
)
|
||||
|
||||
assert response._hidden_params._invocation_arn == async_invoke_response["invocationArn"]
|
||||
assert mock_post.call_args.kwargs["url"].endswith("/async-invoke")
|
||||
assert json.loads(mock_post.call_args.kwargs["data"]) == {
|
||||
"modelId": "twelvelabs.marengo-embed-3-0-v1:0",
|
||||
"modelInput": {
|
||||
"inputType": "video",
|
||||
"video": {
|
||||
"mediaSource": {"s3Location": {"uri": "s3://test-bucket/clip.mp4", "bucketOwner": "123456789012"}},
|
||||
"segmentation": {"method": "fixed", "fixed": {"durationSec": 6}},
|
||||
"embeddingOption": ["visual", "audio"],
|
||||
},
|
||||
},
|
||||
"outputDataConfig": {"s3OutputDataConfig": {"s3Uri": "s3://test-bucket/async-invoke-output/"}},
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_invoke_twelvelabs_embedding_async_with_mock(self):
|
||||
"""Test async invoke embedding with async calls."""
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ from unittest.mock import Mock, patch
|
|||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.llms.bedrock.embed.twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
|
||||
# Mock responses for different embedding models
|
||||
|
|
@ -1059,3 +1060,182 @@ def test_bedrock_embedding_bearer_token_never_runs_the_sigv4_credential_chain(mo
|
|||
|
||||
assert response.data[0]["embedding"] == titan_embedding_response["embedding"]
|
||||
assert mock_post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345"
|
||||
|
||||
|
||||
marengo_3_embedding_response = {"data": [{"embedding": [0.01 * i for i in range(512)]}]}
|
||||
MARENGO_3_DUCK = "data:image/png;base64,ZHVjaw=="
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model,kwargs,expected_body,expected_usage_details",
|
||||
[
|
||||
(
|
||||
"bedrock/us.twelvelabs.marengo-embed-3-0-v1:0",
|
||||
{"input_type": "text"},
|
||||
{"inputType": "text", "text": {"inputText": "a duck on water"}},
|
||||
{"query_count": 1},
|
||||
),
|
||||
(
|
||||
"bedrock/twelvelabs.marengo-embed-3-0-v1:0",
|
||||
{"input_type": "text"},
|
||||
{"inputType": "text", "text": {"inputText": "a duck on water"}},
|
||||
{"query_count": 1},
|
||||
),
|
||||
(
|
||||
"bedrock/us.twelvelabs.marengo-embed-3-0-v1:0",
|
||||
{"input_type": "text_image", "media_source": MARENGO_3_DUCK},
|
||||
{
|
||||
"inputType": "text_image",
|
||||
"text_image": {"inputText": "a duck on water", "mediaSource": {"base64String": "ZHVjaw=="}},
|
||||
},
|
||||
{"query_count": 1, "image_count": 1},
|
||||
),
|
||||
(
|
||||
"bedrock/us.twelvelabs.marengo-embed-3-0-v1:0",
|
||||
{"input_type": "multi_input", "media_sources": {"bird": MARENGO_3_DUCK}},
|
||||
{
|
||||
"inputType": "multi_input",
|
||||
"multi_input": {
|
||||
"inputText": "a duck on water",
|
||||
"mediaSources": [{"name": "bird", "mediaType": "image", "base64String": "ZHVjaw=="}],
|
||||
},
|
||||
},
|
||||
{"query_count": 1, "image_count": 1},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_marengo_3_embedding_sends_the_nested_payload_and_parses_512_dims(
|
||||
model, kwargs, expected_body, expected_usage_details
|
||||
):
|
||||
client = HTTPHandler()
|
||||
|
||||
with patch.object(client, "post") as mock_post:
|
||||
mock_response = Mock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.text = json.dumps(marengo_3_embedding_response)
|
||||
mock_response.json = lambda: json.loads(mock_response.text)
|
||||
mock_post.return_value = mock_response
|
||||
|
||||
response = litellm.embedding(
|
||||
model=model,
|
||||
input="a duck on water",
|
||||
client=client,
|
||||
aws_region_name="us-east-1",
|
||||
api_key="test-bearer-token-12345",
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
assert json.loads(mock_post.call_args.kwargs["data"]) == expected_body
|
||||
assert mock_post.call_args.kwargs["url"].endswith(f"/model/{model.removeprefix('bedrock/').replace(':', '%3A')}/invoke")
|
||||
assert len(response.data[0]["embedding"]) == 512
|
||||
assert response.data[0]["embedding"][:2] == [0.0, 0.01]
|
||||
assert response.usage.prompt_tokens == 0
|
||||
assert response.usage.total_tokens == 0
|
||||
assert response.usage.prompt_tokens_details.model_dump(exclude_none=True) == expected_usage_details
|
||||
|
||||
|
||||
def test_marengo_3_image_embedding_sends_the_media_under_the_image_key():
|
||||
client = HTTPHandler()
|
||||
|
||||
with patch.object(client, "post") as mock_post:
|
||||
mock_response = Mock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.text = json.dumps(marengo_3_embedding_response)
|
||||
mock_response.json = lambda: json.loads(mock_response.text)
|
||||
mock_post.return_value = mock_response
|
||||
|
||||
response = litellm.embedding(
|
||||
model="bedrock/us.twelvelabs.marengo-embed-3-0-v1:0",
|
||||
input=MARENGO_3_DUCK,
|
||||
client=client,
|
||||
aws_region_name="us-east-1",
|
||||
api_key="test-bearer-token-12345",
|
||||
input_type="image",
|
||||
)
|
||||
|
||||
assert json.loads(mock_post.call_args.kwargs["data"]) == {
|
||||
"inputType": "image",
|
||||
"image": {"mediaSource": {"base64String": "ZHVjaw=="}},
|
||||
}
|
||||
assert len(response.data[0]["embedding"]) == 512
|
||||
assert response.data[0]["embedding"][:2] == [0.0, 0.01]
|
||||
assert response.usage.prompt_tokens == 0
|
||||
assert response.usage.prompt_tokens_details.model_dump(exclude_none=True) == {"image_count": 1}
|
||||
|
||||
|
||||
def test_marengo_2_7_embedding_keeps_the_flat_payload():
|
||||
client = HTTPHandler()
|
||||
|
||||
with patch.object(client, "post") as mock_post:
|
||||
mock_response = Mock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.text = json.dumps(twelvelabs_embedding_response)
|
||||
mock_response.json = lambda: json.loads(mock_response.text)
|
||||
mock_post.return_value = mock_response
|
||||
|
||||
response = litellm.embedding(
|
||||
model="bedrock/us.twelvelabs.marengo-embed-2-7-v1:0",
|
||||
input="a duck on water",
|
||||
client=client,
|
||||
aws_region_name="us-east-1",
|
||||
api_key="test-bearer-token-12345",
|
||||
input_type="text",
|
||||
)
|
||||
|
||||
assert json.loads(mock_post.call_args.kwargs["data"]) == {
|
||||
"inputType": "text",
|
||||
"inputText": "a duck on water",
|
||||
"textTruncate": "end",
|
||||
}
|
||||
assert response.data[0]["embedding"] == [0.1, 0.2, 0.3]
|
||||
assert response.usage.prompt_tokens == 0
|
||||
assert response.usage.prompt_tokens_details.model_dump(exclude_none=True) == {"query_count": 1}
|
||||
|
||||
|
||||
def test_marengo_usage_counts_text_requests_and_images_across_a_batch():
|
||||
duck = {"mediaType": "image", "base64String": "ZHVjaw=="}
|
||||
response = TwelveLabsMarengoEmbeddingConfig()._transform_response(
|
||||
response_list=[marengo_3_embedding_response, marengo_3_embedding_response, marengo_3_embedding_response],
|
||||
model="us.twelvelabs.marengo-embed-3-0-v1:0",
|
||||
batch_data=[
|
||||
{"inputType": "text", "text": {"inputText": "a duck"}},
|
||||
{"inputType": "image", "image": {"mediaSource": {"base64String": "ZHVjaw=="}}},
|
||||
{"inputType": "multi_input", "multi_input": {"mediaSources": [{"name": "a", **duck}, {"name": "b", **duck}]}},
|
||||
],
|
||||
)
|
||||
|
||||
assert [item["index"] for item in response.data] == [0, 1, 2]
|
||||
assert response.usage.prompt_tokens == 0
|
||||
assert response.usage.total_tokens == 0
|
||||
assert response.usage.prompt_tokens_details.model_dump(exclude_none=True) == {"query_count": 1, "image_count": 3}
|
||||
|
||||
|
||||
def test_marengo_usage_without_request_data_bills_nothing():
|
||||
response = TwelveLabsMarengoEmbeddingConfig()._transform_response(
|
||||
response_list=[marengo_3_embedding_response], model="us.twelvelabs.marengo-embed-3-0-v1:0"
|
||||
)
|
||||
|
||||
assert len(response.data[0]["embedding"]) == 512
|
||||
assert response.usage.prompt_tokens == 0
|
||||
assert response.usage.prompt_tokens_details is None
|
||||
|
||||
|
||||
def test_marengo_response_items_without_an_embedding_are_skipped():
|
||||
response = TwelveLabsMarengoEmbeddingConfig()._transform_response(
|
||||
response_list=[{"data": [{"embeddingOption": "visual-text", "startSec": 0.0}, {"embedding": [0.1, 0.2, 0.3]}]}],
|
||||
model="us.twelvelabs.marengo-embed-3-0-v1:0",
|
||||
)
|
||||
|
||||
assert [item["embedding"] for item in response.data] == [[0.1, 0.2, 0.3]]
|
||||
assert response.data[0]["index"] == 0
|
||||
|
||||
|
||||
def test_marengo_3_text_image_without_media_source_is_a_bad_request():
|
||||
with pytest.raises(litellm.BadRequestError, match=r"text_image.*media_source"):
|
||||
litellm.embedding(
|
||||
model="bedrock/us.twelvelabs.marengo-embed-3-0-v1:0",
|
||||
input="a duck on water",
|
||||
aws_region_name="us-east-1",
|
||||
api_key="test-bearer-token-12345",
|
||||
input_type="text_image",
|
||||
)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,416 @@
|
|||
import json
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.llms.bedrock.common_utils import BedrockError
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.llms.bedrock.embed.twelvelabs_marengo_3_transformation import (
|
||||
MARENGO_2_7_ONLY_PARAMS,
|
||||
build_marengo_3_request,
|
||||
is_marengo_3_model,
|
||||
)
|
||||
from litellm.llms.bedrock.embed.twelvelabs_marengo_transformation import (
|
||||
TwelveLabsMarengoEmbeddingConfig,
|
||||
drop_params_enabled,
|
||||
)
|
||||
|
||||
MARENGO_3_BASE = "twelvelabs.marengo-embed-3-0-v1:0"
|
||||
MARENGO_3_US = "us.twelvelabs.marengo-embed-3-0-v1:0"
|
||||
MARENGO_27_US = "us.twelvelabs.marengo-embed-2-7-v1:0"
|
||||
DUCK_DATA_URL = "data:image/png;base64,ZHVjaw=="
|
||||
OUTPUT_S3_URI = "s3://out-bucket/marengo/"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model,expected",
|
||||
[
|
||||
(MARENGO_3_BASE, True),
|
||||
(MARENGO_3_US, True),
|
||||
("eu.twelvelabs.marengo-embed-3-0-v1:0", True),
|
||||
("async_invoke/twelvelabs.marengo-embed-3-0-v1:0", True),
|
||||
(MARENGO_27_US, False),
|
||||
("twelvelabs.marengo-embed-2-7-v1:0", False),
|
||||
("twelvelabs.marengo-embed-30-v1:0", False),
|
||||
(None, False),
|
||||
],
|
||||
)
|
||||
def test_is_marengo_3_model(model, expected):
|
||||
assert is_marengo_3_model(model) is expected
|
||||
|
||||
|
||||
def wire(request: object) -> object:
|
||||
return json.loads(json.dumps(request))
|
||||
|
||||
|
||||
def test_text_request_nests_input_text_under_text():
|
||||
assert build_marengo_3_request("a dog on the beach", {"input_type": "text"}) == {
|
||||
"inputType": "text",
|
||||
"text": {"inputText": "a dog on the beach"},
|
||||
}
|
||||
|
||||
|
||||
def test_missing_input_type_defaults_to_text():
|
||||
assert build_marengo_3_request("hello", {})["inputType"] == "text"
|
||||
|
||||
|
||||
def test_camel_case_input_type_wins_over_snake_case():
|
||||
request = build_marengo_3_request(DUCK_DATA_URL, {"inputType": "image", "input_type": "text"})
|
||||
assert request["inputType"] == "image"
|
||||
|
||||
|
||||
def test_image_request_strips_data_url_prefix():
|
||||
assert build_marengo_3_request(DUCK_DATA_URL, {"input_type": "image"}) == {
|
||||
"inputType": "image",
|
||||
"image": {"mediaSource": {"base64String": "ZHVjaw=="}},
|
||||
}
|
||||
|
||||
|
||||
def test_image_request_from_s3_carries_bucket_owner():
|
||||
request = build_marengo_3_request("s3://media/duck.png", {"input_type": "image", "bucketOwner": "123456789012"})
|
||||
assert request == {
|
||||
"inputType": "image",
|
||||
"image": {"mediaSource": {"s3Location": {"uri": "s3://media/duck.png", "bucketOwner": "123456789012"}}},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"input_media,params",
|
||||
[
|
||||
("s3://media/duck.png", {"input_type": "image"}),
|
||||
("s3://media/clip.mp4", {"input_type": "video"}),
|
||||
("a duck", {"input_type": "text_image", "media_source": "s3://media/duck.png"}),
|
||||
("a duck", {"input_type": "multi_input", "media_sources": {"img1": "s3://media/duck.png"}}),
|
||||
],
|
||||
)
|
||||
def test_s3_media_without_bucket_owner_is_rejected_naming_it(input_media, params):
|
||||
with pytest.raises(BedrockError) as excinfo:
|
||||
build_marengo_3_request(input_media, params)
|
||||
assert excinfo.value.status_code == 400
|
||||
assert excinfo.value.message == (
|
||||
"s3:// media requires the 'bucketOwner' parameter, the account id that owns the bucket"
|
||||
)
|
||||
|
||||
|
||||
def test_text_image_request_pairs_text_with_media_source():
|
||||
request = build_marengo_3_request(
|
||||
"a duck", {"input_type": "text_image", "media_source": DUCK_DATA_URL, "output_s3_uri": OUTPUT_S3_URI}
|
||||
)
|
||||
assert request == {
|
||||
"inputType": "text_image",
|
||||
"text_image": {"inputText": "a duck", "mediaSource": {"base64String": "ZHVjaw=="}},
|
||||
}
|
||||
|
||||
|
||||
def test_text_image_request_requires_media_source():
|
||||
with pytest.raises(BedrockError, match=r"text_image.*media_source") as excinfo:
|
||||
build_marengo_3_request("a duck", {"input_type": "text_image"})
|
||||
assert excinfo.value.status_code == 400
|
||||
|
||||
|
||||
def test_multi_input_request_names_each_media_source():
|
||||
request = build_marengo_3_request(
|
||||
"a photo of <@bird> next to <@dog>",
|
||||
{
|
||||
"input_type": "multi_input",
|
||||
"media_sources": {"bird": DUCK_DATA_URL, "dog": "s3://media/dog.png"},
|
||||
"bucketOwner": "123456789012",
|
||||
},
|
||||
)
|
||||
assert wire(request) == {
|
||||
"inputType": "multi_input",
|
||||
"multi_input": {
|
||||
"inputText": "a photo of <@bird> next to <@dog>",
|
||||
"mediaSources": [
|
||||
{"name": "bird", "mediaType": "image", "base64String": "ZHVjaw=="},
|
||||
{
|
||||
"name": "dog",
|
||||
"mediaType": "image",
|
||||
"s3Location": {"uri": "s3://media/dog.png", "bucketOwner": "123456789012"},
|
||||
},
|
||||
],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_multi_input_without_text_omits_input_text():
|
||||
request = build_marengo_3_request("", {"input_type": "multi_input", "media_sources": {"bird": DUCK_DATA_URL}})
|
||||
assert "inputText" not in request["multi_input"]
|
||||
assert request["multi_input"]["mediaSources"][0]["name"] == "bird"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("params", [{"input_type": "multi_input"}, {"input_type": "multi_input", "media_sources": {}}])
|
||||
def test_multi_input_request_requires_media_sources(params):
|
||||
with pytest.raises(BedrockError, match=r"multi_input.*media_sources") as excinfo:
|
||||
build_marengo_3_request("<@bird>", params)
|
||||
assert excinfo.value.status_code == 400
|
||||
|
||||
|
||||
@pytest.mark.parametrize("input_type", ["video", "audio"])
|
||||
def test_timed_media_request_nests_every_option_under_the_media_key(input_type):
|
||||
request = build_marengo_3_request(
|
||||
"s3://media/clip.mp4",
|
||||
{
|
||||
"input_type": input_type,
|
||||
"startSec": 2,
|
||||
"endSec": 12.5,
|
||||
"segmentation": {"method": "dynamic", "dynamic": {"minDurationSec": 4}},
|
||||
"embeddingOption": ["visual", "audio"],
|
||||
"embeddingType": ["fused_embedding"],
|
||||
"embeddingScope": ["clip", "asset"],
|
||||
"inferenceId": "req-42",
|
||||
"bucketOwner": "123456789012",
|
||||
},
|
||||
)
|
||||
assert wire(request) == {
|
||||
"inputType": input_type,
|
||||
input_type: {
|
||||
"mediaSource": {"s3Location": {"uri": "s3://media/clip.mp4", "bucketOwner": "123456789012"}},
|
||||
"startSec": 2.0,
|
||||
"endSec": 12.5,
|
||||
"segmentation": {"method": "dynamic", "dynamic": {"minDurationSec": 4}},
|
||||
"embeddingOption": ["visual", "audio"],
|
||||
"embeddingType": ["fused_embedding"],
|
||||
"embeddingScope": ["clip", "asset"],
|
||||
},
|
||||
"inferenceId": "req-42",
|
||||
}
|
||||
|
||||
|
||||
def test_timed_media_request_without_options_carries_only_the_media_source():
|
||||
request = build_marengo_3_request("s3://media/clip.mp4", {"input_type": "video", "bucketOwner": "123456789012"})
|
||||
assert request["video"] == {
|
||||
"mediaSource": {"s3Location": {"uri": "s3://media/clip.mp4", "bucketOwner": "123456789012"}}
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"params",
|
||||
[
|
||||
{"input_type": "clip"},
|
||||
{"input_type": "video", "embeddingOption": ["visual-text"]},
|
||||
{"input_type": "video", "segmentation": {"method": "fixed", "dynamic": {"minDurationSec": 4}}},
|
||||
{"input_type": "multi_input", "media_sources": ["not", "a", "mapping"]},
|
||||
],
|
||||
)
|
||||
def test_invalid_marengo_3_params_are_rejected_before_the_request_is_sent(params):
|
||||
with pytest.raises(BedrockError, match=r"Invalid Marengo 3\.0 parameters") as excinfo:
|
||||
build_marengo_3_request("s3://media/clip.mp4", params)
|
||||
assert excinfo.value.status_code == 400
|
||||
|
||||
|
||||
def test_config_sends_the_nested_payload_for_marengo_3_and_the_flat_one_for_2_7():
|
||||
nested = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_US)._transform_request(
|
||||
input="hello", inference_params={"input_type": "text"}
|
||||
)
|
||||
flat = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_27_US)._transform_request(
|
||||
input="hello", inference_params={"input_type": "text"}
|
||||
)
|
||||
assert nested == {"inputType": "text", "text": {"inputText": "hello"}}
|
||||
assert flat == {"inputType": "text", "inputText": "hello", "textTruncate": "end"}
|
||||
|
||||
|
||||
def test_config_without_a_model_keeps_the_2_7_payload():
|
||||
request = TwelveLabsMarengoEmbeddingConfig()._transform_request(input="hello", inference_params={})
|
||||
assert request == {"inputType": "text", "inputText": "hello", "textTruncate": "end"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("input_type", ["video", "audio"])
|
||||
def test_marengo_3_video_and_audio_still_require_the_async_route(input_type):
|
||||
with pytest.raises(ValueError, match=f"Input type '{input_type}' requires async_invoke route"):
|
||||
TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_BASE)._transform_request(
|
||||
input="s3://media/clip.mp4", inference_params={"input_type": input_type}
|
||||
)
|
||||
|
||||
|
||||
def test_marengo_3_async_invoke_wraps_the_nested_payload_with_the_base_model_id():
|
||||
request = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_BASE)._transform_request(
|
||||
input="s3://media/clip.mp4",
|
||||
inference_params={
|
||||
"input_type": "video",
|
||||
"embeddingOption": ["visual"],
|
||||
"bucketOwner": "123456789012",
|
||||
"output_s3_uri": OUTPUT_S3_URI,
|
||||
},
|
||||
async_invoke_route=True,
|
||||
model_id="async_invoke%2Ftwelvelabs.marengo-embed-3-0-v1%3A0",
|
||||
output_s3_uri=OUTPUT_S3_URI,
|
||||
)
|
||||
assert wire(request) == {
|
||||
"modelId": MARENGO_3_BASE,
|
||||
"modelInput": {
|
||||
"inputType": "video",
|
||||
"video": {
|
||||
"mediaSource": {"s3Location": {"uri": "s3://media/clip.mp4", "bucketOwner": "123456789012"}},
|
||||
"embeddingOption": ["visual"],
|
||||
},
|
||||
},
|
||||
"outputDataConfig": {"s3OutputDataConfig": {"s3Uri": OUTPUT_S3_URI}},
|
||||
}
|
||||
|
||||
|
||||
def test_marengo_3_async_invoke_requires_an_output_s3_uri():
|
||||
with pytest.raises(ValueError, match="output_s3_uri cannot be empty"):
|
||||
TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_BASE)._transform_request(
|
||||
input="hello",
|
||||
inference_params={"input_type": "text"},
|
||||
async_invoke_route=True,
|
||||
model_id=MARENGO_3_BASE,
|
||||
output_s3_uri="",
|
||||
)
|
||||
|
||||
|
||||
def test_encoding_format_float_no_longer_injects_2_7_embedding_options_for_marengo_3():
|
||||
marengo_3 = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_US).map_openai_params(
|
||||
non_default_params={"encoding_format": "float"}, optional_params={}
|
||||
)
|
||||
marengo_27 = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_27_US).map_openai_params(
|
||||
non_default_params={"encoding_format": "float"}, optional_params={}
|
||||
)
|
||||
assert marengo_3 == {}
|
||||
assert marengo_27 == {"embeddingOption": ["visual-text", "visual-image"]}
|
||||
|
||||
|
||||
def test_marengo_3_only_params_are_forwarded_by_map_openai_params():
|
||||
mapped = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_US).map_openai_params(
|
||||
non_default_params={
|
||||
"input_type": "text_image",
|
||||
"media_source": DUCK_DATA_URL,
|
||||
"media_sources": {"bird": DUCK_DATA_URL},
|
||||
"endSec": 5,
|
||||
"segmentation": {"method": "fixed", "fixed": {"durationSec": 6}},
|
||||
"embeddingType": ["separate_embedding"],
|
||||
"embeddingScope": ["clip"],
|
||||
"inferenceId": "req-1",
|
||||
},
|
||||
optional_params={},
|
||||
)
|
||||
assert mapped == {
|
||||
"inputType": "text_image",
|
||||
"media_source": DUCK_DATA_URL,
|
||||
"media_sources": {"bird": DUCK_DATA_URL},
|
||||
"endSec": 5,
|
||||
"segmentation": {"method": "fixed", "fixed": {"durationSec": 6}},
|
||||
"embeddingType": ["separate_embedding"],
|
||||
"embeddingScope": ["clip"],
|
||||
"inferenceId": "req-1",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"params,problem",
|
||||
[
|
||||
(
|
||||
{"input_type": "clip"},
|
||||
"input_type: Input should be 'text', 'image', 'video', 'audio', 'text_image' or 'multi_input'",
|
||||
),
|
||||
({"input_type": "video", "embeddingOption": "visual"}, "embeddingOption: Input should be a valid tuple"),
|
||||
(
|
||||
{"input_type": "multi_input", "media_sources": ["not", "a", "mapping"]},
|
||||
"media_sources: Input should be a valid dictionary",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_invalid_marengo_3_params_name_the_field_and_the_reason(params, problem):
|
||||
with pytest.raises(BedrockError) as excinfo:
|
||||
build_marengo_3_request("s3://media/clip.mp4", params)
|
||||
assert excinfo.value.message == f"Invalid Marengo 3.0 parameters: {problem}"
|
||||
|
||||
|
||||
MARENGO_2_7_ONLY_VALUES = {"textTruncate": "end", "lengthSec": 5, "useFixedLengthSec": True, "minClipSec": 2}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name", MARENGO_2_7_ONLY_PARAMS)
|
||||
def test_marengo_2_7_only_params_are_rejected_on_3_0_unless_dropped(name):
|
||||
params = {"input_type": "text", name: MARENGO_2_7_ONLY_VALUES[name]}
|
||||
with pytest.raises(BedrockError) as excinfo:
|
||||
build_marengo_3_request("hello", params)
|
||||
assert excinfo.value.status_code == 400
|
||||
assert excinfo.value.message == (
|
||||
f"Marengo 3.0 does not accept the Marengo 2.7 parameters {name}; set drop_params to drop them"
|
||||
)
|
||||
assert build_marengo_3_request("hello", params, drop_params=True) == {
|
||||
"inputType": "text",
|
||||
"text": {"inputText": "hello"},
|
||||
}
|
||||
|
||||
|
||||
def test_marengo_2_7_only_params_are_advertised_only_for_2_7():
|
||||
marengo_3 = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_US).get_supported_openai_params()
|
||||
marengo_27 = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_27_US).get_supported_openai_params()
|
||||
assert set(MARENGO_2_7_ONLY_PARAMS).isdisjoint(marengo_3)
|
||||
assert set(MARENGO_2_7_ONLY_PARAMS) <= set(marengo_27)
|
||||
assert set(marengo_3) <= set(marengo_27)
|
||||
|
||||
|
||||
def test_drop_params_comes_from_the_call_or_the_global(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "drop_params", False)
|
||||
assert drop_params_enabled({}) is False
|
||||
assert drop_params_enabled({"drop_params": True}) is True
|
||||
monkeypatch.setattr(litellm, "drop_params", True)
|
||||
assert drop_params_enabled({}) is True
|
||||
|
||||
|
||||
def test_config_drops_marengo_2_7_only_params_only_when_asked():
|
||||
config = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_US)
|
||||
with pytest.raises(BedrockError, match=r"Marengo 2\.7 parameters textTruncate"):
|
||||
config._transform_request("hello", {"textTruncate": "end"})
|
||||
assert config._transform_request("hello", {"textTruncate": "end"}, drop_params=True) == {
|
||||
"inputType": "text",
|
||||
"text": {"inputText": "hello"},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"params",
|
||||
[
|
||||
{"input_type": "text"},
|
||||
{"input_type": "image"},
|
||||
{"input_type": "text_image", "media_source": DUCK_DATA_URL},
|
||||
{"input_type": "multi_input", "media_sources": {"bird": DUCK_DATA_URL}},
|
||||
],
|
||||
)
|
||||
def test_timed_media_options_are_rejected_on_untimed_input_types_unless_dropped(params):
|
||||
timed = {**params, "startSec": 0, "embeddingOption": ["visual"]}
|
||||
with pytest.raises(BedrockError) as excinfo:
|
||||
build_marengo_3_request(DUCK_DATA_URL, timed)
|
||||
assert excinfo.value.status_code == 400
|
||||
assert excinfo.value.message == (
|
||||
f"Input type '{params['input_type']}' does not accept startSec, embeddingOption; set drop_params to drop them"
|
||||
)
|
||||
assert build_marengo_3_request(DUCK_DATA_URL, timed, drop_params=True) == build_marengo_3_request(
|
||||
DUCK_DATA_URL, params
|
||||
)
|
||||
|
||||
|
||||
def _embed_marengo_3_us(client: HTTPHandler, **params: object):
|
||||
return litellm.embedding(
|
||||
model=f"bedrock/{MARENGO_3_US}",
|
||||
input="hello",
|
||||
client=client,
|
||||
aws_region_name="us-east-1",
|
||||
aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com",
|
||||
api_key="test-bearer-token",
|
||||
**params,
|
||||
)
|
||||
|
||||
|
||||
def test_per_request_drop_params_reaches_the_marengo_3_builder(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "drop_params", False)
|
||||
client = HTTPHandler()
|
||||
with patch.object(client, "post") as mock_post:
|
||||
mock_response = Mock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.text = json.dumps({"data": [{"embedding": [0.1, 0.2]}]})
|
||||
mock_response.json = lambda: json.loads(mock_response.text)
|
||||
mock_post.return_value = mock_response
|
||||
|
||||
with pytest.raises(litellm.BadRequestError, match=r"Marengo 2\.7 parameters textTruncate"):
|
||||
_embed_marengo_3_us(client, textTruncate="end")
|
||||
assert mock_post.call_count == 0
|
||||
|
||||
response = _embed_marengo_3_us(client, textTruncate="end", drop_params=True)
|
||||
|
||||
assert response.data[0]["embedding"] == [0.1, 0.2]
|
||||
assert json.loads(mock_post.call_args.kwargs["data"]) == {"inputType": "text", "text": {"inputText": "hello"}}
|
||||
File diff suppressed because it is too large
Load diff
|
|
@ -2,9 +2,11 @@ import json
|
|||
from litellm._uuid import uuid
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
|
||||
import litellm
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.llms.ollama.completion.transformation import (
|
||||
OllamaConfig,
|
||||
OllamaTextCompletionResponseIterator,
|
||||
|
|
@ -502,3 +504,43 @@ class TestOllamaTextCompletionResponseIterator:
|
|||
assert result["usage"]["prompt_tokens"] == 10
|
||||
assert result["usage"]["completion_tokens"] == 5
|
||||
assert result["usage"]["total_tokens"] == 15
|
||||
|
||||
|
||||
async def test_ollama_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch):
|
||||
image_url = f"https://img.example/{uuid.uuid4()}.png"
|
||||
captured = {}
|
||||
|
||||
def handle(request):
|
||||
captured["body"] = json.loads(request.content)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"model": "llava",
|
||||
"response": "Green",
|
||||
"done": True,
|
||||
"prompt_eval_count": 1,
|
||||
"eval_count": 1,
|
||||
},
|
||||
)
|
||||
|
||||
client = AsyncHTTPHandler()
|
||||
client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
|
||||
|
||||
response = await litellm.acompletion(
|
||||
model="ollama/llava",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What colour is this?"},
|
||||
{"type": "image_url", "image_url": {"url": image_url}},
|
||||
],
|
||||
}
|
||||
],
|
||||
api_base="http://ollama.example:11434",
|
||||
client=client,
|
||||
)
|
||||
|
||||
assert response.choices[0].message.content == "Green"
|
||||
assert async_only_image_fetch.fetched == [image_url]
|
||||
assert captured["body"]["images"] == [async_only_image_fetch.base64_png]
|
||||
|
|
|
|||
|
|
@ -6,11 +6,15 @@ Vertex AI Anthropic models don't support URL sources for images.
|
|||
LiteLLM should convert image URLs to base64 when using Vertex AI Anthropic.
|
||||
"""
|
||||
|
||||
import json
|
||||
import sys
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
|
||||
import litellm
|
||||
from litellm._uuid import uuid
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
anthropic_messages_pt,
|
||||
convert_to_anthropic_tool_result,
|
||||
|
|
@ -371,3 +375,59 @@ class TestToolMessageImageURLHandling:
|
|||
assert item["source"]["type"] == "url"
|
||||
return
|
||||
pytest.fail("Could not find image in tool result")
|
||||
|
||||
|
||||
async def test_vertex_ai_anthropic_async_completion_inlines_https_images_off_the_event_loop(async_only_image_fetch):
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
|
||||
image_url = f"https://img.example/{uuid.uuid4()}.png"
|
||||
captured = {}
|
||||
|
||||
def handle(request):
|
||||
captured["body"] = json.loads(request.content)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-4-6",
|
||||
"content": [{"type": "text", "text": "Green"}],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1},
|
||||
},
|
||||
)
|
||||
|
||||
client = AsyncHTTPHandler()
|
||||
client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
|
||||
vertexai = MagicMock()
|
||||
vertexai.preview.language_models = MagicMock()
|
||||
|
||||
with (
|
||||
patch.dict(sys.modules, {"vertexai": vertexai}),
|
||||
patch.object( # test-quality-ok: litellm.acompletion has no seam for Vertex token minting
|
||||
litellm.main.vertex_partner_models_chat_completion,
|
||||
"_ensure_access_token",
|
||||
return_value=("token", "test-project"),
|
||||
),
|
||||
):
|
||||
response = await litellm.acompletion(
|
||||
model="vertex_ai/claude-sonnet-4-6",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What colour is this?"},
|
||||
{"type": "image_url", "image_url": {"url": image_url}},
|
||||
],
|
||||
}
|
||||
],
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-east5",
|
||||
client=client,
|
||||
)
|
||||
|
||||
assert response.choices[0].message.content == "Green"
|
||||
assert async_only_image_fetch.fetched == [image_url]
|
||||
sources = [part["source"] for part in captured["body"]["messages"][0]["content"] if part["type"] == "image"]
|
||||
assert sources == [{"type": "base64", "media_type": "image/png", "data": async_only_image_fetch.base64_png}]
|
||||
|
|
|
|||
|
|
@ -356,6 +356,74 @@ async def test_watsonx_gpt_oss_uses_async_http_handler():
|
|||
assert result["status"] == "success", "Should return success status"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tokenizer_config_cached", [False, True], ids=["tokenizer_config", "cached_config_jinja"])
|
||||
async def test_watsonx_text_gpt_oss_async_completion_fetches_hf_template_off_the_event_loop(
|
||||
monkeypatch, tokenizer_config_cached
|
||||
):
|
||||
import httpx
|
||||
|
||||
from litellm._uuid import uuid
|
||||
from litellm.litellm_core_utils.prompt_templates import huggingface_template_handler
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
|
||||
hf_model = f"openai/gpt-oss-{uuid.uuid4()}"
|
||||
chat_template = "{% for m in messages %}<|{{ m['role'] }}|>{{ m['content'] }}{% endfor %}"
|
||||
if tokenizer_config_cached:
|
||||
cached_config = {"status": "success", "tokenizer": {"bos_token": None, "eos_token": None}}
|
||||
monkeypatch.setattr(litellm, "known_tokenizer_config", {hf_model: cached_config})
|
||||
expected_fetch = f"https://huggingface.co/{hf_model}/raw/main/chat_template.jinja"
|
||||
else:
|
||||
monkeypatch.setattr(litellm, "known_tokenizer_config", {})
|
||||
expected_fetch = f"https://huggingface.co/{hf_model}/raw/main/tokenizer_config.json"
|
||||
hf_fetched = []
|
||||
captured = {}
|
||||
|
||||
def forbid_sync_client():
|
||||
raise AssertionError("sync HuggingFace fetch ran on the request path")
|
||||
|
||||
async def serve_hf_file(url, **kwargs):
|
||||
hf_fetched.append(url)
|
||||
if url.endswith(".jinja"):
|
||||
return httpx.Response(200, content=chat_template.encode())
|
||||
return httpx.Response(200, json={"chat_template": chat_template, "bos_token": None, "eos_token": None})
|
||||
|
||||
monkeypatch.setattr(huggingface_template_handler, "_get_httpx_client", forbid_sync_client)
|
||||
monkeypatch.setattr(huggingface_template_handler, "get_async_httpx_client", lambda **kwargs: Mock(get=serve_hf_file))
|
||||
|
||||
def handle(request):
|
||||
captured["body"] = json.loads(request.content)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"model_id": hf_model,
|
||||
"results": [
|
||||
{
|
||||
"generated_text": "Hi",
|
||||
"generated_token_count": 1,
|
||||
"input_token_count": 1,
|
||||
"stop_reason": "eos_token",
|
||||
}
|
||||
],
|
||||
},
|
||||
)
|
||||
|
||||
client = AsyncHTTPHandler()
|
||||
client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
|
||||
|
||||
response = await litellm.acompletion(
|
||||
model=f"watsonx_text/{hf_model}",
|
||||
messages=[{"role": "user", "content": "Hi there"}],
|
||||
api_base="https://test-api.watsonx.ai",
|
||||
project_id="test-project-id",
|
||||
token="test-token",
|
||||
client=client,
|
||||
)
|
||||
|
||||
assert response.choices[0].message.content == "Hi"
|
||||
assert hf_fetched == [expected_fetch]
|
||||
assert captured["body"]["input"] == "<|user|>Hi there"
|
||||
|
||||
|
||||
def test_watsonx_chat_completion_with_reasoning_effort(monkeypatch):
|
||||
"""
|
||||
Test that 'reasoning_effort' is correctly passed through to the WatsonX API payload.
|
||||
|
|
|
|||
|
|
@ -7638,6 +7638,75 @@ class TestMCPMetaTraceCarrier:
|
|||
assert _mcp_meta_trace_carrier(SimpleNamespace(meta=only_progress)) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stateful_mcp_tool_call_uses_current_requests_otel_destinations() -> None:
|
||||
from types import SimpleNamespace
|
||||
|
||||
from mcp.server.lowlevel.server import request_ctx
|
||||
from mcp.shared.context import RequestContext
|
||||
|
||||
from litellm.integrations.otel.model.destination import OtelDestination
|
||||
from litellm.integrations.otel.plumbing.context import (
|
||||
request_destinations,
|
||||
reset_request_destinations,
|
||||
set_request_destinations,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
_MCP_DESTINATIONS_SCOPE_KEY,
|
||||
mcp_server_tool_call,
|
||||
set_auth_context,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.tool_registry import (
|
||||
global_mcp_tool_registry,
|
||||
)
|
||||
|
||||
initialized_destination = OtelDestination(endpoint="https://initialize.example", callback_name="langfuse_otel")
|
||||
current_destination = OtelDestination(endpoint="https://current.example", callback_name="arize")
|
||||
server = MCPServer(
|
||||
server_id="otel-context-test",
|
||||
name="otelcontext",
|
||||
transport=MCPTransport.http,
|
||||
allow_all_keys=True,
|
||||
)
|
||||
|
||||
async def observe_destinations() -> str:
|
||||
assert request_destinations() == (current_destination,)
|
||||
return "ok"
|
||||
|
||||
global_mcp_server_manager.registry[server.server_id] = server
|
||||
global_mcp_server_manager.tool_name_to_mcp_server_name_mapping["otelcontext-observe"] = server.name
|
||||
global_mcp_tool_registry.register_tool(
|
||||
name="otelcontext-observe",
|
||||
description="Observe request destinations",
|
||||
input_schema={"type": "object"},
|
||||
handler=observe_destinations,
|
||||
)
|
||||
set_auth_context(None, raw_headers={})
|
||||
destinations_token = set_request_destinations((initialized_destination,))
|
||||
scope = {_MCP_DESTINATIONS_SCOPE_KEY: (current_destination,)}
|
||||
current_request_context = RequestContext(
|
||||
request_id=1,
|
||||
meta=None,
|
||||
session=SimpleNamespace(),
|
||||
lifespan_context=None,
|
||||
request=SimpleNamespace(scope=scope),
|
||||
)
|
||||
request_token = request_ctx.set(current_request_context)
|
||||
try:
|
||||
result = await mcp_server_tool_call("otelcontext-observe", {})
|
||||
assert result.isError is False
|
||||
assert request_destinations() == (initialized_destination,)
|
||||
finally:
|
||||
request_ctx.reset(request_token)
|
||||
reset_request_destinations(destinations_token)
|
||||
global_mcp_tool_registry.tools.pop("otelcontext-observe", None)
|
||||
global_mcp_server_manager.registry.pop(server.server_id, None)
|
||||
global_mcp_server_manager.tool_name_to_mcp_server_name_mapping.pop("otelcontext-observe", None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_allowed_mcp_servers_includes_active_servers_submitted_by_user():
|
||||
"""BYOM submitters can see approved servers they submitted without allow_all_keys."""
|
||||
|
|
|
|||
|
|
@ -7,11 +7,24 @@ Tests that invoke_agent_a2a properly integrates with add_litellm_data_to_request
|
|||
import json
|
||||
import socket
|
||||
import sys
|
||||
from contextlib import ExitStack
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from contextlib import AbstractContextManager, ExitStack
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
AddLiteLLMData = Callable[..., Awaitable[dict[str, object]]]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CapturedAgentCall:
|
||||
request_id: object
|
||||
agent_extra_headers: dict[str, str] | None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invoke_agent_a2a_adds_litellm_data():
|
||||
|
|
@ -364,7 +377,7 @@ def _make_agent_mock(url: str = "http://backend-agent:10001") -> MagicMock:
|
|||
|
||||
|
||||
def _make_request_mock(
|
||||
method: str, params: dict, request_id: object = "req-1"
|
||||
method: str, params: Mapping[str, object], request_id: object = "req-1"
|
||||
) -> MagicMock:
|
||||
req = MagicMock()
|
||||
req.headers = {}
|
||||
|
|
@ -379,7 +392,9 @@ def _make_request_mock(
|
|||
return req
|
||||
|
||||
|
||||
def _base_patches(agent: MagicMock):
|
||||
def _base_patches(
|
||||
agent: MagicMock, add_litellm_data: AddLiteLLMData | None = None
|
||||
) -> list[AbstractContextManager[object]]:
|
||||
return [
|
||||
patch(
|
||||
"litellm.proxy.agent_endpoints.a2a_endpoints._get_agent",
|
||||
|
|
@ -391,7 +406,7 @@ def _base_patches(agent: MagicMock):
|
|||
),
|
||||
patch(
|
||||
"litellm.proxy.common_request_processing.add_litellm_data_to_request",
|
||||
new=AsyncMock(side_effect=_add_proxy_data),
|
||||
new=AsyncMock(side_effect=add_litellm_data or _add_proxy_data),
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.general_settings", {}),
|
||||
patch("litellm.proxy.proxy_server.proxy_config", MagicMock()),
|
||||
|
|
@ -399,84 +414,67 @@ def _base_patches(agent: MagicMock):
|
|||
]
|
||||
|
||||
|
||||
async def _add_proxy_data(data, **kwargs):
|
||||
data["proxy_server_request"] = {
|
||||
"url": "http://localhost:4000",
|
||||
"method": "POST",
|
||||
"headers": {},
|
||||
"body": {},
|
||||
async def _add_proxy_data(data: dict[str, object], **kwargs: object) -> dict[str, object]:
|
||||
return {
|
||||
**data,
|
||||
"proxy_server_request": {"url": "http://localhost:4000", "method": "POST", "headers": {}, "body": {}},
|
||||
"metadata": data.get("metadata", {}),
|
||||
}
|
||||
data.setdefault("metadata", {})
|
||||
return data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("method", ["message/send", "message/stream"])
|
||||
async def test_message_methods_preserve_numeric_zero_request_id(method: str):
|
||||
_HELLO_MESSAGE_PARAMS = {
|
||||
"message": {
|
||||
"role": "user",
|
||||
"parts": [{"kind": "text", "text": "Hello"}],
|
||||
"messageId": "msg-123",
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
async def _invoke_message_method(
|
||||
method: str,
|
||||
mock_request: MagicMock,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
add_litellm_data: AddLiteLLMData | None = None,
|
||||
) -> CapturedAgentCall:
|
||||
from fastapi.responses import JSONResponse
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
class MessageSendParams:
|
||||
def __init__(self, **kwargs):
|
||||
def __init__(self, **kwargs: object) -> None:
|
||||
self.__dict__.update(kwargs)
|
||||
|
||||
class SendMessageRequest:
|
||||
def __init__(self, **kwargs):
|
||||
def __init__(self, **kwargs: object) -> None:
|
||||
self.__dict__.update(kwargs)
|
||||
|
||||
agent = _make_agent_mock()
|
||||
params = {
|
||||
"message": {
|
||||
"role": "user",
|
||||
"parts": [{"kind": "text", "text": "Hello"}],
|
||||
"messageId": "msg-123",
|
||||
}
|
||||
}
|
||||
mock_request = _make_request_mock(method, params, request_id=0)
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1")
|
||||
captured = {}
|
||||
|
||||
async def capture_asend_message(request, **kwargs):
|
||||
captured["request_id"] = request.id
|
||||
response = MagicMock()
|
||||
async def fake_asend_message(request: SendMessageRequest, **kwargs: object) -> MagicMock:
|
||||
response: Final = MagicMock()
|
||||
response.model_dump.return_value = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": request.id,
|
||||
"id": request.__dict__["id"],
|
||||
"result": {"status": "success"},
|
||||
}
|
||||
return response
|
||||
|
||||
async def capture_stream_message(**kwargs):
|
||||
captured["request_id"] = kwargs["request_id"]
|
||||
return JSONResponse({"jsonrpc": "2.0", "id": kwargs["request_id"]})
|
||||
async def fake_stream_message(request_id: object, **kwargs: object) -> JSONResponse:
|
||||
return JSONResponse({"jsonrpc": "2.0", "id": request_id})
|
||||
|
||||
mock_a2a_types = MagicMock()
|
||||
mock_a2a_types: Final = MagicMock()
|
||||
mock_a2a_types.MessageSendParams = MessageSendParams
|
||||
mock_a2a_types.SendMessageRequest = SendMessageRequest
|
||||
is_send: Final = method == "message/send"
|
||||
downstream: Final = AsyncMock(side_effect=fake_asend_message if is_send else fake_stream_message)
|
||||
|
||||
with ExitStack() as stack:
|
||||
for p in _base_patches(agent):
|
||||
for p in _base_patches(_make_agent_mock(), add_litellm_data):
|
||||
stack.enter_context(p)
|
||||
stack.enter_context(patch("litellm.a2a_protocol.main.A2A_SDK_AVAILABLE", True))
|
||||
if method == "message/send":
|
||||
stack.enter_context(
|
||||
patch.dict(
|
||||
sys.modules,
|
||||
{"a2a": MagicMock(), "a2a.types": mock_a2a_types},
|
||||
)
|
||||
)
|
||||
stack.enter_context(
|
||||
patch(
|
||||
"litellm.a2a_protocol.asend_message",
|
||||
new=AsyncMock(side_effect=capture_asend_message),
|
||||
)
|
||||
)
|
||||
if is_send:
|
||||
stack.enter_context(patch.dict(sys.modules, {"a2a": MagicMock(), "a2a.types": mock_a2a_types}))
|
||||
stack.enter_context(patch("litellm.a2a_protocol.asend_message", new=downstream))
|
||||
else:
|
||||
stack.enter_context(
|
||||
patch(
|
||||
"litellm.proxy.agent_endpoints.a2a_endpoints._handle_stream_message",
|
||||
new=AsyncMock(side_effect=capture_stream_message),
|
||||
)
|
||||
patch("litellm.proxy.agent_endpoints.a2a_endpoints._handle_stream_message", new=downstream)
|
||||
)
|
||||
|
||||
from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a
|
||||
|
|
@ -488,7 +486,82 @@ async def test_message_methods_preserve_numeric_zero_request_id(method: str):
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
assert captured["request_id"] == 0
|
||||
kwargs: Final = downstream.call_args.kwargs
|
||||
request_id: Final = kwargs["request"].__dict__["id"] if is_send else kwargs["request_id"]
|
||||
return CapturedAgentCall(request_id=request_id, agent_extra_headers=kwargs.get("agent_extra_headers"))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("method", ["message/send", "message/stream"])
|
||||
async def test_message_methods_preserve_numeric_zero_request_id(method: str):
|
||||
mock_request = _make_request_mock(method, _HELLO_MESSAGE_PARAMS, request_id=0)
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1")
|
||||
|
||||
captured = await _invoke_message_method(method, mock_request, user_api_key_dict)
|
||||
|
||||
assert captured.request_id == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("method", ["message/send", "message/stream"])
|
||||
async def test_message_methods_forward_caller_identity_headers(method: str):
|
||||
mock_request = _make_request_mock(method, _HELLO_MESSAGE_PARAMS)
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="user-abc", team_id="team-xyz")
|
||||
|
||||
captured = await _invoke_message_method(method, mock_request, user_api_key_dict)
|
||||
|
||||
forwarded_headers = captured.agent_extra_headers or {}
|
||||
assert forwarded_headers.get("X-LiteLLM-User-Id") == "user-abc"
|
||||
assert forwarded_headers.get("X-LiteLLM-Team-Id") == "team-xyz"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("method", ["message/send", "message/stream"])
|
||||
async def test_message_methods_caller_identity_headers_cannot_be_spoofed(method: str):
|
||||
mock_request = _make_request_mock(method, _HELLO_MESSAGE_PARAMS)
|
||||
mock_request.headers = {
|
||||
"x-a2a-test-agent-x-litellm-user-id": "attacker-user",
|
||||
"x-a2a-test-agent-x-litellm-team-id": "attacker-team",
|
||||
}
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="real-user", team_id="real-team")
|
||||
|
||||
captured = await _invoke_message_method(method, mock_request, user_api_key_dict)
|
||||
|
||||
forwarded_headers = captured.agent_extra_headers or {}
|
||||
assert (
|
||||
forwarded_headers.get("X-LiteLLM-User-Id") == "real-user"
|
||||
), "authenticated user id must not be overridden by forwarded client headers"
|
||||
assert (
|
||||
forwarded_headers.get("X-LiteLLM-Team-Id") == "real-team"
|
||||
), "authenticated team id must not be overridden by forwarded client headers"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("method", ["message/send", "message/stream"])
|
||||
async def test_message_methods_forward_key_bound_identity_not_pre_call_rewrite(method: str):
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
||||
|
||||
mock_request = _make_request_mock(method, _HELLO_MESSAGE_PARAMS)
|
||||
mock_request.headers = {"X-OpenWebUI-User-Id": "header-mapped-user"}
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="key-user", team_id="key-team")
|
||||
general_settings: Final = {
|
||||
"user_header_mappings": [{"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"}]
|
||||
}
|
||||
|
||||
async def apply_user_header_mapping(data: dict[str, object], **kwargs: object) -> dict[str, object]:
|
||||
LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping(
|
||||
general_settings, user_api_key_dict, dict(mock_request.headers)
|
||||
)
|
||||
return await _add_proxy_data(data, **kwargs)
|
||||
|
||||
captured = await _invoke_message_method(
|
||||
method, mock_request, user_api_key_dict, add_litellm_data=apply_user_header_mapping
|
||||
)
|
||||
|
||||
assert user_api_key_dict.user_id == "header-mapped-user", "precondition: pre-call rewrite ran"
|
||||
forwarded_headers = captured.agent_extra_headers or {}
|
||||
assert forwarded_headers.get("X-LiteLLM-User-Id") == "key-user"
|
||||
assert forwarded_headers.get("X-LiteLLM-Team-Id") == "key-team"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -223,7 +223,7 @@ async def test_static_overrides_dynamic():
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_headers():
|
||||
"""When no headers are configured, agent_extra_headers is None and behaviour is unchanged."""
|
||||
"""When no headers are configured, only the caller identity is forwarded."""
|
||||
mock_agent = _make_mock_agent() # no static_headers or extra_headers
|
||||
mock_request = _make_mock_request()
|
||||
|
||||
|
|
@ -231,7 +231,7 @@ async def test_no_headers():
|
|||
|
||||
call_kwargs = mock_asend.call_args.kwargs
|
||||
headers = call_kwargs.get("agent_extra_headers")
|
||||
assert headers is None
|
||||
assert headers == {"X-LiteLLM-User-Id": "u1"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -303,7 +303,7 @@ async def test_convention_unrelated_prefix_not_forwarded():
|
|||
mock_asend = await _invoke(mock_agent, mock_request, None)
|
||||
|
||||
headers = mock_asend.call_args.kwargs.get("agent_extra_headers")
|
||||
assert headers is None
|
||||
assert headers == {"X-LiteLLM-User-Id": "u1"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -393,7 +393,7 @@ async def test_non_databricks_agent_skips_oauth_resolution():
|
|||
|
||||
mock_resolve.assert_not_called()
|
||||
headers = mock_asend.call_args.kwargs.get("agent_extra_headers")
|
||||
assert headers == {"x-custom": "v"}
|
||||
assert headers == {"x-custom": "v", "X-LiteLLM-User-Id": "u1"}
|
||||
assert "Authorization" not in headers
|
||||
|
||||
|
||||
|
|
@ -477,7 +477,7 @@ async def test_convention_header_blocked_by_case_variant_static():
|
|||
|
||||
headers = mock_asend.call_args.kwargs.get("agent_extra_headers")
|
||||
assert headers is not None
|
||||
assert headers == {"Authorization": "Bearer admin-token"}
|
||||
assert headers == {"Authorization": "Bearer admin-token", "X-LiteLLM-User-Id": "u1"}
|
||||
assert "authorization" not in headers
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -3638,3 +3638,191 @@ def test_agent_registry_route_gate_open_to_non_admin_roles(user_role, method, ro
|
|||
valid_token=valid_token,
|
||||
request_data={},
|
||||
)
|
||||
TEAM_CALLBACK_ROUTES = (
|
||||
"/team/06bda574-5ca9-43d3-beb8-3b23c2f17112/callback",
|
||||
"/team/06bda574-5ca9-43d3-beb8-3b23c2f17112/callback/langfuse",
|
||||
# the routes register team_id with the :path converter, so a team id may
|
||||
# contain a slash
|
||||
"/team/tenant/06bda574/callback",
|
||||
"/team/tenant/06bda574/callback/langfuse",
|
||||
# team_id is a free-form string, so it may also contain a colon
|
||||
"/team/tenant:06bda574/callback",
|
||||
"/team/tenant:06bda574/callback/langfuse",
|
||||
# or both, which is the shape neither a "[^:]+" nor a "[^/]+" expansion
|
||||
# of the placeholder reaches on its own
|
||||
"/team/tenant:acme/prod/callback",
|
||||
"/team/tenant:acme/prod/callback/langfuse",
|
||||
)
|
||||
|
||||
|
||||
def _gate(route, role) -> str:
|
||||
"""Drive the real route gate for a non-proxy-admin caller.
|
||||
|
||||
Reports "allowed" when the gate lets the request through to its handler, and
|
||||
the denial message otherwise, so a caller asserts the verdict as a value
|
||||
instead of on whether an exception escaped.
|
||||
"""
|
||||
user_obj = LiteLLM_UserTable(
|
||||
user_id="team_admin_user",
|
||||
user_email="team-admin@example.com",
|
||||
user_role=role,
|
||||
)
|
||||
request = MagicMock(spec=Request)
|
||||
request.query_params = {}
|
||||
try:
|
||||
RouteChecks.non_proxy_admin_allowed_routes_check(
|
||||
user_obj=user_obj,
|
||||
_user_role=role,
|
||||
route=route,
|
||||
request=request,
|
||||
valid_token=UserAPIKeyAuth(user_id="team_admin_user", user_role=role),
|
||||
request_data={},
|
||||
)
|
||||
except Exception as exc:
|
||||
return f"denied: {exc}"
|
||||
return "allowed"
|
||||
|
||||
|
||||
def test_team_callback_routes_are_self_managed():
|
||||
"""The grant has to come from self_managed_routes specifically.
|
||||
|
||||
That list is the one whose entries carry no role predicate, so the handler
|
||||
decides. Granting the same paths through internal_user_routes instead would
|
||||
look identical for an internal_user while silently denying the org admins and
|
||||
view-only roles that list does not cover.
|
||||
"""
|
||||
for template in (
|
||||
"/team/{team_id:path}/callback",
|
||||
"/team/{team_id:path}/callback/{callback_name}",
|
||||
):
|
||||
assert template in LiteLLMRoutes.self_managed_routes.value
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", TEAM_CALLBACK_ROUTES)
|
||||
@pytest.mark.parametrize(
|
||||
"role",
|
||||
[
|
||||
LitellmUserRoles.INTERNAL_USER.value,
|
||||
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value,
|
||||
LitellmUserRoles.ORG_ADMIN.value,
|
||||
],
|
||||
)
|
||||
def test_team_callback_routes_reach_their_handler_for_non_admins(route, role):
|
||||
"""A team admin manages their own team's logging callbacks, so the route gate
|
||||
must let a non-proxy-admin through to the handler.
|
||||
|
||||
The handler is what authorizes: every team callback endpoint calls
|
||||
_verify_team_access, which admits only a proxy admin, an org admin for the
|
||||
team, or an admin of that team, and 403s everyone else. Before this, the gate
|
||||
rejected the team admin with a 401 naming proxy admin, so the handler's own
|
||||
check was unreachable for them.
|
||||
"""
|
||||
assert _gate(route, role) == "allowed"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"pattern, route, matches",
|
||||
[
|
||||
# a :path placeholder takes what the router's path converter takes
|
||||
("/team/{team_id:path}/callback", "/team/plain/callback", True),
|
||||
("/team/{team_id:path}/callback", "/team/tenant/acme/callback", True),
|
||||
("/team/{team_id:path}/callback", "/team/tenant:acme/callback", True),
|
||||
("/team/{team_id:path}/callback", "/team/tenant:acme/prod/callback", True),
|
||||
# and still has to reach the template's own suffix
|
||||
("/team/{team_id:path}/callback", "/team/tenant:acme/disable_logging", False),
|
||||
# a template with a ":" literal after the placeholder keeps the suffix
|
||||
(
|
||||
"/v1beta/models/{model_name:path}:generateContent",
|
||||
"/v1beta/models/gemini-2.5-flash:generateContent",
|
||||
True,
|
||||
),
|
||||
(
|
||||
"/v1beta/models/{model_name:path}:generateContent",
|
||||
"/v1beta/models/publishers/google/gemini-2.5-flash:generateContent",
|
||||
True,
|
||||
),
|
||||
# the value must not swallow that suffix and match a different verb
|
||||
(
|
||||
"/v1beta/models/{model_name:path}:generateContent",
|
||||
"/v1beta/models/gemini-2.5-flash:countTokens",
|
||||
False,
|
||||
),
|
||||
# a %0A in the value reaches the handler through the path converter, so
|
||||
# the gate has to see it too or DISABLE_ADMIN_ENDPOINTS is bypassable
|
||||
("/v1/mcp/server/{path:path}", "/v1/mcp/server/abc\ndef", True),
|
||||
("/team/{team_id:path}/callback", "/team/ten\nant/callback", True),
|
||||
("/v1beta/models/{model_name:path}:generateContent", "/v1beta/models/gem\nini:generateContent", True),
|
||||
# an ordinary placeholder stays one segment
|
||||
("/team/{team_id}/members/me", "/team/abc/members/me", True),
|
||||
("/team/{team_id}/members/me", "/team/tenant/abc/members/me", False),
|
||||
("/team/{team_id}/members/me", "/team/ab\nc/members/me", True),
|
||||
],
|
||||
)
|
||||
def test_path_placeholder_matches_what_the_router_accepts(pattern, route, matches):
|
||||
"""The gate's placeholder expansion has to agree with the router's.
|
||||
|
||||
A team id may carry a slash, a colon, or both, and the router mounted these
|
||||
paths with the same :path converter, so an id the router routes must not be
|
||||
an id the gate fails to recognize. The one narrowing that stays is a template
|
||||
whose own suffix begins with a colon: there the value stops before it, or
|
||||
":generateContent" would also match a ":countTokens" request.
|
||||
"""
|
||||
assert RouteChecks._route_matches_pattern(route=route, pattern=pattern) is matches
|
||||
|
||||
|
||||
# Every other route the proxy mounts under /team/{team_id}, spelled the way it
|
||||
# is registered. None of them takes a path converter, so none can be reached by
|
||||
# a URL that ends in the callback suffix.
|
||||
PROTECTED_TEAM_ROUTES = (
|
||||
"/team/06bda574-5ca9-43d3-beb8-3b23c2f17112",
|
||||
"/team/06bda574-5ca9-43d3-beb8-3b23c2f17112/disable_logging",
|
||||
"/team/06bda574-5ca9-43d3-beb8-3b23c2f17112/members/me",
|
||||
"/team/06bda574-5ca9-43d3-beb8-3b23c2f17112/member/u-1/reset_spend",
|
||||
# the same routes with the callback suffix spliced in, which is the shape a
|
||||
# caller would craft to make a protected route look self-managed
|
||||
"/team/06bda574/callback/disable_logging/x",
|
||||
"/team/06bda574/callback/member/u-1/reset_spend",
|
||||
"/team/06bda574/callback/members/me",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", PROTECTED_TEAM_ROUTES)
|
||||
def test_the_callback_grant_does_not_reach_another_team_route(route):
|
||||
"""Widening the callback templates must not hand out any neighbouring route.
|
||||
|
||||
The grant is two templates ending in the callback suffix. Every other team
|
||||
route registers an ordinary single-segment placeholder, so no URL the router
|
||||
sends to one of them can end in "/callback" or "/callback/<name>" -- and the
|
||||
gate must agree, or a crafted team id would carry a caller into a handler
|
||||
the grant never covered.
|
||||
"""
|
||||
for template in (
|
||||
"/team/{team_id:path}/callback",
|
||||
"/team/{team_id:path}/callback/{callback_name}",
|
||||
):
|
||||
assert RouteChecks._route_matches_pattern(route=route, pattern=template) is False
|
||||
|
||||
|
||||
def test_team_disable_logging_stays_proxy_admin_only():
|
||||
"""disable_logging was left out of the grant, so it must still be rejected at
|
||||
the gate. It is the one team callback route a team admin cannot reach."""
|
||||
verdict = _gate(
|
||||
"/team/06bda574-5ca9-43d3-beb8-3b23c2f17112/disable_logging",
|
||||
LitellmUserRoles.INTERNAL_USER.value,
|
||||
)
|
||||
|
||||
assert "Only proxy admin" in verdict
|
||||
assert "disable_logging" in verdict
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"route",
|
||||
[
|
||||
"/team/06bda574-5ca9-43d3-beb8-3b23c2f17112",
|
||||
"/team/update",
|
||||
"/team/06bda574-5ca9-43d3-beb8-3b23c2f17112/model/add",
|
||||
],
|
||||
)
|
||||
def test_neighbouring_team_routes_stay_closed(route):
|
||||
"""The grant is the callback paths and nothing else on the team namespace."""
|
||||
assert "Only proxy admin" in _gate(route, LitellmUserRoles.INTERNAL_USER.value)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,145 @@
|
|||
import json
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy._types import ProxyErrorTypes, ProxyException
|
||||
from litellm.proxy.common_utils.openai_error_payload import (
|
||||
error_status_code,
|
||||
openai_error_param,
|
||||
openai_error_type,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"status_code, expected_type",
|
||||
[
|
||||
(400, "invalid_request_error"),
|
||||
(401, "authentication_error"),
|
||||
(403, "permission_error"),
|
||||
(404, "invalid_request_error"),
|
||||
(408, "invalid_request_error"),
|
||||
(422, "invalid_request_error"),
|
||||
(429, "rate_limit_error"),
|
||||
(499, "invalid_request_error"),
|
||||
(500, "internal_server_error"),
|
||||
(502, "internal_server_error"),
|
||||
(503, "internal_server_error"),
|
||||
],
|
||||
)
|
||||
def test_status_code_decides_the_type_when_the_exception_carries_none(status_code: int, expected_type: str):
|
||||
"""A route that raises a bare HTTPException carries no error type, so the status it
|
||||
answered with is the only thing left to name the OpenAI type from."""
|
||||
assert openai_error_type(HTTPException(status_code=status_code, detail="boom"), status_code) == expected_type
|
||||
|
||||
|
||||
def test_a_carried_type_wins_over_the_one_the_status_would_imply():
|
||||
"""A ProxyException raised mid-request already names its own type, and relabelling a
|
||||
402 budget_exceeded as the status map's guess would lose what the client branches on."""
|
||||
carried = ProxyException(
|
||||
message="Budget has been exceeded",
|
||||
type=ProxyErrorTypes.budget_exceeded.value,
|
||||
param=None,
|
||||
code=400,
|
||||
)
|
||||
|
||||
assert openai_error_type(carried, 400) == ProxyErrorTypes.budget_exceeded.value
|
||||
|
||||
|
||||
@pytest.mark.parametrize("carried_type", [None, 400, {"type": "invalid_request_error"}, ["invalid_request_error"]])
|
||||
def test_a_non_string_carried_type_falls_back_to_the_status(carried_type: object):
|
||||
"""OpenAI types error.type as a string, so anything else on the exception is not one and
|
||||
must not reach the wire the way the literal "None" used to."""
|
||||
|
||||
class _Carrier(Exception):
|
||||
type = carried_type
|
||||
|
||||
assert openai_error_type(_Carrier("boom"), 401) == "authentication_error"
|
||||
|
||||
|
||||
def test_the_type_is_never_the_string_none_after_a_json_round_trip():
|
||||
"""The bug this module exists for: json.dumps of a "None" default is indistinguishable
|
||||
from a real type to a client's error handler."""
|
||||
payload = json.loads(
|
||||
json.dumps(
|
||||
{
|
||||
"type": openai_error_type(HTTPException(status_code=400, detail="boom"), 400),
|
||||
"param": openai_error_param(HTTPException(status_code=400, detail="boom")),
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
assert payload == {"type": "invalid_request_error", "param": None}
|
||||
|
||||
|
||||
def test_a_carried_param_names_the_offending_field():
|
||||
carried = ProxyException(message="Invalid purpose", type="invalid_request_error", param="purpose", code=400)
|
||||
|
||||
assert openai_error_param(carried) == "purpose"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exc", [HTTPException(status_code=400, detail="boom"), ValueError("boom"), None])
|
||||
def test_param_is_json_null_when_the_exception_names_no_field(exc: Exception | None):
|
||||
assert openai_error_param(exc) is None
|
||||
|
||||
|
||||
def test_a_non_string_carried_param_is_json_null():
|
||||
class _Carrier(Exception):
|
||||
param = 42
|
||||
|
||||
assert openai_error_param(_Carrier("boom")) is None
|
||||
|
||||
|
||||
def test_a_carried_status_code_wins_over_the_default():
|
||||
assert error_status_code(HTTPException(status_code=429, detail="slow down"), 400) == 429
|
||||
|
||||
|
||||
@pytest.mark.parametrize("default", [400, 500])
|
||||
def test_the_default_status_stands_when_the_exception_carries_none(default: int):
|
||||
assert error_status_code(ValueError("boom"), default) == default
|
||||
|
||||
|
||||
@pytest.mark.parametrize("carried_status", [True, False, "429", None, 429.0])
|
||||
def test_a_non_int_carried_status_falls_back_to_the_default(carried_status: object):
|
||||
"""True is an int in Python but not an HTTP status, and a stringified one would break
|
||||
every caller that compares the code numerically."""
|
||||
|
||||
class _Carrier(Exception):
|
||||
status_code = carried_status
|
||||
|
||||
assert error_status_code(_Carrier("boom"), 500) == 500
|
||||
|
||||
|
||||
def test_a_proxy_exception_keeps_the_status_it_was_raised_with():
|
||||
"""ProxyException stores its status as the string ``code`` rather than ``status_code``,
|
||||
so a route tail that rewraps one used to answer a 4xx rejection as a 500."""
|
||||
rejection = ProxyException(message="session_id is required", type="bad_request_error", param="session_id", code=400)
|
||||
|
||||
assert error_status_code(rejection, 500) == 400
|
||||
|
||||
|
||||
@pytest.mark.parametrize("carried_code", [None, "None", "", "rate_limited", "4xx", 404])
|
||||
def test_a_code_that_is_not_a_decimal_string_falls_back_to_the_default(carried_code: object):
|
||||
"""Only ProxyException's stringified status is a status; ``code`` on anything else
|
||||
(OpenAI's ``invalid_api_key``, a stray int) says nothing about the HTTP answer."""
|
||||
|
||||
class _Carrier(Exception):
|
||||
code = carried_code
|
||||
|
||||
assert error_status_code(_Carrier("boom"), 500) == 500
|
||||
|
||||
|
||||
def test_a_status_code_wins_over_a_stringified_code():
|
||||
class _Carrier(Exception):
|
||||
status_code = 429
|
||||
code = "400"
|
||||
|
||||
assert error_status_code(_Carrier("boom"), 500) == 429
|
||||
|
||||
|
||||
def test_a_status_carried_by_an_exception_drives_the_type_it_reports():
|
||||
"""The two helpers compose at every call site: the status the exception carries is what
|
||||
names its type, not the default the route would have used."""
|
||||
exc = HTTPException(status_code=403, detail="blocked by policy")
|
||||
|
||||
assert openai_error_type(exc, error_status_code(exc, 400)) == "permission_error"
|
||||
|
|
@ -2944,11 +2944,9 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_on_rows_other_worker
|
|||
A `previous_response_id` chained straight off the previous turn reads the DB, so a
|
||||
Responses row cannot sit in this worker's queue until the monitor's next poll.
|
||||
"""
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
prisma = _tool_usage_prisma()
|
||||
PrismaClient.spend_log_flush_requested.clear()
|
||||
prisma.spend_log_flush_requested = asyncio.Event()
|
||||
|
||||
await db_writer._insert_spend_log_to_db(
|
||||
payload={"request_id": "req-1", "call_type": call_type},
|
||||
|
|
@ -2956,8 +2954,7 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_on_rows_other_worker
|
|||
)
|
||||
|
||||
assert prisma.spend_log_transactions == [{"request_id": "req-1", "call_type": call_type}]
|
||||
assert PrismaClient.spend_log_flush_requested.is_set() is expects_flush
|
||||
PrismaClient.spend_log_flush_requested.clear()
|
||||
assert prisma.spend_log_flush_requested.is_set() is expects_flush
|
||||
|
||||
|
||||
def _batch_cost_payload() -> dict:
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ _MANAGED_DB_ENV_VARS = (
|
|||
"IAM_TOKEN_DB_AUTH",
|
||||
"AZURE_POSTGRESQL_AUTH",
|
||||
"DATABASE_DISABLE_PREPARED_STATEMENTS",
|
||||
"DATABASE_MAX_IDLE_CONNECTION_LIFETIME",
|
||||
"DATABASE_URL",
|
||||
"DIRECT_URL",
|
||||
"DATABASE_URL_READ_REPLICA",
|
||||
|
|
@ -111,7 +112,7 @@ def test_assembles_writer_url_when_iam_enabled(monkeypatch):
|
|||
|
||||
assert (
|
||||
os.environ["DATABASE_URL"]
|
||||
== "postgresql://litellm:WRITER_TOKEN@writer.example.com:5432/litellm_db"
|
||||
== "postgresql://litellm:WRITER_TOKEN@writer.example.com:5432/litellm_db?max_idle_connection_lifetime=60"
|
||||
)
|
||||
# Reader was never configured, so it must not have been set.
|
||||
assert "DATABASE_URL_READ_REPLICA" not in os.environ
|
||||
|
|
@ -130,7 +131,9 @@ def test_a_pre_encoded_iam_user_survives_url_assembly(monkeypatch):
|
|||
with _stub_iam_token("WRITER_TOKEN"):
|
||||
assert _apply() is True
|
||||
|
||||
assert os.environ["DATABASE_URL"] == "postgresql://svc%40corp:WRITER_TOKEN@writer.example.com:5432/litellm_db"
|
||||
assert os.environ["DATABASE_URL"] == (
|
||||
"postgresql://svc%40corp:WRITER_TOKEN@writer.example.com:5432/litellm_db?max_idle_connection_lifetime=60"
|
||||
)
|
||||
|
||||
|
||||
def test_an_unreadable_toggle_fails_the_settings_model(monkeypatch):
|
||||
|
|
@ -168,7 +171,7 @@ def test_reader_url_assembled_when_host_set_and_url_unset(monkeypatch):
|
|||
|
||||
assert (
|
||||
os.environ["DATABASE_URL_READ_REPLICA"]
|
||||
== "postgresql://litellm:READER_TOKEN@reader.example.com:5432/litellm_db"
|
||||
== "postgresql://litellm:READER_TOKEN@reader.example.com:5432/litellm_db?max_idle_connection_lifetime=60"
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -191,7 +194,7 @@ def test_reader_url_not_clobbered_when_already_set(monkeypatch):
|
|||
|
||||
assert (
|
||||
os.environ["DATABASE_URL_READ_REPLICA"]
|
||||
== "postgresql://app:secret@reader.example.com:5432/litellm_db"
|
||||
== "postgresql://app:secret@reader.example.com:5432/litellm_db?max_idle_connection_lifetime=60"
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -222,7 +225,8 @@ def test_reader_field_fallbacks_default_to_writer_values(monkeypatch):
|
|||
|
||||
assert (
|
||||
os.environ["DATABASE_URL_READ_REPLICA"]
|
||||
== "postgresql://litellm:READER_TOKEN@reader.example.com:5432/litellm_db?schema=public"
|
||||
== "postgresql://litellm:READER_TOKEN@reader.example.com:5432/litellm_db"
|
||||
"?schema=public&max_idle_connection_lifetime=60"
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -242,7 +246,7 @@ def test_assembles_writer_url_when_azure_entra_enabled(monkeypatch):
|
|||
|
||||
assert os.environ["DATABASE_URL"] == (
|
||||
"postgresql://litellm%40contoso.onmicrosoft.com:ENTRA_TOKEN"
|
||||
"@writer.postgres.database.azure.com:5432/litellm_db"
|
||||
"@writer.postgres.database.azure.com:5432/litellm_db?max_idle_connection_lifetime=60"
|
||||
)
|
||||
assert os.environ["AZURE_POSTGRESQL_AUTH"] == "True"
|
||||
assert "IAM_TOKEN_DB_AUTH" not in os.environ
|
||||
|
|
@ -261,7 +265,7 @@ def test_azure_reader_url_assembled_from_writer_fallbacks(monkeypatch):
|
|||
|
||||
assert os.environ["DATABASE_URL_READ_REPLICA"] == (
|
||||
"postgresql://litellm%40contoso.onmicrosoft.com:ENTRA_TOKEN"
|
||||
"@reader.postgres.database.azure.com:5432/litellm_db?schema=public"
|
||||
"@reader.postgres.database.azure.com:5432/litellm_db?schema=public&max_idle_connection_lifetime=60"
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -357,7 +361,7 @@ def test_assembles_writer_url_from_password(monkeypatch):
|
|||
assert _apply() is True
|
||||
assert (
|
||||
os.environ["DATABASE_URL"]
|
||||
== "postgresql://litellm:s3cr3t@writer.example.com:5432/litellm_db"
|
||||
== "postgresql://litellm:s3cr3t@writer.example.com:5432/litellm_db?max_idle_connection_lifetime=60"
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -370,7 +374,7 @@ def test_writer_password_is_percent_encoded(monkeypatch):
|
|||
assert _apply() is True
|
||||
assert (
|
||||
os.environ["DATABASE_URL"]
|
||||
== "postgresql://litellm:p%40ss%2Fw%3Ard@writer.example.com:5432/litellm_db"
|
||||
== "postgresql://litellm:p%40ss%2Fw%3Ard@writer.example.com:5432/litellm_db?max_idle_connection_lifetime=60"
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -388,7 +392,7 @@ def test_writer_url_not_clobbered_when_already_set(monkeypatch):
|
|||
assert _apply() is False
|
||||
assert (
|
||||
os.environ["DATABASE_URL"]
|
||||
== "postgresql://pinned:url@db.example.com:5432/litellm_db"
|
||||
== "postgresql://pinned:url@db.example.com:5432/litellm_db?max_idle_connection_lifetime=60"
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -400,7 +404,7 @@ def test_writer_url_passwordless(monkeypatch):
|
|||
assert _apply() is True
|
||||
assert (
|
||||
os.environ["DATABASE_URL"]
|
||||
== "postgresql://litellm@writer.example.com:5432/litellm_db"
|
||||
== "postgresql://litellm@writer.example.com:5432/litellm_db?max_idle_connection_lifetime=60"
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -415,7 +419,7 @@ def test_database_username_alias(monkeypatch):
|
|||
assert _apply() is True
|
||||
assert (
|
||||
os.environ["DATABASE_URL"]
|
||||
== "postgresql://litellm:s3cr3t@writer.example.com:5432/litellm_db"
|
||||
== "postgresql://litellm:s3cr3t@writer.example.com:5432/litellm_db?max_idle_connection_lifetime=60"
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -429,7 +433,7 @@ def test_password_reader_falls_back_to_writer_password(monkeypatch):
|
|||
assert _apply() is True
|
||||
assert (
|
||||
os.environ["DATABASE_URL_READ_REPLICA"]
|
||||
== "postgresql://litellm:s3cr3t@reader.example.com:5432/litellm_db"
|
||||
== "postgresql://litellm:s3cr3t@reader.example.com:5432/litellm_db?max_idle_connection_lifetime=60"
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -445,7 +449,7 @@ def test_password_reader_uses_own_credentials(monkeypatch):
|
|||
assert _apply() is True
|
||||
assert (
|
||||
os.environ["DATABASE_URL_READ_REPLICA"]
|
||||
== "postgresql://litellm_ro:ro_pw@reader.example.com:5432/litellm_db"
|
||||
== "postgresql://litellm_ro:ro_pw@reader.example.com:5432/litellm_db?max_idle_connection_lifetime=60"
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -641,19 +645,19 @@ def test_reader_keeps_its_own_options_when_writer_params_are_appended(monkeypatc
|
|||
assert query["connection_limit"] == ["3"]
|
||||
|
||||
|
||||
def test_reader_url_left_alone_when_writer_has_no_params(monkeypatch):
|
||||
def test_reader_url_left_alone_when_nothing_is_missing(monkeypatch):
|
||||
"""No params to inherit must mean the reader URL is not rewritten at all."""
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@writer.example.com:5432/db")
|
||||
monkeypatch.setenv(
|
||||
"DATABASE_URL_READ_REPLICA",
|
||||
"postgresql://u:p@reader.example.com:5432/db?options=-c%20search_path%3Dapp",
|
||||
"postgresql://u:p@reader.example.com:5432/db?options=-c%20search_path%3Dapp&max_idle_connection_lifetime=45",
|
||||
)
|
||||
|
||||
_apply()
|
||||
|
||||
assert (
|
||||
os.environ["DATABASE_URL_READ_REPLICA"]
|
||||
== "postgresql://u:p@reader.example.com:5432/db?options=-c%20search_path%3Dapp"
|
||||
== "postgresql://u:p@reader.example.com:5432/db?options=-c%20search_path%3Dapp&max_idle_connection_lifetime=45"
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -671,7 +675,7 @@ def test_disable_prepared_statements_appends_pgbouncer_to_assembled_writer(monke
|
|||
|
||||
assert _apply() is True
|
||||
assert os.environ["DATABASE_URL"] == (
|
||||
"postgresql://litellm:s3cr3t@writer.example.com:5432/litellm_db?pgbouncer=true"
|
||||
"postgresql://litellm:s3cr3t@writer.example.com:5432/litellm_db?pgbouncer=true&max_idle_connection_lifetime=60"
|
||||
)
|
||||
assert "DIRECT_URL" not in os.environ
|
||||
|
||||
|
|
@ -685,7 +689,9 @@ def test_disable_prepared_statements_appends_pgbouncer_to_pinned_writer(monkeypa
|
|||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@db.example.com:5432/litellm_db")
|
||||
|
||||
assert _apply() is False
|
||||
assert os.environ["DATABASE_URL"] == "postgresql://u:p@db.example.com:5432/litellm_db?pgbouncer=true"
|
||||
assert os.environ["DATABASE_URL"] == (
|
||||
"postgresql://u:p@db.example.com:5432/litellm_db?pgbouncer=true&max_idle_connection_lifetime=60"
|
||||
)
|
||||
|
||||
|
||||
def test_disable_prepared_statements_respects_a_pinned_pgbouncer_value(monkeypatch):
|
||||
|
|
@ -694,7 +700,9 @@ def test_disable_prepared_statements_respects_a_pinned_pgbouncer_value(monkeypat
|
|||
|
||||
_apply()
|
||||
|
||||
assert os.environ["DATABASE_URL"] == "postgresql://u:p@db.example.com:5432/litellm_db?pgbouncer=false"
|
||||
assert os.environ["DATABASE_URL"] == (
|
||||
"postgresql://u:p@db.example.com:5432/litellm_db?pgbouncer=false&max_idle_connection_lifetime=60"
|
||||
)
|
||||
|
||||
|
||||
def test_disable_prepared_statements_applies_to_direct_url(monkeypatch):
|
||||
|
|
@ -704,7 +712,9 @@ def test_disable_prepared_statements_applies_to_direct_url(monkeypatch):
|
|||
|
||||
_apply()
|
||||
|
||||
assert os.environ["DIRECT_URL"] == "postgresql://u:p@direct.example.com:5432/litellm_db?pgbouncer=true"
|
||||
assert os.environ["DIRECT_URL"] == (
|
||||
"postgresql://u:p@direct.example.com:5432/litellm_db?pgbouncer=true&max_idle_connection_lifetime=60"
|
||||
)
|
||||
|
||||
|
||||
def test_reader_inherits_pgbouncer_from_disable_prepared_statements(monkeypatch):
|
||||
|
|
@ -724,7 +734,9 @@ def test_disable_prepared_statements_off_leaves_urls_alone(monkeypatch):
|
|||
|
||||
_apply()
|
||||
|
||||
assert os.environ["DATABASE_URL"] == "postgresql://u:p@db.example.com:5432/litellm_db"
|
||||
assert os.environ["DATABASE_URL"] == (
|
||||
"postgresql://u:p@db.example.com:5432/litellm_db?max_idle_connection_lifetime=60"
|
||||
)
|
||||
|
||||
|
||||
def test_disable_prepared_statements_rejects_an_unreadable_value(monkeypatch):
|
||||
|
|
@ -760,6 +772,7 @@ def test_libpq_verify_full_and_sslrootcert_become_prisma_strict_sslcert(monkeypa
|
|||
"sslmode": ["require"],
|
||||
"sslcert": ["/certs/rds-bundle.pem"],
|
||||
"sslaccept": ["strict"],
|
||||
"max_idle_connection_lifetime": ["60"],
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -768,7 +781,11 @@ def test_libpq_verify_ca_becomes_prisma_strict(monkeypatch):
|
|||
|
||||
_apply()
|
||||
|
||||
assert _query(os.environ["DATABASE_URL"]) == {"sslmode": ["require"], "sslaccept": ["strict"]}
|
||||
assert _query(os.environ["DATABASE_URL"]) == {
|
||||
"sslmode": ["require"],
|
||||
"sslaccept": ["strict"],
|
||||
"max_idle_connection_lifetime": ["60"],
|
||||
}
|
||||
|
||||
|
||||
def test_sslrootcert_alone_turns_on_strict_verification(monkeypatch):
|
||||
|
|
@ -784,6 +801,7 @@ def test_sslrootcert_alone_turns_on_strict_verification(monkeypatch):
|
|||
"sslmode": ["require"],
|
||||
"sslcert": ["/certs/ca.pem"],
|
||||
"sslaccept": ["strict"],
|
||||
"max_idle_connection_lifetime": ["60"],
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -800,11 +818,15 @@ def test_pinned_prisma_ssl_params_win_over_libpq_translation(monkeypatch):
|
|||
"sslmode": ["require"],
|
||||
"sslcert": ["/pinned.pem"],
|
||||
"sslaccept": ["accept_invalid_certs"],
|
||||
"max_idle_connection_lifetime": ["60"],
|
||||
}
|
||||
|
||||
|
||||
def test_prisma_native_ssl_url_is_left_untouched(monkeypatch):
|
||||
url = "postgresql://u:p@db.example.com:5432/litellm_db?sslmode=require&sslcert=/certs/ca.pem&sslaccept=strict"
|
||||
url = (
|
||||
"postgresql://u:p@db.example.com:5432/litellm_db"
|
||||
"?sslmode=require&sslcert=/certs/ca.pem&sslaccept=strict&max_idle_connection_lifetime=60"
|
||||
)
|
||||
monkeypatch.setenv("DATABASE_URL", url)
|
||||
|
||||
_apply()
|
||||
|
|
@ -820,4 +842,84 @@ def test_libpq_ssl_translation_covers_direct_url_and_read_replica(monkeypatch):
|
|||
_apply()
|
||||
|
||||
for env_var in ("DATABASE_URL", "DIRECT_URL", "DATABASE_URL_READ_REPLICA"):
|
||||
assert _query(os.environ[env_var]) == {"sslmode": ["require"], "sslaccept": ["strict"]}, env_var
|
||||
assert _query(os.environ[env_var]) == {
|
||||
"sslmode": ["require"],
|
||||
"sslaccept": ["strict"],
|
||||
"max_idle_connection_lifetime": ["60"],
|
||||
}, env_var
|
||||
|
||||
|
||||
def test_default_idle_lifetime_applied_to_pinned_writer_and_direct_url(monkeypatch):
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@db.example.com:5432/litellm_db")
|
||||
monkeypatch.setenv("DIRECT_URL", "postgresql://u:p@direct.example.com:5432/litellm_db")
|
||||
|
||||
assert _apply() is False
|
||||
|
||||
assert os.environ["DATABASE_URL"] == (
|
||||
"postgresql://u:p@db.example.com:5432/litellm_db?max_idle_connection_lifetime=60"
|
||||
)
|
||||
assert os.environ["DIRECT_URL"] == (
|
||||
"postgresql://u:p@direct.example.com:5432/litellm_db?max_idle_connection_lifetime=60"
|
||||
)
|
||||
|
||||
|
||||
def test_url_pinned_idle_lifetime_wins_over_default_and_env_knob(monkeypatch):
|
||||
monkeypatch.setenv("DATABASE_MAX_IDLE_CONNECTION_LIFETIME", "45")
|
||||
monkeypatch.setenv(
|
||||
"DATABASE_URL", "postgresql://u:p@db.example.com:5432/litellm_db?max_idle_connection_lifetime=300"
|
||||
)
|
||||
|
||||
_apply()
|
||||
|
||||
assert os.environ["DATABASE_URL"] == (
|
||||
"postgresql://u:p@db.example.com:5432/litellm_db?max_idle_connection_lifetime=300"
|
||||
)
|
||||
|
||||
|
||||
def test_env_knob_overrides_default_idle_lifetime(monkeypatch):
|
||||
monkeypatch.setenv("DATABASE_MAX_IDLE_CONNECTION_LIFETIME", "45")
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@db.example.com:5432/litellm_db")
|
||||
monkeypatch.setenv("DIRECT_URL", "postgresql://u:p@direct.example.com:5432/litellm_db")
|
||||
|
||||
_apply()
|
||||
|
||||
assert os.environ["DATABASE_URL"] == (
|
||||
"postgresql://u:p@db.example.com:5432/litellm_db?max_idle_connection_lifetime=45"
|
||||
)
|
||||
assert os.environ["DIRECT_URL"] == (
|
||||
"postgresql://u:p@direct.example.com:5432/litellm_db?max_idle_connection_lifetime=45"
|
||||
)
|
||||
|
||||
|
||||
def test_env_knob_rejects_a_non_integer_value(monkeypatch):
|
||||
monkeypatch.setenv("DATABASE_MAX_IDLE_CONNECTION_LIFETIME", "soon")
|
||||
|
||||
with pytest.raises(ValidationError, match="DATABASE_MAX_IDLE_CONNECTION_LIFETIME"):
|
||||
DatabaseURLSettings.from_env()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("knob", "expected"), [(None, "60"), ("45", "45")])
|
||||
def test_reader_inherits_the_writer_idle_lifetime(monkeypatch, knob, expected):
|
||||
if knob is not None:
|
||||
monkeypatch.setenv("DATABASE_MAX_IDLE_CONNECTION_LIFETIME", knob)
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@writer.example.com:5432/db")
|
||||
monkeypatch.setenv("DATABASE_URL_READ_REPLICA", "postgresql://u:p@reader.example.com:5432/db")
|
||||
|
||||
_apply()
|
||||
|
||||
assert os.environ["DATABASE_URL_READ_REPLICA"] == (
|
||||
f"postgresql://u:p@reader.example.com:5432/db?max_idle_connection_lifetime={expected}"
|
||||
)
|
||||
|
||||
|
||||
def test_reader_keeps_its_own_pinned_idle_lifetime(monkeypatch):
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@writer.example.com:5432/db")
|
||||
monkeypatch.setenv(
|
||||
"DATABASE_URL_READ_REPLICA", "postgresql://u:p@reader.example.com:5432/db?max_idle_connection_lifetime=120"
|
||||
)
|
||||
|
||||
_apply()
|
||||
|
||||
assert os.environ["DATABASE_URL_READ_REPLICA"] == (
|
||||
"postgresql://u:p@reader.example.com:5432/db?max_idle_connection_lifetime=120"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -291,6 +291,56 @@ def test_azure_entra_mint_writes_an_encoded_url_into_the_db_url_env_var(azure_en
|
|||
assert os.environ["DATABASE_URL"] == db_url
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("previous_query", "expected_query"),
|
||||
[
|
||||
("max_idle_connection_lifetime=60", {"max_idle_connection_lifetime": ["60"]}),
|
||||
(
|
||||
"connection_limit=20&pgbouncer=true&max_idle_connection_lifetime=45",
|
||||
{"connection_limit": ["20"], "pgbouncer": ["true"], "max_idle_connection_lifetime": ["45"]},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_token_refresh_keeps_the_connection_params_of_the_url_it_replaces(
|
||||
azure_env, monkeypatch, previous_query, expected_query
|
||||
):
|
||||
old_token = _entra_jwt(60)
|
||||
monkeypatch.setenv(
|
||||
"DATABASE_URL",
|
||||
f"postgresql://litellm%40contoso.onmicrosoft.com:{urllib.parse.quote(old_token, safe='')}"
|
||||
f"@pg.postgres.database.azure.com:5432/litellm_db?{previous_query}",
|
||||
)
|
||||
new_token = _entra_jwt(3600)
|
||||
|
||||
db_url = _azure_wrapper(new_token).get_rds_iam_token()
|
||||
|
||||
assert db_url is not None
|
||||
assert os.environ["DATABASE_URL"] == db_url
|
||||
assert urllib.parse.quote(new_token, safe="") in db_url
|
||||
assert urllib.parse.parse_qs(urllib.parse.urlsplit(db_url).query) == expected_query
|
||||
|
||||
|
||||
def test_token_refresh_keeps_the_reader_url_params_separate_from_the_writer(azure_env, monkeypatch):
|
||||
from litellm.proxy.db.token_auth import IAMEndpoint
|
||||
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://w:t@pg:5432/litellm_db?max_idle_connection_lifetime=45")
|
||||
monkeypatch.setenv(
|
||||
"DATABASE_URL_READ_REPLICA", "postgresql://r:t@replica:5432/litellm_db?max_idle_connection_lifetime=60"
|
||||
)
|
||||
reader = _azure_wrapper(
|
||||
_entra_jwt(3600),
|
||||
db_url_env_var="DATABASE_URL_READ_REPLICA",
|
||||
iam_endpoint=IAMEndpoint(host="replica", port="5432", user="r", name="litellm_db", schema=None),
|
||||
)
|
||||
|
||||
reader_url = reader.get_rds_iam_token()
|
||||
|
||||
assert reader_url is not None
|
||||
assert reader_url.startswith("postgresql://r:")
|
||||
assert urllib.parse.parse_qs(urllib.parse.urlsplit(reader_url).query) == {"max_idle_connection_lifetime": ["60"]}
|
||||
assert os.environ["DATABASE_URL"].endswith("?max_idle_connection_lifetime=45")
|
||||
|
||||
|
||||
def test_azure_entra_refresh_is_scheduled_off_the_jwt_expiry(azure_env):
|
||||
"""Without reading `exp` this falls back to a fixed 600s interval, which silently
|
||||
outlives a token and breaks every reconnect after it lapses (issue #29661)."""
|
||||
|
|
|
|||
|
|
@ -5,12 +5,12 @@ from typing import Any, Dict
|
|||
|
||||
import orjson
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi.testclient import TestClient
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import Response
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.image_endpoints import endpoints
|
||||
|
||||
|
|
@ -167,3 +167,47 @@ def test_image_edit_multipart_n_that_is_not_a_number_is_left_alone(monkeypatch):
|
|||
|
||||
assert response.status_code == 200
|
||||
assert captured["n"] == "two"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_model_the_router_cannot_serve_answers_an_openai_typed_error(monkeypatch: pytest.MonkeyPatch):
|
||||
"""A bare HTTPException carries no type or param, so the tail used to ship the
|
||||
literal string "None" in both fields."""
|
||||
|
||||
async def fake_add_litellm_data_to_request(**kwargs: object) -> object:
|
||||
return kwargs["data"]
|
||||
|
||||
async def fake_pre_call_hook(*, user_api_key_dict: UserAPIKeyAuth, data: dict[str, object], call_type: str) -> dict[str, object]:
|
||||
return data
|
||||
|
||||
async def fake_post_call_failure_hook(**_: object) -> None:
|
||||
return None
|
||||
|
||||
async def failing_route_request(**_: object) -> None:
|
||||
raise HTTPException(
|
||||
status_code=404, detail={"error": "image_generation: Invalid model name passed in model=dall-e-3"}
|
||||
)
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.add_litellm_data_to_request", fake_add_litellm_data_to_request)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", {})
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.proxy_logging_obj",
|
||||
SimpleNamespace(pre_call_hook=fake_pre_call_hook, post_call_failure_hook=fake_post_call_failure_hook),
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.version", "test-version")
|
||||
monkeypatch.setattr("litellm.proxy.image_endpoints.endpoints.route_request", failing_route_request)
|
||||
|
||||
body = orjson.dumps({"model": "dall-e-3", "prompt": "a lighthouse at dusk"})
|
||||
|
||||
async def receive() -> dict[str, object]:
|
||||
return {"type": "http.request", "body": body, "more_body": False}
|
||||
|
||||
request = Request({"type": "http", "method": "POST", "path": "/v1/images/generations", "headers": []}, receive)
|
||||
|
||||
with pytest.raises(ProxyException) as raised:
|
||||
await endpoints.image_generation(request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth())
|
||||
|
||||
assert (raised.value.type, raised.value.param, raised.value.code) == ("invalid_request_error", None, "404")
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ from litellm.proxy._types import (
|
|||
LitellmUserRoles,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.common_utils.callback_config_validation import cross_entry_family_error
|
||||
from litellm.proxy.management_endpoints.team_callback_endpoints import (
|
||||
add_team_callbacks,
|
||||
delete_team_callback,
|
||||
|
|
@ -1443,3 +1444,118 @@ async def test_delete_team_callback_route_accepts_team_ids_containing_slashes():
|
|||
assert response.json()["data"]["success_callbacks"] == ["langsmith"]
|
||||
written = json.loads(mock_prisma.db.litellm_teamtable.update.await_args.kwargs["data"]["metadata"])
|
||||
assert [entry["callback_name"] for entry in written["logging"]] == ["langsmith"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"call_handler",
|
||||
[
|
||||
lambda caller: add_team_callbacks(
|
||||
data=AddTeamCallback(
|
||||
callback_name="langfuse",
|
||||
callback_type="success",
|
||||
callback_vars={"langfuse_public_key": "pk", "langfuse_secret_key": "sk"},
|
||||
),
|
||||
http_request=Mock(spec=Request),
|
||||
team_id="team-does-not-exist",
|
||||
user_api_key_dict=caller,
|
||||
),
|
||||
lambda caller: get_team_callbacks(
|
||||
http_request=Mock(spec=Request),
|
||||
team_id="team-does-not-exist",
|
||||
user_api_key_dict=caller,
|
||||
),
|
||||
lambda caller: delete_team_callback(
|
||||
http_request=Mock(spec=Request),
|
||||
team_id="team-does-not-exist",
|
||||
callback_name="langfuse",
|
||||
user_api_key_dict=caller,
|
||||
),
|
||||
],
|
||||
ids=["add", "get", "delete"],
|
||||
)
|
||||
async def test_unknown_team_is_indistinguishable_from_no_access(call_handler, unauthorized_caller):
|
||||
"""An unauthorized caller must not learn whether a team id exists.
|
||||
|
||||
These routes are reachable by any authenticated caller so a team admin can get
|
||||
as far as the access check, so a distinct "does not exist" would turn them into
|
||||
a probe for valid team ids. The unknown-team response has to match the
|
||||
no-access one exactly, status and body.
|
||||
"""
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_client: # test-quality-ok: the handler imports prisma_client from proxy_server at call time, so there is no seam to inject through
|
||||
mock_client.get_data = AsyncMock(return_value=None)
|
||||
with pytest.raises(HTTPException) as unknown_team:
|
||||
await call_handler(unauthorized_caller)
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_client: # test-quality-ok: the handler imports prisma_client from proxy_server at call time, so there is no seam to inject through
|
||||
mock_client.get_data = AsyncMock(return_value=_team_row())
|
||||
mock_client.db.litellm_teamtable.update = AsyncMock()
|
||||
with patch( # test-quality-ok: _verify_team_access calls this module-level helper directly, so there is no seam to inject through
|
||||
"litellm.proxy.management_endpoints.team_endpoints._is_user_org_admin_for_team",
|
||||
new_callable=AsyncMock,
|
||||
return_value=False,
|
||||
):
|
||||
with pytest.raises(HTTPException) as no_access:
|
||||
await call_handler(unauthorized_caller)
|
||||
|
||||
assert unknown_team.value.status_code == no_access.value.status_code == 403
|
||||
assert unknown_team.value.detail == no_access.value.detail
|
||||
assert "does not exist" not in str(unknown_team.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_admin_still_told_the_team_is_unknown():
|
||||
"""The masking is only for callers who could not have managed the team; a proxy
|
||||
admin keeps the diagnosable error."""
|
||||
admin = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin", api_key="sk-admin")
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_client: # test-quality-ok: the handler imports prisma_client from proxy_server at call time, so there is no seam to inject through
|
||||
mock_client.get_data = AsyncMock(return_value=None)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await get_team_callbacks(
|
||||
http_request=Mock(spec=Request),
|
||||
team_id="team-does-not-exist",
|
||||
user_api_key_dict=admin,
|
||||
)
|
||||
|
||||
assert exc.value.status_code == 404
|
||||
assert "does not exist" in str(exc.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"new_vars, stored, rejected",
|
||||
[
|
||||
# the redirect, in every carrier a caller could pick: an entry naming
|
||||
# only a host, pairing with a key pair written on another entry
|
||||
({"langfuse_host": "http://attacker.invalid"}, [{"langfuse_public_key": "pk", "langfuse_secret_key": "sk"}], True),
|
||||
# the sibling carrier -- langfuse and langfuse_otel are one account
|
||||
({"langfuse_host": "http://attacker.invalid"}, [{"langfuse_host": "https://us.cloud.langfuse.com", "langfuse_secret_key": "sk"}], True),
|
||||
# a destination variable no integration registry lists
|
||||
({"dd_agent_host": "attacker.invalid"}, [{"dd_api_key": "k", "dd_site": "us5.datadoghq.com"}], True),
|
||||
# one entry owning its family end to end is the feature
|
||||
({"langfuse_host": "https://eu.cloud.langfuse.com", "langfuse_public_key": "pk", "langfuse_secret_key": "sk"}, [], False),
|
||||
# a different family alongside an existing one stays fine
|
||||
({"gcs_bucket_name": "bucket"}, [{"langfuse_public_key": "pk", "langfuse_secret_key": "sk"}], False),
|
||||
({"langsmith_api_key": "k"}, [{"dd_api_key": "k"}], False),
|
||||
# variables that configure no backend carry nothing to redirect
|
||||
({"turn_off_message_logging": "true"}, [{"langfuse_secret_key": "sk"}], False),
|
||||
# the same integration registered for a second event: identical values
|
||||
# flatten to the identical dict, so there is nothing to redirect
|
||||
({"langfuse_host": "https://us.cloud.langfuse.com", "langfuse_public_key": "pk", "langfuse_secret_key": "sk"}, [{"langfuse_host": "https://us.cloud.langfuse.com", "langfuse_public_key": "pk", "langfuse_secret_key": "sk"}], False),
|
||||
# the same credential under its other spelling is the same credential
|
||||
({"langfuse_secret": "sk"}, [{"langfuse_public_key": "pk", "langfuse_secret_key": "sk"}], False),
|
||||
# a value the family already holds cannot be moved into another of its
|
||||
# variables either; the exporter would address or authenticate with it
|
||||
({"langfuse_host": "pk"}, [{"langfuse_host": "https://us.cloud.langfuse.com", "langfuse_public_key": "pk"}], True),
|
||||
# the same shape with one value moved is the redirect again
|
||||
({"langfuse_host": "http://attacker.invalid", "langfuse_public_key": "pk", "langfuse_secret_key": "sk"}, [{"langfuse_host": "https://us.cloud.langfuse.com", "langfuse_public_key": "pk", "langfuse_secret_key": "sk"}], True),
|
||||
],
|
||||
)
|
||||
def test_one_entry_owns_a_credential_family(new_vars, stored, rejected):
|
||||
"""A team admin must not be able to redirect a credential they cannot read.
|
||||
|
||||
The stored entries are flattened into one dict before a request reads them,
|
||||
so an entry naming only a destination pairs with a key written elsewhere and
|
||||
carries it to that destination.
|
||||
"""
|
||||
error = cross_entry_family_error(new_vars, stored)
|
||||
assert (error is not None) is rejected
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue