Repository: JoshuaC215/agent-service-toolkit
Pull request #216: Reassemble partial `AIMessage` fragments from streamed tuples
Author: login: alex-feel
Created: 2025-04-29T08:37:40Z | Merged: 2025-04-30T04:19:22Z

I've observed that the model's final response sometimes arrives in multiple fragments instead of a single complete message. This behavior complicates downstream processing, as the system expects a fully assembled message.

Here’s an example of the issue observed in the logs:

<details>

<summary>Log</summary>

```commandLine
2025-04-29 10:54:30 ====== STREAM EVENT ======
2025-04-29 10:54:30 type: <class 'tuple'>
2025-04-29 10:54:30 repr: ('content', "HERE'S THE LONG CONTENT")
2025-04-29 10:54:30 pprint:
2025-04-29 10:54:30 ('content',
2025-04-29 10:54:30  'HERE'S THE LONG CONTENT'
2025-04-29 10:54:30  '\n'
2025-04-29 10:54:30  '...')
2025-04-29 10:54:30 ==========================
2025-04-29 10:54:30 [2025-04-29 07:54:30 +0000] [1] [ERROR] [/app/service/service.py:238] Error parsing message: Unsupported message type: tuple
2025-04-29 10:54:30 [2025-04-29 07:54:30 +0000] [1] [INFO] [/app/service/service.py:224] 
2025-04-29 10:54:30 ====== STREAM EVENT ======
2025-04-29 10:54:30 type: <class 'tuple'>
2025-04-29 10:54:30 repr: ('additional_kwargs', {})
2025-04-29 10:54:30 pprint:
2025-04-29 10:54:30 ('additional_kwargs', {})
2025-04-29 10:54:30 ==========================
2025-04-29 10:54:30 [2025-04-29 07:54:30 +0000] [1] [ERROR] [/app/service/service.py:238] Error parsing message: Unsupported message type: tuple
2025-04-29 10:54:30 [2025-04-29 07:54:30 +0000] [1] [INFO] [/app/service/service.py:224] 
2025-04-29 10:54:30 ====== STREAM EVENT ======
2025-04-29 10:54:30 type: <class 'tuple'>
2025-04-29 10:54:30 repr: ('response_metadata', {'finish_reason': 'stop', 'model_name': 'gpt-4.1-2025-04-14', 'system_fingerprint': 'fp_a1102cf978'})
2025-04-29 10:54:30 pprint:
2025-04-29 10:54:30 ('response_metadata',
2025-04-29 10:54:30  {'finish_reason': 'stop', 'model_name': 'gpt-4.1-2025-04-14', 'system_fingerprint': 'fp_a1102cf978'})
2025-04-29 10:54:30 ==========================
2025-04-29 10:54:30 [2025-04-29 07:54:30 +0000] [1] [ERROR] [/app/service/service.py:238] Error parsing message: Unsupported message type: tuple
2025-04-29 10:54:30 [2025-04-29 07:54:30 +0000] [1] [INFO] [/app/service/service.py:224] 
2025-04-29 10:54:30 ====== STREAM EVENT ======
2025-04-29 10:54:30 type: <class 'tuple'>
2025-04-29 10:54:30 repr: ('type', 'ai')
2025-04-29 10:54:30 pprint:
2025-04-29 10:54:30 ('type', 'ai')
2025-04-29 10:54:30 ==========================
2025-04-29 10:54:30 [2025-04-29 07:54:30 +0000] [1] [ERROR] [/app/service/service.py:238] Error parsing message: Unsupported message type: tuple
2025-04-29 10:54:30 [2025-04-29 07:54:30 +0000] [1] [INFO] [/app/service/service.py:224] 
2025-04-29 10:54:30 ====== STREAM EVENT ======
2025-04-29 10:54:30 type: <class 'tuple'>
2025-04-29 10:54:30 repr: ('name', None)
2025-04-29 10:54:30 pprint:
2025-04-29 10:54:30 ('name', None)
2025-04-29 10:54:30 ==========================
2025-04-29 10:54:30 [2025-04-29 07:54:30 +0000] [1] [ERROR] [/app/service/service.py:238] Error parsing message: Unsupported message type: tuple
2025-04-29 10:54:30 [2025-04-29 07:54:30 +0000] [1] [INFO] [/app/service/service.py:224] 
2025-04-29 10:54:30 ====== STREAM EVENT ======
2025-04-29 10:54:30 type: <class 'tuple'>
2025-04-29 10:54:30 repr: ('id', 'run-79a8baa2-891f-4621-bc7d-08749d308260')
2025-04-29 10:54:30 pprint:
2025-04-29 10:54:30 ('id', 'run-79a8baa2-891f-4621-bc7d-08749d308260')
2025-04-29 10:54:30 ==========================
2025-04-29 10:54:30 [2025-04-29 07:54:30 +0000] [1] [ERROR] [/app/service/service.py:238] Error parsing message: Unsupported message type: tuple
2025-04-29 10:54:30 [2025-04-29 07:54:30 +0000] [1] [INFO] [/app/service/service.py:224] 
2025-04-29 10:54:30 ====== STREAM EVENT ======
2025-04-29 10:54:30 type: <class 'tuple'>
2025-04-29 10:54:30 repr: ('example', False)
2025-04-29 10:54:30 pprint:
2025-04-29 10:54:30 ('example', False)
2025-04-29 10:54:30 ==========================
2025-04-29 10:54:30 [2025-04-29 07:54:30 +0000] [1] [ERROR] [/app/service/service.py:238] Error parsing message: Unsupported message type: tuple
2025-04-29 10:54:30 [2025-04-29 07:54:30 +0000] [1] [INFO] [/app/service/service.py:224] 
2025-04-29 10:54:30 ====== STREAM EVENT ======
2025-04-29 10:54:30 type: <class 'tuple'>
2025-04-29 10:54:30 repr: ('tool_calls', [])
2025-04-29 10:54:30 pprint:
2025-04-29 10:54:30 ('tool_calls', [])
2025-04-29 10:54:30 ==========================
2025-04-29 10:54:30 [2025-04-29 07:54:30 +0000] [1] [ERROR] [/app/service/service.py:238] Error parsing message: Unsupported message type: tuple
2025-04-29 10:54:30 [2025-04-29 07:54:30 +0000] [1] [INFO] [/app/service/service.py:224] 
2025-04-29 10:54:30 ====== STREAM EVENT ======
2025-04-29 10:54:30 type: <class 'tuple'>
2025-04-29 10:54:30 repr: ('invalid_tool_calls', [])
2025-04-29 10:54:30 pprint:
2025-04-29 10:54:30 ('invalid_tool_calls', [])
2025-04-29 10:54:30 ==========================
2025-04-29 10:54:30 [2025-04-29 07:54:30 +0000] [1] [ERROR] [/app/service/service.py:238] Error parsing message: Unsupported message type: tuple
2025-04-29 10:54:30 [2025-04-29 07:54:30 +0000] [1] [INFO] [/app/service/service.py:224] 
2025-04-29 10:54:30 ====== STREAM EVENT ======
2025-04-29 10:54:30 type: <class 'tuple'>
2025-04-29 10:54:30 repr: ('usage_metadata', None)
2025-04-29 10:54:30 pprint:
2025-04-29 10:54:30 ('usage_metadata', None)
2025-04-29 10:54:30 ==========================
2025-04-29 10:54:30 [2025-04-29 07:54:30 +0000] [1] [ERROR] [/app/service/service.py:238] Error parsing message: Unsupported message type: tuple
2025-04-29 10:54:30 [2025-04-29 07:54:30 +0000] [1] [INFO] [/app/service/service.py:224] 
2025-04-29 10:54:30 ====== STREAM EVENT ======
```

