Consumes the next [SendMessageOutput] event.
(&mut self)
| 589 | |
| 590 | /// Consumes the next [SendMessageOutput] event. |
| 591 | async fn next(&mut self) -> Result<Option<ChatResponseStream>, RecvError> { |
| 592 | if let Some(ev) = self.peek.take() { |
| 593 | return Ok(Some(ev)); |
| 594 | } |
| 595 | trace!("Attempting to recv next event"); |
| 596 | let start = std::time::Instant::now(); |
| 597 | let result = self.response.recv().await; |
| 598 | let duration = std::time::Instant::now().duration_since(start); |
| 599 | match result { |
| 600 | Ok(ev) => { |
| 601 | trace!(?ev, "Received new event"); |
| 602 | |
| 603 | // Track metadata about the chunk. |
| 604 | self.time_to_first_chunk |
| 605 | .get_or_insert_with(|| self.request_start_time.elapsed()); |
| 606 | self.time_between_chunks.push(duration); |
| 607 | if let Some(r) = ev.as_ref() { |
| 608 | match r { |
| 609 | ChatResponseStream::AssistantResponseEvent { content } => { |
| 610 | self.received_response_size += content.len(); |
| 611 | }, |
| 612 | ChatResponseStream::ToolUseEvent { input, .. } => { |
| 613 | self.received_response_size += input.as_ref().map(String::len).unwrap_or_default(); |
| 614 | }, |
| 615 | _ => { |
| 616 | warn!(?r, "received unexpected event from the response stream"); |
| 617 | }, |
| 618 | } |
| 619 | } |
| 620 | |
| 621 | Ok(ev) |
| 622 | }, |
| 623 | Err(err) => { |
| 624 | error!(?err, "failed to receive the next event"); |
| 625 | if duration.as_secs() >= 59 { |
| 626 | Err(self.error(RecvErrorKind::StreamTimeout { source: err, duration })) |
| 627 | } else { |
| 628 | Err(self.error(err)) |
| 629 | } |
| 630 | }, |
| 631 | } |
| 632 | } |
| 633 | |
| 634 | /// Helper to create a new [RecvError] populated with the associated request id for the stream. |
| 635 | fn error(&self, source: impl Into<RecvErrorKind>) -> RecvError { |