feat: Relace edit tool (#647)

# Motivation

<!-- Why is this change necessary? -->

# Content

<!-- Please include a summary of the change -->

# Testing

<!-- How was the change tested? -->

# Please check the following before marking your PR as ready for review

- [ ] I have added tests for my changes
- [ ] I have updated the documentation or added new documentation as
needed

---------

Co-authored-by: kopekC <[email protected]>

Author: kopekC

src/codegen/extensions/langchain/agent.py

@@ -13,14 +13,14 @@
 from .tools import (
     CreateFileTool,
     DeleteFileTool,
-    EditFileTool,
     ListDirectoryTool,
     MoveSymbolTool,
+    RelaceEditTool,
     RenameFileTool,
     ReplacementEditTool,
     RevealSymbolTool,
     SearchTool,
-    SemanticEditTool,
+    # SemanticEditTool,
     ViewFileTool,
 )
 
@@ -62,14 +62,15 @@ def create_codebase_agent(
         ViewFileTool(codebase),
         ListDirectoryTool(codebase),
         SearchTool(codebase),
-        EditFileTool(codebase),
+        # EditFileTool(codebase),
         CreateFileTool(codebase),
         DeleteFileTool(codebase),
         RenameFileTool(codebase),
         MoveSymbolTool(codebase),
         RevealSymbolTool(codebase),
-        SemanticEditTool(codebase),
+        # SemanticEditTool(codebase),
         ReplacementEditTool(codebase),
+        RelaceEditTool(codebase),
         # SemanticSearchTool(codebase),
         # =====[ Github Integration ]=====
         # Enable Github integration

src/codegen/extensions/langchain/tools.py

@@ -16,6 +16,7 @@
     linear_search_issues_tool,
 )
 from codegen.extensions.tools.link_annotation import add_links_to_message
+from codegen.extensions.tools.relace_edit import relace_edit
 from codegen.extensions.tools.replacement_edit import replacement_edit
 from codegen.extensions.tools.reveal_symbol import reveal_symbol
 from codegen.extensions.tools.search import search
@@ -37,6 +38,7 @@
     view_file,
     view_pr,
 )
+from ..tools.relace_edit_prompts import RELACE_EDIT_PROMPT
 from ..tools.semantic_edit_prompts import FILE_EDIT_PROMPT
 
 
@@ -729,9 +731,10 @@ def get_workspace_tools(codebase: Codebase) -> list["BaseTool"]:
         RevealSymbolTool(codebase),
         RunBashCommandTool(),  # Note: This tool doesn't need the codebase
         SearchTool(codebase),
-        SemanticEditTool(codebase),
+        # SemanticEditTool(codebase),
         SemanticSearchTool(codebase),
         ViewFileTool(codebase),
+        RelaceEditTool(codebase),
         # Github
         GithubCreatePRTool(codebase),
         GithubCreatePRCommentTool(codebase),
@@ -788,3 +791,47 @@ def _run(
             count=count,
         )
         return result.render()
+
+
+# Brief description for the Relace Edit tool
+_RELACE_EDIT_BRIEF = """Tool for file editing using the Relace Instant Apply API.
+This high-speed code generation engine optimizes for real-time performance at 2000 tokens/second.
+
+Provide an edit snippet that describes the changes you want to make, with helpful comments to indicate unchanged sections, like so:
+```
+// ... keep existing imports ...
+
+// Add new function
+function calculateDiscount(price, discountPercent) {
+  return price * (discountPercent / 100);
+}
+
+// ... keep existing code ...
+```
+
+The API will merge your edit snippet with the existing code to produce the final result.
+The API key will be automatically retrieved from the RELACE_API environment variable.
+"""
+
+
+class RelaceEditInput(BaseModel):
+    """Input for Relace editing."""
+
+    filepath: str = Field(..., description="Path of the file relative to workspace root")
+    edit_snippet: str = Field(..., description=RELACE_EDIT_PROMPT)
+
+
+class RelaceEditTool(BaseTool):
+    """Tool for editing files using the Relace Instant Apply API."""
+
+    name: ClassVar[str] = "relace_edit"
+    description: ClassVar[str] = _RELACE_EDIT_BRIEF
+    args_schema: ClassVar[type[BaseModel]] = RelaceEditInput
+    codebase: Codebase = Field(exclude=True)
+
+    def __init__(self, codebase: Codebase) -> None:
+        super().__init__(codebase=codebase)
+
+    def _run(self, filepath: str, edit_snippet: str) -> str:
+        result = relace_edit(self.codebase, filepath, edit_snippet)
+        return result.render()

src/codegen/extensions/tools/relace_edit.py

@@ -0,0 +1,163 @@
+"""Tool for making edits to files using the Relace Instant Apply API."""
+
+import difflib
+import os
+from typing import ClassVar, Optional
+
+import requests
+from pydantic import Field
+
+from codegen.sdk.core.codebase import Codebase
+
+from .observation import Observation
+from .view_file import add_line_numbers
+
+
+class RelaceEditObservation(Observation):
+    """Response from making edits to a file using Relace Instant Apply API."""
+
+    filepath: str = Field(
+        description="Path to the edited file",
+    )
+    diff: Optional[str] = Field(
+        default=None,
+        description="Unified diff showing the changes made",
+    )
+    new_content: Optional[str] = Field(
+        default=None,
+        description="New content with line numbers",
+    )
+    line_count: Optional[int] = Field(
+        default=None,
+        description="Total number of lines in file",
+    )
+
+    str_template: ClassVar[str] = "Edited file {filepath} using Relace Instant Apply"
+
+
+def generate_diff(original: str, modified: str) -> str:
+    """Generate a unified diff between two strings.
+
+    Args:
+        original: Original content
+        modified: Modified content
+
+    Returns:
+        Unified diff as a string
+    """
+    original_lines = original.splitlines(keepends=True)
+    modified_lines = modified.splitlines(keepends=True)
+
+    diff = difflib.unified_diff(
+        original_lines,
+        modified_lines,
+        fromfile="original",
+        tofile="modified",
+        lineterm="",
+    )
+
+    return "".join(diff)
+
+
+def get_relace_api_key() -> str:
+    """Get the Relace API key from environment variables.
+
+    Returns:
+        The Relace API key
+
+    Raises:
+        ValueError: If the API key is not found
+    """
+    api_key = os.environ.get("RELACE_API")
+    if not api_key:
+        msg = "RELACE_API environment variable not found. Please set it in your .env file."
+        raise ValueError(msg)
+    return api_key
+
+
+def apply_relace_edit(api_key: str, initial_code: str, edit_snippet: str, stream: bool = False) -> str:
+    """Apply an edit using the Relace Instant Apply API.
+
+    Args:
+        api_key: Relace API key
+        initial_code: The existing code to modify
+        edit_snippet: The edit snippet containing the modifications
+        stream: Whether to enable streaming response
+
+    Returns:
+        The merged code
+
+    Raises:
+        Exception: If the API request fails
+    """
+    url = "https://instantapply.endpoint.relace.run/v1/code/apply"
+    headers = {"Content-Type": "application/json", "Authorization": f"Bearer {api_key}"}
+
+    data = {"initialCode": initial_code, "editSnippet": edit_snippet, "stream": stream}
+
+    try:
+        response = requests.post(url, headers=headers, json=data)
+        response.raise_for_status()
+        return response.json()["mergedCode"]
+    except Exception as e:
+        msg = f"Relace API request failed: {e!s}"
+        raise Exception(msg)
+
+
+def relace_edit(codebase: Codebase, filepath: str, edit_snippet: str, api_key: Optional[str] = None) -> RelaceEditObservation:
+    """Edit a file using the Relace Instant Apply API.
+
+    Args:
+        codebase: Codebase object
+        filepath: Path to the file to edit
+        edit_snippet: The edit snippet containing the modifications
+        api_key: Optional Relace API key. If not provided, will be retrieved from environment variables.
+
+    Returns:
+        RelaceEditObservation with the results
+    """
+    try:
+        file = codebase.get_file(filepath)
+    except ValueError:
+        msg = f"File not found: {filepath}"
+        raise FileNotFoundError(msg)
+
+    # Get the original content
+    original_content = file.content
+    original_lines = original_content.split("\n")
+
+    # Get API key if not provided
+    if api_key is None:
+        try:
+            api_key = get_relace_api_key()
+        except ValueError as e:
+            return RelaceEditObservation(
+                status="error",
+                error=str(e),
+                filepath=filepath,
+            )
+
+    # Apply the edit using Relace API
+    try:
+        merged_code = apply_relace_edit(api_key, original_content, edit_snippet)
+    except Exception as e:
+        return RelaceEditObservation(
+            status="error",
+            error=str(e),
+            filepath=filepath,
+        )
+
+    # Generate diff
+    diff = generate_diff(original_content, merged_code)
+
+    # Apply the edit to the file
+    file.edit(merged_code)
+    codebase.commit()
+
+    return RelaceEditObservation(
+        status="success",
+        filepath=filepath,
+        diff=diff,
+        new_content=add_line_numbers(merged_code),
+        line_count=len(merged_code.split("\n")),
+    )

src/codegen/extensions/tools/relace_edit_prompts.py

@@ -0,0 +1,56 @@
+"""Prompts for the Relace edit tool."""
+
+RELACE_EDIT_PROMPT = """Edit a file using the Relace Instant Apply API.
+
+The Relace Instant Apply API is a high-speed code generation engine optimized for real-time performance at 2000 tokens/second. It splits code generation into two specialized steps:
+
+1. Hard Reasoning: Uses SOTA models like Claude for complex code understanding
+2. Fast Integration: Rapidly merges edits into existing code
+
+To use this tool, provide:
+1. The path to the file you want to edit
+2. An edit snippet that describes the changes you want to make
+
+The edit snippet should:
+- Include complete code blocks that will appear in the final output
+- Clearly indicate which parts of the code remain unchanged with comments like "// ... rest of code ..."
+- Maintain correct indentation and code structure
+
+Example edit snippet:
+```
+// ... keep existing imports ...
+
+// Add new function
+function calculateDiscount(price, discountPercent) {
+  return price * (discountPercent / 100);
+}
+
+// ... keep existing code ...
+```
+
+The API will merge your edit snippet with the existing code to produce the final result.
+"""
+
+RELACE_EDIT_SYSTEM_PROMPT = """You are an expert at creating edit snippets for the Relace Instant Apply API.
+
+Your job is to create an edit snippet that describes how to modify the provided existing code according to user specifications.
+
+Follow these guidelines:
+1. Focus only on the MODIFICATION REQUEST, not other aspects of the code
+2. Abbreviate unchanged sections with "// ... rest of headers/sections/code ..." (be descriptive in the comment)
+3. Indicate the location and nature of modifications with comments and ellipses
+4. Preserve indentation and code structure exactly as it should appear in the final code
+5. Do not output lines that will not be in the final code after merging
+6. If removing a section, provide relevant context so it's clear what should be removed
+
+Do NOT provide commentary or explanations - only the code with a focus on the modifications.
+"""
+
+RELACE_EDIT_USER_PROMPT = """EXISTING CODE:
+{initial_code}
+
+MODIFICATION REQUEST:
+{user_instructions}
+
+Create an edit snippet that can be used with the Relace Instant Apply API to implement these changes.
+"""