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