None type not checked before adding UsageMetadata value in AIMessageChunk when using LLM streaming

TheAppCrafter commented 6 days ago

Checked other resources

[X] I added a very descriptive title to this issue.
[X] I searched the LangChain documentation with the integrated search.
[X] I used the GitHub search to find a similar question and didn't find it.
[X] I am sure that this is a bug in LangChain rather than my code.
[X] The bug is not resolved by updating to the latest stable version of LangChain (or the specific integration package).

Example Code

client = ChatOpenAI(
    api_key=API_KEY, 
    base_url=PORTKEY_GATEWAY_URL, 
    streaming=streaming,
    default_headers=portkey_headers, 
    model=api_model_id,
    temperature=options.temperature,
    n=options.n,
    max_tokens=options.maxTokens,
)
messages = [HumanMessage(content='Some question')]
client.stream(messages)

Error Message and Stack Trace (if applicable)

|   File "/Users/user/app/.venv/lib/python3.12/site-packages/langchain_core/language_models/chat_models.py", line 411, in stream
|     raise e
|   File "/Users/user/app/.venv/lib/python3.12/site-packages/langchain_core/language_models/chat_models.py", line 402, in stream
|     generation += chunk
|   File "/Users/user/app/.venv/lib/python3.12/site-packages/langchain_core/outputs/chat_generation.py", line 100, in __add__
|     message=self.message + other.message,
|             ~~~~~~~~~~~~~^~~~~~~~~~~~~~~
|   File "/Users/user/app/.venv/lib/python3.12/site-packages/langchain_core/messages/ai.py", line 308, in __add__
|     return add_ai_message_chunks(self, other)
|            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|   File "/Users/user/app/.venv/lib/python3.12/site-packages/langchain_core/messages/ai.py", line 360, in add_ai_message_chunks
|     usage_metadata_["total_tokens"] += other.usage_metadata["total_tokens"]
| TypeError: unsupported operand type(s) for +=: 'int' and 'NoneType'
+------------------------------------

Description

I'm using the ChatOpenAI class to stream an LLM output or OpenAI compatible API endpoints. In my case when calling an anthropic model (and possibly others) an exception is thrown since other.usage_metadata["total_tokens"] is None.

    # Token usage
    if left.usage_metadata or any(o.usage_metadata is not None for o in others):
        usage_metadata_: UsageMetadata = left.usage_metadata or UsageMetadata(
            input_tokens=0, output_tokens=0, total_tokens=0
        )
        for other in others:
            if other.usage_metadata is not None:
                usage_metadata_["input_tokens"] += other.usage_metadata["input_tokens"]
                usage_metadata_["output_tokens"] += other.usage_metadata[
                    "output_tokens"
                ]
                usage_metadata_["total_tokens"] += other.usage_metadata["total_tokens"]
        usage_metadata: Optional[UsageMetadata] = usage_metadata_
    else:
        usage_metadata = None

I think we should check for None values before attempting to add to the existing UsageMetadata like so:

    # Token usage
    if left.usage_metadata or any(o.usage_metadata is not None for o in others):
        usage_metadata_: UsageMetadata = left.usage_metadata or UsageMetadata(
            input_tokens=0, output_tokens=0, total_tokens=0
        )
        for other in others:
            if other.usage_metadata is not None:
                if other.usage_metadata.get("input_tokens") is not None:
                    usage_metadata_["input_tokens"] += other.usage_metadata["input_tokens"]
                if other.usage_metadata.get("output_tokens") is not None:
                    usage_metadata_["output_tokens"] += other.usage_metadata["output_tokens"]
                if other.usage_metadata.get("total_tokens") is not None:
                    usage_metadata_["total_tokens"] += other.usage_metadata["total_tokens"]
        usage_metadata: Optional[UsageMetadata] = usage_metadata_
    else:
        usage_metadata = None

System Info

langchain-openai version: ^0.1.23 Platform: mac python version: 3.12.0

eyurtsev commented 6 days ago

Could you include full system information?

python -m langchain_core.sys_info

TheAppCrafter commented 6 days ago

Yes here is my system info:

System Information

OS: Darwin OS Version: Darwin Kernel Version 23.5.0: Wed May 1 20:12:58 PDT 2024; root:xnu-10063.121.3~5/RELEASE_ARM64_T6000 Python Version: 3.12.0 (main, Apr 21 2024, 02:41:04) [Clang 15.0.0 (clang-1500.3.9.4)]

keenborder786 commented 5 days ago

@eyurtsev I double checked it and was able to replicate the issue. Have created a small PR for fixing it.

langchain-ai / langchain