Skip to content

Commit c0ed864

Browse files
nuwangeeknuwangeek
committed
Fixed response truncation issue
1 parent 74cb7e5 commit c0ed864

2 files changed

Lines changed: 86 additions & 14 deletions

File tree

‎src/tool_classifier/api_response_formatter.py‎

Lines changed: 44 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
logger = LokiLogger(service_name="api-tool-calling")
2020

2121
_MAX_ITEMS: int = 500
22-
_MAX_RESPONSE_BYTES: int = 50_000
22+
_MAX_RESPONSE_BYTES: int = 150_000
2323

2424

2525
def _get_current_model_name() -> str:
@@ -55,7 +55,7 @@ class APIResponseFormatterSignature(dspy.Signature):
5555
message that no results were found for the query.
5656
- If api_response contains an error field or error status, explain the issue to the user
5757
in a friendly, non-technical way.
58-
- If the data contains more than 20 items, summarize the key highlights rather than
58+
- If the data contains more than 50 items, summarize the key highlights rather than
5959
listing every item. Always mention the total count when summarizing.
6060
- Output must be clean text — no markdown headers (##), no code blocks (```), no raw
6161
JSON. The answer must be ready for direct display to the user.
@@ -375,15 +375,17 @@ async def stream_forward(
375375

376376
stream_started = False
377377
token_count = 0
378-
assembled_answer = ""
378+
accumulated: list[str] = []
379+
final_prediction: dspy.Prediction | None = None
379380
async for chunk in output_stream:
380381
if isinstance(chunk, dspy.streaming.StreamResponse):
381382
if chunk.signature_field_name == "formatted_answer":
382383
stream_started = True
383384
token_count += 1
384-
assembled_answer += chunk.chunk
385+
accumulated.append(chunk.chunk)
385386
yield chunk.chunk
386387
elif isinstance(chunk, dspy.Prediction):
388+
final_prediction = chunk
387389
# dspy.streamify did not stream individual tokens — yield the
388390
# full answer from the final Prediction as a single frame.
389391
if not stream_started:
@@ -394,13 +396,43 @@ async def stream_forward(
394396
"no StreamResponse tokens — yielding full Prediction answer"
395397
)
396398
stream_started = True
397-
assembled_answer = answer
399+
accumulated.append(answer)
398400
yield answer
399401

402+
assembled_answer = "".join(accumulated)
403+
400404
if stream_started and token_count > 0:
401405
logger.debug(
402406
f"APIResponseFormatterModule.stream_forward: streamed {token_count} tokens"
403407
)
408+
# DSPy streaming can drop the last few tokens before EOS.
409+
# The final dspy.Prediction holds the authoritative complete answer.
410+
# Yield any tail that wasn't delivered as StreamResponse chunks.
411+
if final_prediction is not None:
412+
full_answer = getattr(
413+
final_prediction, "formatted_answer", None
414+
)
415+
if full_answer:
416+
streamed_text = assembled_answer
417+
if full_answer.startswith(streamed_text) and len(
418+
full_answer
419+
) > len(streamed_text):
420+
tail = full_answer[len(streamed_text) :]
421+
if tail.strip():
422+
logger.debug(
423+
f"APIResponseFormatterModule.stream_forward: "
424+
f"yielding {len(tail)} missing tail chars from Prediction"
425+
)
426+
assembled_answer += tail
427+
yield tail
428+
elif streamed_text and not full_answer.startswith(
429+
streamed_text
430+
):
431+
logger.warning(
432+
"APIResponseFormatterModule.stream_forward: "
433+
"streamed output is not a prefix of final Prediction; "
434+
"skipping tail reconciliation"
435+
)
404436

405437
if not stream_started:
406438
# Last-resort fallback: blocking forward() — covers cases where
@@ -532,9 +564,14 @@ def _truncate_if_needed(api_response_str: str) -> str:
532564

533565
encoded = api_response_str.encode("utf-8")
534566
if len(encoded) > _MAX_RESPONSE_BYTES:
535-
truncated_str = encoded[:_MAX_RESPONSE_BYTES].decode(
567+
suffix = "\n[NOTE: Response truncated due to size limit]"
568+
suffix_bytes = len(suffix.encode("utf-8"))
569+
truncated_str = encoded[: _MAX_RESPONSE_BYTES - suffix_bytes].decode(
536570
"utf-8", errors="ignore"
537571
)
538-
return truncated_str + "\n[NOTE: Response truncated due to size limit]"
572+
logger.warning(
573+
f"API response exceeded byte limit: {len(encoded)} bytes — truncating to {len(truncated_str.encode('utf-8'))} bytes"
574+
)
575+
return truncated_str + suffix
539576

540577
return api_response_str

‎src/tool_classifier/multi_response_formatter.py‎

Lines changed: 42 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -329,15 +329,17 @@ async def stream_forward_multi(
329329

330330
stream_started = False
331331
token_count = 0
332-
assembled_answer = ""
332+
accumulated: list[str] = []
333+
final_prediction: dspy.Prediction | None = None
333334
async for chunk in output_stream:
334335
if isinstance(chunk, dspy.streaming.StreamResponse):
335336
if chunk.signature_field_name == "unified_answer":
336337
stream_started = True
337338
token_count += 1
338-
assembled_answer += chunk.chunk
339+
accumulated.append(chunk.chunk)
339340
yield chunk.chunk
340341
elif isinstance(chunk, dspy.Prediction):
342+
final_prediction = chunk
341343
# dspy.streamify did not stream individual tokens — yield the
342344
# full answer from the final Prediction as a single frame.
343345
if not stream_started:
@@ -348,14 +350,42 @@ async def stream_forward_multi(
348350
"no StreamResponse tokens — yielding full Prediction answer"
349351
)
350352
stream_started = True
351-
assembled_answer = answer
353+
accumulated.append(answer)
352354
yield answer
353355

356+
assembled_answer = "".join(accumulated)
357+
354358
if stream_started and token_count > 0:
355359
logger.debug(
356360
f"MultiResponseFormatterModule.stream_forward_multi: "
357361
f"streamed {token_count} tokens"
358362
)
363+
# DSPy streaming can drop the last few tokens before EOS.
364+
# The final dspy.Prediction holds the authoritative complete answer.
365+
# Yield any tail that wasn't delivered as StreamResponse chunks.
366+
if final_prediction is not None:
367+
full_answer = getattr(final_prediction, "unified_answer", None)
368+
if full_answer:
369+
streamed_text = assembled_answer
370+
if full_answer.startswith(streamed_text) and len(
371+
full_answer
372+
) > len(streamed_text):
373+
tail = full_answer[len(streamed_text) :]
374+
if tail.strip():
375+
logger.debug(
376+
f"MultiResponseFormatterModule.stream_forward_multi: "
377+
f"yielding {len(tail)} missing tail chars from Prediction"
378+
)
379+
assembled_answer += tail
380+
yield tail
381+
elif streamed_text and not full_answer.startswith(
382+
streamed_text
383+
):
384+
logger.warning(
385+
"MultiResponseFormatterModule.stream_forward_multi: "
386+
"streamed output is not a prefix of final Prediction; "
387+
"skipping tail reconciliation"
388+
)
359389

360390
if not stream_started:
361391
# Last-resort fallback: blocking forward() — covers cases where
@@ -392,6 +422,7 @@ async def stream_forward_multi(
392422
metadata={
393423
"stream_started": stream_started,
394424
"chunk_count": token_count,
425+
"streaming": True,
395426
},
396427
)
397428
except Exception as update_error:
@@ -489,12 +520,16 @@ def _build_results_block(
489520
if total_bytes + section_bytes > _MAX_TOTAL_RESPONSE_BYTES:
490521
remaining = _MAX_TOTAL_RESPONSE_BYTES - total_bytes
491522
if remaining > 0:
523+
suffix = (
524+
"\n[NOTE: Combined results truncated due to total size limit]"
525+
)
526+
suffix_bytes = len(suffix.encode("utf-8"))
492527
encoded = section.encode("utf-8")
493-
truncated = encoded[:remaining].decode("utf-8", errors="ignore")
494-
sections.append(
495-
truncated
496-
+ "\n[NOTE: Combined results truncated due to total size limit]"
528+
truncated = encoded[: max(0, remaining - suffix_bytes)].decode(
529+
"utf-8", errors="ignore"
497530
)
531+
sections.append(truncated + suffix)
532+
total_bytes = _MAX_TOTAL_RESPONSE_BYTES
498533
break
499534

500535
sections.append(section)

0 commit comments

Comments
 (0)