From bb73dd68a6b0bbdd1cc0e01a4bd75daa06f5d445 Mon Sep 17 00:00:00 2001
From: Zachary Lyon <zachary@lyonfamily.org>
Date: Mon, 14 Sep 2026 16:18:47 -0700
Subject: [PATCH] refactor(proxy): make cross-module tinyfish billing hooks
 public

---
 .../tinyfish_passthrough_logging_handler.py                 | 4 ++--
 litellm/proxy/pass_through_endpoints/streaming_handler.py   | 2 +-
 litellm/proxy/pass_through_endpoints/success_handler.py     | 2 +-
 .../test_tinyfish_passthrough_logging_handler.py            | 6 +++---
 4 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py
index 8f1923a3f11..e5b83322a9a 100644
--- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py
+++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py
@@ -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,
diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py
index 2b376bf01e2..69a8a7afc20 100644
--- a/litellm/proxy/pass_through_endpoints/streaming_handler.py
+++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py
@@ -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,
diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py
index e53871a4aae..f3b6ac66515 100644
--- a/litellm/proxy/pass_through_endpoints/success_handler.py
+++ b/litellm/proxy/pass_through_endpoints/success_handler.py
@@ -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(
diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py
index 8c078d3b0c6..1b894850de5 100644
--- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py
+++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py
@@ -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(),