diff --git a/vscode/CHANGELOG.md b/vscode/CHANGELOG.md index f21c4954677..a8b103d3dc0 100644 --- a/vscode/CHANGELOG.md +++ b/vscode/CHANGELOG.md @@ -15,6 +15,7 @@ This is a log of all notable changes to Cody for VS Code. [Unreleased] changes a - Chat: Large file cannot be added via @-mention. [pull/3531](https://github.com/sourcegraph/cody/pull/3531) - Chat: Handle empty chat message input and prevent submission of empty messages. [pull/3554](https://github.com/sourcegraph/cody/pull/3554) +- Chat: Warnings are now displayed correctly for large files in the @-mention file selection list. [pull/3526](https://github.com/sourcegraph/cody/pull/3526) ### Changed diff --git a/vscode/src/chat/chat-view/SimpleChatModel.ts b/vscode/src/chat/chat-view/SimpleChatModel.ts index fa75b655c58..2d11bf61d03 100644 --- a/vscode/src/chat/chat-view/SimpleChatModel.ts +++ b/vscode/src/chat/chat-view/SimpleChatModel.ts @@ -4,6 +4,7 @@ import { type ChatMessage, type ContextItem, type Message, + ModelProvider, type SerializedChatInteraction, type SerializedChatTranscript, errorToChatError, @@ -15,13 +16,17 @@ import type { Repo } from '../../context/repo-fetcher' import { getChatPanelTitle } from './chat-helpers' export class SimpleChatModel { + // The maximum number of characters available in the model's context window. + public readonly maxChars: number constructor( public modelID: string, private messages: ChatMessage[] = [], public readonly sessionID: string = new Date(Date.now()).toUTCString(), private customChatTitle?: string, private selectedRepos?: Repo[] - ) {} + ) { + this.maxChars = ModelProvider.getMaxCharsByModel(this.modelID) + } public isEmpty(): boolean { return this.messages.length === 0 diff --git a/vscode/src/chat/chat-view/SimpleChatPanelProvider.ts b/vscode/src/chat/chat-view/SimpleChatPanelProvider.ts index e33e185ba36..6e6e40b5577 100644 --- a/vscode/src/chat/chat-view/SimpleChatPanelProvider.ts +++ b/vscode/src/chat/chat-view/SimpleChatPanelProvider.ts @@ -46,12 +46,13 @@ import { } from '../../services/utils/codeblock-action-tracker' import { openExternalLinks, openLocalFileWithRange } from '../../services/utils/workspace-action' import { TestSupport } from '../../test-support' -import { countGeneratedCode } from '../utils' +import { countGeneratedCode, getContextWindowLimitInBytes } from '../utils' import type { Span } from '@opentelemetry/api' import { captureException } from '@sentry/core' import type { ContextItemWithContent } from '@sourcegraph/cody-shared/src/codebase-context/messages' import { ModelUsage } from '@sourcegraph/cody-shared/src/models/types' +import { ANSWER_TOKENS, tokensToChars } from '@sourcegraph/cody-shared/src/prompt/constants' import { recordErrorToSpan, tracer } from '@sourcegraph/cody-shared/src/tracing' import type { EnterpriseContextFactory } from '../../context/enterprise-context-factory' import type { Repo } from '../../context/repo-fetcher' @@ -596,12 +597,20 @@ export class SimpleChatPanelProvider implements vscode.Disposable, ChatSession { }) }, } + // Use the number of characters left in the chat model as the limit + // for adding user context files to the chat. + const contextLimit = getContextWindowLimitInBytes( + [...this.chatModel.getMessages()], + // Minus the character limit reserved for the answer token + this.chatModel.maxChars - tokensToChars(ANSWER_TOKENS) + ) try { const items = await getChatContextItemsForMention( query, cancellation.token, - scopedTelemetryRecorder + scopedTelemetryRecorder, + contextLimit ) if (cancellation.token.isCancellationRequested) { return @@ -781,7 +790,7 @@ export class SimpleChatPanelProvider implements vscode.Disposable, ChatSession { ): Promise { const { prompt, newContextUsed, newContextIgnored } = await prompter.makePrompt( this.chatModel, - ModelProvider.getMaxCharsByModel(this.chatModel.modelID) + this.chatModel.maxChars ) // Update UI based on prompt construction diff --git a/vscode/src/chat/context/chatContext.ts b/vscode/src/chat/context/chatContext.ts index ff7adee73dc..10dd0c4c2a7 100644 --- a/vscode/src/chat/context/chatContext.ts +++ b/vscode/src/chat/context/chatContext.ts @@ -19,7 +19,9 @@ export async function getChatContextItemsForMention( telemetryRecorder?: { empty: () => void withType: (type: MentionQuery['type']) => void - } + }, + // The number of characters left in current context window. + maxChars?: number ): Promise { const mentionQuery = parseMentionQuery(query) @@ -35,12 +37,12 @@ export async function getChatContextItemsForMention( const MAX_RESULTS = 20 switch (mentionQuery.type) { case 'empty': - return getOpenTabsContextFile() + return getOpenTabsContextFile(maxChars) case 'symbol': // It would be nice if the VS Code symbols API supports cancellation, but it doesn't return getSymbolContextFiles(mentionQuery.text, MAX_RESULTS) case 'file': - return getFileContextFiles(mentionQuery.text, MAX_RESULTS, cancellationToken) + return getFileContextFiles(mentionQuery.text, MAX_RESULTS, maxChars) case 'url': return (await isURLContextFeatureFlagEnabled()) ? getURLContextItems( diff --git a/vscode/src/chat/utils.ts b/vscode/src/chat/utils.ts index e5c0b954871..1a785ea1116 100644 --- a/vscode/src/chat/utils.ts +++ b/vscode/src/chat/utils.ts @@ -1,4 +1,4 @@ -import type { AuthStatus } from '@sourcegraph/cody-shared' +import type { AuthStatus, ChatMessage } from '@sourcegraph/cody-shared' import { defaultAuthStatus, unauthenticatedStatus } from './protocol' /** @@ -82,3 +82,35 @@ export const countGeneratedCode = (text: string): { lineCount: number; charCount } return count } + +/** + * Counts the total number of bytes used in a list of chat messages. + * + * This function is exported and can be used to calculate the byte usage + * of chat messages for storage/bandwidth purposes. + * + * @param messages - The list of chat messages to count bytes for + * @returns The total number of bytes used in the messages + */ +export function countBytesInChatMessages(messages: ChatMessage[]): number { + if (messages.length === 0) { + return 0 + } + return messages.reduce((acc, msg) => acc + msg.speaker.length + (msg.text?.length || 0) + 3, 0) +} + +/** + * Gets the context window limit in bytes for chat messages, taking into + * account the maximum allowed character count. Returns 0 if the used bytes + * exceeds the limit. + * @param messages - The chat messages + * @param maxChars - The maximum allowed character count + * @returns The context window limit in bytes + */ +export function getContextWindowLimitInBytes(messages: ChatMessage[], maxChars: number): number { + const used = countBytesInChatMessages(messages) + if (used > maxChars) { + return 0 + } + return maxChars - used +} diff --git a/vscode/src/edit/input/get-matching-context.ts b/vscode/src/edit/input/get-matching-context.ts index 67830fddb07..1fbbe62cd38 100644 --- a/vscode/src/edit/input/get-matching-context.ts +++ b/vscode/src/edit/input/get-matching-context.ts @@ -1,6 +1,6 @@ import type { ContextItem, MentionQuery } from '@sourcegraph/cody-shared' -import * as vscode from 'vscode' +import { DEFAULT_FAST_MODEL_CHARS_LIMIT } from '@sourcegraph/cody-shared/src/prompt/constants' import { getFileContextFiles, getSymbolContextFiles } from '../../editor/utils/editor-context' import { getLabelForContextItem } from './utils' @@ -29,11 +29,10 @@ export async function getMatchingContext( } if (mentionQuery.type === 'file') { - const cancellation = new vscode.CancellationTokenSource() const fileResults = await getFileContextFiles( mentionQuery.text, MAX_FUZZY_RESULTS, - cancellation.token + DEFAULT_FAST_MODEL_CHARS_LIMIT ) return fileResults.map(result => ({ key: getLabelForContextItem(result), diff --git a/vscode/src/editor/utils/editor-context.test.ts b/vscode/src/editor/utils/editor-context.test.ts index b6e739a0a76..87518cbfa04 100644 --- a/vscode/src/editor/utils/editor-context.test.ts +++ b/vscode/src/editor/utils/editor-context.test.ts @@ -41,11 +41,7 @@ describe('getFileContextFiles', () => { } async function runSearch(query: string, maxResults: number): Promise<(string | undefined)[]> { - const results = await getFileContextFiles( - query, - maxResults, - new vscode.CancellationTokenSource().token - ) + const results = await getFileContextFiles(query, maxResults) return results.map(f => uriBasename(f.uri)) } @@ -158,13 +154,14 @@ describe('filterLargeFiles', () => { expect(filtered).toEqual([]) }) - it('sets isTooLarge for files exceeding token limit', async () => { + it('sets isTooLarge for files exceeding token limit but under 1MB', async () => { const largeTextFile: ContextItemFile = { uri: vscode.Uri.file('/large-text.txt'), type: 'file', } + const oneByteOverTokenLimit = MAX_CURRENT_FILE_TOKENS * CHARS_PER_TOKEN + 1 vscode.workspace.fs.stat = vi.fn().mockResolvedValueOnce({ - size: MAX_CURRENT_FILE_TOKENS * CHARS_PER_TOKEN + 1, + size: oneByteOverTokenLimit, type: vscode.FileType.File, } as vscode.FileStat) diff --git a/vscode/src/editor/utils/editor-context.ts b/vscode/src/editor/utils/editor-context.ts index c10bad3caa0..f3bbc396527 100644 --- a/vscode/src/editor/utils/editor-context.ts +++ b/vscode/src/editor/utils/editor-context.ts @@ -56,7 +56,7 @@ const throttledFindFiles = throttle(() => findWorkspaceFiles(), 10000) export async function getFileContextFiles( query: string, maxResults: number, - token: vscode.CancellationToken + charsLimit?: number ): Promise { if (!query.trim()) { return [] @@ -116,7 +116,7 @@ export async function getFileContextFiles( // TODO(toolmantim): Add fuzzysort.highlight data to the result so we can show it in the UI - return await filterLargeFiles(sortedResults) + return await filterLargeFiles(sortedResults, charsLimit) } export async function getSymbolContextFiles( @@ -182,11 +182,12 @@ export async function getSymbolContextFiles( * Gets context files for each open editor tab in VS Code. * Filters out large files over 1MB to avoid expensive parsing. */ -export async function getOpenTabsContextFile(): Promise { +export async function getOpenTabsContextFile(charsLimit?: number): Promise { return await filterLargeFiles( getOpenTabsUris() .filter(uri => !isCodyIgnoredFile(uri)) - .flatMap(uri => createContextFileFromUri(uri, ContextItemSource.User, 'file')) + .flatMap(uri => createContextFileFromUri(uri, ContextItemSource.User, 'file')), + charsLimit ) } @@ -253,7 +254,10 @@ function createContextFileRange(selectionRange: vscode.Range): ContextItem['rang * Filters the given context files to remove files larger than 1MB and non-text files. * Sets {@link ContextItemFile.isTooLarge} for files contains more characters than the token limit. */ -export async function filterLargeFiles(contextFiles: ContextItemFile[]): Promise { +export async function filterLargeFiles( + contextFiles: ContextItemFile[], + charsLimit = CHARS_PER_TOKEN * MAX_CURRENT_FILE_TOKENS +): Promise { const filtered = [] for (const cf of contextFiles) { // Remove file larger than 1MB and non-text files @@ -262,13 +266,13 @@ export async function filterLargeFiles(contextFiles: ContextItemFile[]): Promise stat => stat, error => undefined ) - if (fileStat?.type !== vscode.FileType.File || fileStat?.size > 1000000) { + if (cf.type !== 'file' || fileStat?.type !== vscode.FileType.File || fileStat?.size > 1000000) { continue } // Check if file contains more characters than the token limit based on fileStat.size // and set {@link ContextItemFile.isTooLarge} for webview to display file size // warning. - if (fileStat.size > CHARS_PER_TOKEN * MAX_CURRENT_FILE_TOKENS) { + if (fileStat.size > charsLimit) { cf.isTooLarge = true } filtered.push(cf)