Bedrock token count callbacks #20

NAPTlME · 2024-04-22T07:11:23Z

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

on_llm_end, invoke will provide the usage and stop_reason information

…on_chunk Also allowed additional response_body information to be passed into the generation_info (such as token counts and stop reasons)

The on_llm_end call will include all GenerationChunks as well as info on usage and stop_reason (in the llm_output) Added a provider stop code mapping

…sages api and outputs

…from the correct dict

…s api Need the `message_delta` to return in order to get the stop reason and output token count https://docs.anthropic.com/claude/reference/messages-streaming

…ncluding the stop_reason rather than just usage_info Removing the on_llm_end that was added and moving to _call to match with the pattern used for the ChatModel (as this was causing on_llm_end to be called twice)

…d use

…nfo rather than the full chunks to make it compatible with some calls needed from ChatBedrock Updating the usage info to use prompt_tokens and completion_tokens to match those returned from the non streaming route

…being cleared from the llm_output This was returning llm_outputs with no usage information

…reaming and non-streaming generation and passing into the ChatResult

3coins · 2024-04-22T19:00:33Z

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

NAPTlME · 2024-04-22T21:29:35Z

@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.

These appear to be holdovers from the langchain repo's structure and default branch

NAPTlME · 2024-04-23T21:41:08Z

@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 · 2024-04-30T20:43:49Z

@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 · 2024-05-03T18:23:46Z

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

NAPTlME · 2024-05-10T15:23:55Z

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

DanielWhite95 · 2024-05-16T07:34:33Z

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 · 2024-05-16T16:15:37Z

@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 · 2024-05-16T16:18:18Z

@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 · 2024-05-16T20:43:42Z

@3coins is your man!

3coins · 2024-05-16T20:48:59Z

@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 · 2024-05-17T18:26:08Z

@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 <class 'int'>.

/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=<botocore.client.BedrockRuntime object at 0x126949430>, 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 <class 'int'>.

/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=<botocore.client.BedrockRuntime object at 0x1321a7160>, 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 <class 'int'>.

/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=<botocore.client.BedrockRuntime object at 0x126c30c40>, region_name='us-west-2', model_id='anthropi...ing\nThe city and state\n\n\n</tool_description>\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 <class 'int'>.

/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=<botocore.client.BedrockRuntime object at 0x126f0d940>, region_name='us-west-2', model_id='anthropi...ing\nThe city and state\n\n\n</tool_description>\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 <class 'int'>.

/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 <class 'int'>.
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 <class 'int'>.
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 <class 'int'>.
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 <class 'int'>.
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 <class 'int'>.

…ptured on_llm_end.

…hese integer values to prevent errors when merge_dicts is called from langchain_core.utils.

…PTlME/langchain-aws into bedrock-token-count-callbacks

… lists to the message_start message type.

NAPTlME · 2024-05-17T21:25:03Z

@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.

…ion to reduce the code duplication.

