diff --git a/litellm/integrations/callback_configs.json b/litellm/integrations/callback_configs.json index 88f7908e9a..6b30b6b736 100644 --- a/litellm/integrations/callback_configs.json +++ b/litellm/integrations/callback_configs.json @@ -187,6 +187,12 @@ "ui_name": "Sampling Rate", "description": "Sampling rate for logging (0.0 to 1.0, default: 1.0)", "required": false + }, + "langsmith_tenant_id": { + "type": "text", + "ui_name": "Tenant ID", + "description": "LangSmith tenant ID for organization-scoped API keys (required when using org-scoped keys)", + "required": false } }, "description": "Langsmith Logging Integration" diff --git a/litellm/integrations/langsmith.py b/litellm/integrations/langsmith.py index cc9b361b69..570b78f292 100644 --- a/litellm/integrations/langsmith.py +++ b/litellm/integrations/langsmith.py @@ -40,6 +40,7 @@ class LangsmithLogger(CustomBatchLogger): langsmith_project: Optional[str] = None, langsmith_base_url: Optional[str] = None, langsmith_sampling_rate: Optional[float] = None, + langsmith_tenant_id: Optional[str] = None, **kwargs, ): self.flush_lock = asyncio.Lock() @@ -48,6 +49,7 @@ class LangsmithLogger(CustomBatchLogger): langsmith_api_key=langsmith_api_key, langsmith_project=langsmith_project, langsmith_base_url=langsmith_base_url, + langsmith_tenant_id=langsmith_tenant_id, ) self.sampling_rate: float = ( langsmith_sampling_rate @@ -76,6 +78,7 @@ class LangsmithLogger(CustomBatchLogger): langsmith_api_key: Optional[str] = None, langsmith_project: Optional[str] = None, langsmith_base_url: Optional[str] = None, + langsmith_tenant_id: Optional[str] = None, ) -> LangsmithCredentialsObject: _credentials_api_key = langsmith_api_key or os.getenv("LANGSMITH_API_KEY") _credentials_project = ( @@ -86,11 +89,13 @@ class LangsmithLogger(CustomBatchLogger): or os.getenv("LANGSMITH_BASE_URL") or "https://api.smith.langchain.com" ) + _credentials_tenant_id = langsmith_tenant_id or os.getenv("LANGSMITH_TENANT_ID") return LangsmithCredentialsObject( LANGSMITH_API_KEY=_credentials_api_key, LANGSMITH_BASE_URL=_credentials_base_url, LANGSMITH_PROJECT=_credentials_project, + LANGSMITH_TENANT_ID=_credentials_tenant_id, ) def _prepare_log_data( @@ -365,8 +370,11 @@ class LangsmithLogger(CustomBatchLogger): """ langsmith_api_base = credentials["LANGSMITH_BASE_URL"] langsmith_api_key = credentials["LANGSMITH_API_KEY"] + langsmith_tenant_id = credentials.get("LANGSMITH_TENANT_ID") url = self._add_endpoint_to_url(langsmith_api_base, "runs/batch") headers = {"x-api-key": langsmith_api_key} + if langsmith_tenant_id: + headers["x-tenant-id"] = langsmith_tenant_id elements_to_log = [queue_object["data"] for queue_object in queue_objects] try: @@ -418,6 +426,7 @@ class LangsmithLogger(CustomBatchLogger): api_key=credentials["LANGSMITH_API_KEY"], project=credentials["LANGSMITH_PROJECT"], base_url=credentials["LANGSMITH_BASE_URL"], + tenant_id=credentials.get("LANGSMITH_TENANT_ID"), ) if key not in log_queue_by_credentials: @@ -466,6 +475,9 @@ class LangsmithLogger(CustomBatchLogger): langsmith_base_url=standard_callback_dynamic_params.get( "langsmith_base_url", None ), + langsmith_tenant_id=standard_callback_dynamic_params.get( + "langsmith_tenant_id", None + ), ) else: credentials = self.default_credentials @@ -491,13 +503,16 @@ class LangsmithLogger(CustomBatchLogger): def get_run_by_id(self, run_id): langsmith_api_key = self.default_credentials["LANGSMITH_API_KEY"] - langsmith_api_base = self.default_credentials["LANGSMITH_BASE_URL"] + langsmith_tenant_id = self.default_credentials.get("LANGSMITH_TENANT_ID") url = f"{langsmith_api_base}/runs/{run_id}" + headers = {"x-api-key": langsmith_api_key} + if langsmith_tenant_id: + headers["x-tenant-id"] = langsmith_tenant_id response = litellm.module_level_client.get( url=url, - headers={"x-api-key": langsmith_api_key}, + headers=headers, ) return response.json() diff --git a/litellm/types/integrations/langsmith.py b/litellm/types/integrations/langsmith.py index 23f760ecf3..9c026a117f 100644 --- a/litellm/types/integrations/langsmith.py +++ b/litellm/types/integrations/langsmith.py @@ -31,6 +31,7 @@ class LangsmithCredentialsObject(TypedDict): LANGSMITH_API_KEY: Optional[str] LANGSMITH_PROJECT: Optional[str] LANGSMITH_BASE_URL: str + LANGSMITH_TENANT_ID: Optional[str] class LangsmithQueueObject(TypedDict): @@ -52,6 +53,7 @@ class CredentialsKey(NamedTuple): api_key: str project: str base_url: str + tenant_id: Optional[str] @dataclass diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 3eec67d9d2..67b8a290b2 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2677,6 +2677,7 @@ class StandardCallbackDynamicParams(TypedDict, total=False): langsmith_project: Optional[str] langsmith_base_url: Optional[str] langsmith_sampling_rate: Optional[float] + langsmith_tenant_id: Optional[str] # Humanloop dynamic params humanloop_api_key: Optional[str] diff --git a/tests/logging_callback_tests/test_langsmith_unit_test.py b/tests/logging_callback_tests/test_langsmith_unit_test.py index e63ce9f8b3..bde2b94457 100644 --- a/tests/logging_callback_tests/test_langsmith_unit_test.py +++ b/tests/logging_callback_tests/test_langsmith_unit_test.py @@ -47,6 +47,19 @@ async def test_get_credentials_from_env(): credentials = logger.get_credentials_from_env() assert credentials["LANGSMITH_BASE_URL"] == "https://api.smith.langchain.com" + # Test with tenant_id + credentials = logger.get_credentials_from_env( + langsmith_tenant_id="test-tenant-id" + ) + assert credentials["LANGSMITH_TENANT_ID"] == "test-tenant-id" + + # Test tenant_id from environment variable + import os + os.environ["LANGSMITH_TENANT_ID"] = "env-tenant-id" + credentials = logger.get_credentials_from_env() + assert credentials["LANGSMITH_TENANT_ID"] == "env-tenant-id" + del os.environ["LANGSMITH_TENANT_ID"] + @pytest.mark.asyncio async def test_group_batches_by_credentials(): @@ -60,6 +73,7 @@ async def test_group_batches_by_credentials(): "LANGSMITH_API_KEY": "key1", "LANGSMITH_PROJECT": "proj1", "LANGSMITH_BASE_URL": "url1", + "LANGSMITH_TENANT_ID": None, }, ) @@ -69,6 +83,7 @@ async def test_group_batches_by_credentials(): "LANGSMITH_API_KEY": "key1", "LANGSMITH_PROJECT": "proj1", "LANGSMITH_BASE_URL": "url1", + "LANGSMITH_TENANT_ID": None, }, ) @@ -95,6 +110,7 @@ async def test_group_batches_by_credentials_multiple_credentials(): "LANGSMITH_API_KEY": "key1", "LANGSMITH_PROJECT": "proj1", "LANGSMITH_BASE_URL": "url1", + "LANGSMITH_TENANT_ID": None, }, ) @@ -104,6 +120,7 @@ async def test_group_batches_by_credentials_multiple_credentials(): "LANGSMITH_API_KEY": "key2", # Different API key "LANGSMITH_PROJECT": "proj1", "LANGSMITH_BASE_URL": "url1", + "LANGSMITH_TENANT_ID": None, }, ) @@ -113,6 +130,7 @@ async def test_group_batches_by_credentials_multiple_credentials(): "LANGSMITH_API_KEY": "key1", "LANGSMITH_PROJECT": "proj2", # Different project "LANGSMITH_BASE_URL": "url1", + "LANGSMITH_TENANT_ID": None, }, ) @@ -127,6 +145,57 @@ async def test_group_batches_by_credentials_multiple_credentials(): assert len(batch_group.queue_objects) == 1 # Each group should have one object +@pytest.mark.asyncio +async def test_group_batches_by_credentials_with_tenant_id(): + + # Test that different tenant_ids create separate groups + logger = LangsmithLogger(langsmith_api_key="test-key") + + queue_obj1 = LangsmithQueueObject( + data={"test": "data1"}, + credentials={ + "LANGSMITH_API_KEY": "key1", + "LANGSMITH_PROJECT": "proj1", + "LANGSMITH_BASE_URL": "url1", + "LANGSMITH_TENANT_ID": "tenant1", + }, + ) + + queue_obj2 = LangsmithQueueObject( + data={"test": "data2"}, + credentials={ + "LANGSMITH_API_KEY": "key1", + "LANGSMITH_PROJECT": "proj1", + "LANGSMITH_BASE_URL": "url1", + "LANGSMITH_TENANT_ID": "tenant2", # Different tenant_id + }, + ) + + queue_obj3 = LangsmithQueueObject( + data={"test": "data3"}, + credentials={ + "LANGSMITH_API_KEY": "key1", + "LANGSMITH_PROJECT": "proj1", + "LANGSMITH_BASE_URL": "url1", + "LANGSMITH_TENANT_ID": "tenant1", # Same as queue_obj1 + }, + ) + + logger.log_queue = [queue_obj1, queue_obj2, queue_obj3] + + grouped = logger._group_batches_by_credentials() + + # Should have two groups: one for tenant1 (queue_obj1 and queue_obj3), one for tenant2 (queue_obj2) + assert len(grouped) == 2 + for key, batch_group in grouped.items(): + assert isinstance(key, CredentialsKey) + assert key.tenant_id in ["tenant1", "tenant2"] + if key.tenant_id == "tenant1": + assert len(batch_group.queue_objects) == 2 + else: + assert len(batch_group.queue_objects) == 1 + + # Test make_dot_order @pytest.mark.asyncio async def test_make_dot_order(): @@ -201,10 +270,43 @@ async def test_async_send_batch(): call_args = logger.async_httpx_client.post.call_args assert "runs/batch" in call_args[1]["url"] assert "x-api-key" in call_args[1]["headers"] + # tenant_id should not be in headers if not provided + assert "x-tenant-id" not in call_args[1]["headers"] @pytest.mark.asyncio -async def test_langsmith_key_based_logging(mocker): +async def test_async_send_batch_with_tenant_id(): + logger = LangsmithLogger( + langsmith_api_key="test-key", + langsmith_tenant_id="test-tenant-id" + ) + + # Mock the httpx client + mock_response = AsyncMock() + mock_response.status_code = 200 + logger.async_httpx_client = AsyncMock() + logger.async_httpx_client.post.return_value = mock_response + + # Add test data to queue + logger.log_queue = [ + LangsmithQueueObject( + data={"test": "data"}, credentials=logger.default_credentials + ) + ] + + await logger.async_send_batch() + + # Verify the API call includes tenant_id header + logger.async_httpx_client.post.assert_called_once() + call_args = logger.async_httpx_client.post.call_args + assert "runs/batch" in call_args[1]["url"] + assert "x-api-key" in call_args[1]["headers"] + assert "x-tenant-id" in call_args[1]["headers"] + assert call_args[1]["headers"]["x-tenant-id"] == "test-tenant-id" + + +@pytest.mark.asyncio +async def test_langsmith_key_based_logging(): """ In key based logging langsmith_api_key and langsmith_project are passed directly to litellm.acompletion """ @@ -219,10 +321,11 @@ async def test_langsmith_key_based_logging(mocker): mock_response.text = "" mock_async_httpx_handler.post = AsyncMock(return_value=mock_response) - mock_get_client = mocker.patch( + mock_get_client = patch( "litellm.integrations.langsmith.get_async_httpx_client", return_value=mock_async_httpx_handler ) + mock_get_client.start() litellm.set_verbose = True litellm.DEFAULT_FLUSH_INTERVAL_SECONDS = 1 @@ -253,6 +356,8 @@ async def test_langsmith_key_based_logging(mocker): # Check headers contain the correct API key assert call_args[1]["headers"]["x-api-key"] == "fake_key_project2" + # tenant_id should not be in headers if not provided + assert "x-tenant-id" not in call_args[1]["headers"] # Verify the request body contains the expected data request_body = call_args[1]["json"] @@ -344,6 +449,8 @@ async def test_langsmith_key_based_logging(mocker): actual_body["post"][0]["session_name"] == expected_body["post"][0]["session_name"] ) + + mock_get_client.stop() except Exception as e: pytest.fail(f"Error occurred: {e}")