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, sync::LazyLock};
 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, "## Excerpt around the cursor\n")?;
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 = 5))]
43    pub queries: Box<[SearchToolQuery]>,
44}
45
46#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
47pub struct SearchToolQuery {
48    /// A glob pattern to match file paths in the codebase
49    pub glob: String,
50    /// A regular expression to match content within the files matched by the glob pattern
51    pub regex: String,
52}
53
54pub static TOOL_SCHEMA: LazyLock<(serde_json::Value, String)> = LazyLock::new(|| {
55    let schema = schemars::schema_for!(SearchToolInput);
56
57    let description = schema
58        .get("description")
59        .and_then(|description| description.as_str())
60        .unwrap()
61        .to_string();
62
63    (schema.into(), description)
64});
65
66pub const TOOL_NAME: &str = "search";
67
68const SEARCH_INSTRUCTIONS: &str = indoc! {r#"
69    ## Task
70
71    You are part of an edit prediction system in a code editor. Your role is to identify relevant code locations
72    that will serve as context for predicting the next required edit.
73
74    **Your task:**
75    - Analyze the user's recent edits and current cursor context
76    - Use the `search` tool to find code that may be relevant for predicting the next edit
77    - Focus on finding:
78       - Code patterns that might need similar changes based on the recent edits
79       - Functions, variables, types, and constants referenced in the current cursor context
80       - Related implementations, usages, or dependencies that may require consistent updates
81
82    **Important constraints:**
83    - This conversation has exactly 2 turns
84    - You must make ALL search queries in your first response via the `search` tool
85    - All queries will be executed in parallel and results returned together
86    - In the second turn, you will select the most relevant results via the `select` tool.
87"#};
88
89const TOOL_USE_REMINDER: &str = indoc! {"
90    --
91    Use the `search` tool now
92"};