…nfo` test would occasionally fail (returned "Hello! How are you doing?" one time)

3coins

@NAPTlME
Looks good! Thanks for your patience.

NAPTlME added 14 commits April 20, 2024 05:24

Updating _prepare_input_and_invoke to use the on_llm_end callback

7a22060

on_llm_end, invoke will provide the usage and stop_reason information

Consolidated GenerationChunk creation to _stream_response_to_generati…

929e3db

…on_chunk Also allowed additional response_body information to be passed into the generation_info (such as token counts and stop reasons)

Added an on_llm_end call to the streaming route

9b840d1

The on_llm_end call will include all GenerationChunks as well as info on usage and stop_reason (in the llm_output) Added a provider stop code mapping

Updating async streaming options to match streaming functions wrt mes…

5802e53

…sages api and outputs

Updating _stream_response_to_generation_chunk to get the stop reason …

140333b

…from the correct dict

Ending on "message_stop" rather that "content_block_stop" for message…

2f2c8ed

…s api Need the `message_delta` to return in order to get the stop reason and output token count https://docs.anthropic.com/claude/reference/messages-streaming

Renaming to provider_stop_reason_key_map to better convey the intende…

a91b3c7

…d use

changing to output.get() rather than pop() to prevent the usage from …

25e9caf

…being cleared from the llm_output This was returning llm_outputs with no usage information

Passing generation_info through the ChatGenerationChunk from _stream

8006cb5

Getting llm_output (containing usage and stop_reason from both st…

1eeb60f

…reaming and non-streaming generation and passing into the ChatResult

Linting/Formatting/codespell changes

afda911

Updating a function description

02893cc

NAPTlME added 2 commits April 22, 2024 15:21

Line Length for Ruff

794f9e5

Making compatible with python 3.8

b4c5f61

NAPTlME added 3 commits April 22, 2024 19:29

Ruff Formatting

6a870de

Fixing issue with path and branch specification in lint_diff/format_diff

932204a

These appear to be holdovers from the langchain repo's structure and default branch

Addressing linting/formatting issues

4750b80

NAPTlME added 2 commits April 25, 2024 02:52

Typing compatibility for python < 3.10

8a2465c

Merge branch 'langchain-ai:main' into bedrock-token-count-callbacks

64b3578

Merge branch 'main' into bedrock-token-count-callbacks

1526f5a

NAPTlME mentioned this pull request May 8, 2024

Token Counts not used from Bedrock response. #41

Closed

NAPTlME added 7 commits May 17, 2024 14:22

Adding a callback to test token counts on_llm_end

7a7e39f

Adding integration test to verify token counts and stop reason are ca…

98d6d09

…ptured on_llm_end.

Fixing codespell errors

0f681cb

Merge branch 'langchain-ai:main' into bedrock-token-count-callbacks

f1350c4

Updating usage info for input/output tokens to use lists to contain t…

8908cd8

…hese integer values to prevent errors when merge_dicts is called from langchain_core.utils.

Merge branch 'bedrock-token-count-callbacks' of https://github.com/NA…

aa33982

…PTlME/langchain-aws into bedrock-token-count-callbacks

Adding the usage_info transformation in which the integers are put in…

f5e4431

… lists to the message_start message type.

NAPTlME added 2 commits May 17, 2024 16:36

Moving the usage_info nesting of token counts into lists into a funct…

5eace06

…ion to reduce the code duplication.

Ran into an issue where the `test_chat_bedrock_streaming_generation_i…

bfa0871

…nfo` test would occasionally fail (returned "Hello! How are you doing?" one time)

3coins approved these changes May 21, 2024

View reviewed changes

3coins merged commit 4a88b7f into langchain-ai:main May 21, 2024
12 checks passed

3coins mentioned this pull request May 23, 2024

BedrockAnthropicTokenUsageCallbackHandler does not function with langchain_aws #54

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Bedrock token count callbacks #20

Bedrock token count callbacks #20

NAPTlME commented Apr 22, 2024 •

edited

Loading

3coins commented Apr 22, 2024

NAPTlME commented Apr 22, 2024 •

edited

Loading

NAPTlME commented Apr 23, 2024

NAPTlME commented Apr 30, 2024

NAPTlME commented May 3, 2024

NAPTlME commented May 10, 2024

DanielWhite95 commented May 16, 2024

NAPTlME commented May 16, 2024 •

edited

Loading

NAPTlME commented May 16, 2024

efriis commented May 16, 2024

3coins commented May 16, 2024

3coins commented May 17, 2024

---------- coverage: platform darwin, python 3.9.13-final-0 ----------
Name Stmts Miss Cover

NAPTlME commented May 17, 2024

3coins left a comment

Bedrock token count callbacks #20

Bedrock token count callbacks #20

Conversation

NAPTlME commented Apr 22, 2024 • edited Loading

3coins commented Apr 22, 2024

NAPTlME commented Apr 22, 2024 • edited Loading

NAPTlME commented Apr 23, 2024

NAPTlME commented Apr 30, 2024

NAPTlME commented May 3, 2024

NAPTlME commented May 10, 2024

DanielWhite95 commented May 16, 2024

NAPTlME commented May 16, 2024 • edited Loading

NAPTlME commented May 16, 2024

efriis commented May 16, 2024

3coins commented May 16, 2024

3coins commented May 17, 2024

---------- coverage: platform darwin, python 3.9.13-final-0 ---------- Name Stmts Miss Cover

NAPTlME commented May 17, 2024

3coins left a comment

Choose a reason for hiding this comment

NAPTlME commented Apr 22, 2024 •

edited

Loading

NAPTlME commented Apr 22, 2024 •

edited

Loading

NAPTlME commented May 16, 2024 •

edited

Loading

---------- coverage: platform darwin, python 3.9.13-final-0 ----------
Name Stmts Miss Cover