Fix client test

letta/client/client.py

@@ -25,6 +25,7 @@
 from letta.schemas.enums import JobStatus, MessageRole
 from letta.schemas.file import FileMetadata
 from letta.schemas.job import Job
+from letta.schemas.letta_message import LettaMessage
 from letta.schemas.letta_request import LettaRequest
 from letta.schemas.letta_response import LettaResponse, LettaStreamingResponse
 from letta.schemas.llm_config import LLMConfig
@@ -914,7 +915,7 @@ def delete_archival_memory(self, agent_id: str, memory_id: str):
 
     def get_messages(
         self, agent_id: str, before: Optional[str] = None, after: Optional[str] = None, limit: Optional[int] = 1000
-    ) -> List[Message]:
+    ) -> List[LettaMessage]:
         """
         Get messages from an agent with pagination.
 
@@ -925,14 +926,14 @@ def get_messages(
             limit (int): Limit number of messages
 
         Returns:
-            messages (List[Message]): List of messages
+            messages (List[LettaMessage]): List of messages
         """
 
         params = {"before": before, "after": after, "limit": limit, "msg_object": True}
         response = requests.get(f"{self.base_url}/{self.api_prefix}/agents/{agent_id}/messages", params=params, headers=self.headers)
         if response.status_code != 200:
             raise ValueError(f"Failed to get messages: {response.text}")
-        return [Message(**message) for message in response.json()]
+        return [LettaMessage(**message) for message in response.json()]
 
     def send_message(
         self,