fix: serialize Pydantic models (#13)

Move transform code into its own module, since it'll be needed for other
services eventually, and ensure it handles Pydantic models like
`ExtraInfo`.

Author: kenany

flake.lock

@@ -0,0 +1,41 @@
+"""Utilities for converting data structures to JSON-serializable formats."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+
+from pydantic import BaseModel
+
+from ._utils import is_given, is_mapping
+
+
+def _transform_typeddict(data: Mapping[str, object]) -> dict[str, object]:
+    """
+    Transform a TypedDict-like mapping.
+
+    Args:
+        data: A `Mapping` to transform.
+
+    Returns:
+        A new dictionary with transformed values, excluding unset entries.
+    """
+    return {key: transform(value) for key, value in data.items() if is_given(value)}
+
+
+def transform(data: object) -> object:
+    """
+    Transform an object into a JSON-serializable format.
+
+    Args:
+        data: The object to transform.
+
+    Returns:
+        A JSON-serializable representation of the input data.
+    """
+    if is_mapping(data):
+        return _transform_typeddict(data)
+
+    if isinstance(data, BaseModel):
+        return data.model_dump(exclude_unset=True, mode="json")
+
+    return data

src/crowdstrike_aidr/_transform.py

@@ -1,20 +1,15 @@
 from __future__ import annotations
 
-from collections.abc import Mapping
 from typing import Literal
 
 import httpx
 
 from .._client import SyncAPIClient, make_request_options
+from .._transform import transform
 from .._types import Body, Headers, NotGiven, Omit, Query, not_given, omit
-from .._utils import is_given
 from ..models.ai_guard import ExtraInfo, GuardChatCompletionsResponse
 
 
-def _transform_typeddict(data: Mapping[str, object]) -> Mapping[str, object]:
-    return {key: value for key, value in data.items() if is_given(value)}
-
-
 class AIGuard(SyncAPIClient):
     _service_name: str = "aiguard"
 
@@ -83,7 +78,7 @@ def guard_chat_completions(
         """
         return self._post(
             "/v1/guard_chat_completions",
-            body=_transform_typeddict(
+            body=transform(
                 {
                     "guard_input": guard_input,
                     "app_id": app_id,

src/crowdstrike_aidr/services/ai_guard.py

@@ -6,7 +6,7 @@
 import pytest
 
 from crowdstrike_aidr import AIGuard
-from crowdstrike_aidr.models.ai_guard import GuardChatCompletionsResponse
+from crowdstrike_aidr.models.ai_guard import ExtraInfo, GuardChatCompletionsResponse
 
 from .utils import assert_matches_type
 
@@ -19,5 +19,8 @@ def client(request: pytest.FixtureRequest) -> Iterator[AIGuard]:
 
 
 def test_guard_chat_completions(client: AIGuard) -> None:
-    response = client.guard_chat_completions(guard_input={"messages": [{"role": "user", "content": "Hello, world!"}]})
+    response = client.guard_chat_completions(
+        guard_input={"messages": [{"role": "user", "content": "Hello, world!"}]},
+        extra_info=ExtraInfo(app_name="app_name"),
+    )
     assert_matches_type(GuardChatCompletionsResponse, response, path=["guard_chat_completions"])