feat: add message type literal to usage stats (#2297)

Co-authored-by: Caren Thomas <caren@caren-mac.local>
This commit is contained in:
cthomas
2024-12-20 17:13:56 -08:00
committed by GitHub
parent 7876c5d5de
commit 3da62aa4eb
3 changed files with 5 additions and 4 deletions

View File

@@ -59,8 +59,8 @@ def _sse_post(url: str, data: dict, headers: dict) -> Generator[LettaStreamingRe
yield ToolCallMessage(**chunk_data)
elif "tool_return" in chunk_data:
yield ToolReturnMessage(**chunk_data)
elif "usage" in chunk_data:
yield LettaUsageStatistics(**chunk_data["usage"])
elif "step_count" in chunk_data:
yield LettaUsageStatistics(**chunk_data)
else:
raise ValueError(f"Unknown message type in chunk_data: {chunk_data}")