Skip to content

Commit 213b400

Browse files
feat : Embedding configuration enhancement (#185)
1 parent b318558 commit 213b400

File tree

6 files changed

+117
-12
lines changed

6 files changed

+117
-12
lines changed

.changeset/light-laws-grab.md

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
---
2+
"hai-build-code-generator": patch
3+
---
4+
5+
Added Amazon Titan Embed Text V2 model to supported models list
6+
7+
Added configuration support for 'none' as an embedding provider to make it completely optional

src/shared/embeddings.ts

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
export type EmbeddingProvider = "bedrock" | "openai-native" | "openai" | "ollama"
1+
export type EmbeddingProvider = "none" | "bedrock" | "openai-native" | "openai" | "ollama"
22

33
export interface EmbeddingHandlerOptions {
44
modelId?: string
@@ -46,6 +46,11 @@ export const bedrockEmbeddingModels = {
4646
pricePerMillion: 0.0004,
4747
description: "Amazon Titan Text Embeddings model for semantic search and text similarity tasks.",
4848
},
49+
"amazon.titan-embed-text-v2:0": {
50+
maxDimensions: 1024,
51+
pricePerMillion: 0.02,
52+
description: "Amazon Titan Text Embeddings model for semantic search and text similarity tasks with larger context.",
53+
},
4954
} as const
5055

5156
// OpenAI Native
@@ -73,13 +78,17 @@ export const openAiNativeEmbeddingModels = {
7378
} as const
7479

7580
export const embeddingProviderModels = {
81+
none: {},
7682
bedrock: bedrockEmbeddingModels,
7783
"openai-native": openAiNativeEmbeddingModels,
7884
openai: {},
7985
ollama: {},
8086
} as const
8187

8288
export const defaultEmbeddingConfigs: Record<EmbeddingProvider, { defaultModel: string }> = {
89+
none: {
90+
defaultModel: "",
91+
},
8392
bedrock: {
8493
defaultModel: "amazon.titan-embed-text-v1",
8594
},

src/shared/validate.ts

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,8 @@ export function validateApiConfiguration(apiConfiguration?: ApiConfiguration): s
122122
export function validateEmbeddingConfiguration(config?: EmbeddingConfiguration): string | undefined {
123123
if (config) {
124124
switch (config.provider) {
125+
case "none":
126+
break
125127
case "openai-native":
126128
if (!config.openAiNativeApiKey) {
127129
return "You must provide a valid API key."
@@ -136,6 +138,12 @@ export function validateEmbeddingConfiguration(config?: EmbeddingConfiguration):
136138
if (!config.openAiApiKey || !config.openAiBaseUrl || !config.openAiModelId) {
137139
return "You must provide a valid API key, Model ID and base URL."
138140
}
141+
break
142+
case "ollama":
143+
if (!config.ollamaModelId) {
144+
return "You must provide a valid model ID."
145+
}
146+
break
139147
}
140148
}
141149
return undefined

webview-ui/src/components/chat/ChatTextArea.tsx

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -461,6 +461,8 @@ const ChatTextArea = forwardRef<HTMLTextAreaElement, ChatTextAreaProps>(
461461
const expertsButtonRef = useRef<HTMLDivElement>(null)
462462
const [expertsArrowPosition, setExpertsArrowPosition] = useState(0)
463463
const [expertsMenuPosition, setExpertsMenuPosition] = useState(0)
464+
const [isEmbeddingValid, setIsEmbeddingValid] = useState<boolean | null>(null)
465+
const { embeddingConfiguration } = useExtensionState()
464466

465467
// Fetch git commits when Git is selected or when typing a hash
466468
useEffect(() => {
@@ -536,7 +538,9 @@ const ChatTextArea = forwardRef<HTMLTextAreaElement, ChatTextAreaProps>(
536538
}
537539
break
538540
}
539-
541+
case "embeddingConfigValidation":
542+
setIsEmbeddingValid(!!message.bool)
543+
break
540544
default:
541545
console.warn(`Unhandled message type: ${message.type}`)
542546
}
@@ -547,11 +551,16 @@ const ChatTextArea = forwardRef<HTMLTextAreaElement, ChatTextAreaProps>(
547551
// Load experts data on component mount
548552
vscode.postMessage({ type: "loadExperts" })
549553
vscode.postMessage({ type: "loadDefaultExperts" })
554+
if (embeddingConfiguration?.provider === "none") {
555+
setIsEmbeddingValid(null)
556+
return
557+
}
558+
vscode.postMessage({ type: "validateEmbeddingConfig", embeddingConfiguration })
550559

551560
return () => {
552561
window.removeEventListener("message", messageHandler)
553562
}
554-
}, [customExperts, defaultExperts])
563+
}, [customExperts, defaultExperts, embeddingConfiguration])
555564

556565
const queryItems = useMemo(() => {
557566
return [
@@ -1959,12 +1968,18 @@ const ChatTextArea = forwardRef<HTMLTextAreaElement, ChatTextAreaProps>(
19591968
{/* Custom experts without icons */}
19601969
{customExperts.map((expert) => {
19611970
const isProcessing = expert.status === DocumentStatus.PROCESSING
1971+
const isDisabled = isProcessing || (expert.deepCrawl && !isEmbeddingValid)
19621972
return (
19631973
<ExpertItem
19641974
key={expert.name}
19651975
isSelected={selectedExpert?.name === expert.name}
1966-
isDisabled={isProcessing}
1967-
onClick={() => handleExpertSelect(expert)}>
1976+
isDisabled={isDisabled}
1977+
title={
1978+
isDisabled
1979+
? "Expert cannot be accessed due to an invalid embedding configuration."
1980+
: ""
1981+
}
1982+
onClick={() => !isDisabled && handleExpertSelect(expert)}>
19681983
{expert.name}
19691984
{isProcessing && (
19701985
<ExpertTag

webview-ui/src/components/experts/ExpertsView.tsx

Lines changed: 44 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,10 @@ const ExpertsView: React.FC<ExpertsViewProps> = ({ onDone }) => {
7777
window.addEventListener("message", messageHandler)
7878
vscode.postMessage({ type: "loadDefaultExperts" })
7979
vscode.postMessage({ type: "loadExperts" })
80+
if (embeddingConfiguration?.provider === "none") {
81+
setIsEmbeddingValid(null)
82+
return
83+
}
8084
vscode.postMessage({ type: "validateEmbeddingConfig", embeddingConfiguration })
8185
return () => {
8286
window.removeEventListener("message", messageHandler)
@@ -388,7 +392,13 @@ const ExpertsView: React.FC<ExpertsViewProps> = ({ onDone }) => {
388392
text: link.url,
389393
expert: exp.name,
390394
})
391-
}}>
395+
}}
396+
disabled={exp.deepCrawl && !isEmbeddingValid}
397+
title={
398+
exp.deepCrawl && !isEmbeddingValid
399+
? "Document operations are disabled due to invalid embedding configuration"
400+
: undefined
401+
}>
392402
<span className="codicon codicon-refresh" />
393403
</VSCodeButton>
394404
{documentLinkInDeleteConfirmation?.expertName === exp.name &&
@@ -425,7 +435,13 @@ const ExpertsView: React.FC<ExpertsViewProps> = ({ onDone }) => {
425435
expertName: exp.name,
426436
linkUrl: link.url,
427437
})
428-
}}>
438+
}}
439+
disabled={exp.deepCrawl && !isEmbeddingValid}
440+
title={
441+
exp.deepCrawl && !isEmbeddingValid
442+
? "Document operations are disabled due to invalid embedding configuration"
443+
: undefined
444+
}>
429445
<span className="codicon codicon-trash" />
430446
</VSCodeButton>
431447
)}
@@ -538,11 +554,36 @@ const ExpertsView: React.FC<ExpertsViewProps> = ({ onDone }) => {
538554
setInlineEditingDoc({ expertName: exp.name, linkUrl: "" })
539555
setInlineDocLinkError(null) // Reset error when starting inline editing
540556
}}
557+
disabled={exp.deepCrawl && !isEmbeddingValid}
558+
title={
559+
exp.deepCrawl && !isEmbeddingValid
560+
? "Document operations are disabled due to invalid embedding configuration"
561+
: undefined
562+
}
541563
style={{ width: "100%", marginTop: "8px" }}>
542564
+ Add Doc
543565
</VSCodeButton>
544566
)
545567
)}
568+
{exp.deepCrawl && !isEmbeddingValid && (
569+
<div
570+
className="warning-message"
571+
style={{
572+
display: "flex",
573+
alignItems: "center",
574+
gap: "8px",
575+
color: "var(--vscode-editorWarning-foreground)",
576+
fontSize: "11px",
577+
marginTop: "8px",
578+
marginLeft: "20px",
579+
}}>
580+
<i className="codicon codicon-warning" style={{ fontSize: "14px" }} />
581+
<span>
582+
Document operations are disabled due to invalid embedding
583+
configuration
584+
</span>
585+
</div>
586+
)}
546587
</AccordionContainer>
547588
)}
548589
</div>
@@ -779,7 +820,7 @@ const ExpertsView: React.FC<ExpertsViewProps> = ({ onDone }) => {
779820
fontSize: "11px",
780821
}}>
781822
<i className="codicon codicon-warning" style={{ fontSize: "14px" }} />
782-
<span>Valid embedding configuration required for deep crawling</span>
823+
<span>Please provide a valid embedding configuration to enable deep crawling.</span>
783824
</div>
784825
)}
785826
{!isFormReadOnly && (

webview-ui/src/components/settings/EmbeddingOptions.tsx

Lines changed: 29 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import {
99
bedrockeEmbeddingDefaultModelId,
1010
bedrockEmbeddingModels,
1111
azureOpenAIApiVersion,
12+
EmbeddingProvider,
1213
} from "../../../../src/shared/embeddings"
1314
import { useExtensionState } from "../../context/ExtensionStateContext"
1415
import { vscode } from "../../utils/vscode"
@@ -69,13 +70,27 @@ const EmbeddingOptions = ({ showModelOptions, showModelError = true, onValid }:
6970
if (field === "provider") {
7071
// Reset the validation message
7172
setIsEmbeddingValid(null)
73+
74+
// If provider is set to "none", uncheck "Same as LLM API configuration"
75+
if (event.target.value === "none" && buildContextOptions?.useSyncWithApi) {
76+
setBuildContextOptions({
77+
...buildContextOptions,
78+
useSyncWithApi: false,
79+
})
80+
}
7281
}
7382

7483
const newEmbeddingConfiguration = { ...embeddingConfiguration, [field]: event.target.value }
7584
setEmbeddingConfiguration(newEmbeddingConfiguration)
7685
}
7786

7887
useDeepCompareEffect(() => {
88+
// If provider is "none", skip validation and set as valid directly
89+
if (embeddingConfiguration?.provider === "none") {
90+
setIsEmbeddingValid(null)
91+
return
92+
}
93+
7994
const error = validateEmbeddingConfiguration(embeddingConfiguration)
8095

8196
if (error) {
@@ -145,17 +160,20 @@ const EmbeddingOptions = ({ showModelOptions, showModelError = true, onValid }:
145160
useEffect(() => {
146161
setEmbeddingConfiguration({
147162
...embeddingConfiguration,
148-
provider: selectedProvider,
163+
provider: selectedProvider as EmbeddingProvider,
149164
modelId: selectedModelId,
150165
})
151166
// eslint-disable-next-line react-hooks/exhaustive-deps
152167
}, [selectedModelId, selectedProvider])
153168

154169
const availableModels = useMemo(() => {
155-
if (!selectedProvider) {
170+
if (!selectedProvider || selectedProvider === "none") {
156171
return {} as Record<string, EmbeddingModelInfo>
157172
}
158-
return embeddingProviderModels[selectedProvider] as Record<string, EmbeddingModelInfo>
173+
return embeddingProviderModels[selectedProvider as keyof typeof embeddingProviderModels] as Record<
174+
string,
175+
EmbeddingModelInfo
176+
>
159177
}, [selectedProvider])
160178

161179
return (
@@ -170,6 +188,7 @@ const EmbeddingOptions = ({ showModelOptions, showModelError = true, onValid }:
170188
onChange={handleInputChange("provider")}
171189
disabled={isLoading}
172190
style={{ minWidth: 130, position: "relative", width: "100%" }}>
191+
<VSCodeOption value="none">None</VSCodeOption>
173192
<VSCodeOption value="bedrock">AWS Bedrock</VSCodeOption>
174193
<VSCodeOption value="openai-native">OpenAI</VSCodeOption>
175194
<VSCodeOption value="openai">OpenAI Compatible</VSCodeOption>
@@ -476,7 +495,7 @@ const EmbeddingOptions = ({ showModelOptions, showModelError = true, onValid }:
476495
}
477496

478497
export function normalizeEmbeddingConfiguration(embeddingConfiguration?: EmbeddingConfiguration) {
479-
const provider = embeddingConfiguration?.provider || "openai-native"
498+
const provider = embeddingConfiguration?.provider || ("openai-native" as EmbeddingProvider)
480499
const modelId = embeddingConfiguration?.modelId
481500

482501
const getProviderData = (models: Record<string, EmbeddingModelInfo>, defaultId: string) => {
@@ -493,6 +512,12 @@ export function normalizeEmbeddingConfiguration(embeddingConfiguration?: Embeddi
493512
}
494513

495514
switch (provider) {
515+
case "none":
516+
return {
517+
selectedProvider: "none",
518+
selectedModelId: "",
519+
selectedModelInfo: undefined,
520+
}
496521
case "bedrock":
497522
return getProviderData(bedrockEmbeddingModels, bedrockeEmbeddingDefaultModelId)
498523
case "openai-native":

0 commit comments

Comments
 (0)