Bedrock token count callbacks

[NAPTlME]

Updated both the BedrockLLM and ChatBedrock classes to yield token counts and stop reasons upon generation/call. This works for streaming/non-streaming as well as messages vs raw text.

The goal behind this is to take the input/output tokens and stop reasons directly from the Bedrock call and use them in a CallbackHandler on_llm_end.

Example use

from typing import Any
from uuid import UUID

from langchain.callbacks.base import BaseCallbackHandler
from langchain.chains import ConversationChain
from langchain.memory import ConversationTokenBufferMemory

from langchain_core.runnables import RunnableConfig

from langchain_aws.chat_models import ChatBedrock
from langchain_core.outputs import LLMResult

class BedrockHandler(BaseCallbackHandler):

    def __init__(self, initial_text=""):
        self.text = initial_text
        self.input_token_count = 0
        self.output_token_count = 0
        self.stop_reason = None

    def on_llm_new_token(self, token: str, **kwargs):
        self.text += token
        # do something

    def on_llm_end(
        self,
        response: LLMResult,
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **kwargs: Any,
    ) -> Any:
        if response.llm_output is not None:
            self.input_token_count = response.llm_output.get("usage", {}).get("prompt_tokens", None)
            self.output_token_count = response.llm_output.get("usage", {}).get("completion_tokens", None)
            self.stop_reason = response.llm_output.get("stop_reason", None)

llm = ChatBedrock(model_id="anthropic.claude-3-sonnet-20240229-v1:0", streaming=True)

memory = ConversationTokenBufferMemory(llm=llm)

chain = ConversationChain(llm=llm, memory=memory)

callback = BedrockHandler()

input_prompt = "Write an explanation of math in 3 sentences"
stop_sequences = ["\n\nHuman:"]

response = chain.invoke(
    input_prompt,
    RunnableConfig(callbacks=[callback]),
    stop=stop_sequences,
    max_tokens_to_sample = 1024,
)

print(f"Input tokens: {callback.input_token_count}, Output tokens: {callback.output_token_count}, Stop reason: {callback.stop_reason}")

[3coins]

@NAPTlME Thanks for submitting this update. Please fix the lint and test errors from the CI.

[NAPTlME]

@NAPTlME Thanks for submitting this update. Please fix the lint and test errors from the CI.

@3coins Apologies, I'm on Windows, but manually running Ruff and the unit tests (rather than via make). I also don't have the integration tests set up, but have tested with my current project. If anything further fails, I will take a look at my options to run the makefile.

[NAPTlME]

@3coins I went the WSL route to run the makefile. I fixed some references in the makefile that appear to be holdovers from when this was a part of langchain. Currently, all of those checks pass.

The only item I was unable to run was the integration test (I assume due to what I have provisioned using my AWS credentials).

Thanks.

[NAPTlME]

@3coins Hoping to kick off that workflow again. I'm expecting no further issues, but will address them if CI picks anything else up.

[NAPTlME]

@3coins Addressed minor (import and readme) conflicts introduced by recent changes to main.

[NAPTlME]

@3coins Are there any issues with merging this? (Or am I missing any part of your process for contributions?) Thanks.

[DanielWhite95]

Hello, thanks for the suggestion. This could be useful to have in the package. Only one note: I think it should add the token count on each llm_end without resetting it with the new value. This could be useful if used in cases where the application will do different calls to the model with the same callback

[NAPTlME]

@DanielWhite95 Thanks for taking a look at it. For my particular use-case we are logging each transaction to track costs/chargeback (thus the need to overwrite the token counts). We make calls to our logger from on_llm_end as well to log each call and usage. Also for the case using the input+output tokens to know the input tokens associated with the context for the next query.

That said, if you wanted to increment, this is just in the callback example so you would modify the callback handler to do so

    if response.llm_output is not None:
            self.input_token_count += response.llm_output.get("usage", {}).get("prompt_tokens", 0)
            self.output_token_count+ = response.llm_output.get("usage", {}).get("completion_tokens", 0)
            self.stop_reason = response.llm_output.get("stop_reason", None)

[NAPTlME]

@3coins @efriis Hey. Not trying bug you here, but I would like to know if it is possible to contribute to this project. If so, is there anything further you need from me for this PR? Thanks.

[efriis]

@3coins is your man!

[3coins]

