Refactor find-relevant-files command into modular structure with utility files

Author: robertpiosik

apps/editor/src/commands/find-relevant-files-command.ts

@@ -0,0 +1,155 @@
+import * as vscode from 'vscode'
+import { WorkspaceProvider } from '../../context/providers/workspace/workspace-provider'
+import { ModelProvidersManager } from '../../services/model-providers-manager'
+import {
+  LAST_FIND_RELEVANT_FILES_QUERY_STATE_KEY,
+  LAST_FIND_RELEVANT_FILES_SHRINK_STATE_KEY
+} from '../../constants/state-keys'
+import { t } from '@/i18n'
+
+import { get_target_folder_path } from './utils/get-target-folder-path'
+import { prompt_for_instructions } from './utils/prompt-for-instructions'
+import { analyze_workspace_files } from './utils/analyze-workspace-files'
+import { prompt_for_shrink_mode } from './utils/prompt-for-shrink-mode'
+import { prompt_for_config } from './utils/prompt-for-config'
+import { fetch_relevant_files_from_api } from './utils/fetch-relevant-files-from-api'
+import { show_results_and_apply } from './utils/show-results-and-apply'
+
+export const find_relevant_files_command = (
+  workspace_provider: WorkspaceProvider,
+  extension_context: vscode.ExtensionContext
+) => {
+  return vscode.commands.registerCommand(
+    'codeWebChat.findRelevantFiles',
+    async (item?: any) => {
+      const folder_path = await get_target_folder_path(item)
+
+      if (!folder_path) {
+        vscode.window.showErrorMessage(
+          t('command.find-relevant-files.error.no-folder-selected')
+        )
+        return
+      }
+
+      let initial_instructions =
+        extension_context.workspaceState.get<string>(
+          LAST_FIND_RELEVANT_FILES_QUERY_STATE_KEY
+        ) || ''
+
+      while (true) {
+        const instructions = await prompt_for_instructions(initial_instructions)
+        if (!instructions) return
+
+        await extension_context.workspaceState.update(
+          LAST_FIND_RELEVANT_FILES_QUERY_STATE_KEY,
+          instructions
+        )
+        initial_instructions = instructions
+
+        const analysis = await analyze_workspace_files({
+          workspace_provider,
+          folder_path
+        })
+        let go_back_to_input = false
+        while (true) {
+          const should_shrink = extension_context.workspaceState.get<boolean>(
+            LAST_FIND_RELEVANT_FILES_SHRINK_STATE_KEY,
+            false
+          )
+          const shrink_result = await prompt_for_shrink_mode({
+            should_shrink,
+            full_tokens: analysis.full_tokens,
+            shrink_tokens: analysis.shrink_tokens
+          })
+
+          if (shrink_result === 'back') {
+            go_back_to_input = true
+            break
+          }
+          if (shrink_result === 'cancel') return
+
+          await extension_context.workspaceState.update(
+            LAST_FIND_RELEVANT_FILES_SHRINK_STATE_KEY,
+            shrink_result
+          )
+
+          const api_providers_manager = new ModelProvidersManager(
+            extension_context
+          )
+          const configs =
+            await api_providers_manager.get_find_relevant_files_tool_configs()
+
+          if (configs.length === 0) {
+            vscode.commands.executeCommand('codeWebChat.settings')
+            vscode.window.showInformationMessage(
+              t('command.find-relevant-files.error.no-configs')
+            )
+            return
+          }
+
+          let go_back_to_shrink = false
+          while (true) {
+            const tokens_to_process = shrink_result
+              ? analysis.shrink_tokens
+              : analysis.full_tokens
+            const config_result = await prompt_for_config({
+              api_providers_manager,
+              extension_context,
+              configs,
+              tokens_to_process
+            })
+
+            if (config_result === 'back') {
+              go_back_to_shrink = true
+              break
+            }
+            if (config_result === 'cancel') return
+
+            const {
+              config: selected_config,
+              provider,
+              skipped: skipped_config
+            } = config_result
+
+            const api_result = await fetch_relevant_files_from_api(
+              analysis.files_data,
+              shrink_result as boolean,
+              instructions,
+              provider,
+              selected_config
+            )
+
+            if (api_result === 'cancel' || api_result === 'error') return
+            if (api_result === 'error_no_files') {
+              vscode.window.showWarningMessage(
+                t('command.find-relevant-files.error.no-files-found')
+              )
+              go_back_to_input = true
+              break
+            }
+
+            const apply_result = await show_results_and_apply({
+              extracted_files: api_result,
+              analysis,
+              workspace_provider
+            })
+
+            if (apply_result === 'back') {
+              if (skipped_config) {
+                go_back_to_shrink = true
+                break
+              } else continue
+            }
+
+            if (apply_result === 'cancel' || apply_result === 'success') return
+          }
+
+          if (go_back_to_shrink) continue
+          if (go_back_to_input) break
+        }
+
+        if (go_back_to_input) continue
+      }
+    }
+  )
+}

apps/editor/src/commands/find-relevant-files-command/find-relevant-files-command.ts

@@ -0,0 +1 @@
+export * from './find-relevant-files-command'

apps/editor/src/commands/find-relevant-files-command/index.ts

@@ -0,0 +1,64 @@
+import * as vscode from 'vscode'
+import * as path from 'path'
+import * as fs from 'fs'
+import { WorkspaceProvider } from '../../../context/providers/workspace/workspace-provider'
+import { shrink_file } from '../../../context/utils/shrink-file/shrink-file'
+import { t } from '@/i18n'
+
+export interface FileData {
+  file_path: string
+  relative_path: string
+  content: string
+  shrunk_content: string
+}
+
+export interface FileAnalysisResult {
+  full_tokens: number
+  shrink_tokens: number
+  files_data: FileData[]
+  workspace_root: string
+}
+
+export const analyze_workspace_files = async (params: {
+  workspace_provider: WorkspaceProvider
+  folder_path: string
+}): Promise<FileAnalysisResult> => {
+  const all_files = await params.workspace_provider.find_all_files(
+    params.folder_path
+  )
+  const workspace_root =
+    params.workspace_provider.get_workspace_root_for_file(params.folder_path) ||
+    params.folder_path
+
+  let full_tokens = 0
+  let shrink_tokens = 0
+  const files_data: FileData[] = []
+
+  await vscode.window.withProgress(
+    {
+      location: vscode.ProgressLocation.Window,
+      title: t('command.find-relevant-files.progress.analyzing')
+    },
+    async () => {
+      for (const file_path of all_files) {
+        try {
+          const stats = await fs.promises.stat(file_path)
+          if (stats.size > 1024 * 1024) continue
+
+          const content = await fs.promises.readFile(file_path, 'utf8')
+          const shrunk_content = shrink_file(content, path.extname(file_path))
+          const relative_path = path.relative(workspace_root, file_path)
+
+          files_data.push({ file_path, relative_path, content, shrunk_content })
+
+          const token_count =
+            await params.workspace_provider.calculate_file_tokens(file_path)
+          full_tokens += token_count.total
+          shrink_tokens += token_count.shrink
+        } catch (e) {}
+      }
+    }
+  )
+
+  return { full_tokens, shrink_tokens, files_data, workspace_root }
+}

apps/editor/src/commands/find-relevant-files-command/utils/analyze-workspace-files.ts

@@ -0,0 +1,104 @@
+import * as vscode from 'vscode'
+import axios from 'axios'
+import { make_api_request } from '../../../utils/make-api-request'
+import {
+  find_relevant_files_instructions,
+  find_relevant_files_format
+} from '../../../constants/instructions'
+import { apply_reasoning_effort } from '../../../utils/apply-reasoning-effort'
+import { build_user_content } from '../../../utils/build-user-content'
+import { Logger } from '@shared/utils/logger'
+import { FileData } from './analyze-workspace-files'
+import { t } from '@/i18n'
+
+export const fetch_relevant_files_from_api = async (
+  files_data: FileData[],
+  shrink_result: boolean,
+  instructions: string,
+  provider: any,
+  selected_config: any
+): Promise<string[] | 'cancel' | 'error_no_files' | 'error'> => {
+  let xml_files = `<files>\n`
+  for (const file of files_data) {
+    const content_to_use = shrink_result ? file.shrunk_content : file.content
+    xml_files += `<file path="${file.relative_path}">\n<![CDATA[\n${content_to_use}\n]]>\n</file>\n`
+  }
+  xml_files += `</files>`
+
+  const system_instructions_xml = `${find_relevant_files_format}\n${find_relevant_files_instructions}`
+  const part2 = `${system_instructions_xml}\n${instructions}`
+  const user_content = build_user_content({
+    provider_name: provider.name,
+    part1: xml_files,
+    part2,
+    disable_cache: true
+  })
+
+  const messages = [{ role: 'user', content: user_content }]
+  const body: { [key: string]: any } = {
+    messages,
+    model: selected_config.model,
+    temperature: selected_config.temperature
+  }
+
+  apply_reasoning_effort({
+    body,
+    provider,
+    reasoning_effort: selected_config.reasoning_effort
+  })
+
+  const cancel_token_source = axios.CancelToken.source()
+
+  try {
+    const completion_result = await vscode.window.withProgress(
+      {
+        location: vscode.ProgressLocation.Notification,
+        title: t('command.find-relevant-files.progress.finding'),
+        cancellable: true
+      },
+      async (progress, token) => {
+        token.onCancellationRequested(() => {
+          cancel_token_source.cancel(
+            t('command.find-relevant-files.cancel.user')
+          )
+        })
+        progress.report({ message: t('common.progress.waiting-for-server') })
+        return await make_api_request({
+          endpoint_url: provider.base_url,
+          api_key: provider.api_key,
+          body,
+          cancellation_token: cancel_token_source.token,
+          on_chunk: () =>
+            progress.report({ message: t('common.progress.receiving') }),
+          on_thinking_chunk: () =>
+            progress.report({ message: t('common.progress.thinking') })
+        })
+      }
+    )
+
+    if (completion_result) {
+      const match = completion_result.response.match(
+        /<relevant-files>([\s\S]*?)<\/relevant-files>/
+      )
+      const extracted_files: string[] = []
+      if (match && match[1]) {
+        const file_matches = match[1].matchAll(/<file-path>(.*?)<\/file-path>/g)
+        for (const m of file_matches) extracted_files.push(m[1].trim())
+      }
+      return extracted_files.length === 0 ? 'error_no_files' : extracted_files
+    }
+    return 'cancel'
+  } catch (error) {
+    if (!axios.isCancel(error)) {
+      Logger.error({
+        function_name: 'fetch_relevant_files_from_api',
+        message: 'Error finding relevant files',
+        data: error
+      })
+      vscode.window.showErrorMessage(
+        t('command.find-relevant-files.error.finding')
+      )
+    }
+    return 'error'
+  }
+}

apps/editor/src/commands/find-relevant-files-command/utils/fetch-relevant-files-from-api.ts

@@ -0,0 +1,19 @@
+import * as path from 'path'
+import * as fs from 'fs'
+
+export const get_target_folder_path = async (
+  item?: any
+): Promise<string | undefined> => {
+  let folder_path = item?.resourceUri?.fsPath
+  if (folder_path) {
+    try {
+      const stats = await fs.promises.stat(folder_path)
+      if (!stats.isDirectory()) {
+        folder_path = path.dirname(folder_path)
+      }
+    } catch (error) {
+      folder_path = undefined
+    }
+  }
+  return folder_path
+}