fix: convert Usage object to dict for AIMessage (#19)

* fix: convert Usage object to dict for AIMessage * chore: bump version to 0.3.2
ai-forever · Dec 16, 2024 · 270014e · 270014e
1 parent f32a55a
commit 270014e
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 2 deletions.
diff --git 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
@@ -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
@@ -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()