Spaces:
Build error
Build error
Commit ·
823b8dc
1
Parent(s): 970307c
Fix remaining review issues: null safety, compressed count, batch tokens
Browse files- M8: Fix null prompt_tokens_details in streaming.py (2 more instances)
- Issue 2: requests_compressed now counts by tokens_saved > 0 only,
not savings_percent > 0 (which could be negative from pipeline overhead)
- H3: Initialize original_tokens=0 before try block in Google batch handler
so error handler has a valid value
- L4: Google batch handler now looks up model context limit via
openai_provider.get_context_limit() instead of hardcoding 128000
headroom/proxy/cost.py
CHANGED
|
@@ -240,7 +240,7 @@ def build_session_summary(
|
|
| 240 |
if entry.model and "count_tokens" in entry.model:
|
| 241 |
uncompressed_reasons["passthrough"] += 1
|
| 242 |
continue
|
| 243 |
-
if entry.tokens_saved > 0
|
| 244 |
compressed_requests.append(
|
| 245 |
{
|
| 246 |
"savings_pct": round(entry.savings_percent, 1),
|
|
|
|
| 240 |
if entry.model and "count_tokens" in entry.model:
|
| 241 |
uncompressed_reasons["passthrough"] += 1
|
| 242 |
continue
|
| 243 |
+
if entry.tokens_saved > 0:
|
| 244 |
compressed_requests.append(
|
| 245 |
{
|
| 246 |
"savings_pct": round(entry.savings_percent, 1),
|
headroom/proxy/handlers/batch.py
CHANGED
|
@@ -134,9 +134,15 @@ class BatchHandlerMixin:
|
|
| 134 |
continue
|
| 135 |
|
| 136 |
# Apply optimization
|
|
|
|
|
|
|
| 137 |
try:
|
| 138 |
-
#
|
| 139 |
-
context_limit =
|
|
|
|
|
|
|
|
|
|
|
|
|
| 140 |
|
| 141 |
# Use OpenAI pipeline (similar message format after conversion)
|
| 142 |
result = self.openai_pipeline.apply(
|
|
@@ -219,8 +225,7 @@ class BatchHandlerMixin:
|
|
| 219 |
)
|
| 220 |
# Pass through unchanged on failure — count original as optimized
|
| 221 |
compressed_requests.append(batch_req)
|
| 222 |
-
|
| 223 |
-
# Just skip the token accounting for this failed request
|
| 224 |
|
| 225 |
# Update body with compressed requests
|
| 226 |
body["batch"]["input_config"]["requests"]["requests"] = compressed_requests
|
|
|
|
| 134 |
continue
|
| 135 |
|
| 136 |
# Apply optimization
|
| 137 |
+
original_tokens = 0 # Set before try so error handler can use it
|
| 138 |
+
optimized_tokens = 0
|
| 139 |
try:
|
| 140 |
+
# Look up model context limit, fall back to 128K
|
| 141 |
+
context_limit = (
|
| 142 |
+
self.openai_provider.get_context_limit(model)
|
| 143 |
+
if hasattr(self, "openai_provider")
|
| 144 |
+
else 128000
|
| 145 |
+
)
|
| 146 |
|
| 147 |
# Use OpenAI pipeline (similar message format after conversion)
|
| 148 |
result = self.openai_pipeline.apply(
|
|
|
|
| 225 |
)
|
| 226 |
# Pass through unchanged on failure — count original as optimized
|
| 227 |
compressed_requests.append(batch_req)
|
| 228 |
+
total_optimized_tokens += original_tokens # 0 if pipeline never ran
|
|
|
|
| 229 |
|
| 230 |
# Update body with compressed requests
|
| 231 |
body["batch"]["input_config"]["requests"]["requests"] = compressed_requests
|
headroom/proxy/handlers/streaming.py
CHANGED
|
@@ -80,7 +80,7 @@ class StreamingMixin:
|
|
| 80 |
usage["input_tokens"] = chunk_usage.get("prompt_tokens", 0)
|
| 81 |
usage["output_tokens"] = chunk_usage.get("completion_tokens", 0)
|
| 82 |
# OpenAI has cached tokens in prompt_tokens_details
|
| 83 |
-
details = chunk_usage.get("prompt_tokens_details"
|
| 84 |
usage["cache_read_input_tokens"] = details.get("cached_tokens", 0)
|
| 85 |
|
| 86 |
elif provider == "gemini":
|
|
@@ -164,7 +164,7 @@ class StreamingMixin:
|
|
| 164 |
if chunk_usage:
|
| 165 |
usage_found["input_tokens"] = chunk_usage.get("prompt_tokens", 0)
|
| 166 |
usage_found["output_tokens"] = chunk_usage.get("completion_tokens", 0)
|
| 167 |
-
details = chunk_usage.get("prompt_tokens_details"
|
| 168 |
usage_found["cache_read_input_tokens"] = details.get("cached_tokens", 0)
|
| 169 |
|
| 170 |
elif provider == "gemini":
|
|
|
|
| 80 |
usage["input_tokens"] = chunk_usage.get("prompt_tokens", 0)
|
| 81 |
usage["output_tokens"] = chunk_usage.get("completion_tokens", 0)
|
| 82 |
# OpenAI has cached tokens in prompt_tokens_details
|
| 83 |
+
details = chunk_usage.get("prompt_tokens_details") or {}
|
| 84 |
usage["cache_read_input_tokens"] = details.get("cached_tokens", 0)
|
| 85 |
|
| 86 |
elif provider == "gemini":
|
|
|
|
| 164 |
if chunk_usage:
|
| 165 |
usage_found["input_tokens"] = chunk_usage.get("prompt_tokens", 0)
|
| 166 |
usage_found["output_tokens"] = chunk_usage.get("completion_tokens", 0)
|
| 167 |
+
details = chunk_usage.get("prompt_tokens_details") or {}
|
| 168 |
usage_found["cache_read_input_tokens"] = details.get("cached_tokens", 0)
|
| 169 |
|
| 170 |
elif provider == "gemini":
|