mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
Merge branch 'litellm_internal_staging' into litellm_/e2e-test-coverage-c87d3a
This commit is contained in:
commit
9d03b46889
93 changed files with 14567 additions and 1454 deletions
18
.github/workflows/check-ui-api-types.yml
vendored
18
.github/workflows/check-ui-api-types.yml
vendored
|
|
@ -83,6 +83,24 @@ jobs:
|
|||
if: steps.changes.outputs.relevant == 'true'
|
||||
run: uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
|
||||
|
||||
- name: Regenerate the lazy OpenAPI snapshot
|
||||
if: steps.changes.outputs.relevant == 'true'
|
||||
run: uv run --no-sync python -m litellm.proxy._lazy_openapi_snapshot
|
||||
|
||||
- name: Fail if the lazy OpenAPI snapshot is stale
|
||||
if: steps.changes.outputs.relevant == 'true'
|
||||
run: |
|
||||
if ! git diff --exit-code -- litellm/proxy/_lazy_openapi_snapshot.json; then
|
||||
echo "::error file=litellm/proxy/_lazy_openapi_snapshot.json::The lazy OpenAPI snapshot is out of sync with the lazily loaded routes."
|
||||
echo ""
|
||||
echo "A lazily loaded route or model changed without regenerating the snapshot that /openapi.json serves for unloaded features."
|
||||
echo "To fix, run from the repo root:"
|
||||
echo " uv run python -m litellm.proxy._lazy_openapi_snapshot"
|
||||
echo "then run npm run gen:api from ui/litellm-dashboard and commit both files."
|
||||
exit 1
|
||||
fi
|
||||
echo "_lazy_openapi_snapshot.json is in sync with the lazily loaded routes."
|
||||
|
||||
- name: Set up Node.js
|
||||
if: steps.changes.outputs.relevant == 'true'
|
||||
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0
|
||||
|
|
|
|||
|
|
@ -84,7 +84,7 @@
|
|||
"limit": 56
|
||||
},
|
||||
"reportPrivateUsage": {
|
||||
"limit": 1810
|
||||
"limit": 1808
|
||||
},
|
||||
"reportRedeclaration": {
|
||||
"limit": 8
|
||||
|
|
@ -135,7 +135,7 @@
|
|||
"limit": 21
|
||||
},
|
||||
"reportUnusedFunction": {
|
||||
"limit": 139
|
||||
"limit": 138
|
||||
},
|
||||
"reportUnusedImport": {
|
||||
"limit": 544
|
||||
|
|
|
|||
|
|
@ -45,7 +45,7 @@ class CustomBatchLogger(CustomLogger):
|
|||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
async def periodic_flush(self):
|
||||
async def periodic_flush(self) -> None:
|
||||
while True:
|
||||
await asyncio.sleep(self.flush_interval)
|
||||
verbose_logger.debug("CustomLogger periodic flush after %s seconds", self.flush_interval)
|
||||
|
|
|
|||
|
|
@ -149,7 +149,6 @@ class PromptManager:
|
|||
)
|
||||
self.prompts[template_id] = template
|
||||
except Exception:
|
||||
# Optional: print(f"Error loading prompt from JSON: {template_id}")
|
||||
pass
|
||||
|
||||
def _load_prompt_file(self, file_path: str | Path, prompt_id: str) -> PromptTemplate:
|
||||
|
|
|
|||
395
litellm/integrations/newrelic/newrelic_metrics.py
Normal file
395
litellm/integrations/newrelic/newrelic_metrics.py
Normal file
|
|
@ -0,0 +1,395 @@
|
|||
"""
|
||||
New Relic Metric API Integration - sends per-team cost/usage metrics to /metric/v1
|
||||
|
||||
NR Reference API: https://docs.newrelic.com/docs/data-apis/ingest-apis/metric-api/introduction-metric-api/
|
||||
|
||||
`async_log_success_event` / `async_log_failure_event` queue one record per request;
|
||||
at flush the queue is aggregated by (team, model group, model, provider, status)
|
||||
into count/summary metrics. `interval.ms` is the real window between flushes,
|
||||
computed at flush time.
|
||||
|
||||
Team-scoped by construction: the ingest key is injected explicitly and there is
|
||||
deliberately no environment-variable fallback, so a team's metrics are never sent
|
||||
with the proxy operator's credentials (mirrors ``allow_env_credentials=False`` on
|
||||
the Datadog team logger).
|
||||
|
||||
Error policy on flush: 4xx drops the batch (a retry would fail identically; 403
|
||||
is a permanent credential failure), 5xx/network re-queues capped at
|
||||
``max_queue_size`` records with the oldest dropped.
|
||||
|
||||
For batching specific details see CustomBatchLogger class
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import gzip
|
||||
import time
|
||||
import traceback
|
||||
from collections.abc import Mapping
|
||||
from math import ceil
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from httpx import HTTPStatusError, Response
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.custom_batch_logger import CustomBatchLogger
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.types.integrations.newrelic import (
|
||||
NEWRELIC_DEFAULT_REGION,
|
||||
NEWRELIC_METRIC_ATTRIBUTE_MAX_LEN,
|
||||
NEWRELIC_METRIC_COMPLETION_TOKENS,
|
||||
NEWRELIC_METRIC_COST_USD,
|
||||
NEWRELIC_METRIC_ENDPOINT_BY_REGION,
|
||||
NEWRELIC_METRIC_PROMPT_TOKENS,
|
||||
NEWRELIC_METRIC_REQUEST_DURATION_MS,
|
||||
NEWRELIC_METRIC_REQUESTS,
|
||||
NEWRELIC_METRIC_TOTAL_TOKENS,
|
||||
NEWRELIC_METRICS_MAX_BATCH_SIZE,
|
||||
NEWRELIC_METRICS_MAX_DRAIN_PASSES,
|
||||
NEWRELIC_METRICS_MAX_RETRY_QUEUE_SIZE,
|
||||
NewRelicCountMetric,
|
||||
NewRelicMetric,
|
||||
NewRelicMetricCommon,
|
||||
NewRelicMetricEnvelope,
|
||||
NewRelicMetricRecord,
|
||||
NewRelicSummaryMetric,
|
||||
NewRelicSummaryValue,
|
||||
)
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
|
||||
# 408 (request timeout) and 429 (rate limit) are transient client errors the
|
||||
# Metric API expects a retry on, unlike 400/403 which a retry would only repeat.
|
||||
_RETRYABLE_CLIENT_STATUSES: Final = frozenset({408, 429})
|
||||
|
||||
|
||||
def resolve_newrelic_metric_endpoint(newrelic_region: str | None) -> str:
|
||||
if not newrelic_region:
|
||||
return NEWRELIC_METRIC_ENDPOINT_BY_REGION[NEWRELIC_DEFAULT_REGION]
|
||||
endpoint: Final = NEWRELIC_METRIC_ENDPOINT_BY_REGION.get(newrelic_region.lower())
|
||||
if endpoint is None:
|
||||
verbose_logger.warning(
|
||||
"New Relic: unknown newrelic_region %r; supported regions: %s. Using the default (US) endpoint.",
|
||||
newrelic_region,
|
||||
", ".join(sorted(NEWRELIC_METRIC_ENDPOINT_BY_REGION)),
|
||||
)
|
||||
return NEWRELIC_METRIC_ENDPOINT_BY_REGION[NEWRELIC_DEFAULT_REGION]
|
||||
return endpoint
|
||||
|
||||
|
||||
def _metric_record_from_payload(standard_logging_object: StandardLoggingPayload) -> NewRelicMetricRecord:
|
||||
metadata: Final = standard_logging_object.get("metadata")
|
||||
team_id: Final = ((metadata.get("user_api_key_team_id") or metadata.get("team_id")) if metadata else None) or ""
|
||||
team_alias: Final = (
|
||||
(metadata.get("user_api_key_team_alias") or metadata.get("team_alias")) if metadata else None
|
||||
) or ""
|
||||
return NewRelicMetricRecord(
|
||||
team_id=team_id,
|
||||
team_alias=team_alias,
|
||||
model_group=standard_logging_object.get("model_group") or "",
|
||||
model=standard_logging_object.get("model") or "",
|
||||
custom_llm_provider=standard_logging_object.get("custom_llm_provider") or "",
|
||||
status=str(standard_logging_object.get("status") or "success"),
|
||||
response_cost=float(standard_logging_object.get("response_cost") or 0.0),
|
||||
prompt_tokens=int(standard_logging_object.get("prompt_tokens") or 0),
|
||||
completion_tokens=int(standard_logging_object.get("completion_tokens") or 0),
|
||||
total_tokens=int(standard_logging_object.get("total_tokens") or 0),
|
||||
duration_ms=float(standard_logging_object.get("response_time") or 0.0) * 1000.0,
|
||||
)
|
||||
|
||||
|
||||
def _bucket_metrics(bucket_records: tuple[NewRelicMetricRecord, ...]) -> tuple[NewRelicMetric, ...]:
|
||||
first: Final = bucket_records[0]
|
||||
attributes: Final[Mapping[str, str]] = { # mutable-ok: JSON leaf; safe_dumps stringifies MappingProxyType
|
||||
key: value[:NEWRELIC_METRIC_ATTRIBUTE_MAX_LEN]
|
||||
for key, value in (
|
||||
("team_id", first.team_id),
|
||||
("team_alias", first.team_alias),
|
||||
("model_group", first.model_group),
|
||||
("model", first.model),
|
||||
("custom_llm_provider", first.custom_llm_provider),
|
||||
("status", first.status),
|
||||
)
|
||||
if value
|
||||
}
|
||||
durations: Final = tuple(record.duration_ms for record in bucket_records)
|
||||
counts: Final[tuple[tuple[str, float], ...]] = (
|
||||
(NEWRELIC_METRIC_REQUESTS, float(len(bucket_records))),
|
||||
(NEWRELIC_METRIC_COST_USD, sum(record.response_cost for record in bucket_records)),
|
||||
(NEWRELIC_METRIC_PROMPT_TOKENS, float(sum(record.prompt_tokens for record in bucket_records))),
|
||||
(NEWRELIC_METRIC_COMPLETION_TOKENS, float(sum(record.completion_tokens for record in bucket_records))),
|
||||
(NEWRELIC_METRIC_TOTAL_TOKENS, float(sum(record.total_tokens for record in bucket_records))),
|
||||
)
|
||||
count_metrics: Final[tuple[NewRelicMetric, ...]] = tuple(
|
||||
NewRelicCountMetric(name=name, type="count", value=value, attributes=attributes) for name, value in counts
|
||||
)
|
||||
summary_metric: Final = NewRelicSummaryMetric(
|
||||
name=NEWRELIC_METRIC_REQUEST_DURATION_MS,
|
||||
type="summary",
|
||||
value=NewRelicSummaryValue(
|
||||
count=len(durations),
|
||||
sum=sum(durations),
|
||||
min=min(durations),
|
||||
max=max(durations),
|
||||
),
|
||||
attributes=attributes,
|
||||
)
|
||||
return (*count_metrics, summary_metric)
|
||||
|
||||
|
||||
def build_metric_payload(
|
||||
records: tuple[NewRelicMetricRecord, ...],
|
||||
*,
|
||||
window_start: float,
|
||||
now: float,
|
||||
) -> tuple[NewRelicMetricEnvelope, ...]:
|
||||
"""Aggregates records into one Metric API envelope for the flush window."""
|
||||
interval_ms: Final = max(1, int((now - window_start) * 1000))
|
||||
bucket_keys: Final = tuple(dict.fromkeys(record.bucket_key for record in records))
|
||||
metrics: Final = tuple(
|
||||
metric
|
||||
for key in bucket_keys
|
||||
for metric in _bucket_metrics(tuple(record for record in records if record.bucket_key == key))
|
||||
)
|
||||
common: Final[NewRelicMetricCommon] = {
|
||||
"timestamp": int(window_start * 1000),
|
||||
"interval.ms": interval_ms,
|
||||
}
|
||||
return (NewRelicMetricEnvelope(common=common, metrics=metrics),)
|
||||
|
||||
|
||||
class NewRelicMetricsLogger(CustomBatchLogger):
|
||||
def __init__(
|
||||
self,
|
||||
newrelic_api_key: str,
|
||||
newrelic_region: str | None = None,
|
||||
) -> None:
|
||||
if not newrelic_api_key:
|
||||
raise ValueError(
|
||||
"newrelic_api_key is required for NewRelicMetricsLogger; "
|
||||
"team-scoped metrics never fall back to environment credentials"
|
||||
)
|
||||
self.newrelic_api_key: Final = newrelic_api_key
|
||||
self.metric_api_url: Final = resolve_newrelic_metric_endpoint(newrelic_region)
|
||||
self.async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback)
|
||||
self._stopped: bool = False
|
||||
self._drain_lock = asyncio.Lock()
|
||||
asyncio.create_task(self.periodic_flush())
|
||||
self.flush_lock = asyncio.Lock()
|
||||
super().__init__(
|
||||
flush_lock=self.flush_lock,
|
||||
batch_size=NEWRELIC_METRICS_MAX_BATCH_SIZE,
|
||||
max_queue_size=NEWRELIC_METRICS_MAX_RETRY_QUEUE_SIZE,
|
||||
)
|
||||
|
||||
def stop(self) -> None:
|
||||
"""Ends the periodic flush loop; called on DynamicLoggingCache eviction.
|
||||
|
||||
Schedules one final drain of anything still queued, so eviction never
|
||||
silently discards records. Guarded so it can never raise into the
|
||||
cache's eviction path.
|
||||
"""
|
||||
self._stopped = True
|
||||
try:
|
||||
asyncio.get_running_loop().create_task(self._final_drain())
|
||||
except Exception: # noqa: BLE001 # no running loop / shutdown; the periodic loop's final drain still runs
|
||||
verbose_logger.debug("New Relic Metrics: could not schedule final drain on stop()", exc_info=True)
|
||||
|
||||
async def _drain_with_retry(self) -> None:
|
||||
"""Deliver everything queued on a stopped logger, or drop it with a log.
|
||||
|
||||
A stopped logger has no periodic loop left, so every post-stop path
|
||||
funnels through here. ``_drain_lock`` serializes drains: a callback that
|
||||
appends and starts its own drain queues behind the running one instead
|
||||
of racing it. Each pass attempts the whole current queue in
|
||||
``batch_size`` chunks, unlike the periodic path it does not stop at the
|
||||
first failing chunk, so a persistently failing head never starves the
|
||||
tail. Only after ``_MAX_DRAIN_PASSES`` against a permanently failing
|
||||
destination is the remainder dropped, and then only the records that were
|
||||
queued when this drain began, so every dropped record got the full retry
|
||||
budget: a record a callback appended mid-drain is not in that snapshot,
|
||||
so it is left for its own serialized drain rather than dropped after
|
||||
fewer attempts, and is never stranded.
|
||||
"""
|
||||
async with self._drain_lock:
|
||||
attempted: Final = tuple(self.log_queue)
|
||||
for _pass in range(NEWRELIC_METRICS_MAX_DRAIN_PASSES):
|
||||
await self._drain_flush_once()
|
||||
if not self.log_queue:
|
||||
return
|
||||
if _pass < NEWRELIC_METRICS_MAX_DRAIN_PASSES - 1:
|
||||
await asyncio.sleep(2**_pass)
|
||||
async with self.flush_lock:
|
||||
tried_ids: Final = frozenset(id(record) for record in attempted)
|
||||
survivors: Final = tuple(record for record in self.log_queue if id(record) not in tried_ids)
|
||||
dropped: Final = len(self.log_queue) - len(survivors)
|
||||
if dropped:
|
||||
verbose_logger.warning(
|
||||
"New Relic Metrics: dropping %s records after %s drain passes",
|
||||
dropped,
|
||||
NEWRELIC_METRICS_MAX_DRAIN_PASSES,
|
||||
)
|
||||
self.log_queue[:] = list(survivors) # mutable-ok: leave late arrivals for the next serialized drain
|
||||
|
||||
async def _drain_flush_once(self) -> None:
|
||||
"""Attempt every queued record once, in ``batch_size`` chunks, without
|
||||
stopping at the first failing chunk so a persistently failing head does
|
||||
not starve the tail (the periodic ``flush_queue`` deliberately stops
|
||||
instead). Takes the queue under ``flush_lock`` and re-queues only the
|
||||
chunks a 5xx/network error left undelivered, so records a concurrent
|
||||
request appends during the sends survive for the next pass."""
|
||||
async with self.flush_lock:
|
||||
pending: Final = tuple(self.log_queue)
|
||||
window_start: Final = self.last_flush_time
|
||||
self.last_flush_time = time.time()
|
||||
del self.log_queue[:]
|
||||
if not pending:
|
||||
return
|
||||
chunks: Final = tuple(
|
||||
pending[start : start + self.batch_size] for start in range(0, len(pending), self.batch_size)
|
||||
)
|
||||
delivered: Final = tuple([await self._classify_and_send(chunk, window_start) for chunk in chunks])
|
||||
failed: Final = tuple(record for chunk, ok in zip(chunks, delivered) for record in (() if ok else chunk))
|
||||
if failed:
|
||||
self._requeue(failed)
|
||||
|
||||
async def _final_drain(self) -> None:
|
||||
await self._drain_with_retry()
|
||||
|
||||
async def periodic_flush(self) -> None:
|
||||
while not self._stopped:
|
||||
await asyncio.sleep(self.flush_interval)
|
||||
if self._stopped:
|
||||
break
|
||||
await self.flush_queue()
|
||||
await self._final_drain()
|
||||
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None:
|
||||
try:
|
||||
await self._log_async_event(standard_logging_object=kwargs.get("standard_logging_object", None))
|
||||
except Exception as e: # noqa: BLE001 # logging must never break the request path
|
||||
verbose_logger.exception("New Relic Metrics Layer Error - %s\n%s", e, traceback.format_exc())
|
||||
|
||||
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None:
|
||||
try:
|
||||
await self._log_async_event(standard_logging_object=kwargs.get("standard_logging_object", None))
|
||||
except Exception as e: # noqa: BLE001 # logging must never break the request path
|
||||
verbose_logger.exception("New Relic Metrics Layer Error - %s\n%s", e, traceback.format_exc())
|
||||
|
||||
async def _log_async_event(self, standard_logging_object: StandardLoggingPayload | None) -> None:
|
||||
if standard_logging_object is None:
|
||||
raise ValueError("standard_logging_object not found in kwargs")
|
||||
self.log_queue.append(_metric_record_from_payload(standard_logging_object))
|
||||
if self._stopped:
|
||||
# A stopped logger has no periodic loop left; an in-flight callback
|
||||
# that appends after the eviction drain delivers its own record.
|
||||
await self._drain_with_retry()
|
||||
return
|
||||
if len(self.log_queue) >= self.batch_size:
|
||||
await self.flush_queue()
|
||||
|
||||
async def flush_queue(self) -> None:
|
||||
async with self.flush_lock:
|
||||
window_start: Final = self.last_flush_time
|
||||
self.last_flush_time = time.time()
|
||||
queued: Final = len(self.log_queue)
|
||||
if not queued:
|
||||
return
|
||||
verbose_logger.debug("New Relic Metrics: Flushing %s queued records", queued)
|
||||
# Bounded by what is queued now: records appended mid-flush belong to
|
||||
# the next window, and looping until empty would never end under load.
|
||||
for _chunk in range(ceil(queued / self.batch_size)):
|
||||
if not await self.async_send_batch(window_start=window_start):
|
||||
return
|
||||
|
||||
async def async_send_batch(self, window_start: float | None = None) -> bool:
|
||||
"""Sends the oldest ``batch_size`` records only, so a queue grown past that
|
||||
by re-queues cannot breach the Metric API data point cap in one request.
|
||||
Returns False once a chunk fails and is re-queued, so the caller stops."""
|
||||
if not self.log_queue:
|
||||
return False
|
||||
|
||||
batch_to_send: Final[tuple[NewRelicMetricRecord, ...]] = tuple(self.log_queue[: self.batch_size])
|
||||
del self.log_queue[: len(batch_to_send)]
|
||||
|
||||
delivered: Final = await self._classify_and_send(
|
||||
batch_to_send, window_start if window_start is not None else self.last_flush_time
|
||||
)
|
||||
if not delivered:
|
||||
self._requeue(batch_to_send)
|
||||
return delivered
|
||||
|
||||
async def _classify_and_send(self, batch: tuple[NewRelicMetricRecord, ...], window_start: float) -> bool:
|
||||
"""Send one chunk and classify the outcome, never touching the queue.
|
||||
Returns True when the batch is done with (delivered on any 2xx, or a 4xx
|
||||
a retry would only repeat, 403 being a permanent bad-key rejection), and
|
||||
False when a 5xx or network error means the caller should re-queue it.
|
||||
|
||||
``AsyncHTTPHandler.post`` raises ``HTTPStatusError`` on any non-2xx, so a
|
||||
4xx never returns a response here; the status is read off the raised
|
||||
error to keep the client-error path (drop) distinct from 5xx (retry)."""
|
||||
payload: Final = build_metric_payload(records=batch, window_start=window_start, now=time.time())
|
||||
try:
|
||||
status = (
|
||||
await self.async_send_compressed_data(payload)
|
||||
).status_code # rebind-ok: reassigned from the raised HTTPStatusError below
|
||||
except HTTPStatusError as e:
|
||||
status = e.response.status_code
|
||||
except Exception as e: # noqa: BLE001 # transport/network failure re-queues the batch
|
||||
verbose_logger.warning(
|
||||
"New Relic Metrics: network error sending %s records, will retry - %s",
|
||||
len(batch),
|
||||
e,
|
||||
)
|
||||
return False
|
||||
|
||||
if 200 <= status < 300:
|
||||
return True
|
||||
|
||||
if 400 <= status < 500 and status not in _RETRYABLE_CLIENT_STATUSES:
|
||||
verbose_logger.warning(
|
||||
"New Relic Metrics: %s from Metric API%s, dropping %s records.",
|
||||
status,
|
||||
" (permanent credential failure: invalid or revoked team ingest key)" if status == 403 else "",
|
||||
len(batch),
|
||||
)
|
||||
return True
|
||||
|
||||
verbose_logger.warning(
|
||||
"New Relic Metrics: %s from Metric API, will retry %s records",
|
||||
status,
|
||||
len(batch),
|
||||
)
|
||||
return False
|
||||
|
||||
def _requeue(self, batch: tuple[NewRelicMetricRecord, ...]) -> None:
|
||||
"""Prepends ``batch`` in place (never by assignment: records appended by
|
||||
concurrent requests during the flush await must survive), keeping
|
||||
chronological order so the cap drops the oldest records first."""
|
||||
self.log_queue[:0] = batch
|
||||
overflow: Final = len(self.log_queue) - self.max_queue_size
|
||||
if overflow > 0:
|
||||
del self.log_queue[:overflow]
|
||||
verbose_logger.warning(
|
||||
"New Relic Metrics: retry queue exceeded max_queue_size=%s; dropped %s oldest records.",
|
||||
self.max_queue_size,
|
||||
overflow,
|
||||
)
|
||||
|
||||
async def async_send_compressed_data(self, payload: tuple[NewRelicMetricEnvelope, ...]) -> Response:
|
||||
compressed_data: Final = gzip.compress(safe_dumps(payload).encode("utf-8"))
|
||||
headers: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{
|
||||
"Content-Type": "application/json",
|
||||
"Content-Encoding": "gzip",
|
||||
"Api-Key": self.newrelic_api_key,
|
||||
}
|
||||
)
|
||||
return await self.async_client.post(
|
||||
url=self.metric_api_url,
|
||||
data=compressed_data,
|
||||
headers=headers,
|
||||
)
|
||||
90
litellm/integrations/newrelic/newrelic_team_handler.py
Normal file
90
litellm/integrations/newrelic/newrelic_team_handler.py
Normal file
|
|
@ -0,0 +1,90 @@
|
|||
"""
|
||||
New Relic Team Handler
|
||||
|
||||
Used to get the NewRelicMetricsLogger for a given request.
|
||||
Handles Key/Team Based New Relic metrics, following the same pattern as DataDogHandler.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.litellm_logging import StandardCallbackDynamicParams
|
||||
|
||||
from .newrelic_metrics import NewRelicMetricsLogger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import DynamicLoggingCache
|
||||
|
||||
|
||||
class NewRelicLoggingConfig(TypedDict):
|
||||
newrelic_api_key: ReadOnly[str | None]
|
||||
newrelic_region: ReadOnly[str | None]
|
||||
|
||||
|
||||
class NewRelicHandler:
|
||||
@staticmethod
|
||||
def get_newrelic_logger_for_request(
|
||||
standard_callback_dynamic_params: StandardCallbackDynamicParams,
|
||||
in_memory_dynamic_logger_cache: "DynamicLoggingCache",
|
||||
) -> NewRelicMetricsLogger:
|
||||
"""
|
||||
Get a team-scoped NewRelicMetricsLogger for a given request.
|
||||
|
||||
Resolves and caches per-team NewRelicMetricsLogger instances using
|
||||
DynamicLoggingCache, keyed by the team's New Relic credentials. Each unique
|
||||
set of credentials gets its own logger instance with its own batch/flush loop.
|
||||
|
||||
Note: This handler is only called when a team-scoped newrelic_api_key is
|
||||
present. The trace logger for the ``newrelic`` callback (OTel v2 / legacy
|
||||
agent) is managed separately by _init_custom_logger_compatible_class via
|
||||
_in_memory_loggers.
|
||||
"""
|
||||
_credentials: Final = NewRelicHandler.get_dynamic_newrelic_logging_config(
|
||||
standard_callback_dynamic_params=standard_callback_dynamic_params,
|
||||
)
|
||||
|
||||
temp_newrelic_logger = in_memory_dynamic_logger_cache.get_cache(
|
||||
credentials=_credentials, service_name="newrelic"
|
||||
)
|
||||
|
||||
if temp_newrelic_logger is None:
|
||||
temp_newrelic_logger = NewRelicHandler._create_newrelic_logger_from_credentials(
|
||||
credentials=_credentials,
|
||||
in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache,
|
||||
)
|
||||
|
||||
return temp_newrelic_logger
|
||||
|
||||
@staticmethod
|
||||
def _create_newrelic_logger_from_credentials(
|
||||
credentials: NewRelicLoggingConfig,
|
||||
in_memory_dynamic_logger_cache: "DynamicLoggingCache",
|
||||
) -> NewRelicMetricsLogger:
|
||||
newrelic_logger: Final = NewRelicMetricsLogger(
|
||||
newrelic_api_key=credentials.get("newrelic_api_key") or "",
|
||||
newrelic_region=credentials.get("newrelic_region"),
|
||||
)
|
||||
in_memory_dynamic_logger_cache.set_cache(
|
||||
credentials=credentials,
|
||||
service_name="newrelic",
|
||||
logging_obj=newrelic_logger,
|
||||
)
|
||||
verbose_logger.debug("New Relic: Created and cached new NewRelicMetricsLogger for team-scoped credentials")
|
||||
return newrelic_logger
|
||||
|
||||
@staticmethod
|
||||
def get_dynamic_newrelic_logging_config(
|
||||
standard_callback_dynamic_params: StandardCallbackDynamicParams,
|
||||
) -> NewRelicLoggingConfig:
|
||||
return NewRelicLoggingConfig(
|
||||
newrelic_api_key=standard_callback_dynamic_params.get("newrelic_api_key"),
|
||||
newrelic_region=standard_callback_dynamic_params.get("newrelic_region"),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _dynamic_newrelic_credentials_are_passed(
|
||||
standard_callback_dynamic_params: StandardCallbackDynamicParams,
|
||||
) -> bool:
|
||||
return standard_callback_dynamic_params.get("newrelic_api_key") is not None
|
||||
|
|
@ -2341,6 +2341,7 @@ def exception_type(
|
|||
litellm_response_headers: Final = _get_response_headers(original_exception=original_exception)
|
||||
try:
|
||||
error_str = redact_string(str(original_exception)) if _ENABLE_SECRET_REDACTION else str(original_exception)
|
||||
extra_information = ""
|
||||
if model or custom_llm_provider:
|
||||
if hasattr(original_exception, "message"):
|
||||
error_str = (
|
||||
|
|
@ -2357,7 +2358,6 @@ def exception_type(
|
|||
# Common Extra information needed for all providers
|
||||
# We pass num retries, api_base, vertex_deployment etc to the exception here
|
||||
################################################################################
|
||||
extra_information = ""
|
||||
try:
|
||||
_api_base: Final = litellm.get_api_base(model=model, optional_params=extra_kwargs)
|
||||
messages: Final = litellm.get_first_chars_messages(kwargs=completion_kwargs)
|
||||
|
|
|
|||
|
|
@ -613,37 +613,60 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
processed_list: Final[list[str | Callable | CustomLogger]] = []
|
||||
for callback in callback_list:
|
||||
if isinstance(callback, str) and callback in litellm._known_custom_logger_compatible_callbacks:
|
||||
# For callbacks that support team-scoped credentials (e.g. datadog),
|
||||
# pass only the relevant dynamic params as custom_logger_init_args.
|
||||
_custom_logger_init_args: dict | None = None
|
||||
if callback == "datadog":
|
||||
# dd_* params are blocked from standard_callback_dynamic_params
|
||||
# (request-level security); only the proxy-stamped team/key
|
||||
# callback vars are admin-configured and trusted.
|
||||
_custom_logger_init_args = {k: v for k, v in self._trusted_callback_vars if k.startswith("dd_")}
|
||||
|
||||
callback_class = _init_custom_logger_compatible_class(
|
||||
callback,
|
||||
internal_usage_cache=None,
|
||||
llm_router=None,
|
||||
custom_logger_init_args=_custom_logger_init_args,
|
||||
)
|
||||
if callback_class is not None:
|
||||
processed_list.append(callback_class)
|
||||
for callback_instance in self._resolve_dynamic_callback_string(callback):
|
||||
processed_list.append(callback_instance)
|
||||
|
||||
# If processing dynamic_success_callbacks, add to dynamic_async_success_callbacks
|
||||
if dynamic_callbacks_type == "success":
|
||||
if self.dynamic_async_success_callbacks is None:
|
||||
self.dynamic_async_success_callbacks = []
|
||||
self.dynamic_async_success_callbacks.append(callback_class)
|
||||
self.dynamic_async_success_callbacks.append(callback_instance)
|
||||
elif dynamic_callbacks_type == "failure":
|
||||
if self.dynamic_async_failure_callbacks is None:
|
||||
self.dynamic_async_failure_callbacks = []
|
||||
self.dynamic_async_failure_callbacks.append(callback_class)
|
||||
self.dynamic_async_failure_callbacks.append(callback_instance)
|
||||
else:
|
||||
processed_list.append(callback)
|
||||
return processed_list
|
||||
|
||||
def _resolve_dynamic_callback_string(self, callback: str) -> "tuple[CustomLogger, ...]":
|
||||
"""
|
||||
Resolve a known callback name to the logger instance(s) it dispatches to.
|
||||
|
||||
For callbacks that support team-scoped credentials (datadog, newrelic),
|
||||
only the proxy-stamped team/key callback vars are passed as
|
||||
custom_logger_init_args: dd_*/newrelic_* params are blocked from
|
||||
standard_callback_dynamic_params (request-level security), so the
|
||||
trusted-vars channel is the only way credentials reach a per-team logger.
|
||||
"""
|
||||
_trusted_var_prefix: Final = "dd_" if callback == "datadog" else "newrelic_" if callback == "newrelic" else None
|
||||
_custom_logger_init_args: Final[dict | None] = (
|
||||
{k: v for k, v in self._trusted_callback_vars if k.startswith(_trusted_var_prefix)}
|
||||
if _trusted_var_prefix is not None
|
||||
else None
|
||||
)
|
||||
|
||||
callback_class: Final = _init_custom_logger_compatible_class(
|
||||
callback,
|
||||
internal_usage_cache=None,
|
||||
llm_router=None,
|
||||
custom_logger_init_args=_custom_logger_init_args,
|
||||
)
|
||||
if callback_class is None:
|
||||
return ()
|
||||
|
||||
# With team creds, "newrelic" resolves to the per-team METRICS logger;
|
||||
# resolve the name again without creds so the trace logger (OTel v2 /
|
||||
# legacy agent) keeps receiving this request.
|
||||
_newrelic_trace_class: Final = (
|
||||
_init_custom_logger_compatible_class(callback, internal_usage_cache=None, llm_router=None)
|
||||
if callback == "newrelic" and _custom_logger_init_args and _custom_logger_init_args.get("newrelic_api_key")
|
||||
else None
|
||||
)
|
||||
if _newrelic_trace_class is not None and _newrelic_trace_class is not callback_class:
|
||||
return (callback_class, _newrelic_trace_class)
|
||||
return (callback_class,)
|
||||
|
||||
def initialize_standard_callback_dynamic_params(self, kwargs: dict | None = None) -> StandardCallbackDynamicParams:
|
||||
"""
|
||||
Initialize the standard callback dynamic params from the kwargs
|
||||
|
|
@ -4642,6 +4665,19 @@ def _init_custom_logger_compatible_class(
|
|||
_in_memory_loggers.append(gitlab_logger)
|
||||
return gitlab_logger
|
||||
elif logging_integration == "newrelic":
|
||||
if custom_logger_init_args.get("newrelic_api_key"):
|
||||
# Team-scoped credentials: per-team METRICS logger, isolated per
|
||||
# credential set via DynamicLoggingCache. The trace logger for
|
||||
# this name stays on the global path below.
|
||||
from litellm.integrations.newrelic.newrelic_team_handler import (
|
||||
NewRelicHandler,
|
||||
)
|
||||
|
||||
return NewRelicHandler.get_newrelic_logger_for_request(
|
||||
standard_callback_dynamic_params=custom_logger_init_args,
|
||||
in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache,
|
||||
)
|
||||
|
||||
_v2 = _maybe_construct_otel_v2("newrelic", _in_memory_loggers)
|
||||
if _v2 is not None:
|
||||
return _v2
|
||||
|
|
|
|||
|
|
@ -7,7 +7,9 @@ from typing import Any, Final, Literal
|
|||
|
||||
import litellm
|
||||
from litellm.constants import OPENAI_FILE_SEARCH_COST_PER_1K_CALLS
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import _get_web_search_requests
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import (
|
||||
get_web_search_requests_from_usage,
|
||||
)
|
||||
from litellm.types.llms.openai import (
|
||||
FileSearchTool,
|
||||
ResponsesAPIResponse,
|
||||
|
|
@ -368,7 +370,7 @@ class StandardBuiltInToolCostTracking:
|
|||
get_anthropic_web_search_requests_from_response,
|
||||
)
|
||||
|
||||
if usage is not None and (_get_web_search_requests(getattr(usage, "server_tool_use", None)) is not None):
|
||||
if usage is not None and (get_web_search_requests_from_usage(usage) is not None):
|
||||
return usage
|
||||
web_search_requests: Final = get_anthropic_web_search_requests_from_response(response_object)
|
||||
if web_search_requests is None:
|
||||
|
|
@ -416,7 +418,7 @@ class StandardBuiltInToolCostTracking:
|
|||
# Anthropic Claude (direct API and Vertex AI) uses server_tool_use.web_search_requests.
|
||||
# Without this check, Claude ModelResponse always falls through to return False
|
||||
# and _handle_web_search_cost() is never called.
|
||||
if hasattr(usage, "server_tool_use") and _get_web_search_requests(usage.server_tool_use) is not None:
|
||||
if get_web_search_requests_from_usage(usage) is not None:
|
||||
return True
|
||||
# xAI reports usage.server_side_tool_usage_details.web_search_calls; a searched
|
||||
# answer with no url_citation annotations has no other chat-path signal
|
||||
|
|
@ -429,16 +431,12 @@ class StandardBuiltInToolCostTracking:
|
|||
response_object=response_object, output_type="web_search_call"
|
||||
)
|
||||
elif usage is not None:
|
||||
if (
|
||||
hasattr(usage, "server_tool_use")
|
||||
and _get_web_search_requests(usage.server_tool_use) is not None
|
||||
or (
|
||||
hasattr(usage, "prompt_tokens_details")
|
||||
and usage.prompt_tokens_details is not None
|
||||
and isinstance(usage.prompt_tokens_details, PromptTokensDetailsWrapper)
|
||||
and hasattr(usage.prompt_tokens_details, "web_search_requests")
|
||||
and usage.prompt_tokens_details.web_search_requests is not None
|
||||
)
|
||||
if get_web_search_requests_from_usage(usage) is not None or (
|
||||
hasattr(usage, "prompt_tokens_details")
|
||||
and usage.prompt_tokens_details is not None
|
||||
and isinstance(usage.prompt_tokens_details, PromptTokensDetailsWrapper)
|
||||
and hasattr(usage.prompt_tokens_details, "web_search_requests")
|
||||
and usage.prompt_tokens_details.web_search_requests is not None
|
||||
):
|
||||
return True
|
||||
if _usage_reports_server_side_web_search_calls(usage):
|
||||
|
|
|
|||
|
|
@ -72,7 +72,7 @@ def _get_token_detail_value(details: object, key: str) -> int | None:
|
|||
return value if isinstance(value, int) else None
|
||||
|
||||
|
||||
def _get_web_search_requests(server_tool_use: Any) -> int | None:
|
||||
def get_web_search_requests(server_tool_use: Any) -> int | None:
|
||||
"""
|
||||
Tolerantly read ``web_search_requests`` from a ``server_tool_use`` value
|
||||
that may be ``None``, a ``dict``, a ``ServerToolUse`` pydantic instance,
|
||||
|
|
@ -92,6 +92,16 @@ def _get_web_search_requests(server_tool_use: Any) -> int | None:
|
|||
return getattr(server_tool_use, "web_search_requests", None)
|
||||
|
||||
|
||||
def get_web_search_requests_from_usage(usage: Usage) -> int | None:
|
||||
"""Read ``web_search_requests`` from a ``Usage``'s ``server_tool_use``.
|
||||
|
||||
``Usage`` deletes unset optional fields from ``__dict__`` (see
|
||||
``SafeAttributeModel``), so direct attribute access can raise
|
||||
``AttributeError``; ``getattr`` with a default is required here.
|
||||
"""
|
||||
return get_web_search_requests(getattr(usage, "server_tool_use", None))
|
||||
|
||||
|
||||
def _is_above_128k(tokens: float) -> bool:
|
||||
if tokens > 128000:
|
||||
return True
|
||||
|
|
@ -889,11 +899,22 @@ def generic_cost_per_token(
|
|||
total_details: Final = text_tokens + cache_hit + audio_tokens + cache_creation + image_tokens + video_tokens
|
||||
has_double_counting: Final = (cache_hit > 0 or cache_creation > 0) and total_details > usage.prompt_tokens
|
||||
|
||||
if (text_tokens == 0 and prompt_tokens_details["image_count"] == 0) or has_double_counting:
|
||||
text_tokens = usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens - video_tokens
|
||||
if has_double_counting:
|
||||
# cached and per-modality counts are both subsets of prompt_tokens and may overlap, so a
|
||||
# modality can only bill what the cache did not already cover or the overlap is billed twice
|
||||
uncached_budget: Final = max(usage.prompt_tokens - cache_hit - cache_creation, 0)
|
||||
billable_audio: Final = min(audio_tokens, uncached_budget)
|
||||
billable_image: Final = min(image_tokens, uncached_budget - billable_audio)
|
||||
billable_video: Final = min(video_tokens, uncached_budget - billable_audio - billable_image)
|
||||
prompt_tokens_details["audio_tokens"] = billable_audio
|
||||
prompt_tokens_details["image_tokens"] = billable_image
|
||||
prompt_tokens_details["video_tokens"] = billable_video
|
||||
prompt_tokens_details["text_tokens"] = uncached_budget - billable_audio - billable_image - billable_video
|
||||
elif text_tokens == 0 and prompt_tokens_details["image_count"] == 0:
|
||||
# Clamp to zero: inconsistent streaming usage
|
||||
text_tokens = max(text_tokens, 0)
|
||||
prompt_tokens_details["text_tokens"] = text_tokens
|
||||
prompt_tokens_details["text_tokens"] = max(
|
||||
usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens - video_tokens, 0
|
||||
)
|
||||
|
||||
(
|
||||
prompt_base_cost,
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ import json
|
|||
from typing import Any, Final
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import _DEFAULT_TTL_FOR_HTTPX_CLIENTS
|
||||
|
||||
from ...caching import InMemoryCache
|
||||
|
|
@ -46,6 +47,15 @@ class LangfuseInMemoryCache(InMemoryCache):
|
|||
_created_langfuse_logger.Langfuse.flush()
|
||||
_created_langfuse_logger.Langfuse.shutdown()
|
||||
|
||||
# Loggers with a periodic flush task (e.g. NewRelicMetricsLogger) expose
|
||||
# stop() so eviction actually ends the task instead of leaking it.
|
||||
_evicted_stop: Final = getattr(self.cache_dict[key], "stop", None)
|
||||
if callable(_evicted_stop):
|
||||
try:
|
||||
_evicted_stop()
|
||||
except Exception: # noqa: BLE001 # a failing stop() must not block eviction
|
||||
verbose_logger.debug("DynamicLoggingCache: stop() raised during eviction", exc_info=True)
|
||||
|
||||
#########################################################
|
||||
# Call parent class to remove key from cache
|
||||
#########################################################
|
||||
|
|
|
|||
|
|
@ -8,9 +8,9 @@ from typing import TYPE_CHECKING, Final, Optional
|
|||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import (
|
||||
_get_web_search_requests,
|
||||
generic_cost_per_token,
|
||||
get_provider_specific_geo_multiplier,
|
||||
get_web_search_requests_from_usage,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -104,7 +104,7 @@ def get_cost_for_anthropic_web_search(
|
|||
|
||||
if usage is None:
|
||||
return 0.0
|
||||
web_search_requests: Final = _get_web_search_requests(getattr(usage, "server_tool_use", None))
|
||||
web_search_requests: Final = get_web_search_requests_from_usage(usage)
|
||||
if web_search_requests is None:
|
||||
return 0.0
|
||||
|
||||
|
|
|
|||
|
|
@ -99,6 +99,7 @@ from litellm.types.llms.anthropic import (
|
|||
ContextManagementResponse,
|
||||
MessageBlockDelta,
|
||||
MessageDelta,
|
||||
ServerToolUsage,
|
||||
StreamingContentBlockDeltaType,
|
||||
UsageDelta,
|
||||
UsageIteration,
|
||||
|
|
@ -1354,10 +1355,22 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
return explicit_value
|
||||
return cls._first_positive_prompt_tokens_detail_value(usage, ("cache_creation_tokens", "cache_write_tokens"))
|
||||
|
||||
@classmethod
|
||||
def _get_web_search_request_count(cls, usage: Usage) -> int:
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import (
|
||||
get_web_search_requests_from_usage,
|
||||
)
|
||||
|
||||
from_server_tool_use: Final = cls._positive_int(get_web_search_requests_from_usage(usage))
|
||||
if from_server_tool_use > 0:
|
||||
return from_server_tool_use
|
||||
return cls._first_positive_prompt_tokens_detail_value(usage, ("web_search_requests",))
|
||||
|
||||
@classmethod
|
||||
def _translate_openai_usage_to_anthropic_usage_delta(cls, usage: Usage) -> UsageDelta:
|
||||
cache_read_input_tokens: Final = cls._get_cache_read_input_tokens(usage)
|
||||
cache_creation_input_tokens: Final = cls._get_cache_creation_input_tokens(usage)
|
||||
web_search_requests: Final = cls._get_web_search_request_count(usage)
|
||||
input_tokens: Final = max(
|
||||
(usage.prompt_tokens or 0) - cache_read_input_tokens - cache_creation_input_tokens,
|
||||
0,
|
||||
|
|
@ -1371,6 +1384,11 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
usage_delta["cache_creation_input_tokens"] = cache_creation_input_tokens
|
||||
if cache_read_input_tokens > 0:
|
||||
usage_delta["cache_read_input_tokens"] = cache_read_input_tokens
|
||||
if web_search_requests > 0:
|
||||
return UsageDelta(
|
||||
**usage_delta,
|
||||
server_tool_use=ServerToolUsage(web_search_requests=web_search_requests),
|
||||
)
|
||||
return usage_delta
|
||||
|
||||
@classmethod
|
||||
|
|
|
|||
|
|
@ -1434,9 +1434,12 @@ class BaseAWSLLM:
|
|||
data: str | bytes,
|
||||
headers: dict,
|
||||
api_key: str | None = None,
|
||||
supports_bearer_token: bool = True,
|
||||
) -> AWSPreparedRequest:
|
||||
if api_key is not None:
|
||||
aws_bearer_token: str | None = api_key
|
||||
if not supports_bearer_token:
|
||||
aws_bearer_token: str | None = None
|
||||
elif api_key is not None:
|
||||
aws_bearer_token = api_key
|
||||
else:
|
||||
aws_bearer_token = get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
|
||||
|
||||
|
|
|
|||
|
|
@ -138,11 +138,6 @@ class BedrockRerankHandler(BaseAWSLLM):
|
|||
data: dict,
|
||||
optional_params: dict,
|
||||
) -> BedrockPreparedRequest:
|
||||
try:
|
||||
from botocore.auth import SigV4Auth
|
||||
from botocore.awsrequest import AWSRequest
|
||||
except ImportError:
|
||||
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
|
||||
boto3_credentials_info: Final = self._get_boto_credentials_from_optional_params(optional_params, model)
|
||||
|
||||
### SET RUNTIME ENDPOINT ###
|
||||
|
|
@ -153,24 +148,21 @@ class BedrockRerankHandler(BaseAWSLLM):
|
|||
)
|
||||
proxy_endpoint_url = proxy_endpoint_url.replace("bedrock-runtime", "bedrock-agent-runtime")
|
||||
proxy_endpoint_url = f"{proxy_endpoint_url}/rerank"
|
||||
sigv4: Final = SigV4Auth(
|
||||
boto3_credentials_info.credentials,
|
||||
"bedrock",
|
||||
boto3_credentials_info.aws_region_name,
|
||||
)
|
||||
# Make POST Request
|
||||
body: Final = json.dumps(data).encode("utf-8")
|
||||
|
||||
body: Final = json.dumps(data).encode("utf-8")
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if extra_headers is not None:
|
||||
headers = {"Content-Type": "application/json", **extra_headers}
|
||||
request: Final = AWSRequest(method="POST", url=proxy_endpoint_url, data=body, headers=headers)
|
||||
sigv4.add_auth(request)
|
||||
if (
|
||||
extra_headers is not None and "Authorization" in extra_headers
|
||||
): # prevent sigv4 from overwriting the auth header
|
||||
request.headers["Authorization"] = extra_headers["Authorization"]
|
||||
prepped: Final = request.prepare()
|
||||
|
||||
prepped: Final = self.get_request_headers(
|
||||
credentials=boto3_credentials_info.credentials,
|
||||
aws_region_name=boto3_credentials_info.aws_region_name,
|
||||
extra_headers=extra_headers,
|
||||
endpoint_url=proxy_endpoint_url,
|
||||
data=body,
|
||||
headers=headers,
|
||||
supports_bearer_token=False,
|
||||
)
|
||||
|
||||
return BedrockPreparedRequest(
|
||||
endpoint_url=proxy_endpoint_url,
|
||||
|
|
|
|||
|
|
@ -39,28 +39,35 @@ def cost_per_web_search_request(usage: "Usage", model_info: "ModelInfo") -> floa
|
|||
``model_info`` when available, falling back to $0.035 for models not
|
||||
yet updated in the pricing JSON.
|
||||
"""
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import (
|
||||
get_web_search_requests_from_usage,
|
||||
)
|
||||
from litellm.types.utils import PromptTokensDetailsWrapper
|
||||
|
||||
_DEFAULT_COST: Final = 35e-3
|
||||
search_costs: Final = model_info.get("search_context_cost_per_query") or {}
|
||||
_cost: Final = search_costs.get("search_context_size_medium", _DEFAULT_COST)
|
||||
|
||||
number_of_web_search_requests = 0
|
||||
if (
|
||||
usage is not None
|
||||
and usage.prompt_tokens_details is not None
|
||||
and isinstance(usage.prompt_tokens_details, PromptTokensDetailsWrapper)
|
||||
and hasattr(usage.prompt_tokens_details, "web_search_requests")
|
||||
and usage.prompt_tokens_details.web_search_requests is not None
|
||||
):
|
||||
number_of_web_search_requests = usage.prompt_tokens_details.web_search_requests
|
||||
requests_from_prompt_details: Final = (
|
||||
usage.prompt_tokens_details.web_search_requests
|
||||
if (
|
||||
usage is not None
|
||||
and usage.prompt_tokens_details is not None
|
||||
and isinstance(usage.prompt_tokens_details, PromptTokensDetailsWrapper)
|
||||
and hasattr(usage.prompt_tokens_details, "web_search_requests")
|
||||
and usage.prompt_tokens_details.web_search_requests is not None
|
||||
)
|
||||
else None
|
||||
)
|
||||
requests_from_server_tool_use: Final = get_web_search_requests_from_usage(usage)
|
||||
number_of_web_search_requests: Final = requests_from_prompt_details or requests_from_server_tool_use or 0
|
||||
|
||||
# per_prompt billing: clamp to 1 (flat fee per grounded API call)
|
||||
billing_mode: Final = model_info.get("web_search_billing_unit") or "per_prompt"
|
||||
if number_of_web_search_requests > 0 and billing_mode == "per_prompt":
|
||||
number_of_web_search_requests = 1
|
||||
billable_requests: Final = (
|
||||
1 if (number_of_web_search_requests > 0 and billing_mode == "per_prompt") else number_of_web_search_requests
|
||||
)
|
||||
|
||||
return _cost * number_of_web_search_requests
|
||||
return _cost * billable_requests
|
||||
|
||||
|
||||
GOOGLE_MAPS_GROUNDING_DEFAULT_COST_PER_QUERY: Final = 14e-3
|
||||
|
|
|
|||
|
|
@ -46,6 +46,7 @@ from litellm.proxy._types import (
|
|||
)
|
||||
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
|
||||
from litellm.proxy.auth.user_api_key_auth import (
|
||||
_get_bearer_token_or_received_api_key, # pyright: ignore[reportPrivateUsage] # shared x-litellm-api-key parser lives with user_api_key_auth
|
||||
_run_centralized_common_checks,
|
||||
user_api_key_auth,
|
||||
)
|
||||
|
|
@ -429,7 +430,10 @@ class MCPRequestHandler:
|
|||
# An explicit x-litellm-api-key is always a LiteLLM credential, even
|
||||
# for a delegated server, so validate it: identity / spend / rate
|
||||
# limits resolve and any stored upstream token can be forwarded.
|
||||
validated_user_api_key_auth = await user_api_key_auth(api_key=litellm_api_key, request=request)
|
||||
validated_user_api_key_auth = await user_api_key_auth(
|
||||
api_key=f"Bearer {_get_bearer_token_or_received_api_key(litellm_api_key)}",
|
||||
request=request,
|
||||
)
|
||||
elif MCPRequestHandler._target_servers_delegate_auth_to_upstream(
|
||||
path=request_route,
|
||||
mcp_servers=mcp_servers,
|
||||
|
|
|
|||
|
|
@ -239,5 +239,6 @@ def resolve_bridge_envelope(
|
|||
if opened.identity.server_id != expected_server_id:
|
||||
return BridgeEnvelopeInvalid()
|
||||
grant: Final = opened.grant
|
||||
upstream_authorization: Final = f"{grant.token_type} {grant.access_token.get_secret_value()}"
|
||||
authorization_scheme: Final = "Bearer" if grant.token_type.lower() == "bearer" else grant.token_type
|
||||
upstream_authorization: Final = f"{authorization_scheme} {grant.access_token.get_secret_value()}"
|
||||
return BridgeEnvelopeAdmitted(identity=opened.identity, upstream_authorization=SecretStr(upstream_authorization))
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -3,18 +3,27 @@ Per-feature OpenAPI snapshot for lazy-loaded routers.
|
|||
|
||||
The committed JSON is generated by `python -m litellm.proxy._lazy_openapi_snapshot`
|
||||
and consumed at runtime so /openapi.json can show full route info for unloaded
|
||||
features without importing them. No CI job regenerates this file; drift surfaces
|
||||
only indirectly through check-ui-api-types.yml, which rebuilds schema.d.ts from
|
||||
app.openapi() with the committed snapshot injected. After changing any lazily
|
||||
loaded route or this generator, rerun the module and commit the JSON, then run
|
||||
`npm run gen:api` in ui/litellm-dashboard and commit schema.d.ts.
|
||||
features without importing them. check-ui-api-types.yml (mirrored locally by
|
||||
`make check`) regenerates this file and fails when the committed copy differs,
|
||||
then rebuilds schema.d.ts from app.openapi() with the snapshot injected. After
|
||||
changing any lazily loaded route or this generator, rerun the module and commit
|
||||
the JSON, then run `npm run gen:api` in ui/litellm-dashboard and commit schema.d.ts.
|
||||
"""
|
||||
|
||||
import json
|
||||
import re
|
||||
import sys
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastapi import FastAPI
|
||||
|
||||
from litellm.proxy._lazy_features import LazyFeature
|
||||
|
||||
SNAPSHOT_FILE: Final = Path(__file__).parent / "_lazy_openapi_snapshot.json"
|
||||
HTTP_METHOD_SUFFIXES: Final = {
|
||||
|
|
@ -83,51 +92,84 @@ def _normalize_operation_ids(paths: dict[str, dict]) -> None:
|
|||
break
|
||||
|
||||
|
||||
def generate_snapshot() -> dict[str, dict]:
|
||||
class SnapshotFragment(TypedDict):
|
||||
paths: ReadOnly[Mapping[str, Mapping[str, object]]]
|
||||
components: ReadOnly[Mapping[str, Mapping[str, object]]]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SnapshotResult:
|
||||
fragments: Mapping[str, SnapshotFragment]
|
||||
skipped: tuple[str, ...]
|
||||
|
||||
|
||||
def _register_feature(app: "FastAPI", feat: "LazyFeature") -> str | None:
|
||||
import importlib
|
||||
|
||||
try:
|
||||
feat.register_fn(app, importlib.import_module(feat.module_path))
|
||||
except Exception as exc:
|
||||
sys.stderr.write(f"warning: skip {feat.name}: {exc}\n")
|
||||
return feat.name
|
||||
return None
|
||||
|
||||
|
||||
def _feature_fragment(app: "FastAPI", feat: "LazyFeature", used_operation_ids: set[str]) -> SnapshotFragment | None:
|
||||
from fastapi.openapi.utils import get_openapi
|
||||
|
||||
from litellm.proxy.proxy_server import ensure_unique_openapi_operation_ids
|
||||
|
||||
feat_routes: Final = [r for r in app.routes if feat.matches(getattr(r, "path", ""))]
|
||||
if not feat_routes:
|
||||
return None
|
||||
_stabilize_multi_method_route_ids(feat_routes)
|
||||
full: Final = get_openapi(title=app.title, version=app.version, routes=feat_routes)
|
||||
paths: Final = full.get("paths", {})
|
||||
_normalize_operation_ids(paths)
|
||||
# Group all of a feature's routes under one tag.
|
||||
for path_ops in paths.values():
|
||||
for method, op in path_ops.items():
|
||||
if isinstance(op, dict):
|
||||
operation_id = op.get("operationId")
|
||||
if isinstance(operation_id, str):
|
||||
for suffix in HTTP_METHOD_SUFFIXES:
|
||||
if operation_id.endswith(f"_{suffix}"):
|
||||
op["operationId"] = operation_id[: -len(suffix)] + method
|
||||
break
|
||||
op["tags"] = [feat.name]
|
||||
unique: Final = ensure_unique_openapi_operation_ids(full, used_operation_ids)
|
||||
return {
|
||||
"paths": paths,
|
||||
"components": {"schemas": unique.get("components", {}).get("schemas", {})},
|
||||
}
|
||||
|
||||
|
||||
def generate_snapshot() -> SnapshotResult:
|
||||
from litellm.proxy._lazy_features import LAZY_FEATURES
|
||||
from litellm.proxy.proxy_server import app, ensure_unique_openapi_operation_ids
|
||||
from litellm.proxy.proxy_server import app
|
||||
|
||||
for feat in LAZY_FEATURES:
|
||||
try:
|
||||
module = importlib.import_module(feat.module_path)
|
||||
feat.register_fn(app, module)
|
||||
except Exception as exc:
|
||||
sys.stderr.write(f"warning: skip {feat.name}: {exc}\n")
|
||||
|
||||
fragments: Final[dict[str, dict]] = {}
|
||||
skipped: Final = tuple(name for feat in LAZY_FEATURES if (name := _register_feature(app, feat)) is not None)
|
||||
used_operation_ids: Final[set[str]] = set()
|
||||
for feat in LAZY_FEATURES:
|
||||
feat_routes = [r for r in app.routes if feat.matches(getattr(r, "path", ""))]
|
||||
if not feat_routes:
|
||||
continue
|
||||
_stabilize_multi_method_route_ids(feat_routes)
|
||||
full = get_openapi(title=app.title, version=app.version, routes=feat_routes)
|
||||
paths = full.get("paths", {})
|
||||
_normalize_operation_ids(paths)
|
||||
# Group all of a feature's routes under one tag.
|
||||
for path_ops in full.get("paths", {}).values():
|
||||
for method, op in path_ops.items():
|
||||
if isinstance(op, dict):
|
||||
operation_id = op.get("operationId")
|
||||
if isinstance(operation_id, str):
|
||||
for suffix in HTTP_METHOD_SUFFIXES:
|
||||
if operation_id.endswith(f"_{suffix}"):
|
||||
op["operationId"] = operation_id[: -len(suffix)] + method
|
||||
break
|
||||
op["tags"] = [feat.name]
|
||||
full = ensure_unique_openapi_operation_ids(full, used_operation_ids)
|
||||
fragments[feat.name] = {
|
||||
"paths": paths,
|
||||
"components": {"schemas": full.get("components", {}).get("schemas", {})},
|
||||
}
|
||||
return fragments
|
||||
fragments: Final = {
|
||||
feat.name: fragment
|
||||
for feat in LAZY_FEATURES
|
||||
if (fragment := _feature_fragment(app, feat, used_operation_ids)) is not None
|
||||
}
|
||||
return SnapshotResult(fragments=fragments, skipped=skipped)
|
||||
|
||||
|
||||
def main(snapshot_file: Path = SNAPSHOT_FILE, generate: Callable[[], SnapshotResult] = generate_snapshot) -> int:
|
||||
result: Final = generate()
|
||||
if result.skipped:
|
||||
sys.stderr.write(
|
||||
f"error: {len(result.skipped)} feature(s) failed to import, so their fragments would vanish from the "
|
||||
f"snapshot: {', '.join(result.skipped)}\n"
|
||||
)
|
||||
return 1
|
||||
snapshot_file.write_text(json.dumps(result.fragments, indent=2, sort_keys=True) + "\n")
|
||||
sys.stdout.write(f"wrote {len(result.fragments)} feature fragments to {snapshot_file}\n")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
fragments: Final = generate_snapshot()
|
||||
SNAPSHOT_FILE.write_text(json.dumps(fragments, indent=2, sort_keys=True) + "\n")
|
||||
sys.stdout.write(f"wrote {len(fragments)} feature fragments to {SNAPSHOT_FILE}\n")
|
||||
sys.exit(main())
|
||||
|
|
|
|||
|
|
@ -195,7 +195,7 @@ def _models_this_test_can_call(config: RequestComplexityRouterConfig) -> tuple[s
|
|||
model
|
||||
for model in (
|
||||
config.classifier_llm_config.model
|
||||
if config.classifier_type == "llm" and config.classifier_llm_config is not None
|
||||
if config.uses_llm_classifier and config.classifier_llm_config is not None
|
||||
else None,
|
||||
config.embedding_model if config.semantic_keyword_matching else None,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ This is an enterprise feature and requires a premium license.
|
|||
|
||||
import re
|
||||
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
|
||||
from copy import deepcopy
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from itertools import chain
|
||||
|
|
@ -2375,6 +2376,37 @@ async def get_group(
|
|||
raise handle_exception_on_proxy(e)
|
||||
|
||||
|
||||
def _new_team_request_with_defaults(
|
||||
team_id: str,
|
||||
team_alias: str | None,
|
||||
members_with_roles: Sequence[Member],
|
||||
) -> NewTeamRequest:
|
||||
"""Build the SCIM group's team request, applying litellm.default_team_params
|
||||
(including models) the same way SSO auto-created teams do."""
|
||||
default_params: Final = litellm.default_team_params
|
||||
defaults: Final[Mapping[str, object]] = (
|
||||
deepcopy(default_params)
|
||||
if isinstance(default_params, dict)
|
||||
else default_params.model_dump(exclude_none=True)
|
||||
if default_params is not None
|
||||
else {}
|
||||
)
|
||||
default_metadata: Final = defaults.get("metadata")
|
||||
metadata: Final = {
|
||||
**(default_metadata if isinstance(default_metadata, dict) else {}),
|
||||
SCIM_MANAGED_TEAM_METADATA_KEY: True,
|
||||
}
|
||||
return NewTeamRequest.model_validate(
|
||||
{
|
||||
**defaults,
|
||||
"team_id": team_id,
|
||||
"team_alias": team_alias,
|
||||
"members_with_roles": members_with_roles,
|
||||
"metadata": metadata,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@scim_router.post(
|
||||
"/Groups",
|
||||
response_model=SCIMGroup,
|
||||
|
|
@ -2412,11 +2444,10 @@ async def create_group(
|
|||
|
||||
# Create team in database
|
||||
created_team: Final = await new_team(
|
||||
data=NewTeamRequest(
|
||||
data=_new_team_request_with_defaults(
|
||||
team_id=team_id,
|
||||
team_alias=group.displayName,
|
||||
members_with_roles=members_with_roles,
|
||||
metadata={SCIM_MANAGED_TEAM_METADATA_KEY: True},
|
||||
),
|
||||
http_request=Request(scope={"type": "http", "path": "/scim/v2/Groups"}),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
|
|
|
|||
|
|
@ -808,6 +808,15 @@ def normalize_email(email: str | None) -> str | None:
|
|||
return email.lower() if isinstance(email, str) else email
|
||||
|
||||
|
||||
# Ordered highest to lowest privilege
|
||||
LITELLM_USER_ROLE_HIERARCHY: Final = (
|
||||
LitellmUserRoles.PROXY_ADMIN,
|
||||
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
|
||||
LitellmUserRoles.INTERNAL_USER,
|
||||
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
|
||||
)
|
||||
|
||||
|
||||
def determine_role_from_groups(
|
||||
user_groups: list[str],
|
||||
role_mappings: "RoleMappings",
|
||||
|
|
@ -832,19 +841,11 @@ def determine_role_from_groups(
|
|||
# No role mappings configured, return default_role
|
||||
return role_mappings.default_role
|
||||
|
||||
# Role hierarchy (highest to lowest)
|
||||
role_hierarchy: Final = [
|
||||
LitellmUserRoles.PROXY_ADMIN,
|
||||
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
|
||||
LitellmUserRoles.INTERNAL_USER,
|
||||
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
|
||||
]
|
||||
|
||||
# Convert user_groups to a set for efficient lookup
|
||||
user_groups_set: Final = set(user_groups) if isinstance(user_groups, list) else set()
|
||||
|
||||
# Find the highest privilege role the user belongs to
|
||||
for role in role_hierarchy:
|
||||
for role in LITELLM_USER_ROLE_HIERARCHY:
|
||||
if role in role_mappings.roles:
|
||||
role_groups = role_mappings.roles[role]
|
||||
if isinstance(role_groups, list) and user_groups_set.intersection(set(role_groups)):
|
||||
|
|
@ -4236,15 +4237,7 @@ class MicrosoftSSOHandler:
|
|||
verbose_proxy_logger.debug("Extracted app roles from id_token: %s", app_roles)
|
||||
|
||||
# Combine groups and app roles
|
||||
user_role: LitellmUserRoles | None = None
|
||||
if app_roles:
|
||||
# Check if any app role is a valid LitellmUserRoles
|
||||
for role_str in app_roles:
|
||||
role = get_litellm_user_role(role_str)
|
||||
if role is not None:
|
||||
user_role = role
|
||||
verbose_proxy_logger.debug("Found valid LitellmUserRoles '%s' in app_roles", role.value)
|
||||
break
|
||||
user_role: Final = MicrosoftSSOHandler.get_user_role_from_app_roles(app_roles)
|
||||
|
||||
verbose_proxy_logger.debug("Combined team_ids (groups + app roles): %s", user_team_ids)
|
||||
|
||||
|
|
@ -4282,6 +4275,27 @@ class MicrosoftSSOHandler:
|
|||
verbose_proxy_logger.debug("Microsoft SSO OpenID Response: %s", openid_response)
|
||||
return openid_response
|
||||
|
||||
@staticmethod
|
||||
def get_user_role_from_app_roles(
|
||||
app_roles: Sequence[str] | None,
|
||||
) -> LitellmUserRoles | None:
|
||||
"""
|
||||
Resolve the one role LiteLLM stores for a user from their Entra app roles.
|
||||
|
||||
Entra does not guarantee `roles` claim ordering, so a user holding several app
|
||||
roles resolves to the highest privilege one rather than whichever the claim
|
||||
listed first. Roles the hierarchy does not rank (org_admin, team, customer)
|
||||
resolve by name to stay deterministic
|
||||
"""
|
||||
resolved: Final = frozenset(
|
||||
role for role in (get_litellm_user_role(role_str) for role_str in app_roles or ()) if role is not None
|
||||
)
|
||||
if not resolved:
|
||||
return None
|
||||
|
||||
ranked: Final = next((role for role in LITELLM_USER_ROLE_HIERARCHY if role in resolved), None)
|
||||
return ranked if ranked is not None else min(resolved, key=lambda role: role.value)
|
||||
|
||||
@staticmethod
|
||||
def get_app_roles_from_id_token(id_token: str | None) -> list[str]:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1021,19 +1021,7 @@ async def delete_prompt(
|
|||
# Delete versions from the database (scoped to environment if provided)
|
||||
await _prompt_table(prisma_client).delete_many(where=delete_where)
|
||||
|
||||
# Remove matching prompts from memory — scope to environment if provided
|
||||
if environment:
|
||||
prompts_to_delete: Final = [
|
||||
pid
|
||||
for pid, prompt in IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS.items()
|
||||
if get_base_prompt_id(prompt_id=pid) == base_prompt_id and prompt.environment == environment
|
||||
]
|
||||
for pid in prompts_to_delete:
|
||||
del IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS[pid]
|
||||
if pid in IN_MEMORY_PROMPT_REGISTRY.prompt_id_to_custom_prompt:
|
||||
del IN_MEMORY_PROMPT_REGISTRY.prompt_id_to_custom_prompt[pid]
|
||||
else:
|
||||
IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id(base_prompt_id)
|
||||
IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id(base_prompt_id, environment=environment or None)
|
||||
|
||||
env_msg: Final = f" from {environment}" if environment else ""
|
||||
return {"message": f"Prompt {base_prompt_id} deleted successfully{env_msg}"}
|
||||
|
|
|
|||
|
|
@ -195,12 +195,22 @@ class InMemoryPromptRegistry:
|
|||
"""
|
||||
return self.prompt_id_to_custom_prompt.get(prompt_id)
|
||||
|
||||
def delete_prompts_by_base_id(self, base_prompt_id: str) -> list[str]:
|
||||
def remove_prompt(self, prompt_id: str) -> None:
|
||||
import litellm
|
||||
|
||||
self.IN_MEMORY_PROMPTS.pop(prompt_id, None)
|
||||
stale_callback: Final = self.prompt_id_to_custom_prompt.pop(prompt_id, None)
|
||||
if stale_callback is not None:
|
||||
litellm.logging_callback_manager.remove_callback_from_all_lists(stale_callback)
|
||||
|
||||
def delete_prompts_by_base_id(self, base_prompt_id: str, environment: str | None = None) -> list[str]:
|
||||
"""
|
||||
Delete all prompts matching the given base prompt ID from memory.
|
||||
Delete all prompts matching the given base prompt ID from memory, along with their
|
||||
registered callbacks; scoped to one environment when given.
|
||||
|
||||
Args:
|
||||
base_prompt_id: The base prompt ID (without version suffix)
|
||||
environment: When set, only delete prompts deployed to this environment
|
||||
|
||||
Returns:
|
||||
List of prompt IDs that were deleted
|
||||
|
|
@ -208,13 +218,14 @@ class InMemoryPromptRegistry:
|
|||
from litellm.proxy.prompts.prompt_endpoints import get_base_prompt_id
|
||||
|
||||
prompts_to_delete: Final = [
|
||||
pid for pid in self.IN_MEMORY_PROMPTS if get_base_prompt_id(prompt_id=pid) == base_prompt_id
|
||||
pid
|
||||
for pid, prompt in self.IN_MEMORY_PROMPTS.items()
|
||||
if get_base_prompt_id(prompt_id=pid) == base_prompt_id
|
||||
and (environment is None or prompt.environment == environment)
|
||||
]
|
||||
|
||||
for pid in prompts_to_delete:
|
||||
del self.IN_MEMORY_PROMPTS[pid]
|
||||
if pid in self.prompt_id_to_custom_prompt:
|
||||
del self.prompt_id_to_custom_prompt[pid]
|
||||
self.remove_prompt(prompt_id=pid)
|
||||
|
||||
return prompts_to_delete
|
||||
|
||||
|
|
|
|||
|
|
@ -7268,6 +7268,7 @@ class ProxyConfig:
|
|||
return None
|
||||
|
||||
try:
|
||||
prompt_ids_loaded_before_db_read: Final = frozenset(IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS)
|
||||
prompts_in_db: Final[Sequence[object]] = await PromptRepository(prisma_client).table.find_many()
|
||||
parsed_specs: Final[tuple[PromptSpec, ...]] = tuple(
|
||||
spec for row in prompts_in_db if (spec := parse_row(row)) is not None
|
||||
|
|
@ -7290,6 +7291,18 @@ class ProxyConfig:
|
|||
prompt_spec.prompt_id,
|
||||
prompt_sync_error,
|
||||
)
|
||||
# An unparsable row still exists in the DB, so skip the sweep rather than unload its in-memory copy
|
||||
every_row_parsed: Final = len(parsed_specs) == len(prompts_in_db)
|
||||
if every_row_parsed:
|
||||
deleted_db_prompt_ids: Final = tuple(
|
||||
prompt_id
|
||||
for prompt_id in prompt_ids_loaded_before_db_read
|
||||
if (loaded_spec := IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS.get(prompt_id)) is not None
|
||||
and loaded_spec.prompt_info.prompt_type == "db"
|
||||
and prompt_id not in newest_spec_per_id
|
||||
)
|
||||
for deleted_prompt_id in deleted_db_prompt_ids:
|
||||
IN_MEMORY_PROMPT_REGISTRY.remove_prompt(prompt_id=deleted_prompt_id)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug("litellm.proxy.proxy_server.py::ProxyConfig:_init_prompts_in_db - %s", e)
|
||||
|
||||
|
|
|
|||
|
|
@ -18,11 +18,12 @@ from typing import (
|
|||
)
|
||||
|
||||
import fastapi
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, status
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, Response, status
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy._types import ProviderBudgetResponse, ProviderBudgetResponseObject
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
|
@ -54,6 +55,11 @@ router: Final = APIRouter()
|
|||
|
||||
SPEND_LOGS_PAGINATION_COUNT_CAP: Final = 10000
|
||||
|
||||
_INTERNAL_HEALTH_CHECK_API_KEYS: Final = (
|
||||
LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
|
||||
hash_token(token=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME),
|
||||
)
|
||||
|
||||
_RowT = TypeVar("_RowT")
|
||||
|
||||
|
||||
|
|
@ -208,9 +214,18 @@ async def _find_spend_logs(
|
|||
prisma_client: PrismaClient,
|
||||
where: Mapping[str, object],
|
||||
order: Mapping[str, str],
|
||||
take: int,
|
||||
http_response: Response,
|
||||
) -> Sequence[_SupportsModelDump]:
|
||||
"""Read spend log rows as Prisma model instances."""
|
||||
return await _spend_logs_table(prisma_client).find_many(where=where, order=order)
|
||||
"""Read spend log rows as Prisma model instances, capped at ``take`` rows."""
|
||||
rows: Final = await _spend_logs_table(prisma_client).find_many(where=where, order=order, take=take)
|
||||
if len(rows) == take:
|
||||
http_response.headers["x-litellm-spend-logs-truncated"] = "true"
|
||||
verbose_proxy_logger.warning(
|
||||
"/spend/logs result truncated to the %s most recent rows; use /spend/logs/v2 for paginated access",
|
||||
take,
|
||||
)
|
||||
return rows
|
||||
|
||||
|
||||
async def _find_spend_log_row(prisma_client: PrismaClient, request_id: str) -> _SpendLogOwnershipRow | None:
|
||||
|
|
@ -2239,6 +2254,10 @@ async def ui_view_spend_logs(
|
|||
status_filter: str | None = fastapi.Query(
|
||||
default=None, description="Filter logs by status (e.g., success, failure)"
|
||||
),
|
||||
cache_hit_filter: str | None = fastapi.Query(
|
||||
default=None,
|
||||
description="Filter logs by cache state: 'hit' or 'miss'. Miss includes legacy rows with a null/unknown cache state",
|
||||
),
|
||||
model: str | None = fastapi.Query(default=None, description="Filter logs by model"),
|
||||
model_id: str | None = fastapi.Query(
|
||||
default=None,
|
||||
|
|
@ -2259,6 +2278,10 @@ async def ui_view_spend_logs(
|
|||
default="desc",
|
||||
description="Sort order: asc or desc",
|
||||
),
|
||||
exclude_internal_health_checks: bool = fastapi.Query(
|
||||
default=False,
|
||||
description="Exclude LiteLLM internal health check requests from results",
|
||||
),
|
||||
):
|
||||
"""
|
||||
View spend logs with pagination support.
|
||||
|
|
@ -2311,6 +2334,13 @@ async def ui_view_spend_logs(
|
|||
param="sort_order",
|
||||
code=status.HTTP_400_BAD_REQUEST,
|
||||
)
|
||||
if isinstance(cache_hit_filter, str) and cache_hit_filter not in {"hit", "miss"}:
|
||||
raise ProxyException(
|
||||
message=f"Invalid cache_hit_filter: {cache_hit_filter}. Must be one of: hit, miss",
|
||||
type="bad_request",
|
||||
param="cache_hit_filter",
|
||||
code=status.HTTP_400_BAD_REQUEST,
|
||||
)
|
||||
|
||||
try:
|
||||
is_admin_view: Final = _is_admin_view_safe(user_api_key_dict=user_api_key_dict)
|
||||
|
|
@ -2551,6 +2581,16 @@ async def ui_view_spend_logs(
|
|||
sql_params.append(status_filter)
|
||||
p += 1
|
||||
|
||||
if cache_hit_filter == "hit":
|
||||
sql_conditions.append("LOWER(cache_hit) = 'true'")
|
||||
elif cache_hit_filter == "miss":
|
||||
sql_conditions.append("(cache_hit IS NULL OR LOWER(cache_hit) != 'true')")
|
||||
|
||||
if exclude_internal_health_checks:
|
||||
sql_conditions.append(f"api_key NOT IN (${p}, ${p + 1})")
|
||||
sql_params.extend(_INTERNAL_HEALTH_CHECK_API_KEYS)
|
||||
p += 2 # rebind-ok: advances the file's shared $N placeholder counter
|
||||
|
||||
# Spend range
|
||||
if min_spend is not None:
|
||||
sql_conditions.append(f"spend >= ${p}")
|
||||
|
|
@ -2851,6 +2891,7 @@ async def ui_view_request_response_for_request_id(
|
|||
},
|
||||
)
|
||||
async def view_spend_logs(
|
||||
fastapi_response: Response,
|
||||
api_key: str | None = fastapi.Query(
|
||||
default=None,
|
||||
description="Get spend logs based on api key",
|
||||
|
|
@ -2881,6 +2922,8 @@ async def view_spend_logs(
|
|||
[DEPRECATED] This endpoint is not paginated and can cause performance issues.
|
||||
Please use `/spend/logs/v2` instead for paginated access to spend logs.
|
||||
|
||||
Row results are capped at 10,000 most recent entries per response.
|
||||
|
||||
View all spend logs, if request_id is provided, only logs for that request_id will be returned
|
||||
|
||||
When start_date and end_date are provided:
|
||||
|
|
@ -2931,7 +2974,6 @@ async def view_spend_logs(
|
|||
raise Exception(
|
||||
"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
|
||||
)
|
||||
spend_logs = []
|
||||
if (
|
||||
start_date is not None
|
||||
and isinstance(start_date, str)
|
||||
|
|
@ -2970,6 +3012,8 @@ async def view_spend_logs(
|
|||
prisma_client,
|
||||
where=filter_query,
|
||||
order={"startTime": "desc"},
|
||||
take=SPEND_LOGS_PAGINATION_COUNT_CAP,
|
||||
http_response=fastapi_response,
|
||||
)
|
||||
return data
|
||||
|
||||
|
|
@ -3040,14 +3084,12 @@ async def view_spend_logs(
|
|||
if user_id is not None and isinstance(user_id, str):
|
||||
scoped_filter["user"] = user_id
|
||||
|
||||
if not scoped_filter:
|
||||
spend_logs = await prisma_client.get_data(table_name="spend", query_type="find_all")
|
||||
return spend_logs
|
||||
|
||||
data = await _find_spend_logs(
|
||||
prisma_client,
|
||||
where=scoped_filter,
|
||||
order={"startTime": "desc"},
|
||||
take=SPEND_LOGS_PAGINATION_COUNT_CAP,
|
||||
http_response=fastapi_response,
|
||||
)
|
||||
return data
|
||||
|
||||
|
|
|
|||
|
|
@ -178,6 +178,50 @@ response = litellm.completion(
|
|||
|
||||
## Special Behaviors
|
||||
|
||||
### Heuristic-first chaining
|
||||
|
||||
`classifier_type: heuristic_first` runs the local scorer on every request and only calls the LLM
|
||||
classifier for the ones the scorer could not place cheaply. It takes the same classifier settings as
|
||||
`classifier_type: llm`, plus `heuristic_first_max_tier`:
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: smart-router
|
||||
litellm_params:
|
||||
model: auto_router/complexity_router
|
||||
complexity_router_config:
|
||||
classifier_type: heuristic_first
|
||||
heuristic_first_max_tier: SIMPLE
|
||||
classifier_llm_config:
|
||||
model: gpt-4o-mini
|
||||
tiers:
|
||||
SIMPLE: gpt-4o-mini
|
||||
MEDIUM: gpt-4o
|
||||
COMPLEX: claude-sonnet-4
|
||||
REASONING: o1-preview
|
||||
```
|
||||
|
||||
A request short-circuits, meaning it routes on the scorer's own tier with no classifier call, when
|
||||
two things hold: the scorer landed at or below `heuristic_first_max_tier`, and it produced at least
|
||||
one signal. Everything else goes to the classifier, which then decides as it normally would.
|
||||
|
||||
The signal requirement is what keeps this from quietly routing everything to your cheapest model.
|
||||
A prompt where no dimension fires scores exactly 0.0, which is below `simple_medium`, so the score
|
||||
to tier mapping calls it SIMPLE by default rather than by evidence. Around half of general traffic
|
||||
scores that way. Those requests reach the classifier instead, which is the whole reason to configure
|
||||
one. Note the converse too: the score is not a confidence, and a prompt that fires a single weak
|
||||
signal and still lands under the boundary does short-circuit, so a lower threshold buys accuracy and
|
||||
a higher one buys savings.
|
||||
|
||||
`heuristic_first_max_tier` names a built-in tier and may not name the highest one, since that would
|
||||
short-circuit everything and leave the classifier unreachable. Operator-defined tier sets
|
||||
(`tier_definitions`) are not supported here, because the scorer only produces the built-in tiers.
|
||||
When the classifier call fails, the fallback works exactly as it does under `classifier_type: llm`,
|
||||
except that the heuristic outcome is the one already computed rather than a second scoring pass.
|
||||
|
||||
Spend logs record `routing_decision.cause` as `heuristic_first_short_circuit` when the classifier
|
||||
was skipped, and `llm_classifier` when it ran, so the two are told apart per request.
|
||||
|
||||
### Reasoning Override
|
||||
|
||||
If 2+ reasoning markers are detected in the user message, the request is promoted to the REASONING tier even when the weighted score maps lower, so complex reasoning tasks get the appropriate model. The promotion requires the score to reach `reasoning_override_min_score`, which tracks `tier_boundaries.simple_medium` unless set, so stock phrases on an otherwise trivial prompt cannot buy the top tier. Set it to `0` to promote on the markers alone.
|
||||
|
|
|
|||
|
|
@ -719,6 +719,7 @@ class ClassificationOutcome(NamedTuple):
|
|||
"heuristic_scorer",
|
||||
"reasoning_override",
|
||||
"llm_classifier",
|
||||
"heuristic_first_short_circuit",
|
||||
"classifier_plugin",
|
||||
"classifier_fallback",
|
||||
"default_model_fallback",
|
||||
|
|
@ -859,7 +860,7 @@ class ComplexityRouter(CustomLogger):
|
|||
|
||||
# Both are pure functions of the config, so building them per classifier call would
|
||||
# re-run create_model and the schema conversion on every request for the same result.
|
||||
llm_classifier_configured: Final = self.config.classifier_type == "llm" and (
|
||||
llm_classifier_configured: Final = self.config.uses_llm_classifier and (
|
||||
self.config.classifier_llm_config is not None
|
||||
)
|
||||
self._classifier_system_prompt: str | None = (
|
||||
|
|
@ -1237,17 +1238,63 @@ class ComplexityRouter(CustomLogger):
|
|||
"""
|
||||
Classify a prompt by complexity, using the LLM classifier when configured.
|
||||
|
||||
Falls back to the local heuristic scorer if classifier_type is "heuristic". If the LLM call
|
||||
or the classifier plugin fails, times out, or produces no usable tier, the configured
|
||||
fallback_tier wins on a custom tier set, and classifier_fallback otherwise decides between
|
||||
the heuristic scorer and default_model. The outcome's `cause` reports which path actually ran.
|
||||
Falls back to the local heuristic scorer if classifier_type is "heuristic". Under
|
||||
"heuristic_first" the scorer runs first and the classifier is called only for requests it
|
||||
could not place at or below heuristic_first_max_tier. If the LLM call or the classifier
|
||||
plugin fails, times out, or produces no usable tier, the configured fallback_tier wins on a
|
||||
custom tier set, and classifier_fallback otherwise decides between the heuristic scorer and
|
||||
default_model. The outcome's `cause` reports which path actually ran.
|
||||
"""
|
||||
if self.config.classifier_type == "custom":
|
||||
return await self._classify_with_plugin(prompt, system_prompt, request_kwargs, raw_messages)
|
||||
if self.config.classifier_type == "heuristic_first" and self.config.classifier_llm_config is not None:
|
||||
return await self._classify_heuristic_first(prompt, system_prompt, request_kwargs, messages)
|
||||
if self.config.classifier_type != "llm" or self.config.classifier_llm_config is None:
|
||||
tier, score, signals, cause = self._score_and_classify(prompt, system_prompt)
|
||||
return ClassificationOutcome(tier=tier, score=score, signals=signals, cause=cause)
|
||||
return await self._llm_classifier_outcome(prompt, system_prompt, request_kwargs, messages)
|
||||
|
||||
async def _classify_heuristic_first(
|
||||
self,
|
||||
prompt: str,
|
||||
system_prompt: str | None,
|
||||
request_kwargs: dict[str, Any] | None, # mutable-ok: handed to _classify_with_llm as-is
|
||||
messages: Sequence[Mapping[str, object]] | None,
|
||||
) -> ClassificationOutcome:
|
||||
"""Score locally, and only pay for the classifier call when the scorer did not confidently
|
||||
place the request at or below heuristic_first_max_tier.
|
||||
|
||||
Confidence is `signals`, not `score`. A prompt where no dimension fired scores exactly 0.0,
|
||||
which is below simple_medium and so lands SIMPLE by default rather than by evidence, and a
|
||||
threshold check alone would hand that traffic to the cheapest model without ever consulting
|
||||
the classifier. Scores also go negative when simple indicators fire, so a score threshold
|
||||
would reject exactly the trivial prompts this path exists to serve.
|
||||
"""
|
||||
tier, score, signals, cause = self._score_and_classify(prompt, system_prompt)
|
||||
scored: Final = ClassificationOutcome(tier=tier, score=score, signals=signals, cause=cause)
|
||||
threshold: Final = self.config.heuristic_first_max_tier
|
||||
decided_cheaply: Final = (
|
||||
threshold is not None
|
||||
and bool(signals)
|
||||
and self._active_tier_severity(tier) <= self._active_tier_severity(threshold)
|
||||
)
|
||||
if decided_cheaply:
|
||||
return ClassificationOutcome(tier=tier, score=score, signals=signals, cause="heuristic_first_short_circuit")
|
||||
return await self._llm_classifier_outcome(prompt, system_prompt, request_kwargs, messages, scored=scored)
|
||||
|
||||
async def _llm_classifier_outcome(
|
||||
self,
|
||||
prompt: str,
|
||||
system_prompt: str | None,
|
||||
request_kwargs: dict[str, Any] | None, # mutable-ok: handed to _classify_with_llm as-is
|
||||
messages: Sequence[Mapping[str, object]] | None,
|
||||
scored: ClassificationOutcome | None = None,
|
||||
) -> ClassificationOutcome:
|
||||
"""Call the LLM classifier and turn its verdict, or its failure, into an outcome.
|
||||
|
||||
`scored` is the heuristic outcome the caller already computed, which only "heuristic_first"
|
||||
has. It is handed to the failure path so a classifier error does not re-run the scorer.
|
||||
"""
|
||||
try:
|
||||
tier, classifier_cost = await self._classify_with_llm(prompt, system_prompt, request_kwargs, messages)
|
||||
return ClassificationOutcome(
|
||||
|
|
@ -1258,11 +1305,20 @@ class ComplexityRouter(CustomLogger):
|
|||
classifier_cost=classifier_cost,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 -- external LLM call can fail in many distinct ways (timeout, provider error, validation, parse error); any failure must fall back to the configured fallback path
|
||||
return self._classifier_failure_outcome(f"LLM classifier failed ({e})", prompt, system_prompt)
|
||||
return self._classifier_failure_outcome(f"LLM classifier failed ({e})", prompt, system_prompt, scored)
|
||||
|
||||
def _classifier_failure_outcome(self, reason: str, prompt: str, system_prompt: str | None) -> ClassificationOutcome:
|
||||
def _classifier_failure_outcome(
|
||||
self,
|
||||
reason: str,
|
||||
prompt: str,
|
||||
system_prompt: str | None,
|
||||
scored: ClassificationOutcome | None = None,
|
||||
) -> ClassificationOutcome:
|
||||
"""The outcome when the LLM classifier or classifier plugin produced no usable tier:
|
||||
fallback_tier on a custom tier set, classifier_fallback otherwise."""
|
||||
fallback_tier on a custom tier set, classifier_fallback otherwise.
|
||||
|
||||
A caller that already scored the prompt passes `scored` so the heuristic arm returns that
|
||||
verdict instead of running the same scan again on the request path."""
|
||||
fallback_tier: Final = self.config.fallback_tier
|
||||
if fallback_tier is not None:
|
||||
verbose_router_logger.warning("ComplexityRouter: %s, routing to fallback_tier %s", reason, fallback_tier)
|
||||
|
|
@ -1277,6 +1333,8 @@ class ComplexityRouter(CustomLogger):
|
|||
)
|
||||
if self.config.classifier_fallback == "default_model":
|
||||
return self._default_model_fallback_outcome()
|
||||
if scored is not None:
|
||||
return scored
|
||||
tier, score, signals, cause = self._score_and_classify(prompt, system_prompt)
|
||||
return ClassificationOutcome(tier=tier, score=score, signals=signals, cause=cause)
|
||||
|
||||
|
|
|
|||
|
|
@ -38,6 +38,11 @@ class ClassificationRubric(str, Enum):
|
|||
# routers get the calibrated rubric without changing what is already running.
|
||||
DEFAULT_CLASSIFICATION_RUBRIC: Final[ClassificationRubric] = ClassificationRubric.LEGACY
|
||||
|
||||
# The classifier_type values that can call classifier_llm_config.model. Every consumer asking
|
||||
# "is the classifier model a real dependency of this router" resolves it here, including the ones
|
||||
# that only hold the raw config mapping and cannot reach ComplexityRouterConfig.uses_llm_classifier.
|
||||
LLM_CLASSIFIER_TYPES: Final[frozenset[str]] = frozenset({"llm", "heuristic_first"})
|
||||
|
||||
|
||||
TIER_SEVERITY_ORDER: Final[tuple[ComplexityTier, ...]] = (
|
||||
ComplexityTier.SIMPLE,
|
||||
|
|
@ -591,13 +596,30 @@ class ComplexityRouterConfig(BaseModel):
|
|||
)
|
||||
|
||||
# Classifier strategy
|
||||
classifier_type: Literal["heuristic", "llm", "custom"] = Field(
|
||||
classifier_type: Literal["heuristic", "llm", "custom", "heuristic_first"] = Field(
|
||||
default="heuristic",
|
||||
description="Classification strategy: local regex/keyword scoring, an LLM call, or a custom classifier plugin",
|
||||
description=(
|
||||
"Classification strategy: local regex/keyword scoring, an LLM call, a custom classifier "
|
||||
"plugin, or 'heuristic_first', which scores locally and only pays for the LLM classifier "
|
||||
"when the local scorer does not confidently land a cheap tier"
|
||||
),
|
||||
)
|
||||
classifier_llm_config: ClassifierLLMConfig | None = Field(
|
||||
default=None,
|
||||
description="Configuration for the LLM classifier; required when classifier_type is 'llm'",
|
||||
description="Configuration for the LLM classifier; required when classifier_type is 'llm' or 'heuristic_first'",
|
||||
)
|
||||
heuristic_first_max_tier: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"The highest tier the local scorer may decide on its own; required when classifier_type is "
|
||||
"'heuristic_first' and rejected otherwise. A request whose heuristic tier is at or below this "
|
||||
"one skips the LLM classifier and routes straight to that heuristic tier, so the classifier "
|
||||
"call is only paid for on traffic the scorer could not place cheaply. The scorer must also "
|
||||
"have produced at least one signal: a prompt where no dimension fired scores 0.0 and would "
|
||||
"otherwise land SIMPLE by default rather than by evidence, which is how a chained router "
|
||||
"would silently send unclassified traffic to the cheapest model. Names a built-in tier, and "
|
||||
"may not name the highest one, since that would make the LLM classifier unreachable."
|
||||
),
|
||||
)
|
||||
classifier_plugin: ClassifierPlugin | None = Field(
|
||||
default=None,
|
||||
|
|
@ -626,7 +648,7 @@ class ComplexityRouterConfig(BaseModel):
|
|||
"which is what a classifier on some other taxonomy wants: a prompt that grades data "
|
||||
"sensitivity has no use for a complexity score, and scoring one produces a tier unrelated to "
|
||||
"what the operator configured. Requires default_model when set to 'default_model'. Only "
|
||||
"applies when classifier_type is 'llm' or 'custom'."
|
||||
"applies when classifier_type is 'llm', 'custom', or 'heuristic_first'."
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -936,8 +958,8 @@ class ComplexityRouterConfig(BaseModel):
|
|||
|
||||
@model_validator(mode="after")
|
||||
def _validate_classifier_config(self) -> "ComplexityRouterConfig":
|
||||
if self.classifier_type == "llm" and self.classifier_llm_config is None:
|
||||
raise ValueError("classifier_llm_config is required when classifier_type is 'llm'")
|
||||
if self.uses_llm_classifier and self.classifier_llm_config is None:
|
||||
raise ValueError(f"classifier_llm_config is required when classifier_type is {self.classifier_type!r}")
|
||||
if self.classifier_type == "custom" and self.classifier_plugin is None:
|
||||
raise ValueError("classifier_plugin is required when classifier_type is 'custom'")
|
||||
if self.classifier_plugin is not None and self.classifier_type != "custom":
|
||||
|
|
@ -947,6 +969,49 @@ class ComplexityRouterConfig(BaseModel):
|
|||
)
|
||||
return self
|
||||
|
||||
@field_validator("heuristic_first_max_tier", mode="before")
|
||||
@classmethod
|
||||
def _coerce_heuristic_first_max_tier(cls, value: object) -> object:
|
||||
if isinstance(value, ComplexityTier):
|
||||
return value.value
|
||||
if isinstance(value, str):
|
||||
return value.strip()
|
||||
return value
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_heuristic_first_max_tier(self) -> "ComplexityRouterConfig":
|
||||
if self.classifier_type != "heuristic_first":
|
||||
if self.heuristic_first_max_tier is not None:
|
||||
raise ValueError(
|
||||
f"heuristic_first_max_tier is set but classifier_type is {self.classifier_type!r}; "
|
||||
"the local scorer would never gate the classifier. Set classifier_type "
|
||||
"'heuristic_first' or remove heuristic_first_max_tier"
|
||||
)
|
||||
return self
|
||||
threshold: Final = self.heuristic_first_max_tier
|
||||
if threshold is None:
|
||||
raise ValueError(
|
||||
"heuristic_first_max_tier is required when classifier_type is 'heuristic_first': without a "
|
||||
"threshold there is nothing to decide whether a request escalates to the LLM classifier"
|
||||
)
|
||||
names: Final = self.tier_names()
|
||||
if threshold not in names:
|
||||
raise ValueError(
|
||||
f"heuristic_first_max_tier {threshold!r} is not an active tier: it must name one of {', '.join(names)}"
|
||||
)
|
||||
if threshold == names[-1]:
|
||||
raise ValueError(
|
||||
f"heuristic_first_max_tier {threshold} is the highest tier, so every request would short-circuit "
|
||||
"and the LLM classifier would never run; name a lower tier or use classifier_type 'heuristic'"
|
||||
)
|
||||
if threshold not in self.tiers:
|
||||
raise ValueError(
|
||||
f"heuristic_first_max_tier {threshold} has no model configured in tiers; a threshold pointing at "
|
||||
"an unconfigured tier would route short-circuited requests to the default fallback instead of the "
|
||||
"pool the operator intended"
|
||||
)
|
||||
return self
|
||||
|
||||
@field_validator("fallback_tier", "classification_prompt")
|
||||
@classmethod
|
||||
def _reject_blank_optional_text(cls, value: str | None) -> str | None:
|
||||
|
|
@ -969,6 +1034,14 @@ class ComplexityRouterConfig(BaseModel):
|
|||
"""True when the operator replaced the built-in tier set via tier_definitions."""
|
||||
return self.tier_definitions is not None
|
||||
|
||||
@property
|
||||
def uses_llm_classifier(self) -> bool:
|
||||
"""True when this router can call classifier_llm_config.model, so the model is a real
|
||||
dependency: authorized against the caller's key, counted in the health graph, and given a
|
||||
prebuilt rubric. 'heuristic_first' only calls it for traffic the local scorer escalates,
|
||||
which still makes it a dependency on every one of those requests."""
|
||||
return self.classifier_type in LLM_CLASSIFIER_TYPES
|
||||
|
||||
def tier_names(self) -> tuple[str, ...]:
|
||||
"""The active tier names: the defined names, or the built-in set in severity order."""
|
||||
if self.tier_definitions is not None:
|
||||
|
|
@ -1063,7 +1136,7 @@ class ComplexityRouterConfig(BaseModel):
|
|||
)
|
||||
if duplicated:
|
||||
raise ValueError(f"tier_definitions names must be unique (case-insensitive): {', '.join(duplicated)}")
|
||||
if self.classifier_type == "heuristic":
|
||||
if self.classifier_type in ("heuristic", "heuristic_first"):
|
||||
raise ValueError(
|
||||
"tier_definitions requires classifier_type 'llm' or 'custom': the heuristic scorer only "
|
||||
"produces the built-in tiers"
|
||||
|
|
|
|||
|
|
@ -15,6 +15,8 @@ from dataclasses import dataclass
|
|||
from types import MappingProxyType
|
||||
from typing import Final, Literal, TypeAlias
|
||||
|
||||
from litellm.router_strategy.complexity_router.config import LLM_CLASSIFIER_TYPES
|
||||
|
||||
AUTO_ROUTER_MODEL_PREFIX: Final = "auto_router/"
|
||||
|
||||
StrategyRouterKind = Literal["semantic", "complexity", "adaptive", "quality"]
|
||||
|
|
@ -144,7 +146,11 @@ def strategy_router_dependencies(
|
|||
dict.fromkeys(
|
||||
tuple(dep for tier in _mapping(complexity.get("tiers")).values() for dep in _pool(tier, "tier"))
|
||||
+ _named(litellm_params.get("complexity_router_default_model"), "default")
|
||||
+ (_named(classifier.get("model"), "classifier") if complexity.get("classifier_type") == "llm" else ())
|
||||
+ (
|
||||
_named(classifier.get("model"), "classifier")
|
||||
if complexity.get("classifier_type") in LLM_CLASSIFIER_TYPES
|
||||
else ()
|
||||
)
|
||||
+ (
|
||||
_named(complexity.get("embedding_model"), "embedding")
|
||||
if complexity.get("semantic_keyword_matching")
|
||||
|
|
|
|||
|
|
@ -1,3 +1,10 @@
|
|||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal
|
||||
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm.types.integrations.custom_logger import StandardCustomLoggerInitParams
|
||||
|
||||
|
||||
|
|
@ -5,3 +12,110 @@ class NewRelicInitParams(StandardCustomLoggerInitParams):
|
|||
"""
|
||||
Params for initializing a New Relic logger on litellm
|
||||
"""
|
||||
|
||||
|
||||
#: Region -> Metric API endpoint. A fixed table by design: team config picks a
|
||||
#: region enum rather than a free-form endpoint, so callback vars can never
|
||||
#: redirect metrics to an arbitrary host.
|
||||
NEWRELIC_METRIC_ENDPOINT_BY_REGION: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{
|
||||
"us": "https://metric-api.newrelic.com/metric/v1",
|
||||
"eu": "https://metric-api.eu.newrelic.com/metric/v1",
|
||||
}
|
||||
)
|
||||
|
||||
NEWRELIC_DEFAULT_REGION: Final = "us"
|
||||
|
||||
#: Metric API caps a payload at 2000 data points / 1MB compressed; each queued
|
||||
#: record expands to at most 6 metrics, so cap the per-flush record count well
|
||||
#: below that.
|
||||
NEWRELIC_METRICS_MAX_BATCH_SIZE: Final = 250
|
||||
|
||||
#: Hard cap on records retained across failed flushes (5xx/network requeue).
|
||||
#: Beyond this the oldest records are dropped.
|
||||
NEWRELIC_METRICS_MAX_RETRY_QUEUE_SIZE: Final = 10_000
|
||||
# Outer passes over a stopped logger's queue: each pass retries the whole
|
||||
# queue, so records that arrive mid-drain still get attempts before the bounded
|
||||
# terminal drop. Serialized by a per-logger drain lock, so this bounds work.
|
||||
NEWRELIC_METRICS_MAX_DRAIN_PASSES: Final = 3
|
||||
# Metric API caps attribute values; 255 keeps caller-controlled model strings
|
||||
# from inflating the shared batch payload into a 413
|
||||
NEWRELIC_METRIC_ATTRIBUTE_MAX_LEN: Final = 255
|
||||
|
||||
NEWRELIC_METRIC_REQUESTS: Final = "litellm.requests"
|
||||
NEWRELIC_METRIC_COST_USD: Final = "litellm.cost.usd"
|
||||
NEWRELIC_METRIC_PROMPT_TOKENS: Final = "litellm.tokens.prompt"
|
||||
NEWRELIC_METRIC_COMPLETION_TOKENS: Final = "litellm.tokens.completion"
|
||||
NEWRELIC_METRIC_TOTAL_TOKENS: Final = "litellm.tokens.total"
|
||||
NEWRELIC_METRIC_REQUEST_DURATION_MS: Final = "litellm.request.duration_ms"
|
||||
|
||||
|
||||
class NewRelicSummaryValue(TypedDict):
|
||||
"""Value shape of a Metric API ``summary`` data point."""
|
||||
|
||||
count: ReadOnly[int]
|
||||
sum: ReadOnly[float]
|
||||
min: ReadOnly[float]
|
||||
max: ReadOnly[float]
|
||||
|
||||
|
||||
class NewRelicCountMetric(TypedDict):
|
||||
name: ReadOnly[str]
|
||||
type: ReadOnly[Literal["count"]]
|
||||
value: ReadOnly[float]
|
||||
attributes: ReadOnly[Mapping[str, str]]
|
||||
|
||||
|
||||
class NewRelicSummaryMetric(TypedDict):
|
||||
name: ReadOnly[str]
|
||||
type: ReadOnly[Literal["summary"]]
|
||||
value: ReadOnly[NewRelicSummaryValue]
|
||||
attributes: ReadOnly[Mapping[str, str]]
|
||||
|
||||
|
||||
NewRelicMetric = NewRelicCountMetric | NewRelicSummaryMetric
|
||||
|
||||
|
||||
#: ``interval.ms`` has a dot in it, so the functional TypedDict form is required.
|
||||
NewRelicMetricCommon = TypedDict(
|
||||
"NewRelicMetricCommon",
|
||||
{ # mutable-ok: functional TypedDict requires a dict-literal fields argument ("interval.ms" key)
|
||||
"timestamp": ReadOnly[int],
|
||||
"interval.ms": ReadOnly[int],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
class NewRelicMetricEnvelope(TypedDict):
|
||||
"""One element of the Metric API request body (``[{common, metrics}]``)."""
|
||||
|
||||
common: ReadOnly[NewRelicMetricCommon]
|
||||
metrics: ReadOnly[Sequence[NewRelicMetric]]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class NewRelicMetricRecord:
|
||||
"""One request's contribution to the per-flush aggregation."""
|
||||
|
||||
team_id: str
|
||||
team_alias: str
|
||||
model_group: str
|
||||
model: str
|
||||
custom_llm_provider: str
|
||||
status: str
|
||||
response_cost: float
|
||||
prompt_tokens: int
|
||||
completion_tokens: int
|
||||
total_tokens: int
|
||||
duration_ms: float
|
||||
|
||||
@property
|
||||
def bucket_key(self) -> tuple[str, str, str, str, str, str]:
|
||||
return (
|
||||
self.team_id,
|
||||
self.team_alias,
|
||||
self.model_group,
|
||||
self.model,
|
||||
self.custom_llm_provider,
|
||||
self.status,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -502,11 +502,16 @@ class MessageDelta(TypedDict, total=False):
|
|||
stop_reason: str | None
|
||||
|
||||
|
||||
class ServerToolUsage(TypedDict, total=False):
|
||||
web_search_requests: ReadOnly[int]
|
||||
|
||||
|
||||
class UsageDelta(TypedDict, total=False):
|
||||
input_tokens: int
|
||||
output_tokens: int
|
||||
cache_creation_input_tokens: int
|
||||
cache_read_input_tokens: int
|
||||
server_tool_use: ReadOnly[ServerToolUsage]
|
||||
|
||||
|
||||
class AppliedEdit(TypedDict, total=False):
|
||||
|
|
|
|||
|
|
@ -1,11 +1,12 @@
|
|||
from typing import Any, Literal, TypeAlias
|
||||
|
||||
from typing_extensions import NotRequired, TypedDict
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
from litellm.types.llms.anthropic import (
|
||||
AnthropicResponseContentBlockText,
|
||||
AnthropicResponseContentBlockToolUse,
|
||||
ContextManagementResponse,
|
||||
ServerToolUsage,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -71,6 +72,11 @@ class AnthropicUsage(TypedDict, total=False):
|
|||
cache_creation_input_tokens: int
|
||||
cache_read_input_tokens: int
|
||||
|
||||
"""
|
||||
Server-side tool usage (e.g. web search request counts)
|
||||
"""
|
||||
server_tool_use: NotRequired[ReadOnly[ServerToolUsage]]
|
||||
|
||||
|
||||
class AnthropicMessagesResponse(TypedDict, total=False):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -2808,6 +2808,12 @@ RoutingDecisionCause = Literal[
|
|||
# meant anything that filtered `signals` silently changed what the row claimed.
|
||||
"reasoning_override",
|
||||
"llm_classifier",
|
||||
# classifier_type 'heuristic_first': the local scorer produced at least one signal and landed at
|
||||
# or below heuristic_first_max_tier, so it decided the tier and the LLM classifier was never
|
||||
# called. Distinct from "heuristic_scorer", which is a router whose only classifier IS the
|
||||
# scorer, and from "classifier_fallback", which is the scorer running because a call failed:
|
||||
# only this cause means an LLM classifier was configured, reachable, and deliberately skipped.
|
||||
"heuristic_first_short_circuit",
|
||||
# The operator's classifier plugin (classifier_type 'custom') decided the tier.
|
||||
"classifier_plugin",
|
||||
# The LLM classifier or classifier plugin failed on a router with an operator-defined
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@
|
|||
# - tests/e2e Python -> `make lint-e2e-basedpyright` (test-linting.yml's e2e type-check step)
|
||||
# + raw HTTP client ban (test-code-quality.yml's check_e2e_no_raw_requests)
|
||||
# - dashboard -> prettier + eslint + lint budgets (test-litellm-ui-build.yml's frontend-lint)
|
||||
# - proxy/types -> regenerate dashboard API types and fail on drift (check-ui-api-types.yml)
|
||||
# - proxy/types -> regenerate the lazy OpenAPI snapshot and dashboard API types, fail on drift (check-ui-api-types.yml)
|
||||
#
|
||||
# Each block is skipped when no matching files are in scope, so unrelated commits
|
||||
# stay fast. This is intentionally not auto-installed as a git hook (see
|
||||
|
|
@ -244,7 +244,7 @@ fi
|
|||
|
||||
genapi_checks() {
|
||||
local status=0
|
||||
echo "check: checking dashboard API types are in sync (npm run gen:api)"
|
||||
echo "check: checking the lazy OpenAPI snapshot and dashboard API types are in sync (npm run gen:api)"
|
||||
# gen-api-types.mjs imports litellm.proxy.proxy_server, which needs the proxy deps
|
||||
# and an up-to-date Prisma client; check-ui-api-types.yml installs those and runs
|
||||
# prisma generate before gen:api, so mirror that here or a stale client can mask
|
||||
|
|
@ -260,7 +260,14 @@ genapi_checks() {
|
|||
elif ! uv run --no-sync python scripts/prisma_generate_if_needed.py; then
|
||||
echo "✗ Could not regenerate Prisma client (prisma generate failed)." >&2
|
||||
status=1
|
||||
elif ! uv run --no-sync python -m litellm.proxy._lazy_openapi_snapshot; then
|
||||
echo "✗ Could not regenerate the lazy OpenAPI snapshot (python -m litellm.proxy._lazy_openapi_snapshot failed)." >&2
|
||||
status=1
|
||||
elif ( cd ui/litellm-dashboard && LITELLM_PYTHON="uv run --no-sync python" npm run gen:api ); then
|
||||
if ! git diff --quiet -- litellm/proxy/_lazy_openapi_snapshot.json; then
|
||||
echo "✗ The lazy OpenAPI snapshot is stale; regenerated litellm/proxy/_lazy_openapi_snapshot.json. Stage it and commit; re-run make check only if other checks failed too." >&2
|
||||
status=1
|
||||
fi
|
||||
if ! git diff --quiet -- ui/litellm-dashboard/src/lib/http/schema.d.ts; then
|
||||
echo "✗ Dashboard API types are stale; regenerated src/lib/http/schema.d.ts. Stage it and commit; re-run make check only if other checks failed too." >&2
|
||||
status=1
|
||||
|
|
|
|||
|
|
@ -42,7 +42,7 @@ at call time. The provider table below is the source of truth; edit `PROVIDERS`
|
|||
| openai | `openai-realtime` | `openai/gpt-realtime-2` |
|
||||
| azure | `azure-realtime` | `azure/gpt-realtime-2` (GA protocol) |
|
||||
| gemini | `gemini-realtime` | `gemini/gemini-3.1-flash-live-preview` |
|
||||
| vertex_ai | `vertex-realtime` | `vertex_ai/gemini-live-2.5-flash-preview-native-audio-09-2025` |
|
||||
| vertex_ai | `vertex-realtime` | `vertex_ai/gemini-live-2.5-flash-native-audio` |
|
||||
|
||||
Bedrock and xai (`xai/grok-4-1-fast-non-reasoning`) are supported by the proxy but
|
||||
kept commented out in `PROVIDERS` until they pass end-to-end here; re-enable them by
|
||||
|
|
|
|||
|
|
@ -78,7 +78,7 @@ PROVIDERS = (
|
|||
"vertex_ai",
|
||||
"vertex-realtime",
|
||||
LiteLLMParamsBody(
|
||||
model="vertex_ai/gemini-live-2.5-flash-preview-native-audio-09-2025",
|
||||
model="vertex_ai/gemini-live-2.5-flash-native-audio",
|
||||
vertex_location="us-central1",
|
||||
vertex_credentials="os.environ/VERTEXAI_CREDENTIALS",
|
||||
),
|
||||
|
|
|
|||
|
|
@ -308,7 +308,8 @@ class TestGeminiChatCompletions:
|
|||
content=f"Reply with the single word pong. marker={tag}",
|
||||
)
|
||||
],
|
||||
max_tokens=32,
|
||||
max_tokens=64,
|
||||
reasoning_effort="none",
|
||||
),
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -50,11 +50,14 @@ CACHE_WARM_CONSECUTIVE_READS = 3
|
|||
|
||||
|
||||
def _cacheable_system_block(marker: str) -> TextBlock:
|
||||
"""A system prompt comfortably above the 4096-token minimum cacheable size
|
||||
of Haiku 4.5 (the smallest model here), unique per run so no other run's
|
||||
cache entry can satisfy the read."""
|
||||
text = " ".join(
|
||||
f"Reference paragraph {index} for run {marker}." for index in range(300)
|
||||
"""A system prompt at roughly twice the 4096-token minimum cacheable size of
|
||||
Haiku 4.5 (the smallest model here), unique per run so no other run's cache
|
||||
entry can satisfy the read. The marker appears once instead of in every
|
||||
paragraph: repeating it swung the block's size by ~1800 tokens with the
|
||||
marker's own tokenization and left it under the minimum on ~15% of runs, so
|
||||
the system breakpoint went uncached and the priming loop never saw a read."""
|
||||
text = f"Run {marker}.\n" + " ".join(
|
||||
f"Reference paragraph {index}." for index in range(1500)
|
||||
)
|
||||
return TextBlock(text=text, cache_control=CacheControl())
|
||||
|
||||
|
|
@ -101,8 +104,8 @@ def _first_turn_user_text(marker: str) -> str:
|
|||
"""A first user turn heavy enough (hundreds of tokens) that losing its cache
|
||||
entry is unambiguous in the usage numbers, unique per attempt so priming
|
||||
retries never depend on the proxy's response cache behavior."""
|
||||
notes = " ".join(f"Session note {index} for attempt {marker}." for index in range(100))
|
||||
return f"Reply with one word.\n{notes}"
|
||||
notes = " ".join(f"Session note {index}." for index in range(100))
|
||||
return f"Reply with one word. Attempt {marker}.\n{notes}"
|
||||
|
||||
|
||||
class PrimedCache(BaseModel):
|
||||
|
|
|
|||
|
|
@ -70,9 +70,15 @@ def _vertex_params(model: str, location: str) -> LiteLLMParamsBody:
|
|||
|
||||
|
||||
def _cacheable_system_block(marker: str) -> TextBlock:
|
||||
"""A system prompt comfortably above the 1024-token minimum cacheable size,
|
||||
unique per run so no other run's cache entry can satisfy the read."""
|
||||
text = " ".join(f"Reference paragraph {index} for run {marker}." for index in range(300))
|
||||
"""A system prompt at roughly twice the 4096-token minimum cacheable size of
|
||||
Haiku 4.5 (the smallest model here), unique per run so no other run's cache
|
||||
entry can satisfy the read. The marker appears once instead of in every
|
||||
paragraph: repeating it swung the block's size by ~1800 tokens with the
|
||||
marker's own tokenization and left it under the minimum on ~15% of runs, so
|
||||
the system breakpoint went uncached and the priming loop never saw a read."""
|
||||
text = f"Run {marker}.\n" + " ".join(
|
||||
f"Reference paragraph {index}." for index in range(1500)
|
||||
)
|
||||
return TextBlock(text=text, cache_control=CacheControl())
|
||||
|
||||
|
||||
|
|
@ -110,8 +116,8 @@ def _first_turn_user_text(marker: str) -> str:
|
|||
"""A first user turn heavy enough (hundreds of tokens) that losing its cache
|
||||
entry is unambiguous in the usage numbers, unique per attempt so priming
|
||||
retries never depend on the proxy's response cache behavior."""
|
||||
notes = " ".join(f"Session note {index} for attempt {marker}." for index in range(100))
|
||||
return f"Reply with one word.\n{notes}"
|
||||
notes = " ".join(f"Session note {index}." for index in range(100))
|
||||
return f"Reply with one word. Attempt {marker}.\n{notes}"
|
||||
|
||||
|
||||
class PrimedCache(BaseModel):
|
||||
|
|
|
|||
|
|
@ -163,12 +163,6 @@ class TestBudgetManagement:
|
|||
f"/budget/list never included the created budget {budget_id}",
|
||||
)
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason=(
|
||||
"stage red: product gap, /budget/update 500s on any model_max_budget "
|
||||
"(prisma Json arg + unquoted GraphQL interpolation)"
|
||||
)
|
||||
)
|
||||
@pytest.mark.covers("mgmt.budget.update.accepts_model_max_budget")
|
||||
def test_update_accepts_per_model_budgets_including_punctuated_names(
|
||||
self, client: ManagementClient, resources: ResourceManager
|
||||
|
|
|
|||
|
|
@ -2661,7 +2661,7 @@ async def test_run_direct_health_check_drops_only_the_rejected_kwarg(monkeypatch
|
|||
rejected argument alongside working ones would probe deployments the operator opted out."""
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
|
||||
seen: list = []
|
||||
seen: list[tuple[dict[str, str] | None, bool]] = []
|
||||
|
||||
async def fake_perform_health_check(
|
||||
model_list,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,825 @@
|
|||
"""
|
||||
Batching tests for NewRelicMetricsLogger: flush-window interval computation,
|
||||
dimension-bucket aggregation, the 4xx-drop vs 5xx/network-requeue policy, the
|
||||
retry-queue cap, and the stop flag that ends the periodic flush loop.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import gzip
|
||||
import json
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from httpx import HTTPStatusError, Request, Response
|
||||
|
||||
from litellm.integrations.newrelic.newrelic_metrics import (
|
||||
NewRelicMetricsLogger,
|
||||
_bucket_metrics,
|
||||
build_metric_payload,
|
||||
)
|
||||
from litellm.types.integrations.newrelic import (
|
||||
NEWRELIC_METRIC_COMPLETION_TOKENS,
|
||||
NEWRELIC_METRIC_COST_USD,
|
||||
NEWRELIC_METRIC_ENDPOINT_BY_REGION,
|
||||
NEWRELIC_METRIC_PROMPT_TOKENS,
|
||||
NEWRELIC_METRIC_REQUEST_DURATION_MS,
|
||||
NEWRELIC_METRIC_REQUESTS,
|
||||
NEWRELIC_METRIC_TOTAL_TOKENS,
|
||||
NewRelicMetricRecord,
|
||||
)
|
||||
|
||||
|
||||
def _record(
|
||||
team_id="team-a",
|
||||
team_alias=None,
|
||||
model="gpt-4o",
|
||||
model_group=None,
|
||||
status="success",
|
||||
response_cost=0.5,
|
||||
prompt_tokens=10,
|
||||
completion_tokens=20,
|
||||
total_tokens=30,
|
||||
duration_ms=100.0,
|
||||
) -> NewRelicMetricRecord:
|
||||
return NewRelicMetricRecord(
|
||||
team_id=team_id,
|
||||
team_alias=team_alias if team_alias is not None else f"{team_id}-alias",
|
||||
model_group=model_group if model_group is not None else f"{model}-group",
|
||||
model=model,
|
||||
custom_llm_provider="openai",
|
||||
status=status,
|
||||
response_cost=response_cost,
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
total_tokens=total_tokens,
|
||||
duration_ms=duration_ms,
|
||||
)
|
||||
|
||||
|
||||
def _standard_logging_object(team_id="team-a", response_cost=0.25) -> dict:
|
||||
return {
|
||||
"metadata": {"user_api_key_team_id": team_id, "user_api_key_team_alias": f"{team_id}-alias"},
|
||||
"model_group": "gpt-4o-group",
|
||||
"model": "gpt-4o",
|
||||
"custom_llm_provider": "openai",
|
||||
"status": "success",
|
||||
"response_cost": response_cost,
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 20,
|
||||
"total_tokens": 30,
|
||||
"response_time": 0.1,
|
||||
}
|
||||
|
||||
|
||||
def _make_logger(**kwargs) -> NewRelicMetricsLogger:
|
||||
with patch("asyncio.create_task"):
|
||||
return NewRelicMetricsLogger(newrelic_api_key="test-key", **kwargs)
|
||||
|
||||
|
||||
def _response(status_code: int, text: str = "") -> Response:
|
||||
return Response(status_code, request=Request("POST", "https://example.com"), text=text)
|
||||
|
||||
|
||||
def _raises(status_code: int):
|
||||
"""Mock the way AsyncHTTPHandler.post really behaves: raise_for_status() turns
|
||||
every non-2xx into an HTTPStatusError rather than returning the response."""
|
||||
resp = _response(status_code)
|
||||
return AsyncMock(side_effect=HTTPStatusError("err", request=resp.request, response=resp))
|
||||
|
||||
|
||||
def _metrics_by_name(payload, name):
|
||||
return [m for m in payload[0]["metrics"] if m["name"] == name]
|
||||
|
||||
|
||||
class TestBuildMetricPayload:
|
||||
def test_interval_and_timestamp_reflect_flush_window(self):
|
||||
payload = build_metric_payload((_record(),), window_start=1_000.0, now=1_007.5)
|
||||
|
||||
assert payload[0]["common"]["timestamp"] == 1_000_000
|
||||
assert payload[0]["common"]["interval.ms"] == 7_500
|
||||
|
||||
def test_interval_is_at_least_one_ms(self):
|
||||
payload = build_metric_payload((_record(),), window_start=1_000.0, now=1_000.0)
|
||||
|
||||
assert payload[0]["common"]["interval.ms"] == 1
|
||||
|
||||
def test_single_record_metric_values(self):
|
||||
payload = build_metric_payload(
|
||||
(_record(response_cost=0.5, prompt_tokens=10, completion_tokens=20, total_tokens=30, duration_ms=100.0),),
|
||||
window_start=1_000.0,
|
||||
now=1_005.0,
|
||||
)
|
||||
|
||||
by_name = {m["name"]: m for m in payload[0]["metrics"]}
|
||||
assert by_name[NEWRELIC_METRIC_REQUESTS]["value"] == 1.0
|
||||
assert by_name[NEWRELIC_METRIC_REQUESTS]["type"] == "count"
|
||||
assert by_name[NEWRELIC_METRIC_COST_USD]["value"] == 0.5
|
||||
assert by_name[NEWRELIC_METRIC_PROMPT_TOKENS]["value"] == 10.0
|
||||
assert by_name[NEWRELIC_METRIC_COMPLETION_TOKENS]["value"] == 20.0
|
||||
assert by_name[NEWRELIC_METRIC_TOTAL_TOKENS]["value"] == 30.0
|
||||
duration = by_name[NEWRELIC_METRIC_REQUEST_DURATION_MS]
|
||||
assert duration["type"] == "summary"
|
||||
assert duration["value"] == {"count": 1, "sum": 100.0, "min": 100.0, "max": 100.0}
|
||||
assert by_name[NEWRELIC_METRIC_REQUESTS]["attributes"] == {
|
||||
"team_id": "team-a",
|
||||
"team_alias": "team-a-alias",
|
||||
"model_group": "gpt-4o-group",
|
||||
"model": "gpt-4o",
|
||||
"custom_llm_provider": "openai",
|
||||
"status": "success",
|
||||
}
|
||||
|
||||
def test_aggregates_across_dimension_buckets(self):
|
||||
"""Two teams x two models in one queue land in the right bucket sums.
|
||||
|
||||
team_alias and model_group are held constant so bucketing provably keys on
|
||||
team_id and model themselves, not on correlated fields.
|
||||
"""
|
||||
shared = {"team_alias": "shared-alias", "model_group": "shared-group"}
|
||||
records = (
|
||||
_record(team_id="team-a", model="gpt-4o", response_cost=0.1, total_tokens=10, duration_ms=50.0, **shared),
|
||||
_record(team_id="team-a", model="gpt-4o", response_cost=0.2, total_tokens=20, duration_ms=150.0, **shared),
|
||||
_record(
|
||||
team_id="team-a", model="claude-4", response_cost=0.4, total_tokens=40, duration_ms=200.0, **shared
|
||||
),
|
||||
_record(team_id="team-b", model="gpt-4o", response_cost=0.8, total_tokens=80, duration_ms=300.0, **shared),
|
||||
)
|
||||
payload = build_metric_payload(records, window_start=1_000.0, now=1_005.0)
|
||||
|
||||
cost_by_bucket = {
|
||||
(m["attributes"]["team_id"], m["attributes"]["model"]): m["value"]
|
||||
for m in _metrics_by_name(payload, NEWRELIC_METRIC_COST_USD)
|
||||
}
|
||||
assert cost_by_bucket == {
|
||||
("team-a", "gpt-4o"): pytest.approx(0.3),
|
||||
("team-a", "claude-4"): pytest.approx(0.4),
|
||||
("team-b", "gpt-4o"): pytest.approx(0.8),
|
||||
}
|
||||
|
||||
requests_by_bucket = {
|
||||
(m["attributes"]["team_id"], m["attributes"]["model"]): m["value"]
|
||||
for m in _metrics_by_name(payload, NEWRELIC_METRIC_REQUESTS)
|
||||
}
|
||||
assert requests_by_bucket == {
|
||||
("team-a", "gpt-4o"): 2.0,
|
||||
("team-a", "claude-4"): 1.0,
|
||||
("team-b", "gpt-4o"): 1.0,
|
||||
}
|
||||
|
||||
duration_by_bucket = {
|
||||
(m["attributes"]["team_id"], m["attributes"]["model"]): m["value"]
|
||||
for m in _metrics_by_name(payload, NEWRELIC_METRIC_REQUEST_DURATION_MS)
|
||||
}
|
||||
assert duration_by_bucket[("team-a", "gpt-4o")] == {"count": 2, "sum": 200.0, "min": 50.0, "max": 150.0}
|
||||
|
||||
def test_status_is_a_bucket_dimension(self):
|
||||
records = (
|
||||
_record(status="success", response_cost=0.1),
|
||||
_record(status="failure", response_cost=0.0),
|
||||
)
|
||||
payload = build_metric_payload(records, window_start=1_000.0, now=1_005.0)
|
||||
|
||||
statuses = {m["attributes"]["status"] for m in _metrics_by_name(payload, NEWRELIC_METRIC_REQUESTS)}
|
||||
assert statuses == {"success", "failure"}
|
||||
|
||||
def test_empty_attribute_values_are_omitted(self):
|
||||
record = NewRelicMetricRecord(
|
||||
team_id="",
|
||||
team_alias="",
|
||||
model_group="",
|
||||
model="gpt-4o",
|
||||
custom_llm_provider="openai",
|
||||
status="success",
|
||||
response_cost=0.0,
|
||||
prompt_tokens=0,
|
||||
completion_tokens=0,
|
||||
total_tokens=0,
|
||||
duration_ms=0.0,
|
||||
)
|
||||
payload = build_metric_payload((record,), window_start=1_000.0, now=1_005.0)
|
||||
|
||||
attributes = payload[0]["metrics"][0]["attributes"]
|
||||
assert "team_id" not in attributes
|
||||
assert "team_alias" not in attributes
|
||||
assert "model_group" not in attributes
|
||||
|
||||
|
||||
class TestQueueAndFlush:
|
||||
@pytest.mark.asyncio
|
||||
async def test_log_event_queues_record_from_standard_logging_object(self):
|
||||
logger = _make_logger()
|
||||
|
||||
await logger.async_log_success_event(
|
||||
kwargs={"standard_logging_object": _standard_logging_object()},
|
||||
response_obj={},
|
||||
start_time=None,
|
||||
end_time=None,
|
||||
)
|
||||
|
||||
assert len(logger.log_queue) == 1
|
||||
record = logger.log_queue[0]
|
||||
assert record.team_id == "team-a"
|
||||
assert record.response_cost == 0.25
|
||||
assert record.duration_ms == pytest.approx(100.0)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failure_event_queues_record(self):
|
||||
logger = _make_logger()
|
||||
|
||||
slo = _standard_logging_object()
|
||||
slo["status"] = "failure"
|
||||
await logger.async_log_failure_event(
|
||||
kwargs={"standard_logging_object": slo},
|
||||
response_obj={},
|
||||
start_time=None,
|
||||
end_time=None,
|
||||
)
|
||||
|
||||
assert len(logger.log_queue) == 1
|
||||
assert logger.log_queue[0].status == "failure"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_threshold_flush_uses_flush_queue(self):
|
||||
logger = _make_logger()
|
||||
logger.batch_size = 1
|
||||
logger.flush_queue = AsyncMock()
|
||||
|
||||
await logger.async_log_success_event(
|
||||
kwargs={"standard_logging_object": _standard_logging_object()},
|
||||
response_obj={},
|
||||
start_time=None,
|
||||
end_time=None,
|
||||
)
|
||||
|
||||
logger.flush_queue.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flush_queue_updates_last_flush_time_on_success(self):
|
||||
logger = _make_logger()
|
||||
logger.log_queue = [_record()]
|
||||
logger.last_flush_time = 0
|
||||
logger.async_client.post = AsyncMock(return_value=_response(202))
|
||||
|
||||
await logger.flush_queue()
|
||||
|
||||
assert logger.log_queue == []
|
||||
assert logger.last_flush_time > 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flush_advances_window_even_on_requeue(self):
|
||||
# The window start advances every flush cycle so requeued records report
|
||||
# in the next window instead of freezing interval.ms under sustained
|
||||
# failure, and an idle gap never inflates the next batch's window
|
||||
logger = _make_logger()
|
||||
logger.log_queue = [_record()]
|
||||
logger.last_flush_time = 123.0
|
||||
logger.async_client.post = _raises(500)
|
||||
|
||||
await logger.flush_queue()
|
||||
|
||||
assert logger.last_flush_time > 123.0
|
||||
assert len(logger.log_queue) == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sent_payload_window_starts_at_last_flush_time(self):
|
||||
logger = _make_logger()
|
||||
logger.log_queue = [_record()]
|
||||
logger.last_flush_time = 2_000.0
|
||||
logger.async_client.post = AsyncMock(return_value=_response(202))
|
||||
|
||||
with patch("litellm.integrations.newrelic.newrelic_metrics.time.time", return_value=2_010.0):
|
||||
await logger.async_send_batch()
|
||||
|
||||
sent = logger.async_client.post.await_args.kwargs
|
||||
body = json.loads(gzip.decompress(sent["data"]).decode("utf-8"))
|
||||
assert body[0]["common"]["timestamp"] == 2_000_000
|
||||
assert body[0]["common"]["interval.ms"] == 10_000
|
||||
assert sent["headers"]["Api-Key"] == "test-key"
|
||||
assert sent["headers"]["Content-Encoding"] == "gzip"
|
||||
assert sent["url"] == NEWRELIC_METRIC_ENDPOINT_BY_REGION["us"]
|
||||
|
||||
|
||||
class TestBatchSizeCap:
|
||||
@pytest.mark.asyncio
|
||||
async def test_flush_sends_at_most_batch_size_records_per_request(self):
|
||||
"""A queue grown past the batch size by requeues must go out in chunks:
|
||||
one oversized request would breach the Metric API data point cap and get
|
||||
the whole retry backlog dropped as a 4xx."""
|
||||
logger = _make_logger()
|
||||
logger.batch_size = 2
|
||||
logger.log_queue = [_record(model=f"model-{i}") for i in range(5)]
|
||||
logger.async_client.post = AsyncMock(return_value=_response(202))
|
||||
|
||||
await logger.flush_queue()
|
||||
|
||||
sent_counts = [
|
||||
sum(
|
||||
metric["value"]
|
||||
for metric in json.loads(gzip.decompress(call.kwargs["data"]).decode("utf-8"))[0]["metrics"]
|
||||
if metric["name"] == NEWRELIC_METRIC_REQUESTS
|
||||
)
|
||||
for call in logger.async_client.post.await_args_list
|
||||
]
|
||||
assert sent_counts == [2.0, 2.0, 1.0]
|
||||
assert logger.log_queue == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_chunk_stops_the_flush_and_keeps_order(self):
|
||||
"""A 5xx on the first chunk ends the flush instead of hammering the same
|
||||
failing endpoint with the rest of the backlog, and the requeue keeps the
|
||||
records in chronological order."""
|
||||
logger = _make_logger()
|
||||
logger.batch_size = 2
|
||||
records = [_record(model=f"model-{i}") for i in range(5)]
|
||||
logger.log_queue = list(records)
|
||||
logger.async_client.post = _raises(500)
|
||||
|
||||
await logger.flush_queue()
|
||||
|
||||
assert logger.async_client.post.await_count == 1
|
||||
assert logger.log_queue == records
|
||||
|
||||
|
||||
class TestFlushConcurrency:
|
||||
@pytest.mark.asyncio
|
||||
async def test_records_appended_during_flush_await_survive(self):
|
||||
"""A record appended by a concurrent request while the POST is in flight
|
||||
must survive the flush, not be clobbered by a queue replacement."""
|
||||
logger = _make_logger()
|
||||
logger.log_queue = [_record(team_id="team-a")]
|
||||
interleaved = _record(team_id="team-interleaved")
|
||||
|
||||
async def _post_appending_mid_flight(**kwargs):
|
||||
logger.log_queue.append(interleaved)
|
||||
return _response(202)
|
||||
|
||||
logger.async_client.post = AsyncMock(side_effect=_post_appending_mid_flight)
|
||||
|
||||
await logger.async_send_batch()
|
||||
|
||||
assert logger.log_queue == [interleaved]
|
||||
body = json.loads(gzip.decompress(logger.async_client.post.await_args.kwargs["data"]).decode("utf-8"))
|
||||
team_ids = {m["attributes"]["team_id"] for m in body[0]["metrics"]}
|
||||
assert team_ids == {"team-a"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_records_appended_during_failed_flush_await_survive_requeue(self):
|
||||
"""The requeue path must also preserve interleaved records: batch is
|
||||
prepended in place, never assigned over the live queue."""
|
||||
logger = _make_logger()
|
||||
original = _record(team_id="team-a")
|
||||
logger.log_queue = [original]
|
||||
interleaved = _record(team_id="team-interleaved")
|
||||
|
||||
async def _post_appending_mid_flight(**kwargs):
|
||||
logger.log_queue.append(interleaved)
|
||||
raise HTTPStatusError('e', request=_response(500).request, response=_response(500))
|
||||
|
||||
logger.async_client.post = AsyncMock(side_effect=_post_appending_mid_flight)
|
||||
|
||||
await logger.async_send_batch()
|
||||
|
||||
assert logger.log_queue == [original, interleaved]
|
||||
|
||||
|
||||
class TestErrorPolicy:
|
||||
@pytest.mark.asyncio
|
||||
async def test_4xx_drops_batch(self):
|
||||
logger = _make_logger()
|
||||
logger.log_queue = [_record(), _record(team_id="team-b")]
|
||||
logger.async_client.post = AsyncMock(return_value=_response(400, text="bad request"))
|
||||
|
||||
await logger.async_send_batch()
|
||||
|
||||
assert logger.log_queue == []
|
||||
assert logger.async_client.post.await_count == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_403_drops_batch_and_names_permanent_credential_failure(self):
|
||||
logger = _make_logger()
|
||||
logger.log_queue = [_record()]
|
||||
logger.async_client.post = _raises(403)
|
||||
|
||||
with patch("litellm.integrations.newrelic.newrelic_metrics.verbose_logger") as mock_logger:
|
||||
await logger.async_send_batch()
|
||||
|
||||
assert logger.log_queue == []
|
||||
warning_text = " ".join(str(arg) for call in mock_logger.warning.call_args_list for arg in call.args)
|
||||
assert "permanent credential failure" in warning_text
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_5xx_requeues_batch(self):
|
||||
records = [_record(), _record(team_id="team-b")]
|
||||
logger = _make_logger()
|
||||
logger.log_queue = list(records)
|
||||
logger.async_client.post = _raises(500)
|
||||
|
||||
await logger.async_send_batch()
|
||||
|
||||
assert logger.log_queue == records
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_network_error_requeues_batch(self):
|
||||
records = [_record()]
|
||||
logger = _make_logger()
|
||||
logger.log_queue = list(records)
|
||||
logger.async_client.post = AsyncMock(side_effect=ConnectionError("boom"))
|
||||
|
||||
await logger.async_send_batch()
|
||||
|
||||
assert logger.log_queue == records
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_requeue_is_capped_dropping_oldest(self):
|
||||
logger = _make_logger()
|
||||
logger.max_queue_size = 3
|
||||
oldest = _record(team_id="oldest")
|
||||
rest = [_record(team_id=f"team-{i}") for i in range(3)]
|
||||
logger.log_queue = [oldest, *rest]
|
||||
logger.async_client.post = _raises(500)
|
||||
|
||||
await logger.async_send_batch()
|
||||
|
||||
assert logger.log_queue == rest
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_requeued_records_are_resent_with_new_records(self):
|
||||
logger = _make_logger()
|
||||
logger.log_queue = [_record()]
|
||||
logger.async_client.post = _raises(500)
|
||||
|
||||
await logger.async_send_batch()
|
||||
logger.log_queue.append(_record(team_id="team-b"))
|
||||
logger.async_client.post = AsyncMock(return_value=_response(202))
|
||||
|
||||
await logger.async_send_batch()
|
||||
|
||||
sent = logger.async_client.post.await_args.kwargs
|
||||
body = json.loads(gzip.decompress(sent["data"]).decode("utf-8"))
|
||||
team_ids = {m["attributes"]["team_id"] for m in body[0]["metrics"]}
|
||||
assert team_ids == {"team-a", "team-b"}
|
||||
assert logger.log_queue == []
|
||||
|
||||
|
||||
class TestStopFlag:
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_ends_periodic_flush_loop(self):
|
||||
logger = _make_logger()
|
||||
logger.flush_interval = 0.01
|
||||
logger.flush_queue = AsyncMock()
|
||||
|
||||
task = asyncio.create_task(logger.periodic_flush())
|
||||
await asyncio.sleep(0.05)
|
||||
assert not task.done()
|
||||
|
||||
logger.stop()
|
||||
await asyncio.wait_for(task, timeout=1.0)
|
||||
|
||||
assert task.done()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stopped_logger_exits_after_one_final_drain(self):
|
||||
logger = _make_logger()
|
||||
logger.flush_interval = 0.01
|
||||
logger._final_drain = AsyncMock()
|
||||
logger._stopped = True
|
||||
|
||||
await asyncio.wait_for(logger.periodic_flush(), timeout=1.0)
|
||||
|
||||
logger._final_drain.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_eviction_drains_queued_records(self):
|
||||
"""Eviction must post what is already queued, not silently discard it."""
|
||||
from litellm.litellm_core_utils.specialty_caches.dynamic_logging_cache import (
|
||||
DynamicLoggingCache,
|
||||
)
|
||||
|
||||
cache = DynamicLoggingCache()
|
||||
logger = _make_logger()
|
||||
logger.log_queue = [_record(), _record(team_id="team-b")]
|
||||
logger.async_client.post = AsyncMock(return_value=_response(202))
|
||||
credentials = {"newrelic_api_key": "test-key", "newrelic_region": None}
|
||||
cache.set_cache(credentials=credentials, service_name="newrelic", logging_obj=logger)
|
||||
|
||||
key = cache.get_cache_key(args={**credentials, "service_name": "newrelic"})
|
||||
cache.cache._remove_key(key)
|
||||
for _ in range(10):
|
||||
await asyncio.sleep(0)
|
||||
|
||||
logger.async_client.post.assert_awaited_once()
|
||||
body = json.loads(gzip.decompress(logger.async_client.post.await_args.kwargs["data"]).decode("utf-8"))
|
||||
team_ids = {m["attributes"]["team_id"] for m in body[0]["metrics"]}
|
||||
assert team_ids == {"team-a", "team-b"}
|
||||
assert logger.log_queue == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dynamic_logging_cache_eviction_calls_stop(self):
|
||||
from litellm.litellm_core_utils.specialty_caches.dynamic_logging_cache import (
|
||||
DynamicLoggingCache,
|
||||
)
|
||||
|
||||
cache = DynamicLoggingCache()
|
||||
logger = _make_logger()
|
||||
credentials = {"newrelic_api_key": "test-key", "newrelic_region": None}
|
||||
cache.set_cache(credentials=credentials, service_name="newrelic", logging_obj=logger)
|
||||
|
||||
key = cache.get_cache_key(args={**credentials, "service_name": "newrelic"})
|
||||
cache.cache._remove_key(key)
|
||||
|
||||
assert logger._stopped is True
|
||||
assert cache.get_cache(credentials=credentials, service_name="newrelic") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_append_after_eviction_drain_self_flushes():
|
||||
"""An in-flight callback holding an evicted (stopped) logger still delivers
|
||||
its record: with no periodic loop left, the append itself drains."""
|
||||
logger = _make_logger()
|
||||
with patch.object(
|
||||
logger.async_client, "post", new=AsyncMock(return_value=_response(202))
|
||||
) as mock_post:
|
||||
logger.stop()
|
||||
await logger.async_log_success_event(
|
||||
{"standard_logging_object": _standard_logging_object()}, None, None, None
|
||||
)
|
||||
assert mock_post.await_count >= 1, "record appended after stop() must be flushed, not stranded"
|
||||
assert logger.log_queue == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_final_drain_retries_transient_failure_then_delivers():
|
||||
"""A transient 5xx during the eviction drain must not strand the last
|
||||
batch: the final drain retries on its own (no periodic loop is left)."""
|
||||
logger = _make_logger()
|
||||
err = _response(500)
|
||||
responses = [HTTPStatusError('e', request=err.request, response=err), HTTPStatusError('e', request=err.request, response=err), _response(202)]
|
||||
post_mock = AsyncMock(side_effect=responses)
|
||||
with patch.object(logger, "async_client") as client, patch("asyncio.sleep", new=AsyncMock()):
|
||||
client.post = post_mock
|
||||
await logger._log_async_event(standard_logging_object=_standard_logging_object())
|
||||
await logger._final_drain()
|
||||
assert post_mock.await_count == 3
|
||||
assert logger.log_queue == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_final_drain_drops_after_bounded_passes_under_lock():
|
||||
"""A permanently failing destination is retried across bounded passes, then
|
||||
the remainder is dropped under flush_lock and logged, never stranded. A
|
||||
second drain over the now-empty queue is a no-op."""
|
||||
logger = _make_logger()
|
||||
post_mock = _raises(500)
|
||||
with patch.object(logger, "async_client") as client, patch("asyncio.sleep", new=AsyncMock()):
|
||||
client.post = post_mock
|
||||
await logger._log_async_event(standard_logging_object=_standard_logging_object())
|
||||
await logger._final_drain()
|
||||
after_first = post_mock.await_count
|
||||
await logger._final_drain()
|
||||
assert after_first >= 1, "the failing destination was retried before the drop"
|
||||
assert post_mock.await_count == after_first, "second drain over an empty queue is a no-op"
|
||||
assert logger.log_queue == [], "exhausted retries end in a logged drop, not a stranded queue"
|
||||
|
||||
|
||||
def test_attribute_values_bounded_against_payload_bombs():
|
||||
"""A caller-controlled high-entropy model string is truncated in metric
|
||||
attributes so one record cannot inflate the shared batch past the Metric
|
||||
API payload cap and take out other users' metrics."""
|
||||
record = _record(model="m" * 5000)
|
||||
metrics = _bucket_metrics((record,))
|
||||
for metric in metrics:
|
||||
assert len(metric["attributes"]["model"]) == 255
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_idle_gap_does_not_inflate_next_window():
|
||||
"""Empty flush cycles advance the window start, so a burst after idling
|
||||
reports an interval close to the flush cadence, not the whole idle gap."""
|
||||
logger = _make_logger()
|
||||
logger.last_flush_time = 100.0
|
||||
with patch.object(logger, "async_client") as client:
|
||||
client.post = AsyncMock(return_value=_response(202))
|
||||
await logger.flush_queue()
|
||||
assert logger.last_flush_time > 100.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mid_drain_append_delivered_against_healthy_destination():
|
||||
"""A record a callback appends while a drain is running is picked up by a
|
||||
later pass and delivered when the destination is healthy; nothing stranded."""
|
||||
logger = _make_logger()
|
||||
logger.stop()
|
||||
late_record = _record(model="late-model")
|
||||
injected = {"done": False}
|
||||
posted = []
|
||||
|
||||
async def _capture(url, headers=None, content=None, **kw):
|
||||
posted.append(content)
|
||||
if not injected["done"]:
|
||||
injected["done"] = True
|
||||
logger.log_queue.append(late_record)
|
||||
return _response(202)
|
||||
|
||||
with patch.object(logger, "async_client") as client, patch("asyncio.sleep", new=AsyncMock()):
|
||||
client.post = _capture
|
||||
logger.log_queue.append(_record(model="first"))
|
||||
await logger._drain_with_retry()
|
||||
assert logger.log_queue == [], "the mid-drain append was drained too, nothing stranded"
|
||||
assert len(posted) >= 2, "both the original and the mid-drain record were sent"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_drain_attempts_every_chunk_not_just_the_head_under_failure():
|
||||
"""Regression: with more than batch_size records queued on a stopped logger
|
||||
and a persistently failing destination, every record must be attempted before
|
||||
the bounded terminal drop. The periodic path stops at the first failing chunk,
|
||||
so a drain that reused it would drop the un-sent tail (records past the head
|
||||
chunk) as if it had tried them, silently undercounting the team's usage."""
|
||||
logger = _make_logger()
|
||||
logger.stop()
|
||||
logger.batch_size = 2
|
||||
logger.log_queue = [_record(model=f"m{i}") for i in range(5)]
|
||||
sent_models = []
|
||||
|
||||
async def _capture_then_fail(url, data=None, headers=None, **kw):
|
||||
body = json.loads(gzip.decompress(data).decode("utf-8"))
|
||||
sent_models.extend(
|
||||
m["attributes"]["model"] for m in body[0]["metrics"] if m["name"] == NEWRELIC_METRIC_REQUESTS
|
||||
)
|
||||
resp = _response(503)
|
||||
raise HTTPStatusError("err", request=resp.request, response=resp)
|
||||
|
||||
with patch("asyncio.sleep", new=AsyncMock()):
|
||||
logger.async_client.post = _capture_then_fail
|
||||
await logger._drain_with_retry()
|
||||
|
||||
assert set(sent_models) == {"m0", "m1", "m2", "m3", "m4"}, "every chunk, including the tail, was attempted"
|
||||
assert logger.log_queue == [], "the exhausted batch is dropped after bounded passes, nothing stranded"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_drain_delivers_the_tail_once_the_destination_recovers():
|
||||
"""The tail beyond the head chunk must be delivered, not stranded, once a
|
||||
transiently failing destination recovers within the drain's passes."""
|
||||
logger = _make_logger()
|
||||
logger.stop()
|
||||
logger.batch_size = 2
|
||||
logger.log_queue = [_record(model=f"m{i}") for i in range(5)]
|
||||
delivered_models = []
|
||||
posts = {"n": 0}
|
||||
|
||||
async def _fail_first_pass_then_recover(url, data=None, headers=None, **kw):
|
||||
posts["n"] += 1
|
||||
if posts["n"] <= 3: # the first pass's three chunks all fail
|
||||
resp = _response(503)
|
||||
raise HTTPStatusError("err", request=resp.request, response=resp)
|
||||
body = json.loads(gzip.decompress(data).decode("utf-8"))
|
||||
delivered_models.extend(
|
||||
m["attributes"]["model"] for m in body[0]["metrics"] if m["name"] == NEWRELIC_METRIC_REQUESTS
|
||||
)
|
||||
return _response(202)
|
||||
|
||||
with patch("asyncio.sleep", new=AsyncMock()):
|
||||
logger.async_client.post = _fail_first_pass_then_recover
|
||||
await logger._drain_with_retry()
|
||||
|
||||
assert set(delivered_models) == {"m0", "m1", "m2", "m3", "m4"}, "all chunks delivered after recovery"
|
||||
assert logger.log_queue == [], "nothing left stranded once the destination recovered"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_terminal_drop_leaves_untried_late_arrival_for_next_drain():
|
||||
"""Against a permanently failing destination, the terminal drop clears only
|
||||
the records this drain actually tried; a record a callback appends during the
|
||||
final pass, after that pass's snapshot, is left in the queue for its own
|
||||
serialized drain, never wiped un-tried."""
|
||||
logger = _make_logger()
|
||||
logger.stop()
|
||||
from litellm.types.integrations.newrelic import NEWRELIC_METRICS_MAX_DRAIN_PASSES
|
||||
|
||||
late_record = _record(model="late-arrival")
|
||||
posts = {"n": 0}
|
||||
|
||||
async def _fail_and_append_on_final_pass(url, data=None, headers=None, **kw):
|
||||
posts["n"] += 1
|
||||
# One record means one post per pass, so the final pass's post is the
|
||||
# Nth; append then, after the drain has already snapshotted the queue.
|
||||
if posts["n"] == NEWRELIC_METRICS_MAX_DRAIN_PASSES:
|
||||
logger.log_queue.append(late_record)
|
||||
resp = _response(503)
|
||||
raise HTTPStatusError("err", request=resp.request, response=resp)
|
||||
|
||||
with patch("asyncio.sleep", new=AsyncMock()):
|
||||
logger.async_client.post = _fail_and_append_on_final_pass
|
||||
logger.log_queue.append(_record(model="doomed"))
|
||||
await logger._drain_with_retry()
|
||||
assert logger.log_queue == [late_record], "the un-tried late arrival is left for its own drain, not dropped"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_appended_on_an_early_pass_is_not_dropped_short_of_the_retry_budget():
|
||||
"""A record a callback appends during an early drain pass entered the queue
|
||||
after this drain's snapshot, so it has not seen the full retry budget. The
|
||||
terminal drop must clear only records queued when the drain began, leaving
|
||||
the early-pass arrival for its own serialized drain instead of dropping it
|
||||
after fewer than the configured attempts."""
|
||||
logger = _make_logger()
|
||||
logger.stop()
|
||||
early_record = _record(model="early-pass-arrival")
|
||||
posts = {"n": 0}
|
||||
|
||||
async def _fail_and_append_on_first_pass(url, data=None, headers=None, **kw):
|
||||
posts["n"] += 1
|
||||
# One record queued at start means the first pass's post is the 1st;
|
||||
# append during it, before this drain's later passes.
|
||||
if posts["n"] == 1:
|
||||
logger.log_queue.append(early_record)
|
||||
resp = _response(503)
|
||||
raise HTTPStatusError("err", request=resp.request, response=resp)
|
||||
|
||||
with patch("asyncio.sleep", new=AsyncMock()):
|
||||
logger.async_client.post = _fail_and_append_on_first_pass
|
||||
logger.log_queue.append(_record(model="doomed"))
|
||||
await logger._drain_with_retry()
|
||||
assert logger.log_queue == [early_record], "the early-pass arrival is left for its own drain, not dropped short"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_stop_drains_are_serialized():
|
||||
"""A callback that appends to a stopped logger and starts its own drain must
|
||||
queue behind an already-running drain, not race it: otherwise one drain's
|
||||
terminal clear could wipe a record the other is still responsible for.
|
||||
Proven by holding the first drain inside its flush and asserting the second
|
||||
has not entered its own flush until the first releases."""
|
||||
logger = _make_logger()
|
||||
logger._stopped = True # stopped without scheduling a background drain
|
||||
logger.log_queue.append(_record(model="r1"))
|
||||
entered = []
|
||||
release = asyncio.Event()
|
||||
|
||||
async def blocking_flush():
|
||||
entered.append(len(entered) + 1)
|
||||
if len(entered) == 1:
|
||||
await release.wait()
|
||||
logger.log_queue.clear()
|
||||
|
||||
logger._drain_flush_once = blocking_flush
|
||||
t1 = asyncio.create_task(logger._drain_with_retry())
|
||||
await asyncio.sleep(0.02) # let t1 acquire the drain lock and enter flush
|
||||
assert entered == [1], f"first drain did not enter flush: {entered}"
|
||||
t2 = asyncio.create_task(logger._drain_with_retry())
|
||||
await asyncio.sleep(0.02) # t2 must block on the drain lock, not enter flush
|
||||
assert entered == [1], f"second drain raced the first: {entered}"
|
||||
release.set()
|
||||
await asyncio.gather(t1, t2)
|
||||
assert logger.log_queue == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raised_403_is_dropped_not_requeued():
|
||||
"""AsyncHTTPHandler.post raises HTTPStatusError on 4xx, so a 403 (permanent
|
||||
bad key) arrives as an exception, not a response. It must be dropped, never
|
||||
requeued, or a revoked key retries forever."""
|
||||
logger = _make_logger()
|
||||
logger.log_queue.append(_record())
|
||||
logger.async_client.post = _raises(403)
|
||||
await logger.async_send_batch()
|
||||
assert logger.log_queue == [], "a permanent 403 must drop, not requeue"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raised_500_is_requeued():
|
||||
"""A raised 5xx is transient and must be requeued for retry."""
|
||||
logger = _make_logger()
|
||||
record = _record()
|
||||
logger.log_queue.append(record)
|
||||
logger.async_client.post = _raises(503)
|
||||
await logger.async_send_batch()
|
||||
assert logger.log_queue == [record], "a transient 5xx must requeue"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("status", [429, 408])
|
||||
async def test_transient_4xx_is_requeued_not_dropped(status):
|
||||
"""The Metric API returns 429 when it throttles (and 408 on a request
|
||||
timeout); both are transient and expect a retry, so the batch must be
|
||||
requeued rather than permanently dropped like a 400/403."""
|
||||
logger = _make_logger()
|
||||
record = _record()
|
||||
logger.log_queue.append(record)
|
||||
logger.async_client.post = _raises(status)
|
||||
await logger.async_send_batch()
|
||||
assert logger.log_queue == [record], f"a transient {status} must requeue, not drop"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("status", [200, 201, 204])
|
||||
async def test_any_2xx_is_treated_as_delivered_not_requeued(status):
|
||||
"""The Metric API answers 202, but any 2xx means the destination accepted the
|
||||
batch. Treating a non-202 2xx as a failure would re-queue and re-send data
|
||||
New Relic already stored, duplicating the team's metrics until the cap drops."""
|
||||
logger = _make_logger()
|
||||
logger.log_queue.append(_record())
|
||||
logger.async_client.post = AsyncMock(return_value=_response(status))
|
||||
await logger.async_send_batch()
|
||||
assert logger.log_queue == [], f"a {status} success must drop, not requeue and duplicate"
|
||||
|
|
@ -0,0 +1,274 @@
|
|||
"""
|
||||
Tests for team-scoped New Relic metrics callback support.
|
||||
|
||||
Verifies that NewRelicMetricsLogger is instantiated with per-team credentials
|
||||
(newrelic_api_key, newrelic_region) with no environment fallback, and that
|
||||
NewRelicHandler correctly resolves and caches per-team loggers.
|
||||
"""
|
||||
|
||||
import copy
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.integrations.newrelic.newrelic_metrics import NewRelicMetricsLogger
|
||||
from litellm.integrations.newrelic.newrelic_team_handler import NewRelicHandler
|
||||
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
|
||||
TRUSTED_CALLBACK_VARS_FIELD,
|
||||
)
|
||||
from litellm.litellm_core_utils.specialty_caches.dynamic_logging_cache import (
|
||||
DynamicLoggingCache,
|
||||
)
|
||||
from litellm.types.integrations.newrelic import NEWRELIC_METRIC_ENDPOINT_BY_REGION
|
||||
from litellm.types.utils import StandardCallbackDynamicParams
|
||||
|
||||
US_ENDPOINT = NEWRELIC_METRIC_ENDPOINT_BY_REGION["us"]
|
||||
EU_ENDPOINT = NEWRELIC_METRIC_ENDPOINT_BY_REGION["eu"]
|
||||
|
||||
|
||||
class TestNewRelicMetricsLoggerCredentialKwargs:
|
||||
"""The logger takes credentials by injection only; env vars never leak in."""
|
||||
|
||||
def test_init_with_explicit_credentials(self):
|
||||
with patch("asyncio.create_task"):
|
||||
logger = NewRelicMetricsLogger(newrelic_api_key="team_key", newrelic_region="eu")
|
||||
|
||||
assert logger.newrelic_api_key == "team_key"
|
||||
assert logger.metric_api_url == EU_ENDPOINT
|
||||
|
||||
def test_init_defaults_to_us_region(self):
|
||||
with patch("asyncio.create_task"):
|
||||
logger = NewRelicMetricsLogger(newrelic_api_key="team_key")
|
||||
|
||||
assert logger.metric_api_url == US_ENDPOINT
|
||||
|
||||
def test_unknown_region_falls_back_to_us(self):
|
||||
with patch("asyncio.create_task"):
|
||||
logger = NewRelicMetricsLogger(newrelic_api_key="team_key", newrelic_region="mars")
|
||||
|
||||
assert logger.metric_api_url == US_ENDPOINT
|
||||
|
||||
def test_region_is_case_insensitive(self):
|
||||
with patch("asyncio.create_task"):
|
||||
logger = NewRelicMetricsLogger(newrelic_api_key="team_key", newrelic_region="EU")
|
||||
|
||||
assert logger.metric_api_url == EU_ENDPOINT
|
||||
|
||||
def test_init_raises_without_api_key(self):
|
||||
with pytest.raises(ValueError, match="newrelic_api_key"):
|
||||
with patch("asyncio.create_task"):
|
||||
NewRelicMetricsLogger(newrelic_api_key="")
|
||||
|
||||
def test_init_never_falls_back_to_env_license_key(self, monkeypatch):
|
||||
"""A missing team key must fail, never silently reuse the operator's key."""
|
||||
monkeypatch.setenv("NEW_RELIC_LICENSE_KEY", "operator-license-key")
|
||||
|
||||
with pytest.raises(ValueError, match="newrelic_api_key"):
|
||||
with patch("asyncio.create_task"):
|
||||
NewRelicMetricsLogger(newrelic_api_key="")
|
||||
|
||||
|
||||
class TestNewRelicHandler:
|
||||
"""The handler resolves the correct logger per team."""
|
||||
|
||||
def test_creates_team_logger_with_dynamic_credentials(self):
|
||||
cache = DynamicLoggingCache()
|
||||
params = StandardCallbackDynamicParams(newrelic_api_key="team_a_key", newrelic_region="eu")
|
||||
|
||||
with patch("asyncio.create_task"):
|
||||
result = NewRelicHandler.get_newrelic_logger_for_request(
|
||||
standard_callback_dynamic_params=params,
|
||||
in_memory_dynamic_logger_cache=cache,
|
||||
)
|
||||
|
||||
assert result.newrelic_api_key == "team_a_key"
|
||||
assert result.metric_api_url == EU_ENDPOINT
|
||||
|
||||
def test_caches_team_logger(self):
|
||||
cache = DynamicLoggingCache()
|
||||
params = StandardCallbackDynamicParams(newrelic_api_key="team_b_key")
|
||||
|
||||
with patch("asyncio.create_task"):
|
||||
result1 = NewRelicHandler.get_newrelic_logger_for_request(
|
||||
standard_callback_dynamic_params=params,
|
||||
in_memory_dynamic_logger_cache=cache,
|
||||
)
|
||||
result2 = NewRelicHandler.get_newrelic_logger_for_request(
|
||||
standard_callback_dynamic_params=params,
|
||||
in_memory_dynamic_logger_cache=cache,
|
||||
)
|
||||
|
||||
assert result1 is result2
|
||||
|
||||
def test_different_teams_get_different_loggers(self):
|
||||
cache = DynamicLoggingCache()
|
||||
params_a = StandardCallbackDynamicParams(newrelic_api_key="team_a_key")
|
||||
params_b = StandardCallbackDynamicParams(newrelic_api_key="team_b_key")
|
||||
|
||||
with patch("asyncio.create_task"):
|
||||
result_a = NewRelicHandler.get_newrelic_logger_for_request(
|
||||
standard_callback_dynamic_params=params_a,
|
||||
in_memory_dynamic_logger_cache=cache,
|
||||
)
|
||||
result_b = NewRelicHandler.get_newrelic_logger_for_request(
|
||||
standard_callback_dynamic_params=params_b,
|
||||
in_memory_dynamic_logger_cache=cache,
|
||||
)
|
||||
|
||||
assert result_a is not result_b
|
||||
assert result_a.newrelic_api_key == "team_a_key"
|
||||
assert result_b.newrelic_api_key == "team_b_key"
|
||||
|
||||
def test_region_is_part_of_cache_key(self):
|
||||
"""Same key, different region must not share a logger (different endpoints)."""
|
||||
cache = DynamicLoggingCache()
|
||||
|
||||
with patch("asyncio.create_task"):
|
||||
result_us = NewRelicHandler.get_newrelic_logger_for_request(
|
||||
standard_callback_dynamic_params=StandardCallbackDynamicParams(newrelic_api_key="key"),
|
||||
in_memory_dynamic_logger_cache=cache,
|
||||
)
|
||||
result_eu = NewRelicHandler.get_newrelic_logger_for_request(
|
||||
standard_callback_dynamic_params=StandardCallbackDynamicParams(
|
||||
newrelic_api_key="key", newrelic_region="eu"
|
||||
),
|
||||
in_memory_dynamic_logger_cache=cache,
|
||||
)
|
||||
|
||||
assert result_us is not result_eu
|
||||
assert result_us.metric_api_url == US_ENDPOINT
|
||||
assert result_eu.metric_api_url == EU_ENDPOINT
|
||||
|
||||
def test_request_blocked_callback_params_includes_newrelic(self):
|
||||
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
|
||||
_request_blocked_callback_params,
|
||||
)
|
||||
|
||||
assert "newrelic_api_key" in _request_blocked_callback_params
|
||||
assert "newrelic_region" in _request_blocked_callback_params
|
||||
|
||||
|
||||
class TestDynamicCredentialDetection:
|
||||
def test_no_credentials(self):
|
||||
params = StandardCallbackDynamicParams()
|
||||
assert NewRelicHandler._dynamic_newrelic_credentials_are_passed(params) is False
|
||||
|
||||
def test_region_only_is_not_credentials(self):
|
||||
params = StandardCallbackDynamicParams(newrelic_region="eu")
|
||||
assert NewRelicHandler._dynamic_newrelic_credentials_are_passed(params) is False
|
||||
|
||||
def test_api_key_is_credentials(self):
|
||||
params = StandardCallbackDynamicParams(newrelic_api_key="key")
|
||||
assert NewRelicHandler._dynamic_newrelic_credentials_are_passed(params) is True
|
||||
|
||||
|
||||
class TestStandardCallbackDynamicParamsIncludesNewRelic:
|
||||
def test_newrelic_params_in_annotations(self):
|
||||
annotations = StandardCallbackDynamicParams.__annotations__
|
||||
assert "newrelic_api_key" in annotations
|
||||
assert "newrelic_region" in annotations
|
||||
|
||||
|
||||
def _build_logging_obj(kwargs: dict, *, with_newrelic_callback: bool = True):
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
|
||||
with patch("asyncio.create_task"):
|
||||
return Logging(
|
||||
model="gpt-4",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=False,
|
||||
call_type="completion",
|
||||
start_time="2026-01-01",
|
||||
litellm_call_id="test-call-id",
|
||||
function_id="test-func",
|
||||
dynamic_success_callbacks=["newrelic"] if with_newrelic_callback else None,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
|
||||
def _metrics_loggers(logging_obj) -> list[NewRelicMetricsLogger]:
|
||||
return [cb for cb in (logging_obj.dynamic_success_callbacks or []) if isinstance(cb, NewRelicMetricsLogger)]
|
||||
|
||||
|
||||
class TestTeamCallbackFlowPassesNewRelicCredentials:
|
||||
"""
|
||||
newrelic_* credentials reach NewRelicHandler only from the proxy-stamped trusted
|
||||
field. Anything the caller put in the request body must not, or a caller could
|
||||
pair its own newrelic_region with the team's ingest key.
|
||||
"""
|
||||
|
||||
def test_trusted_callback_vars_reach_newrelic_handler(self):
|
||||
logging_obj = _build_logging_obj(
|
||||
{
|
||||
TRUSTED_CALLBACK_VARS_FIELD: {"newrelic_api_key": "team-nr-key-123", "newrelic_region": "eu"},
|
||||
"model": "gpt-4",
|
||||
"litellm_params": {"metadata": {}},
|
||||
}
|
||||
)
|
||||
|
||||
metrics_loggers = _metrics_loggers(logging_obj)
|
||||
assert len(metrics_loggers) == 1, "NewRelicMetricsLogger should be initialized from team callback_vars"
|
||||
assert metrics_loggers[0].newrelic_api_key == "team-nr-key-123"
|
||||
assert metrics_loggers[0].metric_api_url == EU_ENDPOINT
|
||||
|
||||
def test_trace_logger_still_dispatched_alongside_metrics(self):
|
||||
"""The metrics logger must not displace the trace logger for the same name."""
|
||||
logging_obj = _build_logging_obj(
|
||||
{
|
||||
TRUSTED_CALLBACK_VARS_FIELD: {"newrelic_api_key": "team-nr-key-123"},
|
||||
"model": "gpt-4",
|
||||
"litellm_params": {"metadata": {}},
|
||||
}
|
||||
)
|
||||
|
||||
non_metrics = [
|
||||
cb for cb in (logging_obj.dynamic_success_callbacks or []) if not isinstance(cb, NewRelicMetricsLogger)
|
||||
]
|
||||
assert len(non_metrics) == 1, "trace logger (OTel v2 or legacy agent) must remain in the dynamic list"
|
||||
assert len(_metrics_loggers(logging_obj)) == 1
|
||||
async_non_metrics = [
|
||||
cb
|
||||
for cb in (logging_obj.dynamic_async_success_callbacks or [])
|
||||
if not isinstance(cb, NewRelicMetricsLogger)
|
||||
]
|
||||
assert len(async_non_metrics) == 1
|
||||
|
||||
def test_request_kwargs_newrelic_params_are_ignored(self):
|
||||
logging_obj = _build_logging_obj(
|
||||
{
|
||||
"newrelic_api_key": "caller-nr-key",
|
||||
"newrelic_region": "eu",
|
||||
"model": "gpt-4",
|
||||
"litellm_params": {"metadata": {}},
|
||||
}
|
||||
)
|
||||
|
||||
assert _metrics_loggers(logging_obj) == []
|
||||
|
||||
def test_logging_object_stays_deepcopyable(self):
|
||||
logging_obj = _build_logging_obj(
|
||||
{
|
||||
TRUSTED_CALLBACK_VARS_FIELD: {"newrelic_api_key": "team-nr-key-123"},
|
||||
"model": "gpt-4",
|
||||
"litellm_params": {"metadata": {}},
|
||||
},
|
||||
with_newrelic_callback=False,
|
||||
)
|
||||
|
||||
assert copy.deepcopy(logging_obj)._trusted_callback_vars == logging_obj._trusted_callback_vars
|
||||
|
||||
def test_caller_cannot_redirect_team_credentials(self):
|
||||
"""The exfil shape: caller's newrelic_region paired with the team's key."""
|
||||
logging_obj = _build_logging_obj(
|
||||
{
|
||||
TRUSTED_CALLBACK_VARS_FIELD: {"newrelic_api_key": "team-nr-key-123"},
|
||||
"newrelic_region": "eu",
|
||||
"model": "gpt-4",
|
||||
"litellm_params": {"metadata": {}},
|
||||
}
|
||||
)
|
||||
|
||||
metrics_loggers = _metrics_loggers(logging_obj)
|
||||
assert len(metrics_loggers) == 1
|
||||
assert metrics_loggers[0].newrelic_api_key == "team-nr-key-123"
|
||||
assert metrics_loggers[0].metric_api_url == US_ENDPOINT
|
||||
|
|
@ -1552,6 +1552,76 @@ def test_string_cost_values():
|
|||
assert round(completion_cost, 12) == round(expected_completion_cost, 12)
|
||||
|
||||
|
||||
def test_generic_cost_per_token_overlapping_cached_and_image_tokens():
|
||||
"""Some providers report cached_tokens and image_tokens as overlapping subsets of
|
||||
prompt_tokens. Billing each in full charged the overlap twice, once at the cache rate
|
||||
and again at the input rate."""
|
||||
model = "litellm-test-overlapping-cached-image"
|
||||
litellm.register_model(
|
||||
{
|
||||
model: {
|
||||
"litellm_provider": "openai",
|
||||
"mode": "chat",
|
||||
"input_cost_per_token": 1e-6,
|
||||
"cache_read_input_token_cost": 1e-7,
|
||||
"output_cost_per_token": 2e-6,
|
||||
}
|
||||
}
|
||||
)
|
||||
usage = Usage(
|
||||
prompt_tokens=100,
|
||||
completion_tokens=10,
|
||||
total_tokens=110,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(
|
||||
text_tokens=None, cached_tokens=90, image_tokens=80
|
||||
),
|
||||
)
|
||||
|
||||
prompt_cost, completion_cost = generic_cost_per_token(
|
||||
model=model, usage=usage, custom_llm_provider="openai"
|
||||
)
|
||||
|
||||
# 90 cached at 1e-7, the remaining 10 uncached tokens once at 1e-6
|
||||
assert prompt_cost == pytest.approx(90 * 1e-7 + 10 * 1e-6)
|
||||
assert completion_cost == pytest.approx(10 * 2e-6)
|
||||
|
||||
|
||||
def test_generic_cost_per_token_warm_prefix_cache_spanning_text_and_image_tokens():
|
||||
"""xAI reports text_tokens + image_tokens = prompt_tokens with cached_tokens overlapping
|
||||
both, so a warm prefix cache covering the whole image exceeds the text-only count.
|
||||
Observed live on grok-4.6 (issue #37281): the image tokens were billed a second time at
|
||||
the full input rate on top of the cache-read bucket, 0.003500 in vs the provider's own
|
||||
0.001274 bill."""
|
||||
model = "litellm-test-warm-prefix-cache-overlap"
|
||||
litellm.register_model(
|
||||
{
|
||||
model: {
|
||||
"litellm_provider": "openai",
|
||||
"mode": "chat",
|
||||
"input_cost_per_token": 2e-6,
|
||||
"cache_read_input_token_cost": 5e-7,
|
||||
"output_cost_per_token": 6e-6,
|
||||
}
|
||||
}
|
||||
)
|
||||
usage = Usage(
|
||||
prompt_tokens=2461,
|
||||
completion_tokens=440,
|
||||
total_tokens=2901,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(
|
||||
text_tokens=1319, cached_tokens=2432, image_tokens=1142
|
||||
),
|
||||
)
|
||||
|
||||
prompt_cost, completion_cost = generic_cost_per_token(
|
||||
model=model, usage=usage, custom_llm_provider="openai"
|
||||
)
|
||||
|
||||
# 2432 cached at the cache-read rate, the 29 uncached tokens once at the input rate
|
||||
assert prompt_cost == pytest.approx(2432 * 5e-7 + 29 * 2e-6)
|
||||
assert completion_cost == pytest.approx(440 * 6e-6)
|
||||
|
||||
|
||||
def test_calculate_cost_component_with_string_values():
|
||||
"""Test the calculate_cost_component function directly with string cost values."""
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import calculate_cost_component
|
||||
|
|
|
|||
|
|
@ -8,11 +8,10 @@ See https://github.com/BerriAI/litellm/issues/26153.
|
|||
|
||||
import pytest
|
||||
|
||||
|
||||
from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import (
|
||||
StandardBuiltInToolCostTracking,
|
||||
_get_web_search_requests,
|
||||
)
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import get_web_search_requests
|
||||
from litellm.types.utils import ModelResponse, ServerToolUse, Usage
|
||||
|
||||
|
||||
|
|
@ -28,25 +27,25 @@ class _UsageWithDictServerToolUse:
|
|||
|
||||
|
||||
def test_get_web_search_requests_handles_none():
|
||||
assert _get_web_search_requests(None) is None
|
||||
assert get_web_search_requests(None) is None
|
||||
|
||||
|
||||
def test_get_web_search_requests_handles_dict():
|
||||
assert _get_web_search_requests({"web_search_requests": 5}) == 5
|
||||
assert get_web_search_requests({"web_search_requests": 5}) == 5
|
||||
|
||||
|
||||
def test_get_web_search_requests_handles_dict_missing_key():
|
||||
assert _get_web_search_requests({}) is None
|
||||
assert get_web_search_requests({}) is None
|
||||
|
||||
|
||||
def test_get_web_search_requests_handles_pydantic():
|
||||
stu = ServerToolUse(web_search_requests=7)
|
||||
assert _get_web_search_requests(stu) == 7
|
||||
assert get_web_search_requests(stu) == 7
|
||||
|
||||
|
||||
def test_get_web_search_requests_handles_pydantic_with_none_value():
|
||||
stu = ServerToolUse()
|
||||
assert _get_web_search_requests(stu) is None
|
||||
assert get_web_search_requests(stu) is None
|
||||
|
||||
|
||||
def test_response_object_includes_web_search_call_with_dict_server_tool_use():
|
||||
|
|
|
|||
|
|
@ -1002,6 +1002,17 @@ def test_an_exception_without_a_status_is_still_a_connection_error(quiet_excepti
|
|||
)
|
||||
|
||||
|
||||
def test_an_unmapped_exception_with_no_model_or_provider_is_a_connection_error(quiet_exception_mapping):
|
||||
with pytest.raises(litellm.APIConnectionError) as raised:
|
||||
exception_type(
|
||||
model=None,
|
||||
original_exception=ValueError("boom"),
|
||||
custom_llm_provider=None,
|
||||
)
|
||||
|
||||
assert "boom" in raised.value.message
|
||||
|
||||
|
||||
CONTEXT_WINDOW_MESSAGE = "This model's maximum context length is 4096 tokens."
|
||||
CONTENT_POLICY_MESSAGE = (
|
||||
'{"error": {"type": "invalid_request_error", "code": "content_policy_violation"}}'
|
||||
|
|
|
|||
|
|
@ -3997,3 +3997,98 @@ def test_translate_anthropic_messages_to_openai_carries_midturn_system_prompt_ca
|
|||
assert result == [
|
||||
{"role": "system", "content": [{"type": "text", "text": "fix", "prompt_cache_breakpoint": explicit}]}
|
||||
]
|
||||
|
||||
|
||||
def _openai_response_with_usage(usage: Usage) -> ModelResponse:
|
||||
return ModelResponse(
|
||||
id="resp_web_search",
|
||||
model="gemini-3-flash-preview",
|
||||
choices=[
|
||||
Choices(
|
||||
finish_reason="stop",
|
||||
index=0,
|
||||
message=Message(role="assistant", content="searched"),
|
||||
)
|
||||
],
|
||||
usage=usage,
|
||||
)
|
||||
|
||||
|
||||
def test_translate_openai_response_to_anthropic_maps_gemini_web_search_usage():
|
||||
from litellm.types.utils import PromptTokensDetailsWrapper
|
||||
|
||||
usage = Usage(
|
||||
prompt_tokens=385,
|
||||
completion_tokens=566,
|
||||
total_tokens=951,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=385, web_search_requests=2),
|
||||
)
|
||||
|
||||
anthropic_response = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic(
|
||||
response=_openai_response_with_usage(usage)
|
||||
)
|
||||
|
||||
assert anthropic_response["usage"]["server_tool_use"] == {"web_search_requests": 2}
|
||||
|
||||
|
||||
def test_translate_openai_response_to_anthropic_maps_server_tool_use_web_search_usage():
|
||||
from litellm.types.utils import ServerToolUse
|
||||
|
||||
usage = Usage(
|
||||
prompt_tokens=100,
|
||||
completion_tokens=40,
|
||||
total_tokens=140,
|
||||
server_tool_use=ServerToolUse(web_search_requests=3),
|
||||
)
|
||||
|
||||
anthropic_response = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic(
|
||||
response=_openai_response_with_usage(usage)
|
||||
)
|
||||
|
||||
assert anthropic_response["usage"]["server_tool_use"] == {"web_search_requests": 3}
|
||||
|
||||
|
||||
def test_translate_openai_response_to_anthropic_omits_server_tool_use_without_web_search():
|
||||
usage = Usage(prompt_tokens=100, completion_tokens=40, total_tokens=140)
|
||||
|
||||
anthropic_response = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic(
|
||||
response=_openai_response_with_usage(usage)
|
||||
)
|
||||
|
||||
assert "server_tool_use" not in anthropic_response["usage"]
|
||||
|
||||
|
||||
def test_completion_cost_on_translated_anthropic_response_includes_web_search():
|
||||
from litellm.types.utils import PromptTokensDetailsWrapper
|
||||
|
||||
adapter = LiteLLMAnthropicMessagesAdapter()
|
||||
with_search = adapter.translate_openai_response_to_anthropic(
|
||||
response=_openai_response_with_usage(
|
||||
Usage(
|
||||
prompt_tokens=385,
|
||||
completion_tokens=566,
|
||||
total_tokens=951,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=385, web_search_requests=2),
|
||||
)
|
||||
)
|
||||
)
|
||||
without_search = adapter.translate_openai_response_to_anthropic(
|
||||
response=_openai_response_with_usage(Usage(prompt_tokens=385, completion_tokens=566, total_tokens=951))
|
||||
)
|
||||
|
||||
cost_with_search = litellm.completion_cost(
|
||||
completion_response=with_search,
|
||||
model="gemini/gemini-3-flash-preview",
|
||||
call_type="anthropic_messages",
|
||||
)
|
||||
cost_without_search = litellm.completion_cost(
|
||||
completion_response=without_search,
|
||||
model="gemini/gemini-3-flash-preview",
|
||||
call_type="anthropic_messages",
|
||||
)
|
||||
|
||||
per_query_cost = litellm.model_cost["gemini/gemini-3-flash-preview"]["search_context_cost_per_query"][
|
||||
"search_context_size_medium"
|
||||
]
|
||||
assert per_query_cost > 0
|
||||
assert cost_with_search - cost_without_search == pytest.approx(2 * per_query_cost)
|
||||
|
|
|
|||
|
|
@ -17,7 +17,12 @@ from litellm.anthropic_interface import messages
|
|||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.types.utils import Delta, ModelResponse, StreamingChoices
|
||||
from litellm.types.utils import (
|
||||
Delta,
|
||||
ModelResponse,
|
||||
StandardLoggingPayloadErrorInformation,
|
||||
StreamingChoices,
|
||||
)
|
||||
|
||||
|
||||
def test_anthropic_experimental_pass_through_messages_handler():
|
||||
|
|
@ -1292,7 +1297,7 @@ class TestMessagesStreamingSuccessLogging:
|
|||
class _FailureCapture(CustomLogger):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.error_information: List[Dict[str, Any]] = []
|
||||
self.error_information: list[StandardLoggingPayloadErrorInformation] = []
|
||||
|
||||
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
||||
payload = kwargs.get("standard_logging_object") or {}
|
||||
|
|
|
|||
|
|
@ -8,11 +8,8 @@ See https://github.com/BerriAI/litellm/issues/26153.
|
|||
|
||||
import pytest
|
||||
|
||||
|
||||
from litellm.llms.anthropic.cost_calculation import (
|
||||
_get_web_search_requests,
|
||||
get_cost_for_anthropic_web_search,
|
||||
)
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import get_web_search_requests
|
||||
from litellm.llms.anthropic.cost_calculation import get_cost_for_anthropic_web_search
|
||||
from litellm.types.utils import ModelInfo, ServerToolUse
|
||||
|
||||
|
||||
|
|
@ -33,19 +30,19 @@ def _make_model_info(cost_per_query: float = 0.01) -> ModelInfo:
|
|||
|
||||
|
||||
def test_get_web_search_requests_handles_none():
|
||||
assert _get_web_search_requests(None) is None
|
||||
assert get_web_search_requests(None) is None
|
||||
|
||||
|
||||
def test_get_web_search_requests_handles_dict():
|
||||
assert _get_web_search_requests({"web_search_requests": 4}) == 4
|
||||
assert get_web_search_requests({"web_search_requests": 4}) == 4
|
||||
|
||||
|
||||
def test_get_web_search_requests_handles_dict_missing_key():
|
||||
assert _get_web_search_requests({}) is None
|
||||
assert get_web_search_requests({}) is None
|
||||
|
||||
|
||||
def test_get_web_search_requests_handles_pydantic():
|
||||
assert _get_web_search_requests(ServerToolUse(web_search_requests=2)) == 2
|
||||
assert get_web_search_requests(ServerToolUse(web_search_requests=2)) == 2
|
||||
|
||||
|
||||
def test_get_cost_for_anthropic_web_search_with_dict_server_tool_use():
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ import pytest
|
|||
|
||||
import litellm
|
||||
from litellm.llms.bedrock.base_aws_llm import Boto3CredentialsInfo
|
||||
from litellm.llms.bedrock.rerank.handler import BedrockRerankHandler
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
|
||||
# Mock response for Bedrock rerank
|
||||
|
|
@ -402,6 +403,66 @@ def test_bedrock_rerank_extra_headers_and_headers_merge():
|
|||
pytest.fail(f"Failed to merge and forward headers: {str(e)}")
|
||||
|
||||
|
||||
def test_bedrock_rerank_forwarded_headers_excluded_from_sigv4_signature():
|
||||
"""
|
||||
A forwarded header like x-forwarded-for can be rewritten between LiteLLM
|
||||
signing the request and AWS receiving it (e.g. by an intermediate load
|
||||
balancer), which invalidates the signature if that header was part of
|
||||
the signed set. It must still reach Bedrock, just unsigned.
|
||||
"""
|
||||
handler = BedrockRerankHandler()
|
||||
|
||||
prepared_request = handler._prepare_request(
|
||||
model="cohere.rerank-v3-5:0",
|
||||
api_base=None,
|
||||
extra_headers={"x-forwarded-for": "203.0.113.5"},
|
||||
data={"query": test_query, "documents": test_documents},
|
||||
optional_params={
|
||||
"aws_access_key_id": "test-access-key",
|
||||
"aws_secret_access_key": "test-secret-key",
|
||||
"aws_region_name": "us-east-1",
|
||||
},
|
||||
)
|
||||
|
||||
headers = prepared_request["prepped"].headers
|
||||
signed_headers = headers["Authorization"].split("SignedHeaders=")[1].split(",")[0].split(";")
|
||||
|
||||
assert "x-forwarded-for" not in signed_headers, (
|
||||
f"x-forwarded-for must not be part of the SigV4 signature, got SignedHeaders={signed_headers}"
|
||||
)
|
||||
assert headers["x-forwarded-for"] == "203.0.113.5", "forwarded header must still reach Bedrock, unsigned"
|
||||
|
||||
|
||||
def test_bedrock_rerank_signs_with_sigv4_even_when_bedrock_api_key_is_set(monkeypatch):
|
||||
"""
|
||||
Bedrock API keys are only valid for Bedrock and Bedrock Runtime actions, not for
|
||||
Agents for Amazon Bedrock Runtime ones. Rerank is served by bedrock-agent-runtime,
|
||||
so it has to keep signing with SigV4 even when AWS_BEARER_TOKEN_BEDROCK is set.
|
||||
"""
|
||||
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "test-bedrock-api-key")
|
||||
|
||||
handler = BedrockRerankHandler()
|
||||
|
||||
prepared_request = handler._prepare_request(
|
||||
model="cohere.rerank-v3-5:0",
|
||||
api_base=None,
|
||||
extra_headers=None,
|
||||
data={"query": test_query, "documents": test_documents},
|
||||
optional_params={
|
||||
"aws_access_key_id": "test-access-key",
|
||||
"aws_secret_access_key": "test-secret-key",
|
||||
"aws_region_name": "us-east-1",
|
||||
},
|
||||
)
|
||||
|
||||
assert prepared_request["endpoint_url"].startswith("https://bedrock-agent-runtime.")
|
||||
|
||||
authorization = prepared_request["prepped"].headers["Authorization"]
|
||||
assert authorization.startswith("AWS4-HMAC-SHA256"), (
|
||||
f"rerank must sign with SigV4, got Authorization={authorization[:30]}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bedrock_rerank_records_llm_api_duration():
|
||||
"""The bedrock rerank handler must feed httpx timing into the logging obj, so the
|
||||
|
|
|
|||
|
|
@ -84,6 +84,65 @@ def test_no_usage_details():
|
|||
assert cost == 0.0
|
||||
|
||||
|
||||
def _make_server_tool_use_usage(web_search_requests: int) -> Usage:
|
||||
from litellm.types.utils import ServerToolUse
|
||||
|
||||
return Usage(
|
||||
prompt_tokens=100,
|
||||
completion_tokens=50,
|
||||
total_tokens=150,
|
||||
server_tool_use=ServerToolUse(web_search_requests=web_search_requests),
|
||||
)
|
||||
|
||||
|
||||
def test_server_tool_use_fallback_per_query_billing():
|
||||
"""Usage reconstructed from an Anthropic-format response carries the count in
|
||||
server_tool_use, not prompt_tokens_details; per_query billing prices each request."""
|
||||
model_info = {
|
||||
"key": "gemini/gemini-3-flash-preview",
|
||||
"web_search_billing_unit": "per_query",
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_medium": 0.014,
|
||||
},
|
||||
}
|
||||
cost = cost_per_web_search_request(usage=_make_server_tool_use_usage(3), model_info=model_info)
|
||||
assert cost == pytest.approx(0.014 * 3)
|
||||
|
||||
|
||||
def test_server_tool_use_fallback_per_prompt_clamps_to_one():
|
||||
"""per_prompt billing clamps the server_tool_use count to one grounded prompt."""
|
||||
model_info = {
|
||||
"key": "gemini/gemini-2.5-flash",
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_medium": 0.035,
|
||||
},
|
||||
}
|
||||
cost = cost_per_web_search_request(usage=_make_server_tool_use_usage(4), model_info=model_info)
|
||||
assert cost == pytest.approx(0.035 * 1)
|
||||
|
||||
|
||||
def test_prompt_tokens_details_take_precedence_over_server_tool_use():
|
||||
"""The native Gemini field wins when both counts are present."""
|
||||
from litellm.types.utils import ServerToolUse
|
||||
|
||||
model_info = {
|
||||
"key": "gemini/gemini-3-flash-preview",
|
||||
"web_search_billing_unit": "per_query",
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_medium": 0.014,
|
||||
},
|
||||
}
|
||||
usage = Usage(
|
||||
prompt_tokens=100,
|
||||
completion_tokens=50,
|
||||
total_tokens=150,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(web_search_requests=2),
|
||||
server_tool_use=ServerToolUse(web_search_requests=5),
|
||||
)
|
||||
cost = cost_per_web_search_request(usage=usage, model_info=model_info)
|
||||
assert cost == pytest.approx(0.014 * 2)
|
||||
|
||||
|
||||
def _make_maps_usage(google_maps_grounding_requests: int) -> Usage:
|
||||
return Usage(
|
||||
prompt_tokens=100,
|
||||
|
|
|
|||
|
|
@ -1082,11 +1082,38 @@ class TestMCPOAuth2AuthFlow:
|
|||
# LiteLLM key should be used for auth
|
||||
mock_auth.assert_called_once()
|
||||
call_args = mock_auth.call_args
|
||||
assert call_args.kwargs["api_key"] == "sk-litellm-valid-key"
|
||||
assert call_args.kwargs["api_key"] == "Bearer sk-litellm-valid-key"
|
||||
|
||||
# OAuth2 headers should still contain the Authorization token
|
||||
assert oauth2_headers.get("Authorization") == "Bearer atlassian-oauth2-token"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"header_value",
|
||||
[b"sk-litellm-valid-key", b"Bearer sk-litellm-valid-key", b"bearer sk-litellm-valid-key"],
|
||||
)
|
||||
async def test_x_litellm_api_key_survives_bearer_only_strip(self, header_value):
|
||||
from litellm.proxy.auth.user_api_key_auth import _get_bearer_token
|
||||
|
||||
scope = {
|
||||
"type": "http",
|
||||
"method": "POST",
|
||||
"path": "/mcp/some_server",
|
||||
"headers": [(b"x-litellm-api-key", header_value)],
|
||||
}
|
||||
|
||||
async def mock_user_api_key_auth(api_key, request):
|
||||
return UserAPIKeyAuth(api_key=api_key, user_id="test-user")
|
||||
|
||||
with patch( # test-quality-ok: capturing the exact api_key handed to key validation is the regression under test
|
||||
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth",
|
||||
side_effect=mock_user_api_key_auth,
|
||||
) as mock_auth:
|
||||
auth_result, *_rest = await MCPRequestHandler.process_mcp_request(scope)
|
||||
|
||||
mock_auth.assert_called_once()
|
||||
assert _get_bearer_token(api_key=mock_auth.call_args.kwargs["api_key"]) == "sk-litellm-valid-key"
|
||||
assert auth_result.user_id == "test-user"
|
||||
|
||||
async def test_litellm_key_in_authorization_backward_compat(self):
|
||||
"""
|
||||
Backward compatibility: when only Authorization header is present
|
||||
|
|
@ -3007,7 +3034,7 @@ class TestMCPCustomHeaderName:
|
|||
# Verify the mock was called
|
||||
mock_auth.assert_called_once()
|
||||
call_args = mock_auth.call_args
|
||||
assert call_args.kwargs["api_key"] == "test-api-key"
|
||||
assert call_args.kwargs["api_key"] == "Bearer test-api-key"
|
||||
|
||||
def test_get_mcp_server_auth_headers_from_headers(self):
|
||||
"""Test _get_mcp_server_auth_headers_from_headers method"""
|
||||
|
|
@ -6254,7 +6281,7 @@ class TestMCPDcrBridgeDelegateAdmission:
|
|||
) = await MCPRequestHandler.process_mcp_request(scope)
|
||||
|
||||
mock_auth.assert_called_once()
|
||||
assert mock_auth.call_args.kwargs["api_key"] == "sk-explicit-litellm-key"
|
||||
assert mock_auth.call_args.kwargs["api_key"] == "Bearer sk-explicit-litellm-key"
|
||||
# The explicit-key arm admitted; the envelope arm never ran, so no inner token is injected.
|
||||
assert auth_result.user_id == "litellm-key-user"
|
||||
assert mcp_server_auth_headers == {}
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ through the consumer; and no path leaks the upstream token in a repr.
|
|||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
import pytest
|
||||
from pydantic import SecretStr
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import (
|
||||
|
|
@ -188,6 +189,31 @@ def test_resolve_strips_optional_bearer_scheme_before_detection():
|
|||
assert prefixed.upstream_authorization.get_secret_value() == bare.upstream_authorization.get_secret_value()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("token_type", ("bearer", "BEARER", "beArEr"))
|
||||
def test_resolve_canonicalizes_case_insensitive_bearer_token_type(token_type: str):
|
||||
keys = envelope_keys_from_master_key(_MASTER_KEY)
|
||||
grant = UpstreamTokenGrant(access_token=SecretStr(_ACCESS_TOKEN), token_type=token_type, expires_in=600)
|
||||
sealed = mint_envelope(_IDENTITY, grant, keys, _NOW)
|
||||
assert isinstance(sealed, SealedEnvelope)
|
||||
|
||||
result = resolve_bridge_envelope(sealed.token.get_secret_value(), keys, _NOW, _SERVER_ID)
|
||||
|
||||
assert isinstance(result, BridgeEnvelopeAdmitted)
|
||||
assert result.upstream_authorization.get_secret_value() == f"Bearer {_ACCESS_TOKEN}"
|
||||
|
||||
|
||||
def test_resolve_preserves_non_bearer_token_type():
|
||||
keys = envelope_keys_from_master_key(_MASTER_KEY)
|
||||
grant = UpstreamTokenGrant(access_token=SecretStr(_ACCESS_TOKEN), token_type="DPoP", expires_in=600)
|
||||
sealed = mint_envelope(_IDENTITY, grant, keys, _NOW)
|
||||
assert isinstance(sealed, SealedEnvelope)
|
||||
|
||||
result = resolve_bridge_envelope(sealed.token.get_secret_value(), keys, _NOW, _SERVER_ID)
|
||||
|
||||
assert isinstance(result, BridgeEnvelopeAdmitted)
|
||||
assert result.upstream_authorization.get_secret_value() == f"DPoP {_ACCESS_TOKEN}"
|
||||
|
||||
|
||||
def test_resolve_expired_envelope_is_invalid_not_admitted():
|
||||
keys = envelope_keys_from_master_key(_MASTER_KEY)
|
||||
token = _sealed_token(keys, now=_NOW)
|
||||
|
|
|
|||
|
|
@ -4415,6 +4415,69 @@ async def test_create_group_stamps_scim_provenance(mocker, scim_upsert_user_enab
|
|||
assert new_team_mock.call_args.kwargs["data"].metadata == {SCIM_MANAGED_TEAM_METADATA_KEY: True}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("as_pydantic", [False, True])
|
||||
async def test_create_group_applies_default_team_params(
|
||||
mocker: MockerFixture,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
scim_upsert_user_enabled: None,
|
||||
as_pydantic: bool,
|
||||
):
|
||||
"""SCIM-created teams must honor litellm_settings.default_team_params, including
|
||||
models, the same way SSO auto-created teams do."""
|
||||
import litellm
|
||||
from litellm.types.proxy.management_endpoints.ui_sso import DefaultTeamSSOParams
|
||||
|
||||
default_params = {
|
||||
"models": ["no-default-models"],
|
||||
"max_budget": 25.0,
|
||||
"budget_duration": "30d",
|
||||
"tpm_limit": 100,
|
||||
"rpm_limit": 10,
|
||||
}
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"default_team_params",
|
||||
DefaultTeamSSOParams(**default_params) if as_pydantic else default_params,
|
||||
)
|
||||
|
||||
scim_group = SCIMGroup(
|
||||
schemas=["urn:ietf:params:scim:schemas:core:2.0:Group"],
|
||||
id="defaults-group",
|
||||
displayName="Defaults.Apps",
|
||||
members=[],
|
||||
)
|
||||
|
||||
mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable into create_group
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception",
|
||||
AsyncMock(return_value=_member_resolution_prisma(mocker, users=set(), teams=set())),
|
||||
)
|
||||
new_team_mock = mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable into create_group
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.new_team",
|
||||
AsyncMock(return_value=mocker.MagicMock()),
|
||||
)
|
||||
mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable into create_group
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_team_to_scim_group",
|
||||
AsyncMock(return_value=scim_group),
|
||||
)
|
||||
mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable into create_group
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2._recompute_scim_member_roles",
|
||||
AsyncMock(),
|
||||
)
|
||||
|
||||
await create_group(group=scim_group)
|
||||
|
||||
team_request = new_team_mock.call_args.kwargs["data"]
|
||||
assert team_request.models == ["no-default-models"]
|
||||
assert team_request.max_budget == 25.0
|
||||
assert team_request.budget_duration == "30d"
|
||||
assert team_request.tpm_limit == 100
|
||||
assert team_request.rpm_limit == 10
|
||||
assert team_request.team_id == "defaults-group"
|
||||
assert team_request.team_alias == "Defaults.Apps"
|
||||
assert team_request.metadata == {SCIM_MANAGED_TEAM_METADATA_KEY: True}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_group_stamps_scim_provenance(mocker, scim_upsert_user_enabled):
|
||||
"""A PUT full sync adopts a team the identity provider now owns, and the stamp has
|
||||
|
|
|
|||
|
|
@ -1,91 +1,120 @@
|
|||
import jwt
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.management_endpoints.ui_sso import MicrosoftSSOHandler
|
||||
from litellm.proxy.management_endpoints.types import get_litellm_user_role
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
from litellm.proxy.management_endpoints.ui_sso import MicrosoftSSOHandler
|
||||
|
||||
|
||||
def _id_token(**claims) -> str:
|
||||
"""Build a signed id_token carrying the given claims."""
|
||||
payload = {
|
||||
"sub": "user123",
|
||||
"email": "user@company.com",
|
||||
"aud": "litellm-app",
|
||||
"iss": "https://login.microsoftonline.com/tenant-id/v2.0",
|
||||
"exp": 9999999999,
|
||||
**claims,
|
||||
}
|
||||
return jwt.encode(payload, "secret", algorithm="HS256")
|
||||
|
||||
|
||||
def test_extracts_proxy_admin_role_from_jwt():
|
||||
"""Ensure supported app roles like 'proxy_admin' are extracted from the id_token."""
|
||||
payload = {
|
||||
"sub": "user123",
|
||||
"email": "admin@company.com",
|
||||
"app_roles": ["proxy_admin"],
|
||||
"aud": "litellm-app",
|
||||
"iss": "https://login.microsoftonline.com/tenant-id/v2.0",
|
||||
"exp": 9999999999,
|
||||
}
|
||||
token = _id_token(app_roles=["proxy_admin"])
|
||||
|
||||
token = jwt.encode(payload, "secret", algorithm="HS256")
|
||||
roles = MicrosoftSSOHandler.get_app_roles_from_id_token(token)
|
||||
|
||||
assert roles == ["proxy_admin"]
|
||||
|
||||
|
||||
def test_maps_internal_user_role():
|
||||
"""Ensure internal_user role is correctly mapped to LitellmUserRoles."""
|
||||
payload = {
|
||||
"sub": "user456",
|
||||
"email": "user@company.com",
|
||||
"app_roles": ["internal_user"],
|
||||
"aud": "litellm-app",
|
||||
"iss": "https://login.microsoftonline.com/tenant-id/v2.0",
|
||||
"exp": 9999999999,
|
||||
}
|
||||
def test_extracts_app_roles_from_roles_claim():
|
||||
"""Entra emits app role values in the `roles` claim; both spellings are read."""
|
||||
token = _id_token(roles=["internal_user"])
|
||||
|
||||
token = jwt.encode(payload, "secret", algorithm="HS256")
|
||||
roles = MicrosoftSSOHandler.get_app_roles_from_id_token(token)
|
||||
|
||||
# Map to LitellmUserRoles
|
||||
chosen = None
|
||||
for r in roles:
|
||||
mapped = get_litellm_user_role(r)
|
||||
if mapped is not None:
|
||||
chosen = mapped
|
||||
break
|
||||
|
||||
assert chosen == LitellmUserRoles.INTERNAL_USER
|
||||
assert roles == ["internal_user"]
|
||||
|
||||
|
||||
def test_maps_proxy_admin_viewer_role():
|
||||
"""Ensure proxy_admin_viewer role is correctly mapped."""
|
||||
payload = {
|
||||
"sub": "user789",
|
||||
"email": "viewer@company.com",
|
||||
"app_roles": ["proxy_admin_viewer"],
|
||||
"aud": "litellm-app",
|
||||
"iss": "https://login.microsoftonline.com/tenant-id/v2.0",
|
||||
"exp": 9999999999,
|
||||
}
|
||||
|
||||
token = jwt.encode(payload, "secret", algorithm="HS256")
|
||||
roles = MicrosoftSSOHandler.get_app_roles_from_id_token(token)
|
||||
|
||||
chosen = None
|
||||
for r in roles:
|
||||
mapped = get_litellm_user_role(r)
|
||||
if mapped is not None:
|
||||
chosen = mapped
|
||||
break
|
||||
|
||||
assert chosen == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY
|
||||
@pytest.mark.parametrize(
|
||||
"app_roles, expected",
|
||||
[
|
||||
(["proxy_admin"], LitellmUserRoles.PROXY_ADMIN),
|
||||
(["proxy_admin_viewer"], LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY),
|
||||
(["internal_user"], LitellmUserRoles.INTERNAL_USER),
|
||||
(["internal_user_viewer"], LitellmUserRoles.INTERNAL_USER_VIEW_ONLY),
|
||||
# Case-insensitive, matching get_litellm_user_role.
|
||||
(["PROXY_ADMIN_VIEWER"], LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY),
|
||||
# Roles outside the privilege hierarchy still resolve.
|
||||
(["org_admin"], LitellmUserRoles.ORG_ADMIN),
|
||||
],
|
||||
)
|
||||
def test_maps_single_app_role(app_roles, expected):
|
||||
"""A lone app role maps to its LitellmUserRoles equivalent."""
|
||||
assert MicrosoftSSOHandler.get_user_role_from_app_roles(app_roles) == expected
|
||||
|
||||
|
||||
def test_defaults_to_internal_user_viewer_when_no_role():
|
||||
"""Ensure default role is internal_user_viewer when no app role is present."""
|
||||
payload = {
|
||||
"sub": "user_no_role",
|
||||
"email": "noRole@company.com",
|
||||
"aud": "litellm-app",
|
||||
"iss": "https://login.microsoftonline.com/tenant-id/v2.0",
|
||||
"exp": 9999999999,
|
||||
}
|
||||
@pytest.mark.parametrize(
|
||||
"app_roles",
|
||||
[
|
||||
["internal_user", "proxy_admin_viewer"],
|
||||
["proxy_admin_viewer", "internal_user"],
|
||||
],
|
||||
)
|
||||
def test_highest_privilege_role_wins_regardless_of_claim_order(app_roles):
|
||||
"""
|
||||
A user in one group mapped to `internal_user` and another mapped to
|
||||
`proxy_admin_viewer` gets the higher privilege role either way.
|
||||
|
||||
Entra does not guarantee the ordering of the `roles` claim, so the resolved
|
||||
role must not depend on it.
|
||||
"""
|
||||
assert MicrosoftSSOHandler.get_user_role_from_app_roles(app_roles) == (LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"app_roles",
|
||||
[
|
||||
["internal_user", "proxy_admin_viewer", "proxy_admin"],
|
||||
["proxy_admin", "proxy_admin_viewer", "internal_user"],
|
||||
["proxy_admin_viewer", "internal_user", "proxy_admin"],
|
||||
],
|
||||
)
|
||||
def test_proxy_admin_beats_every_other_role(app_roles):
|
||||
"""proxy_admin outranks every other role in the hierarchy, in any claim order."""
|
||||
assert MicrosoftSSOHandler.get_user_role_from_app_roles(app_roles) == LitellmUserRoles.PROXY_ADMIN
|
||||
|
||||
|
||||
def test_unrecognised_app_roles_are_ignored():
|
||||
"""App roles that are not LitellmUserRoles values do not shadow ones that are."""
|
||||
app_roles = ["Some.Custom.Role", "msiam_access", "internal_user"]
|
||||
|
||||
assert MicrosoftSSOHandler.get_user_role_from_app_roles(app_roles) == LitellmUserRoles.INTERNAL_USER
|
||||
|
||||
|
||||
@pytest.mark.parametrize("app_roles", [None, [], ["msiam_access"], ["User"]])
|
||||
def test_returns_none_when_no_role_resolves(app_roles):
|
||||
"""
|
||||
Returning None lets the caller keep the user's stored role or apply
|
||||
default_internal_user_params, rather than forcing a role.
|
||||
"""
|
||||
assert MicrosoftSSOHandler.get_user_role_from_app_roles(app_roles) is None
|
||||
|
||||
|
||||
def test_no_role_claim_yields_no_app_roles():
|
||||
"""An id_token with no role claim produces no app roles, and so no role."""
|
||||
token = _id_token()
|
||||
|
||||
token = jwt.encode(payload, "secret", algorithm="HS256")
|
||||
roles = MicrosoftSSOHandler.get_app_roles_from_id_token(token)
|
||||
|
||||
assert roles == []
|
||||
assert MicrosoftSSOHandler.get_user_role_from_app_roles(roles) is None
|
||||
|
||||
# Default role would be internal_user_viewer
|
||||
default_role = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
|
||||
assert default_role.value == "internal_user_viewer"
|
||||
|
||||
def test_end_to_end_from_id_token_to_role():
|
||||
"""The id_token -> role path resolves the highest privilege role."""
|
||||
token = _id_token(roles=["internal_user", "proxy_admin_viewer"])
|
||||
|
||||
roles = MicrosoftSSOHandler.get_app_roles_from_id_token(token)
|
||||
|
||||
assert MicrosoftSSOHandler.get_user_role_from_app_roles(roles) == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY
|
||||
|
|
|
|||
|
|
@ -79,7 +79,7 @@ async def test_delete_prompt_success():
|
|||
|
||||
# 2. Memory deletion should use base ID
|
||||
mock_registry.delete_prompts_by_base_id.assert_called_once_with(
|
||||
expected_base_id
|
||||
expected_base_id, environment=None
|
||||
)
|
||||
|
||||
assert response == {
|
||||
|
|
@ -150,7 +150,7 @@ async def test_delete_prompt_by_base_id_success():
|
|||
|
||||
# 2. Memory deletion should use base ID
|
||||
mock_registry.delete_prompts_by_base_id.assert_called_once_with(
|
||||
expected_base_id
|
||||
expected_base_id, environment=None
|
||||
)
|
||||
|
||||
assert response == {
|
||||
|
|
@ -158,6 +158,37 @@ async def test_delete_prompt_by_base_id_success():
|
|||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_prompt_environment_scope_reaches_db_and_registry():
|
||||
from litellm.proxy.prompts.prompt_endpoints import delete_prompt
|
||||
|
||||
mock_user_auth = UserAPIKeyAuth(api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_prompttable.delete_many = AsyncMock(return_value=None)
|
||||
|
||||
with patch( # test-quality-ok: stubs the collaborator so the test pins what the endpoint deletes
|
||||
"litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY"
|
||||
) as mock_registry:
|
||||
mock_registry.get_prompt_by_id.return_value = PromptSpec(
|
||||
prompt_id="test_prompt.v2",
|
||||
litellm_params=PromptLiteLLMParams(prompt_id="test_prompt", prompt_integration="dotprompt"),
|
||||
prompt_info=PromptInfo(prompt_type="db"),
|
||||
)
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): # test-quality-ok: proxy_server module global is the endpoint's only injection point
|
||||
response = await delete_prompt(
|
||||
prompt_id="test_prompt.v2",
|
||||
environment="production",
|
||||
user_api_key_dict=mock_user_auth,
|
||||
)
|
||||
|
||||
mock_prisma_client.db.litellm_prompttable.delete_many.assert_called_once_with(
|
||||
where={"prompt_id": "test_prompt", "environment": "production"}
|
||||
)
|
||||
mock_registry.delete_prompts_by_base_id.assert_called_once_with("test_prompt", environment="production")
|
||||
assert response == {"message": "Prompt test_prompt deleted successfully from production"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_prompt_info_by_base_id():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -88,3 +88,55 @@ def test_reload_prompt_keeps_the_old_template_when_the_replacement_fails(isolate
|
|||
assert registry.get_prompt_callback_by_id("greeting.v1") is old_callback
|
||||
assert _served_content(registry) == "begin every reply with AHOY"
|
||||
assert isolated_callbacks == [old_callback]
|
||||
|
||||
|
||||
def _versioned_prompt_spec(version: int, environment: str) -> PromptSpec:
|
||||
return PromptSpec(
|
||||
prompt_id=f"greeting.v{version}",
|
||||
litellm_params=PromptLiteLLMParams(
|
||||
prompt_id="greeting",
|
||||
prompt_integration="dotprompt",
|
||||
prompt_data={"content": f"begin every reply with AHOY v{version}", "metadata": {}},
|
||||
),
|
||||
prompt_info=PromptInfo(prompt_type="db", environment=environment),
|
||||
version=version,
|
||||
environment=environment,
|
||||
)
|
||||
|
||||
|
||||
def test_delete_prompts_by_base_id_removes_the_callbacks_from_litellm_callbacks(isolated_callbacks: list) -> None:
|
||||
registry = InMemoryPromptRegistry()
|
||||
registry.initialize_prompt(prompt=_versioned_prompt_spec(1, "development"))
|
||||
registry.initialize_prompt(prompt=_versioned_prompt_spec(2, "development"))
|
||||
assert len(isolated_callbacks) == 1
|
||||
|
||||
deleted = registry.delete_prompts_by_base_id("greeting")
|
||||
|
||||
assert sorted(deleted) == ["greeting.v1", "greeting.v2"]
|
||||
assert registry.get_prompt_by_id("greeting.v1") is None
|
||||
assert registry.get_prompt_callback_by_id("greeting.v2") is None
|
||||
assert isolated_callbacks == []
|
||||
|
||||
|
||||
def test_delete_prompts_by_base_id_environment_scope_keeps_other_environments(isolated_callbacks: list) -> None:
|
||||
registry = InMemoryPromptRegistry()
|
||||
registry.initialize_prompt(prompt=_versioned_prompt_spec(1, "development"))
|
||||
registry.initialize_prompt(prompt=_versioned_prompt_spec(2, "production"))
|
||||
production_callback = registry.get_prompt_callback_by_id("greeting.v2")
|
||||
|
||||
deleted = registry.delete_prompts_by_base_id("greeting", environment="development")
|
||||
|
||||
assert deleted == ["greeting.v1"]
|
||||
assert registry.get_prompt_by_id("greeting.v1") is None
|
||||
assert registry.get_prompt_by_id("greeting.v2") is not None
|
||||
assert registry.get_prompt_callback_by_id("greeting.v2") is production_callback
|
||||
|
||||
|
||||
def test_remove_prompt_is_a_no_op_for_an_unknown_id(isolated_callbacks: list) -> None:
|
||||
registry = InMemoryPromptRegistry()
|
||||
registry.initialize_prompt(prompt=_versioned_prompt_spec(1, "development"))
|
||||
|
||||
registry.remove_prompt(prompt_id="not_there.v1")
|
||||
|
||||
assert registry.get_prompt_by_id("greeting.v1") is not None
|
||||
assert len(isolated_callbacks) == 1
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import asyncio
|
||||
import collections
|
||||
import datetime
|
||||
import hashlib
|
||||
import json
|
||||
import re
|
||||
from datetime import timezone
|
||||
|
|
@ -96,6 +97,7 @@ def _reconstruct_ui_where_from_sql(sql_query, params):
|
|||
msg = re.search(r"error_message' LIKE \$(\d+)", cond)
|
||||
sess = re.fullmatch(r"session_id LIKE \$(\d+)", cond)
|
||||
status = re.fullmatch(r"status = \$(\d+)", cond)
|
||||
api_key_not_in = re.fullmatch(r"api_key NOT IN \(\$(\d+), \$(\d+)\)", cond)
|
||||
if gte:
|
||||
date_bounds["gte"] = _iso(params[int(gte.group(1)) - 1])
|
||||
elif lte:
|
||||
|
|
@ -104,10 +106,19 @@ def _reconstruct_ui_where_from_sql(sql_query, params):
|
|||
where["OR"] = where.get("OR", []) + [{"multi_team": True}]
|
||||
elif "status = 'success'" in cond:
|
||||
where["OR"] = where.get("OR", []) + [{"status": "success"}]
|
||||
elif cond == "LOWER(cache_hit) = 'true'":
|
||||
where["cache_hit"] = "hit"
|
||||
elif cond == "(cache_hit IS NULL OR LOWER(cache_hit) != 'true')":
|
||||
where["cache_hit"] = "miss"
|
||||
elif sess:
|
||||
where["session_id"] = {"contains": str(params[int(sess.group(1)) - 1]).strip("%")}
|
||||
elif status:
|
||||
where["status"] = {"equals": params[int(status.group(1)) - 1]}
|
||||
elif api_key_not_in:
|
||||
where["api_key_not_in"] = [
|
||||
params[int(api_key_not_in.group(1)) - 1],
|
||||
params[int(api_key_not_in.group(2)) - 1],
|
||||
]
|
||||
elif alias:
|
||||
metadata_conds.append(
|
||||
{
|
||||
|
|
@ -196,6 +207,7 @@ def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=No
|
|||
return MockPrismaClient()
|
||||
|
||||
|
||||
from litellm.constants import LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME
|
||||
from litellm.proxy._types import (
|
||||
LitellmUserRoles,
|
||||
Member,
|
||||
|
|
@ -1256,6 +1268,140 @@ async def test_ui_view_spend_logs_with_team_id(client, monkeypatch):
|
|||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
_HEALTH_CHECK_HASHED_API_KEY = hashlib.sha256(LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME.encode()).hexdigest()
|
||||
|
||||
|
||||
def _spend_logs_with_health_check_rows():
|
||||
now = datetime.datetime.now(timezone.utc).isoformat()
|
||||
return [
|
||||
{
|
||||
"id": "log1",
|
||||
"request_id": "req1",
|
||||
"api_key": "sk-test-key",
|
||||
"user": "test_user_1",
|
||||
"team_id": None,
|
||||
"spend": 0.05,
|
||||
"startTime": now,
|
||||
"model": "gpt-4",
|
||||
},
|
||||
{
|
||||
"id": "log2",
|
||||
"request_id": "req2",
|
||||
"api_key": _HEALTH_CHECK_HASHED_API_KEY,
|
||||
"user": None,
|
||||
"team_id": LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
|
||||
"spend": 0.0,
|
||||
"startTime": now,
|
||||
"model": "gpt-4",
|
||||
},
|
||||
{
|
||||
"id": "log3",
|
||||
"request_id": "req3",
|
||||
"api_key": LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
|
||||
"user": None,
|
||||
"team_id": LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
|
||||
"spend": 0.0,
|
||||
"startTime": now,
|
||||
"model": "gpt-4",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ui_view_spend_logs_exclude_internal_health_checks(client, monkeypatch):
|
||||
mock_spend_logs = _spend_logs_with_health_check_rows()
|
||||
|
||||
def filter_health_checks(where):
|
||||
excluded = where.get("api_key_not_in")
|
||||
if excluded is None:
|
||||
return mock_spend_logs
|
||||
return [log for log in mock_spend_logs if log["api_key"] not in excluded]
|
||||
|
||||
observed_queries = []
|
||||
|
||||
def observe_query(sql_query, params):
|
||||
if 'FROM "LiteLLM_SpendLogs"' in sql_query:
|
||||
observed_queries.append((sql_query, params))
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.prisma_client",
|
||||
make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_health_checks, query_observer=observe_query),
|
||||
)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user"
|
||||
)
|
||||
|
||||
try:
|
||||
start_date, end_date = _default_date_range()
|
||||
response = client.get(
|
||||
"/spend/logs/ui",
|
||||
params={
|
||||
"exclude_internal_health_checks": "true",
|
||||
"start_date": start_date,
|
||||
"end_date": end_date,
|
||||
},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["total"] == 1
|
||||
assert [row["request_id"] for row in data["data"]] == ["req1"]
|
||||
|
||||
page_sql, page_params = next((sql, params) for sql, params in observed_queries if "ORDER BY" in sql)
|
||||
not_in = re.search(r"api_key NOT IN \(\$(\d+), \$(\d+)\)", page_sql)
|
||||
assert not_in is not None
|
||||
assert LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME not in page_sql
|
||||
assert _HEALTH_CHECK_HASHED_API_KEY not in page_sql
|
||||
assert {
|
||||
page_params[int(not_in.group(1)) - 1],
|
||||
page_params[int(not_in.group(2)) - 1],
|
||||
} == {LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, _HEALTH_CHECK_HASHED_API_KEY}
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ui_view_spend_logs_includes_internal_health_checks_by_default(client, monkeypatch):
|
||||
mock_spend_logs = _spend_logs_with_health_check_rows()
|
||||
|
||||
def filter_health_checks(where):
|
||||
excluded = where.get("api_key_not_in")
|
||||
if excluded is None:
|
||||
return mock_spend_logs
|
||||
return [log for log in mock_spend_logs if log["api_key"] not in excluded]
|
||||
|
||||
observed_queries = []
|
||||
|
||||
def observe_query(sql_query, params):
|
||||
if 'FROM "LiteLLM_SpendLogs"' in sql_query:
|
||||
observed_queries.append((sql_query, params))
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.prisma_client",
|
||||
make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_health_checks, query_observer=observe_query),
|
||||
)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user"
|
||||
)
|
||||
|
||||
try:
|
||||
start_date, end_date = _default_date_range()
|
||||
response = client.get(
|
||||
"/spend/logs/ui",
|
||||
params={"start_date": start_date, "end_date": end_date},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["total"] == 3
|
||||
assert [row["request_id"] for row in data["data"]] == ["req1", "req2", "req3"]
|
||||
assert all("NOT IN" not in sql for sql, _ in observed_queries)
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ui_view_spend_logs_internal_user_scoped_without_user_id(
|
||||
client, monkeypatch
|
||||
|
|
@ -2302,6 +2448,96 @@ async def test_ui_view_spend_logs_with_status(client, monkeypatch):
|
|||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ui_view_spend_logs_with_cache_hit_filter(client, monkeypatch):
|
||||
base = {
|
||||
"api_key": "sk-test-key",
|
||||
"user": "test_user_1",
|
||||
"team_id": "team1",
|
||||
"spend": 0.05,
|
||||
"startTime": datetime.datetime.now(timezone.utc).isoformat(),
|
||||
"model": "gpt-4",
|
||||
"status": "success",
|
||||
}
|
||||
mock_spend_logs = [
|
||||
{**base, "id": "log1", "request_id": "req-hit", "cache_hit": "True"},
|
||||
{**base, "id": "log2", "request_id": "req-miss", "cache_hit": "False"},
|
||||
{**base, "id": "log3", "request_id": "req-legacy", "cache_hit": "None"},
|
||||
{**base, "id": "log4", "request_id": "req-null", "cache_hit": None},
|
||||
]
|
||||
|
||||
def filter_by_cache(where):
|
||||
cache_filter = where.get("cache_hit")
|
||||
if cache_filter == "hit":
|
||||
return [log for log in mock_spend_logs if str(log["cache_hit"]).lower() == "true"]
|
||||
if cache_filter == "miss":
|
||||
return [log for log in mock_spend_logs if str(log["cache_hit"]).lower() != "true"]
|
||||
return mock_spend_logs
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.prisma_client",
|
||||
make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_by_cache),
|
||||
)
|
||||
|
||||
start_date, end_date = _default_date_range()
|
||||
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN
|
||||
)
|
||||
try:
|
||||
response = client.get(
|
||||
"/spend/logs/ui",
|
||||
params={
|
||||
"cache_hit_filter": "hit",
|
||||
"start_date": start_date,
|
||||
"end_date": end_date,
|
||||
},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["total"] == 1
|
||||
assert [row["request_id"] for row in data["data"]] == ["req-hit"]
|
||||
|
||||
response = client.get(
|
||||
"/spend/logs/ui",
|
||||
params={
|
||||
"cache_hit_filter": "miss",
|
||||
"start_date": start_date,
|
||||
"end_date": end_date,
|
||||
},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["total"] == 3
|
||||
assert [row["request_id"] for row in data["data"]] == ["req-miss", "req-legacy", "req-null"]
|
||||
|
||||
response = client.get(
|
||||
"/spend/logs/ui",
|
||||
params={
|
||||
"start_date": start_date,
|
||||
"end_date": end_date,
|
||||
},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert response.json()["total"] == 4
|
||||
|
||||
response = client.get(
|
||||
"/spend/logs/ui",
|
||||
params={
|
||||
"cache_hit_filter": "invalid",
|
||||
"start_date": start_date,
|
||||
"end_date": end_date,
|
||||
},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
assert response.status_code == 400
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ui_view_spend_logs_with_model(client, monkeypatch):
|
||||
mock_spend_logs = [
|
||||
|
|
@ -3135,6 +3371,90 @@ async def test_view_spend_logs_summarize_parameter(client, monkeypatch):
|
|||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_view_spend_logs_bounds_row_count(client, monkeypatch):
|
||||
"""Every /spend/logs read path must send take=SPEND_LOGS_PAGINATION_COUNT_CAP to Prisma (LIT-6284)."""
|
||||
captured_find_many_kwargs = []
|
||||
|
||||
class MockDB:
|
||||
def __init__(self):
|
||||
self.litellm_spendlogs = self
|
||||
self.available_rows = 0
|
||||
|
||||
async def find_many(self, *args, **kwargs):
|
||||
captured_find_many_kwargs.append(kwargs)
|
||||
return [{}] * min(kwargs.get("take", 0), self.available_rows)
|
||||
|
||||
class MockPrismaClient:
|
||||
def __init__(self):
|
||||
self.db = MockDB()
|
||||
|
||||
def hash_token(self, token):
|
||||
return f"hashed-{token}"
|
||||
|
||||
mock_prisma_client = MockPrismaClient()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN
|
||||
)
|
||||
start_date = (
|
||||
datetime.datetime.now(timezone.utc) - datetime.timedelta(days=2)
|
||||
).strftime("%Y-%m-%d")
|
||||
end_date = datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d")
|
||||
try:
|
||||
response = client.get(
|
||||
"/spend/logs",
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert (
|
||||
captured_find_many_kwargs[-1].get("take")
|
||||
== spend_management_endpoints.SPEND_LOGS_PAGINATION_COUNT_CAP
|
||||
)
|
||||
assert "x-litellm-spend-logs-truncated" not in response.headers
|
||||
|
||||
response = client.get(
|
||||
"/spend/logs",
|
||||
params={"user_id": "test-user"},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert captured_find_many_kwargs[-1].get("where") == {"user": "test-user"}
|
||||
assert (
|
||||
captured_find_many_kwargs[-1].get("take")
|
||||
== spend_management_endpoints.SPEND_LOGS_PAGINATION_COUNT_CAP
|
||||
)
|
||||
|
||||
response = client.get(
|
||||
"/spend/logs",
|
||||
params={
|
||||
"start_date": start_date,
|
||||
"end_date": end_date,
|
||||
"summarize": "false",
|
||||
},
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert "startTime" in captured_find_many_kwargs[-1].get("where", {})
|
||||
assert (
|
||||
captured_find_many_kwargs[-1].get("take")
|
||||
== spend_management_endpoints.SPEND_LOGS_PAGINATION_COUNT_CAP
|
||||
)
|
||||
|
||||
mock_prisma_client.db.available_rows = (
|
||||
spend_management_endpoints.SPEND_LOGS_PAGINATION_COUNT_CAP
|
||||
)
|
||||
response = client.get(
|
||||
"/spend/logs",
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert len(response.json()) == spend_management_endpoints.SPEND_LOGS_PAGINATION_COUNT_CAP
|
||||
assert response.headers["x-litellm-spend-logs-truncated"] == "true"
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_view_spend_tags(client, monkeypatch):
|
||||
"""Test the /spend/tags endpoint"""
|
||||
|
|
|
|||
|
|
@ -1,8 +1,9 @@
|
|||
import json
|
||||
import sys
|
||||
from types import ModuleType, SimpleNamespace
|
||||
|
||||
from litellm.proxy._lazy_features import LazyFeature
|
||||
from litellm.proxy._lazy_openapi_snapshot import _normalize_operation_ids
|
||||
from litellm.proxy._lazy_openapi_snapshot import SnapshotResult, _normalize_operation_ids, main
|
||||
|
||||
|
||||
def test_generate_snapshot_uses_shared_operation_id_reservations(monkeypatch):
|
||||
|
|
@ -61,7 +62,7 @@ def test_generate_snapshot_uses_shared_operation_id_reservations(monkeypatch):
|
|||
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy_server_module)
|
||||
monkeypatch.setattr("fastapi.openapi.utils.get_openapi", fake_get_openapi)
|
||||
|
||||
fragments = _lazy_openapi_snapshot.generate_snapshot()
|
||||
fragments = _lazy_openapi_snapshot.generate_snapshot().fragments
|
||||
|
||||
assert fragments["feature-a"]["paths"]["/feature-a/items"]["get"]["operationId"] == "shared_operation_id_get"
|
||||
assert fragments["feature-b"]["paths"]["/feature-b/items"]["get"]["operationId"] == "shared_operation_id_get_2"
|
||||
|
|
@ -106,7 +107,7 @@ def test_generate_snapshot_registers_transitively_imported_modules(monkeypatch):
|
|||
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy_server_module)
|
||||
monkeypatch.setattr("fastapi.openapi.utils.get_openapi", fake_get_openapi)
|
||||
|
||||
fragments = _lazy_openapi_snapshot.generate_snapshot()
|
||||
fragments = _lazy_openapi_snapshot.generate_snapshot().fragments
|
||||
|
||||
assert fragments["transitive"]["paths"]["/transitive/items"]["get"]["tags"] == ["transitive"]
|
||||
assert "/v1/{param}/deep/leaf" in fragments["transitive"]["paths"]
|
||||
|
|
@ -144,3 +145,66 @@ def test_normalize_operation_ids_preserves_custom_ids():
|
|||
operations = paths["/proxy/{endpoint}"]
|
||||
assert operations["get"]["operationId"] == "custom_operation"
|
||||
assert operations["post"]["operationId"] == "custom_operation"
|
||||
|
||||
|
||||
def test_generate_snapshot_reports_features_whose_import_fails(monkeypatch):
|
||||
from litellm.proxy import _lazy_openapi_snapshot
|
||||
|
||||
fake_app = SimpleNamespace(title="LiteLLM test", version="0.0.0", routes=[])
|
||||
|
||||
fake_module = ModuleType("fake_importable_feature")
|
||||
monkeypatch.setitem(sys.modules, "fake_importable_feature", fake_module)
|
||||
|
||||
def register_fn(app, module):
|
||||
app.routes.append(SimpleNamespace(path="/importable/items"))
|
||||
|
||||
fake_lazy_features_module = ModuleType("litellm.proxy._lazy_features")
|
||||
fake_lazy_features_module.LAZY_FEATURES = [
|
||||
LazyFeature(
|
||||
name="importable",
|
||||
module_path="fake_importable_feature",
|
||||
path_prefixes=("/importable",),
|
||||
register_fn=register_fn,
|
||||
),
|
||||
LazyFeature(
|
||||
name="broken",
|
||||
module_path="litellm.proxy.this_module_does_not_exist",
|
||||
path_prefixes=("/broken",),
|
||||
),
|
||||
]
|
||||
monkeypatch.setitem(sys.modules, "litellm.proxy._lazy_features", fake_lazy_features_module)
|
||||
|
||||
def fake_get_openapi(title, version, routes):
|
||||
return {"paths": {route.path: {"get": {"operationId": "importable_get"}} for route in routes}}
|
||||
|
||||
fake_proxy_server_module = ModuleType("litellm.proxy.proxy_server")
|
||||
fake_proxy_server_module.app = fake_app
|
||||
fake_proxy_server_module.ensure_unique_openapi_operation_ids = lambda schema, reserved_operation_ids: schema
|
||||
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy_server_module)
|
||||
monkeypatch.setattr("fastapi.openapi.utils.get_openapi", fake_get_openapi)
|
||||
|
||||
result = _lazy_openapi_snapshot.generate_snapshot()
|
||||
|
||||
assert result.skipped == ("broken",)
|
||||
assert sorted(result.fragments) == ["importable"]
|
||||
|
||||
|
||||
def test_main_refuses_to_write_a_snapshot_missing_skipped_features(tmp_path, capsys):
|
||||
snapshot_file = tmp_path / "snapshot.json"
|
||||
result = SnapshotResult(fragments={"importable": {"paths": {}, "components": {"schemas": {}}}}, skipped=("broken",))
|
||||
|
||||
assert main(snapshot_file, generate=lambda: result) == 1
|
||||
assert not snapshot_file.exists()
|
||||
assert "broken" in capsys.readouterr().err
|
||||
|
||||
|
||||
def test_main_writes_sorted_snapshot_when_every_feature_loads(tmp_path):
|
||||
snapshot_file = tmp_path / "snapshot.json"
|
||||
fragments = {
|
||||
"zeta": {"paths": {"/z": {}}, "components": {"schemas": {}}},
|
||||
"alpha": {"paths": {}, "components": {"schemas": {}}},
|
||||
}
|
||||
|
||||
assert main(snapshot_file, generate=lambda: SnapshotResult(fragments=fragments, skipped=())) == 0
|
||||
assert json.loads(snapshot_file.read_text()) == fragments
|
||||
assert snapshot_file.read_text() == json.dumps(fragments, indent=2, sort_keys=True) + "\n"
|
||||
|
|
|
|||
|
|
@ -11492,6 +11492,150 @@ async def test_init_prompts_in_db_serves_the_newest_row_when_environments_collid
|
|||
IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("greeting_env")
|
||||
|
||||
|
||||
def _prompt_db_row(prompt_id: str, litellm_params: str) -> MagicMock:
|
||||
row = MagicMock()
|
||||
row.model_dump.return_value = {
|
||||
"prompt_id": prompt_id,
|
||||
"version": 1,
|
||||
"environment": "development",
|
||||
"created_by": None,
|
||||
"litellm_params": litellm_params,
|
||||
"prompt_info": json.dumps({"prompt_type": "db"}),
|
||||
"created_at": None,
|
||||
"updated_at": None,
|
||||
}
|
||||
return row
|
||||
|
||||
|
||||
def _dotprompt_params(prompt_id: str) -> str:
|
||||
return json.dumps(
|
||||
{
|
||||
"prompt_id": prompt_id,
|
||||
"prompt_integration": "dotprompt",
|
||||
"prompt_data": {"content": "Begin every reply with AHOY", "metadata": {}},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_init_prompts_in_db_unloads_rows_deleted_on_another_worker(monkeypatch):
|
||||
from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
|
||||
prisma_client = MagicMock()
|
||||
try:
|
||||
prisma_client.db.litellm_prompttable.find_many = AsyncMock(
|
||||
return_value=[_prompt_db_row("greeting_del", _dotprompt_params("greeting_del"))]
|
||||
)
|
||||
await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client)
|
||||
assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_del.v1") is not None
|
||||
|
||||
prisma_client.db.litellm_prompttable.find_many = AsyncMock(return_value=[])
|
||||
await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client)
|
||||
|
||||
assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id("greeting_del.v1") is None
|
||||
assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_del.v1") is None
|
||||
assert litellm.callbacks == []
|
||||
finally:
|
||||
IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("greeting_del")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_init_prompts_in_db_keeps_config_prompts_when_their_id_has_no_db_row(monkeypatch):
|
||||
from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
from litellm.types.prompts.init_prompts import PromptInfo, PromptLiteLLMParams, PromptSpec
|
||||
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
|
||||
config_prompt = PromptSpec(
|
||||
prompt_id="greeting_cfg",
|
||||
litellm_params=PromptLiteLLMParams(
|
||||
prompt_id="greeting_cfg",
|
||||
prompt_integration="dotprompt",
|
||||
prompt_data={"content": "Begin every reply with AHOY", "metadata": {}},
|
||||
),
|
||||
prompt_info=PromptInfo(prompt_type="config"),
|
||||
)
|
||||
|
||||
prisma_client = MagicMock()
|
||||
try:
|
||||
IN_MEMORY_PROMPT_REGISTRY.initialize_prompt(prompt=config_prompt)
|
||||
prisma_client.db.litellm_prompttable.find_many = AsyncMock(return_value=[])
|
||||
|
||||
await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client)
|
||||
|
||||
assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_cfg") is not None
|
||||
assert len(litellm.callbacks) == 1
|
||||
finally:
|
||||
IN_MEMORY_PROMPT_REGISTRY.remove_prompt(prompt_id="greeting_cfg")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_init_prompts_in_db_keeps_the_in_memory_copy_when_a_row_fails_to_parse(monkeypatch):
|
||||
from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
|
||||
prisma_client = MagicMock()
|
||||
try:
|
||||
prisma_client.db.litellm_prompttable.find_many = AsyncMock(
|
||||
return_value=[_prompt_db_row("greeting_broken", _dotprompt_params("greeting_broken"))]
|
||||
)
|
||||
await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client)
|
||||
loaded_callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_broken.v1")
|
||||
assert loaded_callback is not None
|
||||
|
||||
prisma_client.db.litellm_prompttable.find_many = AsyncMock(
|
||||
return_value=[_prompt_db_row("greeting_broken", "this is not json")]
|
||||
)
|
||||
await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client)
|
||||
|
||||
assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_broken.v1") is loaded_callback
|
||||
assert litellm.callbacks == [loaded_callback]
|
||||
finally:
|
||||
IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("greeting_broken")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_init_prompts_in_db_keeps_a_prompt_created_while_the_sync_was_reading(monkeypatch):
|
||||
from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
from litellm.types.prompts.init_prompts import PromptInfo, PromptLiteLLMParams, PromptSpec
|
||||
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
|
||||
prisma_client = MagicMock()
|
||||
try:
|
||||
|
||||
async def create_prompt_behind_the_select() -> list:
|
||||
IN_MEMORY_PROMPT_REGISTRY.initialize_prompt(
|
||||
prompt=PromptSpec(
|
||||
prompt_id="greeting_race.v1",
|
||||
litellm_params=PromptLiteLLMParams(
|
||||
prompt_id="greeting_race",
|
||||
prompt_integration="dotprompt",
|
||||
prompt_data={"content": "Begin every reply with AHOY", "metadata": {}},
|
||||
),
|
||||
prompt_info=PromptInfo(prompt_type="db"),
|
||||
)
|
||||
)
|
||||
return []
|
||||
|
||||
prisma_client.db.litellm_prompttable.find_many = AsyncMock(side_effect=create_prompt_behind_the_select)
|
||||
await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client)
|
||||
|
||||
surviving_callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_race.v1")
|
||||
assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id("greeting_race.v1") is not None
|
||||
assert surviving_callback is not None
|
||||
assert litellm.callbacks == [surviving_callback]
|
||||
finally:
|
||||
IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("greeting_race")
|
||||
|
||||
|
||||
class TestEmbeddingsFailureHookRequestData:
|
||||
@pytest.mark.asyncio
|
||||
async def test_failure_hook_gets_post_setup_data_with_logging_obj(self):
|
||||
|
|
|
|||
|
|
@ -1810,9 +1810,7 @@ class TestLLMClassifier:
|
|||
"request_kwargs",
|
||||
[
|
||||
pytest.param({"metadata": {"user_api_key": "sk-abc"}}, id="metadata-bucket"),
|
||||
pytest.param(
|
||||
{"litellm_metadata": {"user_api_key": "sk-abc"}}, id="litellm-metadata-bucket"
|
||||
),
|
||||
pytest.param({"litellm_metadata": {"user_api_key": "sk-abc"}}, id="litellm-metadata-bucket"),
|
||||
pytest.param({}, id="no-caller-context"),
|
||||
pytest.param(None, id="no-request-kwargs"),
|
||||
],
|
||||
|
|
@ -6044,7 +6042,8 @@ class TestContextAwareClassifier:
|
|||
turn = (
|
||||
"We run a multi-region gateway and last night the eu-west pod returned 502s on the "
|
||||
"streaming path only, for thirty minutes, while non-streaming stayed healthy the whole "
|
||||
"window and the cooldown map was mid-failover. " + "Filler sentence to push past the cap. " * 4
|
||||
"window and the cooldown map was mid-failover. "
|
||||
+ "Filler sentence to push past the cap. " * 4
|
||||
+ "Now rewrite the streaming retry path and prove it cannot livelock."
|
||||
)
|
||||
|
||||
|
|
@ -8889,3 +8888,210 @@ async def test_session_pin_survives_json_list_round_trip(mock_router_instance):
|
|||
assert response.model == "shared"
|
||||
assert response.litellm_params == {"reasoning_effort": "low"}
|
||||
assert cache.async_set_cache.call_args.kwargs["value"] == {"model": "shared", "tier": "SIMPLE"}
|
||||
|
||||
|
||||
HEURISTIC_FIRST_TIERS: dict[str, str] = {
|
||||
"SIMPLE": "gpt-4o-mini",
|
||||
"MEDIUM": "gpt-4o",
|
||||
"COMPLEX": "claude-sonnet-4-20250514",
|
||||
"REASONING": "o1-preview",
|
||||
}
|
||||
|
||||
# The scorer maps a weighted score to a tier against these, and PR #37910 is retuning the shipped
|
||||
# defaults, so every heuristic_first test pins them rather than inheriting DEFAULT_TIER_BOUNDARIES.
|
||||
HEURISTIC_FIRST_BOUNDARIES: dict[str, float] = {
|
||||
"simple_medium": 0.15,
|
||||
"medium_complex": 0.35,
|
||||
"complex_reasoning": 0.60,
|
||||
}
|
||||
|
||||
# Scores 0.0 with an empty signals tuple: no dimension fires, so the scorer has no opinion and the
|
||||
# score-to-tier mapping lands SIMPLE purely by default. This is the population the permutation
|
||||
# control measured at ~zero information, and the prompt that must always escalate.
|
||||
NO_SIGNAL_PROMPT = (
|
||||
"A distributed ledger must guarantee linearizability across five regions while tolerating one "
|
||||
"region partition and bounded clock skew. Derive the minimum quorum configuration and prove why "
|
||||
"a smaller quorum violates linearizability."
|
||||
)
|
||||
|
||||
|
||||
def _heuristic_first_router(mock_router_instance, **config_overrides):
|
||||
config = {
|
||||
"tiers": dict(HEURISTIC_FIRST_TIERS),
|
||||
"tier_boundaries": dict(HEURISTIC_FIRST_BOUNDARIES),
|
||||
"classifier_type": "heuristic_first",
|
||||
"heuristic_first_max_tier": "SIMPLE",
|
||||
"classifier_llm_config": {"model": "haiku-classifier", "timeout_ms": 400},
|
||||
**config_overrides,
|
||||
}
|
||||
return ComplexityRouter(
|
||||
model_name="test-complexity-router",
|
||||
litellm_router_instance=mock_router_instance,
|
||||
complexity_router_config=config,
|
||||
)
|
||||
|
||||
|
||||
class TestHeuristicFirstConfig:
|
||||
"""Config validation for classifier_type='heuristic_first'."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"overrides, expected",
|
||||
[
|
||||
({"classifier_llm_config": None}, "classifier_llm_config is required"),
|
||||
({"heuristic_first_max_tier": None}, "heuristic_first_max_tier is required"),
|
||||
({"heuristic_first_max_tier": "REASONING"}, "is the highest tier"),
|
||||
({"heuristic_first_max_tier": "NOPE"}, "is not an active tier"),
|
||||
(
|
||||
{
|
||||
"tiers": {"SIMPLE": "gpt-4o-mini", "COMPLEX": "c", "REASONING": "r"},
|
||||
"heuristic_first_max_tier": "MEDIUM",
|
||||
},
|
||||
"has no model configured in tiers",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_rejects_incoherent_config(self, overrides, expected):
|
||||
config = {
|
||||
"tiers": dict(HEURISTIC_FIRST_TIERS),
|
||||
"classifier_type": "heuristic_first",
|
||||
"heuristic_first_max_tier": "SIMPLE",
|
||||
"classifier_llm_config": {"model": "haiku-classifier"},
|
||||
**overrides,
|
||||
}
|
||||
with pytest.raises(ValidationError, match=expected):
|
||||
ComplexityRouterConfig(**config)
|
||||
|
||||
@pytest.mark.parametrize("classifier_type", ["heuristic", "llm", "custom"])
|
||||
def test_threshold_rejected_on_every_other_classifier_type(self, classifier_type):
|
||||
"""A threshold on a router with no heuristic gate is a silent no-op, so it is refused
|
||||
rather than accepted and ignored."""
|
||||
config: dict[str, object] = {
|
||||
"tiers": dict(HEURISTIC_FIRST_TIERS),
|
||||
"classifier_type": classifier_type,
|
||||
"heuristic_first_max_tier": "SIMPLE",
|
||||
}
|
||||
if classifier_type == "llm":
|
||||
config["classifier_llm_config"] = {"model": "haiku-classifier"}
|
||||
if classifier_type == "custom":
|
||||
config["classifier_plugin"] = _FixedTierClassifier("SIMPLE")
|
||||
with pytest.raises(ValidationError, match="heuristic_first_max_tier is set but classifier_type"):
|
||||
ComplexityRouterConfig(**config)
|
||||
|
||||
def test_custom_tier_set_is_rejected(self):
|
||||
"""The scorer only emits the four built-in tiers, so it cannot gate a replaced tier set."""
|
||||
with pytest.raises(ValidationError, match="tier_definitions requires classifier_type"):
|
||||
ComplexityRouterConfig(
|
||||
classifier_type="heuristic_first",
|
||||
heuristic_first_max_tier="lo",
|
||||
classifier_llm_config={"model": "haiku-classifier"},
|
||||
tier_definitions=[{"name": "lo", "description": "x"}, {"name": "hi", "description": "y"}],
|
||||
tiers={"lo": "gpt-4o-mini", "hi": "gpt-4o"},
|
||||
)
|
||||
|
||||
def test_classifier_model_is_a_dependency(self):
|
||||
"""uses_llm_classifier is what tells the health graph and the routing-test authorizer that
|
||||
the classifier model is really called, so heuristic_first must answer True."""
|
||||
config = ComplexityRouterConfig(
|
||||
tiers=dict(HEURISTIC_FIRST_TIERS),
|
||||
classifier_type="heuristic_first",
|
||||
heuristic_first_max_tier="SIMPLE",
|
||||
classifier_llm_config={"model": "haiku-classifier"},
|
||||
)
|
||||
assert config.uses_llm_classifier is True
|
||||
assert ComplexityRouterConfig(tiers=dict(HEURISTIC_FIRST_TIERS)).uses_llm_classifier is False
|
||||
|
||||
|
||||
class TestHeuristicFirst:
|
||||
"""Behavior of the heuristic-first chain: when the classifier call is skipped, and when it is not."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_signalled_cheap_prompt_short_circuits(self, mock_router_instance):
|
||||
"""A prompt the scorer actually placed at or below the threshold must not reach the LLM."""
|
||||
mock_router_instance.acompletion = AsyncMock()
|
||||
router = _heuristic_first_router(mock_router_instance)
|
||||
outcome = await router.aclassify("thanks so much, appreciate it")
|
||||
mock_router_instance.acompletion.assert_not_called()
|
||||
assert outcome.tier == ComplexityTier.SIMPLE
|
||||
assert outcome.cause == "heuristic_first_short_circuit"
|
||||
assert outcome.score is not None
|
||||
assert outcome.signals
|
||||
assert outcome.classifier_cost is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_signal_prompt_escalates_even_though_it_scores_simple(self, mock_router_instance):
|
||||
"""The core guard. This prompt scores 0.0 and the mapping calls it SIMPLE, which is at the
|
||||
threshold, so a bare tier comparison would short-circuit it to the cheapest model. No
|
||||
dimension fired, so the scorer has no opinion and the classifier must decide."""
|
||||
mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "COMPLEX"}'))
|
||||
router = _heuristic_first_router(mock_router_instance)
|
||||
|
||||
tier, score, signals, _cause = router._score_and_classify(NO_SIGNAL_PROMPT)
|
||||
assert (tier, score, signals) == (ComplexityTier.SIMPLE, 0.0, ())
|
||||
|
||||
outcome = await router.aclassify(NO_SIGNAL_PROMPT)
|
||||
mock_router_instance.acompletion.assert_awaited_once()
|
||||
assert outcome.tier == ComplexityTier.COMPLEX
|
||||
assert outcome.cause == "llm_classifier"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_signalled_prompt_above_threshold_escalates(self, mock_router_instance):
|
||||
"""The scorer had an opinion, but it was above the threshold, so the classifier decides."""
|
||||
mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "REASONING"}'))
|
||||
router = _heuristic_first_router(mock_router_instance)
|
||||
|
||||
tier, _score, signals, _cause = router._score_and_classify("write a python function to reverse a string")
|
||||
assert tier == ComplexityTier.MEDIUM and signals
|
||||
|
||||
outcome = await router.aclassify("write a python function to reverse a string")
|
||||
mock_router_instance.acompletion.assert_awaited_once()
|
||||
assert outcome.tier == ComplexityTier.REASONING
|
||||
assert outcome.cause == "llm_classifier"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raising_threshold_short_circuits_what_it_previously_escalated(self, mock_router_instance):
|
||||
"""The threshold is the knob: the same signalled MEDIUM prompt escalates at SIMPLE and
|
||||
short-circuits at MEDIUM."""
|
||||
mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "REASONING"}'))
|
||||
router = _heuristic_first_router(mock_router_instance, heuristic_first_max_tier="MEDIUM")
|
||||
outcome = await router.aclassify("write a python function to reverse a string")
|
||||
mock_router_instance.acompletion.assert_not_called()
|
||||
assert outcome.tier == ComplexityTier.MEDIUM
|
||||
assert outcome.cause == "heuristic_first_short_circuit"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reasoning_override_never_short_circuits(self, mock_router_instance):
|
||||
"""A reasoning-override prompt lands REASONING, which outranks every legal threshold, so it
|
||||
always reaches the classifier."""
|
||||
mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "MEDIUM"}'))
|
||||
router = _heuristic_first_router(mock_router_instance, heuristic_first_max_tier="COMPLEX")
|
||||
outcome = await router.aclassify(
|
||||
"think step by step and analyze the tradeoffs, then reason through the consequences carefully"
|
||||
)
|
||||
mock_router_instance.acompletion.assert_awaited_once()
|
||||
assert outcome.cause == "llm_classifier"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_classifier_failure_falls_back_to_the_scorer(self, mock_router_instance):
|
||||
"""An escalated request whose classifier call fails still gets the scorer's own verdict,
|
||||
the same way classifier_type='llm' does, rather than erroring out."""
|
||||
mock_router_instance.acompletion = AsyncMock(side_effect=RuntimeError("classifier exploded"))
|
||||
router = _heuristic_first_router(mock_router_instance)
|
||||
expected_tier, expected_score, expected_signals, _cause = router._score_and_classify(NO_SIGNAL_PROMPT)
|
||||
|
||||
outcome = await router.aclassify(NO_SIGNAL_PROMPT)
|
||||
|
||||
assert outcome.tier == expected_tier
|
||||
assert outcome.score == expected_score
|
||||
assert outcome.signals == expected_signals
|
||||
assert outcome.cause == "heuristic_scorer"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_classifier_failure_honors_default_model_fallback(self, mock_router_instance):
|
||||
"""classifier_fallback='default_model' still wins over the heuristic outcome, same as it
|
||||
does for classifier_type='llm'."""
|
||||
mock_router_instance.acompletion = AsyncMock(side_effect=RuntimeError("classifier exploded"))
|
||||
router = _heuristic_first_router(
|
||||
mock_router_instance, classifier_fallback="default_model", default_model="gpt-4o"
|
||||
)
|
||||
outcome = await router.aclassify(NO_SIGNAL_PROMPT)
|
||||
assert outcome.cause == "default_model_fallback"
|
||||
|
|
|
|||
|
|
@ -27,7 +27,7 @@
|
|||
"limit": 0
|
||||
},
|
||||
"LIT010": {
|
||||
"limit": 16619
|
||||
"limit": 16616
|
||||
},
|
||||
"LIT011": {
|
||||
"limit": 5583
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
|
||||
import { fireEvent, render, screen } from "@testing-library/react";
|
||||
import { fireEvent, render, screen, within } from "@testing-library/react";
|
||||
import React from "react";
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
|
|
@ -151,15 +151,18 @@ describe("AutoRouterBenchmarksTab", () => {
|
|||
mockAutoRouters();
|
||||
});
|
||||
|
||||
it("leads with total estimated savings, before the three session-shape metrics", () => {
|
||||
it("leads with total estimated savings, before the four session-shape metrics", () => {
|
||||
mockHook({ data: response([group(), group({ router_name: "gpt-auto" })]) });
|
||||
renderTab();
|
||||
|
||||
const labels = screen
|
||||
.getAllByText(/Total estimated savings|Avg turns per session|Avg session length|Avg tokens per session/)
|
||||
.getAllByText(
|
||||
/Total estimated savings|Avg saved per session|Avg turns per session|Avg session length|Avg tokens per session/,
|
||||
)
|
||||
.map((node) => node.textContent);
|
||||
expect(labels).toEqual([
|
||||
"Total estimated savings",
|
||||
"Avg saved per session",
|
||||
"Avg turns per session",
|
||||
"Avg session length",
|
||||
"Avg tokens per session",
|
||||
|
|
@ -181,13 +184,35 @@ describe("AutoRouterBenchmarksTab", () => {
|
|||
expect(screen.getByText("5.3M")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("pairs the savings with the session count it was earned over", () => {
|
||||
it("pairs the savings with the session count it was earned over, in its own tile", () => {
|
||||
mockHook({ data: response([group(), group({ router_name: "gpt-auto" })]) });
|
||||
renderTab();
|
||||
|
||||
expect(screen.getByText("Avg saved per session")).toBeInTheDocument();
|
||||
expect(screen.getByText("$23.13")).toBeInTheDocument();
|
||||
expect(screen.getByText("across 94 sessions")).toBeInTheDocument();
|
||||
const tile = screen.getByText("Avg saved per session").closest('[data-slot="card"]');
|
||||
if (!tile) throw new Error("expected avg saved per session to render as a metric tile");
|
||||
|
||||
expect(within(tile).getByText("$23.13")).toBeInTheDocument();
|
||||
expect(within(tile).getByText("· 94 sessions")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("exposes each spend row as a term and its value, not as loose text", () => {
|
||||
mockHook({ data: response([group()]) });
|
||||
renderTab();
|
||||
|
||||
const terms = screen.getAllByRole("term").map((node) => node.textContent);
|
||||
const values = screen.getAllByRole("definition").map((node) => node.textContent);
|
||||
expect(terms).toEqual(["Actual auto-router spend", "Estimated spend at highest-tier model"]);
|
||||
expect(values).toEqual(["$359.86", "$2,534.45"]);
|
||||
});
|
||||
|
||||
it("lets both hero columns shrink below their content so a large total cannot clip", () => {
|
||||
const huge = totals({ saved_spend: 123_456_789_012.34 });
|
||||
mockHook({ data: response([group(huge)], huge) });
|
||||
renderTab();
|
||||
|
||||
const figure = screen.getByText("$123,456,789,012.34");
|
||||
const grid = figure.closest('[data-slot="card"]')?.firstElementChild;
|
||||
expect(grid).toHaveClass("md:grid-cols-[minmax(0,1fr)_minmax(0,1fr)]");
|
||||
});
|
||||
|
||||
it("shows a cost increase as a positive delta rather than a saving", () => {
|
||||
|
|
@ -315,7 +340,7 @@ describe("AutoRouterBenchmarksTab", () => {
|
|||
|
||||
expect(screen.getByText("Total estimated savings")).toBeInTheDocument();
|
||||
expect(screen.getAllByText("$0.00")).toHaveLength(4);
|
||||
expect(screen.getByText("across 0 sessions")).toBeInTheDocument();
|
||||
expect(screen.getByText("· 0 sessions")).toBeInTheDocument();
|
||||
expect(screen.getByText("0s")).toBeInTheDocument();
|
||||
expect(screen.getByText(/turns measured/)).toBeInTheDocument();
|
||||
expect(screen.getAllByText("0.0%").length).toBeGreaterThan(0);
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ import AdvancedDatePicker from "@/components/shared/advanced_date_picker";
|
|||
import { Badge } from "@/components/ui/badge";
|
||||
import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
|
||||
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
|
||||
import { Separator } from "@/components/ui/separator";
|
||||
import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table";
|
||||
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
|
||||
import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip";
|
||||
|
|
@ -39,51 +40,51 @@ const Message: React.FC<{ children: React.ReactNode }> = ({ children }) => (
|
|||
<p className="py-8 text-center text-sm text-muted-foreground">{children}</p>
|
||||
);
|
||||
|
||||
const Metric: React.FC<{ label: string; value: string }> = ({ label, value }) => (
|
||||
const Metric: React.FC<{ label: string; value: string; hint?: string }> = ({ label, value, hint }) => (
|
||||
<Card size="sm">
|
||||
<CardHeader>
|
||||
<CardTitle className="text-sm font-normal text-muted-foreground">{label}</CardTitle>
|
||||
</CardHeader>
|
||||
<CardContent>
|
||||
<CardContent className="flex flex-wrap items-baseline gap-2">
|
||||
<p className="text-3xl font-semibold text-foreground">{value}</p>
|
||||
{hint && <p className="text-sm text-muted-foreground">{hint}</p>}
|
||||
</CardContent>
|
||||
</Card>
|
||||
);
|
||||
|
||||
const SpendRow: React.FC<{ label: string; value: string }> = ({ label, value }) => (
|
||||
<dl className="flex items-baseline justify-between gap-6 py-3">
|
||||
<dt className="text-sm text-muted-foreground">{label}</dt>
|
||||
<dd className="text-base font-semibold tabular-nums text-foreground">{value}</dd>
|
||||
</dl>
|
||||
);
|
||||
|
||||
const HeroCard: React.FC<{ view: BenchmarkView }> = ({ view }) => {
|
||||
const stats = view.stats;
|
||||
const cheaper = stats.saved_spend >= 0;
|
||||
return (
|
||||
<Card className="overflow-hidden py-0">
|
||||
<div className="grid md:grid-cols-[1fr_1fr]">
|
||||
<div className="flex flex-col justify-center gap-3 p-6">
|
||||
<p className="text-sm text-muted-foreground">Total estimated savings</p>
|
||||
<div className="flex flex-wrap items-center gap-3">
|
||||
<p className="text-5xl font-semibold tracking-tight text-foreground">{usd(stats.saved_spend)}</p>
|
||||
<div className="grid md:grid-cols-[minmax(0,1fr)_minmax(0,1fr)]">
|
||||
<div className="flex flex-col items-center justify-center gap-2 p-6">
|
||||
<p className="text-xs font-semibold uppercase tracking-wider text-muted-foreground">
|
||||
Total estimated savings
|
||||
</p>
|
||||
<div className="flex flex-wrap items-center justify-center gap-3">
|
||||
<p className="text-6xl font-semibold tracking-tight text-foreground">{usd(stats.saved_spend)}</p>
|
||||
<Badge
|
||||
variant="secondary"
|
||||
className={cheaper ? "bg-success/10 text-success" : "bg-destructive/10 text-destructive"}
|
||||
className={`h-6 px-2.5 text-sm ${cheaper ? "bg-success/10 text-success" : "bg-destructive/10 text-destructive"}`}
|
||||
>
|
||||
{stats.saved_spend !== 0 && (cheaper ? "-" : "+")}
|
||||
{Math.abs(stats.saved_pct).toFixed(0)}%
|
||||
</Badge>
|
||||
</div>
|
||||
<dl className="divide-y text-sm">
|
||||
<div className="flex items-baseline justify-between gap-6 py-3">
|
||||
<dt className="text-muted-foreground">Actual auto-router spend</dt>
|
||||
<dd className="font-medium tabular-nums text-foreground">{usd(stats.spend)}</dd>
|
||||
</div>
|
||||
<div className="flex items-baseline justify-between gap-6 py-3">
|
||||
<dt className="text-muted-foreground">Estimated spend at highest-tier model</dt>
|
||||
<dd className="font-medium tabular-nums text-foreground">{usd(stats.baseline_spend)}</dd>
|
||||
</div>
|
||||
</dl>
|
||||
</div>
|
||||
|
||||
<div className="flex flex-col items-center justify-center gap-2 border-t p-6 md:border-t-0 md:border-l">
|
||||
<p className="text-[11px] uppercase tracking-wide text-muted-foreground">Avg saved per session</p>
|
||||
<p className="text-5xl font-semibold tracking-tight text-foreground">{usd(stats.saved_per_session)}</p>
|
||||
<p className="text-sm text-muted-foreground">across {stats.sessions.toLocaleString()} sessions</p>
|
||||
<div className="flex flex-col justify-center border-t p-6 md:border-t-0 md:border-l">
|
||||
<SpendRow label="Actual auto-router spend" value={usd(stats.spend)} />
|
||||
<Separator />
|
||||
<SpendRow label="Estimated spend at highest-tier model" value={usd(stats.baseline_spend)} />
|
||||
</div>
|
||||
</div>
|
||||
</Card>
|
||||
|
|
@ -239,7 +240,12 @@ const BenchmarksBody: React.FC<BenchmarksBodyProps> = ({ isPending, error, data,
|
|||
|
||||
<TierTurnsChart view={view} autoRouters={autoRouters} />
|
||||
|
||||
<div className="grid grid-cols-1 gap-4 sm:grid-cols-3">
|
||||
<div className="grid grid-cols-1 gap-4 sm:grid-cols-2 lg:grid-cols-4">
|
||||
<Metric
|
||||
label="Avg saved per session"
|
||||
value={usd(stats.saved_per_session)}
|
||||
hint={`· ${stats.sessions.toLocaleString()} sessions`}
|
||||
/>
|
||||
<Metric label="Avg turns per session" value={stats.avg_turns_per_session.toFixed(1)} />
|
||||
<Metric label="Avg session length" value={durationLabel(stats.avg_session_seconds)} />
|
||||
<Metric label="Avg tokens per session" value={formatNumberWithCommas(stats.avg_tokens_per_session, 1, true)} />
|
||||
|
|
|
|||
|
|
@ -54,8 +54,14 @@ const asStringArray = (value: unknown): string[] =>
|
|||
|
||||
const dedupe = (models: string[]): string[] => Array.from(new Set(models));
|
||||
|
||||
const COMPLEXITY_TYPE_LABELS: Record<string, string> = {
|
||||
llm: "LLM Classifier",
|
||||
heuristic_first: "Heuristic first",
|
||||
custom: "Custom classifier",
|
||||
};
|
||||
|
||||
export const complexityTypeLabel = (config: Record<string, unknown>): string =>
|
||||
config.classifier_type === "llm" ? "LLM Classifier" : "Heuristic";
|
||||
(typeof config.classifier_type === "string" && COMPLEXITY_TYPE_LABELS[config.classifier_type]) || "Heuristic";
|
||||
|
||||
interface Presentation {
|
||||
typeLabel: string;
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
|
||||
import { useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations";
|
||||
import { useTeamsTable } from "@/app/(dashboard)/hooks/teams/useTeams";
|
||||
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
|
||||
import {
|
||||
DataTable,
|
||||
DataTableFilterDrawer,
|
||||
|
|
@ -9,14 +10,17 @@ import {
|
|||
DataTableToolbar,
|
||||
} from "@/components/shared/DataTable";
|
||||
import { SearchSelect } from "@/components/shared/SearchSelect";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { Input } from "@/components/ui/input";
|
||||
import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants";
|
||||
import { useDebouncedValue } from "@tanstack/react-pacer/debouncer";
|
||||
import { ColumnFiltersState, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table";
|
||||
import { Download } from "lucide-react";
|
||||
import React, { useCallback, useMemo, useState } from "react";
|
||||
|
||||
import { Team } from "../key_team_helpers/key_list";
|
||||
import { getTeamTableColumns, TEAM_TABLE_HIDDEN_COLUMNS } from "./teamTableColumns";
|
||||
import { exportTeamsToCsv } from "./teamsCsvExport";
|
||||
|
||||
interface TeamsTableProps {
|
||||
userRole: string | null;
|
||||
|
|
@ -49,7 +53,9 @@ export function TeamsTable({ userRole, userID, onSelectTeam, onEditTeam, onDelet
|
|||
const [columnFilters, setColumnFilters] = useState<ColumnFiltersState>([]);
|
||||
const [filtersOpen, setFiltersOpen] = useState(false);
|
||||
const [searchInput, setSearchInput] = useState("");
|
||||
const [isExporting, setIsExporting] = useState(false);
|
||||
const [searchQuery] = useDebouncedValue(searchInput, { wait: DEBOUNCE_WAIT_MS });
|
||||
const { accessToken } = useAuthorized();
|
||||
|
||||
const getFilterValue = useCallback(
|
||||
(columnId: string): string | undefined => {
|
||||
|
|
@ -61,16 +67,19 @@ export function TeamsTable({ userRole, userID, onSelectTeam, onEditTeam, onDelet
|
|||
|
||||
const isAdminView = userRole === "Admin" || userRole === "Admin Viewer";
|
||||
|
||||
const teamListOptions = {
|
||||
organizationID: getFilterValue("org_id"),
|
||||
team_alias: getFilterValue("alias"),
|
||||
teamID: getFilterValue("team_id"),
|
||||
search: searchQuery.trim() || undefined,
|
||||
searchTeamIdMatch: "prefix" as const,
|
||||
userID: isAdminView ? undefined : userID ?? undefined,
|
||||
sortBy: sorting[0]?.id,
|
||||
sortOrder: toSortOrder(sorting),
|
||||
};
|
||||
const teamListOptions = useMemo(
|
||||
() => ({
|
||||
organizationID: getFilterValue("org_id"),
|
||||
team_alias: getFilterValue("alias"),
|
||||
teamID: getFilterValue("team_id"),
|
||||
search: searchQuery.trim() || undefined,
|
||||
searchTeamIdMatch: "prefix" as const,
|
||||
userID: isAdminView ? undefined : userID ?? undefined,
|
||||
sortBy: sorting[0]?.id,
|
||||
sortOrder: toSortOrder(sorting),
|
||||
}),
|
||||
[getFilterValue, searchQuery, isAdminView, userID, sorting],
|
||||
);
|
||||
|
||||
const {
|
||||
data: teamsResponse,
|
||||
|
|
@ -97,6 +106,16 @@ export function TeamsTable({ userRole, userID, onSelectTeam, onEditTeam, onDelet
|
|||
setTablePagination((prev) => ({ ...prev, pageIndex: 0 }));
|
||||
}, []);
|
||||
|
||||
const handleExportCsv = useCallback(async () => {
|
||||
if (!accessToken || isExporting) return;
|
||||
setIsExporting(true);
|
||||
try {
|
||||
await exportTeamsToCsv(accessToken, teamListOptions);
|
||||
} finally {
|
||||
setIsExporting(false);
|
||||
}
|
||||
}, [accessToken, isExporting, teamListOptions]);
|
||||
|
||||
const columns = useMemo(() => {
|
||||
const columnDeps = { organizations, userRole, onSelectTeam, onEditTeam, onDeleteTeam };
|
||||
return getTeamTableColumns(columnDeps);
|
||||
|
|
@ -159,7 +178,18 @@ export function TeamsTable({ userRole, userID, onSelectTeam, onEditTeam, onDelet
|
|||
onOpenFilters={() => setFiltersOpen(true)}
|
||||
filterLabels={FILTER_LABELS}
|
||||
formatFilterValue={formatFilterValue}
|
||||
/>
|
||||
>
|
||||
<Button
|
||||
variant="outline"
|
||||
size="sm"
|
||||
onClick={handleExportCsv}
|
||||
disabled={isExporting}
|
||||
data-testid="teams-export-csv"
|
||||
>
|
||||
<Download />
|
||||
{isExporting ? "Exporting..." : "Export CSV"}
|
||||
</Button>
|
||||
</DataTableToolbar>
|
||||
<DataTableFilterDrawer
|
||||
table={table}
|
||||
open={filtersOpen}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,163 @@
|
|||
import { describe, expect, it, vi } from "vitest";
|
||||
|
||||
import type { TeamsResponse } from "@/app/(dashboard)/hooks/teams/useTeams";
|
||||
|
||||
import type { Team } from "../key_team_helpers/key_list";
|
||||
import {
|
||||
buildTeamsCsv,
|
||||
buildTeamsCsvRows,
|
||||
collectTeamMemberBudgetIds,
|
||||
fetchAllTeams,
|
||||
TEAMS_EXPORT_PAGE_SIZE,
|
||||
} from "./teamsCsvExport";
|
||||
|
||||
const makeTeam = (overrides: Partial<Team>): Team =>
|
||||
({
|
||||
team_id: "team-1",
|
||||
team_alias: "alias-1",
|
||||
models: [],
|
||||
max_budget: null,
|
||||
budget_duration: null,
|
||||
tpm_limit: null,
|
||||
rpm_limit: null,
|
||||
organization_id: "org-1",
|
||||
created_at: "2026-01-01T00:00:00Z",
|
||||
keys: [],
|
||||
members_with_roles: [],
|
||||
spend: 0,
|
||||
...overrides,
|
||||
}) as Team;
|
||||
|
||||
const makePage = (teams: Team[], page: number, totalPages: number): TeamsResponse => ({
|
||||
teams,
|
||||
total: teams.length,
|
||||
page,
|
||||
page_size: TEAMS_EXPORT_PAGE_SIZE,
|
||||
total_pages: totalPages,
|
||||
});
|
||||
|
||||
describe("fetchAllTeams", () => {
|
||||
it("returns the single page without extra requests", async () => {
|
||||
const fetchPage = vi.fn().mockResolvedValue(makePage([makeTeam({ team_id: "a" })], 1, 1));
|
||||
const teams = await fetchAllTeams(fetchPage);
|
||||
expect(teams.map((t) => t.team_id)).toEqual(["a"]);
|
||||
expect(fetchPage).toHaveBeenCalledTimes(1);
|
||||
expect(fetchPage).toHaveBeenCalledWith(1, TEAMS_EXPORT_PAGE_SIZE);
|
||||
});
|
||||
|
||||
it("fetches and concatenates every page in order", async () => {
|
||||
const fetchPage = vi
|
||||
.fn()
|
||||
.mockImplementation(async (page: number) => makePage([makeTeam({ team_id: `team-${page}` })], page, 3));
|
||||
const teams = await fetchAllTeams(fetchPage);
|
||||
expect(teams.map((t) => t.team_id)).toEqual(["team-1", "team-2", "team-3"]);
|
||||
expect(fetchPage).toHaveBeenCalledTimes(3);
|
||||
expect(fetchPage).toHaveBeenCalledWith(2, TEAMS_EXPORT_PAGE_SIZE);
|
||||
expect(fetchPage).toHaveBeenCalledWith(3, TEAMS_EXPORT_PAGE_SIZE);
|
||||
});
|
||||
});
|
||||
|
||||
describe("collectTeamMemberBudgetIds", () => {
|
||||
it("dedupes ids and skips teams without a member budget", () => {
|
||||
const teams = [
|
||||
makeTeam({ team_id: "a", metadata: { team_member_budget_id: "bud-1" } }),
|
||||
makeTeam({ team_id: "b", metadata: { team_member_budget_id: "bud-1" } }),
|
||||
makeTeam({ team_id: "c", metadata: {} }),
|
||||
makeTeam({ team_id: "d", metadata: { team_member_budget_id: "" } }),
|
||||
makeTeam({ team_id: "e" }),
|
||||
makeTeam({ team_id: "f", metadata: { team_member_budget_id: "bud-2" } }),
|
||||
];
|
||||
expect(collectTeamMemberBudgetIds(teams)).toEqual(["bud-1", "bud-2"]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("buildTeamsCsvRows", () => {
|
||||
it("maps configured limits, spend, models, and rate limits", () => {
|
||||
const teamFields: Partial<Team> = {
|
||||
team_id: "team-42",
|
||||
team_alias: "finance",
|
||||
organization_id: "org-9",
|
||||
models: ["gpt-4o", "claude-sonnet-4-5"],
|
||||
max_budget: 250,
|
||||
budget_duration: "30d",
|
||||
budget_reset_at: "2026-02-01T00:00:00Z",
|
||||
spend: 12.5,
|
||||
tpm_limit: 1000,
|
||||
rpm_limit: 50,
|
||||
members_count: 7,
|
||||
keys_count: 3,
|
||||
blocked: false,
|
||||
};
|
||||
const [row] = buildTeamsCsvRows([makeTeam(teamFields)], []);
|
||||
const expectedRow = {
|
||||
"Team Alias": "finance",
|
||||
"Team ID": "team-42",
|
||||
"Organization ID": "org-9",
|
||||
Models: "gpt-4o, claude-sonnet-4-5",
|
||||
"Max Budget (USD)": 250,
|
||||
"Budget Duration": "30d",
|
||||
"Budget Reset At": "2026-02-01T00:00:00Z",
|
||||
"Spend (USD)": 12.5,
|
||||
"TPM Limit": 1000,
|
||||
"RPM Limit": 50,
|
||||
"Team Member Budget (USD)": "",
|
||||
"Team Member Budget Duration": "",
|
||||
"Team Member TPM Limit": "",
|
||||
"Team Member RPM Limit": "",
|
||||
Members: 7,
|
||||
Keys: 3,
|
||||
Blocked: false,
|
||||
"Created At": "2026-01-01T00:00:00Z",
|
||||
};
|
||||
expect(row).toEqual(expectedRow);
|
||||
});
|
||||
|
||||
it("joins team member budget rows by budget id from metadata", () => {
|
||||
const teams = [
|
||||
makeTeam({ team_id: "a", metadata: { team_member_budget_id: "bud-1" } }),
|
||||
makeTeam({ team_id: "b" }),
|
||||
];
|
||||
const rows = buildTeamsCsvRows(teams, [
|
||||
{ budget_id: "bud-1", max_budget: 25, budget_duration: "7d", tpm_limit: 200, rpm_limit: 10 },
|
||||
]);
|
||||
expect(rows[0]["Team Member Budget (USD)"]).toBe(25);
|
||||
expect(rows[0]["Team Member Budget Duration"]).toBe("7d");
|
||||
expect(rows[0]["Team Member TPM Limit"]).toBe(200);
|
||||
expect(rows[0]["Team Member RPM Limit"]).toBe(10);
|
||||
expect(rows[1]["Team Member Budget (USD)"]).toBe("");
|
||||
});
|
||||
|
||||
it("falls back to members_with_roles and keys lengths when counts are absent", () => {
|
||||
const team = makeTeam({
|
||||
members_with_roles: [
|
||||
{ user_id: "u1", role: "admin" },
|
||||
{ user_id: "u2", role: "user" },
|
||||
],
|
||||
keys: [{ token: "t" } as Team["keys"][number]],
|
||||
});
|
||||
const [row] = buildTeamsCsvRows([team], []);
|
||||
expect(row.Members).toBe(2);
|
||||
expect(row.Keys).toBe(1);
|
||||
});
|
||||
});
|
||||
|
||||
describe("buildTeamsCsv", () => {
|
||||
it("produces a header row and quotes values containing commas", () => {
|
||||
const csv = buildTeamsCsv([makeTeam({ team_alias: "sales, emea", models: ["m1", "m2"] })], []);
|
||||
const [header, row] = csv.split("\r\n");
|
||||
expect(header).toBe(
|
||||
"Team Alias,Team ID,Organization ID,Models,Max Budget (USD),Budget Duration,Budget Reset At,Spend (USD)," +
|
||||
"TPM Limit,RPM Limit,Team Member Budget (USD),Team Member Budget Duration,Team Member TPM Limit," +
|
||||
"Team Member RPM Limit,Members,Keys,Blocked,Created At",
|
||||
);
|
||||
expect(row).toContain('"sales, emea"');
|
||||
expect(row).toContain('"m1, m2"');
|
||||
});
|
||||
|
||||
it("neutralizes formula-leading values so spreadsheets render them as text", () => {
|
||||
const csv = buildTeamsCsv([makeTeam({ team_alias: "=SUM(A1:A9)" })], []);
|
||||
const [, row] = csv.split("\r\n");
|
||||
expect(row).toContain('"\'=SUM(A1:A9)"');
|
||||
expect(row).not.toContain("=SUM(A1:A9),");
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,95 @@
|
|||
import Papa from "papaparse";
|
||||
|
||||
import { TeamListCallOptions, TeamsResponse, teamListCall } from "@/app/(dashboard)/hooks/teams/useTeams";
|
||||
|
||||
import { Team } from "../key_team_helpers/key_list";
|
||||
import { apiClient } from "../networking";
|
||||
|
||||
export interface TeamMemberBudget {
|
||||
budget_id: string;
|
||||
max_budget?: number | null;
|
||||
budget_duration?: string | null;
|
||||
tpm_limit?: number | null;
|
||||
rpm_limit?: number | null;
|
||||
}
|
||||
|
||||
export const TEAMS_EXPORT_PAGE_SIZE = 100;
|
||||
|
||||
type FetchTeamsPage = (page: number, pageSize: number) => Promise<TeamsResponse>;
|
||||
|
||||
export const fetchAllTeams = async (fetchPage: FetchTeamsPage): Promise<Team[]> => {
|
||||
const firstPage = await fetchPage(1, TEAMS_EXPORT_PAGE_SIZE);
|
||||
const totalPages = firstPage.total_pages ?? 1;
|
||||
if (totalPages <= 1) return firstPage.teams;
|
||||
|
||||
const remainingPages = await Promise.all(
|
||||
Array.from({ length: totalPages - 1 }, (_, i) => fetchPage(i + 2, TEAMS_EXPORT_PAGE_SIZE)),
|
||||
);
|
||||
return [firstPage, ...remainingPages].flatMap((page) => page.teams);
|
||||
};
|
||||
|
||||
const teamMemberBudgetId = (team: Team): string | null => {
|
||||
const id = team.metadata?.team_member_budget_id;
|
||||
return typeof id === "string" && id.length > 0 ? id : null;
|
||||
};
|
||||
|
||||
export const collectTeamMemberBudgetIds = (teams: Team[]): string[] =>
|
||||
Array.from(new Set(teams.map(teamMemberBudgetId).filter((id): id is string => id !== null)));
|
||||
|
||||
const cell = (value: string | number | boolean | null | undefined): string | number | boolean => value ?? "";
|
||||
|
||||
export const buildTeamsCsvRows = (
|
||||
teams: Team[],
|
||||
budgets: TeamMemberBudget[],
|
||||
): Record<string, string | number | boolean>[] => {
|
||||
const budgetsById = new Map(budgets.map((budget) => [budget.budget_id, budget]));
|
||||
return teams.map((team) => {
|
||||
const budgetId = teamMemberBudgetId(team);
|
||||
const memberBudget = budgetId ? budgetsById.get(budgetId) : undefined;
|
||||
return {
|
||||
"Team Alias": cell(team.team_alias),
|
||||
"Team ID": cell(team.team_id),
|
||||
"Organization ID": cell(team.organization_id),
|
||||
Models: (team.models ?? []).join(", "),
|
||||
"Max Budget (USD)": cell(team.max_budget),
|
||||
"Budget Duration": cell(team.budget_duration),
|
||||
"Budget Reset At": cell(team.budget_reset_at),
|
||||
"Spend (USD)": cell(team.spend),
|
||||
"TPM Limit": cell(team.tpm_limit),
|
||||
"RPM Limit": cell(team.rpm_limit),
|
||||
"Team Member Budget (USD)": cell(memberBudget?.max_budget),
|
||||
"Team Member Budget Duration": cell(memberBudget?.budget_duration),
|
||||
"Team Member TPM Limit": cell(memberBudget?.tpm_limit),
|
||||
"Team Member RPM Limit": cell(memberBudget?.rpm_limit),
|
||||
Members: cell(team.members_count ?? team.members_with_roles?.length),
|
||||
Keys: cell(team.keys_count ?? team.keys?.length),
|
||||
Blocked: cell(team.blocked),
|
||||
"Created At": cell(team.created_at),
|
||||
};
|
||||
});
|
||||
};
|
||||
|
||||
export const buildTeamsCsv = (teams: Team[], budgets: TeamMemberBudget[]): string =>
|
||||
Papa.unparse(buildTeamsCsvRows(teams, budgets), { escapeFormulae: true });
|
||||
|
||||
const downloadCsv = (csv: string, fileName: string): void => {
|
||||
const blob = new Blob([csv], { type: "text/csv;charset=utf-8;" });
|
||||
const url = window.URL.createObjectURL(blob);
|
||||
const a = document.createElement("a");
|
||||
a.href = url;
|
||||
a.download = fileName;
|
||||
document.body.appendChild(a);
|
||||
a.click();
|
||||
document.body.removeChild(a);
|
||||
window.URL.revokeObjectURL(url);
|
||||
};
|
||||
|
||||
export const exportTeamsToCsv = async (accessToken: string, options: TeamListCallOptions): Promise<number> => {
|
||||
const teams = await fetchAllTeams((page, pageSize) => teamListCall(accessToken, page, pageSize, options));
|
||||
const budgetIds = collectTeamMemberBudgetIds(teams);
|
||||
const budgets = budgetIds.length
|
||||
? await apiClient.post<TeamMemberBudget[]>("/budget/info", { accessToken, body: { budgets: budgetIds } })
|
||||
: [];
|
||||
downloadCsv(buildTeamsCsv(teams, budgets), `teams_export_${new Date().toISOString().split("T")[0]}.csv`);
|
||||
return teams.length;
|
||||
};
|
||||
|
|
@ -27,6 +27,10 @@ import {
|
|||
CLASSIFICATION_RUBRIC_KEYS,
|
||||
ClassificationRubric,
|
||||
effectiveTierLabel,
|
||||
heuristicScoringRole,
|
||||
usesLlmClassifier,
|
||||
DEFAULT_HEURISTIC_FIRST_MAX_TIER,
|
||||
HEURISTIC_FIRST_MAX_TIER_KEYS,
|
||||
} from "./ComplexityRouterConfig";
|
||||
|
||||
const DEFAULT_SCORING_EXPLANATION =
|
||||
|
|
@ -49,7 +53,7 @@ const CUSTOM_PROMPT_WITH_DEFAULT_MODEL_FALLBACK =
|
|||
*/
|
||||
const scoringExplanation = (value: ComplexityRouterConfigValue): string => {
|
||||
const usesCustomPrompt =
|
||||
value.classifier_type === "llm" && Boolean(value.classifier_llm_config?.system_prompt?.trim());
|
||||
usesLlmClassifier(value.classifier_type) && Boolean(value.classifier_llm_config?.system_prompt?.trim());
|
||||
if (!usesCustomPrompt) return DEFAULT_SCORING_EXPLANATION;
|
||||
return value.classifier_fallback === "default_model"
|
||||
? CUSTOM_PROMPT_WITH_DEFAULT_MODEL_FALLBACK
|
||||
|
|
@ -148,7 +152,7 @@ const ClassificationMethodConfig: React.FC<ClassificationMethodConfigProps> = ({
|
|||
}) => {
|
||||
const hasDefaultModel = Boolean(defaultModel);
|
||||
const classifierModelMissing =
|
||||
showValidationErrors && value.classifier_type === "llm" && !value.classifier_llm_config?.model;
|
||||
showValidationErrors && usesLlmClassifier(value.classifier_type) && !value.classifier_llm_config?.model;
|
||||
const usesCustomPrompt = Boolean(value.classifier_llm_config?.system_prompt?.trim());
|
||||
const contextBudget = value.classifier_context_budget_chars ?? DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS;
|
||||
const contextBudgetQuotesNothing = contextBudget > 0 && contextBudget < MIN_QUOTED_CONTEXT_TURN_CHARS;
|
||||
|
|
@ -158,29 +162,35 @@ const ClassificationMethodConfig: React.FC<ClassificationMethodConfigProps> = ({
|
|||
const nextValue: ComplexityRouterConfigValue = {
|
||||
...value,
|
||||
classifier_type: classifierType,
|
||||
classifier_llm_config:
|
||||
classifierType === "llm"
|
||||
? value.classifier_llm_config ?? {
|
||||
model: "",
|
||||
timeout_ms: DEFAULT_CLASSIFIER_TIMEOUT_MS,
|
||||
classification_rubric: NEW_CLASSIFIER_CLASSIFICATION_RUBRIC,
|
||||
}
|
||||
classifier_llm_config: usesLlmClassifier(classifierType)
|
||||
? value.classifier_llm_config ?? {
|
||||
model: "",
|
||||
timeout_ms: DEFAULT_CLASSIFIER_TIMEOUT_MS,
|
||||
classification_rubric: NEW_CLASSIFIER_CLASSIFICATION_RUBRIC,
|
||||
}
|
||||
: undefined,
|
||||
classifier_context_window_size: usesLlmClassifier(classifierType)
|
||||
? value.classifier_context_window_size ?? DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE
|
||||
: undefined,
|
||||
classifier_context_budget_chars: usesLlmClassifier(classifierType)
|
||||
? value.classifier_context_budget_chars ?? DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS
|
||||
: undefined,
|
||||
classifier_context_include_assistant_turns: usesLlmClassifier(classifierType)
|
||||
? value.classifier_context_include_assistant_turns
|
||||
: undefined,
|
||||
classifier_fallback: usesLlmClassifier(classifierType) ? value.classifier_fallback : undefined,
|
||||
heuristic_first_max_tier:
|
||||
classifierType === "heuristic_first"
|
||||
? value.heuristic_first_max_tier ?? DEFAULT_HEURISTIC_FIRST_MAX_TIER
|
||||
: undefined,
|
||||
classifier_context_window_size:
|
||||
classifierType === "llm"
|
||||
? value.classifier_context_window_size ?? DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE
|
||||
: undefined,
|
||||
classifier_context_budget_chars:
|
||||
classifierType === "llm"
|
||||
? value.classifier_context_budget_chars ?? DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS
|
||||
: undefined,
|
||||
classifier_context_include_assistant_turns:
|
||||
classifierType === "llm" ? value.classifier_context_include_assistant_turns : undefined,
|
||||
classifier_fallback: classifierType === "llm" ? value.classifier_fallback : undefined,
|
||||
};
|
||||
onChange(nextValue);
|
||||
};
|
||||
|
||||
const handleHeuristicFirstMaxTierChange = (tier: string) => {
|
||||
onChange({ ...value, heuristic_first_max_tier: tier });
|
||||
};
|
||||
|
||||
const handleClassifierModelChange = (model: string) => {
|
||||
onChange({
|
||||
...value,
|
||||
|
|
@ -265,7 +275,7 @@ const ClassificationMethodConfig: React.FC<ClassificationMethodConfigProps> = ({
|
|||
<span>
|
||||
<strong className="font-semibold">Heuristic</strong>{" "}
|
||||
<span className="text-muted-foreground">
|
||||
(default) — rule-based scoring, no API calls, <1ms latency
|
||||
(default), rule-based scoring with no API calls and <1ms latency
|
||||
</span>
|
||||
</span>
|
||||
</Label>
|
||||
|
|
@ -273,13 +283,47 @@ const ClassificationMethodConfig: React.FC<ClassificationMethodConfigProps> = ({
|
|||
<RadioGroupItem value="llm" className="mt-0.5" />
|
||||
<span>
|
||||
<strong className="font-semibold">LLM Classifier</strong>{" "}
|
||||
<span className="text-muted-foreground">— use a model to decide the tier (e.g. a small/fast model)</span>
|
||||
<span className="text-muted-foreground">calls a model to decide the tier (e.g. a small/fast model)</span>
|
||||
</span>
|
||||
</Label>
|
||||
<Label className="items-start font-normal leading-normal">
|
||||
<RadioGroupItem value="heuristic_first" className="mt-0.5" />
|
||||
<span>
|
||||
<strong className="font-semibold">Heuristic first</strong>{" "}
|
||||
<span className="text-muted-foreground">
|
||||
scores locally, and only pays for the classifier when the score does not confidently land a cheap tier
|
||||
</span>
|
||||
</span>
|
||||
</Label>
|
||||
</div>
|
||||
</RadioGroup>
|
||||
|
||||
{value.classifier_type === "llm" && (
|
||||
{value.classifier_type === "heuristic_first" && (
|
||||
<div className="mt-4 space-y-2">
|
||||
<strong className="block font-semibold">Decide locally up to</strong>
|
||||
<Select
|
||||
value={value.heuristic_first_max_tier}
|
||||
onValueChange={(tier: unknown) => handleHeuristicFirstMaxTierChange(tier as string)}
|
||||
>
|
||||
<SelectTrigger className="w-full">
|
||||
<SelectValue />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
{HEURISTIC_FIRST_MAX_TIER_KEYS.map((tier) => (
|
||||
<SelectItem key={tier} value={tier}>
|
||||
{effectiveTierLabel(tier, value.tier_labels)}
|
||||
</SelectItem>
|
||||
))}
|
||||
</SelectContent>
|
||||
</Select>
|
||||
<p className="text-sm text-muted-foreground">
|
||||
A request the scorer places at or below this tier routes there without a classifier call. Anything the
|
||||
scorer places higher, and anything it found no signal for at all, goes to the classifier instead
|
||||
</p>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{usesLlmClassifier(value.classifier_type) && (
|
||||
<div className="mt-4 space-y-3">
|
||||
<div>
|
||||
<strong className="block mb-1 font-semibold">Classifier Model</strong>
|
||||
|
|
@ -459,7 +503,7 @@ const ClassificationMethodConfig: React.FC<ClassificationMethodConfigProps> = ({
|
|||
</div>
|
||||
)}
|
||||
|
||||
{value.classifier_type === "heuristic" && (
|
||||
{heuristicScoringRole(value) !== "never" && (
|
||||
<div className="mt-4">
|
||||
<div className="flex items-center gap-2 mb-1">
|
||||
<strong className="font-semibold">Custom Technical Keywords</strong>
|
||||
|
|
|
|||
|
|
@ -1012,3 +1012,47 @@ describe("ComplexityRouterConfig per-model effort filtering", () => {
|
|||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("ComplexityRouterConfig custom technical keywords", () => {
|
||||
const openClassificationPanel = (value: ComplexityRouterConfigValue) => {
|
||||
renderWithProviders(<ComplexityRouterConfig modelInfo={mockModelInfo} value={value} onChange={vi.fn()} />);
|
||||
fireEvent.click(screen.getByText("Advanced: Classification Method"));
|
||||
};
|
||||
|
||||
const llmConfig = { model: "gpt-3.5-turbo", timeout_ms: 3000 };
|
||||
|
||||
it.each([
|
||||
["heuristic", { ...defaultValue, classifier_type: "heuristic" as const }],
|
||||
[
|
||||
"heuristic_first",
|
||||
{
|
||||
...defaultValue,
|
||||
classifier_type: "heuristic_first" as const,
|
||||
heuristic_first_max_tier: "SIMPLE",
|
||||
classifier_llm_config: llmConfig,
|
||||
},
|
||||
],
|
||||
[
|
||||
"llm falling back to the scorer",
|
||||
{
|
||||
...defaultValue,
|
||||
classifier_type: "llm" as const,
|
||||
classifier_llm_config: llmConfig,
|
||||
classifier_fallback: "heuristic" as const,
|
||||
},
|
||||
],
|
||||
])("offers the keywords on a router whose scorer runs: %s", (_label, value) => {
|
||||
openClassificationPanel(value);
|
||||
expect(screen.getByText("Custom Technical Keywords")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("hides the keywords when the scorer never runs, so they cannot imply an effect they have none", () => {
|
||||
openClassificationPanel({
|
||||
...defaultValue,
|
||||
classifier_type: "llm",
|
||||
classifier_llm_config: llmConfig,
|
||||
classifier_fallback: "default_model",
|
||||
});
|
||||
expect(screen.queryByText("Custom Technical Keywords")).not.toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -95,7 +95,15 @@ export interface ClassifierLLMConfig {
|
|||
system_prompt?: string;
|
||||
}
|
||||
|
||||
export type ClassifierType = "heuristic" | "llm";
|
||||
export type ClassifierType = "heuristic" | "llm" | "heuristic_first";
|
||||
|
||||
/**
|
||||
* Whether this router can call classifier_llm_config.model. Mirrors the backend's
|
||||
* ComplexityRouterConfig.uses_llm_classifier, and is the single gate for every classifier-only
|
||||
* control and payload key, so a new chaining type cannot strip knobs the operator set.
|
||||
*/
|
||||
export const usesLlmClassifier = (classifierType: ClassifierType): boolean =>
|
||||
classifierType === "llm" || classifierType === "heuristic_first";
|
||||
|
||||
export type ClassifierFallback = "heuristic" | "default_model";
|
||||
|
||||
|
|
@ -113,13 +121,14 @@ export type HeuristicScoringRole = "decides" | "fallback_only" | "never";
|
|||
/**
|
||||
* Whether the heuristic scorer runs on this router at all, which is what gates its knobs. An LLM
|
||||
* classifier still falls back to the scorer unless the fallback is the default model, so the gate cannot be
|
||||
* a plain classifier_type check.
|
||||
* a plain classifier_type check. Under heuristic_first the scorer runs first on every request and
|
||||
* decides outright whenever it lands at or below the threshold.
|
||||
*/
|
||||
export const heuristicScoringRoleFor = (
|
||||
classifierType: ClassifierType,
|
||||
classifierFallback: ClassifierFallback | undefined,
|
||||
): HeuristicScoringRole => {
|
||||
if (classifierType === "heuristic") return "decides";
|
||||
if (classifierType === "heuristic" || classifierType === "heuristic_first") return "decides";
|
||||
return (classifierFallback ?? DEFAULT_CLASSIFIER_FALLBACK) === "heuristic" ? "fallback_only" : "never";
|
||||
};
|
||||
|
||||
|
|
@ -142,6 +151,8 @@ export interface ComplexityRouterConfigValue {
|
|||
classifier_context_per_turn_chars?: number;
|
||||
classifier_context_include_assistant_turns?: boolean;
|
||||
classifier_fallback?: ClassifierFallback;
|
||||
/** Highest tier the scorer may decide alone under heuristic_first. Required by that type, rejected by the others. */
|
||||
heuristic_first_max_tier?: string;
|
||||
session_affinity?: boolean;
|
||||
deployment_affinity?: boolean;
|
||||
/** Tier floor for coding-agent plan-mode requests. Unset means detection is off, matching the backend. */
|
||||
|
|
@ -223,6 +234,14 @@ export const TIER_KEYS = Object.keys(TIER_DESCRIPTIONS) as Array<keyof Complexit
|
|||
export const effectiveTierLabel = (tier: keyof ComplexityTiers, tierLabels: ComplexityTierLabels | undefined): string =>
|
||||
tierLabels?.[tier]?.trim() || TIER_DESCRIPTIONS[tier].label;
|
||||
|
||||
export const DEFAULT_HEURISTIC_FIRST_MAX_TIER = "SIMPLE";
|
||||
|
||||
/**
|
||||
* Tiers the heuristic_first threshold may name. The top tier is excluded because it would short
|
||||
* circuit every request and leave the classifier unreachable, which the backend rejects.
|
||||
*/
|
||||
export const HEURISTIC_FIRST_MAX_TIER_KEYS = TIER_KEYS.slice(0, -1);
|
||||
|
||||
const ComplexityRouterConfig: React.FC<ComplexityRouterConfigProps> = ({
|
||||
modelInfo,
|
||||
value,
|
||||
|
|
@ -314,7 +333,7 @@ const ComplexityRouterConfig: React.FC<ComplexityRouterConfigProps> = ({
|
|||
<span className="block mb-4 text-xs text-muted-foreground">
|
||||
Rename a tier to use your own vocabulary in the dashboard and your spend logs. Renaming doesn't change how
|
||||
requests are classified, and callers never see these names.
|
||||
{value.classifier_type === "llm" &&
|
||||
{usesLlmClassifier(value.classifier_type) &&
|
||||
" Your classifier model reads these names, so clearer ones can sharpen its choices."}
|
||||
</span>
|
||||
|
||||
|
|
|
|||
|
|
@ -344,6 +344,7 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
|
|||
tiers: complexityRouterConfig.tiers,
|
||||
defaultModel: complexityRouterConfig.default_model,
|
||||
planModeMinTier: complexityRouterConfig.plan_mode_min_tier,
|
||||
heuristicFirstMaxTier: complexityRouterConfig.heuristic_first_max_tier,
|
||||
tierLabels: complexityRouterConfig.tier_labels,
|
||||
classifierType: complexityRouterConfig.classifier_type,
|
||||
classifierLlmConfig: complexityRouterConfig.classifier_llm_config,
|
||||
|
|
|
|||
|
|
@ -726,3 +726,36 @@ describe("getKeywordTierRulesError orphaned tiers", () => {
|
|||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("heuristic_first", () => {
|
||||
const heuristicFirstParams: BuildComplexityRouterConfigParams = {
|
||||
...baseParams,
|
||||
classifierType: "heuristic_first",
|
||||
heuristicFirstMaxTier: "SIMPLE",
|
||||
classifierLlmConfig: { model: "gpt-4o-mini", timeout_ms: 3000 },
|
||||
classifierContextWindowSize: 5,
|
||||
classifierContextBudgetChars: 4000,
|
||||
classifierFallback: "default_model",
|
||||
};
|
||||
|
||||
it("emits heuristic_first_max_tier", () => {
|
||||
const config = buildComplexityRouterConfig(heuristicFirstParams);
|
||||
expect(config.classifier_type).toBe("heuristic_first");
|
||||
expect(config.heuristic_first_max_tier).toBe("SIMPLE");
|
||||
});
|
||||
|
||||
it("keeps every classifier key the operator set, since heuristic_first still calls the classifier", () => {
|
||||
const config = buildComplexityRouterConfig(heuristicFirstParams);
|
||||
expect(config.classifier_llm_config).toEqual({ model: "gpt-4o-mini", timeout_ms: 3000 });
|
||||
expect(config.classifier_context_window_size).toBe(5);
|
||||
expect(config.classifier_context_budget_chars).toBe(4000);
|
||||
expect(config.classifier_fallback).toBe("default_model");
|
||||
});
|
||||
|
||||
it("omits heuristic_first_max_tier on every other classifier type, which the backend rejects it on", () => {
|
||||
for (const classifierType of ["heuristic", "llm"] as const) {
|
||||
const config = buildComplexityRouterConfig({ ...heuristicFirstParams, classifierType });
|
||||
expect(config.heuristic_first_max_tier).toBeUndefined();
|
||||
}
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ import {
|
|||
TokenThresholds,
|
||||
effectiveTierLabel,
|
||||
heuristicScoringRoleFor,
|
||||
usesLlmClassifier,
|
||||
} from "./ComplexityRouterConfig";
|
||||
|
||||
/**
|
||||
|
|
@ -86,6 +87,7 @@ export interface BuildComplexityRouterConfigParams {
|
|||
classifierContextBudgetChars: number | undefined;
|
||||
classifierContextIncludeAssistantTurns: boolean | undefined;
|
||||
classifierFallback: ClassifierFallback | undefined;
|
||||
heuristicFirstMaxTier: string | undefined;
|
||||
sessionAffinity: boolean;
|
||||
deploymentAffinity: boolean;
|
||||
customTechnicalKeywords: string[];
|
||||
|
|
@ -118,6 +120,7 @@ export interface ComplexityRouterConfigPayload {
|
|||
classifier_context_per_turn_chars?: number;
|
||||
classifier_context_include_assistant_turns?: boolean;
|
||||
classifier_fallback?: ClassifierFallback;
|
||||
heuristic_first_max_tier?: string;
|
||||
session_affinity: boolean;
|
||||
deployment_affinity: boolean;
|
||||
custom_technical_keywords?: string[];
|
||||
|
|
@ -208,7 +211,7 @@ export const getKeywordTierRulesError = (
|
|||
export const getClassifierModelError = (
|
||||
config: Pick<ComplexityRouterConfigValue, "classifier_type" | "classifier_llm_config">,
|
||||
): string | null =>
|
||||
config.classifier_type === "llm" && !config.classifier_llm_config?.model
|
||||
usesLlmClassifier(config.classifier_type) && !config.classifier_llm_config?.model
|
||||
? "Please select a classifier model, or switch back to Heuristic"
|
||||
: null;
|
||||
|
||||
|
|
@ -236,6 +239,7 @@ export const buildComplexityRouterConfig = ({
|
|||
classifierContextBudgetChars,
|
||||
classifierContextIncludeAssistantTurns,
|
||||
classifierFallback,
|
||||
heuristicFirstMaxTier,
|
||||
sessionAffinity,
|
||||
deploymentAffinity,
|
||||
customTechnicalKeywords,
|
||||
|
|
@ -276,18 +280,21 @@ export const buildComplexityRouterConfig = ({
|
|||
...(planModeMinTier?.trim() && { plan_mode_min_tier: planModeMinTier }),
|
||||
...(cleanedTierLabels && { tier_labels: cleanedTierLabels }),
|
||||
classifier_type: classifierType,
|
||||
...(classifierType === "llm" &&
|
||||
...(usesLlmClassifier(classifierType) &&
|
||||
classifierLlmConfig && { classifier_llm_config: normalizeClassifierLlmConfig(classifierLlmConfig) }),
|
||||
...(classifierType === "llm" && classifierFallback !== undefined && { classifier_fallback: classifierFallback }),
|
||||
...(classifierType === "llm" &&
|
||||
...(usesLlmClassifier(classifierType) &&
|
||||
classifierFallback !== undefined && { classifier_fallback: classifierFallback }),
|
||||
...(classifierType === "heuristic_first" &&
|
||||
heuristicFirstMaxTier?.trim() && { heuristic_first_max_tier: heuristicFirstMaxTier }),
|
||||
...(usesLlmClassifier(classifierType) &&
|
||||
classifierContextWindowSize !== undefined && {
|
||||
classifier_context_window_size: classifierContextWindowSize,
|
||||
}),
|
||||
...(classifierType === "llm" &&
|
||||
...(usesLlmClassifier(classifierType) &&
|
||||
classifierContextBudgetChars !== undefined && {
|
||||
classifier_context_budget_chars: classifierContextBudgetChars,
|
||||
}),
|
||||
...(classifierType === "llm" &&
|
||||
...(usesLlmClassifier(classifierType) &&
|
||||
classifierContextIncludeAssistantTurns !== undefined && {
|
||||
classifier_context_include_assistant_turns: classifierContextIncludeAssistantTurns,
|
||||
}),
|
||||
|
|
|
|||
|
|
@ -1,6 +1,11 @@
|
|||
import { describe, expect, it } from "vitest";
|
||||
|
||||
import { buildUpdatedComplexityRouterConfig, type KeywordMatchingState } from "./edit_auto_router_modal";
|
||||
import {
|
||||
MANAGED_COMPLEXITY_ROUTER_KEYS,
|
||||
buildUpdatedComplexityRouterConfig,
|
||||
hydrateComplexityRouterConfig,
|
||||
type KeywordMatchingState,
|
||||
} from "./edit_auto_router_modal";
|
||||
|
||||
const STORED = {
|
||||
tiers: { SIMPLE: ["gpt-4o-mini"], MEDIUM: [], COMPLEX: [], REASONING: [] },
|
||||
|
|
@ -440,3 +445,48 @@ describe("buildUpdatedComplexityRouterConfig tier model params", () => {
|
|||
expect(result).not.toHaveProperty("tier_model_configs");
|
||||
});
|
||||
});
|
||||
|
||||
describe("managed keys survive an untouched open-and-save", () => {
|
||||
// Every managed key is rewritten from form state on save, so one the hydrator forgets is silently
|
||||
// dropped from the saved config. This config sets each managed key to a value that actually
|
||||
// applies, so an untouched open-and-save must return every one of them.
|
||||
const STORED_ALL_MANAGED: Record<string, unknown> = {
|
||||
tiers: { SIMPLE: ["gpt-4o-mini"], MEDIUM: ["gpt-4o"], COMPLEX: ["opus"], REASONING: ["o1"] },
|
||||
tier_model_configs: { REASONING: [{ model_name: "o1", litellm_params: { reasoning_effort: "high" } }] },
|
||||
default_model: "gpt-4o",
|
||||
plan_mode_min_tier: "COMPLEX",
|
||||
tier_labels: { SIMPLE: "Cheap" },
|
||||
classifier_type: "heuristic_first",
|
||||
heuristic_first_max_tier: "SIMPLE",
|
||||
classifier_llm_config: { model: "gpt-4o-mini", timeout_ms: 3000 },
|
||||
classifier_context_window_size: 5,
|
||||
classifier_context_budget_chars: 4000,
|
||||
classifier_context_include_assistant_turns: true,
|
||||
classifier_fallback: "default_model",
|
||||
session_affinity: true,
|
||||
deployment_affinity: false,
|
||||
adaptive: true,
|
||||
adaptive_weights: { quality: 0.4, cost: 0.6 },
|
||||
tier_distance_penalty: 0.25,
|
||||
adaptive_eligible: "all",
|
||||
return_raw_model_name: true,
|
||||
tier_boundaries: { simple_medium: 0.2, medium_complex: 0.4, complex_reasoning: 0.7 },
|
||||
token_thresholds: { simple: 20, complex: 500 },
|
||||
dimension_weights: { tokenCount: 0.1 },
|
||||
reasoning_override_min_score: 0.3,
|
||||
};
|
||||
|
||||
it("carries every managed key through hydrate then save", () => {
|
||||
const hydrated = hydrateComplexityRouterConfig(STORED_ALL_MANAGED, undefined);
|
||||
const saved = buildUpdatedComplexityRouterConfig(STORED_ALL_MANAGED, hydrated);
|
||||
|
||||
const dropped = [...MANAGED_COMPLEXITY_ROUTER_KEYS].filter((key) => saved[key] === undefined);
|
||||
expect(dropped).toEqual([]);
|
||||
});
|
||||
|
||||
it("round-trips the heuristic_first threshold, which save requires and the backend rejects without", () => {
|
||||
const hydrated = hydrateComplexityRouterConfig(STORED_ALL_MANAGED, undefined);
|
||||
expect(hydrated.heuristic_first_max_tier).toBe("SIMPLE");
|
||||
expect(buildUpdatedComplexityRouterConfig(STORED_ALL_MANAGED, hydrated).heuristic_first_max_tier).toBe("SIMPLE");
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -37,6 +37,10 @@ import {
|
|||
hydrateTokenThresholds,
|
||||
} from "../add_model/heuristic_scoring_knobs";
|
||||
import ComplexityRouterConfig, {
|
||||
AdaptiveEligible,
|
||||
AdaptiveRouterWeights,
|
||||
ClassifierLLMConfig,
|
||||
ClassifierType,
|
||||
ComplexityRouterConfigValue,
|
||||
ComplexityTiers,
|
||||
DEFAULT_ADAPTIVE_WEIGHTS,
|
||||
|
|
@ -65,7 +69,101 @@ interface EditAutoRouterModalProps {
|
|||
// Keys this modal rewrites from its own form state on save. Anything absent from this set is
|
||||
// carried through untouched from the stored config, so a key only belongs here once the modal
|
||||
// actually renders a control that can set it.
|
||||
const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([
|
||||
/** The complexity_router_config as it comes back from the proxy, before any hydration. Fields the
|
||||
* hydrators validate themselves stay `unknown`; the ones assigned straight through carry their type. */
|
||||
export interface StoredComplexityRouterConfig {
|
||||
tiers?: Partial<Record<keyof ComplexityTiers, unknown>>;
|
||||
tier_model_configs?: unknown;
|
||||
default_model?: string | null;
|
||||
plan_mode_min_tier?: unknown;
|
||||
heuristic_first_max_tier?: unknown;
|
||||
tier_labels?: unknown;
|
||||
classifier_type?: ClassifierType;
|
||||
classifier_llm_config?: ClassifierLLMConfig;
|
||||
classifier_context_window_size?: unknown;
|
||||
classifier_context_budget_chars?: unknown;
|
||||
classifier_context_include_assistant_turns?: unknown;
|
||||
classifier_fallback?: unknown;
|
||||
tier_boundaries?: unknown;
|
||||
token_thresholds?: unknown;
|
||||
dimension_weights?: unknown;
|
||||
reasoning_override_min_score?: unknown;
|
||||
session_affinity?: unknown;
|
||||
deployment_affinity?: unknown;
|
||||
adaptive?: boolean;
|
||||
adaptive_weights?: AdaptiveRouterWeights;
|
||||
tier_distance_penalty?: number;
|
||||
adaptive_eligible?: AdaptiveEligible;
|
||||
return_raw_model_name?: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* The stored complexity_router_config as form state. Every key in MANAGED_COMPLEXITY_ROUTER_KEYS is
|
||||
* rewritten from this state on save, so a key missing here is silently dropped from the saved config.
|
||||
*/
|
||||
export const hydrateComplexityRouterConfig = (
|
||||
parsedConfig: StoredComplexityRouterConfig,
|
||||
complexityRouterDefaultModel: string | null | undefined,
|
||||
): ComplexityRouterConfigValue => {
|
||||
const hydratedTiers: ComplexityTiers = {
|
||||
SIMPLE: normalizeTierModels(parsedConfig.tiers?.SIMPLE),
|
||||
MEDIUM: normalizeTierModels(parsedConfig.tiers?.MEDIUM),
|
||||
COMPLEX: normalizeTierModels(parsedConfig.tiers?.COMPLEX),
|
||||
REASONING: normalizeTierModels(parsedConfig.tiers?.REASONING),
|
||||
};
|
||||
|
||||
return {
|
||||
tiers: hydratedTiers,
|
||||
tier_model_params: hydrateTierModelParams(parsedConfig.tiers, parsedConfig.tier_model_configs),
|
||||
default_model: hydratePinnedDefaultModel(parsedConfig.default_model, complexityRouterDefaultModel, {
|
||||
tiers: hydratedTiers,
|
||||
}),
|
||||
plan_mode_min_tier:
|
||||
typeof parsedConfig.plan_mode_min_tier === "string" && parsedConfig.plan_mode_min_tier.trim() !== ""
|
||||
? parsedConfig.plan_mode_min_tier
|
||||
: undefined,
|
||||
tier_labels: hydrateTierLabels(parsedConfig.tier_labels),
|
||||
classifier_type: parsedConfig.classifier_type || "heuristic",
|
||||
classifier_llm_config: parsedConfig.classifier_llm_config,
|
||||
classifier_context_window_size:
|
||||
typeof parsedConfig.classifier_context_window_size === "number"
|
||||
? parsedConfig.classifier_context_window_size
|
||||
: undefined,
|
||||
classifier_context_budget_chars:
|
||||
typeof parsedConfig.classifier_context_budget_chars === "number"
|
||||
? parsedConfig.classifier_context_budget_chars
|
||||
: undefined,
|
||||
classifier_context_include_assistant_turns:
|
||||
typeof parsedConfig.classifier_context_include_assistant_turns === "boolean"
|
||||
? parsedConfig.classifier_context_include_assistant_turns
|
||||
: undefined,
|
||||
classifier_fallback:
|
||||
parsedConfig.classifier_fallback === "default_model" || parsedConfig.classifier_fallback === "heuristic"
|
||||
? parsedConfig.classifier_fallback
|
||||
: undefined,
|
||||
heuristic_first_max_tier:
|
||||
typeof parsedConfig.heuristic_first_max_tier === "string" && parsedConfig.heuristic_first_max_tier.trim() !== ""
|
||||
? parsedConfig.heuristic_first_max_tier
|
||||
: undefined,
|
||||
tier_boundaries: hydrateTierBoundaries(parsedConfig.tier_boundaries),
|
||||
token_thresholds: hydrateTokenThresholds(parsedConfig.token_thresholds),
|
||||
dimension_weights: hydrateDimensionWeights(parsedConfig.dimension_weights),
|
||||
reasoning_override_min_score: hydrateReasoningOverrideMinScore(parsedConfig.reasoning_override_min_score),
|
||||
session_affinity:
|
||||
typeof parsedConfig.session_affinity === "boolean" ? parsedConfig.session_affinity : DEFAULT_SESSION_AFFINITY,
|
||||
deployment_affinity:
|
||||
typeof parsedConfig.deployment_affinity === "boolean"
|
||||
? parsedConfig.deployment_affinity
|
||||
: DEFAULT_DEPLOYMENT_AFFINITY,
|
||||
adaptive: parsedConfig.adaptive || false,
|
||||
adaptive_weights: parsedConfig.adaptive_weights,
|
||||
tier_distance_penalty: parsedConfig.tier_distance_penalty,
|
||||
adaptive_eligible: parsedConfig.adaptive_eligible || "all",
|
||||
return_raw_model_name: parsedConfig.return_raw_model_name || false,
|
||||
};
|
||||
};
|
||||
|
||||
export const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([
|
||||
"tiers",
|
||||
"tier_model_configs",
|
||||
"default_model",
|
||||
|
|
@ -77,6 +175,7 @@ const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([
|
|||
"classifier_context_budget_chars",
|
||||
"classifier_context_include_assistant_turns",
|
||||
"classifier_fallback",
|
||||
"heuristic_first_max_tier",
|
||||
"session_affinity",
|
||||
"deployment_affinity",
|
||||
"adaptive",
|
||||
|
|
@ -150,6 +249,7 @@ export const buildUpdatedComplexityRouterConfig = (
|
|||
tiers: value.tiers,
|
||||
defaultModel: value.default_model,
|
||||
planModeMinTier: value.plan_mode_min_tier,
|
||||
heuristicFirstMaxTier: value.heuristic_first_max_tier,
|
||||
tierLabels: value.tier_labels,
|
||||
classifierType: value.classifier_type,
|
||||
classifierLlmConfig: value.classifier_llm_config,
|
||||
|
|
@ -314,62 +414,10 @@ const EditAutoRouterModal: React.FC<EditAutoRouterModalProps> = ({
|
|||
parsedConfig = JSON.parse(parsedConfig);
|
||||
}
|
||||
|
||||
const hydratedTiers: ComplexityTiers = {
|
||||
SIMPLE: normalizeTierModels(parsedConfig.tiers?.SIMPLE),
|
||||
MEDIUM: normalizeTierModels(parsedConfig.tiers?.MEDIUM),
|
||||
COMPLEX: normalizeTierModels(parsedConfig.tiers?.COMPLEX),
|
||||
REASONING: normalizeTierModels(parsedConfig.tiers?.REASONING),
|
||||
};
|
||||
|
||||
const hydratedComplexityRouterConfig: ComplexityRouterConfigValue = {
|
||||
tiers: hydratedTiers,
|
||||
tier_model_params: hydrateTierModelParams(parsedConfig.tiers, parsedConfig.tier_model_configs),
|
||||
default_model: hydratePinnedDefaultModel(
|
||||
parsedConfig.default_model,
|
||||
modelData.litellm_params?.complexity_router_default_model,
|
||||
{ tiers: hydratedTiers },
|
||||
),
|
||||
plan_mode_min_tier:
|
||||
typeof parsedConfig.plan_mode_min_tier === "string" && parsedConfig.plan_mode_min_tier.trim() !== ""
|
||||
? parsedConfig.plan_mode_min_tier
|
||||
: undefined,
|
||||
tier_labels: hydrateTierLabels(parsedConfig.tier_labels),
|
||||
classifier_type: parsedConfig.classifier_type || "heuristic",
|
||||
classifier_llm_config: parsedConfig.classifier_llm_config,
|
||||
classifier_context_window_size:
|
||||
typeof parsedConfig.classifier_context_window_size === "number"
|
||||
? parsedConfig.classifier_context_window_size
|
||||
: undefined,
|
||||
classifier_context_budget_chars:
|
||||
typeof parsedConfig.classifier_context_budget_chars === "number"
|
||||
? parsedConfig.classifier_context_budget_chars
|
||||
: undefined,
|
||||
classifier_context_include_assistant_turns:
|
||||
typeof parsedConfig.classifier_context_include_assistant_turns === "boolean"
|
||||
? parsedConfig.classifier_context_include_assistant_turns
|
||||
: undefined,
|
||||
classifier_fallback:
|
||||
parsedConfig.classifier_fallback === "default_model" || parsedConfig.classifier_fallback === "heuristic"
|
||||
? parsedConfig.classifier_fallback
|
||||
: undefined,
|
||||
tier_boundaries: hydrateTierBoundaries(parsedConfig.tier_boundaries),
|
||||
token_thresholds: hydrateTokenThresholds(parsedConfig.token_thresholds),
|
||||
dimension_weights: hydrateDimensionWeights(parsedConfig.dimension_weights),
|
||||
reasoning_override_min_score: hydrateReasoningOverrideMinScore(parsedConfig.reasoning_override_min_score),
|
||||
session_affinity:
|
||||
typeof parsedConfig.session_affinity === "boolean"
|
||||
? parsedConfig.session_affinity
|
||||
: DEFAULT_SESSION_AFFINITY,
|
||||
deployment_affinity:
|
||||
typeof parsedConfig.deployment_affinity === "boolean"
|
||||
? parsedConfig.deployment_affinity
|
||||
: DEFAULT_DEPLOYMENT_AFFINITY,
|
||||
adaptive: parsedConfig.adaptive || false,
|
||||
adaptive_weights: parsedConfig.adaptive_weights,
|
||||
tier_distance_penalty: parsedConfig.tier_distance_penalty,
|
||||
adaptive_eligible: parsedConfig.adaptive_eligible || "all",
|
||||
return_raw_model_name: parsedConfig.return_raw_model_name || false,
|
||||
};
|
||||
const hydratedComplexityRouterConfig = hydrateComplexityRouterConfig(
|
||||
parsedConfig,
|
||||
modelData.litellm_params?.complexity_router_default_model,
|
||||
);
|
||||
setComplexityRouterConfig(hydratedComplexityRouterConfig);
|
||||
setCustomTechnicalKeywords(
|
||||
Array.isArray(parsedConfig.custom_technical_keywords) ? parsedConfig.custom_technical_keywords : [],
|
||||
|
|
|
|||
|
|
@ -13,6 +13,9 @@ export interface Team {
|
|||
tpm_limit: number | null;
|
||||
rpm_limit: number | null;
|
||||
organization_id: string;
|
||||
metadata?: Record<string, unknown> | null;
|
||||
budget_reset_at?: string | null;
|
||||
blocked?: boolean;
|
||||
created_at: string;
|
||||
updated_at?: string | null;
|
||||
keys: KeyResponse[];
|
||||
|
|
|
|||
|
|
@ -459,6 +459,64 @@ describe("teamInfoCall", () => {
|
|||
});
|
||||
});
|
||||
|
||||
describe("uiSpendLogsCall exclude_internal_health_checks serialization", () => {
|
||||
const originalFetch = global.fetch;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
global.fetch = originalFetch;
|
||||
});
|
||||
|
||||
const mockOkFetch = () => {
|
||||
const mockFetch = vi.fn().mockResolvedValue({
|
||||
ok: true,
|
||||
json: vi.fn().mockResolvedValue({ data: [], total: 0, page: 1, page_size: 50, total_pages: 0 }),
|
||||
} as any);
|
||||
global.fetch = mockFetch as any;
|
||||
return mockFetch;
|
||||
};
|
||||
|
||||
const callWith = (params: Parameters<typeof Networking.uiSpendLogsCall>[0]["params"]) =>
|
||||
Networking.uiSpendLogsCall({
|
||||
accessToken: "token",
|
||||
start_date: "2026-01-01 00:00:00",
|
||||
end_date: "2026-01-02 00:00:00",
|
||||
params,
|
||||
});
|
||||
|
||||
const lastUrl = (mockFetch: ReturnType<typeof vi.fn>) => {
|
||||
const [url] = mockFetch.mock.calls.at(-1) ?? [];
|
||||
return new URL(url as string, "http://example.com");
|
||||
};
|
||||
|
||||
it("appends exclude_internal_health_checks=true when the toggle is on", async () => {
|
||||
const mockFetch = mockOkFetch();
|
||||
|
||||
await callWith({ exclude_internal_health_checks: true });
|
||||
|
||||
expect(lastUrl(mockFetch).searchParams.get("exclude_internal_health_checks")).toBe("true");
|
||||
});
|
||||
|
||||
it("omits exclude_internal_health_checks when the toggle is off", async () => {
|
||||
const mockFetch = mockOkFetch();
|
||||
|
||||
await callWith({ exclude_internal_health_checks: false });
|
||||
|
||||
expect(lastUrl(mockFetch).searchParams.has("exclude_internal_health_checks")).toBe(false);
|
||||
});
|
||||
|
||||
it("omits exclude_internal_health_checks when the param is absent", async () => {
|
||||
const mockFetch = mockOkFetch();
|
||||
|
||||
await callWith({});
|
||||
|
||||
expect(lastUrl(mockFetch).searchParams.has("exclude_internal_health_checks")).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("sessionSpendLogsCall", () => {
|
||||
const originalFetch = global.fetch;
|
||||
|
||||
|
|
|
|||
|
|
@ -2002,6 +2002,7 @@ interface UiSpendLogsParams {
|
|||
user_id?: string;
|
||||
end_user?: string;
|
||||
status_filter?: string;
|
||||
cache_hit_filter?: string;
|
||||
/** Filter by model name (e.g. "gpt-4") */
|
||||
model?: string;
|
||||
/** Filter by model ID (litellm model deployment id) */
|
||||
|
|
@ -2013,6 +2014,7 @@ interface UiSpendLogsParams {
|
|||
sort_order?: "asc" | "desc";
|
||||
min_spend?: number;
|
||||
max_spend?: number;
|
||||
exclude_internal_health_checks?: boolean;
|
||||
}
|
||||
|
||||
interface UiSpendLogsCallOptions {
|
||||
|
|
@ -2047,6 +2049,8 @@ export const uiSpendLogsCall = async ({
|
|||
if (value == null) continue;
|
||||
if (key === "min_spend" || key === "max_spend") {
|
||||
queryParams.append(key, value.toString());
|
||||
} else if (typeof value === "boolean") {
|
||||
if (value) queryParams.append(key, "true");
|
||||
} else if (typeof value === "string" && value !== "") {
|
||||
queryParams.append(key, String(value));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -72,6 +72,20 @@ function describeReasoningOverride(tierLabel: string | undefined, floor: number
|
|||
return `Heuristic, ${tierLabel ?? "REASONING"} override (2 or more reasoning markers, score of at least ${stated})`;
|
||||
}
|
||||
|
||||
const CONSTANT_CAUSE_LABELS: Record<string, string> = {
|
||||
heuristic_scorer: "Heuristic scorer",
|
||||
heuristic_first_short_circuit: "Heuristic scorer, classifier skipped",
|
||||
classifier_plugin: "Custom classifier plugin",
|
||||
semantic_keyword_match: "Semantic keyword match",
|
||||
session_affinity_pin: "Pinned to session",
|
||||
session_affinity_escalation: "Escalated from session pin",
|
||||
quality_tier: "Quality tier mapping",
|
||||
bandit: "Adaptive bandit",
|
||||
default_fallback: "Default model, no route matched",
|
||||
classifier_fallback: "Fallback tier, LLM classifier failed",
|
||||
default_model_fallback: "Default model, LLM classifier failed",
|
||||
};
|
||||
|
||||
function describeCause(decision: RoutingDecision): string {
|
||||
const {
|
||||
cause,
|
||||
|
|
@ -81,35 +95,19 @@ function describeCause(decision: RoutingDecision): string {
|
|||
reasoning_override_min_score: overrideFloor,
|
||||
} = decision;
|
||||
|
||||
const constant = cause ? CONSTANT_CAUSE_LABELS[cause] : undefined;
|
||||
if (constant) return constant;
|
||||
|
||||
switch (cause) {
|
||||
case "heuristic_scorer":
|
||||
return "Heuristic scorer";
|
||||
case "reasoning_override":
|
||||
return describeReasoningOverride(tierLabel, overrideFloor);
|
||||
case "llm_classifier":
|
||||
return classifierModel ? `LLM classifier (${classifierModel})` : "LLM classifier";
|
||||
case "literal_keyword_match":
|
||||
return matchedKeyword ? `Keyword match: "${matchedKeyword}"` : "Keyword match";
|
||||
case "semantic_keyword_match":
|
||||
return "Semantic keyword match";
|
||||
case "plan_mode":
|
||||
return describePlanModeFloor(matchedKeyword);
|
||||
case "session_affinity_pin":
|
||||
return "Pinned to session";
|
||||
case "session_affinity_escalation":
|
||||
return "Escalated from session pin";
|
||||
case "quality_tier":
|
||||
return "Quality tier mapping";
|
||||
case "keyword":
|
||||
return matchedKeyword ? `Keyword match: "${matchedKeyword}"` : "Keyword match";
|
||||
case "bandit":
|
||||
return "Adaptive bandit";
|
||||
case "default_fallback":
|
||||
return "Default model, no route matched";
|
||||
case "classifier_fallback":
|
||||
return "Fallback tier, LLM classifier failed";
|
||||
case "default_model_fallback":
|
||||
return "Default model, LLM classifier failed";
|
||||
case "plan_mode":
|
||||
return describePlanModeFloor(matchedKeyword);
|
||||
default:
|
||||
return cause ?? "Unknown";
|
||||
}
|
||||
|
|
|
|||
|
|
@ -23,6 +23,8 @@ interface LogsTableToolbarProps {
|
|||
onSelectedTimeIntervalChange: (value: { value: number; unit: string }) => void;
|
||||
isLiveTail: boolean;
|
||||
onIsLiveTailChange: (value: boolean) => void;
|
||||
excludeInternalHealthChecks: boolean;
|
||||
onExcludeInternalHealthChecksChange: (value: boolean) => void;
|
||||
onResetToFirstPage: () => void;
|
||||
onResetFilters: () => void;
|
||||
}
|
||||
|
|
@ -38,6 +40,8 @@ export function LogsTableToolbar({
|
|||
onSelectedTimeIntervalChange,
|
||||
isLiveTail,
|
||||
onIsLiveTailChange,
|
||||
excludeInternalHealthChecks,
|
||||
onExcludeInternalHealthChecksChange,
|
||||
onResetToFirstPage,
|
||||
onResetFilters,
|
||||
}: LogsTableToolbarProps) {
|
||||
|
|
@ -125,6 +129,15 @@ export function LogsTableToolbar({
|
|||
<Switch checked={isLiveTail} onCheckedChange={onIsLiveTailChange} aria-label="Live Tail" />
|
||||
</div>
|
||||
|
||||
<div className="flex items-center gap-2">
|
||||
<span className="text-sm font-medium">Hide Health Checks</span>
|
||||
<Switch
|
||||
checked={excludeInternalHealthChecks}
|
||||
onCheckedChange={onExcludeInternalHealthChecksChange}
|
||||
aria-label="Hide Health Checks"
|
||||
/>
|
||||
</div>
|
||||
|
||||
<Button variant="outline" size="sm" onClick={onResetFilters}>
|
||||
Reset Filters
|
||||
</Button>
|
||||
|
|
|
|||
|
|
@ -69,6 +69,7 @@ describe("RequestLogsFilters", () => {
|
|||
for (const label of [
|
||||
"Team ID",
|
||||
"Status",
|
||||
"Cache",
|
||||
"Key Alias",
|
||||
"User ID",
|
||||
"End User",
|
||||
|
|
@ -259,4 +260,37 @@ describe("RequestLogsFilters", () => {
|
|||
|
||||
expect(await screen.findByText(label)).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it.each([
|
||||
["", "All Requests"],
|
||||
["hit", "Cache Hit"],
|
||||
["miss", "Cache Miss"],
|
||||
])("shows the human label on the Cache trigger for %s", async (cacheState, label) => {
|
||||
renderFilters(cacheState === "" ? {} : { [LOG_FILTER_IDS.CACHE_STATUS]: cacheState });
|
||||
|
||||
expect(await screen.findByText(label)).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it.each([
|
||||
["Cache Hit", "hit"],
|
||||
["Cache Miss", "miss"],
|
||||
])("selecting %s sets the cache filter to %s", async (label, expected) => {
|
||||
const user = userEvent.setup();
|
||||
const { set } = renderFilters();
|
||||
|
||||
await user.click(await screen.findByText("All Requests"));
|
||||
await user.click(await screen.findByRole("option", { name: label }));
|
||||
|
||||
expect(set).toHaveBeenCalledWith(LOG_FILTER_IDS.CACHE_STATUS, expected);
|
||||
});
|
||||
|
||||
it("selecting All Requests clears the cache filter", async () => {
|
||||
const user = userEvent.setup();
|
||||
const { set } = renderFilters({ [LOG_FILTER_IDS.CACHE_STATUS]: "hit" });
|
||||
|
||||
await user.click(await screen.findByText("Cache Hit"));
|
||||
await user.click(await screen.findByRole("option", { name: "All Requests" }));
|
||||
|
||||
expect(set).toHaveBeenCalledWith(LOG_FILTER_IDS.CACHE_STATUS, undefined);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -31,6 +31,12 @@ const STATUS_FILTER_ITEMS = [
|
|||
{ value: "success", label: "Success" },
|
||||
{ value: "failure", label: "Failure" },
|
||||
] as const;
|
||||
|
||||
const CACHE_FILTER_ITEMS = [
|
||||
{ value: ALL_VALUE, label: "All Requests" },
|
||||
{ value: "hit", label: "Cache Hit" },
|
||||
{ value: "miss", label: "Cache Miss" },
|
||||
] as const;
|
||||
const PAGE_SIZE = 50;
|
||||
|
||||
const asString = (value: unknown): string => (typeof value === "string" ? value : "");
|
||||
|
|
@ -328,6 +334,27 @@ export function RequestLogsFilters({ get, set, teams, logsWindow }: RequestLogsF
|
|||
</Select>
|
||||
</DataTableFilterField>
|
||||
|
||||
<DataTableFilterField label="Cache">
|
||||
<Select
|
||||
items={CACHE_FILTER_ITEMS}
|
||||
value={valueOf(LOG_FILTER_IDS.CACHE_STATUS) === "" ? ALL_VALUE : valueOf(LOG_FILTER_IDS.CACHE_STATUS)}
|
||||
onValueChange={(next) =>
|
||||
set(LOG_FILTER_IDS.CACHE_STATUS, next === null || next === ALL_VALUE ? undefined : next)
|
||||
}
|
||||
>
|
||||
<SelectTrigger className="w-full">
|
||||
<SelectValue placeholder="All Requests" />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
{CACHE_FILTER_ITEMS.map((item) => (
|
||||
<SelectItem key={item.value} value={item.value}>
|
||||
{item.label}
|
||||
</SelectItem>
|
||||
))}
|
||||
</SelectContent>
|
||||
</Select>
|
||||
</DataTableFilterField>
|
||||
|
||||
<KeyAliasFilterField
|
||||
value={valueOf(LOG_FILTER_IDS.KEY_ALIAS)}
|
||||
onChange={setter(LOG_FILTER_IDS.KEY_ALIAS)}
|
||||
|
|
|
|||
|
|
@ -434,6 +434,44 @@ describe("RequestLogsPanel", () => {
|
|||
});
|
||||
});
|
||||
|
||||
describe("hide health checks", () => {
|
||||
const toggle = () => screen.getByRole("switch", { name: "Hide Health Checks" });
|
||||
|
||||
it("defaults to showing health checks and refetches without them from page 1 when toggled on", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderPanel();
|
||||
|
||||
await waitFor(() => expect(uiSpendLogsCall).toHaveBeenCalled());
|
||||
expect(lastCall()?.params?.exclude_internal_health_checks).toBe(false);
|
||||
expect(toggle()).not.toBeChecked();
|
||||
|
||||
await user.click(toggle());
|
||||
|
||||
await waitFor(() => expect(lastCall()?.params?.exclude_internal_health_checks).toBe(true));
|
||||
expect(lastCall()?.page).toBe(1);
|
||||
expect(toggle()).toBeChecked();
|
||||
expect(sessionStorage.getItem("excludeInternalHealthChecks")).toBe("true");
|
||||
});
|
||||
|
||||
it("restores the persisted toggle from sessionStorage", async () => {
|
||||
sessionStorage.setItem("excludeInternalHealthChecks", "true");
|
||||
renderPanel();
|
||||
|
||||
await waitFor(() => expect(uiSpendLogsCall).toHaveBeenCalled());
|
||||
expect(lastCall()?.params?.exclude_internal_health_checks).toBe(true);
|
||||
expect(toggle()).toBeChecked();
|
||||
});
|
||||
|
||||
it("falls back to showing health checks when the persisted value is malformed", async () => {
|
||||
sessionStorage.setItem("excludeInternalHealthChecks", "{not json");
|
||||
renderPanel();
|
||||
|
||||
await waitFor(() => expect(uiSpendLogsCall).toHaveBeenCalled());
|
||||
expect(lastCall()?.params?.exclude_internal_health_checks).toBe(false);
|
||||
expect(toggle()).not.toBeChecked();
|
||||
});
|
||||
});
|
||||
|
||||
describe("live tail", () => {
|
||||
it("shows the auto-refresh banner on the first page and hides it once stopped", async () => {
|
||||
const user = userEvent.setup();
|
||||
|
|
|
|||
|
|
@ -72,6 +72,14 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID,
|
|||
sessionStorage.setItem("isLiveTail", JSON.stringify(isLiveTail));
|
||||
}, [isLiveTail]);
|
||||
|
||||
const [excludeInternalHealthChecks, setExcludeInternalHealthChecks] = useState<boolean>(
|
||||
() => sessionStorage.getItem("excludeInternalHealthChecks") === "true",
|
||||
);
|
||||
|
||||
useEffect(() => {
|
||||
sessionStorage.setItem("excludeInternalHealthChecks", JSON.stringify(excludeInternalHealthChecks));
|
||||
}, [excludeInternalHealthChecks]);
|
||||
|
||||
const { logsQuery, filteredLogs, allTeams } = useLogFilterLogic({
|
||||
accessToken,
|
||||
token,
|
||||
|
|
@ -80,6 +88,7 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID,
|
|||
columnFilters,
|
||||
activeTab: isActive ? "request logs" : "inactive",
|
||||
isLiveTail,
|
||||
excludeInternalHealthChecks,
|
||||
startTime,
|
||||
endTime,
|
||||
pagination,
|
||||
|
|
@ -219,6 +228,14 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID,
|
|||
setPagination((previous) => ({ ...previous, pageIndex: 0 }));
|
||||
}, []);
|
||||
|
||||
const handleExcludeInternalHealthChecksChange = useCallback(
|
||||
(value: boolean) => {
|
||||
setExcludeInternalHealthChecks(value);
|
||||
resetToFirstPage();
|
||||
},
|
||||
[resetToFirstPage],
|
||||
);
|
||||
|
||||
const handleResetFilters = useCallback(() => {
|
||||
setColumnFilters([]);
|
||||
setStartTime(moment().subtract(24, "hours").format("YYYY-MM-DDTHH:mm"));
|
||||
|
|
@ -313,6 +330,8 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID,
|
|||
onSelectedTimeIntervalChange={setSelectedTimeInterval}
|
||||
isLiveTail={isLiveTail}
|
||||
onIsLiveTailChange={setIsLiveTail}
|
||||
excludeInternalHealthChecks={excludeInternalHealthChecks}
|
||||
onExcludeInternalHealthChecksChange={handleExcludeInternalHealthChecksChange}
|
||||
onResetToFirstPage={resetToFirstPage}
|
||||
onResetFilters={handleResetFilters}
|
||||
/>
|
||||
|
|
|
|||
|
|
@ -47,6 +47,7 @@ const defaultProps = {
|
|||
columnFilters: [] as ColumnFiltersState,
|
||||
activeTab: "request logs",
|
||||
isLiveTail: false,
|
||||
excludeInternalHealthChecks: false,
|
||||
startTime: "2025-01-01T00:00:00",
|
||||
endTime: "2025-01-01T23:59:59",
|
||||
pagination: FIRST_PAGE,
|
||||
|
|
@ -82,6 +83,8 @@ describe("useLogFilterLogic", () => {
|
|||
{ id: LOG_FILTER_IDS.SESSION_ID, value: "sess-1", param: "session_id" },
|
||||
{ id: LOG_FILTER_IDS.END_USER, value: "end-user-1", param: "end_user" },
|
||||
{ id: LOG_FILTER_IDS.STATUS, value: "failure", param: "status_filter" },
|
||||
{ id: LOG_FILTER_IDS.CACHE_STATUS, value: "hit", param: "cache_hit_filter" },
|
||||
{ id: LOG_FILTER_IDS.CACHE_STATUS, value: "miss", param: "cache_hit_filter" },
|
||||
{ id: LOG_FILTER_IDS.MODEL_ID, value: "model-uuid-1", param: "model_id" },
|
||||
{ id: LOG_FILTER_IDS.PUBLIC_MODEL_OR_SEARCH_TOOL, value: "gpt-4o", param: "model" },
|
||||
{ id: LOG_FILTER_IDS.KEY_ALIAS, value: "alias-1", param: "key_alias" },
|
||||
|
|
@ -152,6 +155,7 @@ describe("useLogFilterLogic", () => {
|
|||
["pagination", { pagination: { pageIndex: 1, pageSize: 50 } }],
|
||||
["startTime", { startTime: "2025-02-02T00:00:00" }],
|
||||
["columnFilters", { columnFilters: [{ id: LOG_FILTER_IDS.TEAM_ID, value: "team-2" }] }],
|
||||
["excludeInternalHealthChecks", { excludeInternalHealthChecks: true }],
|
||||
])("refetches when %s changes", async (_label, nextProps) => {
|
||||
const { rerender } = renderHook((props: HookOverrides) => useLogFilterLogic({ ...defaultProps, ...props }), {
|
||||
wrapper,
|
||||
|
|
@ -164,6 +168,22 @@ describe("useLogFilterLogic", () => {
|
|||
});
|
||||
});
|
||||
|
||||
describe("hide health checks toggle", () => {
|
||||
it("passes exclude_internal_health_checks when the toggle is on", async () => {
|
||||
renderFilterHook({ excludeInternalHealthChecks: true });
|
||||
|
||||
await waitFor(() => expect(uiSpendLogsCall).toHaveBeenCalled());
|
||||
expect(lastCallParams()?.params).toMatchObject({ exclude_internal_health_checks: true });
|
||||
});
|
||||
|
||||
it("passes exclude_internal_health_checks as false when the toggle is off", async () => {
|
||||
renderFilterHook();
|
||||
|
||||
await waitFor(() => expect(uiSpendLogsCall).toHaveBeenCalled());
|
||||
expect(lastCallParams()?.params).toMatchObject({ exclude_internal_health_checks: false });
|
||||
});
|
||||
});
|
||||
|
||||
describe("query enablement", () => {
|
||||
it("does not query when the request logs tab is inactive", async () => {
|
||||
renderFilterHook({ activeTab: "audit logs" });
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ export interface PaginatedResponse {
|
|||
export const LOG_FILTER_IDS = {
|
||||
TEAM_ID: "team_id",
|
||||
STATUS: "status",
|
||||
CACHE_STATUS: "cache_hit",
|
||||
KEY_ALIAS: "key_alias",
|
||||
END_USER: "end_user",
|
||||
ERROR_CODE: "error_code",
|
||||
|
|
@ -35,6 +36,7 @@ export const LOG_FILTER_IDS = {
|
|||
export const LOG_FILTER_LABELS: Record<string, string> = {
|
||||
[LOG_FILTER_IDS.TEAM_ID]: "Team ID",
|
||||
[LOG_FILTER_IDS.STATUS]: "Status",
|
||||
[LOG_FILTER_IDS.CACHE_STATUS]: "Cache",
|
||||
[LOG_FILTER_IDS.KEY_ALIAS]: "Key Alias",
|
||||
[LOG_FILTER_IDS.USER_ID]: "User ID",
|
||||
[LOG_FILTER_IDS.END_USER]: "End User",
|
||||
|
|
@ -101,6 +103,7 @@ export function useLogFilterLogic({
|
|||
columnFilters,
|
||||
activeTab,
|
||||
isLiveTail,
|
||||
excludeInternalHealthChecks,
|
||||
startTime,
|
||||
endTime,
|
||||
pagination,
|
||||
|
|
@ -114,6 +117,7 @@ export function useLogFilterLogic({
|
|||
columnFilters: ColumnFiltersState;
|
||||
activeTab: string;
|
||||
isLiveTail: boolean;
|
||||
excludeInternalHealthChecks: boolean;
|
||||
startTime: string;
|
||||
endTime: string;
|
||||
pagination: PaginationState;
|
||||
|
|
@ -137,6 +141,7 @@ export function useLogFilterLogic({
|
|||
columnFilters,
|
||||
sortBy,
|
||||
sortOrder,
|
||||
excludeInternalHealthChecks,
|
||||
],
|
||||
queryFn: async () => {
|
||||
if (!accessToken || !token || !userRole || !userID) {
|
||||
|
|
@ -167,6 +172,7 @@ export function useLogFilterLogic({
|
|||
user_id: userIdFilter,
|
||||
end_user: getFilterValue(columnFilters, LOG_FILTER_IDS.END_USER),
|
||||
status_filter: getFilterValue(columnFilters, LOG_FILTER_IDS.STATUS),
|
||||
cache_hit_filter: getFilterValue(columnFilters, LOG_FILTER_IDS.CACHE_STATUS),
|
||||
model_id: getFilterValue(columnFilters, LOG_FILTER_IDS.MODEL_ID),
|
||||
model: getFilterValue(columnFilters, LOG_FILTER_IDS.PUBLIC_MODEL_OR_SEARCH_TOOL),
|
||||
key_alias: getFilterValue(columnFilters, LOG_FILTER_IDS.KEY_ALIAS),
|
||||
|
|
@ -174,6 +180,7 @@ export function useLogFilterLogic({
|
|||
error_message: getFilterValue(columnFilters, LOG_FILTER_IDS.ERROR_MESSAGE),
|
||||
sort_by: sortBy,
|
||||
sort_order: sortOrder,
|
||||
exclude_internal_health_checks: excludeInternalHealthChecks,
|
||||
},
|
||||
});
|
||||
},
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import {
|
|||
buildPresetPrefill,
|
||||
buildModelAvailability,
|
||||
deploymentRefsFromModelInfo,
|
||||
normalizeModelName,
|
||||
} from "./autorouter_presets";
|
||||
import { DEFAULT_MATCH_THRESHOLD } from "@/components/add_model/SemanticKeywordMatching";
|
||||
import { DEFAULT_ESCALATION_KEYWORDS } from "@/components/add_model/EscalationKeywords";
|
||||
|
|
@ -28,6 +29,29 @@ describe("autorouter_presets", () => {
|
|||
}
|
||||
});
|
||||
|
||||
// buildPresetPrefill resolves every model reference through normalizeModelName, so two spellings
|
||||
// of the same model in one tier (e.g. "claude-sonnet-4-5" and "claude-sonnet-4.5") collapse to one
|
||||
// key. For tier_model_configs that silently drops one model's litellm_params; catch it in the
|
||||
// bundled data itself, since nothing else validates preset authoring.
|
||||
it("never spells the same model two ways within a single tier", () => {
|
||||
for (const preset of getAllPresets()) {
|
||||
const { tiers, tier_model_configs: configs } = preset.complexity_router_config;
|
||||
for (const tier of Object.keys(tiers) as (keyof typeof tiers)[]) {
|
||||
const fromTierList = tiers[tier] ?? [];
|
||||
const fromConfigs = (configs?.[tier] ?? []).map((entry) => entry.model_name);
|
||||
const names = new Set([...fromTierList, ...fromConfigs]);
|
||||
const byNormalized = new Map<string, string[]>();
|
||||
for (const name of names) {
|
||||
const key = normalizeModelName(name);
|
||||
byNormalized.set(key, [...(byNormalized.get(key) ?? []), name]);
|
||||
}
|
||||
for (const spellings of byNormalized.values()) {
|
||||
expect(new Set(spellings).size, `${preset.key}.${tier}: ${spellings.join(", ")}`).toBe(1);
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
it("resolves a preset by its stable JSON key, not its display label", () => {
|
||||
expect(getPresetByKey("anthropic_family")?.label).toBe("Anthropic Family");
|
||||
expect(getPresetByKey("does_not_exist")).toBeUndefined();
|
||||
|
|
@ -544,5 +568,73 @@ describe("autorouter_presets", () => {
|
|||
const prefill = buildPresetPrefill(config, groupsOnly(["claude-sonnet-4.5"]));
|
||||
expect(prefill.complexityRouterConfig.tiers.SIMPLE).toEqual(["claude-sonnet-4.5"]);
|
||||
});
|
||||
|
||||
it("prefills the per-model litellm_params a preset carries in tier_model_configs", () => {
|
||||
const config = {
|
||||
tiers: { SIMPLE: ["gpt-5-nano"], MEDIUM: [], COMPLEX: [], REASONING: ["o3"] },
|
||||
tier_model_configs: {
|
||||
REASONING: [{ model_name: "o3", litellm_params: { reasoning_effort: "high" } }],
|
||||
},
|
||||
classifier_type: "heuristic" as const,
|
||||
session_affinity: false,
|
||||
deployment_affinity: true,
|
||||
};
|
||||
const prefill = buildPresetPrefill(config, groupsOnly(["gpt-5-nano", "o3"]));
|
||||
expect(prefill.complexityRouterConfig.tier_model_params).toEqual({
|
||||
REASONING: { o3: { reasoning_effort: "high" } },
|
||||
});
|
||||
});
|
||||
|
||||
// The params key on the preset's own spelling while the tier entry gets rewritten to the
|
||||
// caller's. Leaving the key alone names a model the tier no longer holds, and
|
||||
// serializeTierModelConfigs then drops the params on submit without saying so.
|
||||
it("rewrites a param key to the same registered spelling its tier entry was rewritten to", () => {
|
||||
const config = {
|
||||
tiers: { SIMPLE: [], MEDIUM: [], COMPLEX: [], REASONING: ["claude-sonnet-4-5"] },
|
||||
tier_model_configs: {
|
||||
REASONING: [{ model_name: "claude-sonnet-4-5", litellm_params: { reasoning_effort: "high" } }],
|
||||
},
|
||||
classifier_type: "heuristic" as const,
|
||||
session_affinity: false,
|
||||
deployment_affinity: true,
|
||||
};
|
||||
const prefill = buildPresetPrefill(config, groupsOnly(["claude-sonnet-4.5"]));
|
||||
expect(prefill.complexityRouterConfig.tier_model_params).toEqual({
|
||||
REASONING: { "claude-sonnet-4.5": { reasoning_effort: "high" } },
|
||||
});
|
||||
});
|
||||
|
||||
// Two spellings of one model in a tier collapse to a single registered key, and one model can
|
||||
// only hold one param set downstream. Merging keeps whatever only one spelling set instead of
|
||||
// dropping that spelling's params wholesale.
|
||||
it("merges rather than drops params when two spellings resolve to the same registered model", () => {
|
||||
const config = {
|
||||
tiers: { SIMPLE: [], MEDIUM: [], COMPLEX: [], REASONING: ["claude-sonnet-4-5", "claude-sonnet-4.5"] },
|
||||
tier_model_configs: {
|
||||
REASONING: [
|
||||
{ model_name: "claude-sonnet-4-5", litellm_params: { reasoning_effort: "high", temperature: 0.2 } },
|
||||
{ model_name: "claude-sonnet-4.5", litellm_params: { reasoning_effort: "low" } },
|
||||
],
|
||||
},
|
||||
classifier_type: "heuristic" as const,
|
||||
};
|
||||
const prefill = buildPresetPrefill(config, groupsOnly(["claude-sonnet-4.5"]));
|
||||
// temperature survives from the spelling that would otherwise have been overwritten;
|
||||
// reasoning_effort, set by both, resolves last-wins.
|
||||
expect(prefill.complexityRouterConfig.tier_model_params).toEqual({
|
||||
REASONING: { "claude-sonnet-4.5": { reasoning_effort: "low", temperature: 0.2 } },
|
||||
});
|
||||
});
|
||||
|
||||
it("leaves tier_model_params undefined for a preset that carries no per-model params", () => {
|
||||
const config = {
|
||||
tiers: { SIMPLE: ["gpt-5-nano"], MEDIUM: [], COMPLEX: [], REASONING: [] },
|
||||
classifier_type: "heuristic" as const,
|
||||
session_affinity: false,
|
||||
deployment_affinity: true,
|
||||
};
|
||||
const prefill = buildPresetPrefill(config, groupsOnly(["gpt-5-nano"]));
|
||||
expect(prefill.complexityRouterConfig.tier_model_params).toBeUndefined();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -8,9 +8,15 @@ import {
|
|||
ClassifierLLMConfig,
|
||||
DEFAULT_SESSION_AFFINITY,
|
||||
DEFAULT_DEPLOYMENT_AFFINITY,
|
||||
usesLlmClassifier,
|
||||
} from "@/components/add_model/ComplexityRouterConfig";
|
||||
import { KeywordTierRule } from "@/components/add_model/KeywordTierRules";
|
||||
import { hydrateKeywordTierRules } from "@/components/add_model/complexity_router_keywords";
|
||||
import {
|
||||
TierModelParams,
|
||||
TierModelParamsByTier,
|
||||
hydrateTierModelParams,
|
||||
} from "@/components/add_model/complexity_router_tiers";
|
||||
import { DEFAULT_ESCALATION_KEYWORDS } from "@/components/add_model/EscalationKeywords";
|
||||
import { DEFAULT_MATCH_THRESHOLD } from "@/components/add_model/SemanticKeywordMatching";
|
||||
import presetsRaw from "@/autorouter_presets.json";
|
||||
|
|
@ -53,7 +59,7 @@ export const getRequiredModels = (
|
|||
// differing only in that separator. Canonicalizing on "-" (the presets' own convention) lets both
|
||||
// spellings match without doing anything looser - two DIFFERENT model names never collide here,
|
||||
// only the punctuation within one version number does.
|
||||
const normalizeModelName = (model: string): string => model.replace(/(\d)\.(\d)/g, "$1-$2");
|
||||
export const normalizeModelName = (model: string): string => model.replace(/(\d)\.(\d)/g, "$1-$2");
|
||||
|
||||
export interface DeploymentModelRef {
|
||||
modelGroup: string;
|
||||
|
|
@ -177,7 +183,7 @@ export const getMissingModelsInPreset = (preset: AutoRouterPreset, availability:
|
|||
// Checks the config actually being built (whether it arrived via a preset prefill or was typed by
|
||||
// hand - the two are indistinguishable once the caller has started editing), not a preset's
|
||||
// original bundled model list. Only counts classifier_llm_config/embedding_model as referenced
|
||||
// when buildComplexityRouterConfig would actually emit them (classifierType === "llm",
|
||||
// when buildComplexityRouterConfig would actually emit them (usesLlmClassifier(classifierType),
|
||||
// semanticMatchingEnabled) - otherwise a dormant selection left over from a toggle no longer in
|
||||
// effect would block submit for a model that was never going to be submitted.
|
||||
export const getReferencedModelsError = (
|
||||
|
|
@ -195,7 +201,7 @@ export const getReferencedModelsError = (
|
|||
{
|
||||
tiers: params.tiers,
|
||||
default_model: params.defaultModel,
|
||||
classifier_llm_config: params.classifierType === "llm" ? params.classifierLlmConfig : undefined,
|
||||
classifier_llm_config: usesLlmClassifier(params.classifierType) ? params.classifierLlmConfig : undefined,
|
||||
embedding_model: params.semanticMatchingEnabled ? params.embeddingModel : undefined,
|
||||
},
|
||||
availability,
|
||||
|
|
@ -243,6 +249,25 @@ export const buildPresetPrefill = (
|
|||
): PresetPrefill => {
|
||||
const resolve = (model: string): string => resolveAvailableModel(model, availability) ?? model;
|
||||
const resolveTier = (models: string[]): string[] => models.map(resolve);
|
||||
// Params key on the model name the preset spells while every tier entry is rewritten to the
|
||||
// caller's registered spelling, so the keys have to be rewritten the same way. Otherwise
|
||||
// serializeTierModelConfigs drops them for naming a model the tier no longer holds.
|
||||
//
|
||||
// Two spellings in one tier can resolve to the same registered model, and one model holds one
|
||||
// param set here and in the payload, so a collision has to collapse. Merge rather than replace:
|
||||
// params only one spelling set still survive, and a key both set resolves last-wins, matching
|
||||
// how hydrateTierModelParams already collapses two entries spelled identically.
|
||||
const resolveParamKeys = (params: TierModelParamsByTier | undefined): TierModelParamsByTier | undefined =>
|
||||
params &&
|
||||
Object.fromEntries(
|
||||
Object.entries(params).map(([tier, byModel]) => [
|
||||
tier,
|
||||
Object.entries(byModel).reduce<Record<string, TierModelParams>>((byResolved, [model, litellmParams]) => {
|
||||
const resolved = resolve(model);
|
||||
return { ...byResolved, [resolved]: { ...byResolved[resolved], ...litellmParams } };
|
||||
}, {}),
|
||||
]),
|
||||
);
|
||||
|
||||
return {
|
||||
complexityRouterConfig: {
|
||||
|
|
@ -252,6 +277,7 @@ export const buildPresetPrefill = (
|
|||
COMPLEX: resolveTier(config.tiers.COMPLEX),
|
||||
REASONING: resolveTier(config.tiers.REASONING),
|
||||
},
|
||||
tier_model_params: resolveParamKeys(hydrateTierModelParams(config.tiers, config.tier_model_configs)),
|
||||
tier_labels: hydrateTierLabels(config.tier_labels),
|
||||
classifier_type: config.classifier_type,
|
||||
classifier_llm_config: config.classifier_llm_config && {
|
||||
|
|
|
|||
2967
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
2967
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
File diff suppressed because it is too large
Load diff
File diff suppressed because one or more lines are too long
Loading…
Add table
Reference in a new issue