@NAPTlME Thanks for making all the updates, I had been sick out of office this past week so could not get to this earlier. Your code looks good overall, would prefer less nested blocks if possible. Also, there seems to be a lot of changes here for me to process, so give me until end of tomorrow to merge this.

Can you do one last update and convert your sample code into an integration/unit test?

[3coins]

@NAPTlME There are a few integration test errors with this PR. Can you check these and make sure these pass. I have added details here.

tests/integration_tests/chat_models/test_bedrock.py ....FFF.......FF

(base) ➜ aws git:(bedrock-token-count-callbacks) poetry run pytest tests/integration_tests/chat_models/test_bedrock.py ======================================================================================== test session starts ========================================================================================= platform darwin -- Python 3.9.13, pytest-7.4.4, pluggy-1.5.0 rootdir: /Users/pijain/projects/langchain-aws-dev/langchain-aws/libs/aws configfile: pyproject.toml plugins: syrupy-4.6.1, anyio-4.3.0, cov-4.1.0, asyncio-0.23.6 asyncio: mode=auto collected 16 items tests/integration_tests/chat_models/test_bedrock.py ....FFF.......FF [100%] ============================================================================================== FAILURES ============================================================================================== ____________________________________________________________________________ test_chat_bedrock_streaming_generation_info _____________________________________________________________________________ @pytest.mark.scheduled def test_chat_bedrock_streaming_generation_info() -> None: """Test that generation info is preserved when streaming.""" class _FakeCallback(FakeCallbackHandler): saved_things: dict = {} def on_llm_end( self, *args: Any, **kwargs: Any, ) -> Any: # Save the generation self.saved_things["generation"] = args[0] callback = _FakeCallback() chat = ChatBedrock( # type: ignore[call-arg] model_id="anthropic.claude-v2", callbacks=[callback], ) > list(chat.stream("hi")) tests/integration_tests/chat_models/test_bedrock.py:97: _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/language_models/chat_models.py:249: in stream raise e /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/language_models/chat_models.py:240: in stream generation += chunk /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/outputs/chat_generation.py:79: in __add__ message=self.message + other.message, /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/messages/ai.py:145: in __add__ response_metadata = merge_dicts( /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/utils/_merge.py:34: in merge_dicts merged[right_k] = merge_dicts(merged[right_k], right_v) _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ left = {'input_tokens': 10, 'output_tokens': 1}, right = {'output_tokens': 6} def merge_dicts(left: Dict[str, Any], right: Dict[str, Any]) -> Dict[str, Any]: """Merge two dicts, handling specific scenarios where a key exists in both dictionaries but has a value of None in 'left'. In such cases, the method uses the value from 'right' for that key in the merged dictionary. Example: If left = {"function_call": {"arguments": None}} and right = {"function_call": {"arguments": "{\n"}} then, after merging, for the key "function_call", the value from 'right' is used, resulting in merged = {"function_call": {"arguments": "{\n"}}. """ merged = left.copy() for right_k, right_v in right.items(): if right_k not in merged: merged[right_k] = right_v elif right_v is not None and merged[right_k] is None: merged[right_k] = right_v elif right_v is None: continue elif type(merged[right_k]) != type(right_v): raise TypeError( f'additional_kwargs["{right_k}"] already exists in this message,' " but with a different type." ) elif isinstance(merged[right_k], str): merged[right_k] += right_v elif isinstance(merged[right_k], dict): merged[right_k] = merge_dicts(merged[right_k], right_v) elif isinstance(merged[right_k], list): merged[right_k] = merge_lists(merged[right_k], right_v) elif merged[right_k] == right_v: continue else: > raise TypeError( f"Additional kwargs key {right_k} already exists in left dict and " f"value has unsupported type {type(merged[right_k])}." ) E TypeError: Additional kwargs key output_tokens already exists in left dict and value has unsupported type . /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/utils/_merge.py:40: TypeError _______________________________________________________________________________________ test_bedrock_streaming _______________________________________________________________________________________ chat = ChatBedrock(client=, region_name='us-west-2', model_id='anthropic.claude-v2', model_kwargs={'temperature': 0}) @pytest.mark.scheduled def test_bedrock_streaming(chat: ChatBedrock) -> None: """Test streaming tokens from OpenAI.""" full = None for token in chat.stream("I'm Pickle Rick"): > full = token if full is None else full + token # type: ignore[operator] tests/integration_tests/chat_models/test_bedrock.py:109: _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/messages/ai.py:145: in __add__ response_metadata = merge_dicts( /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/utils/_merge.py:34: in merge_dicts merged[right_k] = merge_dicts(merged[right_k], right_v) _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ left = {'input_tokens': 13, 'output_tokens': 1}, right = {'output_tokens': 46} def merge_dicts(left: Dict[str, Any], right: Dict[str, Any]) -> Dict[str, Any]: """Merge two dicts, handling specific scenarios where a key exists in both dictionaries but has a value of None in 'left'. In such cases, the method uses the value from 'right' for that key in the merged dictionary. Example: If left = {"function_call": {"arguments": None}} and right = {"function_call": {"arguments": "{\n"}} then, after merging, for the key "function_call", the value from 'right' is used, resulting in merged = {"function_call": {"arguments": "{\n"}}. """ merged = left.copy() for right_k, right_v in right.items(): if right_k not in merged: merged[right_k] = right_v elif right_v is not None and merged[right_k] is None: merged[right_k] = right_v elif right_v is None: continue elif type(merged[right_k]) != type(right_v): raise TypeError( f'additional_kwargs["{right_k}"] already exists in this message,' " but with a different type." ) elif isinstance(merged[right_k], str): merged[right_k] += right_v elif isinstance(merged[right_k], dict): merged[right_k] = merge_dicts(merged[right_k], right_v) elif isinstance(merged[right_k], list): merged[right_k] = merge_lists(merged[right_k], right_v) elif merged[right_k] == right_v: continue else: > raise TypeError( f"Additional kwargs key {right_k} already exists in left dict and " f"value has unsupported type {type(merged[right_k])}." ) E TypeError: Additional kwargs key output_tokens already exists in left dict and value has unsupported type . /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/utils/_merge.py:40: TypeError ________________________________________________________________________________________ test_bedrock_astream ________________________________________________________________________________________ chat = ChatBedrock(client=, region_name='us-west-2', model_id='anthropic.claude-v2', model_kwargs={'temperature': 0}) @pytest.mark.scheduled async def test_bedrock_astream(chat: ChatBedrock) -> None: """Test streaming tokens from OpenAI.""" > async for token in chat.astream("I'm Pickle Rick"): tests/integration_tests/chat_models/test_bedrock.py:118: _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/language_models/chat_models.py:319: in astream raise e /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/language_models/chat_models.py:312: in astream generation += chunk /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/outputs/chat_generation.py:79: in __add__ message=self.message + other.message, /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/messages/ai.py:145: in __add__ response_metadata = merge_dicts( /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/utils/_merge.py:34: in merge_dicts merged[right_k] = merge_dicts(merged[right_k], right_v) _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ left = {'input_tokens': 13, 'output_tokens': 1}, right = {'output_tokens': 46} def merge_dicts(left: Dict[str, Any], right: Dict[str, Any]) -> Dict[str, Any]: """Merge two dicts, handling specific scenarios where a key exists in both dictionaries but has a value of None in 'left'. In such cases, the method uses the value from 'right' for that key in the merged dictionary. Example: If left = {"function_call": {"arguments": None}} and right = {"function_call": {"arguments": "{\n"}} then, after merging, for the key "function_call", the value from 'right' is used, resulting in merged = {"function_call": {"arguments": "{\n"}}. """ merged = left.copy() for right_k, right_v in right.items(): if right_k not in merged: merged[right_k] = right_v elif right_v is not None and merged[right_k] is None: merged[right_k] = right_v elif right_v is None: continue elif type(merged[right_k]) != type(right_v): raise TypeError( f'additional_kwargs["{right_k}"] already exists in this message,' " but with a different type." ) elif isinstance(merged[right_k], str): merged[right_k] += right_v elif isinstance(merged[right_k], dict): merged[right_k] = merge_dicts(merged[right_k], right_v) elif isinstance(merged[right_k], list): merged[right_k] = merge_lists(merged[right_k], right_v) elif merged[right_k] == right_v: continue else: > raise TypeError( f"Additional kwargs key {right_k} already exists in left dict and " f"value has unsupported type {type(merged[right_k])}." ) E TypeError: Additional kwargs key output_tokens already exists in left dict and value has unsupported type . /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/utils/_merge.py:40: TypeError ___________________________________________________________________________ test_function_call_invoke_with_system_astream ____________________________________________________________________________ chat = ChatBedrock(client=, region_name='us-west-2', model_id='anthropi...ing\nThe city and state\n\n\n\n") @pytest.mark.scheduled async def test_function_call_invoke_with_system_astream(chat: ChatBedrock) -> None: class GetWeather(BaseModel): location: str = Field(..., description="The city and state") llm_with_tools = chat.bind_tools([GetWeather]) messages = [ SystemMessage(content="anwser only in french"), HumanMessage(content="what is the weather like in San Francisco"), ] > for chunk in llm_with_tools.stream(messages): tests/integration_tests/chat_models/test_bedrock.py:207: _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/language_models/chat_models.py:249: in stream raise e /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/language_models/chat_models.py:240: in stream generation += chunk /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/outputs/chat_generation.py:79: in __add__ message=self.message + other.message, /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/messages/ai.py:145: in __add__ response_metadata = merge_dicts( /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/utils/_merge.py:34: in merge_dicts merged[right_k] = merge_dicts(merged[right_k], right_v) _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ left = {'input_tokens': 205, 'output_tokens': 1}, right = {'output_tokens': 115} def merge_dicts(left: Dict[str, Any], right: Dict[str, Any]) -> Dict[str, Any]: """Merge two dicts, handling specific scenarios where a key exists in both dictionaries but has a value of None in 'left'. In such cases, the method uses the value from 'right' for that key in the merged dictionary. Example: If left = {"function_call": {"arguments": None}} and right = {"function_call": {"arguments": "{\n"}} then, after merging, for the key "function_call", the value from 'right' is used, resulting in merged = {"function_call": {"arguments": "{\n"}}. """ merged = left.copy() for right_k, right_v in right.items(): if right_k not in merged: merged[right_k] = right_v elif right_v is not None and merged[right_k] is None: merged[right_k] = right_v elif right_v is None: continue elif type(merged[right_k]) != type(right_v): raise TypeError( f'additional_kwargs["{right_k}"] already exists in this message,' " but with a different type." ) elif isinstance(merged[right_k], str): merged[right_k] += right_v elif isinstance(merged[right_k], dict): merged[right_k] = merge_dicts(merged[right_k], right_v) elif isinstance(merged[right_k], list): merged[right_k] = merge_lists(merged[right_k], right_v) elif merged[right_k] == right_v: continue else: > raise TypeError( f"Additional kwargs key {right_k} already exists in left dict and " f"value has unsupported type {type(merged[right_k])}." ) E TypeError: Additional kwargs key output_tokens already exists in left dict and value has unsupported type . /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/utils/_merge.py:40: TypeError __________________________________________________________________________ test_function_call_invoke_without_system_astream __________________________________________________________________________ chat = ChatBedrock(client=, region_name='us-west-2', model_id='anthropi...ing\nThe city and state\n\n\n\n") @pytest.mark.scheduled async def test_function_call_invoke_without_system_astream(chat: ChatBedrock) -> None: class GetWeather(BaseModel): location: str = Field(..., description="The city and state") llm_with_tools = chat.bind_tools([GetWeather]) messages = [HumanMessage(content="what is the weather like in San Francisco")] > for chunk in llm_with_tools.stream(messages): tests/integration_tests/chat_models/test_bedrock.py:220: _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/language_models/chat_models.py:249: in stream raise e /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/language_models/chat_models.py:240: in stream generation += chunk /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/outputs/chat_generation.py:79: in __add__ message=self.message + other.message, /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/messages/ai.py:145: in __add__ response_metadata = merge_dicts( /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/utils/_merge.py:34: in merge_dicts merged[right_k] = merge_dicts(merged[right_k], right_v) _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ left = {'input_tokens': 197, 'output_tokens': 1}, right = {'output_tokens': 183} def merge_dicts(left: Dict[str, Any], right: Dict[str, Any]) -> Dict[str, Any]: """Merge two dicts, handling specific scenarios where a key exists in both dictionaries but has a value of None in 'left'. In such cases, the method uses the value from 'right' for that key in the merged dictionary. Example: If left = {"function_call": {"arguments": None}} and right = {"function_call": {"arguments": "{\n"}} then, after merging, for the key "function_call", the value from 'right' is used, resulting in merged = {"function_call": {"arguments": "{\n"}}. """ merged = left.copy() for right_k, right_v in right.items(): if right_k not in merged: merged[right_k] = right_v elif right_v is not None and merged[right_k] is None: merged[right_k] = right_v elif right_v is None: continue elif type(merged[right_k]) != type(right_v): raise TypeError( f'additional_kwargs["{right_k}"] already exists in this message,' " but with a different type." ) elif isinstance(merged[right_k], str): merged[right_k] += right_v elif isinstance(merged[right_k], dict): merged[right_k] = merge_dicts(merged[right_k], right_v) elif isinstance(merged[right_k], list): merged[right_k] = merge_lists(merged[right_k], right_v) elif merged[right_k] == right_v: continue else: > raise TypeError( f"Additional kwargs key {right_k} already exists in left dict and " f"value has unsupported type {type(merged[right_k])}." ) E TypeError: Additional kwargs key output_tokens already exists in left dict and value has unsupported type . /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/utils/_merge.py:40: TypeError ========================================================================================== warnings summary ========================================================================================== tests/integration_tests/chat_models/test_bedrock.py::test_chat_bedrock /Users/pijain/Library/Caches/pypoetry/virtualenvs/langchain-aws-eH7P7gjZ-py3.9/lib/python3.9/site-packages/langchain_core/_api/deprecation.py:119: LangChainDeprecationWarning: The method `BaseChatModel.__call__` was deprecated in langchain-core 0.1.7 and will be removed in 0.2.0. Use invoke instead. warn_deprecated( -- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html ---------- coverage: platform darwin, python 3.9.13-final-0 ---------- Name Stmts Miss Cover --------------------------------------------------------------- langchain_aws/__init__.py 6 0 100% langchain_aws/chat_models/__init__.py 2 0 100% langchain_aws/chat_models/bedrock.py 201 78 61% langchain_aws/embeddings/__init__.py 2 0 100% langchain_aws/embeddings/bedrock.py 79 53 33% langchain_aws/function_calling.py 39 6 85% langchain_aws/graphs/__init__.py 2 0 100% langchain_aws/graphs/neptune_graph.py 156 129 17% langchain_aws/graphs/neptune_rdf_graph.py 126 126 0% langchain_aws/llms/__init__.py 3 0 100% langchain_aws/llms/bedrock.py 387 185 52% langchain_aws/llms/sagemaker_endpoint.py 115 73 37% langchain_aws/retrievers/__init__.py 3 0 100% langchain_aws/retrievers/bedrock.py 47 25 47% langchain_aws/retrievers/kendra.py 183 87 52% langchain_aws/utils.py 18 12 33% --------------------------------------------------------------- TOTAL 1369 774 43% ======================================================================================== slowest 5 durations ========================================================================================= 5.80s call tests/integration_tests/chat_models/test_bedrock.py::test_function_call_invoke_without_system_astream 5.75s call tests/integration_tests/chat_models/test_bedrock.py::test_function_call_invoke_without_system 3.53s call tests/integration_tests/chat_models/test_bedrock.py::test_function_call_invoke_with_system 3.45s call tests/integration_tests/chat_models/test_bedrock.py::test_function_call_invoke_with_system_astream 1.90s call tests/integration_tests/chat_models/test_bedrock.py::test_bedrock_astream ====================================================================================== short test summary info ======================================================================================= FAILED tests/integration_tests/chat_models/test_bedrock.py::test_chat_bedrock_streaming_generation_info - TypeError: Additional kwargs key output_tokens already exists in left dict and value has unsupported type . FAILED tests/integration_tests/chat_models/test_bedrock.py::test_bedrock_streaming - TypeError: Additional kwargs key output_tokens already exists in left dict and value has unsupported type . FAILED tests/integration_tests/chat_models/test_bedrock.py::test_bedrock_astream - TypeError: Additional kwargs key output_tokens already exists in left dict and value has unsupported type . FAILED tests/integration_tests/chat_models/test_bedrock.py::test_function_call_invoke_with_system_astream - TypeError: Additional kwargs key output_tokens already exists in left dict and value has unsupported type . FAILED tests/integration_tests/chat_models/test_bedrock.py::test_function_call_invoke_without_system_astream - TypeError: Additional kwargs key output_tokens already exists in left dict and value has unsupported type .

[NAPTlME]

@3coins Thanks for the update. Hope you are feeling better.

It looks like merge_dicts doesn't allow for int types. I looked into updating this to merge integers by addition, but I see that would break some tests in which integers are used in a nominal fashion.

To get around this, I am now putting "usage" token counts into lists and summing them when combining.

Feel free to kick off the workflow. I believe everything is good. Let me know if there are any further areas you feel need to be modified.