From 47ded80b64fa8bd5d3d2f8cab0fe17fd66689019 Mon Sep 17 00:00:00 2001
From: ccurme <chester.curme@gmail.com>
Date: Tue, 15 Apr 2025 15:22:58 -0400
Subject: [PATCH] ollama[patch]: fix generation info (#30863)

https://github.com/langchain-ai/langchain/pull/30778 (not released)
broke all invocation modes of ChatOllama (intent was to remove
`"message"` from `generation_info`, but we turned `generation_info` into
`stream_resp["message"]`), resulting in validation errors.
---
 .../ollama/langchain_ollama/chat_models.py    | 22 ++++++++++---------
 1 file changed, 12 insertions(+), 10 deletions(-)

diff --git a/libs/partners/ollama/langchain_ollama/chat_models.py b/libs/partners/ollama/langchain_ollama/chat_models.py
index a0b47bb5b33..879a3cccf49 100644
--- a/libs/partners/ollama/langchain_ollama/chat_models.py
+++ b/libs/partners/ollama/langchain_ollama/chat_models.py
@@ -719,6 +719,11 @@ class ChatOllama(BaseChatModel):
         is_thinking = False
         for stream_resp in self._create_chat_stream(messages, stop, **kwargs):
             if not isinstance(stream_resp, str):
+                if stream_resp.get("done") is True:
+                    generation_info = dict(stream_resp)
+                    _ = generation_info.pop("message", None)
+                else:
+                    generation_info = None
                 chunk = ChatGenerationChunk(
                     message=AIMessageChunk(
                         content=(
@@ -732,11 +737,7 @@ class ChatOllama(BaseChatModel):
                         ),
                         tool_calls=_get_tool_calls_from_response(stream_resp),
                     ),
-                    generation_info=(
-                        dict(stream_resp).pop("message", None)
-                        if stream_resp.get("done") is True
-                        else None
-                    ),
+                    generation_info=generation_info,
                 )
                 if chunk.generation_info and (
                     model := chunk.generation_info.get("model")
@@ -773,6 +774,11 @@ class ChatOllama(BaseChatModel):
         is_thinking = False
         async for stream_resp in self._acreate_chat_stream(messages, stop, **kwargs):
             if not isinstance(stream_resp, str):
+                if stream_resp.get("done") is True:
+                    generation_info = dict(stream_resp)
+                    _ = generation_info.pop("message", None)
+                else:
+                    generation_info = None
                 chunk = ChatGenerationChunk(
                     message=AIMessageChunk(
                         content=(
@@ -786,11 +792,7 @@ class ChatOllama(BaseChatModel):
                         ),
                         tool_calls=_get_tool_calls_from_response(stream_resp),
                     ),
-                    generation_info=(
-                        dict(stream_resp).pop("message", None)
-                        if stream_resp.get("done") is True
-                        else None
-                    ),
+                    generation_info=generation_info,
                 )
                 if chunk.generation_info and (
                     model := chunk.generation_info.get("model")