mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-28 06:23:07 +00:00
9fbeebc988
* feat(extraction): add Clairic extraction types and Cloudflare provider config * feat(extraction): add the /v1/extraction route, router wiring, and Clairic billing * fix(extraction): refuse or drop stream, pass unknown keys, run guardrails, keep the requested model name Clairic does not stream, so stream=true answers 400 and is dropped and billed under drop_params. Keys LiteLLM does not know pass through to Cloudflare. Guardrails scan state, field instructions and text answers. Frozen decision and extraction responses now carry the requested model name. List fields must name their item type. /health/test_connection and the Add Model page accept the extraction mode. * fix(extraction): keep upstream model names on decisions, scan images, keep typed values, refuse OpenAI params Scope the frozen-response model stamp to extraction so /v1/decisions and /v1/systemone return the upstream name as on main. Pass request images to guardrails as data URLs and drop answer spans when a response guardrail rewrites state_text. Field spec values keep their JSON type, so a wrong type answers 400 instead of being coerced. OpenAI generation params answer 400 unless drop_params. The Model Hub shows a /v1/extraction snippet for extraction models. * fix(extraction): send only Clairic's request fields, count extraction tokens, keep typed answers typed, drop spans a request guardrail made stale, and honor --max_tokens * fix(extraction): scan number answers in response guardrails, rewrite guardrail fields in linear time, and leave --max_tokens as on main * fix(extraction): scan float answers as whole numbers so fraction digits never trip a response guardrail * fix(extraction): round only mean and median for response guardrails so fractional values keep their decimal text * fix(extraction): count state and fields for budget reservations and scan estimate quantiles in response guardrails * fix(extraction): send each repeated text to guardrails once and keep the context-window budget reservation * test(extraction): type request_data in the guardrail test helpers
61 lines
2.2 KiB
Python
61 lines
2.2 KiB
Python
import asyncio
|
|
from typing import Final, Literal
|
|
|
|
import pytest
|
|
|
|
from litellm.integrations.custom_guardrail import CustomGuardrail
|
|
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
|
from litellm.llms.cloudflare.extraction.guardrail_translation.handler import ExtractionHandler
|
|
from litellm.types.extraction import ExtractionResponse
|
|
from litellm.types.utils import GenericGuardrailAPIInputs
|
|
|
|
MANY_FIELDS: Final = 2_000
|
|
|
|
|
|
class _MaskingGuardrail(CustomGuardrail):
|
|
def __init__(self) -> None:
|
|
super().__init__(guardrail_name="masking")
|
|
|
|
async def apply_guardrail(
|
|
self,
|
|
inputs: GenericGuardrailAPIInputs,
|
|
request_data: dict[str, object],
|
|
input_type: Literal["request", "response"],
|
|
logging_obj: LiteLLMLoggingObj | None = None,
|
|
) -> GenericGuardrailAPIInputs:
|
|
return GenericGuardrailAPIInputs(texts=[text.replace("secret", "[MASKED]") for text in inputs.get("texts", [])])
|
|
|
|
|
|
@pytest.mark.benchmark
|
|
def test_extraction_request_guardrail_over_many_field_instructions() -> None:
|
|
data: Final = {
|
|
"model": "invoice-reader",
|
|
"fields": {
|
|
f"field_{index}": {"type": "text", "instructions": f"secret {index}"} for index in range(MANY_FIELDS)
|
|
},
|
|
}
|
|
|
|
result: Final = asyncio.run(ExtractionHandler().process_input_messages(data, _MaskingGuardrail()))
|
|
|
|
assert result["fields"] == {
|
|
f"field_{index}": {"type": "text", "instructions": f"[MASKED] {index}"} for index in range(MANY_FIELDS)
|
|
}
|
|
|
|
|
|
@pytest.mark.benchmark
|
|
def test_extraction_response_guardrail_over_many_text_answers() -> None:
|
|
response: Final = ExtractionResponse.model_validate(
|
|
{
|
|
"model": "invoice-reader",
|
|
"answers": {f"field_{index}": {"type": "text", "value": f"secret {index}"} for index in range(MANY_FIELDS)},
|
|
"truncated": False,
|
|
"usage": {"input_tokens": 52, "output_tokens": 0},
|
|
}
|
|
)
|
|
|
|
result: Final = asyncio.run(ExtractionHandler().process_output_response(response, _MaskingGuardrail()))
|
|
|
|
assert isinstance(result, ExtractionResponse)
|
|
assert {name: answer.value for name, answer in result.answers.items()} == {
|
|
f"field_{index}": f"[MASKED] {index}" for index in range(MANY_FIELDS)
|
|
}
|