feat: Incorporate source code

Author: kgilpin

packages/cli/src/cmds/ask/ask.ts

@@ -1,25 +1,15 @@
 import { warn } from 'console';
+import OpenAI from 'openai';
+import { ChatCompletionMessageParam } from 'openai/resources';
+import { readFile } from 'fs/promises';
+import { join } from 'path';
+import { AppMapFilter, Event, Metadata, buildAppMap } from '@appland/models';
+import { Action, Specification, buildDiagram, nodeName } from '@appland/sequence-diagram';
+
 import { handleWorkingDirectory } from '../../lib/handleWorkingDirectory';
 import { locateAppMapDir } from '../../lib/locateAppMapDir';
 import { exists, verbose } from '../../utils';
-import OpenAI from 'openai';
-import { ChatCompletionMessageParam } from 'openai/resources';
 import Context from '../../inspect/context';
-import { EventMatch } from '../../search/types';
-import { readFile, writeFile } from 'fs/promises';
-import { join } from 'path';
-import { AppMapFilter, Metadata, buildAppMap } from '@appland/models';
-import { loadSequenceDiagram } from '../compare/loadSequenceDiagram';
-import {
-  Action,
-  Diagram,
-  FormatType,
-  Specification,
-  buildDiagram,
-  format,
-  nodeName,
-} from '@appland/sequence-diagram';
-import { match } from 'assert';
 
 export const command = 'ask <question>';
 export const describe =
@@ -46,12 +36,12 @@ function buildOpenAI(): OpenAI {
 }
 
 type SearchDiagramsParam = {
-  keywords: string[];
+  keyword: string;
 };
 
