feat(chat): support auto-sync active selection in current Notebook fi…

…le (#3537) * feat(chat): support auto-sync active selection in current Notebook file * [autofix.ci] apply automated fixes * update * update * [autofix.ci] apply automated fixes * update * update * update * update * update * revert * [autofix.ci] apply automated fixes * update * update * update * [autofix.ci] apply automated fixes * update: type * update: remove --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: Zhiming Ma <[email protected]>
TabbyML · Dec 28, 2024 · 3f7af32 · 3f7af32
1 parent fcb1783
commit 3f7af32
Show file tree

Hide file tree

Showing 5 changed files with 174 additions and 17 deletions.
diff --git a/clients/vscode/src/chat/WebviewHelper.ts b/clients/vscode/src/chat/WebviewHelper.ts
@@ -41,6 +41,7 @@ import {
   vscodePositionToChatPanelPosition,
   vscodeRangeToChatPanelPositionRange,
   chatPanelLocationToVSCodeRange,
+  isSupportedSchemeForActiveSelection,
 } from "./utils";
 
 export class WebviewHelper {
@@ -263,11 +264,6 @@ export class WebviewHelper {
     }
   }
 
-  public isSupportedSchemeForActiveSelection(scheme: string) {
-    const supportedSchemes = ["file", "untitled"];
-    return supportedSchemes.includes(scheme);
-  }
-
   public async syncActiveSelectionToChatPanel(context: EditorContext | null) {
     try {
       await this.client?.updateActiveSelection(context);
@@ -346,7 +342,7 @@ export class WebviewHelper {
   }
 
   public async syncActiveSelection(editor: TextEditor | undefined) {
-    if (!editor || !this.isSupportedSchemeForActiveSelection(editor.document.uri.scheme)) {
+    if (!editor || !isSupportedSchemeForActiveSelection(editor.document.uri.scheme)) {
       await this.syncActiveSelectionToChatPanel(null);
       return;
     }
@@ -374,7 +370,7 @@ export class WebviewHelper {
 
     window.onDidChangeTextEditorSelection((e) => {
       // This listener only handles text files.
-      if (!this.isSupportedSchemeForActiveSelection(e.textEditor.document.uri.scheme)) {
+      if (!isSupportedSchemeForActiveSelection(e.textEditor.document.uri.scheme)) {
         return;
       }
       this.syncActiveSelection(e.textEditor);

diff --git a/clients/vscode/src/chat/utils.ts b/clients/vscode/src/chat/utils.ts
@@ -1,22 +1,50 @@
 import path from "path";
 import { Position as VSCodePosition, Range as VSCodeRange, Uri, workspace } from "vscode";
-import type { Filepath, Position as ChatPanelPosition, LineRange, PositionRange, Location } from "tabby-chat-panel";
+import type {
+  Filepath,
+  Position as ChatPanelPosition,
+  LineRange,
+  PositionRange,
+  Location,
+  FilepathInGitRepository,
+} from "tabby-chat-panel";
 import type { GitProvider } from "../git/GitProvider";
 import { getLogger } from "../logger";
 
 const logger = getLogger("chat/utils");
 
+enum Schemes {
+  file = "file",
+  untitled = "untitled",
+  vscodeNotebookCell = "vscode-notebook-cell",
+  vscodeVfs = "vscode-vfs",
+  vscodeUserdata = "vscode-userdata",
+}
+
+export function isSupportedSchemeForActiveSelection(scheme: string): boolean {
+  const supportedSchemes: string[] = [Schemes.file, Schemes.untitled, Schemes.vscodeNotebookCell];
+  return supportedSchemes.includes(scheme);
+}
+
 export function localUriToChatPanelFilepath(uri: Uri, gitProvider: GitProvider): Filepath {
+  let uriFilePath = uri.toString(true);
+  if (uri.scheme === Schemes.vscodeNotebookCell) {
+    const notebook = parseNotebookCellUri(uri);
+    if (notebook) {
+      // add fragment `#cell={number}` to filepath
+      uriFilePath = uri.with({ scheme: notebook.notebook.scheme, fragment: `cell=${notebook.handle}` }).toString(true);
+    }
+  }
+
   const workspaceFolder = workspace.getWorkspaceFolder(uri);
 
   let repo = gitProvider.getRepository(uri);
   if (!repo && workspaceFolder) {
     repo = gitProvider.getRepository(workspaceFolder.uri);
   }
   const gitRemoteUrl = repo ? gitProvider.getDefaultRemoteUrl(repo) : undefined;
-
   if (repo && gitRemoteUrl) {
-    const relativeFilePath = path.relative(repo.rootUri.toString(true), uri.toString(true));
+    const relativeFilePath = path.relative(repo.rootUri.toString(true), uriFilePath);
     if (!relativeFilePath.startsWith("..")) {
       return {
         kind: "git",
@@ -28,13 +56,28 @@ export function localUriToChatPanelFilepath(uri: Uri, gitProvider: GitProvider):
 
   return {
     kind: "uri",
-    uri: uri.toString(true),
+    uri: uriFilePath,
   };
 }
 
+function isJupyterNotebookFilepath(filepath: Filepath): boolean {
+  const _filepath = filepath.kind === "uri" ? filepath.uri : filepath.filepath;
+  const extname = path.extname(_filepath);
+  return extname.startsWith(".ipynb");
+}
+
 export function chatPanelFilepathToLocalUri(filepath: Filepath, gitProvider: GitProvider): Uri | null {
+  const isNotebook = isJupyterNotebookFilepath(filepath);
+
   if (filepath.kind === "uri") {
     try {
+      if (isNotebook) {
+        const handle = chatPanelFilePathToNotebookCellHandle(filepath.uri);
+        if (typeof handle === "number") {
+          return generateLocalNotebookCellUri(Uri.parse(filepath.uri), handle);
+        }
+      }
+
       return Uri.parse(filepath.uri, true);
     } catch (e) {
       // FIXME(@icycodes): this is a hack for uri is relative filepaths in workspaces
@@ -46,13 +89,53 @@ export function chatPanelFilepathToLocalUri(filepath: Filepath, gitProvider: Git
   } else if (filepath.kind === "git") {
     const localGitRoot = gitProvider.findLocalRootUriByRemoteUrl(filepath.gitUrl);
     if (localGitRoot) {
+      // handling for Jupyter Notebook (.ipynb) files
+      if (isNotebook) {
+        return chatPanelFilepathToVscodeNotebookCellUri(localGitRoot, filepath);
+      }
+
       return Uri.joinPath(localGitRoot, filepath.filepath);
     }
   }
   logger.warn(`Invalid filepath params.`, filepath);
   return null;
 }
 
+function chatPanelFilepathToVscodeNotebookCellUri(root: Uri, filepath: FilepathInGitRepository): Uri | null {
+  if (filepath.kind !== "git") {
+    logger.warn(`Invalid filepath params.`, filepath);
+    return null;
+  }
+
+  const filePathUri = Uri.parse(filepath.filepath);
+  const notebookUri = Uri.joinPath(root, filePathUri.path);
+
+  const handle = chatPanelFilePathToNotebookCellHandle(filepath.filepath);
+  if (typeof handle === "undefined") {
+    logger.warn(`Invalid filepath params.`, filepath);
+    return null;
+  }
+  return generateLocalNotebookCellUri(notebookUri, handle);
+}
+
+function chatPanelFilePathToNotebookCellHandle(filepath: string): number | undefined {
+  let handle: number | undefined;
+
+  const fileUri = Uri.parse(filepath);
+  const fragment = fileUri.fragment;
+  const searchParams = new URLSearchParams(fragment);
+  if (searchParams.has("cell")) {
+    const cellString = searchParams.get("cell")?.toString() || "";
+    handle = parseInt(cellString, 10);
+  }
+
+  if (typeof handle === "undefined" || isNaN(handle)) {
+    return undefined;
+  }
+
+  return handle;
+}
+
 export function vscodePositionToChatPanelPosition(position: VSCodePosition): ChatPanelPosition {
   return {
     line: position.line + 1,
@@ -103,3 +186,34 @@ export function chatPanelLocationToVSCodeRange(location: Location | undefined):
   logger.warn(`Invalid location params.`, location);
   return null;
 }
+
+const nb_lengths = ["W", "X", "Y", "Z", "a", "b", "c", "d", "e", "f"];
+const nb_padRegexp = new RegExp(`^[${nb_lengths.join("")}]+`);
+const nb_radix = 7;
+export function parseNotebookCellUri(cell: Uri): { notebook: Uri; handle: number } | undefined {
+  if (cell.scheme !== Schemes.vscodeNotebookCell) {
+    return undefined;
+  }
+
+  const idx = cell.fragment.indexOf("s");
+  if (idx < 0) {
+    return undefined;
+  }
+
+  const handle = parseInt(cell.fragment.substring(0, idx).replace(nb_padRegexp, ""), nb_radix);
+  const _scheme = Buffer.from(cell.fragment.substring(idx + 1), "base64").toString("utf-8");
+  if (isNaN(handle)) {
+    return undefined;
+  }
+  return {
+    handle,
+    notebook: cell.with({ scheme: _scheme, fragment: "" }),
+  };
+}
+
+export function generateLocalNotebookCellUri(notebook: Uri, handle: number): Uri {
+  const s = handle.toString(nb_radix);
+  const p = s.length < nb_lengths.length ? nb_lengths[s.length - 1] : "z";
+  const fragment = `${p}${s}s${Buffer.from(notebook.scheme).toString("base64")}`;
+  return notebook.with({ scheme: Schemes.vscodeNotebookCell, fragment });
+}
diff --git a/ee/tabby-ui/components/chat/chat-panel.tsx b/ee/tabby-ui/components/chat/chat-panel.tsx
@@ -13,7 +13,11 @@ import { useChatStore } from '@/lib/stores/chat-store'
 import { useMutation } from '@/lib/tabby/gql'
 import { setThreadPersistedMutation } from '@/lib/tabby/query'
 import type { Context } from '@/lib/types'
-import { cn, getTitleFromMessages } from '@/lib/utils'
+import {
+  cn,
+  getTitleFromMessages,
+  resolveFileNameForDisplay
+} from '@/lib/utils'
 import { Badge } from '@/components/ui/badge'
 import { Button } from '@/components/ui/button'
 import {
@@ -344,7 +348,6 @@ function ContextLabel({
   context: Context
   className?: string
 }) {
-  const [fileName] = context.filepath.split('/').slice(-1)
   const line = context.range
     ? context.range.start === context.range.end
       ? `:${context.range.start}`
@@ -353,7 +356,7 @@ function ContextLabel({
 
   return (
     <span className={cn('truncate', className)}>
-      {fileName}
+      {resolveFileNameForDisplay(context.filepath)}
       {!!context.range && <span className="text-muted-foreground">{line}</span>}
     </span>
   )

diff --git a/ee/tabby-ui/components/chat/code-references.tsx b/ee/tabby-ui/components/chat/code-references.tsx
@@ -2,7 +2,7 @@ import React, { forwardRef, useEffect, useState } from 'react'
 import { isNil } from 'lodash-es'
 
 import { RelevantCodeContext } from '@/lib/types'
-import { cn } from '@/lib/utils'
+import { cn, resolveFileNameForDisplay } from '@/lib/utils'
 import {
   Tooltip,
   TooltipContent,
@@ -151,7 +151,6 @@ function ContextItem({
     !isNil(context.range?.end) &&
     context.range.start < context.range.end
   const pathSegments = context.filepath.split('/')
-  const fileName = pathSegments[pathSegments.length - 1]
   const path = pathSegments.slice(0, pathSegments.length - 1).join('/')
   const scores = context?.extra?.scores
   const onTooltipOpenChange = (v: boolean) => {
@@ -178,7 +177,7 @@ function ContextItem({
           <div className="flex items-center gap-1 overflow-hidden">
             <IconFile className="shrink-0" />
             <div className="flex-1 truncate" title={context.filepath}>
-              <span>{fileName}</span>
+              <span>{resolveFileNameForDisplay(context.filepath)}</span>
               {context.range ? (
                 <>
                   {context.range?.start && (

diff --git a/ee/tabby-ui/lib/utils/chat.ts b/ee/tabby-ui/lib/utils/chat.ts
@@ -105,3 +105,48 @@ export function checkSourcesAvailability(
 
   return { hasCodebaseSource, hasDocumentSource }
 }
+
+/**
+ * url e.g #cell=1
+ * @param fragment
+ * @returns
+ */
+function parseNotebookCellUriFragment(fragment: string) {
+  if (!fragment) return undefined
+  try {
+    const searchParams = new URLSearchParams(fragment)
+    const cellString = searchParams.get('cell')?.toString()
+    if (!cellString) {
+      return undefined
+    }
+
+    const handle = parseInt(cellString, 10)
+
+    if (isNaN(handle)) {
+      return undefined
+    }
+    return {
+      handle
+    }
+  } catch (error) {
+    return undefined
+  }
+}
+
+export function resolveFileNameForDisplay(uri: string) {
+  let url: URL
+  try {
+    url = new URL(uri)
+  } catch (e) {
+    url = new URL(uri, 'file://')
+  }
+  const filename = url.pathname.split('/').pop() || ''
+  const extname = filename.includes('.') ? `.${filename.split('.').pop()}` : ''
+  const isNotebook = extname.startsWith('.ipynb')
+  const hash = url.hash ? url.hash.substring(1) : ''
+  const cell = parseNotebookCellUriFragment(hash)
+  if (isNotebook && cell) {
+    return `${filename} · Cell ${(cell.handle || 0) + 1}`
+  }
+  return filename
+}