vtcode_core/commands/
ask.rs

1//! Ask command implementation - single prompt without tools
2
3use crate::cli::input_hardening::validate_agent_safe_text;
4use crate::config::types::AgentConfig;
5use crate::llm::collect_single_response;
6use crate::llm::factory::{ProviderConfig, create_provider_with_config, infer_provider_from_model};
7use crate::llm::provider::{LLMRequest, Message};
8use crate::prompts::system::lightweight_instruction_text;
9use anyhow::Result;
10use std::io::IsTerminal;
11use std::sync::Arc;
12
13/// Handle the ask command - single prompt without tools
14#[cfg_attr(feature = "profiling", hotpath::measure)]
15pub async fn handle_ask_command(
16    config: AgentConfig,
17    prompt: Vec<String>,
18    options: crate::cli::AskCommandOptions,
19) -> Result<()> {
20    let prompt_text = prompt.join(" ");
21    validate_agent_safe_text("prompt", &prompt_text)?;
22
23    if config.verbose {
24        tracing::debug!(model = %config.model, prompt = %prompt_text, "Sending prompt");
25    }
26
27    let request = LLMRequest {
28        messages: Arc::new(vec![Message::user(prompt_text)]),
29        system_prompt: Some(Arc::from(lightweight_instruction_text())),
30        model: config.model.clone(),
31        ..Default::default()
32    };
33    let provider_name = if config.provider.trim().is_empty() {
34        infer_provider_from_model(&request.model)
35            .map(|provider| provider.to_string())
36            .ok_or_else(|| anyhow::anyhow!("Cannot determine provider for model: {}", request.model))?
37    } else {
38        config.provider.to_lowercase()
39    };
40    let provider = create_provider_with_config(
41        &provider_name,
42        ProviderConfig {
43            api_key: Some(config.api_key.clone()),
44            openai_chatgpt_auth: config.openai_chatgpt_auth.clone(),
45            copilot_auth: None,
46            base_url: None,
47            model: Some(request.model.clone()),
48            prompt_cache: None,
49            timeouts: None,
50            openai: None,
51            anthropic: None,
52            openrouter: None,
53            model_behavior: config.model_behavior.clone(),
54            workspace_root: Some(config.workspace.clone()),
55        },
56    )?;
57    let backend_kind = provider.name().to_string();
58    let response = collect_single_response(provider.as_ref(), request).await?;
59    let response_model = if response.model.is_empty() {
60        config.model.clone()
61    } else {
62        response.model.clone()
63    };
64
65    // Handle output based on format preference
66    if let Some(crate::cli::args::AskOutputFormat::Json) = options.output_format {
67        // Build a comprehensive JSON structure
68        let output = serde_json::json!({
69            "response": response,
70            "provider": {
71                "kind": backend_kind,
72                "model": response_model,
73            }
74        });
75        use std::io::Write;
76        let mut stdout = std::io::stdout().lock();
77        serde_json::to_writer_pretty(&mut stdout, &output)?;
78        writeln!(stdout)?;
79    } else {
80        use std::io::Write;
81        let mut stdout = std::io::stdout().lock();
82        if is_pipe_output() {
83            if let Some(code_only) = extract_code_only(response.content_text()) {
84                write!(stdout, "{code_only}")?;
85            } else {
86                writeln!(stdout, "{}", response.content_text())?;
87            }
88        } else {
89            // Print the response content directly (default behavior)
90            writeln!(stdout, "{}", response.content_text())?;
91        }
92    }
93
94    Ok(())
95}
96
97fn is_pipe_output() -> bool {
98    !std::io::stdout().is_terminal()
99}
100
101fn extract_code_only(text: &str) -> Option<String> {
102    let blocks = extract_code_fence_blocks(text);
103    let block = select_best_code_block(&blocks)?;
104    let mut output = block.lines.join("\n");
105    if !output.ends_with('\n') {
106        output.push('\n');
107    }
108    Some(output)
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114    use crate::llm::collect_single_response;
115    use crate::llm::provider::{
116        FinishReason, LLMError, LLMNormalizedStream, LLMProvider, LLMResponse, LLMStream, LLMStreamEvent,
117        NormalizedStreamEvent, ReasoningSource, Usage,
118    };
119    use async_trait::async_trait;
120    use futures::stream;
121
122    #[derive(Clone)]
123    struct StreamingOnlyProvider;
124
125    #[async_trait]
126    impl LLMProvider for StreamingOnlyProvider {
127        fn name(&self) -> &str {
128            "test"
129        }
130
131        fn supports_streaming(&self) -> bool {
132            true
133        }
134
135        fn supports_non_streaming(&self, _model: &str) -> bool {
136            false
137        }
138
139        async fn generate(&self, _request: LLMRequest) -> Result<LLMResponse, LLMError> {
140            panic!("generate should not be called for streaming-only provider")
141        }
142
143        async fn stream(&self, _request: LLMRequest) -> Result<LLMStream, LLMError> {
144            Ok(Box::pin(stream::iter(vec![
145                Ok(LLMStreamEvent::Token { delta: "hello ".to_string() }),
146                Ok(LLMStreamEvent::Token { delta: "world".to_string() }),
147                Ok(LLMStreamEvent::Completed {
148                    response: Box::new(LLMResponse {
149                        content: None,
150                        model: "gpt-5.2-codex".to_string(),
151                        tool_calls: None,
152                        usage: None,
153                        finish_reason: FinishReason::Stop,
154                        reasoning: None,
155                        reasoning_details: None,
156                        organization_id: None,
157                        request_id: None,
158                        tool_references: Vec::new(),
159                        compaction: None,
160                    }),
161                }),
162            ])))
163        }
164
165        fn supported_models(&self) -> Vec<String> {
166            vec!["gpt-5.2-codex".to_string()]
167        }
168
169        fn validate_request(&self, _request: &LLMRequest) -> Result<(), LLMError> {
170            Ok(())
171        }
172    }
173
174    #[tokio::test]
175    async fn collect_single_response_uses_stream_when_non_streaming_is_unsupported() {
176        let provider = StreamingOnlyProvider;
177        let response = collect_single_response(
178            &provider,
179            LLMRequest {
180                model: "gpt-5.2-codex".to_string(),
181                ..Default::default()
182            },
183        )
184        .await
185        .expect("stream collection should succeed");
186
187        assert_eq!(response.content.as_deref(), Some("hello world"));
188    }
189
190    #[derive(Clone)]
191    struct NormalizedOnlyProvider;
192
193    #[async_trait]
194    impl LLMProvider for NormalizedOnlyProvider {
195        fn name(&self) -> &str {
196            "test"
197        }
198
199        fn supports_streaming(&self) -> bool {
200            true
201        }
202
203        fn supports_non_streaming(&self, _model: &str) -> bool {
204            false
205        }
206
207        async fn generate(&self, _request: LLMRequest) -> Result<LLMResponse, LLMError> {
208            panic!("generate should not be called for streaming-only provider")
209        }
210
211        async fn stream(&self, _request: LLMRequest) -> Result<LLMStream, LLMError> {
212            panic!("legacy stream should not be used when normalized stream is available")
213        }
214
215        async fn stream_normalized(&self, _request: LLMRequest) -> Result<LLMNormalizedStream, LLMError> {
216            Ok(Box::pin(stream::iter(vec![
217                Ok(NormalizedStreamEvent::TextDelta { delta: "hello ".to_string() }),
218                Ok(NormalizedStreamEvent::ReasoningDelta {
219                    delta: "thinking ".to_string(),
220                    source: ReasoningSource::ProviderSummary,
221                }),
222                Ok(NormalizedStreamEvent::Usage {
223                    usage: Usage {
224                        prompt_tokens: 10,
225                        completion_tokens: 2,
226                        total_tokens: 12,
227                        cached_prompt_tokens: None,
228                        cache_creation_tokens: None,
229                        cache_read_tokens: None,
230                        iterations: None,
231                    },
232                }),
233                Ok(NormalizedStreamEvent::Done {
234                    response: Box::new(LLMResponse {
235                        content: None,
236                        model: "gpt-5.2-codex".to_string(),
237                        tool_calls: None,
238                        usage: None,
239                        finish_reason: FinishReason::Stop,
240                        reasoning: None,
241                        reasoning_details: None,
242                        organization_id: None,
243                        request_id: None,
244                        tool_references: Vec::new(),
245                        compaction: None,
246                    }),
247                }),
248            ])))
249        }
250
251        fn supported_models(&self) -> Vec<String> {
252            vec!["gpt-5.2-codex".to_string()]
253        }
254
255        fn validate_request(&self, _request: &LLMRequest) -> Result<(), LLMError> {
256            Ok(())
257        }
258    }
259
260    #[tokio::test]
261    async fn collect_single_response_prefers_normalized_stream() {
262        let provider = NormalizedOnlyProvider;
263        let response = collect_single_response(
264            &provider,
265            LLMRequest {
266                model: "gpt-5.2-codex".to_string(),
267                ..Default::default()
268            },
269        )
270        .await
271        .expect("normalized stream collection should succeed");
272
273        assert_eq!(response.content.as_deref(), Some("hello "));
274        assert_eq!(response.reasoning.as_deref(), Some("thinking "));
275        assert_eq!(response.usage.as_ref().map(|usage| usage.total_tokens), Some(12));
276    }
277}
278
279fn extract_code_fence_blocks(text: &str) -> Vec<CodeFenceBlock> {
280    let mut blocks = Vec::new();
281    let mut current_language: Option<String> = None;
282    let mut current_lines: Vec<String> = Vec::new();
283
284    for raw_line in text.lines() {
285        let trimmed_start = raw_line.trim_start();
286        if let Some(rest) = trimmed_start.strip_prefix("```") {
287            let rest_clean = rest.trim_matches('\r');
288            let rest_trimmed = rest_clean.trim();
289            if current_language.is_some() {
290                if rest_trimmed.is_empty() {
291                    let language = current_language.take().and_then(|lang| {
292                        let cleaned = lang.trim_matches(|ch| matches!(ch, '"' | '\'' | '`'));
293                        let cleaned = cleaned.trim();
294                        if cleaned.is_empty() {
295                            None
296                        } else {
297                            Some(cleaned.to_string())
298                        }
299                    });
300                    let block_lines = std::mem::take(&mut current_lines);
301                    blocks.push(CodeFenceBlock { language, lines: block_lines });
302                    continue;
303                }
304            } else {
305                let token = rest_trimmed.split_whitespace().next().unwrap_or_default();
306                let normalized = token.trim_matches(|ch| matches!(ch, '"' | '\'' | '`')).trim();
307                current_language = Some(normalized.to_ascii_lowercase());
308                current_lines.clear();
309                continue;
310            }
311        }
312
313        if current_language.is_some() {
314            current_lines.push(raw_line.trim_end_matches('\r').to_string());
315        }
316    }
317
318    blocks
319}
320
321fn select_best_code_block(blocks: &[CodeFenceBlock]) -> Option<&CodeFenceBlock> {
322    let mut best = None;
323    let mut best_score = (0usize, 0u8);
324    for block in blocks {
325        let score = score_code_block(block);
326        if score > best_score {
327            best_score = score;
328            best = Some(block);
329        }
330    }
331    best
332}
333
334fn score_code_block(block: &CodeFenceBlock) -> (usize, u8) {
335    let line_count = block.lines.iter().filter(|line| !line.trim().is_empty()).count();
336    let has_language = block.language.as_ref().is_some_and(|lang| !lang.trim().is_empty());
337    (line_count, if has_language { 1 } else { 0 })
338}
339
340#[derive(Debug, Clone, PartialEq, Eq)]
341struct CodeFenceBlock {
342    language: Option<String>,
343    lines: Vec<String>,
344}
vtcode_core/commands/ask.rs

vtcode_core/commands/
ask.rs
