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                        iterations: None,
239                    },
240                }),
241                Ok(NormalizedStreamEvent::Done {
242                    response: Box::new(LLMResponse {
243                        content: None,
244                        model: "gpt-5.2-codex".to_string(),
245                        tool_calls: None,
246                        usage: None,
247                        finish_reason: FinishReason::Stop,
248                        reasoning: None,
249                        reasoning_details: None,
250                        organization_id: None,
251                        request_id: None,
252                        tool_references: Vec::new(),
253                        compaction: None,
254                    }),
255                }),
256            ])))
257        }
258
259        fn supported_models(&self) -> Vec<String> {
260            vec!["gpt-5.2-codex".to_string()]
261        }
262
263        fn validate_request(&self, _request: &LLMRequest) -> Result<(), LLMError> {
264            Ok(())
265        }
266    }
267
268    #[tokio::test]
269    async fn collect_single_response_prefers_normalized_stream() {
270        let provider = NormalizedOnlyProvider;
271        let response = collect_single_response(
272            &provider,
273            LLMRequest {
274                model: "gpt-5.2-codex".to_string(),
275                ..Default::default()
276            },
277        )
278        .await
279        .expect("normalized stream collection should succeed");
280
281        assert_eq!(response.content.as_deref(), Some("hello "));
282        assert_eq!(response.reasoning.as_deref(), Some("thinking "));
283        assert_eq!(
284            response.usage.as_ref().map(|usage| usage.total_tokens),
285            Some(12)
286        );
287    }
288}
289
290fn extract_code_fence_blocks(text: &str) -> Vec<CodeFenceBlock> {
291    let mut blocks = Vec::new();
292    let mut current_language: Option<String> = None;
293    let mut current_lines: Vec<String> = Vec::new();
294
295    for raw_line in text.lines() {
296        let trimmed_start = raw_line.trim_start();
297        if let Some(rest) = trimmed_start.strip_prefix("```") {
298            let rest_clean = rest.trim_matches('\r');
299            let rest_trimmed = rest_clean.trim();
300            if current_language.is_some() {
301                if rest_trimmed.is_empty() {
302                    let language = current_language.take().and_then(|lang| {
303                        let cleaned = lang.trim_matches(|ch| matches!(ch, '"' | '\'' | '`'));
304                        let cleaned = cleaned.trim();
305                        if cleaned.is_empty() {
306                            None
307                        } else {
308                            Some(cleaned.to_string())
309                        }
310                    });
311                    let block_lines = std::mem::take(&mut current_lines);
312                    blocks.push(CodeFenceBlock {
313                        language,
314                        lines: block_lines,
315                    });
316                    continue;
317                }
318            } else {
319                let token = rest_trimmed.split_whitespace().next().unwrap_or_default();
320                let normalized = token
321                    .trim_matches(|ch| matches!(ch, '"' | '\'' | '`'))
322                    .trim();
323                current_language = Some(normalized.to_ascii_lowercase());
324                current_lines.clear();
325                continue;
326            }
327        }
328
329        if current_language.is_some() {
330            current_lines.push(raw_line.trim_end_matches('\r').to_string());
331        }
332    }
333
334    blocks
335}
336
337fn select_best_code_block(blocks: &[CodeFenceBlock]) -> Option<&CodeFenceBlock> {
338    let mut best = None;
339    let mut best_score = (0usize, 0u8);
340    for block in blocks {
341        let score = score_code_block(block);
342        if score > best_score {
343            best_score = score;
344            best = Some(block);
345        }
346    }
347    best
348}
349
350fn score_code_block(block: &CodeFenceBlock) -> (usize, u8) {
351    let line_count = block
352        .lines
353        .iter()
354        .filter(|line| !line.trim().is_empty())
355        .count();
356    let has_language = block
357        .language
358        .as_ref()
359        .is_some_and(|lang| !lang.trim().is_empty());
360    (line_count, if has_language { 1 } else { 0 })
361}
362
363#[derive(Debug, Clone, PartialEq, Eq)]
364struct CodeFenceBlock {
365    language: Option<String>,
366    lines: Vec<String>,
367}
vtcode_core/commands/ask.rs

vtcode_core/commands/
ask.rs
