mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-27 01:22:18 +00:00
fix(ci): stop stale CI reds, keep unit tests off the host env, retry CyberArk policy conflicts (#43294)
* fix(ci): stop five stale or flaky CI reds and retry CyberArk policy-load conflicts The Langfuse redaction unit test exports to a local OTLP capture instead of polling Langfuse Cloud through a recorded lookup. The passthrough worker-kill test only requires spend rows for requests the surviving worker served. The spend-routes sweep treats the intentional /spend/capture_rate 503 as expected. CyberArk retries a 409 policy load in Python, Rust and the e2e Conjur helper instead of reading it as "variable exists". The integration egress guard now matches the script's own cgroup, so it no longer blocks the CircleCI agent, which runs as the same user. * fix(ci): keep the policy-load backoff typed as float * fix(ci): retry CyberArk policy loads without blocking the event loop and tighten the worker-kill and Langfuse tests * fix(secrets): load CyberArk policy one request at a time per manager * test(secrets): pin that non-conflict CyberArk policy failures are not retried * test(unit): run tests/unit with only an allowlisted host environment CircleCI's unit job inherits every project env var, so real provider keys, REDIS_HOST, DATABASE_URL and AWS or Azure credentials reached tests that assume none are set. Locally, litellm's import-time load_dotenv did the same from any .env up the tree. The unit conftest now drops every variable outside a small allowlist and disables dotenv before litellm is imported. * test(e2e): name a failed search and the stuck batch status instead of misattributing them The websearch session test read an empty web_search_tool_result_error block as a successful search, so a failing search tool surfaced as a session billing bug. The batch cancellation timeout now reports the last status the proxy returned. * fix(ci): scrub the host environment per unit test instead of for the whole pytest process GHA shards run tests/unit next to other suites in one process, so the import-time scrub deleted MCP_TEST_PEER_PYTHON before tests/mcp_tests read it and the MCP upstream fell back to the SDK2 interpreter. The two websearch tests that called OpenAI and Perplexity live are removed: tests/unit no longer sees their keys. * fix(ci): scrub only the host variables present before litellm is imported The per-test scrub also deleted TIKTOKEN_CACHE_DIR, which litellm sets at import to its bundled encodings, so tokenizer paths tried to download them and hit the socket guard. The prisma setup test now passes its own database URL instead of reading one another test leaked into the process environment. * fix(ci): stop the order-dependent unit reds and settle logging tasks on their own queue LoggingWorker marked a task done on whichever queue was current when the callback finished, so a callback that outlived an event-loop change raised "task_done() called too many times" or undercounted the new loop's queue. It now settles the queue the task came from. The rest are test isolation fixes for failures that only appeared when another file ran first on the same xdist worker: a replaced user_api_key_cache, breaker metrics unregistered by prometheus tests, semantic_router's health-check filter on uvicorn.access, logging tasks carried over from bedrock tests, a Router-written model_cost entry, and a stray post captured by the langflow test. The token counter check now asserts bounded chunking instead of wall-clock time. * test(e2e/ui): wait for the logout redirect before visiting a protected page Logout revokes the session server-side before clearing cookies and navigating, so an immediate page.goto either ran with the cookie still set or was aborted by the logout redirect (net::ERR_ABORTED). * test(unit): restore the prometheus metrics config per test and settle logs carried from earlier tests in the a2a cost tests * test(router): pin the router clock in the usage counter tests so a minute rollover cannot empty the read * test(e2e/ui): wait for logout to clear the token cookie instead of for a login redirect * test(integration/mcp): answer the model-info probe another test's proxy sends to the model double
This commit is contained in:
yuneng-jiang
GitHub
parent
31678a1dbc
commit
14f4c34c61
32 changed files with 561 additions and 287 deletions
|
|
@ -26,6 +26,7 @@ guard_created=false
|
|||
guard_installed=false
|
||||
guard6_created=false
|
||||
guard6_installed=false
|
||||
egress_cgroup=litellm-integration
|
||||
cleanup() {
|
||||
original_status=$?
|
||||
trap - EXIT INT TERM
|
||||
|
|
@ -47,14 +48,14 @@ cleanup() {
|
|||
fi
|
||||
done
|
||||
if [ "$guard_installed" = true ]; then
|
||||
sudo iptables -D OUTPUT -m owner --uid-owner "$(id -u)" -j integration_only || original_status=1
|
||||
sudo iptables -D OUTPUT -m cgroup --path "$egress_cgroup" -j integration_only || original_status=1
|
||||
fi
|
||||
if [ "$guard_created" = true ]; then
|
||||
sudo iptables -F integration_only || original_status=1
|
||||
sudo iptables -X integration_only || original_status=1
|
||||
fi
|
||||
if [ "$guard6_installed" = true ]; then
|
||||
sudo ip6tables -D OUTPUT -m owner --uid-owner "$(id -u)" -j integration_only || original_status=1
|
||||
sudo ip6tables -D OUTPUT -m cgroup --path "$egress_cgroup" -j integration_only || original_status=1
|
||||
fi
|
||||
if [ "$guard6_created" = true ]; then
|
||||
sudo ip6tables -F integration_only || original_status=1
|
||||
|
|
@ -100,6 +101,8 @@ if [ "$mode" = parity ]; then
|
|||
export INTEGRATION_ROUTING=capture
|
||||
fi
|
||||
|
||||
sudo mkdir -p "/sys/fs/cgroup/$egress_cgroup"
|
||||
echo "$$" | sudo tee "/sys/fs/cgroup/$egress_cgroup/cgroup.procs" > /dev/null
|
||||
sudo iptables -N integration_only
|
||||
guard_created=true
|
||||
sudo iptables -A integration_only -o lo -j ACCEPT
|
||||
|
|
@ -109,13 +112,13 @@ for service in postgres-db redis-cache; do
|
|||
sudo iptables -A integration_only -d "$address" -j ACCEPT
|
||||
done
|
||||
sudo iptables -A integration_only -j REJECT
|
||||
sudo iptables -I OUTPUT 1 -m owner --uid-owner "$(id -u)" -j integration_only
|
||||
sudo iptables -I OUTPUT 1 -m cgroup --path "$egress_cgroup" -j integration_only
|
||||
guard_installed=true
|
||||
sudo ip6tables -N integration_only
|
||||
guard6_created=true
|
||||
sudo ip6tables -A integration_only -o lo -j ACCEPT
|
||||
sudo ip6tables -A integration_only -j REJECT
|
||||
sudo ip6tables -I OUTPUT 1 -m owner --uid-owner "$(id -u)" -j integration_only
|
||||
sudo ip6tables -I OUTPUT 1 -m cgroup --path "$egress_cgroup" -j integration_only
|
||||
guard6_installed=true
|
||||
|
||||
if curl --noproxy '*' --connect-timeout 2 -s http://198.51.100.1 >/dev/null 2>&1; then
|
||||
|
|
|
|||
|
|
@ -48,6 +48,7 @@ pub struct CyberArkSecretManager {
|
|||
token: Cache<(), SecretValue>,
|
||||
secrets: SecretCache<String, SecretValue>,
|
||||
authentication_lock: Arc<tokio::sync::Mutex<()>>,
|
||||
policy_load_lock: Arc<tokio::sync::Mutex<()>>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ impl CyberArkSecretManager {
|
|||
token,
|
||||
secrets,
|
||||
authentication_lock: Arc::new(tokio::sync::Mutex::new(())),
|
||||
policy_load_lock: Arc::new(tokio::sync::Mutex::new(())),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,8 @@
|
|||
use super::*;
|
||||
|
||||
const POLICY_LOAD_ATTEMPTS: u32 = 5;
|
||||
const POLICY_LOAD_RETRY_DELAY: std::time::Duration = std::time::Duration::from_millis(200);
|
||||
|
||||
impl CyberArkSecretManager {
|
||||
pub async fn async_write_secret(
|
||||
&self,
|
||||
|
|
@ -105,37 +108,43 @@ impl CyberArkSecretManager {
|
|||
"- !variable {}\n",
|
||||
serde_json::to_string(name).expect("serializing a string cannot fail")
|
||||
);
|
||||
let response = with_timeout(
|
||||
self.client
|
||||
.post(policy_url)
|
||||
.header("Authorization", authorization)
|
||||
.header("Content-Type", "application/x-yaml")
|
||||
.body(body),
|
||||
context,
|
||||
)
|
||||
.send()
|
||||
.await;
|
||||
match response {
|
||||
Ok(response) if response.status().is_success() => {}
|
||||
Ok(response)
|
||||
if matches!(
|
||||
response.status(),
|
||||
reqwest::StatusCode::CONFLICT | reqwest::StatusCode::UNPROCESSABLE_ENTITY
|
||||
) =>
|
||||
{
|
||||
litellm_tracing::debug!(
|
||||
"CyberArk variable policy already exists or conflicts: {}",
|
||||
response.status()
|
||||
);
|
||||
}
|
||||
Ok(response) => {
|
||||
litellm_tracing::warn!(
|
||||
"Could not ensure CyberArk variable exists: {}",
|
||||
response.status()
|
||||
);
|
||||
}
|
||||
Err(error) => {
|
||||
litellm_tracing::warn!("Error ensuring CyberArk variable exists: {error}");
|
||||
let _policy_load = self.policy_load_lock.lock().await;
|
||||
for attempt in 0..POLICY_LOAD_ATTEMPTS {
|
||||
let response = with_timeout(
|
||||
self.client
|
||||
.post(policy_url.clone())
|
||||
.header("Authorization", authorization.clone())
|
||||
.header("Content-Type", "application/x-yaml")
|
||||
.body(body.clone()),
|
||||
context,
|
||||
)
|
||||
.send()
|
||||
.await;
|
||||
match response {
|
||||
Ok(response)
|
||||
if response.status() == reqwest::StatusCode::CONFLICT
|
||||
&& attempt + 1 < POLICY_LOAD_ATTEMPTS =>
|
||||
{
|
||||
tokio::time::sleep(POLICY_LOAD_RETRY_DELAY * 2_u32.pow(attempt)).await;
|
||||
}
|
||||
Ok(response) if response.status().is_success() => return,
|
||||
Ok(response) if response.status() == reqwest::StatusCode::UNPROCESSABLE_ENTITY => {
|
||||
litellm_tracing::debug!(
|
||||
"CyberArk variable policy was rejected as unprocessable"
|
||||
);
|
||||
return;
|
||||
}
|
||||
Ok(response) => {
|
||||
litellm_tracing::warn!(
|
||||
"Could not ensure CyberArk variable exists: {}",
|
||||
response.status()
|
||||
);
|
||||
return;
|
||||
}
|
||||
Err(error) => {
|
||||
litellm_tracing::warn!("Error ensuring CyberArk variable exists: {error}");
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ async fn rejected_write_token_is_reauthenticated_once() {
|
|||
let server = MockServer::start().await;
|
||||
mount_auth(&server, 2).await;
|
||||
Mock::given(path("/policies/acct/policy/root"))
|
||||
.respond_with(ResponseTemplate::new(409))
|
||||
.respond_with(ResponseTemplate::new(201))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
|
@ -45,7 +45,6 @@ async fn rejected_write_token_is_reauthenticated_once() {
|
|||
|
||||
#[rstest]
|
||||
#[case::created(201)]
|
||||
#[case::already_exists(409)]
|
||||
#[case::unprocessable(422)]
|
||||
#[case::server_error(500)]
|
||||
#[tokio::test]
|
||||
|
|
@ -81,13 +80,81 @@ async fn writes_tolerate_policy_status_and_cache_value(#[case] policy_status: u1
|
|||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn policy_load_conflict_is_retried_before_the_value_write() {
|
||||
let server = MockServer::start().await;
|
||||
mount_auth(&server, 1).await;
|
||||
let policy_loads = Arc::new(AtomicUsize::new(0));
|
||||
let policy_loads_for_response = Arc::clone(&policy_loads);
|
||||
Mock::given(path("/policies/acct/policy/root"))
|
||||
.respond_with(move |_: &Request| {
|
||||
if policy_loads_for_response.fetch_add(1, Ordering::SeqCst) < 2 {
|
||||
ResponseTemplate::new(409)
|
||||
} else {
|
||||
ResponseTemplate::new(201)
|
||||
}
|
||||
})
|
||||
.expect(3)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let policy_loads_at_value_write = Arc::clone(&policy_loads);
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/secrets/acct/variable/key"))
|
||||
.respond_with(move |_: &Request| {
|
||||
if policy_loads_at_value_write.load(Ordering::SeqCst) == 3 {
|
||||
ResponseTemplate::new(201)
|
||||
} else {
|
||||
ResponseTemplate::new(404)
|
||||
}
|
||||
})
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let manager = manager(&server, Duration::from_secs(60));
|
||||
|
||||
manager
|
||||
.async_write_secret("key", &SecretValue::new("v"), None)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn concurrent_writes_load_policy_one_at_a_time() {
|
||||
let server = MockServer::start().await;
|
||||
mount_auth(&server, 1).await;
|
||||
Mock::given(path("/policies/acct/policy/root"))
|
||||
.respond_with(ResponseTemplate::new(201).set_delay(Duration::from_millis(100)))
|
||||
.expect(4)
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("POST"))
|
||||
.respond_with(ResponseTemplate::new(201))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let manager = manager(&server, Duration::from_secs(60));
|
||||
let started = std::time::Instant::now();
|
||||
|
||||
let value = SecretValue::new("v");
|
||||
let results = tokio::join!(
|
||||
manager.async_write_secret("key-0", &value, None),
|
||||
manager.async_write_secret("key-1", &value, None),
|
||||
manager.async_write_secret("key-2", &value, None),
|
||||
manager.async_write_secret("key-3", &value, None),
|
||||
);
|
||||
|
||||
assert!(results.0.is_ok() && results.1.is_ok() && results.2.is_ok() && results.3.is_ok());
|
||||
assert!(started.elapsed() >= Duration::from_millis(400));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn failed_value_write_is_not_cached() {
|
||||
let server = MockServer::start().await;
|
||||
mount_auth(&server, 1).await;
|
||||
Mock::given(path("/policies/acct/policy/root"))
|
||||
.respond_with(ResponseTemplate::new(409))
|
||||
.respond_with(ResponseTemplate::new(201))
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(path("/secrets/acct/variable/key"))
|
||||
|
|
|
|||
|
|
@ -165,7 +165,9 @@ class LoggingWorker:
|
|||
if self._worker_task is None or self._worker_task.done():
|
||||
self._worker_task = asyncio.create_task(self._worker_loop())
|
||||
|
||||
async def _process_log_task(self, task: LoggingTask, sem: asyncio.Semaphore):
|
||||
async def _process_log_task(
|
||||
self, task: LoggingTask, sem: asyncio.Semaphore, queue: "asyncio.Queue[LoggingTask]"
|
||||
) -> None:
|
||||
"""Runs the logging task and handles cleanup. Releases semaphore when done."""
|
||||
try:
|
||||
if self._queue is not None:
|
||||
|
|
@ -182,7 +184,7 @@ class LoggingWorker:
|
|||
verbose_logger.exception("LoggingWorker error: %s", e)
|
||||
finally:
|
||||
self._untrack_dequeued(task)
|
||||
self._queue.task_done()
|
||||
queue.task_done()
|
||||
finally:
|
||||
# Always release semaphore, even if queue is None
|
||||
sem.release()
|
||||
|
|
@ -219,7 +221,8 @@ class LoggingWorker:
|
|||
async def _worker_loop(self) -> None:
|
||||
"""Main worker loop that gets tasks and schedules them to run concurrently."""
|
||||
try:
|
||||
if self._queue is None or self._sem is None:
|
||||
queue: Final = self._queue
|
||||
if queue is None or self._sem is None:
|
||||
return
|
||||
|
||||
while True:
|
||||
|
|
@ -227,10 +230,10 @@ class LoggingWorker:
|
|||
# unbounded growth of waiting tasks
|
||||
await self._sem.acquire()
|
||||
try:
|
||||
task = await self._queue.get()
|
||||
task = await queue.get()
|
||||
self._track_dequeued(task)
|
||||
# Track each spawned coroutine so we can cancel on shutdown.
|
||||
processing_task = asyncio.create_task(self._process_log_task(task, self._sem))
|
||||
processing_task = asyncio.create_task(self._process_log_task(task, self._sem, queue))
|
||||
self._running_tasks.add(processing_task)
|
||||
processing_task.add_done_callback(self._running_tasks.discard)
|
||||
except Exception:
|
||||
|
|
@ -497,7 +500,8 @@ class LoggingWorker:
|
|||
"""
|
||||
Clear the queue with a maximum time limit.
|
||||
"""
|
||||
if self._queue is None:
|
||||
queue: Final = self._queue
|
||||
if queue is None:
|
||||
return
|
||||
|
||||
start_time: Final = asyncio.get_event_loop().time()
|
||||
|
|
@ -509,7 +513,7 @@ class LoggingWorker:
|
|||
break
|
||||
|
||||
try:
|
||||
task = self._queue.get_nowait()
|
||||
task = queue.get_nowait()
|
||||
# Await the coroutine to properly execute and avoid "never awaited" warnings
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
|
|
@ -522,7 +526,7 @@ class LoggingWorker:
|
|||
finally:
|
||||
# Clear reference to prevent memory leaks
|
||||
task = None
|
||||
self._queue.task_done() # If you're using join() elsewhere
|
||||
queue.task_done()
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import asyncio
|
||||
import base64
|
||||
import os
|
||||
from typing import Any, Final
|
||||
|
|
@ -10,6 +11,7 @@ import litellm
|
|||
from litellm._logging import verbose_logger
|
||||
from litellm.caching import InMemoryCache
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
_get_httpx_client,
|
||||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
|
|
@ -20,6 +22,9 @@ from litellm.rust_bridge.secret_manager import resolve_native_provider_reader, r
|
|||
from .base_secret_manager import BaseSecretManager, raise_if_unsafe_secret_name
|
||||
from .main import str_to_bool
|
||||
|
||||
CYBERARK_POLICY_LOAD_ATTEMPTS: Final = 5
|
||||
CYBERARK_POLICY_LOAD_RETRY_DELAY_SECONDS: Final = 0.2
|
||||
|
||||
|
||||
class CyberArkSecretManager(BaseSecretManager):
|
||||
def __init__(self):
|
||||
|
|
@ -30,6 +35,7 @@ class CyberArkSecretManager(BaseSecretManager):
|
|||
self.conjur_account = os.getenv("CYBERARK_ACCOUNT", "default")
|
||||
self.conjur_username = os.getenv("CYBERARK_USERNAME", "admin")
|
||||
self.conjur_api_key = os.getenv("CYBERARK_API_KEY", "")
|
||||
self._policy_load_lock: Final = asyncio.Lock()
|
||||
|
||||
# Optional config for certificate-based auth
|
||||
self.tls_cert_path = os.getenv("CYBERARK_CLIENT_CERT", "")
|
||||
|
|
@ -118,7 +124,7 @@ class CyberArkSecretManager(BaseSecretManager):
|
|||
token: Final = self._authenticate()
|
||||
return {"Authorization": f'Token token="{token}"'}
|
||||
|
||||
def _ensure_variable_exists(self, secret_name: str) -> None:
|
||||
async def _ensure_variable_exists(self, secret_name: str, async_client: AsyncHTTPHandler) -> None:
|
||||
"""
|
||||
Ensure a variable exists in CyberArk Conjur by creating a policy entry if needed.
|
||||
|
||||
|
|
@ -134,27 +140,33 @@ class CyberArkSecretManager(BaseSecretManager):
|
|||
policy_yaml: Final = f"- !variable {quoted_name}\n"
|
||||
|
||||
try:
|
||||
client: Final = _get_httpx_client(params={"ssl_verify": self.ssl_verify})
|
||||
resp: Final = client.client.post(
|
||||
policy_url,
|
||||
headers={
|
||||
**self._get_request_headers(),
|
||||
"Content-Type": "application/x-yaml",
|
||||
},
|
||||
content=policy_yaml,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
verbose_logger.debug("Created policy entry for variable: %s", secret_name)
|
||||
except httpx.HTTPStatusError as e:
|
||||
# Variable might already exist, which is fine
|
||||
if e.response.status_code in [409, 422]:
|
||||
verbose_logger.debug("Variable %s already exists or policy conflict (expected)", secret_name)
|
||||
else:
|
||||
verbose_logger.warning(
|
||||
"Could not ensure variable exists: %s - %s", e.response.status_code, e.response.text
|
||||
)
|
||||
async with self._policy_load_lock:
|
||||
resp: Final = await self._load_variable_policy(async_client, policy_url, policy_yaml)
|
||||
except Exception as e:
|
||||
verbose_logger.warning("Error ensuring variable exists: %s", e)
|
||||
return
|
||||
if resp.is_success:
|
||||
verbose_logger.debug("Created policy entry for variable: %s", secret_name)
|
||||
elif resp.status_code == 422:
|
||||
verbose_logger.debug("Variable %s policy was rejected as unprocessable", secret_name)
|
||||
else:
|
||||
verbose_logger.warning("Could not ensure variable exists: %s - %s", resp.status_code, resp.text)
|
||||
|
||||
async def _load_variable_policy(
|
||||
self, async_client: AsyncHTTPHandler, policy_url: str, policy_yaml: str, attempt: int = 0
|
||||
) -> httpx.Response:
|
||||
resp: Final = await async_client.client.post(
|
||||
policy_url,
|
||||
headers={
|
||||
**self._get_request_headers(),
|
||||
"Content-Type": "application/x-yaml",
|
||||
},
|
||||
content=policy_yaml,
|
||||
)
|
||||
if resp.status_code != 409 or attempt + 1 == CYBERARK_POLICY_LOAD_ATTEMPTS:
|
||||
return resp
|
||||
await asyncio.sleep(CYBERARK_POLICY_LOAD_RETRY_DELAY_SECONDS * (1 << attempt))
|
||||
return await self._load_variable_policy(async_client, policy_url, policy_yaml, attempt + 1)
|
||||
|
||||
def get_url(self, secret_name: str) -> str:
|
||||
"""
|
||||
|
|
@ -303,7 +315,7 @@ class CyberArkSecretManager(BaseSecretManager):
|
|||
|
||||
try:
|
||||
# Ensure the variable exists in the policy first
|
||||
self._ensure_variable_exists(secret_name)
|
||||
await self._ensure_variable_exists(secret_name, async_client)
|
||||
|
||||
# Now set the secret value
|
||||
url: Final = self.get_url(secret_name)
|
||||
|
|
|
|||
|
|
@ -121,7 +121,8 @@ def cleanup_batch(
|
|||
if current.status == "cancelling" and not needs_terminal_state:
|
||||
return
|
||||
assert clock() < deadline, (
|
||||
f"Batch {batch_id} cancellation did not finish within {BATCH_CANCEL_TIMEOUT_SECONDS}s"
|
||||
f"Batch {batch_id} cancellation did not finish within {BATCH_CANCEL_TIMEOUT_SECONDS}s, "
|
||||
f"last status {current.status}"
|
||||
)
|
||||
wait(BATCH_CANCEL_POLL_SECONDS)
|
||||
|
||||
|
|
|
|||
|
|
@ -210,6 +210,7 @@ class TestBatchCancellation:
|
|||
with pytest.raises(ExceptionGroup) as caught:
|
||||
manager.teardown()
|
||||
assert "cancellation did not finish" in str(caught.value.exceptions[0])
|
||||
assert "last status cancelling" in str(caught.value.exceptions[0])
|
||||
client.calls.assert_done()
|
||||
|
||||
@pytest.mark.parametrize("status", ["completed", "failed", "expired", "cancelled"])
|
||||
|
|
|
|||
|
|
@ -74,6 +74,8 @@ SPEND_ROUTES = (
|
|||
|
||||
_SPEND_PREFIXES = ("/spend", "/global/spend", "/global/activity")
|
||||
|
||||
_CAPTURE_RATE_ROUTE: Final = "/spend/capture_rate"
|
||||
|
||||
# Served from the MonthlyGlobalSpend / DailyTagSpend / Last30d* views, which the
|
||||
# proxy creates in the background once the schema migrations have landed, so on a
|
||||
# fresh database they can 500 for a while after the proxy starts serving.
|
||||
|
|
@ -120,7 +122,7 @@ def test_schema_listed_spend_routes_are_responsive(client: SpendClient) -> None:
|
|||
and "{" not in path
|
||||
and any(path.startswith(prefix) for prefix in _SPEND_PREFIXES)
|
||||
]
|
||||
extras = [path for path in discovered if path not in SPEND_ROUTES]
|
||||
extras = [path for path in discovered if path not in (*SPEND_ROUTES, _CAPTURE_RATE_ROUTE)]
|
||||
|
||||
params = _date_range()
|
||||
results = [(path, client.probe(path, params=params)) for path in extras]
|
||||
|
|
@ -132,3 +134,11 @@ def test_schema_listed_spend_routes_are_responsive(client: SpendClient) -> None:
|
|||
if not result.healthy
|
||||
]
|
||||
assert not offenders, "non-responsive schema spend routes:\n" + "\n".join(offenders)
|
||||
|
||||
|
||||
def test_capture_rate_reports_or_names_the_missing_billing_key(client: SpendClient) -> None:
|
||||
result: Final = client.probe(_CAPTURE_RATE_ROUTE, params=_date_range())
|
||||
print(f"{_CAPTURE_RATE_ROUTE} -> {result.status_code}\n{result.body[:600]}")
|
||||
assert result.status_code == 200 or (result.status_code == 503 and "OPENAI_ADMIN_KEY is not set" in result.body), (
|
||||
f"{_CAPTURE_RATE_ROUTE} -> {result.status_code}\n{result.body[:600]}"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -12,13 +12,14 @@ Needs a proxy booted with the callback and a real search backend, which
|
|||
``gateway/stage_mirror_ci_config.yml`` carries as the ``e2e-search`` Perplexity tool.
|
||||
"""
|
||||
|
||||
from typing import Final
|
||||
from typing import Final, Literal
|
||||
|
||||
import pytest
|
||||
from e2e_config import unique_marker
|
||||
from e2e_http import unwrap
|
||||
from lifecycle import ResourceManager
|
||||
from models import (
|
||||
AnthropicContentBlock,
|
||||
AnthropicMessagesBody,
|
||||
AnthropicWebSearchTool,
|
||||
ChatMessage,
|
||||
|
|
@ -26,6 +27,7 @@ from models import (
|
|||
SpendLogRow,
|
||||
)
|
||||
from proxy_client import ProxyClient
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
|
|
@ -37,6 +39,18 @@ def _has_search_row(rows: list[SpendLogRow]) -> bool:
|
|||
return any(row.call_type == SEARCH_CALL_TYPE for row in rows)
|
||||
|
||||
|
||||
class _SearchResultError(BaseModel):
|
||||
type: Literal["web_search_tool_result_error"]
|
||||
error_code: str
|
||||
|
||||
|
||||
def _search_error_code(block: AnthropicContentBlock) -> str | None:
|
||||
try:
|
||||
return _SearchResultError.model_validate((block.model_extra or {}).get("content")).error_code
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
class TestWebSearchInterceptionSession:
|
||||
@pytest.mark.covers(
|
||||
"quota_management.spend_tracking.websearch_interception.bills_under_request_session",
|
||||
|
|
@ -79,6 +93,15 @@ class TestWebSearchInterceptionSession:
|
|||
f"precondition: the turn never ran an intercepted search, so there is no search row to attribute. "
|
||||
f"blocks={block_types}"
|
||||
)
|
||||
search_errors: Final = tuple(
|
||||
code
|
||||
for block in response.content or ()
|
||||
if block.type == "web_search_tool_result" and (code := _search_error_code(block)) is not None
|
||||
)
|
||||
assert not search_errors, (
|
||||
f"precondition: the e2e-search tool failed upstream ({search_errors}), so no {SEARCH_CALL_TYPE} row is "
|
||||
"billed at all; check the proxy's search tool credentials before reading this as a session bug"
|
||||
)
|
||||
|
||||
rows: Final = proxy.poll_logs_for_session(session_id, min_rows=2, predicate=_has_search_row)
|
||||
by_call_type: Final = {row.call_type or "" for row in rows}
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ from __future__ import annotations
|
|||
|
||||
import base64
|
||||
import os
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Final, Literal
|
||||
from urllib.parse import quote
|
||||
|
|
@ -25,6 +26,9 @@ DEFAULT_USERNAME: Final = "admin"
|
|||
|
||||
SYSTEM: Final = "cyberark"
|
||||
|
||||
_POLICY_LOAD_ATTEMPTS: Final = 5
|
||||
_POLICY_LOAD_RETRY_DELAY_SECONDS: Final = 0.2
|
||||
|
||||
_START_HINT: Final = (
|
||||
f"Start one with `bash tests/e2e/secret_manager/backend.sh up {SYSTEM}`, which writes the env for "
|
||||
f"the proxy (booted from gateway/secret_manager_{SYSTEM}_ci_config.yml) and for the tests"
|
||||
|
|
@ -72,13 +76,20 @@ class Conjur:
|
|||
def _secret_url(self, name: str) -> str:
|
||||
return f"{self.base_url}/secrets/{self.account}/variable/{quote(name, safe='')}"
|
||||
|
||||
def _update_root_policy(self, method: Literal["POST", "PATCH"], policy: str, action: str) -> None:
|
||||
def _load_root_policy(self, method: Literal["POST", "PATCH"], policy: str, attempt: int = 0) -> ExternalWrite:
|
||||
result: Final = send_text_external(
|
||||
method,
|
||||
f"{self.base_url}/policies/{self.account}/policy/root",
|
||||
headers=self._headers(content_type="application/x-yaml"),
|
||||
content=policy,
|
||||
)
|
||||
if result.status_code != 409 or attempt + 1 == _POLICY_LOAD_ATTEMPTS:
|
||||
return result
|
||||
time.sleep(_POLICY_LOAD_RETRY_DELAY_SECONDS * (1 << attempt))
|
||||
return self._load_root_policy(method, policy, attempt + 1)
|
||||
|
||||
def _update_root_policy(self, method: Literal["POST", "PATCH"], policy: str, action: str) -> None:
|
||||
result: Final = self._load_root_policy(method, policy)
|
||||
self._fail_unless_reached(result, action)
|
||||
if not result.ok:
|
||||
pytest.fail(f"Conjur refused to {action}: HTTP {result.status_code} {result.body[:300]}")
|
||||
|
|
|
|||
|
|
@ -24,6 +24,11 @@ test.describe("Logout", () => {
|
|||
// Click Logout — the handler clears the auth cookie and navigates via
|
||||
// window.location.href = PROXY_LOGOUT_URL (empty string in the e2e env).
|
||||
await popup.getByRole("button", { name: "Logout" }).click();
|
||||
await expect
|
||||
.poll(async () => (await page.context().cookies()).filter((c) => c.name === "token").length, {
|
||||
timeout: 15_000,
|
||||
})
|
||||
.toBe(0);
|
||||
|
||||
// The cookie is now gone — visiting a protected page must redirect to /ui/login.
|
||||
await page.goto("/ui?page=llm-playground", { waitUntil: "domcontentloaded" });
|
||||
|
|
|
|||
|
|
@ -47,6 +47,8 @@ def _model_double(tool: str) -> Callable[[Request], Reply]:
|
|||
arguments: Final = json.dumps(ADD)
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
if request.method == "GET" and request.target.endswith("/models"):
|
||||
return _json({"object": "list", "data": []})
|
||||
body: Final = json.loads(request.body)
|
||||
assert isinstance(body, dict), request.body
|
||||
done: Final = _has_tool_result(body)
|
||||
|
|
@ -192,7 +194,9 @@ class Rig:
|
|||
)
|
||||
|
||||
def upstream_tools(self) -> tuple[tuple[str, ...], ...]:
|
||||
return tuple(_tool_names(json.loads(request.body)) for request in self.wire.drain())
|
||||
return tuple(
|
||||
_tool_names(json.loads(request.body)) for request in self.wire.drain() if request.method == "POST"
|
||||
)
|
||||
|
||||
def final_text(self, body: Mapping[str, object]) -> str:
|
||||
if self.surface == "chat":
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ import asyncio
|
|||
import json
|
||||
import re
|
||||
import signal
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
|
|
@ -51,6 +52,16 @@ def _error_information(call_id: str) -> dict[str, JsonValue]:
|
|||
return object_value(parsed["error_information"])
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Served:
|
||||
response: httpx.Response
|
||||
client_port: int
|
||||
|
||||
|
||||
def _spend_rows(call_id: str) -> list[dict[str, JsonValue]]:
|
||||
return read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (call_id,))
|
||||
|
||||
|
||||
def _single_spend_row(call_id: str) -> None:
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (call_id,)),
|
||||
|
|
@ -62,19 +73,23 @@ def _single_spend_row(call_id: str) -> None:
|
|||
|
||||
async def _fire_burst(
|
||||
base_url: str, key: str, count: int, *, tolerate_transport_errors: bool = False
|
||||
) -> tuple[httpx.Response, ...]:
|
||||
async def one(client: httpx.AsyncClient, index: int) -> httpx.Response:
|
||||
) -> tuple[_Served, ...]:
|
||||
async def one(client: httpx.AsyncClient, index: int) -> _Served:
|
||||
if index % 3 == 0:
|
||||
path: Final = "/gemini/v1beta/models/nope-9:generateContent"
|
||||
elif index % 3 == 1:
|
||||
path = "/gemini/v1beta/models/nope-9:streamGenerateContent?alt=sse"
|
||||
else:
|
||||
path = "/gemini/v1beta/models/healthy-model:streamGenerateContent?alt=sse"
|
||||
return await client.post(
|
||||
async with client.stream(
|
||||
"POST",
|
||||
path,
|
||||
json=_GENERATE_CONTENT,
|
||||
headers={"Authorization": f"Bearer {key}", "x-goog-api-key": key},
|
||||
)
|
||||
) as response:
|
||||
client_port: Final = int(response.extensions["network_stream"].get_extra_info("client_addr")[1])
|
||||
await response.aread()
|
||||
return _Served(response=response, client_port=client_port)
|
||||
|
||||
async with httpx.AsyncClient(base_url=base_url, timeout=30, trust_env=False) as client:
|
||||
results: Final = await asyncio.gather(
|
||||
|
|
@ -82,7 +97,7 @@ async def _fire_burst(
|
|||
)
|
||||
for result in results:
|
||||
assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result)
|
||||
return tuple(result for result in results if isinstance(result, httpx.Response))
|
||||
return tuple(result for result in results if isinstance(result, _Served))
|
||||
|
||||
|
||||
async def test_passthrough_upstream_outage_mid_burst_still_logs_errors_once(gateway: Gateway, tmp_path: Path) -> None:
|
||||
|
|
@ -97,7 +112,7 @@ async def test_passthrough_upstream_outage_mid_burst_still_logs_errors_once(gate
|
|||
burst: Final = asyncio.create_task(_fire_burst(str(candidate.client.base_url), candidate.key, 30))
|
||||
await asyncio.to_thread(eventually, lambda: wire.received.qsize(), lambda size: size >= 10, 30)
|
||||
with wire_server(_chaos_reply, port=port):
|
||||
responses: Final = await burst
|
||||
responses: Final = tuple(served.response for served in await burst)
|
||||
assert len(responses) == 30
|
||||
for response in responses:
|
||||
assert response.status_code in (200, 404, 500, 502), response.status_code
|
||||
|
|
@ -131,10 +146,15 @@ async def test_passthrough_worker_sigkill_leaves_sibling_serving_and_logging(gat
|
|||
_fire_burst(str(candidate.client.base_url), candidate.key, 20, tolerate_transport_errors=True)
|
||||
)
|
||||
await asyncio.to_thread(eventually, lambda: wire.received.qsize(), lambda size: size >= 5, 30)
|
||||
psutil.Process(workers[0]).send_signal(signal.SIGKILL)
|
||||
responses: Final = await burst
|
||||
for response in responses:
|
||||
assert response.status_code in (200, 404, 500, 502), response.status_code
|
||||
victim: Final = psutil.Process(workers[0])
|
||||
victim.suspend()
|
||||
victim_ports: Final = frozenset(
|
||||
connection.raddr.port for connection in victim.net_connections(kind="tcp") if connection.raddr
|
||||
)
|
||||
victim.send_signal(signal.SIGKILL)
|
||||
served: Final = await burst
|
||||
for item in served:
|
||||
assert item.response.status_code in (200, 404, 500, 502), item.response.status_code
|
||||
follow_up: Final = candidate.request(
|
||||
"POST",
|
||||
"/gemini/v1beta/models/nope-9:generateContent",
|
||||
|
|
@ -143,8 +163,12 @@ async def test_passthrough_worker_sigkill_leaves_sibling_serving_and_logging(gat
|
|||
)
|
||||
assert follow_up.status_code == 404, follow_up.text
|
||||
assert follow_up.json() == json.loads(_NOT_FOUND_BODY), follow_up.text
|
||||
for response in responses:
|
||||
if "x-litellm-call-id" in response.headers:
|
||||
_single_spend_row(response.headers["x-litellm-call-id"])
|
||||
logged: Final = tuple(item for item in served if "x-litellm-call-id" in item.response.headers)
|
||||
survivor_served: Final = tuple(item for item in logged if item.client_port not in victim_ports)
|
||||
assert survivor_served, [item.client_port for item in logged]
|
||||
for item in survivor_served:
|
||||
_single_spend_row(item.response.headers["x-litellm-call-id"])
|
||||
for item in logged:
|
||||
assert len(_spend_rows(item.response.headers["x-litellm-call-id"])) <= 1, item.response.headers
|
||||
error_information: Final = _error_information(follow_up.headers["x-litellm-call-id"])
|
||||
assert "not found for this scripted upstream" in str(error_information["error_message"]), follow_up.text
|
||||
|
|
|
|||
|
|
@ -86,7 +86,8 @@ async def test_cyberark_write_secret_rejects_yaml_injection():
|
|||
"team/user@example.com",
|
||||
],
|
||||
)
|
||||
def test_cyberark_ensure_variable_exists_escapes_yaml_metacharacters(secret_name):
|
||||
@pytest.mark.asyncio
|
||||
async def test_cyberark_ensure_variable_exists_escapes_yaml_metacharacters(secret_name):
|
||||
"""
|
||||
Regression test: _ensure_variable_exists must escape secret_name (not just
|
||||
denylist-check it) so the policy body always parses back to exactly one
|
||||
|
|
@ -95,19 +96,21 @@ def test_cyberark_ensure_variable_exists_escapes_yaml_metacharacters(secret_name
|
|||
with patch("litellm.proxy.proxy_server.premium_user", True):
|
||||
captured = {}
|
||||
|
||||
def _capture_post(url, headers=None, content=None):
|
||||
async def _capture_post(url, headers=None, content=None):
|
||||
captured["content"] = content
|
||||
return create_mock_response(status_code=201, text="")
|
||||
|
||||
mock_sync_client = MagicMock()
|
||||
mock_sync_client.client.post.side_effect = _capture_post
|
||||
mock_sync_client.client.post.return_value = create_mock_response(status_code=200, text="mock-token")
|
||||
mock_async_client = MagicMock()
|
||||
mock_async_client.client.post.side_effect = _capture_post
|
||||
|
||||
with patch(
|
||||
"litellm.secret_managers.cyberark_secret_manager._get_httpx_client",
|
||||
return_value=mock_sync_client,
|
||||
):
|
||||
cyberark_manager = CyberArkSecretManager()
|
||||
cyberark_manager._ensure_variable_exists(secret_name)
|
||||
await cyberark_manager._ensure_variable_exists(secret_name, mock_async_client)
|
||||
|
||||
policy_yaml = captured["content"]
|
||||
parsed = yaml.compose(policy_yaml)
|
||||
|
|
|
|||
|
|
@ -5,6 +5,10 @@ import logging
|
|||
import os
|
||||
from typing import Any, Optional
|
||||
from unittest.mock import MagicMock, patch
|
||||
import threading
|
||||
from http.server import BaseHTTPRequestHandler, HTTPServer
|
||||
|
||||
from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest
|
||||
|
||||
logging.basicConfig(level=logging.DEBUG)
|
||||
|
||||
|
|
@ -206,53 +210,91 @@ def create_async_task(**completion_kwargs):
|
|||
return asyncio.create_task(litellm.acompletion(**completion_args))
|
||||
|
||||
|
||||
def _otlp_capture(exports: list[bytes]) -> type[BaseHTTPRequestHandler]:
|
||||
class OtlpCapture(BaseHTTPRequestHandler):
|
||||
def do_POST(self):
|
||||
exports.append(self.rfile.read(int(self.headers.get("content-length", 0))))
|
||||
self.send_response(200)
|
||||
self.end_headers()
|
||||
|
||||
def do_GET(self):
|
||||
self.send_response(200)
|
||||
self.send_header("content-type", "application/json")
|
||||
self.end_headers()
|
||||
self.wfile.write(b"{}")
|
||||
|
||||
def log_message(self, *args):
|
||||
pass
|
||||
|
||||
return OtlpCapture
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def local_langfuse():
|
||||
exports: list[bytes] = []
|
||||
server = HTTPServer(("127.0.0.1", 0), _otlp_capture(exports))
|
||||
threading.Thread(target=server.serve_forever, daemon=True).start()
|
||||
yield f"http://127.0.0.1:{server.server_port}", exports
|
||||
server.shutdown()
|
||||
|
||||
|
||||
def _exported_spans(exports: list[bytes]):
|
||||
for body in exports:
|
||||
for resource_spans in ExportTraceServiceRequest.FromString(body).resource_spans:
|
||||
for scope_spans in resource_spans.scope_spans:
|
||||
yield from scope_spans.spans
|
||||
|
||||
|
||||
def _exported_attributes(exports: list[bytes], trace_id: str) -> list[dict[str, str]]:
|
||||
return [
|
||||
{attribute.key: attribute.value.string_value for attribute in span.attributes}
|
||||
for span in _exported_spans(list(exports))
|
||||
if span.trace_id.hex() == trace_id
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("stream", [False, True])
|
||||
@pytest.mark.flaky(retries=12, delay=2)
|
||||
async def test_langfuse_logging_without_request_response(stream, langfuse_client):
|
||||
try:
|
||||
from litellm._uuid import uuid
|
||||
async def test_langfuse_logging_without_request_response(stream, local_langfuse, monkeypatch):
|
||||
from litellm._uuid import uuid
|
||||
|
||||
_unique_trace_name = f"litellm-test-{str(uuid.uuid4())}"
|
||||
litellm.set_verbose = True
|
||||
litellm.turn_off_message_logging = True
|
||||
litellm.success_callback = ["langfuse"]
|
||||
response = await create_async_task(
|
||||
model="gpt-3.5-turbo",
|
||||
stream=stream,
|
||||
metadata={"trace_id": _unique_trace_name},
|
||||
)
|
||||
print(response)
|
||||
if stream:
|
||||
async for chunk in response:
|
||||
print(chunk)
|
||||
langfuse_host, exports = local_langfuse
|
||||
prompt = f"prompt-{uuid.uuid4()}"
|
||||
answer = f"answer-{uuid.uuid4()}"
|
||||
trace_name = f"litellm-test-{uuid.uuid4()}"
|
||||
monkeypatch.setattr(litellm, "turn_off_message_logging", True)
|
||||
monkeypatch.setattr(litellm, "success_callback", ["langfuse"])
|
||||
response = await litellm.acompletion(
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
mock_response=answer,
|
||||
stream=stream,
|
||||
metadata={"trace_id": trace_name},
|
||||
langfuse_public_key=f"pk-lf-{trace_name}",
|
||||
langfuse_secret_key="sk-lf-local",
|
||||
langfuse_host=langfuse_host,
|
||||
)
|
||||
if stream:
|
||||
async for _ in response:
|
||||
pass
|
||||
|
||||
langfuse_client.flush()
|
||||
generations: list[dict[str, str]] = []
|
||||
for _ in range(60):
|
||||
generations = [
|
||||
attributes
|
||||
for attributes in _exported_attributes(exports, resolve_trace_id(trace_name))
|
||||
if attributes.get("langfuse.observation.type") == "generation"
|
||||
]
|
||||
if generations:
|
||||
break
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
for _ in range(30):
|
||||
_trace_data = langfuse_client.api.observations.get_many(
|
||||
trace_id=resolve_trace_id(_unique_trace_name),
|
||||
type="GENERATION",
|
||||
fields="core,io",
|
||||
).data
|
||||
if _trace_data:
|
||||
break
|
||||
await asyncio.sleep(3)
|
||||
|
||||
print(f"_trace_data: {_trace_data}")
|
||||
assert json.loads(_trace_data[0].input) == {
|
||||
"messages": [{"content": "redacted-by-litellm", "role": "user"}]
|
||||
}
|
||||
assert json.loads(_trace_data[0].output) == {
|
||||
"role": "assistant",
|
||||
"content": "redacted-by-litellm",
|
||||
"function_call": None,
|
||||
"tool_calls": None,
|
||||
"provider_specific_fields": None,
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
pytest.fail(f"An exception occurred - {e}")
|
||||
assert len(generations) == 1, generations
|
||||
assert json.loads(generations[0]["langfuse.observation.input"]) == {
|
||||
"messages": [{"content": "redacted-by-litellm", "role": "user"}]
|
||||
}
|
||||
assert json.loads(generations[0]["langfuse.observation.output"])["content"] == "redacted-by-litellm"
|
||||
assert all(prompt.encode() not in body and answer.encode() not in body for body in exports)
|
||||
|
||||
|
||||
# Get the current directory of the file being run
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ import os
|
|||
import traceback
|
||||
from dotenv import load_dotenv
|
||||
from fastapi import Request
|
||||
from datetime import datetime
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from litellm import Router
|
||||
import pytest
|
||||
|
|
@ -971,11 +971,18 @@ def _rpm_tpm_router(model_id: str) -> Router:
|
|||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def router_minute_pinned(monkeypatch):
|
||||
pinned = datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc)
|
||||
monkeypatch.setattr("litellm.router.get_utc_datetime", lambda: pinned)
|
||||
|
||||
|
||||
def _ratelimit_headers(response: ModelResponse | CustomStreamWrapper) -> dict[str, int]:
|
||||
return {k: v for k, v in response._hidden_params["additional_headers"].items() if k.startswith("x-ratelimit-")}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("router_minute_pinned")
|
||||
async def test_acompletion_headers_read_post_increment_counter_and_count_once():
|
||||
router = _rpm_tpm_router("lit-3058-async")
|
||||
|
||||
|
|
@ -1018,6 +1025,7 @@ async def test_acompletion_wildcard_route_headers_and_counter_use_resolved_deplo
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("router_minute_pinned")
|
||||
async def test_acompletion_stream_counts_request_before_headers_and_tokens_once_on_completion():
|
||||
router = _rpm_tpm_router("lit-3058-stream")
|
||||
|
||||
|
|
|
|||
|
|
@ -10,6 +10,12 @@ import pytest
|
|||
|
||||
import litellm
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
|
||||
|
||||
async def _reset_callbacks_and_settle_pending_logs() -> None:
|
||||
litellm.logging_callback_manager._reset_all_callbacks()
|
||||
await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0)
|
||||
|
||||
|
||||
def _make_send_message_request(request_id: str, user_text: str = "Hello"):
|
||||
|
|
@ -129,7 +135,7 @@ async def test_asend_message_uses_cost_per_query(monkeypatch):
|
|||
from litellm.a2a_protocol import asend_message
|
||||
|
||||
# Setup logger
|
||||
litellm.logging_callback_manager._reset_all_callbacks()
|
||||
await _reset_callbacks_and_settle_pending_logs()
|
||||
cost_logger = CostLogger()
|
||||
monkeypatch.setattr(litellm, "callbacks", [cost_logger])
|
||||
|
||||
|
|
@ -164,7 +170,7 @@ async def test_asend_message_uses_cost_per_query_from_litellm_params_dict(monkey
|
|||
"""
|
||||
from litellm.a2a_protocol import asend_message
|
||||
|
||||
litellm.logging_callback_manager._reset_all_callbacks()
|
||||
await _reset_callbacks_and_settle_pending_logs()
|
||||
cost_logger = CostLogger()
|
||||
monkeypatch.setattr(litellm, "callbacks", [cost_logger])
|
||||
|
||||
|
|
@ -225,7 +231,7 @@ async def test_asend_message_uses_input_output_cost_per_token(monkeypatch):
|
|||
from litellm.a2a_protocol import asend_message
|
||||
|
||||
# Setup logger
|
||||
litellm.logging_callback_manager._reset_all_callbacks()
|
||||
await _reset_callbacks_and_settle_pending_logs()
|
||||
token_cost_logger = TokenAndCostLogger()
|
||||
monkeypatch.setattr(litellm, "callbacks", [token_cost_logger])
|
||||
|
||||
|
|
@ -299,7 +305,7 @@ async def test_asend_message_passes_agent_id_to_callback(monkeypatch):
|
|||
from litellm.a2a_protocol import asend_message
|
||||
|
||||
# Setup logger
|
||||
litellm.logging_callback_manager._reset_all_callbacks()
|
||||
await _reset_callbacks_and_settle_pending_logs()
|
||||
agent_id_logger = AgentIdLogger()
|
||||
monkeypatch.setattr(litellm, "callbacks", [agent_id_logger])
|
||||
|
||||
|
|
@ -359,7 +365,7 @@ async def test_asend_message_streaming_propagates_metadata():
|
|||
from litellm.a2a_protocol import asend_message_streaming
|
||||
|
||||
# Setup logger
|
||||
litellm.logging_callback_manager._reset_all_callbacks()
|
||||
await _reset_callbacks_and_settle_pending_logs()
|
||||
metadata_logger = MetadataLogger()
|
||||
litellm.logging_callback_manager.add_litellm_async_success_callback(metadata_logger)
|
||||
|
||||
|
|
@ -406,7 +412,7 @@ async def test_asend_message_streaming_triggers_callbacks():
|
|||
from litellm.a2a_protocol import asend_message_streaming
|
||||
|
||||
# Setup logger - must use logging_callback_manager to properly register
|
||||
litellm.logging_callback_manager._reset_all_callbacks()
|
||||
await _reset_callbacks_and_settle_pending_logs()
|
||||
callback_logger = AgentIdLogger()
|
||||
litellm.logging_callback_manager.add_litellm_async_success_callback(callback_logger)
|
||||
litellm.logging_callback_manager.add_litellm_success_callback(callback_logger)
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ import asyncio
|
|||
import time
|
||||
from collections.abc import Iterator
|
||||
from datetime import timedelta
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -1023,7 +1024,12 @@ async def test_breaker_metrics_track_state_and_failure_class():
|
|||
from redis.exceptions import ConnectionError as RedisConnectionError
|
||||
from redis.exceptions import TimeoutError as RedisTimeoutError
|
||||
|
||||
from litellm.caching.redis_cache import RedisCircuitBreaker, is_redis_timeout_failure
|
||||
from litellm.caching.redis_cache import RedisCircuitBreaker, _breaker_metrics, is_redis_timeout_failure
|
||||
|
||||
metrics: Final = _breaker_metrics()
|
||||
for collector in (metrics._state_gauge, metrics._transitions, metrics._failures):
|
||||
if collector is not None and collector not in REGISTRY._collector_to_names:
|
||||
REGISTRY.register(collector)
|
||||
|
||||
def sample(name, labels=None):
|
||||
return REGISTRY.get_sample_value(name, labels) or 0.0
|
||||
|
|
|
|||
|
|
@ -12,6 +12,32 @@ import httpx
|
|||
import pytest
|
||||
from pytest_socket import enable_socket, socket_allow_hosts
|
||||
|
||||
HOST_ENVIRONMENT_ALLOWLIST: Final = frozenset(
|
||||
(
|
||||
"PATH",
|
||||
"HOME",
|
||||
"USER",
|
||||
"LOGNAME",
|
||||
"TMPDIR",
|
||||
"TEMP",
|
||||
"TMP",
|
||||
"LANG",
|
||||
"LC_ALL",
|
||||
"LC_CTYPE",
|
||||
"TZ",
|
||||
"VIRTUAL_ENV",
|
||||
"LITELLM_LOCAL_MODEL_COST_MAP",
|
||||
"TIKTOKEN_CACHE_DIR",
|
||||
)
|
||||
)
|
||||
HOST_ENVIRONMENT_ALLOWED_PREFIXES: Final = ("PYTEST_", "PYTHON", "COV_CORE_", "COVERAGE_")
|
||||
HOST_ONLY_ENVIRONMENT: Final = frozenset(
|
||||
name
|
||||
for name in os.environ
|
||||
if name not in HOST_ENVIRONMENT_ALLOWLIST and not name.startswith(HOST_ENVIRONMENT_ALLOWED_PREFIXES)
|
||||
)
|
||||
|
||||
os.environ["PYTHON_DOTENV_DISABLED"] = "1"
|
||||
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
|
||||
|
||||
import litellm # noqa: E402 # litellm reads LITELLM_LOCAL_MODEL_COST_MAP at import
|
||||
|
|
@ -170,6 +196,8 @@ def isolated_aws_config_files(tmp_path_factory: pytest.TempPathFactory) -> tuple
|
|||
def isolate_host_environment(isolated_aws_config_files: tuple[Path, Path]) -> Iterator[None]:
|
||||
credentials, config = isolated_aws_config_files
|
||||
with pytest.MonkeyPatch.context() as environment:
|
||||
for name in HOST_ONLY_ENVIRONMENT:
|
||||
environment.delenv(name, raising=False)
|
||||
environment.setenv("AWS_SHARED_CREDENTIALS_FILE", str(credentials))
|
||||
environment.setenv("AWS_CONFIG_FILE", str(config))
|
||||
environment.setenv("AWS_EC2_METADATA_DISABLED", "true")
|
||||
|
|
|
|||
|
|
@ -477,7 +477,7 @@ def test_valid_configuration_passes_validation():
|
|||
# ==============================================================================
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_prometheus_exclude_settings():
|
||||
"""Restore the global exclude settings after each test so they don't leak."""
|
||||
prev_metrics = litellm.prometheus_exclude_metrics
|
||||
|
|
|
|||
|
|
@ -5,7 +5,6 @@ Tests the end-to-end flow of websearch_interception callback with
|
|||
litellm.acompletion() for transparent server-side web search execution.
|
||||
"""
|
||||
|
||||
import os
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
|
@ -37,75 +36,6 @@ def websearch_logger():
|
|||
return WebSearchInterceptionLogger(enabled_providers=[LlmProviders.OPENAI, LlmProviders.MINIMAX])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.skipif(
|
||||
os.environ.get("OPENAI_API_KEY") is None,
|
||||
reason="OPENAI_API_KEY not set",
|
||||
)
|
||||
async def test_websearch_chat_completion_with_openai():
|
||||
"""Test websearch interception with OpenAI chat completions API.
|
||||
|
||||
This test verifies that:
|
||||
1. Model calls litellm_web_search tool
|
||||
2. Server executes web search automatically
|
||||
3. Server makes follow-up request with search results
|
||||
4. User gets final answer without tool_calls
|
||||
"""
|
||||
# Configure WebSearch interception
|
||||
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
|
||||
websearch_logger = WebSearchInterceptionLogger(enabled_providers=[LlmProviders.OPENAI])
|
||||
litellm.callbacks = [websearch_logger]
|
||||
|
||||
try:
|
||||
response = await litellm.acompletion(
|
||||
model="gpt-4o-mini", # Use cheaper model for testing
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": "What's the weather in San Francisco today?",
|
||||
}
|
||||
],
|
||||
tools=[
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "litellm_web_search",
|
||||
"description": "Search the web for information",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {
|
||||
"type": "string",
|
||||
"description": "Search query",
|
||||
}
|
||||
},
|
||||
"required": ["query"],
|
||||
},
|
||||
},
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
# Verify response structure
|
||||
assert isinstance(response, ModelResponse)
|
||||
assert response.choices[0].message.content is not None
|
||||
assert len(response.choices[0].message.content) > 0
|
||||
|
||||
# If agentic loop worked, we should NOT have tool_calls in final response
|
||||
# (they should have been executed and replaced with final answer)
|
||||
if hasattr(response.choices[0].message, "tool_calls"):
|
||||
# If tool_calls exist, it means agentic loop didn't run
|
||||
# This could happen if search tool is not configured
|
||||
pytest.skip("Agentic loop did not execute - search tool may not be configured")
|
||||
|
||||
# Verify we got a meaningful response
|
||||
assert response.choices[0].finish_reason in ["stop", "end_turn"]
|
||||
|
||||
finally:
|
||||
# Restore original callbacks
|
||||
litellm.callbacks = original_callbacks
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websearch_chat_completion_hook_detection():
|
||||
"""Test that websearch hook correctly detects tool calls in response."""
|
||||
|
|
@ -321,61 +251,6 @@ async def test_websearch_json_serialization_fix():
|
|||
assert arguments_str != "{'query': 'weather in SF'}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.skipif(
|
||||
os.environ.get("OPENAI_API_KEY") is None or os.environ.get("PERPLEXITY_API_KEY") is None,
|
||||
reason="OPENAI_API_KEY or PERPLEXITY_API_KEY not set",
|
||||
)
|
||||
async def test_websearch_streaming_conversion():
|
||||
"""Test that streaming requests are converted to non-streaming for web search.
|
||||
|
||||
When stream=True is passed with web search tools, the handler should:
|
||||
1. Convert stream=True to stream=False for initial request
|
||||
2. Execute web search
|
||||
3. Convert final response back to streaming
|
||||
"""
|
||||
websearch_logger = WebSearchInterceptionLogger(
|
||||
enabled_providers=[LlmProviders.OPENAI], search_tool_name="perplexity-search"
|
||||
)
|
||||
litellm.callbacks = [websearch_logger]
|
||||
|
||||
try:
|
||||
response = await litellm.acompletion(
|
||||
model="gpt-4o-mini",
|
||||
messages=[{"role": "user", "content": "What's the latest AI news?"}],
|
||||
tools=[
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "litellm_web_search",
|
||||
"description": "Search the web",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"query": {"type": "string"}},
|
||||
},
|
||||
},
|
||||
}
|
||||
],
|
||||
stream=True,
|
||||
)
|
||||
|
||||
# Response should be a streaming iterator
|
||||
chunks = []
|
||||
async for chunk in response:
|
||||
chunks.append(chunk)
|
||||
|
||||
# Verify we got streaming chunks
|
||||
assert len(chunks) > 0
|
||||
|
||||
# Verify chunks have expected structure
|
||||
for chunk in chunks:
|
||||
assert hasattr(chunk, "choices")
|
||||
assert len(chunk.choices) > 0
|
||||
|
||||
finally:
|
||||
litellm.callbacks = []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_maybe_run_chat_completion_agentic_loop_calls_chat_completion_hook():
|
||||
"""Regression test: maybe_run_chat_completion_agentic_loop must call
|
||||
|
|
|
|||
|
|
@ -180,6 +180,39 @@ class TestLoggingWorker:
|
|||
|
||||
assert sorted(fired) == ["first", "second"]
|
||||
|
||||
def test_callback_finishing_after_loop_change_settles_only_its_own_queue(self):
|
||||
worker = LoggingWorker(timeout=1.0, max_queue_size=10, concurrency=1)
|
||||
fired = []
|
||||
|
||||
async def marker(name, delay=0.0):
|
||||
await asyncio.sleep(delay)
|
||||
fired.append(name)
|
||||
|
||||
async def start_slow_callback():
|
||||
worker.ensure_initialized_and_enqueue(marker("slow", delay=0.05))
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
async def log_on_second_loop():
|
||||
for name in ("b1", "b2", "b3"):
|
||||
worker.ensure_initialized_and_enqueue(marker(name))
|
||||
for _ in range(2):
|
||||
await asyncio.sleep(0)
|
||||
|
||||
first_loop = asyncio.new_event_loop()
|
||||
try:
|
||||
first_loop.run_until_complete(start_slow_callback())
|
||||
first_loop_tasks = tuple(asyncio.all_tasks(first_loop))
|
||||
asyncio.run(log_on_second_loop())
|
||||
first_loop.run_until_complete(asyncio.sleep(0.1))
|
||||
failures = [
|
||||
task.exception() for task in first_loop_tasks if task.done() and not task.cancelled() and task.exception()
|
||||
]
|
||||
finally:
|
||||
first_loop.close()
|
||||
|
||||
assert failures == []
|
||||
assert "slow" in fired
|
||||
|
||||
@pytest.mark.parametrize("stranded", ["still_queued", "dequeued_never_started"])
|
||||
def test_flush_on_new_loop_drains_tasks_stranded_on_previous_loop(self, stranded):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ import litellm.constants
|
|||
from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_COUNTS
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
from litellm.litellm_core_utils.token_counter import (
|
||||
_encoding_count,
|
||||
_get_exact_count_function,
|
||||
_get_extrapolating_count_function,
|
||||
_get_tiktoken_count_function,
|
||||
|
|
@ -79,15 +80,17 @@ def test_token_counter_basic():
|
|||
)
|
||||
|
||||
|
||||
def test_token_counter_large_repeated_text_is_fast():
|
||||
messages = [{"role": "user", "content": [{"type": "text", "text": "A" * 1024 * 1024}]}]
|
||||
def test_token_counter_large_repeated_text_is_encoded_in_bounded_chunks():
|
||||
text_length: Final = 1024 * 1024
|
||||
messages: Final = [{"role": "user", "content": [{"type": "text", "text": "A" * text_length}]}]
|
||||
|
||||
start_time = time.perf_counter()
|
||||
tokens = token_counter_new(model="us.anthropic.claude-sonnet-4-6", messages=messages)
|
||||
elapsed = time.perf_counter() - start_time
|
||||
with patch("litellm.litellm_core_utils.token_counter._encoding_count", wraps=_encoding_count) as encoding_count:
|
||||
tokens: Final = token_counter_new(model="us.anthropic.claude-sonnet-4-6", messages=messages)
|
||||
|
||||
assert elapsed < 2, f"Token counting took too long: {elapsed:.2f}s"
|
||||
encoded_lengths: Final = tuple(len(call.args[1]) for call in encoding_count.call_args_list)
|
||||
assert tokens > 0
|
||||
assert sum(encoded_lengths) >= text_length
|
||||
assert max(encoded_lengths) <= litellm.constants.TIKTOKEN_ENCODE_MAX_CHUNK_SIZE_CHARS
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
|
|||
|
|
@ -222,7 +222,8 @@ def test_langflow_extra_body_cannot_inject_tweaks_into_run_payload():
|
|||
|
||||
def fake_post(*args, **kwargs):
|
||||
body = kwargs.get("data")
|
||||
posted_bodies.append(json.loads(body) if isinstance(body, str) else body)
|
||||
if str(kwargs.get("url", "")).startswith("http://example.com"):
|
||||
posted_bodies.append(json.loads(body) if isinstance(body, (str, bytes)) else body)
|
||||
resp = MagicMock(spec=httpx.Response)
|
||||
resp.status_code = 200
|
||||
resp.json.return_value = {"outputs": [{"outputs": [{"results": {"message": {"text": "hi"}}}]}]}
|
||||
|
|
|
|||
|
|
@ -3717,7 +3717,7 @@ async def test_auth_vertex_ai_route(prisma_client):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_api_key_auth_db_unavailable():
|
||||
async def test_user_api_key_auth_db_unavailable(monkeypatch):
|
||||
"""
|
||||
Test that user_api_key_auth handles DB connection failures appropriately when:
|
||||
1. DB connection fails during token validation
|
||||
|
|
@ -3747,7 +3747,7 @@ async def test_user_api_key_auth_db_unavailable():
|
|||
|
||||
# Set up test environment
|
||||
setattr(litellm.proxy.proxy_server, "prisma_client", MockPrismaClient())
|
||||
setattr(litellm.proxy.proxy_server, "user_api_key_cache", MockDualCache())
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", MockDualCache())
|
||||
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
|
||||
setattr(
|
||||
litellm.proxy.proxy_server,
|
||||
|
|
@ -3777,7 +3777,7 @@ async def test_user_api_key_auth_db_unavailable():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_api_key_auth_db_unavailable_not_allowed():
|
||||
async def test_user_api_key_auth_db_unavailable_not_allowed(monkeypatch):
|
||||
"""
|
||||
Test that user_api_key_auth raises an exception when:
|
||||
This is default behavior
|
||||
|
|
@ -3808,7 +3808,7 @@ async def test_user_api_key_auth_db_unavailable_not_allowed():
|
|||
|
||||
# Set up test environment
|
||||
setattr(litellm.proxy.proxy_server, "prisma_client", MockPrismaClient())
|
||||
setattr(litellm.proxy.proxy_server, "user_api_key_cache", MockDualCache())
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", MockDualCache())
|
||||
setattr(litellm.proxy.proxy_server, "general_settings", {})
|
||||
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import os
|
||||
import traceback
|
||||
from typing import Final
|
||||
from unittest import mock
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
|
@ -35,6 +36,7 @@ from fastapi import FastAPI
|
|||
from fastapi.testclient import TestClient
|
||||
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.proxy_server import ( # Replace with the actual module where your FastAPI router is defined
|
||||
app,
|
||||
initialize,
|
||||
|
|
@ -418,7 +420,7 @@ def test_chat_completion_forward_llm_provider_auth_headers(
|
|||
|
||||
@mock_patch_acompletion()
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_disable_guardrails(mock_acompletion, client_no_auth):
|
||||
async def test_team_disable_guardrails(mock_acompletion, client_no_auth, monkeypatch):
|
||||
"""
|
||||
If team not allowed to turn on/off guardrails
|
||||
|
||||
|
|
@ -438,8 +440,9 @@ async def test_team_disable_guardrails(mock_acompletion, client_no_auth):
|
|||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.proxy_server import hash_token, user_api_key_cache
|
||||
from litellm.proxy.proxy_server import hash_token
|
||||
|
||||
user_api_key_cache: Final = UserApiKeyCache()
|
||||
_team_id = "1234"
|
||||
user_key = "sk-12345678"
|
||||
|
||||
|
|
@ -459,7 +462,7 @@ async def test_team_disable_guardrails(mock_acompletion, client_no_auth):
|
|||
user_api_key_cache.set_cache(key=hash_token(user_key), value=valid_token)
|
||||
user_api_key_cache.set_cache(key="team_id:{}".format(_team_id), value=team_obj)
|
||||
|
||||
setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
|
||||
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
|
||||
setattr(litellm.proxy.proxy_server, "prisma_client", "hello-world")
|
||||
|
||||
|
|
@ -481,10 +484,11 @@ from tests.unit.proxy.test_custom_callback_input import CompletionCustomHandler
|
|||
|
||||
|
||||
@mock_patch_acompletion()
|
||||
def test_custom_logger_failure_handler(mock_acompletion, client_no_auth):
|
||||
def test_custom_logger_failure_handler(mock_acompletion, client_no_auth, monkeypatch):
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.proxy_server import hash_token, user_api_key_cache
|
||||
from litellm.proxy.proxy_server import hash_token
|
||||
|
||||
user_api_key_cache: Final = UserApiKeyCache()
|
||||
rpm_limit = 0
|
||||
|
||||
mock_api_key = "sk-my-test-key"
|
||||
|
|
@ -501,7 +505,7 @@ def test_custom_logger_failure_handler(mock_acompletion, client_no_auth):
|
|||
litellm.callbacks = [mock_logger, mock_logger_unit_tests]
|
||||
proxy_logging_obj._init_litellm_callbacks(llm_router=None)
|
||||
|
||||
setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
|
||||
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
|
||||
setattr(litellm.proxy.proxy_server, "prisma_client", "FAKE-VAR")
|
||||
setattr(litellm.proxy.proxy_server, "proxy_logging_obj", proxy_logging_obj)
|
||||
|
|
@ -1296,7 +1300,7 @@ async def test_create_team_member_add(prisma_client, new_member_method): # noqa
|
|||
@pytest.mark.parametrize("team_route", ["/team/member_add", "/team/member_delete"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_team_member_add_team_admin_user_api_key_auth(
|
||||
prisma_client, team_member_role, team_route # noqa: F811 # pytest fixture, not a redefinition
|
||||
prisma_client, team_member_role, team_route, monkeypatch # noqa: F811 # pytest fixture, not a redefinition
|
||||
):
|
||||
import time
|
||||
|
||||
|
|
@ -1307,9 +1311,10 @@ async def test_create_team_member_add_team_admin_user_api_key_auth(
|
|||
ProxyException,
|
||||
hash_token,
|
||||
user_api_key_auth,
|
||||
user_api_key_cache,
|
||||
)
|
||||
|
||||
user_api_key_cache: Final = UserApiKeyCache()
|
||||
|
||||
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
|
||||
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
|
||||
setattr(litellm, "max_internal_user_budget", 10)
|
||||
|
|
@ -1335,7 +1340,7 @@ async def test_create_team_member_add_team_admin_user_api_key_auth(
|
|||
|
||||
user_api_key_cache.set_cache(key="team_id:{}".format(_team_id), value=team_obj)
|
||||
|
||||
setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
|
||||
|
||||
## TEST IF TEAM ADMIN ALLOWED TO CALL /MEMBER_ADD ENDPOINT
|
||||
import json
|
||||
|
|
@ -2349,7 +2354,7 @@ async def test_proxy_server_prisma_setup():
|
|||
mock_client.db = mock_db
|
||||
|
||||
prisma_client = await ProxyStartupEvent._setup_prisma_client(
|
||||
database_url=os.getenv("DATABASE_URL"),
|
||||
database_url="postgresql://user:pass@localhost:5432/litellm",
|
||||
proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -14120,6 +14120,7 @@ class TestContextWindowEscalation:
|
|||
assert result.routing_decision["context_escalated"] is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("local_model_cost_map")
|
||||
@pytest.mark.parametrize(
|
||||
"deployments,tiers,expected_model",
|
||||
[
|
||||
|
|
|
|||
|
|
@ -1,7 +1,9 @@
|
|||
import asyncio
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Final, TypedDict, cast
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
|
|
@ -107,6 +109,86 @@ async def test_async_write_matches_parity_fixture(monkeypatch: pytest.MonkeyPatc
|
|||
assert value_route.calls.last.request.content == b"v"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_async_write_retries_policy_load_conflict(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
fixture: Final = _fixture()
|
||||
manager: Final = _configure_manager(monkeypatch, fixture)
|
||||
secret: Final = fixture["secrets"][0]
|
||||
endpoint: Final = fixture["endpoint"]
|
||||
_respond(respx.post(endpoint + fixture["authenticate_path"]), content=fixture["token_json"].encode())
|
||||
policy_route: Final = respx.post(endpoint + fixture["policy_path"]).mock(
|
||||
side_effect=[httpx.Response(409), httpx.Response(409), httpx.Response(201)]
|
||||
)
|
||||
value_route: Final = respx.post(endpoint + secret["path"]).mock(
|
||||
side_effect=lambda _: httpx.Response(201 if policy_route.call_count == 3 else 404)
|
||||
)
|
||||
|
||||
result: Final = await manager.async_write_secret(secret["name"], "v") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # legacy secret manager API is untyped
|
||||
|
||||
assert policy_route.call_count == 3
|
||||
assert value_route.call_count == 1
|
||||
assert result["status"] == "success"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"policy_outcome",
|
||||
[422, 500, httpx.ConnectError("conjur unreachable")],
|
||||
ids=["unprocessable", "server_error", "unreachable"],
|
||||
)
|
||||
@respx.mock
|
||||
async def test_async_write_does_not_retry_non_conflict_policy_failures(
|
||||
monkeypatch: pytest.MonkeyPatch, policy_outcome: int | httpx.ConnectError
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
fixture: Final = _fixture()
|
||||
manager: Final = _configure_manager(monkeypatch, fixture)
|
||||
secret: Final = fixture["secrets"][0]
|
||||
endpoint: Final = fixture["endpoint"]
|
||||
_respond(respx.post(endpoint + fixture["authenticate_path"]), content=fixture["token_json"].encode())
|
||||
policy_route: Final = respx.post(endpoint + fixture["policy_path"])
|
||||
if isinstance(policy_outcome, int):
|
||||
_respond(policy_route, status_code=policy_outcome)
|
||||
else:
|
||||
policy_route.mock(side_effect=policy_outcome)
|
||||
value_route: Final = _respond(respx.post(endpoint + secret["path"]), status_code=201)
|
||||
|
||||
await manager.async_write_secret(secret["name"], "v") # pyright: ignore[reportUnknownMemberType] # legacy secret manager API is untyped
|
||||
|
||||
assert policy_route.call_count == 1
|
||||
assert value_route.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_concurrent_async_writes_load_policy_one_at_a_time(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
fixture: Final = _fixture()
|
||||
manager: Final = _configure_manager(monkeypatch, fixture)
|
||||
endpoint: Final = fixture["endpoint"]
|
||||
_respond(respx.post(endpoint + fixture["authenticate_path"]), content=fixture["token_json"].encode())
|
||||
in_flight: Final = asyncio.Semaphore(1)
|
||||
|
||||
async def load_policy(_: httpx.Request) -> httpx.Response:
|
||||
if in_flight.locked():
|
||||
return httpx.Response(409)
|
||||
async with in_flight:
|
||||
await asyncio.sleep(0.05)
|
||||
return httpx.Response(201)
|
||||
|
||||
policy_route: Final = respx.post(endpoint + fixture["policy_path"]).mock(side_effect=load_policy)
|
||||
respx.post(url__startswith=endpoint + "/secrets/").respond(status_code=201) # pyright: ignore[reportUnknownMemberType] # respx route stubs leave response builder partially unknown
|
||||
|
||||
results: Final = await asyncio.gather(
|
||||
*(manager.async_write_secret(f"concurrent-{index}", "v") for index in range(4)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # legacy secret manager API is untyped
|
||||
)
|
||||
|
||||
assert policy_route.call_count == 4
|
||||
assert [result["status"] for result in results] == ["success"] * 4
|
||||
|
||||
|
||||
def test_missing_credentials_raise_value_error(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True)
|
||||
for name in (
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ import sys
|
|||
import time
|
||||
from io import StringIO
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
from typing import Final, List
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel, computed_field
|
||||
|
|
@ -1584,6 +1584,8 @@ def _emit_access_line(full_path: str) -> str:
|
|||
handler = logging.StreamHandler(stream)
|
||||
handler.setFormatter(AccessFormatter('%(client_addr)s - "%(request_line)s" %(status_code)s', use_colors=False))
|
||||
saved_level, saved_propagate = logger.level, logger.propagate
|
||||
saved_filters: Final = logger.filters[:]
|
||||
logger.filters = [f for f in saved_filters if type(f).__module__.split(".")[0] == "litellm"]
|
||||
logger.addHandler(handler)
|
||||
logger.setLevel(logging.INFO)
|
||||
logger.propagate = False
|
||||
|
|
@ -1593,6 +1595,7 @@ def _emit_access_line(full_path: str) -> str:
|
|||
logger.removeHandler(handler)
|
||||
logger.setLevel(saved_level)
|
||||
logger.propagate = saved_propagate
|
||||
logger.filters = saved_filters
|
||||
return stream.getvalue()
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import litellm
|
|||
from litellm.cost_calculator import default_video_cost_calculator
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
|
||||
from litellm.llms.gemini.videos.transformation import GeminiVideoConfig
|
||||
|
|
@ -988,6 +989,7 @@ class TestVideoLogging:
|
|||
"""
|
||||
custom_logger = self.TestVideoLogger()
|
||||
litellm.logging_callback_manager._reset_all_callbacks()
|
||||
await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0)
|
||||
litellm.callbacks = [custom_logger]
|
||||
|
||||
# Mock video generation response
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue