dnrops
/
rig
mirror of https://github.com/0xplaygrounds/rig
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

feat: change StreamingCompletionResponse to resemble CompletionResponse

This commit is contained in:
yavens 2025-04-15 09:06:31 -04:00
parent cb4be5c6da
commit 9689693468
1 changed files with 13 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
rig-core/src/streaming.rs
View File

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