Skip to main content

vtcode_llm/
single_response.rs

1#[cfg(feature = "profiling")]
2use super::provider::LLMNormalizedStream;
3use super::provider::{LLMError, LLMProvider, LLMRequest, LLMResponse, NormalizedStreamEvent};
4use futures::StreamExt as _;
5
6#[cfg_attr(feature = "profiling", hotpath::measure)]
7pub async fn collect_single_response(
8    provider: &(impl LLMProvider + ?Sized),
9    request: LLMRequest,
10) -> Result<LLMResponse, LLMError> {
11    if provider.supports_non_streaming(&request.model) {
12        #[cfg(feature = "profiling")]
13        return hotpath::future!(provider.generate(request), label = "llm_non_streaming").await;
14        #[cfg(not(feature = "profiling"))]
15        return provider.generate(request).await;
16    }
17
18    // `stream_normalized` is the fallback transport for this request, so make
19    // the selected wire mode explicit for adapters that inspect the request
20    // instead of unconditionally setting it in their transport layer.
21    let mut request = request;
22    request.stream = true;
23
24    #[cfg(feature = "profiling")]
25    let mut stream: LLMNormalizedStream =
26        hotpath::future!(provider.stream_normalized(request), label = "llm_streaming").await?;
27    #[cfg(not(feature = "profiling"))]
28    let mut stream = provider.stream_normalized(request).await?;
29    let mut streamed_content = String::with_capacity(4096);
30    let mut streamed_reasoning = String::with_capacity(1024);
31    let mut streamed_usage = None;
32    let mut completed = None;
33
34    while let Some(event) = stream.next().await {
35        match event? {
36            NormalizedStreamEvent::TextDelta { delta } => streamed_content.push_str(&delta),
37            NormalizedStreamEvent::ReasoningDelta { delta, source } if source.is_public_summary() => {
38                streamed_reasoning.push_str(&delta)
39            }
40            NormalizedStreamEvent::ReasoningDelta { .. } => {}
41            NormalizedStreamEvent::ReasoningStage { .. } => {}
42            NormalizedStreamEvent::ToolCallStart { .. } | NormalizedStreamEvent::ToolCallDelta { .. } => {}
43            NormalizedStreamEvent::Usage { usage } => streamed_usage = Some(usage),
44            NormalizedStreamEvent::Done { response } => {
45                completed = Some(*response);
46                break;
47            }
48        }
49    }
50
51    let mut response = completed.ok_or_else(|| LLMError::Provider {
52        message: format!("{} stream ended without a completed response", provider.name()),
53        metadata: None,
54    })?;
55    if response.usage.is_none() {
56        response.usage = streamed_usage;
57    }
58    if response.content.as_deref().unwrap_or_default().is_empty() && !streamed_content.is_empty() {
59        response.content = Some(streamed_content);
60    }
61    if response.reasoning.is_none() && !streamed_reasoning.is_empty() {
62        response.reasoning = Some(streamed_reasoning);
63    }
64    Ok(response)
65}