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:
James Liounis 2026-08-21 16:08:28 -04:00
parent 1ed4ae981a
commit 52e7ea4332
10 changed files with 600 additions and 1 deletions

View 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"]

View 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

View file

@ -466,6 +466,7 @@ class LiteLLMRoutes(enum.Enum):
"/openai_passthrough",
"/assemblyai",
"/eu.assemblyai",
"/parallel_ai",
"/vllm",
"/mistral",
"/milvus",

View file

@ -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"],

View file

@ -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,
},
}

View file

@ -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,

View file

@ -36,7 +36,7 @@
"limit": 177
},
"B008": {
"limit": 503
"limit": 505
},
"B009": {
"limit": 59

View file

@ -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

View file

@ -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

View file

@ -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