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')} /> <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] ## Custom Callback Class [Async]
Use this when you want to run custom callbacks in `python` 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 = ( default_soft_budget: float = (
50.0 # by default all litellm proxy keys have a soft budget of 50.0 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_finish_reasons = ["stop", "length", "function_call", "content_filter", "null"]
_openai_completion_params = [ _openai_completion_params = [
"functions", "functions",

View file

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

View file

@ -1,9 +1,16 @@
import asyncio import asyncio
import os
import sys
from unittest.mock import Mock from unittest.mock import Mock
import pytest import pytest
from fastapi import Request 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._types import UserAPIKeyAuth
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
from litellm.types.utils import SupportedCacheControls from litellm.types.utils import SupportedCacheControls
@ -13,7 +20,7 @@ from litellm.types.utils import SupportedCacheControls
def mock_request(monkeypatch): def mock_request(monkeypatch):
mock_request = Mock(spec=Request) mock_request = Mock(spec=Request)
mock_request.query_params = {} # Set mock query_params to an empty dictionary mock_request.query_params = {} # Set mock query_params to an empty dictionary
mock_request.headers = {} mock_request.headers = {"traceparent": "test_traceparent"}
monkeypatch.setattr( monkeypatch.setattr(
"litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request", mock_request "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 "metadata" in data
assert "litellm_metadata" not 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