diff --git a/langfuse/_client/get_client.py b/langfuse/_client/get_client.py index ff619095e..8bcdecd40 100644 --- a/langfuse/_client/get_client.py +++ b/langfuse/_client/get_client.py @@ -33,6 +33,28 @@ def _set_current_public_key(public_key: Optional[str]) -> Iterator[None]: _current_public_key.reset(token) +def _create_client_from_instance( + instance: "LangfuseResourceManager", public_key: Optional[str] = None +) -> Langfuse: + """Create a Langfuse client from a resource manager instance with all settings preserved.""" + return Langfuse( + public_key=public_key or instance.public_key, + secret_key=instance.secret_key, + host=instance.host, + tracing_enabled=instance.tracing_enabled, + environment=instance.environment, + timeout=instance.timeout, + flush_at=instance.flush_at, + flush_interval=instance.flush_interval, + release=instance.release, + media_upload_thread_count=instance.media_upload_thread_count, + sample_rate=instance.sample_rate, + mask=instance.mask, + blocked_instrumentation_scopes=instance.blocked_instrumentation_scopes, + additional_headers=instance.additional_headers, + ) + + def get_client(*, public_key: Optional[str] = None) -> Langfuse: """Get or create a Langfuse client instance. @@ -93,12 +115,7 @@ def get_client(*, public_key: Optional[str] = None) -> Langfuse: # Initialize with the credentials bound to the instance # This is important if the original instance was instantiated # via constructor arguments - return Langfuse( - public_key=instance.public_key, - secret_key=instance.secret_key, - host=instance.host, - tracing_enabled=instance.tracing_enabled, - ) + return _create_client_from_instance(instance) else: # Multiple clients exist but no key specified - disable tracing @@ -126,9 +143,4 @@ def get_client(*, public_key: Optional[str] = None) -> Langfuse: ) # target_instance is guaranteed to be not None at this point - return Langfuse( - public_key=public_key, - secret_key=target_instance.secret_key, - host=target_instance.host, - tracing_enabled=target_instance.tracing_enabled, - ) + return _create_client_from_instance(target_instance, public_key) diff --git a/langfuse/_client/resource_manager.py b/langfuse/_client/resource_manager.py index e0e3cbadc..70ed5b17c 100644 --- a/langfuse/_client/resource_manager.py +++ b/langfuse/_client/resource_manager.py @@ -162,6 +162,17 @@ def _initialize_instance( self.tracing_enabled = tracing_enabled self.host = host self.mask = mask + self.environment = environment + + # Store additional client settings for get_client() to use + self.timeout = timeout + self.flush_at = flush_at + self.flush_interval = flush_interval + self.release = release + self.media_upload_thread_count = media_upload_thread_count + self.sample_rate = sample_rate + self.blocked_instrumentation_scopes = blocked_instrumentation_scopes + self.additional_headers = additional_headers # OTEL Tracer if tracing_enabled: diff --git a/tests/test_resource_manager.py b/tests/test_resource_manager.py new file mode 100644 index 000000000..fa6eb56bf --- /dev/null +++ b/tests/test_resource_manager.py @@ -0,0 +1,79 @@ +"""Test the LangfuseResourceManager and get_client() function.""" + +from langfuse import Langfuse +from langfuse._client.get_client import get_client +from langfuse._client.resource_manager import LangfuseResourceManager + + +def test_get_client_preserves_all_settings(): + """Test that get_client() preserves environment and all client settings.""" + with LangfuseResourceManager._lock: + LangfuseResourceManager._instances.clear() + + settings = { + "environment": "test-env", + "release": "v1.2.3", + "timeout": 30, + "flush_at": 100, + "sample_rate": 0.8, + "additional_headers": {"X-Custom": "value"}, + } + + original_client = Langfuse(**settings) + retrieved_client = get_client() + + assert retrieved_client._environment == settings["environment"] + + assert retrieved_client._resources is not None + rm = retrieved_client._resources + assert rm.environment == settings["environment"] + assert rm.timeout == settings["timeout"] + assert rm.sample_rate == settings["sample_rate"] + assert rm.additional_headers == settings["additional_headers"] + + original_client.shutdown() + + +def test_get_client_multiple_clients_preserve_different_settings(): + """Test that get_client() preserves different settings for multiple clients.""" + # Settings for client A + settings_a = { + "public_key": "pk-comprehensive-a", + "secret_key": "sk-comprehensive-a", + "environment": "env-a", + "release": "release-a", + "timeout": 10, + "sample_rate": 0.5, + } + + # Settings for client B + settings_b = { + "public_key": "pk-comprehensive-b", + "secret_key": "sk-comprehensive-b", + "environment": "env-b", + "release": "release-b", + "timeout": 20, + "sample_rate": 0.9, + } + + client_a = Langfuse(**settings_a) + client_b = Langfuse(**settings_b) + + # Get clients via get_client() + retrieved_a = get_client(public_key="pk-comprehensive-a") + retrieved_b = get_client(public_key="pk-comprehensive-b") + + # Verify each client preserves its own settings + assert retrieved_a._environment == settings_a["environment"] + assert retrieved_b._environment == settings_b["environment"] + + if retrieved_a._resources and retrieved_b._resources: + assert retrieved_a._resources.timeout == settings_a["timeout"] + assert retrieved_b._resources.timeout == settings_b["timeout"] + assert retrieved_a._resources.sample_rate == settings_a["sample_rate"] + assert retrieved_b._resources.sample_rate == settings_b["sample_rate"] + assert retrieved_a._resources.release == settings_a["release"] + assert retrieved_b._resources.release == settings_b["release"] + + client_a.shutdown() + client_b.shutdown()