mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-16 04:29:31 +00:00
0ebec123fd
* chore(typing): validate provider payloads at the boundary to clear 464 basedpyright Any errors Replaces untyped reads of decoded JSON in 48 provider, guardrail and proxy modules with one pydantic validation at the point the payload enters the module, so every value read afterwards carries a real type. reportAny 6111 -> 5663 and reportExplicitAny 1416 -> 1400. No basedpyright, ruff-strict or type-discipline rule is higher repo-wide or in any single file, and no cast, ignore or suppression marker was added. * chore(typing): replace 47 Any annotations that every caller already satisfies with object Each annotation here is a parameter or return type whose value is only stored, compared, logged or forwarded, so object describes it exactly and no function body changes. A site was kept only if no basedpyright, ruff-strict or type-discipline rule rose in any file that references it. reportExplicitAny 1400 -> 1353 and reportAny 5663 -> 5655. * chore(typing): type upstream responses and loose annotations across providers, proxy and integrations Second pass of the Any reduction. Decoded upstream JSON is validated once at the provider boundary and read through typed views afterwards, annotations that every caller already satisfies move from Any to a concrete type, and the nine search transformations validate each result item on its own so error text no longer carries the item index. No cast, no type: ignore, no suppression marker and no budget file edit * ci(unit): run the brave, google_pse, nlp_cloud and opensandbox tests in the providers shard
110 lines
4.2 KiB
Python
110 lines
4.2 KiB
Python
import asyncio
|
|
import json
|
|
import re
|
|
|
|
import httpx
|
|
import pytest
|
|
import respx
|
|
|
|
import litellm
|
|
from litellm.integrations.argilla import ArgillaLogger
|
|
|
|
ARGILLA_BASE_URL = "https://argilla.example.test"
|
|
USER_MESSAGE = {"role": "user", "content": "where is my order?"}
|
|
|
|
|
|
async def _cancel_background_tasks() -> None:
|
|
background_tasks = [task for task in asyncio.all_tasks() if task is not asyncio.current_task()]
|
|
for task in background_tasks:
|
|
task.cancel()
|
|
await asyncio.gather(*background_tasks, return_exceptions=True)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("transformation_object", "standard_logging_object", "expected_fields"),
|
|
[
|
|
pytest.param(
|
|
{"prompt": "messages", "answer": "response"},
|
|
{"messages": [USER_MESSAGE], "response": "it ships tomorrow"},
|
|
{"prompt": [USER_MESSAGE], "answer": "it ships tomorrow"},
|
|
id="message-list-and-text-response",
|
|
),
|
|
pytest.param(
|
|
{"answer": "response", "prompt": "messages"},
|
|
{"messages": USER_MESSAGE, "response": {"choices": [{"message": {"content": "it ships tomorrow"}}]}},
|
|
{"answer": "it ships tomorrow", "prompt": [USER_MESSAGE]},
|
|
id="single-message-and-chat-completion-response",
|
|
),
|
|
pytest.param(
|
|
{"prompt": "messages", "context": "messages", "answer": "response", "label": "response"},
|
|
{"messages": [USER_MESSAGE, USER_MESSAGE], "response": {"choices": [{"message": {}}]}},
|
|
{
|
|
"prompt": [USER_MESSAGE, USER_MESSAGE],
|
|
"context": [USER_MESSAGE, USER_MESSAGE],
|
|
"answer": "",
|
|
"label": "",
|
|
},
|
|
id="payload-field-reused-by-several-argilla-fields",
|
|
),
|
|
],
|
|
)
|
|
async def test_argilla_logger_sends_the_record_shaped_by_the_configured_field_mapping(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
respx_mock: respx.MockRouter,
|
|
transformation_object: dict[str, str],
|
|
standard_logging_object: dict[str, object],
|
|
expected_fields: dict[str, object],
|
|
) -> None:
|
|
monkeypatch.setattr(litellm, "argilla_transformation_object", transformation_object)
|
|
respx_mock.get(f"{ARGILLA_BASE_URL}/api/v1/me/datasets", params={"name": "support-chats"}).mock(
|
|
return_value=httpx.Response(200, json={"items": [{"id": "dataset-403"}]})
|
|
)
|
|
bulk_route = respx_mock.post(f"{ARGILLA_BASE_URL}/api/v1/datasets/dataset-403/records/bulk").mock(
|
|
return_value=httpx.Response(200, json={})
|
|
)
|
|
logger = ArgillaLogger(
|
|
argilla_api_key="argilla-key",
|
|
argilla_dataset_name="support-chats",
|
|
argilla_base_url=ARGILLA_BASE_URL,
|
|
batch_size=1,
|
|
)
|
|
try:
|
|
logger.log_success_event({"standard_logging_object": standard_logging_object}, None, None, None)
|
|
finally:
|
|
await _cancel_background_tasks()
|
|
|
|
sent_request = bulk_route.calls.last.request
|
|
assert json.loads(sent_request.content) == [{"fields": expected_fields}]
|
|
assert sent_request.headers["X-Argilla-Api-Key"] == "argilla-key"
|
|
assert logger.argilla_transformation_object is transformation_object
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("transformation_object", "message"),
|
|
[
|
|
pytest.param(
|
|
None,
|
|
"'litellm.argilla_transformation_object' is required, to log your payload to Argilla.",
|
|
id="mapping-not-configured",
|
|
),
|
|
pytest.param(
|
|
["messages", "response"],
|
|
"'argilla_transformation_object' must be a dictionary, to log your payload to Argilla.",
|
|
id="mapping-is-a-list",
|
|
),
|
|
pytest.param(
|
|
{"prompt": "messages", "spend": "response_cost"},
|
|
"All values in argilla_transformation_object must be a key in SUPPORTED_PAYLOAD_FIELDS, "
|
|
"response_cost is not a valid key.",
|
|
id="mapping-names-an-unsupported-payload-field",
|
|
),
|
|
],
|
|
)
|
|
def test_argilla_logger_refuses_to_start_without_a_supported_field_mapping(
|
|
monkeypatch: pytest.MonkeyPatch, transformation_object: object, message: str
|
|
) -> None:
|
|
monkeypatch.setattr(litellm, "argilla_transformation_object", transformation_object)
|
|
|
|
with pytest.raises(Exception, match=re.escape(message)):
|
|
ArgillaLogger(argilla_api_key="argilla-key")
|