From f82b8534f7102659a02bc6c2947d9a745e63da7a Mon Sep 17 00:00:00 2001
From: Eugene Yurtsev <eyurtsev@gmail.com>
Date: Fri, 30 May 2025 15:49:50 -0400
Subject: [PATCH] x

---
 .../langchain_openai/chat_models/base.py      | 14 ++++-
 .../chat_models/test_responses_api.py         | 59 +++++++++++++++++++
 2 files changed, 71 insertions(+), 2 deletions(-)

diff --git a/libs/partners/openai/langchain_openai/chat_models/base.py b/libs/partners/openai/langchain_openai/chat_models/base.py
index 50cb985a572..8dd83796f05 100644
--- a/libs/partners/openai/langchain_openai/chat_models/base.py
+++ b/libs/partners/openai/langchain_openai/chat_models/base.py
@@ -3057,7 +3057,8 @@ def _construct_responses_api_payload(
             if tool["type"] == "image_generation" and "partial_images" in tool:
                 raise NotImplementedError(
                     "Partial image generation is not yet supported "
-                    "via the LangChain ChatOpenAI client. "
+                    "via the LangChain ChatOpenAI client. Please "
+                    "drop the 'partial_images' key from the image_generation tool."
                 )
         payload["tools"] = new_tools
     if tool_choice := payload.pop("tool_choice", None):
@@ -3216,7 +3217,16 @@ def _construct_responses_api_input(messages: Sequence[BaseMessage]) -> list:
                     pass
             input_.extend(code_interpreter_calls)
             input_.extend(mcp_calls)
-            input_.extend(image_generation_calls)
+
+            # A previous image generation call can be referenced by ID
+
+            input_.extend(
+                [
+                    {"type": "image_generation", "id": image_generation_call["id"]}
+                    for image_generation_call in image_generation_calls
+                ]
+            )
+
             msg["content"] = msg.get("content") or []
             if lc_msg.additional_kwargs.get("refusal"):
                 if isinstance(msg["content"], str):
diff --git a/libs/partners/openai/tests/integration_tests/chat_models/test_responses_api.py b/libs/partners/openai/tests/integration_tests/chat_models/test_responses_api.py
index 165a756761c..a38dcca5759 100644
--- a/libs/partners/openai/tests/integration_tests/chat_models/test_responses_api.py
+++ b/libs/partners/openai/tests/integration_tests/chat_models/test_responses_api.py
@@ -511,3 +511,62 @@ def test_image_generation_streaming() -> None:
     # _check_response(complete_ai_message)
     tool_output = complete_ai_message.additional_kwargs["tool_outputs"][0]
     assert set(tool_output.keys()).issubset(expected_keys)
+
+
+def test_image_generation_multi_turn() -> None:
+    """Test multi-turn editing of image generation by passing in history."""
+    # Test multi-turn
+    llm = ChatOpenAI(model="gpt-4.1", use_responses_api=True)
+    # Test invocation
+    tool = {
+        "type": "image_generation",
+        # For testing purposes let's keep the quality low, so the test runs faster.
+        "quality": "low",
+    }
+    llm_with_tools = llm.bind_tools([tool])
+
+    chat_history = [{"role": "user", "content": "Make a picture of a fuzzy cat"}]
+    ai_message = llm_with_tools.invoke(chat_history)
+    _check_response(ai_message)
+    tool_output = ai_message.additional_kwargs["tool_outputs"][0]
+
+    # Example tool output for an image
+    # {
+    #     "background": "opaque",
+    #     "id": "ig_683716a8ddf0819888572b20621c7ae4029ec8c11f8dacf8",
+    #     "output_format": "png",
+    #     "quality": "high",
+    #     "revised_prompt": "A fluffy, fuzzy cat sitting calmly, with soft fur, bright "
+    #     "eyes, and a cute, friendly expression. The background is "
+    #     "simple and light to emphasize the cat's texture and "
+    #     "fluffiness.",
+    #     "size": "1024x1024",
+    #     "status": "completed",
+    #     "type": "image_generation_call",
+    #     "result": # base64 encode image data
+    # }
+
+    expected_keys = {
+        "id",
+        "background",
+        "output_format",
+        "quality",
+        "result",
+        "revised_prompt",
+        "size",
+        "status",
+        "type",
+    }
+
+    assert set(tool_output.keys()).issubset(expected_keys)
+
+    chat_history.extend(
+        ai_message,  # AI message with tool output
+        # New request
+        {"role": "user", "content": "Now make it realistic."},
+    )
+
+    ai_message2 = llm_with_tools.invoke(chat_history)
+    _check_response(ai_message2)
+    tool_output2 = ai_message2.additional_kwargs["tool_outputs"][0]
+    assert set(tool_output2.keys()).issubset(expected_keys)