Skip to content

Commit 358fbd2

Browse files
authored
feat(client): allow setting additional headers (#1248)
1 parent acfc5a0 commit 358fbd2

File tree

4 files changed

+211
-8
lines changed

4 files changed

+211
-8
lines changed

langfuse/_client/client.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,7 @@ class Langfuse:
105105
sample_rate (Optional[float]): Sampling rate for traces (0.0 to 1.0). Defaults to 1.0 (100% of traces are sampled). Can also be set via LANGFUSE_SAMPLE_RATE environment variable.
106106
mask (Optional[MaskFunction]): Function to mask sensitive data in traces before sending to the API.
107107
blocked_instrumentation_scopes (Optional[List[str]]): List of instrumentation scope names to block from being exported to Langfuse. Spans from these scopes will be filtered out before being sent to the API. Useful for filtering out spans from specific libraries or frameworks. For exported spans, you can see the instrumentation scope name in the span metadata in Langfuse (`metadata.scope.name`)
108+
additional_headers (Optional[Dict[str, str]]): Additional headers to include in all API requests and OTLPSpanExporter requests. These headers will be merged with default headers. Note: If httpx_client is provided, additional_headers must be set directly on your custom httpx_client as well.
108109
109110
Example:
110111
```python
@@ -163,6 +164,7 @@ def __init__(
163164
sample_rate: Optional[float] = None,
164165
mask: Optional[MaskFunction] = None,
165166
blocked_instrumentation_scopes: Optional[List[str]] = None,
167+
additional_headers: Optional[Dict[str, str]] = None,
166168
):
167169
self._host = host or os.environ.get(LANGFUSE_HOST, "https://cloud.langfuse.com")
168170
self._environment = environment or os.environ.get(LANGFUSE_TRACING_ENVIRONMENT)
@@ -225,6 +227,7 @@ def __init__(
225227
mask=mask,
226228
tracing_enabled=self._tracing_enabled,
227229
blocked_instrumentation_scopes=blocked_instrumentation_scopes,
230+
additional_headers=additional_headers,
228231
)
229232
self._mask = self._resources.mask
230233

langfuse/_client/resource_manager.py

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,7 @@ def __new__(
9494
mask: Optional[MaskFunction] = None,
9595
tracing_enabled: Optional[bool] = None,
9696
blocked_instrumentation_scopes: Optional[List[str]] = None,
97+
additional_headers: Optional[Dict[str, str]] = None,
9798
) -> "LangfuseResourceManager":
9899
if public_key in cls._instances:
99100
return cls._instances[public_key]
@@ -119,6 +120,7 @@ def __new__(
119120
if tracing_enabled is not None
120121
else True,
121122
blocked_instrumentation_scopes=blocked_instrumentation_scopes,
123+
additional_headers=additional_headers,
122124
)
123125

124126
cls._instances[public_key] = instance
@@ -142,6 +144,7 @@ def _initialize_instance(
142144
mask: Optional[MaskFunction] = None,
143145
tracing_enabled: bool = True,
144146
blocked_instrumentation_scopes: Optional[List[str]] = None,
147+
additional_headers: Optional[Dict[str, str]] = None,
145148
):
146149
self.public_key = public_key
147150
self.secret_key = secret_key
@@ -163,6 +166,7 @@ def _initialize_instance(
163166
flush_at=flush_at,
164167
flush_interval=flush_interval,
165168
blocked_instrumentation_scopes=blocked_instrumentation_scopes,
169+
additional_headers=additional_headers,
166170
)
167171
tracer_provider.add_span_processor(langfuse_processor)
168172

@@ -179,7 +183,13 @@ def _initialize_instance(
179183
## use connection pools with limited capacity. Creating multiple instances
180184
## could exhaust the OS's maximum number of available TCP sockets (file descriptors),
181185
## leading to connection errors.
182-
self.httpx_client = httpx_client or httpx.Client(timeout=timeout)
186+
if httpx_client is not None:
187+
self.httpx_client = httpx_client
188+
else:
189+
# Create a new httpx client with additional_headers if provided
190+
client_headers = additional_headers if additional_headers else {}
191+
self.httpx_client = httpx.Client(timeout=timeout, headers=client_headers)
192+
183193
self.api = FernLangfuse(
184194
base_url=host,
185195
username=self.public_key,

langfuse/_client/span_processor.py

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313

1414
import base64
1515
import os
16-
from typing import List, Optional
16+
from typing import Dict, List, Optional
1717

1818
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
1919
from opentelemetry.sdk.trace import ReadableSpan
@@ -56,6 +56,7 @@ def __init__(
5656
flush_at: Optional[int] = None,
5757
flush_interval: Optional[float] = None,
5858
blocked_instrumentation_scopes: Optional[List[str]] = None,
59+
additional_headers: Optional[Dict[str, str]] = None,
5960
):
6061
self.public_key = public_key
6162
self.blocked_instrumentation_scopes = (
@@ -78,14 +79,20 @@ def __init__(
7879
f"{public_key}:{secret_key}".encode("utf-8")
7980
).decode("ascii")
8081

82+
# Prepare default headers
83+
default_headers = {
84+
"Authorization": basic_auth_header,
85+
"x_langfuse_sdk_name": "python",
86+
"x_langfuse_sdk_version": langfuse_version,
87+
"x_langfuse_public_key": public_key,
88+
}
89+
90+
# Merge additional headers if provided
91+
headers = {**default_headers, **(additional_headers or {})}
92+
8193
langfuse_span_exporter = OTLPSpanExporter(
8294
endpoint=f"{host}/api/public/otel/v1/traces",
83-
headers={
84-
"Authorization": basic_auth_header,
85-
"x_langfuse_sdk_name": "python",
86-
"x_langfuse_sdk_version": langfuse_version,
87-
"x_langfuse_public_key": public_key,
88-
},
95+
headers=headers,
8996
timeout=timeout,
9097
)
9198

Lines changed: 183 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,183 @@
1+
"""Simplified tests for additional_headers functionality in Langfuse client.
2+
3+
This module tests that additional headers are properly configured in the HTTP clients.
4+
"""
5+
6+
import httpx
7+
8+
from langfuse._client.client import Langfuse
9+
10+
11+
class TestAdditionalHeadersSimple:
12+
"""Simple test suite for additional_headers functionality."""
13+
14+
def teardown_method(self):
15+
"""Clean up after each test to avoid singleton interference."""
16+
from langfuse._client.resource_manager import LangfuseResourceManager
17+
18+
LangfuseResourceManager.reset()
19+
20+
def test_httpx_client_has_additional_headers_when_none_provided(self):
21+
"""Test that additional headers are set in httpx client when no custom client is provided."""
22+
additional_headers = {
23+
"X-Custom-Header": "custom-value",
24+
"X-Another-Header": "another-value",
25+
}
26+
27+
langfuse = Langfuse(
28+
public_key="test-public-key",
29+
secret_key="test-secret-key",
30+
host="https://mock-host.com",
31+
additional_headers=additional_headers,
32+
tracing_enabled=False, # Disable tracing to avoid OTEL setup
33+
)
34+
35+
# Verify the httpx client has the additional headers
36+
assert (
37+
langfuse._resources.httpx_client.headers["X-Custom-Header"]
38+
== "custom-value"
39+
)
40+
assert (
41+
langfuse._resources.httpx_client.headers["X-Another-Header"]
42+
== "another-value"
43+
)
44+
45+
def test_custom_httpx_client_with_additional_headers_ignores_additional_headers(
46+
self,
47+
):
48+
"""Test that when additional headers are provided with custom client, additional headers are ignored."""
49+
# Create a custom httpx client with headers
50+
existing_headers = {"X-Existing-Header": "existing-value"}
51+
custom_client = httpx.Client(headers=existing_headers)
52+
53+
additional_headers = {
54+
"X-Custom-Header": "custom-value",
55+
"X-Another-Header": "another-value",
56+
}
57+
58+
langfuse = Langfuse(
59+
public_key="test-public-key",
60+
secret_key="test-secret-key",
61+
host="https://mock-host.com",
62+
httpx_client=custom_client,
63+
additional_headers=additional_headers,
64+
tracing_enabled=False,
65+
)
66+
67+
# Verify the original client is used (same instance)
68+
assert langfuse._resources.httpx_client is custom_client
69+
70+
# Verify existing headers are preserved and additional headers are NOT added
71+
assert (
72+
langfuse._resources.httpx_client.headers["x-existing-header"]
73+
== "existing-value"
74+
)
75+
76+
# Additional headers should NOT be present
77+
assert "x-custom-header" not in langfuse._resources.httpx_client.headers
78+
assert "x-another-header" not in langfuse._resources.httpx_client.headers
79+
80+
def test_custom_httpx_client_without_additional_headers_preserves_client(self):
81+
"""Test that when no additional headers are provided, the custom client is preserved."""
82+
# Create a custom httpx client with headers
83+
existing_headers = {"X-Existing-Header": "existing-value"}
84+
custom_client = httpx.Client(headers=existing_headers)
85+
86+
langfuse = Langfuse(
87+
public_key="test-public-key",
88+
secret_key="test-secret-key",
89+
host="https://mock-host.com",
90+
httpx_client=custom_client,
91+
additional_headers=None, # No additional headers
92+
tracing_enabled=False,
93+
)
94+
95+
# Note: The client instance might be different due to Fern API wrapper behavior,
96+
# but the important thing is that the headers are preserved
97+
# Verify existing headers are preserved
98+
assert (
99+
langfuse._resources.httpx_client.headers["x-existing-header"]
100+
== "existing-value"
101+
)
102+
103+
def test_none_additional_headers_works(self):
104+
"""Test that passing None for additional_headers works without errors."""
105+
langfuse = Langfuse(
106+
public_key="test-public-key",
107+
secret_key="test-secret-key",
108+
host="https://mock-host.com",
109+
additional_headers=None,
110+
tracing_enabled=False,
111+
)
112+
113+
# Verify client was created successfully
114+
assert langfuse is not None
115+
assert langfuse._resources is not None
116+
assert langfuse._resources.httpx_client is not None
117+
118+
def test_empty_additional_headers_works(self):
119+
"""Test that passing an empty dict for additional_headers works."""
120+
langfuse = Langfuse(
121+
public_key="test-public-key",
122+
secret_key="test-secret-key",
123+
host="https://mock-host.com",
124+
additional_headers={},
125+
tracing_enabled=False,
126+
)
127+
128+
# Verify client was created successfully
129+
assert langfuse is not None
130+
assert langfuse._resources is not None
131+
assert langfuse._resources.httpx_client is not None
132+
133+
def test_span_processor_has_additional_headers_in_otel_exporter(self):
134+
"""Test that span processor includes additional headers in OTEL exporter."""
135+
from langfuse._client.span_processor import LangfuseSpanProcessor
136+
137+
additional_headers = {
138+
"X-Custom-Trace-Header": "trace-value",
139+
"X-Override-Default": "override-value",
140+
}
141+
142+
# Create span processor with additional headers
143+
processor = LangfuseSpanProcessor(
144+
public_key="test-public-key",
145+
secret_key="test-secret-key",
146+
host="https://mock-host.com",
147+
additional_headers=additional_headers,
148+
)
149+
150+
# Get the OTLP span exporter to check its headers
151+
exporter = processor.span_exporter
152+
153+
# Verify additional headers are in the exporter's headers
154+
assert exporter._headers["X-Custom-Trace-Header"] == "trace-value"
155+
assert exporter._headers["X-Override-Default"] == "override-value"
156+
157+
# Verify default headers are still present
158+
assert "Authorization" in exporter._headers
159+
assert "x_langfuse_sdk_name" in exporter._headers
160+
assert "x_langfuse_public_key" in exporter._headers
161+
162+
# Check that our override worked
163+
assert exporter._headers["X-Override-Default"] == "override-value"
164+
165+
def test_span_processor_none_additional_headers_works(self):
166+
"""Test that span processor works with None additional headers."""
167+
from langfuse._client.span_processor import LangfuseSpanProcessor
168+
169+
# Create span processor without additional headers
170+
processor = LangfuseSpanProcessor(
171+
public_key="test-public-key",
172+
secret_key="test-secret-key",
173+
host="https://mock-host.com",
174+
additional_headers=None,
175+
)
176+
177+
# Get the OTLP span exporter
178+
exporter = processor.span_exporter
179+
180+
# Verify default headers are present
181+
assert "Authorization" in exporter._headers
182+
assert "x_langfuse_sdk_name" in exporter._headers
183+
assert "x_langfuse_public_key" in exporter._headers

0 commit comments

Comments
 (0)