retrieval_prompt.rs

 1use anyhow::Result;
 2use cloud_llm_client::predict_edits_v3::{self, Excerpt};
 3use indoc::indoc;
 4use schemars::JsonSchema;
 5use serde::{Deserialize, Serialize};
 6use std::fmt::Write;
 7
 8use crate::{push_events, write_codeblock};
 9
10pub fn build_prompt(request: predict_edits_v3::PlanContextRetrievalRequest) -> Result<String> {
11    let mut prompt = SEARCH_INSTRUCTIONS.to_string();
12
13    if !request.events.is_empty() {
14        writeln!(&mut prompt, "## User Edits\n")?;
15        push_events(&mut prompt, &request.events);
16    }
17
18    writeln!(&mut prompt, "## Cursor context")?;
19    write_codeblock(
20        &request.excerpt_path,
21        &[Excerpt {
22            start_line: request.excerpt_line_range.start,
23            text: request.excerpt.into(),
24        }],
25        &[],
26        request.cursor_file_max_row,
27        true,
28        &mut prompt,
29    );
30
31    writeln!(&mut prompt, "{TOOL_USE_REMINDER}")?;
32
33    Ok(prompt)
34}
35
36/// Search for relevant code
37///
38/// For the best results, run multiple queries at once with a single invocation of this tool.
39#[derive(Clone, Deserialize, Serialize, JsonSchema)]
40pub struct SearchToolInput {
41    /// An array of queries to run for gathering context relevant to the next prediction
42    #[schemars(length(max = 3))]
43    pub queries: Box<[SearchToolQuery]>,
44}
45
46/// Search for relevant code by path, syntax hierarchy, and content.
47#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
48pub struct SearchToolQuery {
49    /// 1. A glob pattern to match file paths in the codebase to search in.
50    pub glob: String,
51    /// 2. Regular expressions to match syntax nodes **by their first line** and hierarchy.
52    ///
53    /// Subsequent regexes match nodes within the full content of the nodes matched by the previous regexes.
54    ///
55    /// Example: Searching for a `User` class
56    ///     ["class\s+User"]
57    ///
58    /// Example: Searching for a `get_full_name` method under a `User` class
59    ///     ["class\s+User", "def\sget_full_name"]
60    ///
61    /// Skip this field to match on content alone.
62    #[schemars(length(max = 3))]
63    #[serde(default)]
64    pub syntax_node: Vec<String>,
65    /// 3. An optional regular expression to match the final content that should appear in the results.
66    ///
67    /// - Content will be matched within all lines of the matched syntax nodes.
68    /// - If syntax node regexes are provided, this field can be skipped to include as much of the node itself as possible.
69    /// - If no syntax node regexes are provided, the content will be matched within the entire file.
70    pub content: Option<String>,
71}
72
73pub const TOOL_NAME: &str = "search";
74
75const SEARCH_INSTRUCTIONS: &str = indoc! {r#"
76    You are part of an edit prediction system in a code editor.
77    Your role is to search for code that will serve as context for predicting the next edit.
78
79    - Analyze the user's recent edits and current cursor context
80    - Use the `search` tool to find code that is relevant for predicting the next edit
81    - Focus on finding:
82       - Code patterns that might need similar changes based on the recent edits
83       - Functions, variables, types, and constants referenced in the current cursor context
84       - Related implementations, usages, or dependencies that may require consistent updates
85       - How items defined in the cursor excerpt are used or altered
86    - You will not be able to filter results or perform subsequent queries, so keep searches as targeted as possible
87    - Use `syntax_node` parameter whenever you're looking for a particular type, class, or function
88    - Avoid using wildcard globs if you already know the file path of the content you're looking for
89"#};
90
91const TOOL_USE_REMINDER: &str = indoc! {"
92    --
93    Analyze the user's intent in one to two sentences, then call the `search` tool.
94"};