feat(chat): refactor provideAtInfo to provideSymbolAtInfo and provide…

…FileAtInfo with enhanced content retrieval
TabbyML · Dec 20, 2024 · eff9450 · eff9450
1 parent 34f8616
commit eff9450
Show file tree

Hide file tree

Showing 4 changed files with 98 additions and 63 deletions.
diff --git a/clients/tabby-chat-panel/src/index.ts b/clients/tabby-chat-panel/src/index.ts
@@ -372,11 +372,17 @@ export interface ClientApiMethods {
   readWorkspaceGitRepositories?: () => Promise<GitRepository[]>
 
   /**
-   * Return a AtInfo List with kind of file
+   * Return a SymbolAtInfo List with kind of file
    * @param kind passing what kind of At info client want to get
-   * @returns AtInfo array
+   * @returns SymbolAtInfo array
    */
-  provideAtInfo?: (kind: 'symbol' | 'file', opts?: AtInputOpts) => Promise<AtInfo[] | null>
+  provideSymbolAtInfo?: (opts?: AtInputOpts) => Promise<SymbolAtInfo[] | null>
+
+  getSymbolAtInfoContent?: (info: SymbolAtInfo) => Promise<string | null>
+
+  provideFileAtInfo?: (opts?: AtInputOpts) => Promise<FileAtInfo[] | null>
+
+  getFileAtInfoContent?: (info: FileAtInfo) => Promise<string | null>
 }
 
 export interface ClientApi extends ClientApiMethods {

diff --git a/clients/vscode/src/chat/WebviewHelper.ts b/clients/vscode/src/chat/WebviewHelper.ts
@@ -28,9 +28,9 @@ import type {
   SymbolInfo,
   FileLocation,
   GitRepository,
-  AtInfo,
-  AtKind,
   AtInputOpts,
+  SymbolAtInfo,
+  FileAtInfo,
 } from "tabby-chat-panel";
 import { TABBY_CHAT_PANEL_API_VERSION } from "tabby-chat-panel";
 import hashObject from "object-hash";
@@ -49,9 +49,9 @@ import {
   vscodeRangeToChatPanelPositionRange,
   chatPanelLocationToVSCodeRange,
   getAllowedSymbolKinds,
-  vscodeSymbolToAtInfo,
+  vscodeSymbolToAtInfo as vscodeSymbolToSymbolAtInfo,
   isDocumentSymbol,
-  uriToAtInfo,
+  uriToFileAtInfo,
 } from "./utils";
 
 export class WebviewHelper {
@@ -739,71 +739,96 @@ export class WebviewHelper {
         }
         return infoList;
       },
-      provideAtInfo: async (kind: AtKind, opts?: AtInputOpts): Promise<AtInfo[] | null> => {
-        // Common parameters
+      provideSymbolAtInfo: async (opts?: AtInputOpts): Promise<SymbolAtInfo[] | null> => {
         const maxResults = opts?.limit || 50;
         const query = opts?.query?.toLowerCase();
 
-        switch (kind) {
-          case "symbol": {
-            const editor = window.activeTextEditor;
-            if (!editor) return null;
-            const document = editor.document;
+        const editor = window.activeTextEditor;
+        if (!editor) return null;
+        const document = editor.document;
 
-            // Try document symbols first
-            const documentSymbols = await commands.executeCommand<DocumentSymbol[] | SymbolInformation[]>(
-              "vscode.executeDocumentSymbolProvider",
-              document.uri,
-            );
+        // Try document symbols first
+        const documentSymbols = await commands.executeCommand<DocumentSymbol[] | SymbolInformation[]>(
+          "vscode.executeDocumentSymbolProvider",
+          document.uri,
+        );
 
-            let results: AtInfo[] = [];
+        let results: SymbolAtInfo[] = [];
 
-            if (documentSymbols && documentSymbols.length > 0) {
-              const processSymbol = (symbol: DocumentSymbol | SymbolInformation) => {
-                if (results.length >= maxResults) return;
+        if (documentSymbols && documentSymbols.length > 0) {
+          const processSymbol = (symbol: DocumentSymbol | SymbolInformation) => {
+            if (results.length >= maxResults) return;
 
-                const symbolName = symbol.name.toLowerCase();
-                if (query && !symbolName.includes(query)) return;
+            const symbolName = symbol.name.toLowerCase();
+            if (query && !symbolName.includes(query)) return;
 
-                if (getAllowedSymbolKinds().includes(symbol.kind)) {
-                  results.push(vscodeSymbolToAtInfo(symbol, document.uri, this.gitProvider));
-                }
-                if (isDocumentSymbol(symbol)) {
-                  symbol.children.forEach(processSymbol);
-                }
-              };
-              documentSymbols.forEach(processSymbol);
+            if (getAllowedSymbolKinds().includes(symbol.kind)) {
+              results.push(vscodeSymbolToSymbolAtInfo(symbol, document.uri, this.gitProvider));
             }
-
-            // Try workspace symbols if no document symbols found
-            if (results.length === 0 && query) {
-              const workspaceSymbols = await commands.executeCommand<SymbolInformation[]>(
-                "vscode.executeWorkspaceSymbolProvider",
-                query,
-              );
-
-              if (workspaceSymbols) {
-                results = workspaceSymbols
-                  .filter((symbol) => getAllowedSymbolKinds().includes(symbol.kind))
-                  .slice(0, maxResults)
-                  .map((symbol) => vscodeSymbolToAtInfo(symbol, symbol.location.uri, this.gitProvider));
-              }
+            if (isDocumentSymbol(symbol)) {
+              symbol.children.forEach(processSymbol);
             }
+          };
+          documentSymbols.forEach(processSymbol);
+        }
 
-            return results.length > 0 ? results : null;
-          }
-          case "file": {
-            // hack way to only get prefix match
-            const globPattern = query ? `**/${query}*` : "**/*";
-            try {
-              const files = await workspace.findFiles(globPattern, null, maxResults);
-              return files.map((uri) => uriToAtInfo(uri, this.gitProvider));
-            } catch (error) {
-              this.logger.error("Failed to find files:", error);
-              return null;
-            }
+        // Try workspace symbols if no document symbols found
+        if (results.length === 0 && query) {
+          const workspaceSymbols = await commands.executeCommand<SymbolInformation[]>(
+            "vscode.executeWorkspaceSymbolProvider",
+            query,
+          );
+
+          if (workspaceSymbols) {
+            results = workspaceSymbols
+              .filter((symbol) => getAllowedSymbolKinds().includes(symbol.kind))
+              .slice(0, maxResults)
+              .map((symbol) => vscodeSymbolToSymbolAtInfo(symbol, symbol.location.uri, this.gitProvider));
           }
         }
+
+        return results.length > 0 ? results : null;
+      },
+
+      provideFileAtInfo: async (opts?: AtInputOpts): Promise<FileAtInfo[] | null> => {
+        const maxResults = opts?.limit || 50;
+        const query = opts?.query?.toLowerCase();
+
+        const globPattern = query ? `**/${query}*` : "**/*";
+        try {
+          const files = await workspace.findFiles(globPattern, null, maxResults);
+          return files.map((uri) => uriToFileAtInfo(uri, this.gitProvider));
+        } catch (error) {
+          this.logger.error("Failed to find files:", error);
+          return null;
+        }
+      },
+      getSymbolAtInfoContent: async (info: SymbolAtInfo): Promise<string | null> => {
+        try {
+          const uri = chatPanelFilepathToLocalUri(info.location.filepath, this.gitProvider);
+          if (!uri) return null;
+
+          const document = await workspace.openTextDocument(uri);
+          const range = chatPanelLocationToVSCodeRange(info.location.location);
+          if (!range) return null;
+
+          return document.getText(range);
+        } catch (error) {
+          this.logger.error("Failed to get symbol content:", error);
+          return null;
+        }
+      },
+      getFileAtInfoContent: async (info: FileAtInfo): Promise<string | null> => {
+        try {
+          const uri = chatPanelFilepathToLocalUri(info.filepath, this.gitProvider);
+          if (!uri) return null;
+
+          const document = await workspace.openTextDocument(uri);
+          return document.getText();
+        } catch (error) {
+          this.logger.error("Failed to get file content:", error);
+          return null;
+        }
       },
     });
   }

diff --git a/clients/vscode/src/chat/chatPanel.ts b/clients/vscode/src/chat/chatPanel.ts
@@ -36,7 +36,10 @@ export function createClient(webview: Webview, api: ClientApiMethods): ServerApi
       lookupSymbol: api.lookupSymbol,
       openInEditor: api.openInEditor,
       readWorkspaceGitRepositories: api.readWorkspaceGitRepositories,
-      provideAtInfo: api.provideAtInfo,
+      provideSymbolAtInfo: api.provideSymbolAtInfo,
+      getSymbolAtInfoContent: api.getSymbolAtInfoContent,
+      provideFileAtInfo: api.provideFileAtInfo,
+      getFileAtInfoContent: api.getFileAtInfoContent,
     },
   });
 }
diff --git a/clients/vscode/src/chat/utils.ts b/clients/vscode/src/chat/utils.ts
@@ -14,7 +14,8 @@ import type {
   LineRange,
   PositionRange,
   Location,
-  AtInfo,
+  SymbolAtInfo,
+  FileAtInfo,
 } from "tabby-chat-panel";
 import type { GitProvider } from "../git/GitProvider";
 import { getLogger } from "../logger";
@@ -136,7 +137,7 @@ export function vscodeSymbolToAtInfo(
   symbol: DocumentSymbol | SymbolInformation,
   documentUri: Uri,
   gitProvider: GitProvider,
-): AtInfo {
+): SymbolAtInfo {
   if (isDocumentSymbol(symbol)) {
     return {
       atKind: "symbol",
@@ -157,7 +158,7 @@ export function vscodeSymbolToAtInfo(
   };
 }
 
-export function uriToAtInfo(uri: Uri, gitProvider: GitProvider): AtInfo {
+export function uriToFileAtInfo(uri: Uri, gitProvider: GitProvider): FileAtInfo {
   return {
     atKind: "file",
     name: path.basename(uri.fsPath),