Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py: 28%
190 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1import asyncio
2import json
3import os
4import time
5import urllib.parse
6from collections.abc import Mapping, Sequence
7from datetime import datetime
8from types import MappingProxyType
9from typing import Final, NamedTuple
10from urllib.parse import urlparse
12import httpx
13from pydantic import TypeAdapter, ValidationError
15from litellm._logging import verbose_proxy_logger
16from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
17from litellm.litellm_core_utils.litellm_logging import (
18 get_standard_logging_object_payload,
19)
20from litellm.llms.custom_httpx.http_handler import (
21 AsyncHTTPHandler,
22 get_async_httpx_client,
23)
24from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
25from litellm.types.llms.custom_http import httpxSpecialProvider
26from litellm.types.passthrough_endpoints.tinyfish import (
27 TINYFISH_AGENT_DEFAULT_API_BASE,
28 TINYFISH_DEFAULT_COST_PER_STEP,
29 TINYFISH_MAX_CONSECUTIVE_POLL_FAILURES,
30 TINYFISH_MAX_POLLING_SECONDS,
31 TINYFISH_MODEL_NAME,
32 TINYFISH_POLLING_INTERVAL_SECONDS,
33 TINYFISH_TERMINAL_RUN_STATUSES,
34 TinyfishRun,
35)
36from litellm.types.utils import StandardPassThroughResponseObject
38_RUN_ADAPTER: Final = TypeAdapter(TinyfishRun)
40_EMPTY_KWARGS: Final[Mapping[str, object]] = MappingProxyType({})
43class _TinyfishLoggingPayload(NamedTuple):
44 result: StandardPassThroughResponseObject
45 kwargs: Mapping[str, object]
47 def as_handler_result(self) -> PassThroughEndpointLoggingTypedDict:
48 handler_result: Final[PassThroughEndpointLoggingTypedDict] = {
49 "result": self.result,
50 "kwargs": {**self.kwargs},
51 }
52 return handler_result
55# asyncio tasks are weakly referenced by the loop; hold them until done or they can vanish mid-poll
56_BACKGROUND_BILLING_TASKS: Final[set["asyncio.Task[None]"]] = set() # mutable-ok: task registry
59def _register_billing_task(task: "asyncio.Task[None]") -> None:
60 _BACKGROUND_BILLING_TASKS.add(task)
61 task.add_done_callback(_BACKGROUND_BILLING_TASKS.discard)
62 task.add_done_callback(_warn_if_cancelled)
65def _warn_if_cancelled(task: "asyncio.Task[None]") -> None:
66 # CancelledError bypasses the poller's exception handler, so shutdown-time charge loss must be logged here
67 if task.cancelled():
68 verbose_proxy_logger.warning("TinyFish passthrough: billing poller cancelled mid-poll; the run may go unbilled")
71_SSE_POLLER_SPAWNED_KEY: Final = "tinyfish_sse_poller_spawned"
74def mark_sse_poller_spawned(logging_obj: LiteLLMLoggingObj) -> None:
75 logging_obj.model_call_details[_SSE_POLLER_SPAWNED_KEY] = True # rebind-ok: request-scoped scratch dict
78def sse_poller_spawned(logging_obj: LiteLLMLoggingObj) -> bool:
79 return logging_obj.model_call_details.get(_SSE_POLLER_SPAWNED_KEY) is True
82def run_id_from_sse_frames(frames: bytes) -> str | None:
83 return _run_id_from_sse_chunks(frames.decode("utf-8", errors="replace").splitlines())
86def resolve_tinyfish_agent_api_base() -> str:
87 raw: Final = (os.getenv("TINYFISH_AGENT_API_BASE") or TINYFISH_AGENT_DEFAULT_API_BASE).rstrip("/")
88 # a schemeless override would silently break both routing and billing (urlparse hostname becomes None)
89 return raw if "://" in raw else f"https://{raw}"
92def resolve_tinyfish_cost_per_step() -> float:
93 raw: Final = os.getenv("TINYFISH_COST_PER_STEP")
94 if raw is None:
95 return TINYFISH_DEFAULT_COST_PER_STEP
96 try:
97 return float(raw)
98 except ValueError:
99 verbose_proxy_logger.warning(
100 "TINYFISH_COST_PER_STEP=%r is not a number; using the default rate %s",
101 raw,
102 TINYFISH_DEFAULT_COST_PER_STEP,
103 )
104 return TINYFISH_DEFAULT_COST_PER_STEP
107def is_tinyfish_agent_url(url: str) -> bool:
108 hostname: Final = urlparse(url).hostname
109 return hostname is not None and hostname == urlparse(resolve_tinyfish_agent_api_base()).hostname
112def _parse_run(payload: object) -> TinyfishRun | None:
113 try:
114 return _RUN_ADAPTER.validate_python(payload)
115 except ValidationError as e:
116 verbose_proxy_logger.warning("TinyFish passthrough: unexpected run object shape: %s", e)
117 return None
120def _run_cost(run: TinyfishRun | None) -> float | None:
121 if run is None:
122 return None
123 # TinyFish only invoices COMPLETED runs, so FAILED/CANCELLED runs must charge the team $0
124 if run.get("status") != "COMPLETED":
125 return None
126 num_of_steps: Final = run.get("num_of_steps")
127 if num_of_steps is None:
128 return None
129 return num_of_steps * resolve_tinyfish_cost_per_step()
132class TinyFishPassthroughLoggingHandler:
133 @staticmethod
134 def should_log_request(request_method: str, url_route: str) -> bool:
135 """Only run submissions are billed; GET /v1/runs* polling and cancels never write spend rows."""
136 return request_method == "POST" and "/v1/automation/" in urlparse(url_route).path
138 @staticmethod
139 def is_run_async_route(url_route: str) -> bool:
140 return urlparse(url_route).path.endswith("/v1/automation/run-async")
142 @staticmethod
143 def tinyfish_passthrough_handler(
144 httpx_response: httpx.Response,
145 response_body: Mapping[str, object] | None,
146 logging_obj: LiteLLMLoggingObj,
147 url_route: str,
148 result: str,
149 start_time: datetime,
150 end_time: datetime,
151 cache_hit: bool,
152 request_body: Mapping[str, object],
153 **kwargs: object, # kwargs-ok: the passthrough logging dispatch forwards shared logging kwargs to every handler
154 ) -> PassThroughEndpointLoggingTypedDict:
155 """Bill a blocking POST /v1/automation/run: the response is the terminal run object."""
156 try:
157 run: Final = _parse_run(response_body) if response_body is not None else None
158 handler_payload: Final = TinyFishPassthroughLoggingHandler._build_logging_payload(
159 run=run,
160 logging_obj=logging_obj,
161 result=result,
162 start_time=start_time,
163 end_time=end_time,
164 kwargs=kwargs,
165 ).as_handler_result()
166 except Exception as e: # noqa: BLE001 # billing/logging must never break the relayed request
167 verbose_proxy_logger.exception("Error in TinyFish passthrough logging handler: %s", e)
168 fallback_payload: Final[PassThroughEndpointLoggingTypedDict] = {
169 "result": StandardPassThroughResponseObject(response=result),
170 "kwargs": kwargs,
171 }
172 return fallback_payload
173 return handler_payload
175 @staticmethod
176 def start_async_run_billing(
177 response_body: Mapping[str, object] | None,
178 logging_obj: LiteLLMLoggingObj,
179 result: str,
180 start_time: datetime,
181 cache_hit: bool,
182 **kwargs: object, # kwargs-ok: shared logging kwargs, replayed into _handle_logging when the run finishes
183 ) -> None:
184 """Bill POST /v1/automation/run-async once, when the polled run turns terminal."""
185 submitted: Final = _parse_run(response_body) if response_body is not None else None
186 run_id: Final = submitted.get("run_id") if submitted is not None else None
187 if not run_id:
188 verbose_proxy_logger.warning(
189 "TinyFish passthrough: run-async response carried no run_id; logging the request without cost"
190 )
191 task: Final = asyncio.create_task(
192 TinyFishPassthroughLoggingHandler._poll_and_log(
193 run_id=run_id,
194 logging_obj=logging_obj,
195 result=result,
196 start_time=start_time,
197 cache_hit=cache_hit,
198 kwargs=kwargs,
199 )
200 )
201 _register_billing_task(task)
203 @staticmethod
204 def start_sse_run_billing(
205 run_id: str,
206 litellm_logging_obj: LiteLLMLoggingObj,
207 start_time: datetime,
208 client: AsyncHTTPHandler | None = None,
209 ) -> None:
210 """Bill POST /v1/automation/run-sse once via a detached poller that outlives client disconnects."""
211 mark_sse_poller_spawned(litellm_logging_obj)
212 task: Final = asyncio.create_task(
213 TinyFishPassthroughLoggingHandler._poll_and_log(
214 run_id=run_id,
215 logging_obj=litellm_logging_obj,
216 result="",
217 start_time=start_time,
218 cache_hit=litellm_logging_obj.model_call_details.get("cache_hit") is True,
219 kwargs=_EMPTY_KWARGS,
220 client=client,
221 )
222 )
223 _register_billing_task(task)
225 @staticmethod
226 async def _poll_and_log(
227 run_id: str | None,
228 logging_obj: LiteLLMLoggingObj,
229 result: str,
230 start_time: datetime,
231 cache_hit: bool,
232 kwargs: Mapping[str, object],
233 client: AsyncHTTPHandler | None = None,
234 ) -> None:
235 from ..pass_through_endpoints import pass_through_endpoint_logging
237 try:
238 run: Final = (
239 await TinyFishPassthroughLoggingHandler._poll_until_terminal(run_id, client) if run_id else None
240 )
241 run_end_time: Final = datetime.now() # noqa: DTZ005 # naive to match the start_time stamped by pass_through_request
242 payload: Final = TinyFishPassthroughLoggingHandler._build_logging_payload(
243 run=run,
244 logging_obj=logging_obj,
245 result=result,
246 start_time=start_time,
247 end_time=run_end_time,
248 kwargs=kwargs,
249 )
250 await pass_through_endpoint_logging._handle_logging( # pyright: ignore[reportPrivateUsage] # shared passthrough logging dispatcher, same access as the assemblyai handler
251 logging_obj=logging_obj,
252 standard_logging_response_object=payload.result,
253 result=result,
254 start_time=start_time,
255 end_time=run_end_time,
256 cache_hit=cache_hit,
257 **payload.kwargs,
258 )
259 except Exception as e: # noqa: BLE001 # billing/logging must never break the relayed request
260 verbose_proxy_logger.exception("[Non blocking logging error] TinyFish run-async billing failed: %s", e)
262 @staticmethod
263 async def _poll_until_terminal(
264 run_id: str,
265 client: AsyncHTTPHandler | None = None,
266 poll_interval_seconds: float = TINYFISH_POLLING_INTERVAL_SECONDS,
267 ) -> TinyfishRun | None:
268 deadline: Final = time.monotonic() + TINYFISH_MAX_POLLING_SECONDS
269 last_run: TinyfishRun | None = None # rebind-ok: poll-loop state
270 consecutive_failures = 0 # rebind-ok: poll-loop state
271 while time.monotonic() < deadline:
272 run = await TinyFishPassthroughLoggingHandler._fetch_run(run_id, client)
273 if run is None:
274 # a single transient poll failure must not drop the run's charge
275 consecutive_failures += 1
276 if consecutive_failures >= TINYFISH_MAX_CONSECUTIVE_POLL_FAILURES:
277 verbose_proxy_logger.warning(
278 "TinyFish passthrough: giving up on run %s after %s consecutive poll failures; "
279 "logging the request without cost",
280 run_id,
281 consecutive_failures,
282 )
283 return last_run
284 else:
285 consecutive_failures = 0
286 last_run = run
287 if (run.get("status") or "") in TINYFISH_TERMINAL_RUN_STATUSES:
288 return run
289 await asyncio.sleep(poll_interval_seconds)
290 verbose_proxy_logger.warning(
291 "TinyFish passthrough: run %s not terminal after %ss; logging the request without cost",
292 run_id,
293 TINYFISH_MAX_POLLING_SECONDS,
294 )
295 return last_run
297 @staticmethod
298 async def _fetch_run(run_id: str, client: AsyncHTTPHandler | None = None) -> TinyfishRun | None:
299 from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
300 passthrough_endpoint_router,
301 )
303 api_key: Final = passthrough_endpoint_router.get_credentials(custom_llm_provider="tinyfish", region_name=None)
304 if api_key is None:
305 verbose_proxy_logger.warning("TinyFish passthrough: no API key available to poll run %s", run_id)
306 return None
307 if any(c in run_id for c in ("/", "\\", "#", "?")) or ".." in run_id:
308 verbose_proxy_logger.warning("TinyFish passthrough: invalid run_id %r", run_id)
309 return None
310 safe_run_id: Final = urllib.parse.quote(run_id, safe="")
311 resolved_client: Final = client or get_async_httpx_client(
312 llm_provider=httpxSpecialProvider.PassThroughEndpoint,
313 params={"timeout": 30.0}, # mutable-ok: get_async_httpx_client takes a plain dict of client params
314 )
315 try:
316 # screenshots=none keeps the poll payload small (no per-step screenshot URLs needed)
317 response: Final = await resolved_client.get(
318 f"{resolve_tinyfish_agent_api_base()}/v1/runs/{safe_run_id}?screenshots=none",
319 headers={"X-API-Key": api_key}, # mutable-ok: httpx headers= takes a plain dict
320 )
321 if not (200 <= response.status_code < 300):
322 verbose_proxy_logger.warning(
323 "TinyFish passthrough: GET /v1/runs/%s returned %s", safe_run_id, response.status_code
324 )
325 return None
326 payload: Final[object] = response.json() # any-ok: httpx Response.json() -> Any
327 return _parse_run(payload)
328 except Exception as e: # noqa: BLE001 # billing/logging must never break the relayed request
329 verbose_proxy_logger.warning("[Non blocking logging error] TinyFish run fetch failed: %s", e)
330 return None
332 @staticmethod
333 async def handle_logging_tinyfish_collected_chunks(
334 litellm_logging_obj: LiteLLMLoggingObj,
335 url_route: str,
336 start_time: datetime,
337 all_chunks: Sequence[str],
338 end_time: datetime,
339 client: AsyncHTTPHandler | None = None,
340 ) -> PassThroughEndpointLoggingTypedDict:
341 """Fallback for run-sse streams with no poller: logs the request, pricing via one GET if a run_id parses."""
342 try:
343 run_id: Final = _run_id_from_sse_chunks(all_chunks)
344 if run_id is None:
345 verbose_proxy_logger.warning(
346 "TinyFish passthrough: no run_id in SSE stream; logging the request without cost"
347 )
348 run: Final = await TinyFishPassthroughLoggingHandler._fetch_run(run_id, client) if run_id else None
349 payload: Final = TinyFishPassthroughLoggingHandler._build_logging_payload(
350 run=run,
351 logging_obj=litellm_logging_obj,
352 result="",
353 start_time=start_time,
354 end_time=end_time,
355 kwargs=_EMPTY_KWARGS,
356 ).as_handler_result()
357 except Exception as e: # noqa: BLE001 # billing/logging must never break the relayed request
358 verbose_proxy_logger.exception("Error in TinyFish SSE passthrough logging handler: %s", e)
359 fallback_payload: Final[PassThroughEndpointLoggingTypedDict] = {
360 "result": StandardPassThroughResponseObject(response=""),
361 "kwargs": {},
362 }
363 return fallback_payload
364 return payload
366 @staticmethod
367 def _build_logging_payload(
368 run: TinyfishRun | None,
369 logging_obj: LiteLLMLoggingObj,
370 result: str,
371 start_time: datetime,
372 end_time: datetime,
373 kwargs: Mapping[str, object],
374 ) -> _TinyfishLoggingPayload:
375 response_cost: Final = _run_cost(run)
376 updated_kwargs: Final = { # mutable-ok: the logging pipeline requires a plain kwargs dict
377 **kwargs,
378 "model": TINYFISH_MODEL_NAME,
379 "custom_llm_provider": "tinyfish",
380 "response_cost": response_cost,
381 # spend rows key on this as request_id; without it every poller-billed row is a NULL-key collision
382 "litellm_call_id": logging_obj.litellm_call_id,
383 # the poller paths pass no request kwargs, so SLO attribution (key hash, team, tags) needs the stored params
384 "litellm_params": kwargs.get("litellm_params")
385 or logging_obj.model_call_details.get("litellm_params")
386 or {}, # mutable-ok: the logging pipeline requires a plain kwargs dict
387 }
388 logging_obj.model_call_details.update(
389 model=TINYFISH_MODEL_NAME,
390 custom_llm_provider="tinyfish",
391 response_cost=response_cost,
392 )
394 logged_response: Final = StandardPassThroughResponseObject(
395 response=json.dumps(run) if run is not None else result
396 )
397 standard_logging_object: Final = get_standard_logging_object_payload(
398 kwargs=updated_kwargs,
399 init_response_obj=logged_response,
400 start_time=start_time,
401 end_time=end_time,
402 logging_obj=logging_obj,
403 status="success",
404 )
405 return _TinyfishLoggingPayload(
406 result=logged_response,
407 kwargs=MappingProxyType({**updated_kwargs, "standard_logging_object": standard_logging_object}),
408 )
411def _run_id_from_sse_chunks(all_chunks: Sequence[str]) -> str | None:
412 for line in all_chunks:
413 if not line.startswith("data:"):
414 continue
415 try:
416 event_payload: object = json.loads(line[5:].strip()) # any-ok: json.loads -> Any
417 except json.JSONDecodeError:
418 continue
419 event = _parse_run(event_payload)
420 if event is None:
421 continue
422 run_id = event.get("run_id")
423 if run_id:
424 return run_id
425 return None