From 27ee61645dc89e2d80bc403e1352d3292623c992 Mon Sep 17 00:00:00 2001 From: Jon Nolen Date: Fri, 29 Dec 2023 16:41:47 -0500 Subject: [PATCH] core: Update messages/__init__.py to account for AIMessageChunk which breaks message history runnable. (#15327) - **Description:** fix parse issue for AIMessageChunk when using - **Issue:** https://github.com/langchain-ai/langchain/issues/14511 - **Dependencies:** none - **Twitter handle:** none Taken from this fix: https://github.com/gpt-engineer-org/gpt-engineer/issues/804#issuecomment-1769853850 Please make sure your PR is passing linting and testing before submitting. Run `make format`, `make lint` and `make test` from the root of the package you've modified to check this locally. See contribution guidelines for more information on how to write/run tests, lint, etc: https://python.langchain.com/docs/contributing/ If you're adding a new integration, please include: 1. a test for the integration, preferably unit tests that do not rely on network access, 2. an example notebook showing its use. It lives in `docs/docs/integrations` directory. If no one reviews your PR within a few days, please @-mention one of @baskaryan, @eyurtsev, @hwchase17. --------- Co-authored-by: Harrison Chase --- libs/core/langchain_core/messages/__init__.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/libs/core/langchain_core/messages/__init__.py b/libs/core/langchain_core/messages/__init__.py index cc8bd4d21a..44444c9d53 100644 --- a/libs/core/langchain_core/messages/__init__.py +++ b/libs/core/langchain_core/messages/__init__.py @@ -82,6 +82,16 @@ def _message_from_dict(message: dict) -> BaseMessage: return FunctionMessage(**message["data"]) elif _type == "tool": return ToolMessage(**message["data"]) + elif _type == "AIMessageChunk": + return AIMessageChunk(**message["data"]) + elif _type == "HumanMessageChunk": + return HumanMessageChunk(**message["data"]) + elif _type == "FunctionMessageChunk": + return FunctionMessageChunk(**message["data"]) + elif _type == "ToolMessageChunk": + return ToolMessageChunk(**message["data"]) + elif _type == "SystemMessageChunk": + return SystemMessageChunk(**message["data"]) else: raise ValueError(f"Got unexpected message type: {_type}")