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)]
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 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}