chopratejas commited on
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 and entry.savings_percent > 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
- # Default context limit for most models
139
- context_limit = 128000
 
 
 
 
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
- # original_tokens may be unbound if pipeline failed before producing result
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":