fix bug with ws reasoning (#514)

Signed-off-by: Jessie Frazelle <[email protected]>

Author: jessfraz

generate/type_generators.py

@@ -557,48 +557,37 @@ def generate_one_of_type(path: str, name: str, schema: dict, data: dict):
                 flattened_schema["required"] = required_fields
 
             extra_imports = None
-            if not alias_imports_needed:
-                extra_imports = [
-                    "from pydantic import model_serializer, model_validator\n"
-                ]
-                alias_imports_needed = True
-
+            extra_body_lines = []
             wrapped_key = outer_name
-            extra_body_lines = [
-                '    @model_validator(mode="before")',
-                "    @classmethod",
-                "    def _unwrap(cls, data):",
-                "        if isinstance(data, dict) and '"
-                + wrapped_key
-                + "' in data and isinstance(data['"
-                + wrapped_key
-                + "'], dict):",
-                "            return data['" + wrapped_key + "']",
-                "        return data",
-            ]
 
             if wrap_entire_payload:
+                if not alias_imports_needed:
+                    extra_imports = [
+                        "from pydantic import model_serializer, model_validator\n"
+                    ]
+                    alias_imports_needed = True
+
                 extra_body_lines.extend(
                     [
-                        '    @model_serializer(mode="wrap")',
-                        "    def _wrap(self, handler, info):",
-                        "        payload = handler(self, info)",
-                        "        return {'" + wrapped_key + "': payload}",
+                        '    @model_validator(mode="before")',
+                        "    @classmethod",
+                        "    def _unwrap(cls, data):",
+                        "        if isinstance(data, dict) and '"
+                        + wrapped_key
+                        + "' in data and isinstance(data['"
+                        + wrapped_key
+                        + "'], dict):",
+                        "            return data['" + wrapped_key + "']",
+                        "        return data",
                     ]
                 )
-            else:
+
                 extra_body_lines.extend(
                     [
                         '    @model_serializer(mode="wrap")',
                         "    def _wrap(self, handler, info):",
                         "        payload = handler(self, info)",
-                        "        if isinstance(payload, dict) and '"
-                        + wrapped_key
-                        + "' in payload:",
-                        "            value = payload['" + wrapped_key + "']",
-                        "        else:",
-                        "            value = payload",
-                        "        return {'" + wrapped_key + "': value}",
+                        "        return {'" + wrapped_key + "': payload}",
                     ]
                 )
 

kittycad/models/ml_copilot_server_message.py

@@ -156,30 +156,6 @@ class Reasoning(KittyCadBaseModel):
 
     reasoning: ReasoningMessage
 
-    @model_validator(mode="before")
-    @classmethod
-    def _unwrap(cls, data):
-        if (
-            isinstance(data, dict)
-            and "reasoning" in data
-            and isinstance(data["reasoning"], dict)
-        ):
-            return data["reasoning"]
-
-        return data
-
-    @model_serializer(mode="wrap")
-    def _wrap(self, handler, info):
-        payload = handler(self, info)
-
-        if isinstance(payload, dict) and "reasoning" in payload:
-            value = payload["reasoning"]
-
-        else:
-            value = payload
-
-        return {"reasoning": value}
-
 
 class Replay(KittyCadBaseModel):
     """Replay containing raw bytes for previously-saved messages for a conversation. Includes server messages and client `User` messages.

kittycad/tests/test_ml_reasoning_ws_regression.py

@@ -0,0 +1,56 @@
+import json
+import uuid
+from typing import cast
+
+from websockets.sync.client import ClientConnection as ClientConnectionSync
+
+from kittycad import WebSocketMlReasoningWs
+from kittycad.models.ml_copilot_server_message import Reasoning, SessionData
+from kittycad.models.reasoning_message import OptionText
+
+
+class FakeWS:
+    def __init__(self, messages):
+        self._messages = iter(messages)
+
+    def recv(self, timeout=60):  # pragma: no cover - tiny helper
+        try:
+            return next(self._messages)
+        except StopIteration as exc:
+            raise AssertionError("unexpected recv() after messages exhausted") from exc
+
+
+def test_ml_reasoning_ws_recv_parses_reasoning_messages():
+    cache_buster = uuid.uuid4().hex
+    reasoning_content = (
+        f":mag: Querying relevant KCL code examples... cache-buster-{cache_buster}"
+    )
+
+    fake_ws = FakeWS(
+        [
+            json.dumps({"session_data": {"api_call_id": "abc123"}}),
+            json.dumps(
+                {
+                    "reasoning": {
+                        "type": "text",
+                        "content": reasoning_content,
+                    }
+                }
+            ),
+        ]
+    )
+
+    websocket = cast(
+        WebSocketMlReasoningWs, WebSocketMlReasoningWs.__new__(WebSocketMlReasoningWs)
+    )
+    websocket.ws = cast(ClientConnectionSync, fake_ws)
+
+    session_message = websocket.recv()
+    assert isinstance(session_message.root, SessionData)
+    assert session_message.root.api_call_id == "abc123"
+
+    reasoning_message = websocket.recv()
+    assert isinstance(reasoning_message.root, Reasoning)
+    resolved_reasoning = reasoning_message.root.reasoning.root
+    assert isinstance(resolved_reasoning, OptionText)
+    assert resolved_reasoning.content == reasoning_content

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "kittycad"
-version = "1.2.3"
+version = "1.2.4"
 description = "A client library for accessing KittyCAD"
 authors = []
 readme = "README.md"