Skip to content

Commit

Permalink
fix(server): solve the chat error (#720)
Browse files Browse the repository at this point in the history
  • Loading branch information
RaoHai authored Feb 5, 2025
2 parents ab40b6b + c0ac11e commit 406973b
Showing 1 changed file with 13 additions and 9 deletions.
22 changes: 13 additions & 9 deletions server/agent/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -156,13 +156,17 @@ async def run_stream_chat(self, input_data: ChatData) -> AsyncIterator[Dict]:
output = event.get("data", {}).get("output", {})
generations = output.get("generations", [])
if generations and len(generations) > 0:
content = generations[0][0].get("message", {}).get("usage_metadata")
if content:
yield {
"id": event["run_id"],
"type": "usage",
**content,
}
generation = generations[0][0]
message = generation.get("message")
if message:
content = message.usage_metadata
if content:
yield {
"id": event["run_id"],
"type": "usage",
**content,
}

elif kind == "on_tool_start":
children_value = event["data"].get("input", {})
yield {
Expand Down Expand Up @@ -216,9 +220,9 @@ async def run_stream_chat(self, input_data: ChatData) -> AsyncIterator[Dict]:

except Exception as e:
if isinstance(e, APIError):
yield { "status": "error", "message": e.body }
yield {"status": "error", "message": e.body}
else:
yield { "status": "error", "message": str(e) }
yield {"status": "error", "message": str(e)}

async def run_chat(self, input_data: ChatData) -> str:
try:
Expand Down

0 comments on commit 406973b

Please sign in to comment.