mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
df13d59306
* fix(logging): skip sync success callbacks for internal sub-calls Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(langfuse): assert no retry sleep instead of a wall-clock budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: add types to deflake regression coverage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(tests): avoid unrelated formatting changes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
733 lines
30 KiB
Python
733 lines
30 KiB
Python
"""
|
|
Tests for the RAG query pipeline in litellm/rag/main.py.
|
|
|
|
The RAG pipeline forwards its kwargs (including the parent litellm_logging_obj)
|
|
into @client-decorated sub-calls (vector store search, completion). Each logging
|
|
object allows exactly one async_success event, so if sub-calls are not marked as
|
|
internal, the vector store search consumes the slot first and the LLM
|
|
completion's usage/cost is never logged (spend tracking and budget enforcement
|
|
are bypassed). These tests pin the invariant that the single billing event for
|
|
aquery carries the completion response with real usage and cost.
|
|
"""
|
|
|
|
import asyncio
|
|
import datetime
|
|
import json
|
|
from collections.abc import Callable, Mapping
|
|
from concurrent.futures import Future
|
|
from types import MappingProxyType
|
|
from typing import Final, cast
|
|
from unittest.mock import patch
|
|
|
|
import httpx
|
|
import pytest
|
|
import respx
|
|
from pydantic import ValidationError
|
|
|
|
import litellm
|
|
import litellm.utils as litellm_utils
|
|
from litellm._internal_context import is_internal_call
|
|
from litellm.integrations.custom_logger import CustomLogger
|
|
from litellm.litellm_core_utils import litellm_logging
|
|
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
|
from litellm.types.utils import CallTypes, ModelResponse
|
|
|
|
|
|
async def _drain_logging_worker() -> None:
|
|
"""Run every queued logging task to completion on the current event loop.
|
|
|
|
The success event is delivered through the fire-and-forget GLOBAL_LOGGING_WORKER
|
|
singleton, whose queue survives across tests. start() rebinds any tasks left over
|
|
from a previous test's event loop onto the current one, and flush() waits until
|
|
the queue is fully processed, so tests neither miss their own event nor observe
|
|
a neighbour's
|
|
"""
|
|
await asyncio.sleep(0)
|
|
GLOBAL_LOGGING_WORKER.start()
|
|
await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0)
|
|
|
|
|
|
class RecordingLogger(CustomLogger):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.success_events: Final[list[dict[str, object]]] = []
|
|
self.sync_success_events: Final[list[dict[str, object]]] = []
|
|
|
|
async def async_log_success_event(
|
|
self,
|
|
kwargs: Mapping[str, object],
|
|
response_obj: object,
|
|
start_time: datetime.datetime,
|
|
end_time: datetime.datetime,
|
|
) -> None:
|
|
self.success_events.append({"kwargs": kwargs, "response_obj": response_obj})
|
|
|
|
def log_success_event(
|
|
self,
|
|
kwargs: Mapping[str, object],
|
|
response_obj: object,
|
|
start_time: datetime.datetime,
|
|
end_time: datetime.datetime,
|
|
) -> None:
|
|
self.sync_success_events.append({"kwargs": kwargs, "response_obj": response_obj})
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("use_router", [False, True])
|
|
async def test_aquery_single_billing_event_carries_completion_usage_and_cost(use_router):
|
|
"""
|
|
litellm.aquery must produce exactly one success event, and that event must
|
|
carry the LLM completion (a ModelResponse with non-zero usage and cost),
|
|
not the vector store search response. The proxy always passes a router, so
|
|
both the router and non-router completion branches are pinned.
|
|
"""
|
|
await _drain_logging_worker()
|
|
recording_logger = RecordingLogger()
|
|
original_callbacks = litellm.callbacks
|
|
litellm.callbacks = [recording_logger]
|
|
|
|
router_kwargs = {}
|
|
if use_router:
|
|
router_kwargs["router"] = litellm.Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "gpt-4o-mini",
|
|
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"},
|
|
}
|
|
]
|
|
)
|
|
|
|
try:
|
|
response = await litellm.aquery(
|
|
model="gpt-4o-mini",
|
|
messages=[{"role": "user", "content": "What is the secret project codename?"}],
|
|
retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"},
|
|
mock_response="The secret project codename is AZURE-FALCON-42.",
|
|
**router_kwargs,
|
|
)
|
|
|
|
assert isinstance(response, ModelResponse)
|
|
assert is_internal_call.get() is False
|
|
|
|
await _drain_logging_worker()
|
|
finally:
|
|
litellm.callbacks = original_callbacks
|
|
|
|
assert len(recording_logger.success_events) == 1
|
|
event = recording_logger.success_events[0]
|
|
|
|
response_obj = event["response_obj"]
|
|
assert isinstance(response_obj, ModelResponse)
|
|
assert response_obj.usage.total_tokens > 0
|
|
|
|
standard_logging_object = event["kwargs"]["standard_logging_object"]
|
|
assert standard_logging_object["call_type"] == "aquery"
|
|
assert standard_logging_object["total_tokens"] > 0
|
|
assert standard_logging_object["prompt_tokens"] > 0
|
|
assert standard_logging_object["completion_tokens"] > 0
|
|
assert standard_logging_object["response_cost"] > 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("use_router", [False, True])
|
|
async def test_aquery_vector_store_search_sub_call_logs_no_sync_success_event_on_the_parent(
|
|
use_router: bool, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
await _drain_logging_worker()
|
|
recording_logger: Final = RecordingLogger()
|
|
|
|
class InlineExecutor:
|
|
def submit(
|
|
self,
|
|
fn: Callable[..., object],
|
|
*args: object,
|
|
**kwargs: object,
|
|
) -> Future[object]:
|
|
future: Final = Future[object]()
|
|
future.set_result(fn(*args, **kwargs))
|
|
return future
|
|
|
|
monkeypatch.setattr(litellm_utils, "executor", InlineExecutor())
|
|
monkeypatch.setattr(litellm_logging, "executor", InlineExecutor())
|
|
monkeypatch.setattr(litellm, "callbacks", [recording_logger])
|
|
|
|
router_kwargs: Final = (
|
|
{
|
|
"router": litellm.Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "gpt-4o-mini",
|
|
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"},
|
|
}
|
|
]
|
|
)
|
|
}
|
|
if use_router
|
|
else {}
|
|
)
|
|
|
|
response: Final = await litellm.aquery(
|
|
model="gpt-4o-mini",
|
|
messages=[{"role": "user", "content": "What is the secret project codename?"}],
|
|
retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"},
|
|
mock_response="The secret project codename is AZURE-FALCON-42.",
|
|
**router_kwargs,
|
|
)
|
|
assert isinstance(response, ModelResponse), "aquery should return its completion response"
|
|
assert is_internal_call.get() is False, "aquery should restore the internal-call context"
|
|
await _drain_logging_worker()
|
|
|
|
assert len(recording_logger.success_events) == 1, "aquery should run one async success callback"
|
|
event: Final = recording_logger.success_events[0]
|
|
response_obj: Final = event["response_obj"]
|
|
assert isinstance(response_obj, ModelResponse), "the async callback should receive the completion response"
|
|
assert response_obj.usage.total_tokens > 0, "the async callback response should include completion usage"
|
|
|
|
event_kwargs: Final = cast(Mapping[str, object], event["kwargs"])
|
|
standard_logging_object: Final = cast(Mapping[str, object], event_kwargs["standard_logging_object"])
|
|
assert cast(str, standard_logging_object["call_type"]) == "aquery", "the async event should be for aquery"
|
|
assert cast(int, standard_logging_object["total_tokens"]) > 0, (
|
|
"the async event should include total completion tokens"
|
|
)
|
|
assert cast(int, standard_logging_object["prompt_tokens"]) > 0, "the async event should include prompt tokens"
|
|
assert cast(int, standard_logging_object["completion_tokens"]) > 0, (
|
|
"the async event should include completion tokens"
|
|
)
|
|
assert cast(float, standard_logging_object["response_cost"]) > 0, "the async event should include completion cost"
|
|
sync_response_types: Final = [
|
|
type(event["response_obj"]).__name__ for event in recording_logger.sync_success_events
|
|
]
|
|
assert recording_logger.sync_success_events == [], (
|
|
"an internal sub-call must not log on the parent's logging object; "
|
|
f"sync success event response types: {sync_response_types}"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aquery_response_hidden_params_carry_completion_cost():
|
|
"""
|
|
The aquery response must expose the completion's response_cost via hidden
|
|
params, so the proxy can return the x-litellm-response-cost header.
|
|
"""
|
|
response = await litellm.aquery(
|
|
model="gpt-4o-mini",
|
|
messages=[{"role": "user", "content": "hello"}],
|
|
retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"},
|
|
mock_response="hi there",
|
|
)
|
|
|
|
await _drain_logging_worker()
|
|
|
|
assert isinstance(response, ModelResponse)
|
|
response_cost = response._hidden_params.get("response_cost")
|
|
assert response_cost is not None
|
|
assert response_cost > 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aquery_billed_cost_includes_priced_vector_store_search():
|
|
"""
|
|
When the vector store provider prices search calls (e.g. per-query cost),
|
|
that cost must be folded into the aquery billing instead of being dropped
|
|
with the suppressed sub-call event.
|
|
"""
|
|
await _drain_logging_worker()
|
|
recording_logger = RecordingLogger()
|
|
original_callbacks = litellm.callbacks
|
|
litellm.callbacks = [recording_logger]
|
|
|
|
try:
|
|
with patch("litellm.rag.main.vector_store_search_cost", return_value=(0.002, 0.0)):
|
|
response = await litellm.aquery(
|
|
model="gpt-4o-mini",
|
|
messages=[{"role": "user", "content": "hello"}],
|
|
retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"},
|
|
mock_response="hi there",
|
|
)
|
|
|
|
await _drain_logging_worker()
|
|
finally:
|
|
litellm.callbacks = original_callbacks
|
|
|
|
assert isinstance(response, ModelResponse)
|
|
total_cost = response._hidden_params.get("response_cost")
|
|
assert total_cost is not None
|
|
assert total_cost > 0.002
|
|
|
|
assert len(recording_logger.success_events) == 1
|
|
standard_logging_object = recording_logger.success_events[0]["kwargs"]["standard_logging_object"]
|
|
assert standard_logging_object["response_cost"] == total_cost
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aquery_with_rerank_bills_once_and_folds_rerank_cost():
|
|
"""
|
|
When rerank is enabled, its sub-call must run under the internal-call
|
|
context (no standalone billing event) and its cost must be folded into
|
|
the single aquery billing event.
|
|
"""
|
|
from litellm.types.rerank import RerankResponse
|
|
|
|
await _drain_logging_worker()
|
|
recording_logger = RecordingLogger()
|
|
original_callbacks = litellm.callbacks
|
|
litellm.callbacks = [recording_logger]
|
|
rerank_seen = {}
|
|
|
|
async def fake_arerank(**kwargs):
|
|
rerank_seen["internal"] = is_internal_call.get()
|
|
rerank_result = RerankResponse(id="rr_1", results=[{"index": 0, "relevance_score": 0.9}], meta={})
|
|
rerank_result._hidden_params["response_cost"] = 0.001
|
|
return rerank_result
|
|
|
|
try:
|
|
with patch("litellm.arerank", side_effect=fake_arerank):
|
|
response = await litellm.aquery(
|
|
model="gpt-4o-mini",
|
|
messages=[{"role": "user", "content": "hello"}],
|
|
retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"},
|
|
rerank={"enabled": True, "model": "cohere/rerank-english-v3.0", "top_n": 1},
|
|
mock_response="hi there",
|
|
)
|
|
|
|
await _drain_logging_worker()
|
|
finally:
|
|
litellm.callbacks = original_callbacks
|
|
|
|
assert rerank_seen["internal"] is True
|
|
assert is_internal_call.get() is False
|
|
|
|
assert isinstance(response, ModelResponse)
|
|
total_cost = response._hidden_params.get("response_cost")
|
|
assert total_cost is not None
|
|
assert total_cost > 0.001
|
|
|
|
assert len(recording_logger.success_events) == 1
|
|
standard_logging_object = recording_logger.success_events[0]["kwargs"]["standard_logging_object"]
|
|
assert standard_logging_object["call_type"] == "aquery"
|
|
assert standard_logging_object["response_cost"] == total_cost
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aquery_streaming_bills_sub_call_costs_into_final_event():
|
|
"""
|
|
On the streaming path the response cost is computed from the assembled
|
|
chunks after the pipeline returns, so there is no response object to fold
|
|
sub-call costs into. The pipeline must instead carry the accumulated
|
|
search and rerank cost through the logging object so the single streamed
|
|
billing event includes it; otherwise a caller passing stream=true incurs
|
|
priced vector search and rerank costs that never reach spend tracking.
|
|
"""
|
|
from litellm.types.rerank import RerankResponse
|
|
|
|
await _drain_logging_worker()
|
|
recording_logger = RecordingLogger()
|
|
original_callbacks = litellm.callbacks
|
|
litellm.callbacks = [recording_logger]
|
|
rerank_seen = {}
|
|
|
|
async def fake_arerank(**kwargs):
|
|
rerank_seen["internal"] = is_internal_call.get()
|
|
rerank_result = RerankResponse(id="rr_1", results=[{"index": 0, "relevance_score": 0.9}], meta={})
|
|
rerank_result._hidden_params["response_cost"] = 0.001
|
|
return rerank_result
|
|
|
|
try:
|
|
with (
|
|
patch("litellm.rag.main.vector_store_search_cost", return_value=(0.002, 0.0)),
|
|
patch("litellm.arerank", side_effect=fake_arerank),
|
|
):
|
|
response = await litellm.aquery(
|
|
model="gpt-4o-mini",
|
|
messages=[{"role": "user", "content": "hello"}],
|
|
retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"},
|
|
rerank={"enabled": True, "model": "cohere/rerank-english-v3.0", "top_n": 1},
|
|
mock_response="hi there",
|
|
stream=True,
|
|
)
|
|
async for _ in response:
|
|
pass
|
|
|
|
await _drain_logging_worker()
|
|
finally:
|
|
litellm.callbacks = original_callbacks
|
|
|
|
assert rerank_seen["internal"] is True
|
|
assert is_internal_call.get() is False
|
|
|
|
assert len(recording_logger.success_events) == 1
|
|
standard_logging_object = recording_logger.success_events[0]["kwargs"]["standard_logging_object"]
|
|
assert standard_logging_object["call_type"] == "aquery"
|
|
assert standard_logging_object["response_cost"] >= 0.003
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("retrieval_config_json", "top_level_filter_json", "expected_filter_json"),
|
|
(
|
|
(
|
|
'{"vector_store_id":"vs_test_123","custom_llm_provider":"openai","top_k":50,'
|
|
'"retrieval_filter":{"equals":{"key":"tenant","value":"retrieval"}}}',
|
|
None,
|
|
'{"equals":{"key":"tenant","value":"retrieval"}}',
|
|
),
|
|
(
|
|
'{"vector_store_id":"vs_test_123","custom_llm_provider":"openai","top_k":50,'
|
|
'"filters":{"equals":{"key":"tenant","value":"alias"}}}',
|
|
None,
|
|
'{"equals":{"key":"tenant","value":"alias"}}',
|
|
),
|
|
(
|
|
'{"vector_store_id":"vs_test_123","custom_llm_provider":"openai","top_k":50}',
|
|
'{"equals":{"key":"tenant","value":"top-level"}}',
|
|
'{"equals":{"key":"tenant","value":"top-level"}}',
|
|
),
|
|
(
|
|
'{"vector_store_id":"vs_test_123","custom_llm_provider":"openai","top_k":50,'
|
|
'"retrieval_filter":{"equals":{"key":"tenant","value":"retrieval"}},'
|
|
'"filters":{"equals":{"key":"tenant","value":"alias"}}}',
|
|
'{"equals":{"key":"tenant","value":"top-level"}}',
|
|
'{"equals":{"key":"tenant","value":"retrieval"}}',
|
|
),
|
|
(
|
|
'{"vector_store_id":"vs_test_123","custom_llm_provider":"openai","top_k":50}',
|
|
None,
|
|
None,
|
|
),
|
|
),
|
|
)
|
|
async def test_aquery_forwards_filters_to_vector_store_search(
|
|
retrieval_config_json: str,
|
|
top_level_filter_json: str | None,
|
|
expected_filter_json: str | None,
|
|
monkeypatch,
|
|
):
|
|
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
|
retrieval_config: Final = json.loads(retrieval_config_json)
|
|
top_level_filter: Final = json.loads(top_level_filter_json) if top_level_filter_json is not None else None
|
|
expected_filter: Final = json.loads(expected_filter_json) if expected_filter_json is not None else None
|
|
|
|
with respx.mock(assert_all_called=True) as respx_mock:
|
|
search_route: Final = respx_mock.post("https://example.com/v1/vector_stores/vs_test_123/search").mock(
|
|
return_value=httpx.Response(
|
|
200,
|
|
content='{"object":"vector_store.search_results.page","search_query":"q","data":[]}',
|
|
)
|
|
)
|
|
respx_mock.post("https://example.com/v1/chat/completions").mock(
|
|
return_value=httpx.Response(
|
|
200,
|
|
content=(
|
|
'{"id":"chatcmpl-test","object":"chat.completion","created":1,"model":"gpt-4o-mini",'
|
|
'"choices":[{"index":0,"message":{"role":"assistant","content":"answer"},"finish_reason":"stop"}],'
|
|
'"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}'
|
|
),
|
|
)
|
|
)
|
|
response: Final = await litellm.aquery(
|
|
model="openai/gpt-4o-mini",
|
|
messages=json.loads('[{"role":"user","content":"most frequent causes of low nicotine"}]'),
|
|
retrieval_config=retrieval_config,
|
|
filters=top_level_filter,
|
|
api_key="sk-test",
|
|
api_base="https://example.com/v1",
|
|
)
|
|
request_body: Final = json.loads(search_route.calls.last.request.content)
|
|
|
|
assert isinstance(response, ModelResponse)
|
|
assert response.choices[0].message.content == "answer"
|
|
assert request_body["query"] == "most frequent causes of low nicotine"
|
|
assert request_body.get("filters") == expected_filter
|
|
assert request_body["max_num_results"] == 50
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aquery_forwards_provider_retrieval_config_and_router_to_search():
|
|
"""
|
|
Regression: provider-specific retrieval_config keys (aws_region_name,
|
|
embedding_model, vector_bucket_name, ...) and the router must be forwarded
|
|
to the vector store search call. Pre-fix they were silently dropped, so
|
|
/v1/rag/query failed with provider config errors (e.g. S3 Vectors
|
|
"aws_region_name is required") even when the caller supplied them.
|
|
"""
|
|
from unittest.mock import AsyncMock
|
|
|
|
from litellm.types.vector_stores import VectorStoreSearchResponse
|
|
|
|
router = litellm.Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "gpt-4o-mini",
|
|
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"},
|
|
}
|
|
]
|
|
)
|
|
|
|
fake_search = AsyncMock(
|
|
return_value=VectorStoreSearchResponse(
|
|
object="vector_store.search_results.page", search_query="q", data=[]
|
|
)
|
|
)
|
|
with patch("litellm.vector_stores.asearch", new=fake_search): # test-quality-ok: asearch is the boundary the forwarding contract under test targets
|
|
response = await litellm.aquery(
|
|
model="gpt-4o-mini",
|
|
messages=[{"role": "user", "content": "hello"}],
|
|
retrieval_config={
|
|
"vector_store_id": "bkt:idx",
|
|
"custom_llm_provider": "s3_vectors",
|
|
"top_k": 5,
|
|
"aws_region_name": "eu-west-1",
|
|
"embedding_model": "my-embed",
|
|
"vector_bucket_name": "bkt",
|
|
},
|
|
router=router,
|
|
mock_response="hi",
|
|
)
|
|
|
|
assert isinstance(response, ModelResponse)
|
|
fake_search.assert_awaited_once()
|
|
search_kwargs = fake_search.await_args.kwargs
|
|
assert search_kwargs["vector_store_id"] == "bkt:idx"
|
|
assert search_kwargs["custom_llm_provider"] == "s3_vectors"
|
|
assert search_kwargs["max_num_results"] == 5
|
|
assert search_kwargs["router"] is router
|
|
# provider-specific extras forwarded
|
|
assert search_kwargs["aws_region_name"] == "eu-west-1"
|
|
assert search_kwargs["embedding_model"] == "my-embed"
|
|
assert search_kwargs["vector_bucket_name"] == "bkt"
|
|
# consumed keys are not duplicated into the spread
|
|
assert "top_k" not in search_kwargs
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aquery_minimal_retrieval_config_forwards_no_extras():
|
|
"""
|
|
A minimal retrieval_config must not leak consumed keys (or invent extras)
|
|
into the vector store search call.
|
|
"""
|
|
from unittest.mock import AsyncMock
|
|
|
|
from litellm.types.vector_stores import VectorStoreSearchResponse
|
|
|
|
fake_search = AsyncMock(
|
|
return_value=VectorStoreSearchResponse(
|
|
object="vector_store.search_results.page", search_query="q", data=[]
|
|
)
|
|
)
|
|
with patch("litellm.vector_stores.asearch", new=fake_search): # test-quality-ok: asearch is the boundary the forwarding contract under test targets
|
|
await litellm.aquery(
|
|
model="gpt-4o-mini",
|
|
messages=[{"role": "user", "content": "hello"}],
|
|
retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"},
|
|
mock_response="hi",
|
|
)
|
|
|
|
fake_search.assert_awaited_once()
|
|
search_kwargs = fake_search.await_args.kwargs
|
|
assert search_kwargs["vector_store_id"] == "vs_test_123"
|
|
assert search_kwargs["custom_llm_provider"] == "openai"
|
|
assert search_kwargs["router"] is None
|
|
leaked = {"top_k", "filters", "retrieval_filter", "aws_region_name", "embedding_model", "vector_bucket_name"}
|
|
assert not (leaked & set(search_kwargs.keys()))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aquery_does_not_forward_connection_override_keys_to_search():
|
|
"""
|
|
Only allowlisted retrieval_config keys may reach the vector store search
|
|
call. Caller-controlled connection overrides (api_base, api_key, arbitrary
|
|
extras) must be dropped, otherwise a caller could redirect store
|
|
credentials to an attacker-chosen host.
|
|
"""
|
|
from unittest.mock import AsyncMock
|
|
|
|
from litellm.types.vector_stores import VectorStoreSearchResponse
|
|
|
|
fake_search = AsyncMock(
|
|
return_value=VectorStoreSearchResponse(
|
|
object="vector_store.search_results.page", search_query="q", data=[]
|
|
)
|
|
)
|
|
with patch("litellm.vector_stores.asearch", new=fake_search): # test-quality-ok: asearch is the boundary the forwarding contract under test targets
|
|
await litellm.aquery(
|
|
model="gpt-4o-mini",
|
|
messages=[{"role": "user", "content": "hello"}],
|
|
retrieval_config={
|
|
"vector_store_id": "bkt:idx",
|
|
"custom_llm_provider": "s3_vectors",
|
|
"aws_region_name": "eu-west-1",
|
|
"api_base": "https://attacker.example.com",
|
|
"api_key": "attacker-key",
|
|
"arbitrary_extra": "nope",
|
|
},
|
|
mock_response="hi",
|
|
)
|
|
|
|
fake_search.assert_awaited_once()
|
|
search_kwargs = fake_search.await_args.kwargs
|
|
assert search_kwargs["aws_region_name"] == "eu-west-1"
|
|
blocked = {"api_base", "api_key", "arbitrary_extra"}
|
|
assert not (blocked & set(search_kwargs.keys()))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aquery_forwards_vector_store_params_to_search_but_not_completion():
|
|
"""
|
|
Regression for LIT-6773: the server-trusted vector_store_params (a managed
|
|
store's litellm_params) must reach the search call wholesale, including the
|
|
connection keys the caller allowlist blocks, while the caller's own
|
|
retrieval_config overrides stay blocked, the caller's top-level api_key and
|
|
api_base stay on the completion only, and the completion never inherits the
|
|
store's connection params.
|
|
"""
|
|
from unittest.mock import AsyncMock
|
|
|
|
from litellm.types.vector_stores import VectorStoreSearchResponse
|
|
|
|
fake_search = AsyncMock(
|
|
return_value=VectorStoreSearchResponse(
|
|
object="vector_store.search_results.page", search_query="q", data=[]
|
|
)
|
|
)
|
|
fake_completion = AsyncMock(
|
|
return_value=ModelResponse(
|
|
id="chatcmpl-test",
|
|
choices=[{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
|
|
model="gpt-4o-mini",
|
|
)
|
|
)
|
|
with (
|
|
patch("litellm.vector_stores.asearch", new=fake_search), # test-quality-ok: the search boundary under test
|
|
patch("litellm.acompletion", new=fake_completion), # test-quality-ok: the completion boundary under test
|
|
):
|
|
await litellm.aquery(
|
|
model="gpt-4o-mini",
|
|
messages=[{"role": "user", "content": "hello"}],
|
|
api_key="sk-llm-key",
|
|
api_base="https://llm.example.com",
|
|
retrieval_config={
|
|
"vector_store_id": "customer_kb",
|
|
"custom_llm_provider": "milvus",
|
|
"api_base": "https://attacker.example.com",
|
|
"api_key": "attacker-key",
|
|
},
|
|
vector_store_params={
|
|
"vector_store_id": "customer_kb",
|
|
"custom_llm_provider": "milvus",
|
|
"api_base": "http://127.0.0.1:19530",
|
|
"api_key": "root:Milvus",
|
|
"milvus_text_field": "book_intro_text",
|
|
"outputFields": ["book_intro_text"],
|
|
},
|
|
)
|
|
|
|
fake_search.assert_awaited_once()
|
|
search_kwargs = fake_search.await_args.kwargs
|
|
assert search_kwargs["vector_store_id"] == "customer_kb"
|
|
assert search_kwargs["custom_llm_provider"] == "milvus"
|
|
assert search_kwargs["api_base"] == "http://127.0.0.1:19530"
|
|
assert search_kwargs["api_key"] == "root:Milvus"
|
|
assert search_kwargs["milvus_text_field"] == "book_intro_text"
|
|
assert search_kwargs["outputFields"] == ["book_intro_text"]
|
|
fake_completion.assert_awaited_once()
|
|
completion_kwargs = fake_completion.await_args.kwargs
|
|
assert completion_kwargs["api_key"] == "sk-llm-key"
|
|
assert completion_kwargs["api_base"] == "https://llm.example.com"
|
|
assert not ({"milvus_text_field", "outputFields"} & set(completion_kwargs))
|
|
|
|
|
|
_UNDECODABLE_FILE: Final = {"filename": "notes.txt", "content": "x"}
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("ingest_options", "expected_provider"),
|
|
[
|
|
({"vector_store": {"custom_llm_provider": "bedrock"}}, "bedrock"),
|
|
({"vector_store": MappingProxyType({"custom_llm_provider": "bedrock"})}, "bedrock"),
|
|
({"vector_store": {"vector_store_id": "vs_1", 7: "ignored"}}, None),
|
|
({}, None),
|
|
],
|
|
)
|
|
def test_ingest_failure_is_attributed_to_the_vector_store_provider(
|
|
ingest_options: dict[str, object], expected_provider: str | None
|
|
) -> None:
|
|
with pytest.raises(litellm.APIConnectionError, match="Invalid base64-encoded string") as raised:
|
|
litellm.ingest(ingest_options=ingest_options, file=_UNDECODABLE_FILE)
|
|
|
|
assert raised.value.llm_provider == expected_provider
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("ingest_options", "expected_provider"),
|
|
[
|
|
({"vector_store": {"custom_llm_provider": "bedrock"}}, "bedrock"),
|
|
({"vector_store": MappingProxyType({"custom_llm_provider": "bedrock"})}, "bedrock"),
|
|
({"vector_store": {"vector_store_id": "vs_1", 7: "ignored"}}, None),
|
|
({}, None),
|
|
],
|
|
)
|
|
async def test_aingest_failure_is_attributed_to_the_vector_store_provider(
|
|
ingest_options: dict[str, object], expected_provider: str | None
|
|
) -> None:
|
|
with pytest.raises(litellm.APIConnectionError, match="Invalid base64-encoded string") as raised:
|
|
await litellm.aingest(ingest_options=ingest_options, file=_UNDECODABLE_FILE)
|
|
|
|
assert raised.value.llm_provider == expected_provider
|
|
|
|
|
|
@pytest.mark.parametrize("vector_store", [None, "openai", ["openai"], [{"api_key": "sk-test"}]])
|
|
def test_ingest_failure_with_a_vector_store_that_is_not_a_mapping_raises_a_validation_error(
|
|
vector_store: object,
|
|
) -> None:
|
|
with pytest.raises(ValidationError) as raised:
|
|
litellm.ingest(ingest_options={"vector_store": vector_store}, file=_UNDECODABLE_FILE)
|
|
|
|
assert "sk-test" not in str(raised.value)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("vector_store", [None, "openai", ["openai"], [{"api_key": "sk-test"}]])
|
|
async def test_aingest_failure_with_a_vector_store_that_is_not_a_mapping_raises_a_validation_error(
|
|
vector_store: object,
|
|
) -> None:
|
|
with pytest.raises(ValidationError) as raised:
|
|
await litellm.aingest(ingest_options={"vector_store": vector_store}, file=_UNDECODABLE_FILE)
|
|
|
|
assert "sk-test" not in str(raised.value)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("provider", [None, 5])
|
|
async def test_aquery_with_a_provider_that_is_not_a_string_bills_only_the_completion(provider: object) -> None:
|
|
messages: Final = [{"role": "user", "content": "hello"}]
|
|
default_provider_response: Final = await litellm.aquery(
|
|
model="gpt-4o-mini",
|
|
messages=messages,
|
|
retrieval_config={"vector_store_id": "vs_test_123"},
|
|
mock_response="hi there",
|
|
)
|
|
|
|
response: Final = await litellm.aquery(
|
|
model="gpt-4o-mini",
|
|
messages=messages,
|
|
retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": provider},
|
|
mock_response="hi there",
|
|
)
|
|
|
|
await _drain_logging_worker()
|
|
|
|
assert response._hidden_params["response_cost"] > 0
|
|
assert response._hidden_params["response_cost"] == default_provider_response._hidden_params["response_cost"]
|
|
|
|
|
|
def test_rag_call_types_are_registered():
|
|
"""
|
|
query/aquery/ingest/aingest are @client-decorated entry points, so their
|
|
function names must resolve to CallTypes members (deployment hooks and
|
|
call-type driven logic silently no-op for unregistered call types).
|
|
"""
|
|
assert CallTypes("query") is CallTypes.query
|
|
assert CallTypes("aquery") is CallTypes.aquery
|
|
assert CallTypes("ingest") is CallTypes.ingest
|
|
assert CallTypes("aingest") is CallTypes.aingest
|