opendev_repl/
query_enhancer.rs

1//! Query enhancement and message preparation for the REPL.
2//!
3//! Mirrors `opendev/repl/query_enhancer.py`.
4//!
5//! Responsibilities:
6//! - Strip `@` references from queries while injecting file contents
7//! - Prepare the full message list for LLM API calls (system prompt,
8//!   session history, multimodal content, playbook context)
9
10use regex::Regex;
11use serde_json::Value;
12use std::path::PathBuf;
13use tracing::warn;
14
15use crate::file_injector::{FileContentInjector, ImageBlock};
16
17// ---------------------------------------------------------------------------
18// Constants
19// ---------------------------------------------------------------------------
20
21/// Default thinking-on instruction text (injected when thinking is visible).
22const THINKING_ON_INSTRUCTION: &str = "Use your thinking/reasoning capabilities to work through complex problems step by step. \
23     Show your reasoning process.";
24
25/// Default thinking-off instruction text (injected when thinking is hidden).
26const THINKING_OFF_INSTRUCTION: &str =
27    "Proceed directly with your response without showing internal reasoning.";
28
29// ---------------------------------------------------------------------------
30// QueryEnhancer
31// ---------------------------------------------------------------------------
32
33/// Handles query enhancement (@ file injection) and message preparation.
34pub struct QueryEnhancer {
35    /// Working directory for resolving relative `@` paths.
36    working_dir: PathBuf,
37}
38
39impl QueryEnhancer {
40    /// Create a new enhancer rooted at `working_dir`.
41    pub fn new(working_dir: PathBuf) -> Self {
42        Self { working_dir }
43    }
44
45    /// Enhance a query by injecting `@`-referenced file contents.
46    ///
47    /// Returns `(enhanced_query, image_blocks)` where:
48    /// - `enhanced_query` has `@` markers stripped and file content appended
49    /// - `image_blocks` contains base64-encoded images for multimodal calls
50    pub fn enhance_query(&self, query: &str) -> (String, Vec<ImageBlock>) {
51        let injector = FileContentInjector::new(self.working_dir.clone());
52        let result = injector.inject_content(query);
53
54        // Strip @ references from the query text.
55        // Pattern 1: Quoted paths @"path with spaces"
56        let quoted_re = Regex::new(r#"@"([^"]+)""#).expect("valid regex");
57        let enhanced = quoted_re.replace_all(query, "$1").to_string();
58
59        // Pattern 2: Unquoted paths (but not emails like user@example.com)
60        let unquoted_re = Regex::new(r"(?:^|\s)@([a-zA-Z0-9_./\-]+)").expect("valid regex");
61        let enhanced = unquoted_re
62            .replace_all(&enhanced, |caps: &regex::Captures| {
63                // Preserve the leading whitespace (or start-of-string) that was matched
64                let full = caps.get(0).unwrap().as_str();
65                let path = &caps[1];
66                if full.starts_with(char::is_whitespace) {
67                    format!("{}{}", &full[..full.len() - path.len() - 1], path)
68                } else {
69                    path.to_string()
70                }
71            })
72            .to_string();
73
74        // Append injected text content if any
75        let enhanced = if result.text_content.is_empty() {
76            enhanced
77        } else {
78            format!("{}\n\n{}", enhanced, result.text_content)
79        };
80
81        (enhanced, result.image_blocks)
82    }
83
84    /// Prepare the full message list for an LLM API call.
85    ///
86    /// # Arguments
87    ///
88    /// * `query` - Original user query (before enhancement)
89    /// * `enhanced_query` - Query after `@` processing
90    /// * `system_prompt` - Base system prompt text
91    /// * `session_messages` - Existing conversation messages (if any)
92    /// * `image_blocks` - Multimodal image blocks from enhancement
93    /// * `thinking_visible` - Whether thinking mode is visible to the user
94    /// * `playbook_context` - Optional learned-strategies text to append
95    ///
96    /// # Returns
97    ///
98    /// A `Vec<Value>` of message objects ready for the LLM API.
99    #[allow(clippy::too_many_arguments)]
100    pub fn prepare_messages(
101        &self,
102        query: &str,
103        enhanced_query: &str,
104        system_prompt: &str,
105        session_messages: Option<&[Value]>,
106        image_blocks: &[ImageBlock],
107        thinking_visible: bool,
108        playbook_context: Option<&str>,
109    ) -> Vec<Value> {
110        // Start with session messages or empty vec
111        let mut messages: Vec<Value> = match session_messages {
112            Some(msgs) => msgs.to_vec(),
113            None => Vec::new(),
114        };
115
116        // If the query was enhanced, replace the last user message content
117        if enhanced_query != query {
118            for msg in messages.iter_mut().rev() {
119                if msg.get("role").and_then(|r| r.as_str()) == Some("user") {
120                    msg["content"] = Value::String(enhanced_query.to_string());
121                    break;
122                }
123            }
124        }
125
126        // Build final system content
127        let mut system_content = system_prompt.to_string();
128
129        // Replace {thinking_instruction} placeholder
130        if system_content.contains("{thinking_instruction}") {
131            let thinking_text = if thinking_visible {
132                THINKING_ON_INSTRUCTION
133            } else {
134                THINKING_OFF_INSTRUCTION
135            };
136            system_content = system_content.replace("{thinking_instruction}", thinking_text);
137        }
138
139        // Append playbook context if present
140        if let Some(playbook) = playbook_context
141            && !playbook.is_empty()
142        {
143            system_content = format!(
144                "{}\n\n## Learned Strategies\n{}",
145                system_content.trim_end(),
146                playbook
147            );
148        }
149
150        // Insert or update system message at position 0
151        if messages.is_empty() || messages[0].get("role").and_then(|r| r.as_str()) != Some("system")
152        {
153            messages.insert(
154                0,
155                serde_json::json!({
156                    "role": "system",
157                    "content": system_content,
158                }),
159            );
160        } else {
161            messages[0]["content"] = Value::String(system_content);
162        }
163
164        // Handle multimodal content (images)
165        if !image_blocks.is_empty() {
166            for msg in messages.iter_mut().rev() {
167                if msg.get("role").and_then(|r| r.as_str()) == Some("user") {
168                    let current_content = msg
169                        .get("content")
170                        .and_then(|c| c.as_str())
171                        .unwrap_or("")
172                        .to_string();
173
174                    let mut multimodal: Vec<Value> = vec![serde_json::json!({
175                        "type": "text",
176                        "text": current_content,
177                    })];
178
179                    for block in image_blocks {
180                        multimodal.push(serde_json::json!({
181                            "type": "image",
182                            "source": {
183                                "type": "base64",
184                                "media_type": block.media_type,
185                                "data": block.data,
186                            }
187                        }));
188                    }
189
190                    msg["content"] = Value::Array(multimodal);
191                    break;
192                }
193            }
194        }
195
196        // Estimate tokens and warn if large
197        let total_chars: usize = messages
198            .iter()
199            .map(|m| {
200                m.get("content")
201                    .map(|c| match c {
202                        Value::String(s) => s.len(),
203                        other => other.to_string().len(),
204                    })
205                    .unwrap_or(0)
206            })
207            .sum();
208        let estimated_tokens = total_chars / 4;
209        if estimated_tokens > 100_000 {
210            warn!(
211                messages = messages.len(),
212                estimated_tokens, "Large context detected"
213            );
214        }
215
216        messages
217    }
218
219    /// Format a debug summary of a message list.
220    pub fn format_messages_summary(messages: &[Value], max_preview: usize) -> String {
221        if messages.is_empty() {
222            return "0 messages".to_string();
223        }
224
225        let mut summary_parts = Vec::new();
226        for msg in messages {
227            let role = msg
228                .get("role")
229                .and_then(|r| r.as_str())
230                .unwrap_or("unknown");
231            let content = msg.get("content");
232
233            let preview = match content {
234                Some(Value::String(s)) => {
235                    if s.len() > max_preview {
236                        format!("{}...", &s[..max_preview])
237                    } else {
238                        s.clone()
239                    }
240                }
241                Some(Value::Array(arr)) => {
242                    format!("[{} blocks]", arr.len())
243                }
244                Some(other) => {
245                    let s = other.to_string();
246                    if s.len() > max_preview {
247                        format!("{}...", &s[..max_preview])
248                    } else {
249                        s
250                    }
251                }
252                None => String::new(),
253            };
254
255            summary_parts.push(format!("{}: {}", role, preview));
256        }
257
258        format!("{} messages: {}", messages.len(), summary_parts.join(" | "))
259    }
260}
261
262// ---------------------------------------------------------------------------
263// Tests
264// ---------------------------------------------------------------------------
265
266#[cfg(test)]
267#[path = "query_enhancer_tests.rs"]
268mod tests;
opendev_repl/query_enhancer.rs

opendev_repl/
query_enhancer.rs
