mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-14 04:09:09 +00:00
d80f8c28ca
* refactor(types): prove runtime types at harness, search, rag and client boundaries Replace Any with adapter-validated types in the litellm.agent() harness, the search provider transformations, RAG ingestion and query, the vector store pre-call hook and registry, the galileo and opik logging integrations and the proxy client CLI. Each boundary gets unit tests for well-formed and malformed payloads. * chore(typing): prove types at more provider boundaries and restore search transformations Second pass over a2a, embedding, rerank, image, audio and small provider modules. Search transformations go back to their previous form because validating their response bodies would change the proxy status for malformed upstream bodies from 400 to 500. * refactor(types): prove types at logging, files, rerank, image, audio and management boundaries Replace Any with validated or annotation-only types in 37 more files: logging integrations, token counters, provider files/rerank/image generation/audio transcription transformations, pass-through logging handlers and management endpoints. No proxy HTTP status or error type changes. * refactor(types): prove types at repository, spend, files and router boundaries Replace Any with repository table accessors, validated mappings and annotation-only types in 33 more files: Prisma repositories, the enterprise batch and responses cost checkers, budget reservation, files endpoints, management endpoints, the policy registry, the adaptive and complexity routers and the secret managers. No proxy HTTP status or error type changes. * test(types): run the aiohttp transformation test in-process and cover repository row conversion The aiohttp chat transformation test no longer starts a server. It feeds the transformation a response whose json() returns the body under test. The proxy unit shards now exercise stored model rows whose params are JSON strings and the object permission create and update paths.
96 lines
3.2 KiB
Python
96 lines
3.2 KiB
Python
import json
|
|
from collections.abc import Iterator
|
|
|
|
import httpx
|
|
import pytest
|
|
import respx
|
|
from pydantic import ValidationError
|
|
|
|
import litellm
|
|
from litellm.llms.together_ai.rerank.handler import TogetherAIRerank
|
|
|
|
API_BASE = "https://api.together.example/v1"
|
|
RERANKED = {
|
|
"id": "rr-1",
|
|
"results": [
|
|
{"index": 1, "relevance_score": 0.9, "document": {"text": "Paris"}},
|
|
{"index": 0, "relevance_score": 0.1},
|
|
],
|
|
"usage": {"total_tokens": 7},
|
|
}
|
|
NON_OBJECT_BODIES = [7, "sensitive-document", [{"id": "sensitive-document"}]]
|
|
|
|
|
|
@pytest.fixture
|
|
def rerank_route(monkeypatch: pytest.MonkeyPatch) -> Iterator[respx.Route]:
|
|
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
|
with respx.mock as router:
|
|
yield router.post(f"{API_BASE}/rerank")
|
|
|
|
|
|
def _rerank(*, is_async: bool):
|
|
return TogetherAIRerank().rerank(
|
|
model="Salesforce/Llama-Rank-V1",
|
|
api_key="sk-test",
|
|
api_base=API_BASE,
|
|
query="capital of france",
|
|
documents=["Berlin", {"text": "Paris"}],
|
|
top_n=1,
|
|
_is_async=is_async,
|
|
)
|
|
|
|
|
|
def _assert_is_the_reranking(response: litellm.RerankResponse, route: respx.Route) -> None:
|
|
assert response.id == "rr-1"
|
|
assert response.results == [
|
|
{"index": 1, "relevance_score": 0.9, "document": {"text": "Paris"}},
|
|
{"index": 0, "relevance_score": 0.1},
|
|
]
|
|
assert response.meta == {"billed_units": {"total_tokens": 7}, "tokens": {}}
|
|
assert route.calls.last.request.headers["authorization"] == "Bearer sk-test"
|
|
assert json.loads(route.calls.last.request.content) == {
|
|
"model": "Salesforce/Llama-Rank-V1",
|
|
"query": "capital of france",
|
|
"top_n": 1,
|
|
"documents": ["Berlin", {"text": "Paris"}],
|
|
"return_documents": True,
|
|
}
|
|
|
|
|
|
def test_rerank_returns_the_upstream_ranking(rerank_route: respx.Route):
|
|
rerank_route.mock(return_value=httpx.Response(200, json=RERANKED))
|
|
|
|
_assert_is_the_reranking(_rerank(is_async=False), rerank_route)
|
|
|
|
|
|
async def test_async_rerank_returns_the_upstream_ranking(rerank_route: respx.Route):
|
|
rerank_route.mock(return_value=httpx.Response(200, json=RERANKED))
|
|
|
|
_assert_is_the_reranking(await _rerank(is_async=True), rerank_route)
|
|
|
|
|
|
@pytest.mark.parametrize("payload", NON_OBJECT_BODIES)
|
|
def test_rerank_rejects_non_object_bodies(rerank_route: respx.Route, payload: object):
|
|
rerank_route.mock(return_value=httpx.Response(200, json=payload))
|
|
|
|
with pytest.raises(ValidationError) as exc_info:
|
|
_rerank(is_async=False)
|
|
|
|
assert "sensitive-document" not in str(exc_info.value)
|
|
|
|
|
|
@pytest.mark.parametrize("payload", NON_OBJECT_BODIES)
|
|
async def test_async_rerank_rejects_non_object_bodies(rerank_route: respx.Route, payload: object):
|
|
rerank_route.mock(return_value=httpx.Response(200, json=payload))
|
|
|
|
with pytest.raises(ValidationError) as exc_info:
|
|
await _rerank(is_async=True)
|
|
|
|
assert "sensitive-document" not in str(exc_info.value)
|
|
|
|
|
|
def test_rerank_without_results_is_a_value_error(rerank_route: respx.Route):
|
|
rerank_route.mock(return_value=httpx.Response(200, json={"id": "rr-1"}))
|
|
|
|
with pytest.raises(ValueError, match="No results found"):
|
|
_rerank(is_async=False)
|