scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-12 03:48:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 57
litellm/tests/unit/integrations/test_langsmith.py
devin-ai-integration[bot] dfe4df8c8a
test: move whole-unit legacy test files into tests/unit and delete dead skips (#44807) 
* test: delete unconditionally skipped legacy tests

* test: move whole-unit legacy test files into tests/unit

* test: keep moved legacy tests free of import-time global state

* test: keep the module-level invocation scan pointed at tests/local_testing

* ci: drop the agent_testing CircleCI job emptied by the move

* test: fix moved-test isolation and router coverage

* test: add Tinyfish search package marker

* test: isolate moved tests from logger state leaks

* test: isolate Helicone logging fixture state

* test: isolate Vertex pass-through credentials between moved tests

* test: cancel S3 periodic flush tasks started by moved tests

* ci: restore CircleCI assistant test selection after move

* test: prevent Datadog datetime import shadowing

* ci: drop the litellm_assistants_api_testing CircleCI job emptied by the move

* test: deduplicate imports in rebased unit tests

* test: remove duplicate passthrough router patch import

* test: remove moved legacy source files after rebase

* test: align moved tests with rebased main

* test: carry main's legacy-file edits into moved destinations

* test: make the moved cost map fallback tests assert the fetch and the backup

The four fallback cases only checked the result was non-empty, so they still
passed with integrity validation disabled. They now inject a mock client, assert
one fetch happened, and assert the result is exactly the local backup with the
fallback reason recorded.

---------

Co-authored-by: yuneng <yuneng@berri.ai>
2026-10-07 05:02:28 +00:00

646 lines
23 KiB
Python
Raw Blame History

import asyncio
import importlib
import json
import os
from collections.abc import AsyncIterator
from datetime import datetime
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
import pytest_asyncio
import litellm
from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE
from litellm.integrations.langsmith import CredentialsKey, LangsmithLogger, LangsmithQueueObject
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
@pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown")
@pytest.mark.asyncio
async def test_get_credentials_from_env_does_not_use_env_for_dynamic_base_url(
monkeypatch,
):
monkeypatch.setenv("LANGSMITH_API_KEY", "global-key")
monkeypatch.setenv("LANGSMITH_PROJECT", "global-project")
monkeypatch.setenv("LANGSMITH_TENANT_ID", "global-tenant")
logger = LangsmithLogger(
langsmith_api_key="default-key",
langsmith_project="default-project",
langsmith_base_url="https://default.example",
)
credentials = logger.get_credentials_from_env(
langsmith_base_url="https://attacker.example",
allow_env_credentials=False,
)
assert credentials["LANGSMITH_API_KEY"] is None
assert credentials["LANGSMITH_PROJECT"] == "litellm-completion"
assert credentials["LANGSMITH_BASE_URL"] == "https://attacker.example"
assert credentials["LANGSMITH_TENANT_ID"] is None
@pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown")
@pytest.mark.asyncio
async def test_dynamic_langsmith_base_url_does_not_inherit_default_api_key(
monkeypatch,
):
monkeypatch.setenv("LANGSMITH_API_KEY", "global-key")
logger = LangsmithLogger(
langsmith_api_key="default-key",
langsmith_project="default-project",
langsmith_base_url="https://default.example",
)
credentials = logger._get_credentials_to_use_for_request(
kwargs={"standard_callback_dynamic_params": {"langsmith_base_url": "https://attacker.example"}}
)
assert credentials["LANGSMITH_API_KEY"] is None
assert credentials["LANGSMITH_BASE_URL"] == "https://attacker.example"
@pytest.fixture()
def _vcr_outcome_gate(request, vcr):
install_live_call_probe(request, vcr)
yield
record_vcr_outcome(request, vcr)
@pytest_asyncio.fixture(loop_scope="function")
async def drain_logging_worker(isolate_litellm_state: None) -> AsyncIterator[None]:
yield
await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=LOGGING_WORKER_DRAIN_TIMEOUT_SECONDS)
LOGGING_WORKER_DRAIN_TIMEOUT_SECONDS: Final = LOGGING_WORKER_MAX_TIME_PER_COROUTINE + 5.0
@pytest.fixture(scope="function")
def isolate_litellm_state():
"""
Per-function isolation fixture.
Resets litellm state to the true defaults captured at conftest import time,
then restores after the test. This prevents module-level mutations (e.g.
`litellm.num_retries = 3` at the top of test_langfuse_e2e_test.py) from
leaking across tests within the same xdist worker.
"""
from litellm.litellm_core_utils import litellm_logging as ll_logging
from litellm.proxy.management_helpers import audit_logs as ll_audit_logs
if hasattr(litellm, "in_memory_llm_clients_cache"):
litellm.in_memory_llm_clients_cache.flush_cache()
ll_logging._in_memory_loggers.clear()
ll_audit_logs._audit_log_callback_cache.clear()
for attr in _LIST_ATTRS:
if attr in _DEFAULTS:
default = _DEFAULTS[attr]
setattr(litellm, attr, default.copy() if isinstance(default, list) else default)
for attr in _SCALAR_ATTRS:
if attr in _DEFAULTS:
setattr(litellm, attr, _DEFAULTS[attr])
yield
if hasattr(litellm, "in_memory_llm_clients_cache"):
litellm.in_memory_llm_clients_cache.flush_cache()
ll_logging._in_memory_loggers.clear()
ll_audit_logs._audit_log_callback_cache.clear()
for attr in _LIST_ATTRS:
if attr in _DEFAULTS:
default = _DEFAULTS[attr]
setattr(litellm, attr, default.copy() if isinstance(default, list) else default)
for attr in _SCALAR_ATTRS:
if attr in _DEFAULTS:
setattr(litellm, attr, _DEFAULTS[attr])
_LIST_ATTRS = (
"callbacks",
"success_callback",
"failure_callback",
"_async_success_callback",
"_async_failure_callback",
"service_callback",
"pre_call_rules",
"post_call_rules",
)
_SCALAR_ATTRS = (
"set_verbose",
"cache",
"num_retries",
"num_retries_per_request",
"turn_off_message_logging",
"redact_messages_in_exceptions",
"redact_user_api_key_info",
"s3_callback_params",
"s3_audit_callback_params",
"datadog_params",
"vector_store_registry",
)
_DEFAULTS: dict = {}
@pytest.fixture(scope="module")
def setup_and_teardown():
"""
Module-scoped setup. Reloads litellm only in single-process mode
(skipped under xdist to avoid cross-worker interference).
"""
import litellm
worker_id = os.environ.get("PYTEST_XDIST_WORKER", None)
if worker_id is None:
importlib.reload(litellm)
try:
if hasattr(litellm, "proxy") and hasattr(litellm.proxy, "proxy_server"):
import litellm.proxy.proxy_server
importlib.reload(litellm.proxy.proxy_server)
except Exception:
pass
if hasattr(litellm, "in_memory_llm_clients_cache"):
litellm.in_memory_llm_clients_cache.flush_cache()
yield
# Test get_credentials_from_env
@pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown")
@pytest.mark.asyncio
async def test_get_credentials_from_env():
# Test with direct parameters
logger = LangsmithLogger(
langsmith_api_key="test-key",
langsmith_project="test-project",
langsmith_base_url="http://test-url",
)
credentials = logger.get_credentials_from_env(
langsmith_api_key="custom-key",
langsmith_project="custom-project",
langsmith_base_url="http://custom-url",
)
assert credentials["LANGSMITH_API_KEY"] == "custom-key"
assert credentials["LANGSMITH_PROJECT"] == "custom-project"
assert credentials["LANGSMITH_BASE_URL"] == "http://custom-url"
# assert that the default api base is used if not provided
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
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.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown")
@pytest.mark.asyncio
async def test_group_batches_by_credentials():
logger = LangsmithLogger(langsmith_api_key="test-key")
# Create test queue objects
queue_obj1 = LangsmithQueueObject(
data={"test": "data1"},
credentials={
"LANGSMITH_API_KEY": "key1",
"LANGSMITH_PROJECT": "proj1",
"LANGSMITH_BASE_URL": "url1",
"LANGSMITH_TENANT_ID": None,
},
)
queue_obj2 = LangsmithQueueObject(
data={"test": "data2"},
credentials={
"LANGSMITH_API_KEY": "key1",
"LANGSMITH_PROJECT": "proj1",
"LANGSMITH_BASE_URL": "url1",
"LANGSMITH_TENANT_ID": None,
},
)
logger.log_queue = [queue_obj1, queue_obj2]
grouped = logger._group_batches_by_credentials()
# Check grouping
assert len(grouped) == 1 # Should have one group since credentials are same
key = list(grouped.keys())[0]
assert isinstance(key, CredentialsKey)
assert len(grouped[key].queue_objects) == 2
@pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown")
@pytest.mark.asyncio
async def test_group_batches_by_credentials_multiple_credentials():
# Test with multiple different credentials
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": None,
},
)
queue_obj2 = LangsmithQueueObject(
data={"test": "data2"},
credentials={
"LANGSMITH_API_KEY": "key2", # Different API key
"LANGSMITH_PROJECT": "proj1",
"LANGSMITH_BASE_URL": "url1",
"LANGSMITH_TENANT_ID": None,
},
)
queue_obj3 = LangsmithQueueObject(
data={"test": "data3"},
credentials={
"LANGSMITH_API_KEY": "key1",
"LANGSMITH_PROJECT": "proj2", # Different project
"LANGSMITH_BASE_URL": "url1",
"LANGSMITH_TENANT_ID": None,
},
)
logger.log_queue = [queue_obj1, queue_obj2, queue_obj3]
grouped = logger._group_batches_by_credentials()
# Check grouping
assert len(grouped) == 3 # Should have three groups since credentials differ
for key, batch_group in grouped.items():
assert isinstance(key, CredentialsKey)
assert len(batch_group.queue_objects) == 1 # Each group should have one object
@pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown")
@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.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown")
@pytest.mark.asyncio
async def test_make_dot_order():
logger = LangsmithLogger(langsmith_api_key="test-key")
run_id = "729cff0e-f30c-4336-8b79-45d6b61c64b4"
dot_order = logger.make_dot_order(run_id)
print("dot_order=", dot_order)
# Check format: YYYYMMDDTHHMMSSfffZ + run_id
# Check the timestamp portion (first 23 characters)
timestamp_part = dot_order[:-36] # 36 is length of run_id
assert len(timestamp_part) == 22
assert timestamp_part[8] == "T" # Check T separator
assert timestamp_part[-1] == "Z" # Check Z suffix
# Verify timestamp format
try:
# Parse the timestamp portion (removing the Z)
datetime.strptime(timestamp_part[:-1], "%Y%m%dT%H%M%S%f")
except ValueError:
pytest.fail("Timestamp portion is not in correct format")
# Verify run_id portion
assert dot_order[-36:] == run_id
# Test is_serializable
@pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown")
@pytest.mark.asyncio
async def test_is_serializable():
from pydantic import BaseModel
from litellm.integrations.langsmith import is_serializable
# Test basic types
assert is_serializable("string") is True
assert is_serializable(123) is True
assert is_serializable({"key": "value"}) is True
# Test non-serializable types
async def async_func():
pass
assert is_serializable(async_func) is False
class TestModel(BaseModel):
field: str
assert is_serializable(TestModel(field="test")) is False
@pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown")
@pytest.mark.asyncio
async def test_async_send_batch():
logger = LangsmithLogger(langsmith_api_key="test-key")
# 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
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"]
# tenant_id should not be in headers if not provided
assert "x-tenant-id" not in call_args[1]["headers"]
@pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown")
@pytest.mark.asyncio
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.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown")
@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
"""
try:
# Mock the httpx post request
# We need to mock get_async_httpx_client to return a mock AsyncHTTPHandler
# because LangsmithLogger creates its own instance
mock_async_httpx_handler = AsyncMock()
mock_response = MagicMock() # Use MagicMock for response to allow sync methods
mock_response.status_code = 200
mock_response.raise_for_status = MagicMock() # raise_for_status is sync in httpx
mock_response.text = ""
mock_async_httpx_handler.post = AsyncMock(return_value=mock_response)
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
litellm.callbacks = [LangsmithLogger()]
response = await litellm.acompletion(
model="gpt-4.1-mini",
messages=[{"role": "user", "content": "Test message"}],
max_tokens=10,
temperature=0.2,
mock_response="This is a mock response",
langsmith_api_key="fake_key_project2",
langsmith_project="fake_project2",
)
print("Waiting for logs to be flushed to Langsmith.....")
await asyncio.sleep(3)
print("done sleeping 3 seconds...")
# Verify the post request was made with correct parameters
mock_async_httpx_handler.post.assert_called_once()
call_args = mock_async_httpx_handler.post.call_args
print("call_args", call_args)
# Check URL contains /runs/batch
assert "/runs/batch" in call_args[1]["url"]
# 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"]
assert call_args[1]["headers"]["Content-Type"] == "application/json"
# Verify the request body contains the expected data
request_body = json.loads(call_args[1]["content"])
assert "post" in request_body
assert len(request_body["post"]) == 1
# EXPECTED BODY
expected_body = {
"post": [
{
"name": "LLMRun",
"run_type": "llm",
"inputs": {
"id": "chatcmpl-82699ee4-7932-4fc0-9585-76abc8caeafa",
"call_type": "acompletion",
"model": "gpt-4.1-mini",
"messages": [{"role": "user", "content": "Test message"}],
"model_parameters": {
"temperature": 0.2,
"max_tokens": 10,
},
},
"outputs": {
"id": "chatcmpl-82699ee4-7932-4fc0-9585-76abc8caeafa",
"model": "gpt-4.1-mini",
"choices": [
{
"finish_reason": "stop",
"index": 0,
"message": {
"content": "This is a mock response",
"role": "assistant",
"tool_calls": None,
"function_call": None,
},
}
],
"usage": {
"completion_tokens": 20,
"prompt_tokens": 10,
"total_tokens": 30,
},
},
"session_name": "fake_project2",
}
]
}
# Print both bodies for debugging
actual_body = json.loads(call_args[1]["content"])
print("\nExpected body:")
print(json.dumps(expected_body, indent=2))
print("\nActual body:")
print(json.dumps(actual_body, indent=2))
assert len(actual_body["post"]) == 1
# Assert only the critical parts we care about
assert actual_body["post"][0]["name"] == expected_body["post"][0]["name"]
assert actual_body["post"][0]["run_type"] == expected_body["post"][0]["run_type"]
assert actual_body["post"][0]["inputs"]["messages"] == expected_body["post"][0]["inputs"]["messages"]
assert (
actual_body["post"][0]["inputs"]["model_parameters"]
== expected_body["post"][0]["inputs"]["model_parameters"]
)
assert actual_body["post"][0]["outputs"]["choices"] == expected_body["post"][0]["outputs"]["choices"]
assert (
actual_body["post"][0]["outputs"]["usage"]["completion_tokens"]
== expected_body["post"][0]["outputs"]["usage"]["completion_tokens"]
)
assert (
actual_body["post"][0]["outputs"]["usage"]["prompt_tokens"]
== expected_body["post"][0]["outputs"]["usage"]["prompt_tokens"]
)
assert (
actual_body["post"][0]["outputs"]["usage"]["total_tokens"]
== expected_body["post"][0]["outputs"]["usage"]["total_tokens"]
)
assert 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}")
@pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown")
@pytest.mark.asyncio
async def test_langsmith_queue_logging():
try:
# Initialize LangsmithLogger
test_langsmith_logger = LangsmithLogger()
litellm.callbacks = [test_langsmith_logger]
test_langsmith_logger.batch_size = 6
litellm.set_verbose = True
# Make multiple calls to ensure we don't hit the batch size
for _ in range(5):
response = await litellm.acompletion(
model="gpt-4.1-mini",
messages=[{"role": "user", "content": "Test message"}],
max_tokens=10,
temperature=0.2,
mock_response="This is a mock response",
)
# Poll for async callbacks to complete (up to 10s)
for _ in range(20):
if len(test_langsmith_logger.log_queue) >= 5:
break
await asyncio.sleep(0.5)
# Check that logs are in the queue
assert len(test_langsmith_logger.log_queue) == 5
# Now make calls to exceed the batch size
for _ in range(3):
response = await litellm.acompletion(
model="gpt-4.1-mini",
messages=[{"role": "user", "content": "Test message"}],
max_tokens=10,
temperature=0.2,
mock_response="This is a mock response",
)
# Poll for flush to complete (up to 10s)
for _ in range(20):
if len(test_langsmith_logger.log_queue) < 5:
break
await asyncio.sleep(0.5)
print("Length of langsmith log queue: {}".format(len(test_langsmith_logger.log_queue)))
# Check that the queue was flushed after exceeding batch size
assert len(test_langsmith_logger.log_queue) < 5
# Clean up
for cb in litellm.callbacks:
if isinstance(cb, LangsmithLogger):
await cb.async_httpx_client.client.aclose()
except Exception as e:
pytest.fail(f"Error occurred: {e}")
Reference in a new issue View git blame Copy permalink