Merge pull request #4995 from BerriAI/litellm_otel_fwd_traceparent

[Fix-OTEL Proxy] Only forward traceparent to llm api when `litellm.forward_traceparent_to_llm_provider=True`
This commit is contained in:
Ishaan Jaff 2024-08-01 11:57:35 -07:00 committed by GitHub
commit c87f97bf6d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 74 additions and 11 deletions

View file

@ -714,6 +714,23 @@ Search for Trace=`80e1afed08e019fc1110464cfa66635c` on your OTEL Collector
<Image img={require('../../img/otel_parent.png')} />
### Forwarding `Traceparent HTTP Header` to LLM APIs
Use this if you want to forward the traceparent headers to your self hosted LLMs like vLLM
Set `forward_traceparent_to_llm_provider: True` in your `config.yaml`. This will forward the `traceparent` header to your LLM API
:::warning
Only use this for self hosted LLMs, this can cause Bedrock, VertexAI calls to fail
:::
```yaml
litellm_settings:
forward_traceparent_to_llm_provider: True
```
## Custom Callback Class [Async]
Use this when you want to run custom callbacks in `python`

View file

@ -165,6 +165,7 @@ budget_duration: Optional[str] = (
default_soft_budget: float = (
50.0 # by default all litellm proxy keys have a soft budget of 50.0
)
forward_traceparent_to_llm_provider: bool = False
_openai_finish_reasons = ["stop", "length", "function_call", "content_filter", "null"]
_openai_completion_params = [
"functions",

View file

@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Any, Dict, Optional
from fastapi import Request
import litellm
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm.proxy._types import CommonProxyErrors, TeamCallbackMetadata, UserAPIKeyAuth
from litellm.types.utils import SupportedCacheControls
@ -250,13 +251,15 @@ def _add_otel_traceparent_to_data(data: dict, request: Request):
# if user is not use OTEL don't send extra_headers
# relevant issue: https://github.com/BerriAI/litellm/issues/4448
return
if request.headers:
if "traceparent" in request.headers:
# we want to forward this to the LLM Provider
# Relevant issue: https://github.com/BerriAI/litellm/issues/4419
# pass this in extra_headers
if "extra_headers" not in data:
data["extra_headers"] = {}
_exra_headers = data["extra_headers"]
if "traceparent" not in _exra_headers:
_exra_headers["traceparent"] = request.headers["traceparent"]
if litellm.forward_traceparent_to_llm_provider is True:
if request.headers:
if "traceparent" in request.headers:
# we want to forward this to the LLM Provider
# Relevant issue: https://github.com/BerriAI/litellm/issues/4419
# pass this in extra_headers
if "extra_headers" not in data:
data["extra_headers"] = {}
_exra_headers = data["extra_headers"]
if "traceparent" not in _exra_headers:
_exra_headers["traceparent"] = request.headers["traceparent"]

View file

@ -1,9 +1,16 @@
import asyncio
import os
import sys
from unittest.mock import Mock
import pytest
from fastapi import Request
import litellm
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
from litellm.types.utils import SupportedCacheControls
@ -13,7 +20,7 @@ from litellm.types.utils import SupportedCacheControls
def mock_request(monkeypatch):
mock_request = Mock(spec=Request)
mock_request.query_params = {} # Set mock query_params to an empty dictionary
mock_request.headers = {}
mock_request.headers = {"traceparent": "test_traceparent"}
monkeypatch.setattr(
"litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request", mock_request
)
@ -60,3 +67,38 @@ async def test_add_litellm_data_to_request_non_thread_endpoint(endpoint, mock_re
assert "metadata" in data
assert "litellm_metadata" not in data
# test adding traceparent
@pytest.mark.parametrize(
"endpoint", ["/chat/completions", "/v1/completions", "/completions"]
)
@pytest.mark.asyncio
async def test_traceparent_not_added_by_default(endpoint, mock_request):
"""
This tests that traceparent is not forwarded in the extra_headers
We had an incident where bedrock calls were failing because traceparent was forwarded
"""
from litellm.integrations.opentelemetry import OpenTelemetry
otel_logger = OpenTelemetry()
setattr(litellm.proxy.proxy_server, "open_telemetry_logger", otel_logger)
mock_request.url.path = endpoint
user_api_key_dict = UserAPIKeyAuth(
api_key="test_api_key", user_id="test_user_id", org_id="test_org_id"
)
proxy_config = Mock()
data = {}
await add_litellm_data_to_request(
data, mock_request, user_api_key_dict, proxy_config
)
print("DATA: ", data)
_extra_headers = data.get("extra_headers") or {}
assert "traceparent" not in _extra_headers