vtcode_llm/
single_response.rs1#[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 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}