-type DiagramEvent = {
-  name: string;
-  id: number;
+type ActionInfo = {
+  eventIds?: string;
+  location?: string;
 };
 
 type SearchDiagramResult = {
@@ -67,15 +57,29 @@ type DiagramDetailsResult = {
   summary: string;
 };
 
+type Reference = {
+  id: number;
+  diagramId: string;
+  sourceLocation?: string;
+  eventIds?: number[];
+};
+
 class Ask {
   constructor(public appmapDir: string) {}
 
-  async searchDiagrams(keywords: string[]): Promise<SearchDiagramResult[] | undefined> {
+  async searchDiagrams(keyword: string): Promise<SearchDiagramResult[] | undefined> {
     const regexpOfGlob = (glob: string) =>
       glob.replace(/\./g, '\\.').replace(/\*/g, '.*').replace(/\?/g, '.');
 
     // TODO: Use proper keyword search
-    const searchTerms = keywords.map((keyword) => keyword.split(' ')).flat();
+    const searchTerms = [keyword];
+    // .map((keyword) => keyword.split(' '))
+    // .flat()
+    // .map((kw) => {
+    //   if (kw.startsWith('+')) return kw.substring(1);
+    //   if (kw.startsWith('-')) return kw.substring(1);
+    //   return kw;
+    // });
     const matchesForKeyword = new Array<Set<string>>();
     for (const keyword of searchTerms) {
       matchesForKeyword.push(new Set<string>());
@@ -85,31 +89,9 @@ class Ask {
       if (!context.codeObjectMatches) continue;
 
       const stats = await context.buildStats();
-      /*
-      const eventName = (eventMatch: EventMatch): string => {
-        if (eventMatch.event.sqlQuery) return eventMatch.event.sqlQuery;
-        else if (eventMatch.event.httpServerRequest)
-          return eventMatch.event.httpServerRequest.normalized_path;
-
-        return [eventMatch.event.definedClass, eventMatch.event.methodId].join('#');
-      };
-      */
-
       stats.eventMatches.forEach((eventMatch) => {
         const diagramId = eventMatch.appmap;
         matchesForKeyword[matchesForKeyword.length - 1].add(diagramId);
-        /*
-        const event: DiagramEvent = {
-          name: eventName(eventMatch),
-          id: eventMatch.event.id,
-        };
-        const eventMatches = eventMatchesByDiagramId.get(diagramId);
-        if (eventMatches) {
-          eventMatches.push(event);
-        } else {
-          eventMatchesByDiagramId.set(diagramId, [event]);
-        }
-        */
       });
     }
 
@@ -127,10 +109,6 @@ class Ask {
     }
 
     return Array.from(matchingDiagrams).map((diagramId) => ({ diagramId }));
-    // return Array.from(eventMatchesByDiagramId.entries()).map(([diagramId, matchingEvents]) => ({
-    //   diagramId,
-    //   matchingEvents,
-    // }));
   }
 }
 
@@ -148,9 +126,9 @@ export const handler = async (argv: any) => {
 
   async function searchDiagrams(paramStr: string): Promise<SearchDiagramResult[]> {
     const params = JSON.parse(paramStr) as SearchDiagramsParam;
-    let { keywords } = params;
-    warn(`Searching for diagrams matching "${keywords}"`);
-    const eventMatches = await ask.searchDiagrams(keywords);
+    let { keyword } = params;
+    warn(`Searching for diagrams matching "${keyword}"`);
+    const eventMatches = await ask.searchDiagrams(keyword);
     if (!eventMatches) return [];
 
     return eventMatches;
@@ -178,6 +156,10 @@ export const handler = async (argv: any) => {
 
       filter.declutter.limitRootEvents.on = true;
       const filteredAppMap = filter.filter(prunedAppMap, []);
+      const eventsById = filteredAppMap.events.reduce((map, event) => {
+        map.set(event.id, event);
+        return map;
+      }, new Map<number, Event>());
       const specification = Specification.build(filteredAppMap, { loops: true });
       const diagram = buildDiagram(appmapFile, filteredAppMap, specification);
 
@@ -186,8 +168,28 @@ export const handler = async (argv: any) => {
       const collectAction = (action: Action, depth = 0) => {
         if (depth > maxDepth) return;
 
+        const actionInfo: ActionInfo = {};
+        if (action.eventIds.length > 0) {
+          actionInfo.eventIds = action.eventIds.join(',');
+          const co = eventsById.get(action.eventIds[0])?.codeObject;
+          if (co) {
+            actionInfo.location = co.location;
+          } else {
+            warn(`No code object for event ${action.eventIds[0]}`);
+          }
+        }
+        const actionInfoStr = Object.keys(actionInfo)
+          .sort()
+          .map((key) => {
+            const value = actionInfo[key];
+            return `${key}=${value}`;
+          })
+          .join(',');
+
         const indent = '  '.repeat(depth);
-        diagramText.push(`${indent}${nodeName(action)}(events:${action.eventIds.join(', ')}})`);
+        diagramText.push(
+          `${indent}${nodeName(action)}${actionInfoStr !== '' ? ` (${actionInfoStr})` : ''}`
+        );
         if (action.children) {
           action.children.forEach((child) => collectAction(child, depth + 1));
         }
@@ -198,19 +200,60 @@ export const handler = async (argv: any) => {
     return result;
   }
 
+  async function lookupSourceCode(locationStr: string): Promise<string | undefined> {
+    const params = JSON.parse(locationStr) as { location: string };
+
+    const [path, lineno] = params.location.split(':');
+    if (!(await exists(path))) return;
+
+    const fileContent = await readFile(path, 'utf-8');
+    if (!lineno) return fileContent;
+
+    const languageRegexMap: Record<string, RegExp> = {
+      '.rb': new RegExp(`def\\s+\\w+.*?\\n(.*?\\n)*?^end\\b`, 'gm'),
+      '.java': new RegExp(
+        `(?:public|private|protected)?\\s+(?:static\\s+)?(?:final\\s+)?(?:synchronized\\s+)?(?:abstract\\s+)?(?:native\\s+)?(?:strictfp\\s+)?(?:transient\\s+)?(?:volatile\\s+)?(?:\\w+\\s+)*\\w+\\s+\\w+\\s*\\([^)]*\\)\\s*(?:throws\\s+\\w+(?:,\\s*\\w+)*)?\\s*\\{(?:[^{}]*\\{[^{}]*\\})*[^{}]*\\}`,
+        'gm'
+      ),
+      '.py': new RegExp(`def\\s+\\w+.*?:\\n(.*?\\n)*?`, 'gm'),
+      '.js': new RegExp(
+        `(?:async\\s+)?function\\s+\\w+\\s*\\([^)]*\\)\\s*\\{(?:[^{}]*\\{[^{}]*\\})*[^{}]*\\}`,
+        'gm'
+      ),
+    };
+
+    const extension = path.substring(path.lastIndexOf('.'));
+    const regex = languageRegexMap[extension];
+
+    if (regex) {
+      const match = regex.exec(fileContent);
+      if (match) {
+        const lines = match[0].split('\n');
+        const startLine = parseInt(lineno, 10);
+        const endLine = startLine + lines.length - 1;
+        if (startLine <= endLine) {
+          return lines.slice(startLine - 1, endLine).join('\n');
+        }
+      }
+    }
+  }
+
   const question = argv.question;
 
   const systemMessages: ChatCompletionMessageParam[] = [
     'You are an assistant that answers questions about the design and architecture of code.',
     'You answer these questions by accessing a knowledge base of sequence diagrams.',
     'Each sequence diagram conists of a series of events, such as function calls, HTTP server requests, SQL queries, etc.',
-    'You are given a question about the code, and you have to (a) Locate the sequence diagram that best answers the question (b) Focus the sequence diagram on the sub-section of the diagram that best answers the question.',
-    'You use the function calling interface to choose the best diagram and configure the diagram view to reveal the answer.',
-    'Always try multiple keywords, and collate the results together',
-    'If a diagram keyword search returns no diagrams, try other keywords.',
-    'The first function you should call is the "showPlan" function, to which you will provide a markdown document describing your general strategy for answering the question.',
-    'In your response, include links to the diagram files, with anchor tags to events that are particular interest',
-    'In your response, include links to source code files, with anchor tags to the lines of code that are of particular interest',
+    'First, call the  "showPlan" function with a Markdown document that describes your strategy for answering the question.',
+    'Then, call the "searchDiagrams" function with an array of keywords to search for in the diagrams.',
+    'If "searchDiagrams" fails, continue calling "searchDiagrams" with different keywords and synonyms of keywords until you find a match.',
+    'Call "searchDiagrams" with only one keyword at a time',
+    'Next, call the "getDiagramDetails" function with an array of diagram ids to get details about the diagrams.',
+    'Evaluate which diagrams are most relevant to the question.',
+    'Enhance your answer by using "lookupSourceCode" function to get the source code for the most relevant functions.',
+    'Finally, respond with a Markdown document that summarizes the diagrams and answers the question.',
+    'Subsequent mentions of the function should use backticks but should not be links.',
+    'Never emit phrases like "note that the actual behavior may vary between different applications"',
   ].map((msg) => ({
     content: msg,
     role: 'system',
@@ -248,12 +291,8 @@ export const handler = async (argv: any) => {
         parameters: {
           type: 'object',
           properties: {
-            keywords: {
-              type: 'array',
-              description: 'Array of keywords to find in the diagrams',
-              items: {
-                type: 'string',
-              },
+            keyword: {
+              type: 'string',
             },
           },
         },
@@ -274,6 +313,19 @@ export const handler = async (argv: any) => {
           },
         },
       },
+      {
+        function: lookupSourceCode,
+        description: `Get the source code for a specific function.`,
+        parameters: {
+          type: 'object',
+          properties: {
+            location: {
+              type: 'string',
+              description: `Source code location in the format <path>[:<line number>]. Line number can be omitted if it's not known`,
+            },
+          },
+        },
+      },
     ],
   });
 
@@ -296,12 +348,4 @@ export const handler = async (argv: any) => {
     return;
   }
   console.log(response);
-
-  // const context = new Context(appmapDir, codeObjectId);
-  // await context.findCodeObjects();
-  // if (context.codeObjectMatches?.length === 0) {
-  //   return yargs.exit(1, new Error(`Code object '${context.codeObjectId}' not found`));
-  // }
-  // await context.buildStats();
-  // console.log(JSON.stringify(context.stats, null, 2));
 };