Initial progress on automatic application and diagnostics

Author: camelid

accelerant/chat.py

@@ -10,9 +10,17 @@
     ChatCompletionSystemMessageParam,
     ParsedChatCompletion,
 )
-from pydantic import BaseModel
 from rich import print as rprint
 
+from accelerant.chat_interface import (
+    ProjectAnalysis,
+    OptimizationSuite,
+)
+from accelerant.lsp import (
+    request_document_diagnostics,
+    syncexec,
+)
+from accelerant.patch import apply_simultaneous_suggestions
 from accelerant.perf import PerfData
 from accelerant.project import Project
 from accelerant.tools import (
@@ -24,28 +32,6 @@
 from perfparser import LineLoc
 
 
-class RegionAnalysis(BaseModel):
-    filename: str
-    line: int
-    performanceAnalysis: str
-
-
-class ProjectAnalysis(BaseModel):
-    regions: List[RegionAnalysis]
-
-
-class OptimizationSuggestion(BaseModel):
-    filename: str
-    startLine: int
-    endLine: int
-    newCode: str
-
-
-class OptimizationSuite(BaseModel):
-    highLevelSummary: str
-    suggestions: List[OptimizationSuggestion]
-
-
 def optimize_locations(
     project: Project, locs: List[LineLoc], perf_data: Optional[PerfData], model_id: str
 ) -> str:
@@ -117,9 +103,22 @@ def optimize_locations(
         ),
         {"role": "user", "content": sugg_req_msg},
     ]
-    sugg_str, _ = run_chat(
+    sugg_str, opt_suite = run_chat(
         messages, client, model_id, tool_runner, response_format=OptimizationSuite
     )
+    assert opt_suite
+    apply_simultaneous_suggestions(opt_suite.suggestions, project)
+    # TODO: collect diagnostics from all files, then feed them back to LLM
+    # FIXME: how to get diagnostics from files that depend on these files
+    # but weren't themselves changed?
+    print(
+        syncexec(
+            project.lsp(),
+            request_document_diagnostics(
+                project.lsp().language_server, opt_suite.suggestions[0].filename
+            ),
+        )
+    )
     return sugg_str
 
 

accelerant/chat_interface.py

@@ -0,0 +1,24 @@
+from typing import List
+from pydantic import BaseModel
+
+
+class RegionAnalysis(BaseModel):
+    filename: str
+    line: int
+    performanceAnalysis: str
+
+
+class ProjectAnalysis(BaseModel):
+    regions: List[RegionAnalysis]
+
+
+class OptimizationSuggestion(BaseModel):
+    filename: str
+    startLine: int
+    endLine: int
+    newCode: str
+
+
+class OptimizationSuite(BaseModel):
+    highLevelSummary: str
+    suggestions: List[OptimizationSuggestion]

accelerant/lsp.py

@@ -5,20 +5,15 @@
 from multilspy import LanguageServer, SyncLanguageServer
 from multilspy.multilspy_utils import PathUtils
 
-from accelerant.project import Project
-
 
 async def request_definition_full(
     lsp: LanguageServer, relative_file_path: str, line: int, column: int
 ):
     with lsp.open_file(relative_file_path):
+        uri = relpath_to_uri(relative_file_path, lsp.repository_root_path)
         response = await lsp.server.send.definition(
             {
-                "textDocument": {
-                    "uri": Path(
-                        str(PurePath(lsp.repository_root_path, relative_file_path))
-                    ).as_uri()
-                },
+                "textDocument": {"uri": uri},
                 "position": {
                     "line": line,
                     "character": column,
@@ -28,9 +23,22 @@ async def request_definition_full(
         return response
 
 
-def extract_relative_path(uri: str, project: Project) -> str:
+async def request_document_diagnostics(lsp: LanguageServer, relpath: str):
+    with lsp.open_file(relpath):
+        uri = relpath_to_uri(relpath, lsp.repository_root_path)
+        lsp.server.notify.did_save_text_document({"textDocument": {"uri": uri}})
+        return await lsp.server.send.text_document_diagnostic(
+            {"textDocument": {"uri": uri}}
+        )
+
+
+def relpath_to_uri(relpath: str, root: str) -> str:
+    return Path(str(PurePath(root, relpath))).as_uri()
+
+
+def uri_to_relpath(uri: str, root: str) -> str:
     absolutePath = PathUtils.uri_to_path(uri)
-    relativePath = str(PurePath(os.path.relpath(absolutePath, project._root)))
+    relativePath = str(PurePath(os.path.relpath(absolutePath, root)))
     return relativePath
 
 

accelerant/patch.py

@@ -0,0 +1,45 @@
+from pathlib import PurePath
+from typing import List
+
+from accelerant.chat_interface import OptimizationSuggestion
+from accelerant.project import Project
+
+
+def apply_simultaneous_suggestions(
+    suggestions: List[OptimizationSuggestion],
+    project: Project,
+):
+    """
+    Apply suggestions that were created to be applied simultaneously.
+    In other words, they don't account for changes in line numbers caused by
+    other suggestions.
+    """
+    by_file: dict[str, List[OptimizationSuggestion]] = {}
+    for sugg in suggestions:
+        if sugg.filename not in by_file:
+            by_file[sugg.filename] = []
+        by_file[sugg.filename].append(sugg)
+
+    for relpath, suggs in by_file.items():
+        abspath = PurePath(project._root, relpath)
+        # FIXME: handle conflicting suggestions (i.e. that overlap)
+        suggs = list(sorted(suggs, key=lambda s: s.startLine))
+        with open(abspath, "r") as f:
+            old_lines = enumerate(f.readlines(), start=1)
+
+        print(abspath, suggs)
+        with open(abspath, "w") as f:
+            skip_until_after = 0
+            for old_num, old_line in old_lines:
+                if old_num <= skip_until_after:
+                    continue
+                if suggs and suggs[0].startLine == old_num:
+                    print("apply")
+                    sugg = suggs[0]
+                    suggs = suggs[1:]
+                    f.write(sugg.newCode)
+                    if not sugg.newCode.endswith("\n"):
+                        f.write("\n")
+                    skip_until_after = sugg.endLine
+                else:
+                    f.write(old_line)

accelerant/tools.py

@@ -5,7 +5,7 @@
 import openai
 from openai.types.chat import ChatCompletionToolParam
 
-from accelerant.lsp import request_definition_full, syncexec, extract_relative_path
+from accelerant.lsp import request_definition_full, syncexec, uri_to_relpath
 from accelerant.util import find_symbol, truncate_for_llm
 from accelerant.project import Project
 
@@ -109,7 +109,7 @@ def convert_lsp_loc(r: dict, p: Project) -> dict:
     if "relativePath" in r:
         filename = r["relativePath"]
     else:
-        filename = extract_relative_path(r["targetUri"], p)
+        filename = uri_to_relpath(r["targetUri"], str(p._root))
     if "targetRange" in r:
         range = r["targetRange"]
     else: