Add printed out responses for easier debugging from tests

This commit is contained in:
Matt Zhou
2024-10-04 15:19:40 -07:00
parent 776d7dd6e8
commit 6bcec854d6
6 changed files with 75 additions and 15 deletions

View File

@@ -1,3 +1,4 @@
import json
from typing import List, Union
from pydantic import BaseModel, Field
@@ -23,6 +24,16 @@ class LettaResponse(BaseModel):
messages: Union[List[Message], List[LettaMessage]] = Field(..., description="The messages returned by the agent.")
usage: LettaUsageStatistics = Field(..., description="The usage statistics of the agent.")
def __str__(self):
return json.dumps(
{
"messages": [message.model_dump() for message in self.messages],
# Assume `Message` and `LettaMessage` have a `dict()` method
"usage": self.usage.model_dump(), # Assume `LettaUsageStatistics` has a `dict()` method
},
indent=4,
)
# The streaming response is either [DONE], [DONE_STEP], [DONE], an error, or a LettaMessage
LettaStreamingResponse = Union[LettaMessage, MessageStreamStatus]