mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
fix(otel): backport auth spans and callback merge to rc/1.101.0
Cherry-pick PR #40335, restoring the Datadog auth span and last-wins callback credential merging.
(cherry picked from commit 43a1b2992a)
This commit is contained in:
parent
f0b6d66b84
commit
bcddffcbdc
4 changed files with 162 additions and 6 deletions
|
|
@ -2844,7 +2844,6 @@ async def _authorize_authenticated_request(
|
|||
return None
|
||||
|
||||
|
||||
@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.
|
||||
|
||||
|
|
@ -2882,6 +2881,7 @@ def _seed_request_destinations(user_api_key_dict: UserAPIKeyAuth, request: Reque
|
|||
verbose_proxy_logger.debug("OTel V2: tenant destination resolution failed: %s", exc)
|
||||
|
||||
|
||||
@tracer.wrap()
|
||||
async def user_api_key_auth(
|
||||
request: Request,
|
||||
api_key: str = fastapi.Security(api_key_header),
|
||||
|
|
|
|||
|
|
@ -1040,7 +1040,9 @@ def resolve_tenant_otel_destinations(
|
|||
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.
|
||||
the operator's exporter is left alone. Its ``callback_vars`` still take part in the
|
||||
merge for a backend another entry made eligible, so the destination carries the
|
||||
same credentials the runtime parser resolves for that request.
|
||||
|
||||
A backend the request disabled dynamically, through the key's
|
||||
``litellm_disabled_callbacks`` or the ``x-litellm-disable-callbacks`` header in
|
||||
|
|
@ -1068,12 +1070,13 @@ def resolve_tenant_otel_destinations(
|
|||
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)
|
||||
for name in dict.fromkeys(
|
||||
callback.callback_name for callback in callbacks if callback.callback_type != "failure"
|
||||
)
|
||||
if (
|
||||
destination := destination_for(
|
||||
name,
|
||||
|
|
|
|||
|
|
@ -2,7 +2,9 @@
|
|||
|
||||
import contextvars
|
||||
import time
|
||||
from base64 import b64encode
|
||||
from collections.abc import Mapping
|
||||
from functools import reduce
|
||||
from types import MappingProxyType
|
||||
|
||||
import pytest
|
||||
|
|
@ -49,8 +51,11 @@ from litellm.integrations.otel.presets.destinations import (
|
|||
destination_for,
|
||||
)
|
||||
from litellm.integrations.otel.presets.langfuse import langfuse_preset
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.litellm_pre_call_utils import resolve_tenant_otel_destinations
|
||||
from litellm.proxy._types import AddTeamCallback, UserAPIKeyAuth
|
||||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
convert_key_logging_metadata_to_callback,
|
||||
resolve_tenant_otel_destinations,
|
||||
)
|
||||
from litellm.types.utils import StandardCallbackDynamicParams
|
||||
|
||||
LANGFUSE_DEST = OtelDestination(
|
||||
|
|
@ -1822,6 +1827,26 @@ class TestTenantConfigAgreement:
|
|||
|
||||
assert [d.endpoint for d in destinations] == ["http://key.local/api/public/otel"]
|
||||
|
||||
def test_a_failure_entry_still_wins_the_merge_next_to_a_success_entry(self):
|
||||
entries = [
|
||||
{**self._entry("http://team.local"), "callback_type": "success"},
|
||||
{
|
||||
**self._entry("http://key.local", langfuse_public_key="pk-failure", langfuse_secret_key="sk-failure"),
|
||||
"callback_type": "failure",
|
||||
},
|
||||
]
|
||||
runtime = reduce(
|
||||
lambda merged, entry: convert_key_logging_metadata_to_callback(AddTeamCallback(**entry), merged),
|
||||
entries,
|
||||
None,
|
||||
)
|
||||
|
||||
destinations = resolve_tenant_otel_destinations(UserAPIKeyAuth(team_metadata={"logging": entries}))
|
||||
|
||||
assert runtime.callback_vars["langfuse_host"] == "http://key.local"
|
||||
assert [d.endpoint for d in destinations] == ["http://key.local/api/public/otel"]
|
||||
assert destinations[0].headers["Authorization"] == f"Basic {b64encode(b'pk-failure:sk-failure').decode()}"
|
||||
|
||||
@pytest.fixture
|
||||
def premium(self, monkeypatch):
|
||||
from litellm.proxy import proxy_server
|
||||
|
|
|
|||
|
|
@ -1,7 +1,12 @@
|
|||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from contextlib import contextmanager
|
||||
from datetime import datetime, timedelta
|
||||
from pathlib import Path
|
||||
from textwrap import dedent
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import ANY, AsyncMock, MagicMock, patch
|
||||
|
||||
|
|
@ -6872,3 +6877,126 @@ class TestLitellmReceivedAtStamping:
|
|||
|
||||
assert result == earlier
|
||||
assert request.state.litellm_received_at == earlier
|
||||
|
||||
|
||||
_RECORDING_DDTRACE = dedent(
|
||||
'''
|
||||
import functools
|
||||
import inspect
|
||||
|
||||
|
||||
class _Span:
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *exc):
|
||||
return None
|
||||
|
||||
|
||||
class _Tracer:
|
||||
def __init__(self):
|
||||
self.spans = []
|
||||
|
||||
def wrap(self, name=None, **kwargs):
|
||||
def decorator(f):
|
||||
span_name = name or f"{f.__module__}.{f.__name__}"
|
||||
if inspect.iscoroutinefunction(f):
|
||||
|
||||
@functools.wraps(f)
|
||||
async def async_wrapped(*args, **kw):
|
||||
self.spans.append(span_name)
|
||||
return await f(*args, **kw)
|
||||
|
||||
return async_wrapped
|
||||
|
||||
@functools.wraps(f)
|
||||
def wrapped(*args, **kw):
|
||||
self.spans.append(span_name)
|
||||
return f(*args, **kw)
|
||||
|
||||
return wrapped
|
||||
|
||||
return decorator
|
||||
|
||||
def trace(self, name, **kwargs):
|
||||
return _Span()
|
||||
|
||||
def current_span(self):
|
||||
return None
|
||||
|
||||
def current_root_span(self):
|
||||
return None
|
||||
|
||||
|
||||
tracer = _Tracer()
|
||||
'''
|
||||
)
|
||||
|
||||
_DDTRACE_AUTH_PROBE = dedent(
|
||||
'''
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
import ddtrace
|
||||
from fastapi import Request
|
||||
from starlette.datastructures import URL
|
||||
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
proxy_server.master_key = "sk-probe"
|
||||
|
||||
|
||||
async def auth(api_key):
|
||||
request = Request(scope={"type": "http", "headers": [], "method": "POST", "path": "/chat/completions"})
|
||||
request._url = URL(url="/chat/completions")
|
||||
try:
|
||||
await user_api_key_auth(
|
||||
request=request,
|
||||
api_key=api_key,
|
||||
azure_api_key_header="",
|
||||
anthropic_api_key_header=None,
|
||||
google_ai_studio_api_key_header=None,
|
||||
azure_apim_header=None,
|
||||
custom_litellm_key_header=None,
|
||||
)
|
||||
return "accepted"
|
||||
except ProxyException:
|
||||
return "rejected"
|
||||
|
||||
|
||||
async def main():
|
||||
outcomes = [await auth("Bearer sk-probe"), await auth("Bearer sk-wrong")]
|
||||
print(json.dumps({"outcomes": outcomes, "spans": ddtrace.tracer.spans}))
|
||||
|
||||
|
||||
asyncio.run(main())
|
||||
'''
|
||||
)
|
||||
|
||||
|
||||
def test_user_api_key_auth_opens_a_datadog_span_for_accepted_and_rejected_keys(tmp_path: Path):
|
||||
stub_root = tmp_path / "site"
|
||||
(stub_root / "ddtrace").mkdir(parents=True)
|
||||
(stub_root / "ddtrace" / "__init__.py").write_text(_RECORDING_DDTRACE)
|
||||
probe = tmp_path / "probe.py"
|
||||
probe.write_text(_DDTRACE_AUTH_PROBE)
|
||||
repo_root = Path(litellm.__file__).resolve().parent.parent
|
||||
env = {
|
||||
**os.environ,
|
||||
"USE_DDTRACE": "true",
|
||||
"PYTHONPATH": os.pathsep.join(
|
||||
[str(stub_root), str(repo_root)] + [p for p in (os.environ.get("PYTHONPATH"),) if p]
|
||||
),
|
||||
}
|
||||
|
||||
result = subprocess.run(
|
||||
[sys.executable, str(probe)], env=env, cwd=repo_root, capture_output=True, text=True, check=False
|
||||
)
|
||||
|
||||
assert result.returncode == 0, result.stderr[-4000:]
|
||||
report = json.loads(result.stdout.strip().splitlines()[-1])
|
||||
assert report["outcomes"] == ["accepted", "rejected"]
|
||||
auth_span = "litellm.proxy.auth.user_api_key_auth.user_api_key_auth"
|
||||
assert [span for span in report["spans"] if span == auth_span] == [auth_span, auth_span]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue