| 1 | //! FIM (Fill-in-the-Middle) edit tool. |
| 2 | //! |
| 3 | //! Reads a file, finds `prefix_anchor` and `suffix_anchor`, calls the |
| 4 | //! DeepSeek `/beta/completions` FIM endpoint, and writes the generated |
| 5 | //! middle content back into the file. |
| 6 | |
| 7 | use std::fs; |
| 8 | |
| 9 | use async_trait::async_trait; |
| 10 | use serde_json::{Value, json}; |
| 11 | use thiserror::Error; |
| 12 | |
| 13 | use crate::client::DeepSeekClient; |
| 14 | |
| 15 | use super::spec::{ |
| 16 | ApprovalRequirement, ToolCapability, ToolContext, ToolError, ToolResult, ToolSpec, |
| 17 | optional_u64, required_str, |
| 18 | }; |
| 19 | |
| 20 | /// Result of a FIM edit operation |
| 21 | #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] |
| 22 | pub struct FimEditResult { |
| 23 | pub success: bool, |
| 24 | pub path: String, |
| 25 | pub generated_text: String, |
| 26 | pub prefix_end: usize, |
| 27 | pub suffix_start: usize, |
| 28 | pub message: String, |
| 29 | } |
| 30 | |
| 31 | /// Tool for performing Fill-in-the-Middle edits via the DeepSeek FIM API. |
| 32 | pub struct FimEditTool { |
| 33 | pub client: Option<DeepSeekClient>, |
| 34 | pub model: String, |
| 35 | } |
| 36 | |
| 37 | impl FimEditTool { |
| 38 | #[must_use] |
| 39 | pub fn new(client: Option<DeepSeekClient>, model: String) -> Self { |
| 40 | Self { client, model } |
| 41 | } |
| 42 | } |
| 43 | |
| 44 | // === Errors === |
| 45 | |
| 46 | #[derive(Debug, Error)] |
| 47 | enum FimError { |
| 48 | #[error("Prefix anchor not found in file: '{0}'")] |
| 49 | PrefixNotFound(String), |
| 50 | #[error("Suffix anchor not found after prefix anchor: '{0}'")] |
| 51 | SuffixNotFound(String), |
| 52 | #[error("Prefix and suffix anchors overlap (suffix starts at {0}, prefix ends at {1})")] |
| 53 | AnchorsOverlap(usize, usize), |
| 54 | #[error("FIM API call failed: {0}")] |
| 55 | ApiFailed(String), |
| 56 | } |
| 57 | |
| 58 | #[async_trait] |
| 59 | impl ToolSpec for FimEditTool { |
| 60 | fn name(&self) -> &'static str { |
| 61 | "fim_edit" |
| 62 | } |
| 63 | |
| 64 | fn description(&self) -> &'static str { |
| 65 | "Edit a file using Fill-in-the-Middle (FIM) completion. Provide a file path, \ |
| 66 | prefix_anchor (text that appears before the section to replace), and \ |
| 67 | suffix_anchor (text that appears after the section to replace). The tool \ |
| 68 | calls DeepSeek's FIM endpoint to generate replacement content." |
| 69 | } |
| 70 | |
| 71 | fn input_schema(&self) -> Value { |
| 72 | json!({ |
| 73 | "type": "object", |
| 74 | "properties": { |
| 75 | "path": { |
| 76 | "type": "string", |
| 77 | "description": "Path to the file to edit (relative to workspace)" |
| 78 | }, |
| 79 | "prefix_anchor": { |
| 80 | "type": "string", |
| 81 | "description": "Text anchor marking the end of the prefix. Everything up to and including this anchor is kept as-is before the generated middle." |
| 82 | }, |
| 83 | "suffix_anchor": { |
| 84 | "type": "string", |
| 85 | "description": "Text anchor marking the start of the suffix. Everything from this anchor onward is kept as-is after the generated middle." |
| 86 | }, |
| 87 | "max_tokens": { |
| 88 | "type": "integer", |
| 89 | "description": "Maximum tokens to generate (default: 1024)" |
| 90 | } |
| 91 | }, |
| 92 | "required": ["path", "prefix_anchor", "suffix_anchor"] |
| 93 | }) |
| 94 | } |
| 95 | |
| 96 | fn capabilities(&self) -> Vec<ToolCapability> { |
| 97 | vec![ |
| 98 | ToolCapability::ReadOnly, |
| 99 | ToolCapability::WritesFiles, |
| 100 | ToolCapability::RequiresApproval, |
| 101 | ] |
| 102 | } |
| 103 | |
| 104 | fn approval_requirement(&self) -> ApprovalRequirement { |
| 105 | ApprovalRequirement::Suggest |
| 106 | } |
| 107 | |
| 108 | async fn execute(&self, input: Value, context: &ToolContext) -> Result<ToolResult, ToolError> { |
| 109 | let path = required_str(&input, "path")?; |
| 110 | let prefix_anchor = required_str(&input, "prefix_anchor")?; |
| 111 | let suffix_anchor = required_str(&input, "suffix_anchor")?; |
| 112 | let max_tokens = optional_u64(&input, "max_tokens", 1024); |
| 113 | |
| 114 | // 1. Read the file |
| 115 | let resolved = context.resolve_path(path)?; |
| 116 | let content = fs::read_to_string(&resolved).map_err(|e| { |
| 117 | ToolError::execution_failed(format!("Failed to read {}: {}", resolved.display(), e)) |
| 118 | })?; |
| 119 | |
| 120 | // 2. Find prefix anchor |
| 121 | let prefix_pos = content.find(prefix_anchor).ok_or_else(|| { |
| 122 | ToolError::execution_failed( |
| 123 | FimError::PrefixNotFound(prefix_anchor.to_string()).to_string(), |
| 124 | ) |
| 125 | })?; |
| 126 | let prefix_end = prefix_pos + prefix_anchor.len(); |
| 127 | |
| 128 | // 3. Find suffix anchor (after prefix anchor) |
| 129 | let suffix_pos = content[prefix_end..].find(suffix_anchor).ok_or_else(|| { |
| 130 | ToolError::execution_failed( |
| 131 | FimError::SuffixNotFound(suffix_anchor.to_string()).to_string(), |
| 132 | ) |
| 133 | })?; |
| 134 | let suffix_start = prefix_end + suffix_pos; |
| 135 | |
| 136 | // 4. Validate anchors don't overlap |
| 137 | if suffix_start < prefix_end { |
| 138 | return Err(ToolError::execution_failed( |
| 139 | FimError::AnchorsOverlap(suffix_start, prefix_end).to_string(), |
| 140 | )); |
| 141 | } |
| 142 | |
| 143 | // 5. Extract prefix and suffix for the FIM API |
| 144 | let fim_prompt = content[..prefix_end].to_string(); |
| 145 | let fim_suffix = content[suffix_start..].to_string(); |
| 146 | |
| 147 | // 6. Call FIM API |
| 148 | let generated_text = match self.client.as_ref() { |
| 149 | Some(client) => client |
| 150 | .fim_completion(&self.model, &fim_prompt, &fim_suffix, max_tokens as u32) |
| 151 | .await |
| 152 | .map_err(|e| { |
| 153 | ToolError::execution_failed(FimError::ApiFailed(e.to_string()).to_string()) |
| 154 | })?, |
| 155 | None => { |
| 156 | return Err(ToolError::execution_failed( |
| 157 | "FIM API client not available".to_string(), |
| 158 | )); |
| 159 | } |
| 160 | }; |
| 161 | |
| 162 | // 7. Build the new content and write it back |
| 163 | let generated_len = generated_text.len(); |
| 164 | let new_content = format!("{}{}{}", fim_prompt, generated_text, fim_suffix); |
| 165 | fs::write(&resolved, &new_content).map_err(|e| { |
| 166 | ToolError::execution_failed(format!("Failed to write {}: {}", resolved.display(), e)) |
| 167 | })?; |
| 168 | |
| 169 | let result = FimEditResult { |
| 170 | success: true, |
| 171 | path: path.to_string(), |
| 172 | generated_text, |
| 173 | prefix_end, |
| 174 | suffix_start, |
| 175 | message: format!( |
| 176 | "FIM edit applied to `{}`. Generated {} chars between prefix_anchor end (byte {}) and suffix_anchor start (byte {}).", |
| 177 | path, generated_len, prefix_end, suffix_start, |
| 178 | ), |
| 179 | }; |
| 180 | |
| 181 | ToolResult::json(&result).map_err(|e| ToolError::execution_failed(e.to_string())) |
| 182 | } |
| 183 | } |
| 184 |