scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-03 02:22:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

refactor(proxy): make cross-module tinyfish billing hooks public

Browse source
This commit is contained in:
Zachary Lyon 2026-09-14 16:18:47 -07:00
parent 6692946d70
commit bb73dd68a6
4 changed files with 7 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py
View file

@ -98,7 +98,7 @@ def _run_cost(run: TinyfishRun | None) -> float | None:
class TinyFishPassthroughLoggingHandler:
@staticmethod
def _should_log_request(request_method: str, url_route: str) -> bool:
def should_log_request(request_method: str, url_route: str) -> bool:
"""Only run submissions are billed; GET /v1/runs* polling and cancels never write spend rows."""
return request_method == "POST" and "/v1/automation/" in urlparse(url_route).path
@ -258,7 +258,7 @@ class TinyFishPassthroughLoggingHandler:
return None
@staticmethod
async def _handle_logging_tinyfish_collected_chunks(
async def handle_logging_tinyfish_collected_chunks(
litellm_logging_obj: LiteLLMLoggingObj,
url_route: str,
start_time: datetime,

2
litellm/proxy/pass_through_endpoints/streaming_handler.py
View file

@ -278,7 +278,7 @@ class PassThroughStreamingHandler:
# num_of_steps, so pricing needs an async GET /v1/runs/{id} after the stream.
if endpoint_type == EndpointType.TINYFISH:
tinyfish_payload: Final = (
await TinyFishPassthroughLoggingHandler._handle_logging_tinyfish_collected_chunks(
await TinyFishPassthroughLoggingHandler.handle_logging_tinyfish_collected_chunks(
litellm_logging_obj=litellm_logging_obj,
url_route=url_route,
start_time=start_time,

2
litellm/proxy/pass_through_endpoints/success_handler.py
View file

@ -322,7 +322,7 @@ class PassThroughEndpointLogging:
if self.is_tinyfish_route(url_route, custom_llm_provider):
# GET /v1/runs* polling and cancels never write spend rows; run-async bills once,
# from a background poller that re-enters _handle_logging at run completion.
if not TinyFishPassthroughLoggingHandler._should_log_request(httpx_response.request.method, url_route):
if not TinyFishPassthroughLoggingHandler.should_log_request(httpx_response.request.method, url_route):
return
if TinyFishPassthroughLoggingHandler.is_run_async_route(url_route):
TinyFishPassthroughLoggingHandler.start_async_run_billing(

6
tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py
View file

@ -77,7 +77,7 @@ class TestBillingGate:
],
)
def test_only_run_submissions_are_billed(self, method, url, expected):
assert TinyFishPassthroughLoggingHandler._should_log_request(method, url) is expected
assert TinyFishPassthroughLoggingHandler.should_log_request(method, url) is expected
def test_polling_writes_no_spend_row(self, tinyfish_env):
logging_obj = _make_logging_obj()
@ -208,7 +208,7 @@ class TestSseBilling:
)
payload = asyncio.run(
TinyFishPassthroughLoggingHandler._handle_logging_tinyfish_collected_chunks(
TinyFishPassthroughLoggingHandler.handle_logging_tinyfish_collected_chunks(
litellm_logging_obj=logging_obj,
url_route="https://agent.tinyfish.ai/v1/automation/run-sse",
start_time=datetime.now(),
@ -226,7 +226,7 @@ class TestSseBilling:
fake_client = _FakeClient(payloads=[{}])
payload = asyncio.run(
TinyFishPassthroughLoggingHandler._handle_logging_tinyfish_collected_chunks(
TinyFishPassthroughLoggingHandler.handle_logging_tinyfish_collected_chunks(
litellm_logging_obj=_make_logging_obj(),
url_route="https://agent.tinyfish.ai/v1/automation/run-sse",
start_time=datetime.now(),