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
14pub async fn handle_ask_command(
15    config: AgentConfig,
16    prompt: Vec<String>,
17    options: crate::cli::AskCommandOptions,
18) -> Result<()> {
19    let prompt_text = prompt.join(" ");
20    validate_agent_safe_text("prompt", &prompt_text)?;
21
22    if config.verbose {
23        tracing::debug!(model = %config.model, prompt = %prompt_text, "Sending prompt");
24    }
25
26    let request = LLMRequest {
27        messages: vec![Message::user(prompt_text)],
28        system_prompt: Some(Arc::new(lightweight_instruction_text())),
29        model: config.model.clone(),
30        ..Default::default()
31    };
32    let provider_name = if config.provider.trim().is_empty() {
33        infer_provider_from_model(&request.model)
34            .map(|provider| provider.to_string())
35            .ok_or_else(|| {
36                anyhow::anyhow!("Cannot determine provider for model: {}", request.model)
37            })?
38    } else {
39        config.provider.to_lowercase()
40    };
41    let provider = create_provider_with_config(
42        &provider_name,
43        ProviderConfig {
44            api_key: Some(config.api_key.clone()),
45            openai_chatgpt_auth: config.openai_chatgpt_auth.clone(),
46            copilot_auth: None,
47            base_url: None,
48            model: Some(request.model.clone()),
49            prompt_cache: None,
50            timeouts: None,
51            openai: None,
52            anthropic: 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,
117        LLMStreamEvent, NormalizedStreamEvent, 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 {
146                    delta: "hello ".to_string(),
147                }),
148                Ok(LLMStreamEvent::Token {
149                    delta: "world".to_string(),
150                }),
151                Ok(LLMStreamEvent::Completed {
152                    response: Box::new(LLMResponse {
153                        content: None,
154                        model: "gpt-5.2-codex".to_string(),
155                        tool_calls: None,
156                        usage: None,
157                        finish_reason: FinishReason::Stop,
158                        reasoning: None,
159                        reasoning_details: None,
160                        organization_id: None,
161                        request_id: None,
162                        tool_references: Vec::new(),
163                        compaction: None,
164                    }),
165                }),
166            ])))
167        }
168
169        fn supported_models(&self) -> Vec<String> {
170            vec!["gpt-5.2-codex".to_string()]
171        }
172
173        fn validate_request(&self, _request: &LLMRequest) -> Result<(), LLMError> {
174            Ok(())
175        }
176    }
177
178    #[tokio::test]
179    async fn collect_single_response_uses_stream_when_non_streaming_is_unsupported() {
180        let provider = StreamingOnlyProvider;
181        let response = collect_single_response(
182            &provider,
183            LLMRequest {
184                model: "gpt-5.2-codex".to_string(),
185                ..Default::default()
186            },
187        )
188        .await
189        .expect("stream collection should succeed");
190
191        assert_eq!(response.content.as_deref(), Some("hello world"));
192    }
193
194    #[derive(Clone)]
195    struct NormalizedOnlyProvider;
196
197    #[async_trait]
198    impl LLMProvider for NormalizedOnlyProvider {
199        fn name(&self) -> &str {
200            "test"
201        }
202
203        fn supports_streaming(&self) -> bool {
204            true
205        }
206
207        fn supports_non_streaming(&self, _model: &str) -> bool {
208            false
209        }
210
211        async fn generate(&self, _request: LLMRequest) -> Result<LLMResponse, LLMError> {
212            panic!("generate should not be called for streaming-only provider")
213        }
214
215        async fn stream(&self, _request: LLMRequest) -> Result<LLMStream, LLMError> {
216            panic!("legacy stream should not be used when normalized stream is available")
217        }
218
219        async fn stream_normalized(
220            &self,
221            _request: LLMRequest,
222        ) -> Result<LLMNormalizedStream, LLMError> {
223            Ok(Box::pin(stream::iter(vec![
224                Ok(NormalizedStreamEvent::TextDelta {
225                    delta: "hello ".to_string(),
226                }),
227                Ok(NormalizedStreamEvent::ReasoningDelta {
228                    delta: "thinking ".to_string(),
229                }),
230                Ok(NormalizedStreamEvent::Usage {
231                    usage: Usage {
232                        prompt_tokens: 10,
233                        completion_tokens: 2,
234                        total_tokens: 12,
235                        cached_prompt_tokens: None,
236                        cache_creation_tokens: None,
237                        cache_read_tokens: None,
238                    },
239                }),
240                Ok(NormalizedStreamEvent::Done {
241                    response: Box::new(LLMResponse {
242                        content: None,
243                        model: "gpt-5.2-codex".to_string(),
244                        tool_calls: None,
245                        usage: None,
246                        finish_reason: FinishReason::Stop,
247                        reasoning: None,
248                        reasoning_details: None,
249                        organization_id: None,
250                        request_id: None,
251                        tool_references: Vec::new(),
252                        compaction: None,
253                    }),
254                }),
255            ])))
256        }
257
258        fn supported_models(&self) -> Vec<String> {
259            vec!["gpt-5.2-codex".to_string()]
260        }
261
262        fn validate_request(&self, _request: &LLMRequest) -> Result<(), LLMError> {
263            Ok(())
264        }
265    }
266
267    #[tokio::test]
268    async fn collect_single_response_prefers_normalized_stream() {
269        let provider = NormalizedOnlyProvider;
270        let response = collect_single_response(
271            &provider,
272            LLMRequest {
273                model: "gpt-5.2-codex".to_string(),
274                ..Default::default()
275            },
276        )
277        .await
278        .expect("normalized stream collection should succeed");
279
280        assert_eq!(response.content.as_deref(), Some("hello "));
281        assert_eq!(response.reasoning.as_deref(), Some("thinking "));
282        assert_eq!(
283            response.usage.as_ref().map(|usage| usage.total_tokens),
284            Some(12)
285        );
286    }
287}
288
289fn extract_code_fence_blocks(text: &str) -> Vec<CodeFenceBlock> {
290    let mut blocks = Vec::new();
291    let mut current_language: Option<String> = None;
292    let mut current_lines: Vec<String> = Vec::new();
293
294    for raw_line in text.lines() {
295        let trimmed_start = raw_line.trim_start();
296        if let Some(rest) = trimmed_start.strip_prefix("```") {
297            let rest_clean = rest.trim_matches('\r');
298            let rest_trimmed = rest_clean.trim();
299            if current_language.is_some() {
300                if rest_trimmed.is_empty() {
301                    let language = current_language.take().and_then(|lang| {
302                        let cleaned = lang.trim_matches(|ch| matches!(ch, '"' | '\'' | '`'));
303                        let cleaned = cleaned.trim();
304                        if cleaned.is_empty() {
305                            None
306                        } else {
307                            Some(cleaned.to_string())
308                        }
309                    });
310                    let block_lines = std::mem::take(&mut current_lines);
311                    blocks.push(CodeFenceBlock {
312                        language,
313                        lines: block_lines,
314                    });
315                    continue;
316                }
317            } else {
318                let token = rest_trimmed.split_whitespace().next().unwrap_or_default();
319                let normalized = token
320                    .trim_matches(|ch| matches!(ch, '"' | '\'' | '`'))
321                    .trim();
322                current_language = Some(normalized.to_ascii_lowercase());
323                current_lines.clear();
324                continue;
325            }
326        }
327
328        if current_language.is_some() {
329            current_lines.push(raw_line.trim_end_matches('\r').to_string());
330        }
331    }
332
333    blocks
334}
335
336fn select_best_code_block(blocks: &[CodeFenceBlock]) -> Option<&CodeFenceBlock> {
337    let mut best = None;
338    let mut best_score = (0usize, 0u8);
339    for block in blocks {
340        let score = score_code_block(block);
341        if score > best_score {
342            best_score = score;
343            best = Some(block);
344        }
345    }
346    best
347}
348
349fn score_code_block(block: &CodeFenceBlock) -> (usize, u8) {
350    let line_count = block
351        .lines
352        .iter()
353        .filter(|line| !line.trim().is_empty())
354        .count();
355    let has_language = block
356        .language
357        .as_ref()
358        .is_some_and(|lang| !lang.trim().is_empty());
359    (line_count, if has_language { 1 } else { 0 })
360}
361
362#[derive(Debug, Clone, PartialEq, Eq)]
363struct CodeFenceBlock {
364    language: Option<String>,
365    lines: Vec<String>,
366}
vtcode_core/commands/ask.rs

vtcode_core/commands/
ask.rs
