From 80f05bf8a9fef046b30d80c878dc04142db83a48 Mon Sep 17 00:00:00 2001
From: cthomas <caren@letta.com>
Date: Mon, 6 Jan 2025 10:05:03 -0800
Subject: [PATCH] fix: add message union type ref for LettaResponse (#507)

---
 letta/schemas/letta_response.py | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/letta/schemas/letta_response.py b/letta/schemas/letta_response.py
index 10a1b380..d7019280 100644
--- a/letta/schemas/letta_response.py
+++ b/letta/schemas/letta_response.py
@@ -29,19 +29,20 @@ class LettaResponse(BaseModel):
         json_schema_extra={
             "items": {
                 "oneOf": [
-                    {"x-ref-name": "SystemMessage"},
-                    {"x-ref-name": "UserMessage"},
-                    {"x-ref-name": "ReasoningMessage"},
-                    {"x-ref-name": "ToolCallMessage"},
-                    {"x-ref-name": "ToolReturnMessage"},
-                    {"x-ref-name": "AssistantMessage"},
+                    {"$ref": "#/components/schemas/SystemMessage-Output"},
+                    {"$ref": "#/components/schemas/UserMessage-Output"},
+                    {"$ref": "#/components/schemas/ReasoningMessage"},
+                    {"$ref": "#/components/schemas/ToolCallMessage"},
+                    {"$ref": "#/components/schemas/ToolReturnMessage"},
+                    {"$ref": "#/components/schemas/AssistantMessage-Output"},
                 ],
                 "discriminator": {"propertyName": "message_type"},
             }
         },
     )
     usage: LettaUsageStatistics = Field(
-        ..., description="The usage statistics of the agent.", json_schema_extra={"x-ref-name": "LettaUsageStatistics"}
+        ...,
+        description="The usage statistics of the agent.",
     )
 
     def __str__(self):