</details>

To address this, this PR adds logic to correctly reassemble incoming fragmented tuples into a proper `AIMessage`. If the tuple represents a `content` or `tool_calls` field, we now construct an `AIMessage` accordingly; otherwise, metadata tuples are skipped.


Diff:
diff --git a/src/service/service.py b/src/service/service.py
index 24ab976a..733d5262 100644
--- a/src/service/service.py
+++ b/src/service/service.py
@@ -1,3 +1,4 @@
+import inspect
 import json
 import logging
 import warnings
@@ -221,7 +222,29 @@ async def message_generator(
             if stream_mode == "custom":
                 new_messages = [event]
 
+            # LangGraph streaming may emit tuples: (field_name, field_value)
+            # e.g. ('content', <str>), ('tool_calls', [ToolCall,...]), ('additional_kwargs', {...}), etc.
+            # We accumulate only supported fields into `parts` and skip unsupported metadata.
+            # More info at: https://langchain-ai.github.io/langgraph/cloud/how-tos/stream_messages/
+            processed_messages = []
+            current_message: dict[str, Any] = {}
             for message in new_messages:
+                if isinstance(message, tuple):
+                    key, value = message
+                    # Store parts in temporary dict
+                    current_message[key] = value
+                else:
+                    # Add complete message if we have one in progress
+                    if current_message:
+                        processed_messages.append(_create_ai_message(current_message))
+                        current_message = {}
+                    processed_messages.append(message)
+
+            # Add any remaining message parts
+            if current_message:
+                processed_messages.append(_create_ai_message(current_message))
+
+            for message in processed_messages:
                 try:
                     chat_message = langchain_to_chat_message(message)
                     chat_message.run_id = str(run_id)
@@ -257,6 +280,13 @@ async def message_generator(
         yield "data: [DONE]\n\n"
 
 
+def _create_ai_message(parts: dict) -> AIMessage:
+    sig = inspect.signature(AIMessage)
+    valid_keys = set(sig.parameters)
+    filtered = {k: v for k, v in parts.items() if k in valid_keys}
+    return AIMessage(**filtered)
+
+
 def _sse_response_example() -> dict[int | str, Any]:
     return {
         status.HTTP_200_OK: {
diff --git a/tests/service/test_service_streaming.py b/tests/service/test_service_streaming.py
new file mode 100644
index 00000000..d478fea9
--- /dev/null
+++ b/tests/service/test_service_streaming.py
@@ -0,0 +1,58 @@
+import pytest
+from langchain_core.messages import AIMessage
+
+from service.service import _create_ai_message
+
+
+@pytest.mark.parametrize(
+    "parts, expected",
+    [
+        # 1) Basic content + tool_calls
+        (
+            {"content": "Hello", "tool_calls": []},
+            {"content": "Hello", "tool_calls": []},
+        ),
+        # 2) Unknown keys are ignored
+        (
+            {"content": "Test", "foobar": 123, "tool_calls": []},
+            {"content": "Test", "tool_calls": []},
+        ),
+        # 3) Extra valid AIMessage params (id, type) pass through
+        (
+            {
+                "content": "Hey",
+                "id": "abc-123",
+                "type": "ai",
+                "tool_calls": [],
+            },
+            {"content": "Hey", "id": "abc-123", "type": "ai", "tool_calls": []},
+        ),
+    ],
+)
+def test_create_ai_message_filters_and_passes_through(parts, expected):
+    """
+    _create_ai_message should:
+      - Drop unknown keys ("foobar")
+      - Preserve keys that match AIMessage signature
+      - Use the final value for duplicate keys in the parts dict
+    """
+    msg: AIMessage = _create_ai_message(parts)
+    for key, val in expected.items():
+        assert getattr(msg, key) == val
+
+
+def test_create_ai_message_missing_required_content_raises():
+    """
+    AIMessage requires 'content'; if missing, _create_ai_message should
+    bubble up the TypeError from the constructor.
+    """
+    with pytest.raises(TypeError):
+        _create_ai_message({"tool_calls": []})
+
+
+def test_create_ai_message_empty_dict_raises():
+    """
+    Completely empty parts should also fail to construct an AIMessage.
+    """
+    with pytest.raises(TypeError):
+        _create_ai_message({})


Reviews:
user: login: JoshuaC215
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-29T13:30:45Z
user: login: alex-feel
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T19:42:23Z
user: login: JoshuaC215
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-30T04:15:18Z
user: login: JoshuaC215
state: APPROVED
author_association: OWNER
submitted_at: 2025-04-30T04:19:01Z

Review comments:
diff_hunk: @@ -257,6 +276,21 @@ async def message_generator(
         yield "data: [DONE]\n\n"
 
 
+def _create_ai_message(parts: dict) -> AIMessage:
path: src/service/service.py
user: login: JoshuaC215
body: I’m slightly nervous that this will break in the future if new attributes are added or defaults change. I wonder if there’s another way to do this that’s more automatic? Could the dict be passed as **kwargs ?
created_at: 2025-04-29T13:30:45Z
author_association: OWNER
diff_hunk: @@ -257,6 +276,21 @@ async def message_generator(
         yield "data: [DONE]\n\n"
 
 
+def _create_ai_message(parts: dict) -> AIMessage:
path: src/service/service.py
user: login: alex-feel
body: Could you please check the current option?
created_at: 2025-04-29T19:42:23Z
author_association: CONTRIBUTOR
diff_hunk: @@ -257,6 +276,21 @@ async def message_generator(
         yield "data: [DONE]\n\n"
 
 
+def _create_ai_message(parts: dict) -> AIMessage:
path: src/service/service.py
user: login: JoshuaC215
body: Nice!
created_at: 2025-04-30T04:15:18Z
author_association: OWNER

PR comments:
user: login: JoshuaC215
created_at: 2025-04-29T13:29:08Z
body: Hey, thanks for tracking this down and providing a reasonable fix. Which model / LLM provider were you seeing this behavior for? Is it a documented langchain behavior? 🤔

May want to document in the in-line comments so it’s more obvious.

Also would be nice to include a test for this, do you feel comfortable adding that? I can do it later if not
author_association: OWNER
user: login: alex-feel
created_at: 2025-04-29T20:00:14Z
body: @JoshuaC215 I believe this behavior is described here: https://langchain-ai.github.io/langgraph/cloud/how-tos/stream_messages/.  

I've just pushed the latest changes — could you please take a look and let me know if everything aligns with your expectations?
author_association: CONTRIBUTOR
user: login: codecov[bot]
created_at: 2025-04-30T04:17:29Z
body: ## [Codecov](https://app.codecov.io/gh/JoshuaC215/agent-service-toolkit/pull/216?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joshua+Carroll) Report
Attention: Patch coverage is `72.22222%` with `5 lines` in your changes missing coverage. Please review.
> Project coverage is 67.54%. Comparing base [(`cbc51aa`)](https://app.codecov.io/gh/JoshuaC215/agent-service-toolkit/commit/cbc51aa89ab356538801dd3083821c22dff47213?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joshua+Carroll) to head [(`ab4bb60`)](https://app.codecov.io/gh/JoshuaC215/agent-service-toolkit/commit/ab4bb604b47fbfbbdbd431ebbd1f80d0d3057025?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joshua+Carroll).
> Report is 2 commits behind head on main.

| [Files with missing lines](https://app.codecov.io/gh/JoshuaC215/agent-service-toolkit/pull/216?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joshua+Carroll) | Patch % | Lines |
|---|---|---|
| [src/service/service.py](https://app.codecov.io/gh/JoshuaC215/agent-service-toolkit/pull/216?src=pr&el=tree&filepath=src%2Fservice%2Fservice.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joshua+Carroll#diff-c3JjL3NlcnZpY2Uvc2VydmljZS5weQ==) | 72.22% | [5 Missing :warning: ](https://app.codecov.io/gh/JoshuaC215/agent-service-toolkit/pull/216?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joshua+Carroll) |

| [Files with missing lines](https://app.codecov.io/gh/JoshuaC215/agent-service-toolkit/pull/216?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joshua+Carroll) | Coverage Δ | |
|---|---|---|
| [src/service/service.py](https://app.codecov.io/gh/JoshuaC215/agent-service-toolkit/pull/216?src=pr&el=tree&filepath=src%2Fservice%2Fservice.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joshua+Carroll#diff-c3JjL3NlcnZpY2Uvc2VydmljZS5weQ==) | `81.76% <72.22%> (-1.06%)` | :arrow_down: |

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE