diff --git a/langfuse/_client/client.py b/langfuse/_client/client.py index 791de5a3f..3a8a902c7 100644 --- a/langfuse/_client/client.py +++ b/langfuse/_client/client.py @@ -105,6 +105,7 @@ class Langfuse: 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. mask (Optional[MaskFunction]): Function to mask sensitive data in traces before sending to the API. 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`) + 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. Example: ```python @@ -163,6 +164,7 @@ def __init__( sample_rate: Optional[float] = None, mask: Optional[MaskFunction] = None, blocked_instrumentation_scopes: Optional[List[str]] = None, + additional_headers: Optional[Dict[str, str]] = None, ): self._host = host or os.environ.get(LANGFUSE_HOST, "https://cloud.langfuse.com") self._environment = environment or os.environ.get(LANGFUSE_TRACING_ENVIRONMENT) @@ -225,6 +227,7 @@ def __init__( mask=mask, tracing_enabled=self._tracing_enabled, blocked_instrumentation_scopes=blocked_instrumentation_scopes, + additional_headers=additional_headers, ) self._mask = self._resources.mask diff --git a/langfuse/_client/resource_manager.py b/langfuse/_client/resource_manager.py index 3adac223f..041d2ec05 100644 --- a/langfuse/_client/resource_manager.py +++ b/langfuse/_client/resource_manager.py @@ -94,6 +94,7 @@ def __new__( mask: Optional[MaskFunction] = None, tracing_enabled: Optional[bool] = None, blocked_instrumentation_scopes: Optional[List[str]] = None, + additional_headers: Optional[Dict[str, str]] = None, ) -> "LangfuseResourceManager": if public_key in cls._instances: return cls._instances[public_key] @@ -119,6 +120,7 @@ def __new__( if tracing_enabled is not None else True, blocked_instrumentation_scopes=blocked_instrumentation_scopes, + additional_headers=additional_headers, ) cls._instances[public_key] = instance @@ -142,6 +144,7 @@ def _initialize_instance( mask: Optional[MaskFunction] = None, tracing_enabled: bool = True, blocked_instrumentation_scopes: Optional[List[str]] = None, + additional_headers: Optional[Dict[str, str]] = None, ): self.public_key = public_key self.secret_key = secret_key @@ -163,6 +166,7 @@ def _initialize_instance( flush_at=flush_at, flush_interval=flush_interval, blocked_instrumentation_scopes=blocked_instrumentation_scopes, + additional_headers=additional_headers, ) tracer_provider.add_span_processor(langfuse_processor) @@ -179,7 +183,13 @@ def _initialize_instance( ## use connection pools with limited capacity. Creating multiple instances ## could exhaust the OS's maximum number of available TCP sockets (file descriptors), ## leading to connection errors. - self.httpx_client = httpx_client or httpx.Client(timeout=timeout) + if httpx_client is not None: + self.httpx_client = httpx_client + else: + # Create a new httpx client with additional_headers if provided + client_headers = additional_headers if additional_headers else {} + self.httpx_client = httpx.Client(timeout=timeout, headers=client_headers) + self.api = FernLangfuse( base_url=host, username=self.public_key, diff --git a/langfuse/_client/span_processor.py b/langfuse/_client/span_processor.py index 13451b6de..ca8fb9b5a 100644 --- a/langfuse/_client/span_processor.py +++ b/langfuse/_client/span_processor.py @@ -13,7 +13,7 @@ import base64 import os -from typing import List, Optional +from typing import Dict, List, Optional from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter from opentelemetry.sdk.trace import ReadableSpan @@ -56,6 +56,7 @@ def __init__( flush_at: Optional[int] = None, flush_interval: Optional[float] = None, blocked_instrumentation_scopes: Optional[List[str]] = None, + additional_headers: Optional[Dict[str, str]] = None, ): self.public_key = public_key self.blocked_instrumentation_scopes = ( @@ -78,14 +79,20 @@ def __init__( f"{public_key}:{secret_key}".encode("utf-8") ).decode("ascii") + # Prepare default headers + default_headers = { + "Authorization": basic_auth_header, + "x_langfuse_sdk_name": "python", + "x_langfuse_sdk_version": langfuse_version, + "x_langfuse_public_key": public_key, + } + + # Merge additional headers if provided + headers = {**default_headers, **(additional_headers or {})} + langfuse_span_exporter = OTLPSpanExporter( endpoint=f"{host}/api/public/otel/v1/traces", - headers={ - "Authorization": basic_auth_header, - "x_langfuse_sdk_name": "python", - "x_langfuse_sdk_version": langfuse_version, - "x_langfuse_public_key": public_key, - }, + headers=headers, timeout=timeout, ) diff --git a/tests/test_additional_headers_simple.py b/tests/test_additional_headers_simple.py new file mode 100644 index 000000000..4298a3ba5 --- /dev/null +++ b/tests/test_additional_headers_simple.py @@ -0,0 +1,183 @@ +"""Simplified tests for additional_headers functionality in Langfuse client. + +This module tests that additional headers are properly configured in the HTTP clients. +""" + +import httpx + +from langfuse._client.client import Langfuse + + +class TestAdditionalHeadersSimple: + """Simple test suite for additional_headers functionality.""" + + def teardown_method(self): + """Clean up after each test to avoid singleton interference.""" + from langfuse._client.resource_manager import LangfuseResourceManager + + LangfuseResourceManager.reset() + + def test_httpx_client_has_additional_headers_when_none_provided(self): + """Test that additional headers are set in httpx client when no custom client is provided.""" + additional_headers = { + "X-Custom-Header": "custom-value", + "X-Another-Header": "another-value", + } + + langfuse = Langfuse( + public_key="test-public-key", + secret_key="test-secret-key", + host="https://mock-host.com", + additional_headers=additional_headers, + tracing_enabled=False, # Disable tracing to avoid OTEL setup + ) + + # Verify the httpx client has the additional headers + assert ( + langfuse._resources.httpx_client.headers["X-Custom-Header"] + == "custom-value" + ) + assert ( + langfuse._resources.httpx_client.headers["X-Another-Header"] + == "another-value" + ) + + def test_custom_httpx_client_with_additional_headers_ignores_additional_headers( + self, + ): + """Test that when additional headers are provided with custom client, additional headers are ignored.""" + # Create a custom httpx client with headers + existing_headers = {"X-Existing-Header": "existing-value"} + custom_client = httpx.Client(headers=existing_headers) + + additional_headers = { + "X-Custom-Header": "custom-value", + "X-Another-Header": "another-value", + } + + langfuse = Langfuse( + public_key="test-public-key", + secret_key="test-secret-key", + host="https://mock-host.com", + httpx_client=custom_client, + additional_headers=additional_headers, + tracing_enabled=False, + ) + + # Verify the original client is used (same instance) + assert langfuse._resources.httpx_client is custom_client + + # Verify existing headers are preserved and additional headers are NOT added + assert ( + langfuse._resources.httpx_client.headers["x-existing-header"] + == "existing-value" + ) + + # Additional headers should NOT be present + assert "x-custom-header" not in langfuse._resources.httpx_client.headers + assert "x-another-header" not in langfuse._resources.httpx_client.headers + + def test_custom_httpx_client_without_additional_headers_preserves_client(self): + """Test that when no additional headers are provided, the custom client is preserved.""" + # Create a custom httpx client with headers + existing_headers = {"X-Existing-Header": "existing-value"} + custom_client = httpx.Client(headers=existing_headers) + + langfuse = Langfuse( + public_key="test-public-key", + secret_key="test-secret-key", + host="https://mock-host.com", + httpx_client=custom_client, + additional_headers=None, # No additional headers + tracing_enabled=False, + ) + + # Note: The client instance might be different due to Fern API wrapper behavior, + # but the important thing is that the headers are preserved + # Verify existing headers are preserved + assert ( + langfuse._resources.httpx_client.headers["x-existing-header"] + == "existing-value" + ) + + def test_none_additional_headers_works(self): + """Test that passing None for additional_headers works without errors.""" + langfuse = Langfuse( + public_key="test-public-key", + secret_key="test-secret-key", + host="https://mock-host.com", + additional_headers=None, + tracing_enabled=False, + ) + + # Verify client was created successfully + assert langfuse is not None + assert langfuse._resources is not None + assert langfuse._resources.httpx_client is not None + + def test_empty_additional_headers_works(self): + """Test that passing an empty dict for additional_headers works.""" + langfuse = Langfuse( + public_key="test-public-key", + secret_key="test-secret-key", + host="https://mock-host.com", + additional_headers={}, + tracing_enabled=False, + ) + + # Verify client was created successfully + assert langfuse is not None + assert langfuse._resources is not None + assert langfuse._resources.httpx_client is not None + + def test_span_processor_has_additional_headers_in_otel_exporter(self): + """Test that span processor includes additional headers in OTEL exporter.""" + from langfuse._client.span_processor import LangfuseSpanProcessor + + additional_headers = { + "X-Custom-Trace-Header": "trace-value", + "X-Override-Default": "override-value", + } + + # Create span processor with additional headers + processor = LangfuseSpanProcessor( + public_key="test-public-key", + secret_key="test-secret-key", + host="https://mock-host.com", + additional_headers=additional_headers, + ) + + # Get the OTLP span exporter to check its headers + exporter = processor.span_exporter + + # Verify additional headers are in the exporter's headers + assert exporter._headers["X-Custom-Trace-Header"] == "trace-value" + assert exporter._headers["X-Override-Default"] == "override-value" + + # Verify default headers are still present + assert "Authorization" in exporter._headers + assert "x_langfuse_sdk_name" in exporter._headers + assert "x_langfuse_public_key" in exporter._headers + + # Check that our override worked + assert exporter._headers["X-Override-Default"] == "override-value" + + def test_span_processor_none_additional_headers_works(self): + """Test that span processor works with None additional headers.""" + from langfuse._client.span_processor import LangfuseSpanProcessor + + # Create span processor without additional headers + processor = LangfuseSpanProcessor( + public_key="test-public-key", + secret_key="test-secret-key", + host="https://mock-host.com", + additional_headers=None, + ) + + # Get the OTLP span exporter + exporter = processor.span_exporter + + # Verify default headers are present + assert "Authorization" in exporter._headers + assert "x_langfuse_sdk_name" in exporter._headers + assert "x_langfuse_public_key" in exporter._headers