Skip to content

[Bug]: SageMaker LLM streaming silently corrupts generated text (lstrip/rstrip used as affix removal) #22221

Description

@ritsth

Bug Description

IOHandler.deserialize_streaming_output in llama-index-llms-sagemaker-endpoint strips the [{"generated_text":" / "}] JSON envelope with str.lstrip / str.rstrip:

response_str = (
    response.decode("utf-8").lstrip('[{"generated_text":"').rstrip('"}]')
)

lstrip/rstrip treat their argument as a set of characters, not a literal prefix/suffix. After the envelope is consumed, stripping continues into the generated text itself: any leading characters from the set (t, g, d, a, e, n, r, x, _, :, ...) and any trailing ", }, ] are silently removed. Since a large fraction of English responses start with words like "the", "a", "data", "great"..., streamed SageMaker output is routinely truncated with no error. Text ending in an escaped quote crashes the subsequent json.loads instead.

Version

llama-index-llms-sagemaker-endpoint 0.5.0 (current main)

Steps to Reproduce

No AWS needed; the function is pure bytes -> str:

from llama_index.llms.sagemaker_endpoint.utils import IOHandler

handler = IOHandler()
print(handler.deserialize_streaming_output(b'[{"generated_text":"the answer is 42"}]'))
# 'he answer is 42'   (leading "t" eaten)
print(handler.deserialize_streaming_output(b'[{"generated_text":"great and neat"}]'))
# ' and neat'         ("great" eaten)
print(handler.deserialize_streaming_output(b'[{"generated_text":"data engineering"}]'))
# ' engineering'      ("data" eaten)

Expected: the generated text unchanged.

Relevant Logs/Tracebacks

No response

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions