From 9c0a9c0c2a2a728af024131ed261556d672fd6ad Mon Sep 17 00:00:00 2001
From: Eoin O'Shaughnessy <e.oshaughnessy@f5.com>
Date: Wed, 16 Jul 2025 12:14:36 +0100
Subject: [PATCH] feat: support tool calls passthrough

---
 src/f5_ai_gateway_sdk/request_input.py   | 46 +++++++++++++++++++++++-
 src/f5_ai_gateway_sdk/response_output.py |  4 ++-
 tests/test_request_input.py              | 40 +++++++++++++++++++++
 3 files changed, 88 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_request_input.py

diff --git a/src/f5_ai_gateway_sdk/request_input.py b/src/f5_ai_gateway_sdk/request_input.py
index 9b47bf6..83172ed 100644
--- a/src/f5_ai_gateway_sdk/request_input.py
+++ b/src/f5_ai_gateway_sdk/request_input.py
@@ -7,7 +7,15 @@
 
 from io import StringIO
 
-from pydantic import BaseModel
+from typing import Any, Self
+
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    field_serializer,
+    PrivateAttr,
+)
+from pydantic.functional_validators import model_validator
 
 from f5_ai_gateway_sdk.multipart_fields import INPUT_NAME
 from f5_ai_gateway_sdk.multipart_response import MultipartResponseField
@@ -40,9 +48,43 @@ class Message(BaseModel):
     """
 
     __autoclass_content__ = "class"
+    model_config = ConfigDict(extra="allow")
 
     content: str
     role: str = MessageRole.USER
+    _content_parsed_as_null: bool = PrivateAttr(default=False)
+
+    # messages may have null content when
+    # containing tool_calls
+    # this tracks that case in order to allow
+    # returning in the same format without the
+    # SDK user needing to handle None on content
+    @model_validator(mode="before")
+    @classmethod
+    def track_null_content(cls, data: Any) -> Any:
+        if isinstance(data, dict) and data.get("content") is None:
+            # Store this info in the data itself so it survives validation
+            data["__content_parsed_as_null__"] = True
+            data["content"] = ""
+        return data
+
+    @model_validator(mode="after")
+    def set_null_flag(self) -> Self:
+        # Check if the original data indicated null content
+        if hasattr(self, "__content_parsed_as_null__") or getattr(
+            self, "__content_parsed_as_null__", False
+        ):
+            self._content_parsed_as_null = True
+            # Remove the temporary tracking field now that we've set the private attribute
+            if hasattr(self, "__content_parsed_as_null__"):
+                delattr(self, "__content_parsed_as_null__")
+        return self
+
+    @field_serializer("content")
+    def serialize_content(self, content: str):
+        if self._content_parsed_as_null and len(content) == 0:
+            return None
+        return content
 
 
 class RequestInput(BaseModel):
@@ -68,6 +110,8 @@ class RequestInput(BaseModel):
     """
 
     __autoclass_content__ = "class"
+    model_config = ConfigDict(extra="allow")
+
     messages: list[Message]
 
     def to_multipart_field(self) -> MultipartResponseField:
diff --git a/src/f5_ai_gateway_sdk/response_output.py b/src/f5_ai_gateway_sdk/response_output.py
index 5afe871..ebe7962 100644
--- a/src/f5_ai_gateway_sdk/response_output.py
+++ b/src/f5_ai_gateway_sdk/response_output.py
@@ -5,7 +5,7 @@
 LICENSE file in the root directory of this source tree.
 """
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 from f5_ai_gateway_sdk.multipart_fields import RESPONSE_NAME
 from f5_ai_gateway_sdk.multipart_response import MultipartResponseField
@@ -21,6 +21,7 @@ class Choice(BaseModel):
     """
 
     __autoclass_content__ = "class"
+    model_config = ConfigDict(extra="allow")
 
     message: Message
 
@@ -47,6 +48,7 @@ class ResponseOutput(BaseModel):
     """
 
     __autoclass_content__ = "class"
+    model_config = ConfigDict(extra="allow")
 
     choices: list[Choice]
     """A list of ``Choice`` objects."""
diff --git a/tests/test_request_input.py b/tests/test_request_input.py
new file mode 100644
index 0000000..9e039c5
--- /dev/null
+++ b/tests/test_request_input.py
@@ -0,0 +1,40 @@
+"""
+Copyright (c) F5, Inc.
+
+This source code is licensed under the Apache License Version 2.0 found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+from f5_ai_gateway_sdk.request_input import RequestInput
+import pytest
+
+
+@pytest.mark.parametrize(
+    "content_value,should_have_null",
+    [
+        ("null", True),
+        ('"Hello world"', False),
+    ],
+)
+def test_maintains_content_and_excludes_tracking_fields(
+    content_value, should_have_null
+):
+    """
+    Test that messages with both null and non-null content are properly handled
+    during parsing and serialization, and that no internal tracking fields
+    are exposed in the serialized result.
+
+    This verifies that the SDK properly handles:
+    - Messages with null content (common in tool call scenarios)
+    - Messages with regular string content
+    - Additional fields (tool_calls) are persisted
+    - Internal implementation details remain hidden from serialized output
+    """
+    data = f'{{"messages":[{{"role":"user","content":{content_value},"tool_calls":[{{"id":"call_abc"}}]}}]}}'
+
+    parsed = RequestInput.model_validate_json(data)
+    serialized = parsed.model_dump_json()
+
+    assert ("null" in serialized) == should_have_null
+    assert "tool_calls" in serialized
+    assert "content_parsed_as_null" not in serialized