mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
feat(parallel_ai): add a native Extract pass-through with cost tracking
Parallel's /v1/extract endpoint has no equivalent in any unified LiteLLM surface: it takes a URL list and returns per-URL excerpts and full content, with partial success reported as a parallel errors array. Forcing it into the search or responses shape would lose that response. Adds it as a pass-through at /parallel_ai/v1/extract so the request and response cross unchanged, with the gateway supplying the provider credential. The success handler prices the call from the URLs the provider reports as billed, falling back to the requested URL count when usage is absent, so extract traffic is attributed rather than logged at $0. Adding a route adds one FastAPI Depends() default, so the B008 ceiling moves by one.
This commit is contained in:
parent
1ed4ae981a
commit
52e7ea4332
10 changed files with 600 additions and 1 deletions
6
litellm/llms/parallel_ai/extract/__init__.py
Normal file
6
litellm/llms/parallel_ai/extract/__init__.py
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
from litellm.llms.parallel_ai.extract.cost_calculator import (
|
||||
PARALLEL_AI_EXTRACT_MODEL,
|
||||
parallel_ai_extract_cost,
|
||||
)
|
||||
|
||||
__all__ = ["PARALLEL_AI_EXTRACT_MODEL", "parallel_ai_extract_cost"]
|
||||
73
litellm/llms/parallel_ai/extract/cost_calculator.py
Normal file
73
litellm/llms/parallel_ai/extract/cost_calculator.py
Normal file
|
|
@ -0,0 +1,73 @@
|
|||
from typing import Annotated, Final
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, StrictInt, StrictStr, ValidationError
|
||||
|
||||
PARALLEL_AI_EXTRACT_COST_PER_URL: Final = 0.001
|
||||
PARALLEL_AI_EXTRACT_MODEL: Final = "parallel_ai/extract"
|
||||
PARALLEL_AI_EXTRACT_USAGE_SKU: Final = "sku_extract_excerpts"
|
||||
|
||||
|
||||
class _ParallelAIExtractUsageItem(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
name: StrictStr
|
||||
count: Annotated[StrictInt, Field(ge=0)]
|
||||
|
||||
|
||||
class _ParallelAIExtractUsageName(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
name: StrictStr
|
||||
|
||||
|
||||
class _ParallelAIExtractBillingResponse(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
usage: tuple[object, ...] | None = None
|
||||
|
||||
|
||||
class _ParallelAIExtractBillingRequest(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
urls: tuple[StrictStr, ...] = ()
|
||||
|
||||
|
||||
def _usage_url_count(response_body: object) -> int | None:
|
||||
try:
|
||||
parsed: Final = _ParallelAIExtractBillingResponse.model_validate(response_body)
|
||||
except ValidationError:
|
||||
return None
|
||||
if parsed.usage is None:
|
||||
return None
|
||||
|
||||
target_items: Final = tuple(item for item in parsed.usage if _usage_name(item) == PARALLEL_AI_EXTRACT_USAGE_SKU)
|
||||
if not target_items:
|
||||
return 0
|
||||
|
||||
try:
|
||||
usage_items: Final = tuple(_ParallelAIExtractUsageItem.model_validate(item) for item in target_items)
|
||||
except ValidationError:
|
||||
return None
|
||||
return sum(item.count for item in usage_items)
|
||||
|
||||
|
||||
def _usage_name(usage_item: object) -> str | None:
|
||||
try:
|
||||
parsed: Final = _ParallelAIExtractUsageName.model_validate(usage_item)
|
||||
except ValidationError:
|
||||
return None
|
||||
return parsed.name
|
||||
|
||||
|
||||
def _request_url_count(request_body: object) -> int:
|
||||
try:
|
||||
parsed: Final = _ParallelAIExtractBillingRequest.model_validate(request_body)
|
||||
except ValidationError:
|
||||
return 0
|
||||
return len(parsed.urls)
|
||||
|
||||
|
||||
def parallel_ai_extract_cost(request_body: object, response_body: object) -> float:
|
||||
usage_url_count: Final = _usage_url_count(response_body)
|
||||
billed_url_count: Final = usage_url_count if usage_url_count is not None else _request_url_count(request_body)
|
||||
return billed_url_count * PARALLEL_AI_EXTRACT_COST_PER_URL
|
||||
|
|
@ -466,6 +466,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/openai_passthrough",
|
||||
"/assemblyai",
|
||||
"/eu.assemblyai",
|
||||
"/parallel_ai",
|
||||
"/vllm",
|
||||
"/mistral",
|
||||
"/milvus",
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ from litellm.constants import (
|
|||
BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES,
|
||||
)
|
||||
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
|
||||
from litellm.llms.parallel_ai.common_utils import resolve_parallel_ai_credentials
|
||||
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
|
|
@ -310,6 +311,52 @@ async def cohere_proxy_route(
|
|||
return received_value
|
||||
|
||||
|
||||
def _parallel_ai_extract_url(api_base: str) -> str:
|
||||
trimmed: Final = api_base.rstrip("/")
|
||||
if trimmed.endswith("/v1/extract"):
|
||||
return trimmed
|
||||
return f"{trimmed.removesuffix('/v1')}/v1/extract"
|
||||
|
||||
|
||||
@router.post(
|
||||
"/parallel_ai/v1/extract",
|
||||
tags=["Parallel AI Pass-through", "pass-through"],
|
||||
)
|
||||
async def parallel_ai_extract_proxy_route(
|
||||
request: Request,
|
||||
fastapi_response: Response,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
deployment_api_key: Final = passthrough_endpoint_router.get_credentials(
|
||||
custom_llm_provider="parallel_ai",
|
||||
region_name=None,
|
||||
)
|
||||
api_base, api_key = resolve_parallel_ai_credentials(
|
||||
api_base=None,
|
||||
api_key=deployment_api_key,
|
||||
)
|
||||
if api_key is None:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail="PARALLEL_AI_API_KEY or PARALLEL_API_KEY is required for the Parallel AI Extract pass-through.",
|
||||
)
|
||||
|
||||
endpoint_func: Final = create_pass_through_route(
|
||||
endpoint="/parallel_ai/v1/extract",
|
||||
target=_parallel_ai_extract_url(api_base),
|
||||
custom_headers={
|
||||
"Content-Type": "application/json",
|
||||
"x-api-key": api_key,
|
||||
},
|
||||
custom_llm_provider="parallel_ai",
|
||||
)
|
||||
return await endpoint_func(
|
||||
request,
|
||||
fastapi_response,
|
||||
user_api_key_dict,
|
||||
)
|
||||
|
||||
|
||||
@router.api_route(
|
||||
"/vllm/{endpoint:path}",
|
||||
methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
|
||||
|
|
|
|||
|
|
@ -0,0 +1,51 @@
|
|||
import json
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.parallel_ai.extract.cost_calculator import (
|
||||
PARALLEL_AI_EXTRACT_MODEL,
|
||||
parallel_ai_extract_cost,
|
||||
)
|
||||
from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
|
||||
from litellm.types.utils import StandardPassThroughResponseObject
|
||||
|
||||
|
||||
class ParallelAIPassthroughLoggingHandler:
|
||||
@staticmethod
|
||||
def is_extract_route(url_route: str, custom_llm_provider: str | None) -> bool:
|
||||
path: Final = urlparse(url_route).path.rstrip("/")
|
||||
return custom_llm_provider == "parallel_ai" and path.endswith("/v1/extract")
|
||||
|
||||
@staticmethod
|
||||
def parallel_ai_extract_handler(
|
||||
response_body: Mapping[str, object],
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
request_body: Mapping[str, object],
|
||||
**kwargs: object, # kwargs-ok: the passthrough logging dispatch forwards shared logging kwargs to every handler
|
||||
) -> PassThroughEndpointLoggingTypedDict:
|
||||
"""
|
||||
Prices a Parallel AI Extract call from the URLs the provider reports as
|
||||
billed (falling back to the requested URL count) and records model,
|
||||
provider, and cost on the logging payload.
|
||||
"""
|
||||
response_cost: Final = parallel_ai_extract_cost(
|
||||
request_body=request_body,
|
||||
response_body=response_body,
|
||||
)
|
||||
logging_obj.model_call_details.update(
|
||||
model=PARALLEL_AI_EXTRACT_MODEL,
|
||||
custom_llm_provider="parallel_ai",
|
||||
response_cost=response_cost,
|
||||
)
|
||||
|
||||
return {
|
||||
"result": StandardPassThroughResponseObject(response=json.dumps(response_body)),
|
||||
"kwargs": { # mutable-ok: the logging pipeline requires a plain kwargs dict
|
||||
**kwargs,
|
||||
"model": PARALLEL_AI_EXTRACT_MODEL,
|
||||
"custom_llm_provider": "parallel_ai",
|
||||
"response_cost": response_cost,
|
||||
},
|
||||
}
|
||||
|
|
@ -27,6 +27,9 @@ from .llm_provider_handlers.cursor_passthrough_logging_handler import (
|
|||
from .llm_provider_handlers.gemini_passthrough_logging_handler import (
|
||||
GeminiPassthroughLoggingHandler,
|
||||
)
|
||||
from .llm_provider_handlers.parallel_ai_passthrough_logging_handler import (
|
||||
ParallelAIPassthroughLoggingHandler,
|
||||
)
|
||||
from .llm_provider_handlers.vertex_passthrough_logging_handler import (
|
||||
VertexPassthroughLoggingHandler,
|
||||
)
|
||||
|
|
@ -221,6 +224,16 @@ class PassThroughEndpointLogging:
|
|||
standard_logging_response_object = openai_passthrough_logging_handler_result["result"]
|
||||
kwargs = openai_passthrough_logging_handler_result["kwargs"]
|
||||
|
||||
elif ParallelAIPassthroughLoggingHandler.is_extract_route(url_route, custom_llm_provider):
|
||||
parallel_ai_result = ParallelAIPassthroughLoggingHandler.parallel_ai_extract_handler(
|
||||
response_body=response_body if isinstance(response_body, dict) else {},
|
||||
logging_obj=logging_obj,
|
||||
request_body=request_body,
|
||||
**kwargs,
|
||||
)
|
||||
standard_logging_response_object = parallel_ai_result["result"]
|
||||
kwargs = parallel_ai_result["kwargs"] # rebind-ok: every dispatch branch reassigns kwargs
|
||||
|
||||
elif self.is_cursor_route(url_route, custom_llm_provider):
|
||||
cursor_passthrough_logging_handler_result = CursorPassthroughLoggingHandler.cursor_passthrough_handler(
|
||||
httpx_response=httpx_response,
|
||||
|
|
|
|||
|
|
@ -36,7 +36,7 @@
|
|||
"limit": 177
|
||||
},
|
||||
"B008": {
|
||||
"limit": 503
|
||||
"limit": 505
|
||||
},
|
||||
"B009": {
|
||||
"limit": 59
|
||||
|
|
|
|||
|
|
@ -0,0 +1,90 @@
|
|||
import pytest
|
||||
|
||||
from litellm.llms.parallel_ai.extract.cost_calculator import parallel_ai_extract_cost
|
||||
|
||||
|
||||
def test_extract_cost_prefers_provider_usage_over_requested_urls() -> None:
|
||||
cost = parallel_ai_extract_cost(
|
||||
request_body={"urls": ["https://example.com/1", "https://example.com/2", "https://example.com/3"]},
|
||||
response_body={"usage": [{"name": "sku_extract_excerpts", "count": 2}]},
|
||||
)
|
||||
|
||||
assert cost == pytest.approx(0.002)
|
||||
|
||||
|
||||
def test_extract_cost_sums_repeated_usage_skus() -> None:
|
||||
cost = parallel_ai_extract_cost(
|
||||
request_body={"urls": ["https://example.com/1"]},
|
||||
response_body={
|
||||
"usage": [
|
||||
{"name": "sku_extract_excerpts", "count": 1},
|
||||
{"name": "unrelated_sku", "count": 9},
|
||||
{"name": "sku_extract_excerpts", "count": 2},
|
||||
]
|
||||
},
|
||||
)
|
||||
|
||||
assert cost == pytest.approx(0.003)
|
||||
|
||||
|
||||
def test_extract_cost_ignores_malformed_unrelated_usage_skus() -> None:
|
||||
cost = parallel_ai_extract_cost(
|
||||
request_body={"urls": ["https://example.com/1", "https://example.com/2"]},
|
||||
response_body={
|
||||
"usage": [
|
||||
{"name": "sku_extract_excerpts", "count": 1},
|
||||
{"name": "unrelated_sku", "count": "not-an-integer"},
|
||||
]
|
||||
},
|
||||
)
|
||||
|
||||
assert cost == pytest.approx(0.001)
|
||||
|
||||
|
||||
def test_extract_cost_treats_usage_without_extract_sku_as_unbilled() -> None:
|
||||
cost = parallel_ai_extract_cost(
|
||||
request_body={"urls": ["https://example.com/1", "https://example.com/2"]},
|
||||
response_body={"usage": [{"name": "unrelated_sku", "count": 2}]},
|
||||
)
|
||||
|
||||
assert cost == 0.0
|
||||
|
||||
|
||||
def test_extract_cost_treats_empty_usage_as_unbilled() -> None:
|
||||
cost = parallel_ai_extract_cost(
|
||||
request_body={"urls": ["https://example.com/1", "https://example.com/2"]},
|
||||
response_body={"usage": []},
|
||||
)
|
||||
|
||||
assert cost == 0.0
|
||||
|
||||
|
||||
def test_extract_cost_falls_back_to_requested_url_count_without_usage() -> None:
|
||||
cost = parallel_ai_extract_cost(
|
||||
request_body={"urls": ["https://example.com/1", "https://example.com/2"]},
|
||||
response_body={"results": []},
|
||||
)
|
||||
|
||||
assert cost == pytest.approx(0.002)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("invalid_count", [True, -1, "2"])
|
||||
def test_extract_cost_falls_back_when_provider_usage_is_invalid(invalid_count: object) -> None:
|
||||
cost = parallel_ai_extract_cost(
|
||||
request_body={"urls": ["https://example.com/1", "https://example.com/2"]},
|
||||
response_body={"usage": [{"name": "sku_extract_excerpts", "count": invalid_count}]},
|
||||
)
|
||||
|
||||
assert cost == pytest.approx(0.002)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"request_body",
|
||||
[
|
||||
{},
|
||||
{"urls": "https://example.com"},
|
||||
{"urls": ["https://example.com", 42]},
|
||||
],
|
||||
)
|
||||
def test_extract_cost_does_not_guess_from_invalid_request_urls(request_body: object) -> None:
|
||||
assert parallel_ai_extract_cost(request_body=request_body, response_body={}) == 0.0
|
||||
|
|
@ -0,0 +1,185 @@
|
|||
"""Gateway coverage for Parallel AI's Extract pass-through."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Iterator
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
PARALLEL_EXTRACT_URL: Final = "https://api.parallel.ai/v1/extract"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client() -> TestClient:
|
||||
return TestClient(proxy_server.app, raise_server_exceptions=False)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def auth_as() -> Iterator[None]:
|
||||
async def _authorized_request() -> UserAPIKeyAuth:
|
||||
return UserAPIKeyAuth(
|
||||
api_key="hashed-sk-test",
|
||||
user_id="parallel-test-user",
|
||||
)
|
||||
|
||||
previous: Final = proxy_server.app.dependency_overrides.get(user_api_key_auth)
|
||||
proxy_server.app.dependency_overrides[user_api_key_auth] = _authorized_request
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
if previous is None:
|
||||
proxy_server.app.dependency_overrides.pop(user_api_key_auth, None)
|
||||
else:
|
||||
proxy_server.app.dependency_overrides[user_api_key_auth] = previous
|
||||
|
||||
|
||||
def _parallel_extract_body() -> dict[str, object]:
|
||||
return {
|
||||
"extract_id": "extract_parallel_gateway",
|
||||
"results": [
|
||||
{
|
||||
"url": "https://example.com/parallel",
|
||||
"title": "Parallel result",
|
||||
"publish_date": "2026-08-14",
|
||||
"excerpts": ["Focused excerpt"],
|
||||
"full_content": "# Full content",
|
||||
}
|
||||
],
|
||||
"errors": [
|
||||
{
|
||||
"url": "https://example.com/unavailable",
|
||||
"error_type": "fetch_error",
|
||||
"http_status_code": 503,
|
||||
"content": "Upstream unavailable",
|
||||
}
|
||||
],
|
||||
"warnings": None,
|
||||
"usage": [{"name": "sku_extract_excerpts", "count": 2}],
|
||||
"session_id": "session_parallel_gateway",
|
||||
}
|
||||
|
||||
|
||||
def _parallel_router() -> Router:
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "parallel-gateway",
|
||||
"litellm_params": {
|
||||
"model": "parallel_ai/parallel",
|
||||
"api_key": "parallel-responses-key",
|
||||
"use_in_pass_through": True,
|
||||
},
|
||||
}
|
||||
],
|
||||
num_retries=0,
|
||||
)
|
||||
|
||||
|
||||
def test_parallel_extract_gateway_route(client, auth_as, monkeypatch, respx_mock):
|
||||
"""The native Extract route preserves Parallel's V1 request and partial-success response."""
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
|
||||
monkeypatch.setattr(proxy_server, "llm_router", _parallel_router())
|
||||
upstream_route = respx_mock.post(PARALLEL_EXTRACT_URL).respond(json=_parallel_extract_body())
|
||||
request_body = {
|
||||
"urls": [
|
||||
"https://example.com/parallel",
|
||||
"https://example.com/unavailable",
|
||||
],
|
||||
"objective": "Find the integration details",
|
||||
"search_queries": ["Parallel integration"],
|
||||
"max_chars_total": 50000,
|
||||
"session_id": "session_parallel_gateway",
|
||||
"client_model": "gpt-5.4",
|
||||
"advanced_settings": {
|
||||
"fetch_policy": {
|
||||
"max_age_seconds": 3600,
|
||||
"timeout_seconds": 30,
|
||||
"disable_cache_fallback": False,
|
||||
},
|
||||
"excerpt_settings": {"max_chars_per_result": 5000},
|
||||
"full_content": {"max_chars_per_result": 50000},
|
||||
},
|
||||
}
|
||||
|
||||
response = client.post("/parallel_ai/v1/extract", json=request_body)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json() == _parallel_extract_body()
|
||||
assert upstream_route.called
|
||||
|
||||
upstream_request = upstream_route.calls.last.request
|
||||
assert upstream_request.headers["x-api-key"] == "parallel-responses-key"
|
||||
assert "authorization" not in upstream_request.headers
|
||||
assert json.loads(upstream_request.content) == request_body
|
||||
|
||||
|
||||
def test_parallel_extract_route_is_classified_as_an_llm_api_route() -> None:
|
||||
assert RouteChecks.is_llm_api_route(route="/parallel_ai/v1/extract") is True
|
||||
|
||||
|
||||
def test_parallel_extract_gateway_uses_environment_configuration(client, auth_as, monkeypatch, respx_mock):
|
||||
custom_url = "https://parallel-proxy.example.com/v1/extract"
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
|
||||
monkeypatch.setattr(proxy_server, "llm_router", None)
|
||||
monkeypatch.delenv("PARALLEL_AI_API_KEY", raising=False)
|
||||
monkeypatch.setenv("PARALLEL_API_KEY", "parallel-env-key")
|
||||
monkeypatch.setenv("PARALLEL_AI_API_BASE", "https://parallel-proxy.example.com/v1")
|
||||
upstream_route = respx_mock.post(custom_url).respond(json=_parallel_extract_body())
|
||||
|
||||
response = client.post(
|
||||
"/parallel_ai/v1/extract",
|
||||
json={"urls": ["https://example.com/parallel"]},
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert upstream_route.called
|
||||
assert upstream_route.calls.last.request.headers["x-api-key"] == "parallel-env-key"
|
||||
|
||||
|
||||
def test_parallel_extract_gateway_requires_a_provider_key(client, auth_as, monkeypatch) -> None:
|
||||
monkeypatch.setattr(proxy_server, "llm_router", None)
|
||||
monkeypatch.delenv("PARALLEL_AI_API_KEY", raising=False)
|
||||
monkeypatch.delenv("PARALLEL_API_KEY", raising=False)
|
||||
|
||||
response = client.post(
|
||||
"/parallel_ai/v1/extract",
|
||||
json={"urls": ["https://example.com/parallel"]},
|
||||
)
|
||||
|
||||
assert response.status_code == 500
|
||||
assert response.json()["detail"] == (
|
||||
"PARALLEL_AI_API_KEY or PARALLEL_API_KEY is required for the Parallel AI Extract pass-through."
|
||||
)
|
||||
|
||||
|
||||
def test_parallel_extract_gateway_preserves_validation_errors(client, auth_as, monkeypatch, respx_mock):
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
|
||||
monkeypatch.setattr(proxy_server, "llm_router", _parallel_router())
|
||||
error_body = {
|
||||
"error": {
|
||||
"type": "validation_error",
|
||||
"message": "urls must contain at most 20 items",
|
||||
}
|
||||
}
|
||||
respx_mock.post(PARALLEL_EXTRACT_URL).respond(status_code=422, json=error_body)
|
||||
|
||||
response = client.post(
|
||||
"/parallel_ai/v1/extract",
|
||||
json={"urls": [f"https://example.com/{index}" for index in range(21)]},
|
||||
)
|
||||
|
||||
assert response.status_code == 422
|
||||
assert response.json() == error_body
|
||||
|
|
@ -0,0 +1,133 @@
|
|||
from datetime import datetime
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from litellm.llms.parallel_ai.extract.cost_calculator import PARALLEL_AI_EXTRACT_MODEL
|
||||
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.parallel_ai_passthrough_logging_handler import (
|
||||
ParallelAIPassthroughLoggingHandler,
|
||||
)
|
||||
from litellm.proxy.pass_through_endpoints.success_handler import PassThroughEndpointLogging
|
||||
|
||||
|
||||
def test_extract_route_detection_requires_parallel_provider() -> None:
|
||||
assert ParallelAIPassthroughLoggingHandler.is_extract_route(
|
||||
"https://api.parallel.ai/v1/extract",
|
||||
"parallel_ai",
|
||||
)
|
||||
assert not ParallelAIPassthroughLoggingHandler.is_extract_route(
|
||||
"https://api.parallel.ai/v1/extract",
|
||||
None,
|
||||
)
|
||||
assert not ParallelAIPassthroughLoggingHandler.is_extract_route(
|
||||
"https://api.parallel.ai/v1/search",
|
||||
"parallel_ai",
|
||||
)
|
||||
|
||||
|
||||
def test_extract_handler_sets_usage_aware_cost_and_model() -> None:
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.litellm_call_id = "parallel-extract-call"
|
||||
logging_obj.model_call_details = {}
|
||||
response_body = {
|
||||
"extract_id": "extract_test",
|
||||
"results": [],
|
||||
"errors": [],
|
||||
"usage": [{"name": "sku_extract_excerpts", "count": 2}],
|
||||
"session_id": "session_test",
|
||||
}
|
||||
|
||||
result = ParallelAIPassthroughLoggingHandler.parallel_ai_extract_handler(
|
||||
response_body=response_body,
|
||||
logging_obj=logging_obj,
|
||||
request_body={"urls": ["https://example.com/1", "https://example.com/2"]},
|
||||
)
|
||||
|
||||
assert result["kwargs"]["model"] == PARALLEL_AI_EXTRACT_MODEL
|
||||
assert result["kwargs"]["custom_llm_provider"] == "parallel_ai"
|
||||
assert result["kwargs"]["response_cost"] == 0.002
|
||||
assert logging_obj.model_call_details["model"] == PARALLEL_AI_EXTRACT_MODEL
|
||||
assert logging_obj.model_call_details["custom_llm_provider"] == "parallel_ai"
|
||||
assert logging_obj.model_call_details["response_cost"] == 0.002
|
||||
|
||||
|
||||
def test_success_handler_dispatches_parallel_extract_billing() -> None:
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.litellm_call_id = "parallel-extract-dispatch"
|
||||
logging_obj.model_call_details = {}
|
||||
response_body = {
|
||||
"extract_id": "extract_dispatch",
|
||||
"results": [],
|
||||
"errors": [],
|
||||
"usage": [{"name": "sku_extract_excerpts", "count": 1}],
|
||||
"session_id": "session_dispatch",
|
||||
}
|
||||
response = httpx.Response(
|
||||
200,
|
||||
json=response_body,
|
||||
request=httpx.Request("POST", "https://api.parallel.ai/v1/extract"),
|
||||
)
|
||||
|
||||
normalized = PassThroughEndpointLogging().normalize_llm_passthrough_logging_payload(
|
||||
httpx_response=response,
|
||||
response_body=response_body,
|
||||
request_body={"urls": ["https://example.com"]},
|
||||
logging_obj=logging_obj,
|
||||
url_route="https://api.parallel.ai/v1/extract",
|
||||
result="",
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
cache_hit=False,
|
||||
custom_llm_provider="parallel_ai",
|
||||
)
|
||||
|
||||
assert normalized["standard_logging_response_object"] is not None
|
||||
assert normalized["kwargs"]["model"] == PARALLEL_AI_EXTRACT_MODEL
|
||||
assert normalized["kwargs"]["response_cost"] == 0.001
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_success_handler_sends_extract_cost_to_async_loggers() -> None:
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.litellm_call_id = "parallel-extract-logging"
|
||||
logging_obj.model_call_details = {}
|
||||
logging_obj.dispatch_success_handlers = AsyncMock()
|
||||
response_body = {
|
||||
"extract_id": "extract_logging",
|
||||
"results": [],
|
||||
"errors": [],
|
||||
"usage": [{"name": "sku_extract_excerpts", "count": 2}],
|
||||
"session_id": "session_logging",
|
||||
}
|
||||
response = httpx.Response(
|
||||
200,
|
||||
json=response_body,
|
||||
request=httpx.Request("POST", "https://api.parallel.ai/v1/extract"),
|
||||
)
|
||||
|
||||
await PassThroughEndpointLogging().pass_through_async_success_handler(
|
||||
httpx_response=response,
|
||||
response_body=response_body,
|
||||
logging_obj=logging_obj,
|
||||
url_route="https://api.parallel.ai/v1/extract",
|
||||
result="",
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
cache_hit=False,
|
||||
request_body={"urls": ["https://example.com/1", "https://example.com/2"]},
|
||||
passthrough_logging_payload={
|
||||
"url": "https://api.parallel.ai/v1/extract",
|
||||
"request_body": {"urls": ["https://example.com/1", "https://example.com/2"]},
|
||||
"request_method": "POST",
|
||||
"cost_per_request": None,
|
||||
},
|
||||
custom_llm_provider="parallel_ai",
|
||||
)
|
||||
|
||||
logging_obj.dispatch_success_handlers.assert_awaited_once()
|
||||
dispatched_kwargs = logging_obj.dispatch_success_handlers.await_args.kwargs
|
||||
assert dispatched_kwargs["model"] == PARALLEL_AI_EXTRACT_MODEL
|
||||
assert dispatched_kwargs["custom_llm_provider"] == "parallel_ai"
|
||||
assert dispatched_kwargs["response_cost"] == 0.002
|
||||
assert dispatched_kwargs["prefer_async_handlers"] is True
|
||||
Loading…
Add table
Reference in a new issue