diff --git a/rig-core/src/streaming.rs b/rig-core/src/streaming.rs index 603d483..100f699 100644 --- a/rig-core/src/streaming.rs +++ b/rig-core/src/streaming.rs @@ -11,7 +11,8 @@ use crate::agent::Agent; use crate::completion::{ - CompletionError, CompletionModel, CompletionRequest, CompletionRequestBuilder, Message, + CompletionError, CompletionModel, CompletionRequest, CompletionRequestBuilder, + CompletionResponse, Message, }; use crate::message::{AssistantContent, ToolCall, ToolFunction}; use crate::OneOrMany; @@ -52,7 +53,7 @@ pub struct StreamingCompletionResponse { tool_calls: Vec, /// The final aggregated message from the stream /// contains all text and tool calls generated - pub message: Message, + pub choice: OneOrMany, /// The final response from the stream, may be `None` /// if the provider didn't yield it during the stream pub response: Option, @@ -64,12 +65,21 @@ impl StreamingCompletionResponse { inner, text: "".to_string(), tool_calls: vec![], - message: Message::assistant(""), + choice: OneOrMany::one(AssistantContent::text("")), response: None, } } } +impl Into>> for StreamingCompletionResponse { + fn into(self) -> CompletionResponse> { + CompletionResponse { + choice: self.choice, + raw_response: self.response, + } + } +} + impl Stream for StreamingCompletionResponse { type Item = Result;