refactor: patching openapi schema

letta/schemas/letta_request.py

@@ -1,13 +1,13 @@
-from typing import List, Union
+from typing import List
 
 from pydantic import BaseModel, Field
 
 from letta.constants import DEFAULT_MESSAGE_TOOL, DEFAULT_MESSAGE_TOOL_KWARG
-from letta.schemas.message import Message, MessageCreate
+from letta.schemas.message import MessageCreate
 
 
 class LettaRequest(BaseModel):
-    messages: Union[List[MessageCreate], List[Message]] = Field(..., description="The messages to be sent to the agent.")
+    messages: List[MessageCreate] = Field(..., description="The messages to be sent to the agent.")
     run_async: bool = Field(default=False, description="Whether to asynchronously send the messages to the agent.")  # TODO: implement
 
     stream_steps: bool = Field(

letta/schemas/letta_response.py

@@ -3,7 +3,7 @@
 import re
 from typing import List, Union
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, RootModel
 
 from letta.schemas.enums import MessageStreamStatus
 from letta.schemas.letta_message import LettaMessage, LettaMessageUnion
@@ -14,6 +14,20 @@
 # TODO: consider moving into own file
 
 
+class MessageListResponse(RootModel):
+    root: List[Message]
+
+    class Config:
+        title = "MessageListResponse"  # Explicit title for Stainless
+
+
+class LettaMessageListResponse(RootModel):
+    root: List[LettaMessageUnion]
+
+    class Config:
+        title = "LettaMessageListResponse"  # Explicit title for Stainless
+
+
 class LettaResponse(BaseModel):
     """
     Response object from an agent interaction, consisting of the new messages generated by the agent and usage statistics.

letta/server/rest_api/routers/v1/agents.py

@@ -8,13 +8,13 @@
 from letta.constants import DEFAULT_MESSAGE_TOOL, DEFAULT_MESSAGE_TOOL_KWARG
 from letta.schemas.agent import AgentState, CreateAgent, UpdateAgentState
 from letta.schemas.enums import MessageStreamStatus
-from letta.schemas.letta_message import (
-    LegacyLettaMessage,
-    LettaMessage,
-    LettaMessageUnion,
-)
+from letta.schemas.letta_message import LegacyLettaMessage, LettaMessage
 from letta.schemas.letta_request import LettaRequest
-from letta.schemas.letta_response import LettaResponse
+from letta.schemas.letta_response import (
+    LettaMessageListResponse,
+    LettaResponse,
+    MessageListResponse,
+)
 from letta.schemas.memory import (
     ArchivalMemorySummary,
     BasicBlockMemory,
@@ -310,7 +310,11 @@ def delete_agent_archival_memory(
     return JSONResponse(status_code=status.HTTP_200_OK, content={"message": f"Memory id={memory_id} successfully deleted"})
 
 
-@router.get("/{agent_id}/messages", response_model=Union[List[Message], List[LettaMessageUnion]], operation_id="list_agent_messages")
+@router.get(
+    "/{agent_id}/messages",
+    response_model=Union[MessageListResponse, LettaMessageListResponse],
+    operation_id="list_agent_messages",
+)
 def get_agent_messages(
     agent_id: str,
     server: "SyncServer" = Depends(get_letta_server),
@@ -372,8 +376,14 @@ def update_message(
         200: {
             "description": "Successful response",
             "content": {
-                "application/json": {"$ref": "#/components/schemas/LettaResponse"},  # Use model_json_schema() instead of model directly
-                "text/event-stream": {"description": "Server-Sent Events stream"},
+                "application/json": {
+                    "schema": {
+                        "$ref": "#/components/schemas/LettaResponse",
+                    },
+                },
+                "text/event-stream": {
+                    "description": "Server-Sent Events stream",
+                },
             },
         }
     },