Update token counting mechanism to include detailed breakdowns for files, instructions, and formats

- to 'token count' message, add keys 'edit instructions count' 'ask instructions count'. Now TokenCountMessage should re-send whenever instructions change.
- remove `apps/editor/src/views/prompt/backend/message-handlers/handle-get-token-count.ts` if redundant
- to properly calculate [Fragment], you need to temporaroly replace symbols
- TokenCountMessage should also include 'edit format instructions count'
- place separator below 'code at cursor' items
- [Fragment] should be suffixed with 'token_count', not just 'count'
- rename 'token count' to 'selected files token count'

└── apps
    └── editor
        ├── package.json
        └── src
            ├── constants
            │   └── edit-format-instructions.ts
            ├── context
            │   └── context-initialization.ts
            └── views
                └── prompt
                    ├── backend
                    │   ├── message-handlers
                    │   │   ├── handle-autofill.ts
                    │   │   └── index.ts
                    │   ├── prompt-view-provider.ts
                    │   └── utils
                    │       └── symbols
                    │           ├── clipboard-paths
                    │           │   └── replace-clipboard-paths-symbol.ts
                    │           ├── fragment
                    │           │   └── replace-fragment-symbol.ts
                    │           ├── git
                    │           │   └── replace-git-symbols.ts
                    │           ├── image
                    │           │   └── replace-image-symbol.ts
                    │           ├── pasted-text
                    │           │   └── replace-pasted-text-symbol.ts
                    │           ├── replace-symbols.ts
                    │           ├── saved-context
                    │           │   └── replace-saved-context-symbol.ts
                    │           ├── selection
                    │           │   └── replace-selection-symbol.ts
                    │           ├── skill
                    │           │   └── replace-skill-symbol.ts
                    │           └── website
                    │               └── replace-website-symbol.ts
                    ├── frontend
                    │   ├── Main
                    │   │   ├── Main.tsx
                    │   │   └── MainView
                    │   │       └── MainView.tsx
                    │   ├── Prompt.tsx
                    │   └── hooks
                    │       └── use-prompt
                    │           └── use-prompt.ts
                    └── types
                        └── messages.ts

Author: robertpiosik

apps/editor/package.json

@@ -460,11 +460,11 @@
         },
         {
           "command": "codeWebChat.selectReferencingFiles",
-          "group": "1_action@7"
+          "group": "2_action@1"
         },
         {
           "command": "codeWebChat.selectDefinitionFile",
-          "group": "1_action@8"
+          "group": "2_action@2"
         }
       ],
       "commandPalette": [

apps/editor/src/views/prompt/backend/message-handlers/handle-get-token-count.ts

@@ -66,4 +66,3 @@ export * from './handle-delete-api-configuration'
 export * from './handle-update-last-used-web-configuration'
 export * from './handle-request-return-home'
 export * from './handle-pick-tasks-workspace'
-export * from './handle-get-token-count'

apps/editor/src/views/prompt/backend/message-handlers/index.ts

@@ -79,8 +79,7 @@ import {
   handle_delete_api_configuration,
   handle_update_last_used_web_configuration,
   handle_request_return_home,
-  handle_pick_tasks_workspace,
-  handle_get_token_count
+  handle_pick_tasks_workspace
 } from './message-handlers'
 import { handle_update_api_configuration } from './message-handlers/handle-update-api-configuration'
 import { handle_pick_model_provider } from './message-handlers/handle-pick-model-provider'
@@ -114,6 +113,13 @@ import { SharedContextState } from '@/context/shared-context-state'
 import { webview_html } from '@/views/shared/utils/webview-html'
 import { get_selected_files } from '@/context/helpers/get-selected-files'
 import { normalize_path } from '@/utils/normalize-path'
+import { replace_symbols } from './utils/symbols/replace-symbols'
+import {
+  EDIT_FORMAT_INSTRUCTIONS_WHOLE,
+  EDIT_FORMAT_INSTRUCTIONS_TRUNCATED,
+  EDIT_FORMAT_INSTRUCTIONS_SEARCH_REPLACE,
+  EDIT_FORMAT_INSTRUCTIONS_DIFF
+} from '@/constants/edit-format-instructions'
 
 export class PromptViewProvider implements vscode.WebviewViewProvider {
   public readonly extension_uri: vscode.Uri
@@ -158,6 +164,7 @@ export class PromptViewProvider implements vscode.WebviewViewProvider {
     workspace_name?: string
   }[] = []
   public api_call_abort_controller: AbortController | null = null
+  public current_selected_files_token_count: number = 0
   public prompt_view_api_calls_manager!: PromptViewApiCallsManager
   public response_history: ResponseHistoryItem[] = []
   public message_listeners: ((message: BackendMessage) => void)[] = []
@@ -398,12 +405,9 @@ export class PromptViewProvider implements vscode.WebviewViewProvider {
     )
 
     token_count_emitter.on('token-count-updated', (token_count: number) => {
+      this.current_selected_files_token_count = token_count
       if (this.webview_view) {
-        this.send_message({
-          command: 'TOKEN_COUNT_UPDATED',
-          token_count
-        })
-
+        this.send_token_count()
         this.send_selected_files()
       }
     })
@@ -536,6 +540,50 @@ export class PromptViewProvider implements vscode.WebviewViewProvider {
     controllers.forEach((item) => item.controller.abort('Preview finished.'))
   }
 
+  public async send_token_count() {
+    const edit_instructions_result = await replace_symbols({
+      instruction: this.current_edit_files_instruction,
+      extension_context: this.extension_context,
+      workspace_provider: this.workspace_provider,
+      remove_images: true
+    })
+    const ask_instructions_result = await replace_symbols({
+      instruction: this.current_ask_about_context_instruction,
+      extension_context: this.extension_context,
+      workspace_provider: this.workspace_provider,
+      remove_images: true
+    })
+
+    let formatted_system_instructions = ''
+    const edit_format_instructions = {
+      whole: EDIT_FORMAT_INSTRUCTIONS_WHOLE,
+      truncated: EDIT_FORMAT_INSTRUCTIONS_TRUNCATED,
+      'search-replace': EDIT_FORMAT_INSTRUCTIONS_SEARCH_REPLACE,
+      diff: EDIT_FORMAT_INSTRUCTIONS_DIFF
+    }[this.edit_format]
+    if (edit_format_instructions) {
+      formatted_system_instructions = `# Output formatting\n\n${edit_format_instructions}`
+    }
+
+    this.send_message({
+      command: 'TOKEN_COUNT_UPDATED',
+      selected_files_token_count: this.current_selected_files_token_count,
+      edit_instructions_token_count: Math.ceil(
+        (edit_instructions_result.instruction.length +
+          edit_instructions_result.skill_definitions.length) /
+          4
+      ),
+      ask_instructions_token_count: Math.ceil(
+        (ask_instructions_result.instruction.length +
+          ask_instructions_result.skill_definitions.length) /
+          4
+      ),
+      edit_format_instructions_token_count: Math.ceil(
+        formatted_system_instructions.length / 4
+      )
+    })
+  }
+
   public send_selected_files() {
     const file_paths = get_selected_files({
       workspace_provider: this.workspace_provider,
@@ -612,6 +660,7 @@ export class PromptViewProvider implements vscode.WebviewViewProvider {
             handle_get_instructions(this)
           } else if (message.command == 'SAVE_INSTRUCTIONS') {
             await handle_save_instructions(this, message)
+            this.send_token_count()
           } else if (message.command == 'GET_CONNECTION_STATUS') {
             handle_get_connection_status(this)
           } else if (message.command == 'GET_WEB_CONFIGURATIONS') {
@@ -650,8 +699,10 @@ export class PromptViewProvider implements vscode.WebviewViewProvider {
             await handle_delete_web_configuration(message)
           } else if (message.command == 'UNDO') {
             await handle_undo(this)
+            this.send_token_count()
           } else if (message.command == 'APPLY_RESPONSE_FROM_HISTORY') {
             await handle_apply_response_from_history(message)
+            this.send_token_count()
           } else if (message.command == 'EXECUTE_COMMAND') {
             await vscode.commands.executeCommand(message.command_id)
           } else if (message.command == 'MAKE_API_CALL') {
@@ -800,7 +851,7 @@ export class PromptViewProvider implements vscode.WebviewViewProvider {
           } else if (message.command == 'PICK_TASKS_WORKSPACE') {
             await handle_pick_tasks_workspace(this, message)
           } else if (message.command == 'GET_TOKEN_COUNT') {
-            await handle_get_token_count(this)
+            this.send_token_count()
           }
         } catch (error: any) {
           Logger.error({
@@ -907,5 +958,6 @@ export class PromptViewProvider implements vscode.WebviewViewProvider {
       edit_files: this.edit_files_instructions,
       caret_position: this.caret_position
     })
+    this.send_token_count()
   }
 }

apps/editor/src/views/prompt/backend/prompt-view-provider.ts

@@ -68,7 +68,7 @@ type Props = {
   on_tabs_reorder: (new_order: number[]) => void
   missing_web_configuration?: boolean
   voice_input_push_to_talk: boolean
-  token_count: number
+  selected_files_token_count: number
   bottom_spacer_height?: number
 }
 
@@ -487,7 +487,7 @@ export const Main: React.FC<Props> = (props) => {
       on_quick_action_click={handle_quick_action_click}
       current_selection={props.current_selection}
       chat_history={current_history || []}
-      token_count={props.token_count}
+      selected_files_token_count={props.selected_files_token_count}
       web_prompt_type={props.web_prompt_type}
       api_prompt_type={props.api_prompt_type}
       on_web_prompt_type_change={props.on_web_prompt_type_change}

apps/editor/src/views/prompt/frontend/Main/Main.tsx

@@ -51,7 +51,7 @@ type Props = {
   currently_open_file_path?: string
   current_selection?: SelectionState | null
   chat_history: string[]
-  token_count: number
+  selected_files_token_count: number
   web_prompt_type: WebPromptType
   api_prompt_type: ApiPromptType
   on_web_prompt_type_change: (prompt_type: WebPromptType) => void
@@ -329,7 +329,7 @@ export const MainView: React.FC<Props> = (props) => {
             on_tabs_reorder={props.on_tabs_reorder}
             warning={warning}
             voice_input_push_to_talk={props.voice_input_push_to_talk}
-            token_count={props.token_count}
+            token_count={props.selected_files_token_count}
             currently_open_file_path={props.currently_open_file_path}
             translations={{
               voice_input: t('prompt-field.voice-input'),

apps/editor/src/views/prompt/frontend/Main/MainView/MainView.tsx

@@ -73,7 +73,7 @@ export const Prompt = () => {
     handle_tabs_reorder,
     voice_input_push_to_talk,
     is_modern_ui,
-    token_count
+    selected_files_token_count
   } = use_panel(vscode)
 
   const {
@@ -349,7 +349,7 @@ export const Prompt = () => {
                 on_tab_delete={handle_tab_delete}
                 on_tabs_reorder={handle_tabs_reorder}
                 voice_input_push_to_talk={voice_input_push_to_talk}
-                token_count={token_count}
+                selected_files_token_count={selected_files_token_count}
               />
             </div>
             <div

apps/editor/src/views/prompt/frontend/Prompt.tsx

@@ -25,7 +25,8 @@ export const use_prompt = (vscode: any) => {
   const [chat_input_focus_key, set_chat_input_focus_key] = useState(0)
   const [chat_input_focus_and_select_key, set_chat_input_focus_and_select_key] =
     useState(0)
-  const [token_count, set_token_count] = useState<number>(0)
+  const [selected_files_token_count, set_selected_files_token_count] =
+    useState<number>(0)
 
   const {
     ask_about_context_instructions,
@@ -119,7 +120,7 @@ export const use_prompt = (vscode: any) => {
       } else if (message.command == 'VOICE_INPUT_PUSH_TO_TALK') {
         set_voice_input_push_to_talk(message.enabled)
       } else if (message.command == 'TOKEN_COUNT_UPDATED') {
-        set_token_count(message.token_count)
+        set_selected_files_token_count(message.selected_files_token_count)
       } else if (message.command == 'IS_MODERN_UI') {
         set_is_modern_ui(message.is_modern_ui)
       }
@@ -236,6 +237,6 @@ export const use_prompt = (vscode: any) => {
     handle_tabs_reorder,
     voice_input_push_to_talk,
     is_modern_ui,
-    token_count
+    selected_files_token_count
   }
 }

apps/editor/src/views/prompt/frontend/hooks/use-prompt/use-prompt.ts

@@ -660,7 +660,10 @@ export interface ResponseHistoryMessage extends BaseMessage {
 
 export interface TokenCountMessage extends BaseMessage {
   command: 'TOKEN_COUNT_UPDATED'
-  token_count: number
+  selected_files_token_count: number
+  edit_instructions_token_count: number
+  ask_instructions_token_count: number
+  edit_format_instructions_token_count: number
 }
 
 export interface WebConfigurationUpdatedMessage extends BaseMessage {