Setting streaming to True doesn't seems to make any difference

[wuodar]

Hi,

If I try to run following code:

import boto3
from langchain_aws import ChatBedrock
from typing import Any
from langchain.callbacks.base import BaseCallbackHandler


class CallbackHandler(BaseCallbackHandler):
    def on_llm_new_token(self, token: str, **kwargs: Any) -> Any:
        print(token, end="|")


session = boto3.Session(profile_name='sandbox')
bedrock_client = session.client(service_name='bedrock-runtime', region_name='us-east-1')
model_id = "anthropic.claude-3-5-sonnet-20240620-v1:0"
model = ChatBedrock(model_id=model_id, client=bedrock_client, streaming=True)
model.invoke("Hello, how are you?", config={"callbacks": [CallbackHandler()]})

# output:
# AIMessage(content="Hello! As an AI language model, I don't have feelings, but I'm functioning well and ready to assist you. How can I help you today?", additional_kwargs={'usage': {'prompt_tokens': 13, 'completion_tokens': 35, 'total_tokens': 48}, 'stop_reason': 'end_turn', 'model_id': 'anthropic.claude-3-5-sonnet-20240620-v1:0'}, response_metadata={'usage': {'prompt_tokens': 13, 'completion_tokens': 35, 'total_tokens': 48}, 'stop_reason': 'end_turn', 'model_id': 'anthropic.claude-3-5-sonnet-20240620-v1:0'}, id='run-0e73f036-39a4-402f-8e1a-438d1b9bd03b-0', usage_metadata={'input_tokens': 13, 'output_tokens': 35, 'total_tokens': 48})

The on_llm_new_token won't be called, even though I set streaming to True.

But when i run this:

import boto3
from langchain_aws import ChatBedrock
from typing import Any
from langchain.callbacks.base import BaseCallbackHandler


class CallbackHandler(BaseCallbackHandler):
    def on_llm_new_token(self, token: str, **kwargs: Any) -> Any:
        print(token, end="|")


session = boto3.Session(profile_name='sandbox')
bedrock_client = session.client(service_name='bedrock-runtime', region_name='us-east-1')
model_id = "anthropic.claude-3-5-sonnet-20240620-v1:0"
model = ChatBedrock(model_id=model_id, client=bedrock_client, streaming=True)
for event in model.stream("Hello, how are you?", config={"callbacks": [CallbackHandler()]}):
    pass

# output: 
# |Hello|! As| an AI| language| model, I don|'t have| feelings|,| but I'm| functioning| well| an|d ready| to assist| you.| How can| I help| you today?|||

It works.

My question is why is that the case? I'm using langchain 3.0.0, before that, calling invoke with streamingequal to True worked. Now it looks like I need to call stream even if I set streaming parameter.

It's confusing. Can anyone explain why is that the case?

[shabie]

I can look into it.

[valedica-core]

Here seems the problem 0fe4dd9#diff-7b555324692cbefa47fc582ade684c4faa5d083023137e2839f37aff0401aa4eL832

---

It was probably done on purpose. I fixed it anyway in my case by extending langchain_core.tracers.event_stream._AstreamEventsCallbackHandler in my CallbackHandler

[shabie]

Same issue mentioned here too: #134 (comment)

[3coins]

@wuodar
For streaming, you don't need to use any callbacks any more. You can directly access the chunks with stream method. Here is a slight modification to your code that should work.

model = ChatBedrock(model_id=model_id, client=bedrock_client)
for chunk in model.stream("Hello, how are you?"):
    print(chunk.content, end="", flush=True)

For more info on streaming support in LangChain, see here.
https://python.langchain.com/v0.2/docs/how_to/streaming/

[roie-cye]

Here seems the problem 0fe4dd9#diff-7b555324692cbefa47fc582ade684c4faa5d083023137e2839f37aff0401aa4eL832

It was probably done on purpose. I fixed it anyway in my case by extending langchain_core.tracers.event_stream._AstreamEventsCallbackHandler in my CallbackHandler

could you share your fix?

[wuodar]

@wuodar

For streaming, you don't need to use any callbacks any more. You can directly access the chunks with stream method. Here is a slight modification to your code that should work.

model = ChatBedrock(model_id=model_id, client=bedrock_client)

for chunk in model.stream("Hello, how are you?"):

    print(chunk.content, end="", flush=True)

For more info on streaming support in LangChain, see here.

https://python.langchain.com/v0.2/docs/how_to/streaming/

Actually, I treat callback handler as a central place to forward messages to end user (e.g. WebSocket connection) so I'd like to stick to it

[valedica-core]

Here seems the problem 0fe4dd9#diff-7b555324692cbefa47fc582ade684c4faa5d083023137e2839f37aff0401aa4eL832
It was probably done on purpose. I fixed it anyway in my case by extending langchain_core.tracers.event_stream._AstreamEventsCallbackHandler in my CallbackHandler

could you share your fix?

from langchain_core.tracers.event_stream._AstreamEventsCallbackHandler

class CallbackHandler(_AstreamEventsCallbackHandler):
    def on_llm_new_token(self, token: str, **kwargs: Any) -> Any:
        print(token, end="|")
...
await model.ainvoke("Hello, how are you?", config={"callbacks": [CallbackHandler()]})