From 52e7ea43325a9b1ebcea1340ec98046242c09c19 Mon Sep 17 00:00:00 2001 From: James Liounis Date: Fri, 21 Aug 2026 16:08:28 -0400 Subject: [PATCH] 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. --- litellm/llms/parallel_ai/extract/__init__.py | 6 + .../parallel_ai/extract/cost_calculator.py | 73 +++++++ litellm/proxy/_types.py | 1 + .../llm_passthrough_endpoints.py | 47 +++++ ...parallel_ai_passthrough_logging_handler.py | 51 +++++ .../pass_through_endpoints/success_handler.py | 13 ++ ruff-strict-budget.json | 2 +- .../extract/test_cost_calculator.py | 90 +++++++++ .../test_parallel_ai_extract_gateway.py | 185 ++++++++++++++++++ ...parallel_ai_passthrough_logging_handler.py | 133 +++++++++++++ 10 files changed, 600 insertions(+), 1 deletion(-) create mode 100644 litellm/llms/parallel_ai/extract/__init__.py create mode 100644 litellm/llms/parallel_ai/extract/cost_calculator.py create mode 100644 litellm/proxy/pass_through_endpoints/llm_provider_handlers/parallel_ai_passthrough_logging_handler.py create mode 100644 tests/test_litellm/llms/parallel_ai/extract/test_cost_calculator.py create mode 100644 tests/test_litellm/llms/parallel_ai/test_parallel_ai_extract_gateway.py create mode 100644 tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_parallel_ai_passthrough_logging_handler.py diff --git a/litellm/llms/parallel_ai/extract/__init__.py b/litellm/llms/parallel_ai/extract/__init__.py new file mode 100644 index 00000000000..b3d6df3959b --- /dev/null +++ b/litellm/llms/parallel_ai/extract/__init__.py @@ -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"] diff --git a/litellm/llms/parallel_ai/extract/cost_calculator.py b/litellm/llms/parallel_ai/extract/cost_calculator.py new file mode 100644 index 00000000000..7ec090e2029 --- /dev/null +++ b/litellm/llms/parallel_ai/extract/cost_calculator.py @@ -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 diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e51a3138d2d..546bddab883 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -466,6 +466,7 @@ class LiteLLMRoutes(enum.Enum): "/openai_passthrough", "/assemblyai", "/eu.assemblyai", + "/parallel_ai", "/vllm", "/mistral", "/milvus", diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 7ce41c1d5b6..3ce62b2f7f7 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -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"], diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/parallel_ai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/parallel_ai_passthrough_logging_handler.py new file mode 100644 index 00000000000..1d98e647702 --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/parallel_ai_passthrough_logging_handler.py @@ -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, + }, + } diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index c38566375f4..1553364bd4b 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -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, diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index a990f7c3830..7b3b67146c8 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -36,7 +36,7 @@ "limit": 177 }, "B008": { - "limit": 503 + "limit": 505 }, "B009": { "limit": 59 diff --git a/tests/test_litellm/llms/parallel_ai/extract/test_cost_calculator.py b/tests/test_litellm/llms/parallel_ai/extract/test_cost_calculator.py new file mode 100644 index 00000000000..83d8532da14 --- /dev/null +++ b/tests/test_litellm/llms/parallel_ai/extract/test_cost_calculator.py @@ -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 diff --git a/tests/test_litellm/llms/parallel_ai/test_parallel_ai_extract_gateway.py b/tests/test_litellm/llms/parallel_ai/test_parallel_ai_extract_gateway.py new file mode 100644 index 00000000000..4df5f3d5333 --- /dev/null +++ b/tests/test_litellm/llms/parallel_ai/test_parallel_ai_extract_gateway.py @@ -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 diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_parallel_ai_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_parallel_ai_passthrough_logging_handler.py new file mode 100644 index 00000000000..431ab51b181 --- /dev/null +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_parallel_ai_passthrough_logging_handler.py @@ -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