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

[wuodar]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example 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})

Description

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?

System Info

langchain==0.3.0
langchain-core==0.3.1
langchain-aws==0.2.0
langchain-community==0.3.0

[dosubot[bot]]

The behavior you're observing is due to the way the invoke and stream methods are implemented in the ChatBedrock class. Setting streaming to True in the invoke method does not inherently trigger the on_llm_new_token callback, whereas using the stream method does.

Here's a detailed explanation:

1. invoke Method:
   The invoke method in the ChatBedrock class handles the streaming parameter by checking its value. If streaming is True, it streams the response by calling the _stream method and concatenates the chunks of text received. However, it does not directly trigger the on_llm_new_token callback.

   def _call(
       self,
       prompt: str,
       stop: Optional[List[str]] = None,
       run_manager: Optional[CallbackManagerForLLMRun] = None,
       **kwargs: Any,
   ) -> str:
       """Call out to Bedrock service model.
   
       Args:
           prompt: The prompt to pass into the model.
           stop: Optional list of stop words to use when generating.
   
       Returns:
           The string generated by the model.
   
       Example:
           .. code-block:: python
   
               response = llm.invoke("Tell me a joke.")
       """
   
       if self.streaming:
           completion = ""
           for chunk in self._stream(
               prompt=prompt, stop=stop, run_manager=run_manager, **kwargs
           ):
               completion += chunk.text
           return completion
   
       text, _ = self._prepare_input_and_invoke(
           prompt=prompt, stop=stop, run_manager=run_manager, **kwargs
       )
       return text

2. stream Method:
   The stream method is specifically designed to handle streaming and trigger the on_llm_new_token callback. This method yields chunks of the response as they are received, allowing the callback to be invoked for each new token.

   def _stream(
       self,
       prompt: str,
       stop: Optional[List[str]] = None,
       run_manager: Optional[CallbackManagerForLLMRun] = None,
       **kwargs: Any,
   ) -> Iterator[GenerationChunk]:
       """Call out to Bedrock service with streaming.
   
       Args:
           prompt (str): The prompt to pass into the model
           stop (Optional[List[str]], optional): Stop sequences. These will
               override any stop sequences in the `model_kwargs` attribute.
               Defaults to None.
           run_manager (Optional[CallbackManagerForLLMRun], optional): Callback
               run managers used to process the output. Defaults to None.
   
       Returns:
           Iterator[GenerationChunk]: Generator that yields the streamed responses.
   
       Yields:
           Iterator[GenerationChunk]: Responses from the model.
       """

In summary, the invoke method does not inherently trigger the on_llm_new_token callback even with streaming=True, whereas the stream method is designed to handle streaming and trigger the callback appropriately ^[1].

For your use case, you should use the stream method if you need the on_llm_new_token callback to be triggered.

To continue talking to Dosu, mention @dosu.