Skip to main content

vtcode_core/llm/
single_response.rs

1use super::provider::{LLMError, LLMProvider, LLMRequest, LLMResponse, NormalizedStreamEvent};
2use futures::StreamExt as _;
3
4pub async fn collect_single_response(
5    provider: &(impl LLMProvider + ?Sized),
6    request: LLMRequest,
7) -> Result<LLMResponse, LLMError> {
8    if provider.supports_non_streaming(&request.model) {
9        return provider.generate(request).await;
10    }
11
12    let mut stream = provider.stream_normalized(request).await?;
13    let mut streamed_content = String::new();
14    let mut streamed_reasoning = String::new();
15    let mut streamed_usage = None;
16    let mut completed = None;
17
18    while let Some(event) = stream.next().await {
19        match event? {
20            NormalizedStreamEvent::TextDelta { delta } => streamed_content.push_str(&delta),
21            NormalizedStreamEvent::ReasoningDelta { delta } => streamed_reasoning.push_str(&delta),
22            NormalizedStreamEvent::ToolCallStart { .. }
23            | NormalizedStreamEvent::ToolCallDelta { .. } => {}
24            NormalizedStreamEvent::Usage { usage } => streamed_usage = Some(usage),
25            NormalizedStreamEvent::Done { response } => {
26                completed = Some(*response);
27                break;
28            }
29        }
30    }
31
32    let mut response = completed.ok_or_else(|| LLMError::Provider {
33        message: format!(
34            "{} stream ended without a completed response",
35            provider.name()
36        ),
37        metadata: None,
38    })?;
39    if response.usage.is_none() {
40        response.usage = streamed_usage;
41    }
42    if response.content.as_deref().unwrap_or_default().is_empty() && !streamed_content.is_empty() {
43        response.content = Some(streamed_content);
44    }
45    if response.reasoning.is_none() && !streamed_reasoning.is_empty() {
46        response.reasoning = Some(streamed_reasoning);
47    }
48    Ok(response)
49}