From 270014ed07b5114b7f931c6d188cde7955d7d6c6 Mon Sep 17 00:00:00 2001 From: Dmitry Labazkin Date: Mon, 16 Dec 2024 12:24:51 +0300 Subject: [PATCH] fix: convert Usage object to dict for AIMessage (#19) * fix: convert Usage object to dict for AIMessage * chore: bump version to 0.3.2 --- libs/gigachat/langchain_gigachat/chat_models/gigachat.py | 5 ++++- libs/gigachat/pyproject.toml | 2 +- libs/gigachat/tests/unit_tests/test_gigachat.py | 6 ++++++ 3 files changed, 11 insertions(+), 2 deletions(-) diff --git a/libs/gigachat/langchain_gigachat/chat_models/gigachat.py b/libs/gigachat/langchain_gigachat/chat_models/gigachat.py index 4855898..444ff58 100644 --- a/libs/gigachat/langchain_gigachat/chat_models/gigachat.py +++ b/libs/gigachat/langchain_gigachat/chat_models/gigachat.py @@ -407,7 +407,10 @@ def _create_chat_result(self, response: Any) -> ChatResult: generations.append(gen) if self.verbose: logger.warning("Giga response: %s", message.content) - llm_output = {"token_usage": response.usage, "model_name": response.model} + llm_output = { + "token_usage": response.usage.dict(), + "model_name": response.model, + } return ChatResult(generations=generations, llm_output=llm_output) def _generate( diff --git a/libs/gigachat/pyproject.toml b/libs/gigachat/pyproject.toml index c97ebd9..dd11503 100644 --- a/libs/gigachat/pyproject.toml +++ b/libs/gigachat/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langchain-gigachat" -version = "0.3.1" +version = "0.3.2" description = "An integration package connecting GigaChat and LangChain" authors = [] readme = "README.md" diff --git a/libs/gigachat/tests/unit_tests/test_gigachat.py b/libs/gigachat/tests/unit_tests/test_gigachat.py index 204eab4..3832672 100644 --- a/libs/gigachat/tests/unit_tests/test_gigachat.py +++ b/libs/gigachat/tests/unit_tests/test_gigachat.py @@ -375,3 +375,9 @@ def test_structured_output() -> None: } ], } + + +def test_ai_message_json_serialization(patch_gigachat: None) -> None: + llm = GigaChat() + response = llm.invoke("hello") + response.model_dump_json()