ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

added the option of adding langsmith tenant id in the env (#18623)

Browse Source
This commit is contained in:
Shivam Rawat 2026-01-05 11:49:27 -08:00 committed by GitHub
parent 196509cbb1
commit 8c21fcb957
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 135 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
litellm/integrations/callback_configs.json
View File

@ -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"

19
litellm/integrations/langsmith.py
View File

@ -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()

2
litellm/types/integrations/langsmith.py
View File

@ -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

1
litellm/types/utils.py
View File

@ -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]

111
tests/logging_callback_tests/test_langsmith_unit_test.py
View File

@ -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}")