From 7141e5909dca34eac2b93d7d9164ef85110a3977 Mon Sep 17 00:00:00 2001
From: Caren Thomas <caren@letta.com>
Date: Thu, 26 Dec 2024 14:22:10 -0800
Subject: [PATCH] remove lettaresponse spec hardcode and add references to
 pydantic

---
 letta/schemas/letta_response.py | 28 ++++++++++++++++++++++++++--
 letta/server/rest_api/app.py    |  3 ---
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/letta/schemas/letta_response.py b/letta/schemas/letta_response.py
index c6a1e8be..d12b2dc7 100644
--- a/letta/schemas/letta_response.py
+++ b/letta/schemas/letta_response.py
@@ -23,8 +23,32 @@ class LettaResponse(BaseModel):
         usage (LettaUsageStatistics): The usage statistics
     """
 
-    messages: List[LettaMessageUnion] = Field(..., description="The messages returned by the agent.")
-    usage: LettaUsageStatistics = Field(..., description="The usage statistics of the agent.")
+    messages: List[LettaMessageUnion] = Field(
+        ..., 
+        description="The messages returned by the agent.",
+        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"}
+                ],
+                "discriminator": {
+                    "propertyName": "message_type"
+                }
+            }
+        }
+    )
+    usage: LettaUsageStatistics = Field(
+        ..., 
+        description="The usage statistics of the agent.",
+        json_schema_extra={
+            "x-ref-name": "LettaUsageStatistics"
+        }
+    )
 
     def __str__(self):
         return json_dumps(
diff --git a/letta/server/rest_api/app.py b/letta/server/rest_api/app.py
index 8e2b7fb7..a8d31df4 100644
--- a/letta/server/rest_api/app.py
+++ b/letta/server/rest_api/app.py
@@ -68,9 +68,6 @@ def generate_openapi_schema(app: FastAPI):
     openai_docs["info"]["title"] = "OpenAI Assistants API"
     letta_docs["paths"] = {k: v for k, v in letta_docs["paths"].items() if not k.startswith("/openai")}
     letta_docs["info"]["title"] = "Letta API"
-    letta_docs["components"]["schemas"]["LettaResponse"] = {
-        "properties": LettaResponse.model_json_schema(ref_template="#/components/schemas/LettaResponse/properties/{model}")["$defs"]
-    }
 
     # Split the API docs into Letta API, and OpenAI Assistants compatible API
     for name, docs in [