perf: add all file imports to LLM prompt

xeol-io · May 6, 2024 · b2ee156 · b2ee156
1 parent b235592
commit b2ee156
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 51 deletions.
diff --git a/packages/bumpgen-core/src/index.ts b/packages/bumpgen-core/src/index.ts
@@ -158,16 +158,15 @@ const bumpgen = ({
         >();
 
         for (const err of errs) {
-          const affectedNodes =
-            graphService.dependency.getNodesInFileWithinRange(
-              dependencyGraph,
-              projectRoot,
-              {
-                filePath: err.path,
-                startLine: err.line,
-                endLine: err.line,
-              },
-            );
+          const affectedNodes = graphService.dependency.getNodes(
+            dependencyGraph,
+            projectRoot,
+            {
+              filePath: err.path,
+              startLine: err.line,
+              endLine: err.line,
+            },
+          );
 
           if (affectedNodes.length === 0) {
             console.debug(
@@ -276,11 +275,12 @@ const bumpgen = ({
             },
           );
 
-          const importContext = graphService.dependency.getReferencingNodes(
+          const importContext = graphService.dependency.getNodes(
             graph.dependency,
+            args.projectRoot,
             {
-              id: planNode.id,
-              relationships: ["importDeclaration"],
+              filePath: depGraphNode.path,
+              kinds: ["ImportDeclaration", "ImportSpecifier"],
             },
           );
 

diff --git a/packages/bumpgen-core/src/services/graph/dependency.ts b/packages/bumpgen-core/src/services/graph/dependency.ts
@@ -2,7 +2,7 @@ import path from "path";
 import { unique } from "radash";
 
 import type { DependencyGraph } from "../../models/graph";
-import type { Relationship } from "../../models/graph/dependency";
+import type { Kind, Relationship } from "../../models/graph/dependency";
 
 export const createDependencyGraphService = () => {
   return {
@@ -45,38 +45,37 @@ export const createDependencyGraphService = () => {
     getNodeById: (graph: DependencyGraph, { id }: { id: string }) => {
       return graph.getNodeAttributes(id);
     },
-    getNodes: (graph: DependencyGraph) => {
-      return graph.nodes().map((node) => graph.getNodeAttributes(node));
-    },
-    getNodesInFile: (
-      graph: DependencyGraph,
-      { filePath }: { filePath: string },
-    ) => {
-      return graph
-        .filterNodes((_, attrs) => {
-          return attrs.path === filePath;
-        })
-        .map((node) => graph.getNodeAttributes(node));
-    },
-    getNodesInFileWithinRange: (
+    getNodes: (
       graph: DependencyGraph,
       projectRoot: string,
       {
         filePath,
         startLine,
         endLine,
-      }: { filePath: string; startLine: number; endLine: number },
+        kinds,
+      }: {
+        filePath?: string;
+        startLine?: number;
+        endLine?: number;
+        kinds?: Kind[];
+      },
     ) => {
-      const fullPath = path.join(projectRoot, filePath);
+      const fullPath = filePath
+        ? filePath.startsWith(projectRoot)
+          ? filePath
+          : path.join(projectRoot, filePath)
+        : null;
       return graph
         .filterNodes((_, attrs) => {
           return (
-            attrs.path === fullPath &&
-            attrs.startLine <= startLine &&
-            attrs.endLine >= endLine
+            (fullPath ? attrs.path === fullPath : true) &&
+            (startLine ? attrs.startLine <= startLine : true) &&
+            (endLine ? attrs.endLine >= endLine : true) &&
+            (kinds ? kinds.includes(attrs.kind) : true)
           );
         })
-        .map((node) => graph.getNodeAttributes(node));
+        .map((node) => graph.getNodeAttributes(node))
+        .sort((a, b) => a.startLine - b.startLine);
     },
   };
 };

diff --git a/packages/bumpgen-core/src/services/language/typescript/index.ts b/packages/bumpgen-core/src/services/language/typescript/index.ts
@@ -156,7 +156,6 @@ export const makeTypescriptService = (
 
         return {
           source: "ts-morph",
-
           tree: project,
         };
       },

diff --git a/packages/bumpgen-core/src/services/llm/openai.ts b/packages/bumpgen-core/src/services/llm/openai.ts
@@ -1,7 +1,5 @@
 import type { OpenAI } from "openai";
-import { unique } from "radash";
 
-// import type { ContextSearchResponse } from "../../clients/sourcegraph/responses";
 import type { DependencyGraphNode } from "../../models/graph/dependency";
 import type { PlanGraphNode } from "../../models/graph/plan";
 import type { LLMContext } from "../../models/llm";
@@ -20,16 +18,14 @@ const makePlanNodeMessage = (
   importContext: DependencyGraphNode[],
   bumpedPackage: string,
 ) => {
-  const importMessages = unique(
-    importContext.map((context) => {
-      return context.block;
-    }),
-  );
+  const importMessages = importContext.map((context) => {
+    return context.block;
+  });
   return {
     role: "user" as const,
     content: [
       `I'm upgrading the package '${bumpedPackage}' and my code is failing. You might need to modify the code or the imports. Look at the errors below and think step-by-step about what the errors mean and how to fix the code.\n`,
-      `<relevant_imports path=${planNode.path}>\n${importMessages.join("\n")}\n</relevant_imports>`,
+      `<file_imports path=${planNode.path}>\n${importMessages.join("\n")}\n</file_imports>`,
       `<code \n  path="${planNode.path}"\n>`,
       `${planNode.block}`,
       "</code>\n",
@@ -47,17 +43,17 @@ const makePlanNodeMessage = (
 
 const makeExternalDependencyContextMessage = (
   pkg: string,
-  importContext: (DependencyGraphNode & {
+  externalImportContext: (DependencyGraphNode & {
     typeSignature: string;
   })[],
 ) => {
-  const typeSignatures = importContext
+  const typeSignatures = externalImportContext
     .filter((imp) => imp.typeSignature !== "")
     .map((imp) => {
       return `<import \n  statement="${imp.block}"\n>\n${imp.typeSignature}\n</import>`;
     });
 
-  const exports = importContext
+  const exports = externalImportContext
     .filter(
       (
         imp,
@@ -67,9 +63,7 @@ const makeExternalDependencyContextMessage = (
       } => !!imp.external,
     )
     .flatMap((imp) => {
-      return imp.external.exports.map((exp) => {
-        return exp;
-      });
+      return imp.external.exports;
     });
 
   if (typeSignatures.length === 0 && exports.length === 0) {
@@ -235,7 +229,7 @@ export const createOpenAIService = (openai: OpenAI) => {
             "- Maintain all hardcoded values as is.",
             "- Avoid adding comments within the code.",
             "- Refrain from using explicit type casting.",
-            "- Only show the specific lines of code that have been changed or need modification, without including unchanged surrounding code.",
+            "- Only show the specific lines of code that have been changed or need modification, without including unchanged surrounding code except to disambiguate the edited line.",
             "- Keep all existing variable, function, and class names unchanged.",
           ].join("\n"),
         };
@@ -290,7 +284,7 @@ export const createOpenAIService = (openai: OpenAI) => {
                           oldCode: {
                             type: "string",
                             description:
-                              "The old lines of code. Be sure to add lines before and after to disambiguate the change.",
+                              "The old lines of code. If needed, add lines before and after to disambiguate the change from other lines.",
                           },
                           newCode: {
                             type: "string",
-Original file line number
+Diff line change
@@ Expand Up / @@ -156,7 +156,6 @@ export const makeTypescriptService = ( @@
             return {
               source: "ts-morph",
               tree: project,
             };
           },
@@ Expand Down @@