Skip to content

Commit cad51ac

Browse files
PriNovaPriNova
committed
feat: Centralize model feature configuration and remove Google Search toggle
* Use client-side config for model features and badges - Replace hardcoded feature flags with model.tags and clientSideConfig - Add SEARCH and IMAGE badge display based on model configuration - Update ModelSelectField to use model.tags for BYOK status * Remove Google Search toggle functionality - Remove isGoogleSearchEnabled parameter from interfaces and components - Clean up related code in ChatBuilder, HumanMessageEditor, and Transcript - Remove chat/google-search message protocol - Delete isGeminiFlash2Model utility function * Simplify feature management in Google Chat client - Determine tools array from model's clientSideConfig - Conditionally render UploadImageButton based on googleImage config - Update Model type to include tags property This consolidates feature flag logic within model configuration for easier maintenance.
1 parent 8196610 commit cad51ac

File tree

15 files changed

+45
-72
lines changed

15 files changed

+45
-72
lines changed

‎lib/shared/src/llm-providers/clients.ts

+1-2
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@ export async function useCustomChatClient({
1010
completionsEndpoint,
1111
logger,
1212
signal,
13-
isGoogleSearchEnabled,
1413
}: ChatNetworkClientParams): Promise<boolean> {
1514
const model = modelsService.getModelByID(params.model ?? '')
1615
if (!model || !isCustomModel(model)) {
@@ -28,7 +27,7 @@ export async function useCustomChatClient({
2827
const client = clientMap[model.provider.toLowerCase()]
2928

3029
if (client) {
31-
await client({ params, cb, completionsEndpoint, logger, signal, isGoogleSearchEnabled })
30+
await client({ params, cb, completionsEndpoint, logger, signal })
3231
return true
3332
}
3433

‎lib/shared/src/llm-providers/google/chat-client.ts

+3-2
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import type { GeminiChatMessage, GeminiCompletionResponse } from '.'
22
import type { ChatNetworkClientParams } from '..'
33
import {
4+
type Model,
45
ModelUsage,
56
contextFiltersProvider,
67
firstValueFrom,
@@ -82,8 +83,8 @@ export async function googleChatClient({
8283
})
8384
}
8485
}
85-
86-
const tools = params.googleSearch ? [{ google_search: {} }] : []
86+
const hasSearch = (model as Model).clientSideConfig?.options?.googleSearch
87+
const tools = hasSearch ? [{ google_search: {} }] : []
8788
const configs = isGeminiThinkModel ? { thinkingConfig: { includeThoughts: true } } : {}
8889

8990
const body = {

‎lib/shared/src/llm-providers/google/utils.ts

-3
Original file line numberDiff line numberDiff line change
@@ -38,9 +38,6 @@ export async function constructGeminiChatMessages(messages: Message[]): Promise<
3838
return geminiMessages.filter((_, i, arr) => i !== arr.length - 1 || arr[i].role !== 'model')
3939
}
4040

41-
export const isGeminiFlash2Model = (model: Model): boolean =>
42-
model?.tags.includes(ModelTag.BYOK) && model?.id.includes('gemini-2.0-flash')
43-
4441
export const isGeminiThinkingModel = (model: Model | undefined | typeof pendingOperation): boolean =>
4542
Boolean(
4643
model &&

‎lib/shared/src/llm-providers/index.ts

-1
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@ export interface ChatNetworkClientParams {
1818
completionsEndpoint: string
1919
logger?: CompletionLogger
2020
signal?: AbortSignal
21-
isGoogleSearchEnabled?: boolean
2221
}
2322

2423
export type ChatNetworkClient = (params: ChatNetworkClientParams) => Promise<void>

‎lib/shared/src/models/sync.ts

+3-3
Original file line numberDiff line numberDiff line change
@@ -437,9 +437,9 @@ function getModelsFromVSCodeConfiguration({
437437
}: PickResolvedConfiguration<{ configuration: 'devModels' }>): Model[] {
438438
return (
439439
devModels?.map(m => {
440-
const isGeminiFlash = m?.model.includes('gemini-2.0-flash')
440+
//const isGeminiFlash = m?.model.includes('gemini-2.0-flash')
441441
const baseTags = [ModelTag.BYOK, ModelTag.Experimental, ModelTag.Local]
442-
const tags = isGeminiFlash ? [...baseTags, ModelTag.Vision] : [...baseTags]
442+
//const tags = isGeminiFlash ? [...baseTags, ModelTag.Vision] : [...baseTags]
443443
return createModel({
444444
id: `${m.provider}/${m.model}`,
445445
usage: [ModelUsage.Chat, ModelUsage.Edit],
@@ -452,7 +452,7 @@ function getModelsFromVSCodeConfiguration({
452452
apiEndpoint: m.apiEndpoint,
453453
options: m.options,
454454
},
455-
tags: tags,
455+
tags: baseTags,
456456
title: m.title,
457457
})
458458
}) ?? []

‎lib/shared/src/sourcegraph-api/completions/types.ts

-1
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,6 @@ export interface CompletionParameters {
5656
model?: string
5757
stream?: boolean
5858
images?: ImageData[]
59-
googleSearch?: boolean
6059
// Configuration for a Predicted Output, which can greatly improve response
6160
// times when large parts of the model response are known ahead of time.
6261
// https://platform.openai.com/docs/guides/latency-optimization#use-predicted-outputs

‎vscode/package.json

+1-1
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
"name": "cody-ai",
44
"private": true,
55
"displayName": "Cody: AI Code Assistant",
6-
"version": "1.69.0+4",
6+
"version": "1.71.0+0",
77
"publisher": "sourcegraph",
88
"license": "Apache-2.0",
99
"icon": "resources/sourcegraph.png",

‎vscode/src/chat/chat-view/ChatBuilder.ts

+1-20
Original file line numberDiff line numberDiff line change
@@ -96,8 +96,7 @@ export class ChatBuilder {
9696
public readonly sessionID: string = new Date(Date.now()).toUTCString(),
9797
private messages: ChatMessage[] = [],
9898
private customChatTitle?: string,
99-
private images: ImageData[] = [],
100-
private isGoogleSearchEnabled = false
99+
private images: ImageData[] = []
101100
) {}
102101

103102
/** An observable that emits whenever the {@link ChatBuilder}'s chat changes. */
@@ -389,24 +388,6 @@ export class ChatBuilder {
389388
// Default to jpeg if unknown
390389
return 'image/jpeg'
391390
}
392-
393-
/**
394-
* Sets the Google search toggle to enabled.
395-
*/
396-
public async setGoogleSearchToggle(): Promise<void> {
397-
this.isGoogleSearchEnabled = true
398-
}
399-
400-
/**
401-
* Retrieves the current state of the Google search toggle and resets it to disabled.
402-
*
403-
* @returns The previous state of the Google search toggle, indicating whether it was enabled or disabled.
404-
*/
405-
public getAndResetGoogleSearchToggle(): boolean {
406-
const isGoogleSearchEnabled = this.isGoogleSearchEnabled
407-
this.isGoogleSearchEnabled = false
408-
return isGoogleSearchEnabled
409-
}
410391
}
411392

412393
function messageToSerializedChatInteraction(

‎vscode/src/chat/chat-view/ChatController.ts

-3
Original file line numberDiff line numberDiff line change
@@ -561,9 +561,6 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
561561
await this.chatBuilder.addImage(message.image)
562562
break
563563
}
564-
case 'chat/google-search': {
565-
await this.chatBuilder.setGoogleSearchToggle()
566-
}
567564
}
568565
}
569566

‎vscode/src/chat/chat-view/handlers/ChatHandler.ts

-1
Original file line numberDiff line numberDiff line change
@@ -155,7 +155,6 @@ export class ChatHandler implements AgentHandler {
155155
model,
156156
maxTokensToSample: contextWindow.output,
157157
images: chatBuilder.getAndResetImages(),
158-
googleSearch: chatBuilder.getAndResetGoogleSearchToggle(),
159158
} as CompletionParameters
160159

161160
// Set stream param only when the model is disabled for streaming.

‎vscode/src/chat/protocol.ts

-4
Original file line numberDiff line numberDiff line change
@@ -177,9 +177,6 @@ export type WebviewMessage =
177177
command: 'openRelativeFile'
178178
uri: Uri
179179
}
180-
| {
181-
command: 'chat/google-search'
182-
}
183180

184181
export interface SmartApplyResult {
185182
taskId: FixupTaskID
@@ -248,7 +245,6 @@ export interface WebviewSubmitMessage extends WebviewContextMessage {
248245
manuallySelectedIntent?: ChatMessage['intent'] | undefined | null
249246
traceparent?: string | undefined | null
250247
steps?: ProcessingStep[] | undefined | null
251-
isGoogleSearchEnabled?: boolean
252248
}
253249

254250
interface WebviewEditMessage extends WebviewContextMessage {

‎vscode/webviews/chat/Transcript.tsx

+1-6
Original file line numberDiff line numberDiff line change
@@ -373,7 +373,6 @@ const TranscriptInteraction: FC<TranscriptInteractionProps> = memo(props => {
373373
} else {
374374
submitHumanMessage({
375375
...commonProps,
376-
isGoogleSearchEnabled: isGoogleSearchEnabled,
377376
})
378377
}
379378
},
@@ -383,7 +382,6 @@ const TranscriptInteraction: FC<TranscriptInteractionProps> = memo(props => {
383382
isLastSentInteraction,
384383
lastEditorRef,
385384
manuallySelectedIntent,
386-
isGoogleSearchEnabled,
387385
]
388386
)
389387
const onEditSubmit = useCallback(
@@ -536,7 +534,7 @@ const TranscriptInteraction: FC<TranscriptInteractionProps> = memo(props => {
536534
}, [humanMessage, assistantMessage, isContextLoading])
537535

538536
const onHumanMessageSubmit = useCallback(
539-
(intent?: ChatMessage['intent'], isGoogleSearchEnabled?: boolean) => {
537+
(intent?: ChatMessage['intent']) => {
540538
if (humanMessage.isUnsentFollowup) {
541539
return onFollowupSubmit(intent)
542540
}
@@ -695,12 +693,10 @@ function submitHumanMessage({
695693
editorValue,
696694
manuallySelectedIntent,
697695
traceparent,
698-
isGoogleSearchEnabled,
699696
}: {
700697
editorValue: SerializedPromptEditorValue
701698
manuallySelectedIntent?: ChatMessage['intent']
702699
traceparent: string
703-
isGoogleSearchEnabled: boolean
704700
}): void {
705701
getVSCodeAPI().postMessage({
706702
command: 'submit',
@@ -709,7 +705,6 @@ function submitHumanMessage({
709705
contextItems: editorValue.contextItems.map(deserializeContextItem),
710706
manuallySelectedIntent,
711707
traceparent,
712-
isGoogleSearchEnabled,
713708
})
714709
focusLastHumanMessageEditor()
715710
}

‎vscode/webviews/chat/cells/messageCell/human/editor/HumanMessageEditor.tsx

+4-5
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ export const HumanMessageEditor: FunctionComponent<{
6666

6767
onEditorFocusChange?: (focused: boolean) => void
6868
onChange?: (editorState: SerializedPromptEditorValue) => void
69-
onSubmit: (intent?: ChatMessage['intent'], isGoogleSearchEnabled?: boolean) => void
69+
onSubmit: (intent?: ChatMessage['intent']) => void
7070
onStop: () => void
7171

7272
isFirstInteraction?: boolean
@@ -223,16 +223,16 @@ export const HumanMessageEditor: FunctionComponent<{
223223
}
224224
setImageFile(undefined)
225225
processImage()
226-
const processGoogleSearch = async () => {
226+
/* const processGoogleSearch = async () => {
227227
if (isGoogleSearchEnabled) {
228228
getVSCodeAPI().postMessage({
229229
command: 'chat/google-search',
230230
})
231231
}
232232
}
233-
processGoogleSearch()
233+
processGoogleSearch() */
234234

235-
parentOnSubmit(intent, isGoogleSearchEnabled)
235+
parentOnSubmit(intent)
236236

237237
telemetryRecorder.recordEvent('cody.humanMessageEditor', 'submit', {
238238
metadata: {
@@ -257,7 +257,6 @@ export const HumanMessageEditor: FunctionComponent<{
257257
isSent,
258258
imageFile,
259259
setImageFile,
260-
isGoogleSearchEnabled,
261260
]
262261
)
263262

‎vscode/webviews/chat/cells/messageCell/human/editor/toolbar/Toolbar.tsx

+3-8
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@ import { PromptSelectField } from '../../../../../../components/promptSelectFiel
88
import { Checkbox } from '../../../../../../components/shadcn/ui/checkbox'
99
import toolbarStyles from '../../../../../../components/shadcn/ui/toolbar.module.css'
1010
import { useActionSelect } from '../../../../../../prompts/PromptsTab'
11-
import { isGeminiFlash2Model } from '../../../../../../utils/modelUtils'
1211
import { useClientConfig } from '../../../../../../utils/useClientConfig'
1312
import { AddContextButton } from './AddContextButton'
1413
import { SubmitButton, type SubmitButtonState } from './SubmitButton'
@@ -85,10 +84,6 @@ export const Toolbar: FunctionComponent<{
8584
[onGapClick]
8685
)
8786

88-
const isGoogleModel = useCallback((model: Model) => {
89-
return isGeminiFlash2Model(model)
90-
}, [])
91-
9287
return (
9388
// biome-ignore lint/a11y/useKeyWithClickEvents: only relevant to click areas
9489
<menu
@@ -105,7 +100,7 @@ export const Toolbar: FunctionComponent<{
105100
>
106101
<div className="tw-flex tw-items-center">
107102
{/* Can't use tw-gap-1 because the popover creates an empty element when open. */}
108-
{isGoogleModel(models[0]) && (
103+
{models[0]?.clientSideConfig?.options?.googleImage && (
109104
<UploadImageButton
110105
className="tw-opacity-60"
111106
imageFile={imageFile}
@@ -124,7 +119,7 @@ export const Toolbar: FunctionComponent<{
124119
userInfo={userInfo}
125120
focusEditor={focusEditor}
126121
className="tw-mr-1"
127-
supportsImageUpload={isGoogleModel(models[0])}
122+
supportsImageUpload={models[0]?.clientSideConfig?.options?.googleImage}
128123
/>
129124

130125
{tokenCount !== undefined &&
@@ -139,7 +134,7 @@ export const Toolbar: FunctionComponent<{
139134
)}
140135
</div>
141136
<div className="tw-flex tw-items-center tw-gap-2">
142-
{isGoogleModel(models[0]) && (
137+
{models[0]?.clientSideConfig?.options?.googleSearch && (
143138
<div className="tw-flex tw-items-center">
144139
<Checkbox
145140
id="google-search-toggle"

‎vscode/webviews/components/modelSelectField/ModelSelectField.tsx

+28-12
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@ import { BookOpenIcon, BrainIcon, BuildingIcon, ExternalLinkIcon } from 'lucide-
55
import { type FunctionComponent, type ReactNode, useCallback, useMemo } from 'react'
66
import type { UserAccountInfo } from '../../Chat'
77
import { getVSCodeAPI } from '../../utils/VSCodeApi'
8-
import { isGeminiFlash2Model } from '../../utils/modelUtils'
98
import { useTelemetryRecorder } from '../../utils/telemetry'
109
import { chatModelIconComponent } from '../ChatModelIcon'
1110
import { Badge } from '../shadcn/ui/badge'
@@ -376,17 +375,36 @@ const ModelTitleWithIcon: React.FC<{
376375
)
377376
) : null}
378377
<span className={clsx('tw-flex-grow', styles.modelName)}>
379-
{isGeminiFlash2Model(model) ? getModelTitle(model.title) : model.title}
378+
{model.tags.includes(ModelTag.BYOK) ? getModelTitle(model.title) : model.title}
380379
</span>
381380
{modelBadge && (
382-
<Badge
383-
variant="secondary"
384-
className={clsx(styles.badge, {
385-
'tw-opacity-75': modelAvailability === 'needs-cody-pro',
386-
})}
387-
>
388-
{model.tags.includes(ModelTag.BYOK) ? 'BYOK' : modelBadge}
389-
</Badge>
381+
<>
382+
<Badge
383+
variant="secondary"
384+
className={clsx(styles.badge, {
385+
'tw-opacity-75': modelAvailability === 'needs-cody-pro',
386+
})}
387+
>
388+
{model.tags.includes(ModelTag.BYOK) &&
389+
model?.clientSideConfig?.options?.googleSearch
390+
? 'SEARCH'
391+
: !model.tags.includes(ModelTag.BYOK)
392+
? modelBadge
393+
: ''}
394+
</Badge>
395+
396+
{model.tags.includes(ModelTag.BYOK) &&
397+
model?.clientSideConfig?.options?.googleImage && (
398+
<Badge
399+
variant="secondary"
400+
className={clsx(styles.badge, {
401+
'tw-opacity-75': modelAvailability === 'needs-cody-pro',
402+
})}
403+
>
404+
IMAGE
405+
</Badge>
406+
)}
407+
</>
390408
)}
391409
</span>
392410
)
@@ -419,7 +437,6 @@ const getModelDropDownUIGroup = (model: Model): string => {
419437
if (model.tags.includes(ModelTag.Speed)) return ModelUIGroup.Speed
420438
if (model.tags.includes(ModelTag.Ollama)) return ModelUIGroup.Ollama
421439
if (model.tags.includes(ModelTag.BYOK)) return ModelUIGroup.BYOK
422-
//if (model.tags.includes(ModelTag.Vision)) return ModelUIGroup.Vision
423440
return ModelUIGroup.Other
424441
}
425442

@@ -431,7 +448,6 @@ const optionByGroup = (
431448
ModelUIGroup.Balanced,
432449
ModelUIGroup.Speed,
433450
ModelUIGroup.BYOK,
434-
//ModelUIGroup.Vision,
435451
ModelUIGroup.Ollama,
436452
ModelUIGroup.Other,
437453
]

0 commit comments

Comments
 (0)