Skip to content

Commit 23d6832

Browse files
committed
feat: Add custom model dimension support and improve codebase indexing
- Add custom embedding dimension configuration for unknown models - Fix Ollama API implementation to use correct prompt-based structure - Enhance OpenAI embedder with base URL and dimension support - Improve UI with provider-specific settings and better model management - Add comprehensive VSCode settings for codebase indexing feature **Core Changes:** - Add `codebaseIndexEmbedderDimension` field to schemas and config - Implement proper dimension validation and fallback logic - Support manual dimension override for custom/unknown models **API Fixes:** - Fix Ollama embedder to use individual `prompt` requests instead of batch `input` - Correct response parsing to use `embedding` (singular) field - Add OpenAI base URL support for proxies and Azure endpoints - Implement dimension parameter for OpenAI embedding requests **UI Improvements:** - Reorganize settings with provider-specific sections - Replace model dropdown with flexible text input - Add dimension input field with validation - Improve placeholder text and help descriptions - Add OpenAI base URL configuration field **Configuration:** - Add VSCode settings for all codebase indexing options - Enhance config manager with proper dimension handling - Improve service factory with better error messages - Add comprehensive validation for custom models
1 parent fe9137b commit 23d6832

File tree

14 files changed

+295
-83
lines changed

14 files changed

+295
-83
lines changed

src/core/webview/webviewMessageHandler.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1327,6 +1327,7 @@ export const webviewMessageHandler = async (provider: ClineProvider, message: We
13271327
codebaseIndexEmbedderProvider: "openai",
13281328
codebaseIndexEmbedderBaseUrl: "",
13291329
codebaseIndexEmbedderModelId: "",
1330+
codebaseIndexEmbedderDimension: null,
13301331
}
13311332
await updateGlobalState("codebaseIndexConfig", codebaseIndexConfig)
13321333

src/exports/roo-code.d.ts

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,7 @@ type GlobalSettings = {
8686
codebaseIndexEmbedderProvider?: ("openai" | "ollama") | undefined
8787
codebaseIndexEmbedderBaseUrl?: string | undefined
8888
codebaseIndexEmbedderModelId?: string | undefined
89+
codebaseIndexEmbedderDimension?: (number | null) | undefined
8990
}
9091
| undefined
9192
alwaysAllowWrite?: boolean | undefined
@@ -865,6 +866,7 @@ type IpcMessage =
865866
codebaseIndexEmbedderProvider?: ("openai" | "ollama") | undefined
866867
codebaseIndexEmbedderBaseUrl?: string | undefined
867868
codebaseIndexEmbedderModelId?: string | undefined
869+
codebaseIndexEmbedderDimension?: (number | null) | undefined
868870
}
869871
| undefined
870872
alwaysAllowWrite?: boolean | undefined
@@ -1377,6 +1379,7 @@ type TaskCommand =
13771379
codebaseIndexEmbedderProvider?: ("openai" | "ollama") | undefined
13781380
codebaseIndexEmbedderBaseUrl?: string | undefined
13791381
codebaseIndexEmbedderModelId?: string | undefined
1382+
codebaseIndexEmbedderDimension?: (number | null) | undefined
13801383
}
13811384
| undefined
13821385
alwaysAllowWrite?: boolean | undefined

src/exports/types.ts

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,7 @@ type GlobalSettings = {
8686
codebaseIndexEmbedderProvider?: ("openai" | "ollama") | undefined
8787
codebaseIndexEmbedderBaseUrl?: string | undefined
8888
codebaseIndexEmbedderModelId?: string | undefined
89+
codebaseIndexEmbedderDimension?: (number | null) | undefined
8990
}
9091
| undefined
9192
alwaysAllowWrite?: boolean | undefined
@@ -879,6 +880,7 @@ type IpcMessage =
879880
codebaseIndexEmbedderProvider?: ("openai" | "ollama") | undefined
880881
codebaseIndexEmbedderBaseUrl?: string | undefined
881882
codebaseIndexEmbedderModelId?: string | undefined
883+
codebaseIndexEmbedderDimension?: (number | null) | undefined
882884
}
883885
| undefined
884886
alwaysAllowWrite?: boolean | undefined
@@ -1393,6 +1395,7 @@ type TaskCommand =
13931395
codebaseIndexEmbedderProvider?: ("openai" | "ollama") | undefined
13941396
codebaseIndexEmbedderBaseUrl?: string | undefined
13951397
codebaseIndexEmbedderModelId?: string | undefined
1398+
codebaseIndexEmbedderDimension?: (number | null) | undefined
13961399
}
13971400
| undefined
13981401
alwaysAllowWrite?: boolean | undefined

src/package.json

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -313,6 +313,57 @@
313313
"type": "string",
314314
"default": "",
315315
"description": "%settings.customStoragePath.description%"
316+
},
317+
"roo-cline.codebaseIndexEnabled": {
318+
"type": "boolean",
319+
"default": false,
320+
"description": "Enable codebase indexing for semantic search.",
321+
"scope": "resource"
322+
},
323+
"roo-cline.codebaseIndexEmbedderProvider": {
324+
"type": "string",
325+
"enum": [
326+
"openai",
327+
"ollama"
328+
],
329+
"default": "openai",
330+
"description": "Select the embedding provider (OpenAI or Ollama).",
331+
"scope": "resource"
332+
},
333+
"roo-cline.codebaseIndexEmbedderBaseUrl": {
334+
"type": "string",
335+
"default": "",
336+
"description": "Optional: Base URL for the selected embedding provider API (e.g., for proxies, Azure OpenAI, or self-hosted Ollama).",
337+
"scope": "resource"
338+
},
339+
"roo-cline.codebaseIndexEmbedderModelId": {
340+
"type": "string",
341+
"default": "",
342+
"description": "Optional: Specify a custom embedding model ID. Leave empty to use the provider's default (e.g., text-embedding-3-small for OpenAI).",
343+
"scope": "resource"
344+
},
345+
"roo-cline.codebaseIndexEmbedderDimension": {
346+
"type": [
347+
"number",
348+
"null"
349+
],
350+
"default": null,
351+
"description": "Optional: Specify the embedding dimension for your custom model ID if it's not automatically recognized. Required only if using an unknown model ID.",
352+
"scope": "resource"
353+
},
354+
"roo-cline.codebaseIndexQdrantUrl": {
355+
"type": "string",
356+
"default": "http://localhost:6333",
357+
"description": "URL for the Qdrant vector database instance.",
358+
"scope": "resource"
359+
},
360+
"roo-cline.codebaseIndexSearchMinScore": {
361+
"type": "number",
362+
"default": 0.4,
363+
"minimum": 0,
364+
"maximum": 1,
365+
"description": "Minimum similarity score (0-1) for codebase search results.",
366+
"scope": "resource"
316367
}
317368
}
318369
}

src/schemas/index.ts

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -232,6 +232,7 @@ export const codebaseIndexConfigSchema = z.object({
232232
codebaseIndexEmbedderProvider: z.enum(["openai", "ollama"]).optional(),
233233
codebaseIndexEmbedderBaseUrl: z.string().optional(),
234234
codebaseIndexEmbedderModelId: z.string().optional(),
235+
codebaseIndexEmbedderDimension: z.number().nullish(),
235236
})
236237

237238
export type CodebaseIndexConfig = z.infer<typeof codebaseIndexConfigSchema>
@@ -244,7 +245,7 @@ export const codebaseIndexModelsSchema = z.object({
244245
export type CodebaseIndexModels = z.infer<typeof codebaseIndexModelsSchema>
245246

246247
export const codebaseIndexProviderSchema = z.object({
247-
codeIndexOpenAiKey: z.string().optional(),
248+
codeIndexOpenAiKey: z.string().optional(),
248249
codeIndexQdrantApiKey: z.string().optional(),
249250
})
250251

@@ -661,7 +662,7 @@ export const providerSettingsSchema = z.object({
661662
...groqSchema.shape,
662663
...chutesSchema.shape,
663664
...litellmSchema.shape,
664-
...codebaseIndexProviderSchema.shape
665+
...codebaseIndexProviderSchema.shape,
665666
})
666667

667668
export type ProviderSettings = z.infer<typeof providerSettingsSchema>

src/services/code-index/config-manager.ts

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ export class CodeIndexConfigManager {
1717
private qdrantUrl?: string = "http://localhost:6333"
1818
private qdrantApiKey?: string
1919
private searchMinScore?: number
20+
private dimension?: number // Added for custom model dimension
2021

2122
constructor(private readonly contextProxy: ContextProxy) {}
2223

@@ -35,6 +36,7 @@ export class CodeIndexConfigManager {
3536
qdrantUrl?: string
3637
qdrantApiKey?: string
3738
searchMinScore?: number
39+
dimension?: number // Add dimension to return type
3840
}
3941
requiresRestart: boolean
4042
}> {
@@ -47,6 +49,7 @@ export class CodeIndexConfigManager {
4749
ollamaBaseUrl: this.ollamaOptions?.ollamaBaseUrl,
4850
qdrantUrl: this.qdrantUrl,
4951
qdrantApiKey: this.qdrantApiKey,
52+
dimension: this.dimension, // Added dimension to snapshot
5053
}
5154

5255
let codebaseIndexConfig = this.contextProxy?.getGlobalState("codebaseIndexConfig") ?? {
@@ -56,31 +59,42 @@ export class CodeIndexConfigManager {
5659
codebaseIndexEmbedderProvider: "openai",
5760
codebaseIndexEmbedderBaseUrl: "",
5861
codebaseIndexEmbedderModelId: "",
62+
codebaseIndexEmbedderDimension: null,
5963
}
6064

65+
// Destructure known properties, access dimension defensively
6166
const {
6267
codebaseIndexEnabled,
6368
codebaseIndexQdrantUrl,
6469
codebaseIndexEmbedderProvider,
6570
codebaseIndexEmbedderBaseUrl,
6671
codebaseIndexEmbedderModelId,
6772
} = codebaseIndexConfig
73+
const codebaseIndexEmbedderDimension = (codebaseIndexConfig as any)?.codebaseIndexEmbedderDimension
6874

6975
const openAiKey = this.contextProxy?.getSecret("codeIndexOpenAiKey") ?? ""
7076
const qdrantApiKey = this.contextProxy?.getSecret("codeIndexQdrantApiKey") ?? ""
7177

7278
this.isEnabled = codebaseIndexEnabled || false
7379
this.qdrantUrl = codebaseIndexQdrantUrl
7480
this.qdrantApiKey = qdrantApiKey ?? ""
75-
this.openAiOptions = { openAiNativeApiKey: openAiKey }
7681
this.searchMinScore = SEARCH_MIN_SCORE
7782

7883
this.embedderProvider = codebaseIndexEmbedderProvider === "ollama" ? "ollama" : "openai"
84+
this.openAiOptions = {
85+
openAiNativeApiKey: openAiKey,
86+
openAiBaseUrl: this.embedderProvider === "openai" ? codebaseIndexEmbedderBaseUrl : undefined,
87+
}
7988
this.modelId = codebaseIndexEmbedderModelId || undefined
8089

8190
this.ollamaOptions = {
8291
ollamaBaseUrl: codebaseIndexEmbedderBaseUrl,
8392
}
93+
// Parse and store dimension, ensuring it's a positive number or undefined
94+
this.dimension =
95+
typeof codebaseIndexEmbedderDimension === "number" && codebaseIndexEmbedderDimension > 0
96+
? codebaseIndexEmbedderDimension
97+
: undefined
8498

8599
return {
86100
configSnapshot: previousConfigSnapshot,
@@ -94,6 +108,7 @@ export class CodeIndexConfigManager {
94108
qdrantUrl: this.qdrantUrl,
95109
qdrantApiKey: this.qdrantApiKey,
96110
searchMinScore: this.searchMinScore,
111+
dimension: this.dimension,
97112
},
98113
requiresRestart: this._didConfigChangeRequireRestart(previousConfigSnapshot),
99114
}
@@ -103,7 +118,6 @@ export class CodeIndexConfigManager {
103118
* Checks if the service is properly configured based on the embedder type.
104119
*/
105120
public isConfigured(): boolean {
106-
107121
if (this.embedderProvider === "openai") {
108122
const openAiKey = this.openAiOptions?.openAiNativeApiKey
109123
const qdrantUrl = this.qdrantUrl
@@ -160,6 +174,9 @@ export class CodeIndexConfigManager {
160174
if (prev.qdrantUrl !== this.qdrantUrl || prev.qdrantApiKey !== this.qdrantApiKey) {
161175
return true
162176
}
177+
178+
// Check dimension change
179+
if (prev.dimension !== this.dimension) return true
163180
}
164181

165182
return false
@@ -179,6 +196,7 @@ export class CodeIndexConfigManager {
179196
qdrantUrl: this.qdrantUrl,
180197
qdrantApiKey: this.qdrantApiKey,
181198
searchMinScore: this.searchMinScore,
199+
dimension: this.dimension,
182200
}
183201
}
184202

src/services/code-index/embedders/ollama.ts

Lines changed: 36 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -22,50 +22,55 @@ export class CodeIndexOllamaEmbedder implements IEmbedder {
2222
*/
2323
async createEmbeddings(texts: string[], model?: string): Promise<EmbeddingResponse> {
2424
const modelToUse = model || this.defaultModelId
25-
const url = `${this.baseUrl}/api/embed` // Endpoint as specified
25+
const url = `${this.baseUrl}/api/embed`
2626

2727
try {
28-
// Note: Standard Ollama API uses 'prompt' for single text, not 'input' for array.
29-
// Implementing based on user's specific request structure.
30-
const response = await fetch(url, {
31-
method: "POST",
32-
headers: {
33-
"Content-Type": "application/json",
34-
},
35-
body: JSON.stringify({
28+
// Ollama API processes one text at a time using 'prompt' field
29+
const embeddings: number[][] = []
30+
31+
for (const text of texts) {
32+
const requestBody = {
3633
model: modelToUse,
37-
input: texts, // Using 'input' as requested
38-
}),
39-
})
34+
prompt: text,
35+
}
36+
37+
const response = await fetch(url, {
38+
method: "POST",
39+
headers: {
40+
"Content-Type": "application/json",
41+
},
42+
body: JSON.stringify(requestBody),
43+
})
4044

41-
if (!response.ok) {
42-
let errorBody = "Could not read error body"
43-
try {
44-
errorBody = await response.text()
45-
} catch (e) {
46-
// Ignore error reading body
45+
if (!response.ok) {
46+
let errorBody = "Could not read error body"
47+
try {
48+
errorBody = await response.text()
49+
} catch (e) {
50+
// Ignore error reading body
51+
}
52+
throw new Error(
53+
`Ollama API request failed with status ${response.status} ${response.statusText}: ${errorBody}`,
54+
)
4755
}
48-
throw new Error(
49-
`Ollama API request failed with status ${response.status} ${response.statusText}: ${errorBody}`,
50-
)
51-
}
5256

53-
const data = await response.json()
57+
const data = await response.json()
58+
59+
// Extract embedding using 'embedding' key (singular)
60+
const embedding = data.embedding
61+
if (!embedding || !Array.isArray(embedding)) {
62+
throw new Error(
63+
'Invalid response structure from Ollama API: "embedding" array not found or not an array.',
64+
)
65+
}
5466

55-
// Extract embeddings using 'embeddings' key as requested
56-
const embeddings = data.embeddings
57-
if (!embeddings || !Array.isArray(embeddings)) {
58-
throw new Error(
59-
'Invalid response structure from Ollama API: "embeddings" array not found or not an array.',
60-
)
67+
embeddings.push(embedding)
6168
}
6269

6370
return {
6471
embeddings: embeddings,
6572
}
6673
} catch (error: any) {
67-
// Log the original error for debugging purposes
68-
console.error("Ollama embedding failed:", error)
6974
// Re-throw a more specific error for the caller
7075
throw new Error(`Ollama embedding failed: ${error.message}`)
7176
}

src/services/code-index/embedders/openai.ts

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,15 +15,21 @@ import {
1515
export class OpenAiEmbedder extends OpenAiNativeHandler implements IEmbedder {
1616
private embeddingsClient: OpenAI
1717
private readonly defaultModelId: string
18+
private readonly configuredDimension?: number
1819

1920
/**
2021
* Creates a new OpenAI embedder
21-
* @param options API handler options
22+
* @param options API handler options including optional model ID and dimension
2223
*/
23-
constructor(options: ApiHandlerOptions & { openAiEmbeddingModelId?: string }) {
24+
constructor(options: ApiHandlerOptions & { openAiEmbeddingModelId?: string; embeddingDimension?: number }) {
2425
super(options)
2526
const apiKey = this.options.openAiNativeApiKey ?? "not-provided"
26-
this.embeddingsClient = new OpenAI({ apiKey })
27+
const baseURL = this.options.openAiBaseUrl
28+
this.configuredDimension = options.embeddingDimension
29+
this.embeddingsClient = new OpenAI({
30+
apiKey,
31+
...(baseURL && { baseURL }),
32+
})
2733
this.defaultModelId = options.openAiEmbeddingModelId || "text-embedding-3-small"
2834
}
2935

@@ -98,10 +104,15 @@ export class OpenAiEmbedder extends OpenAiNativeHandler implements IEmbedder {
98104
): Promise<{ embeddings: number[][]; usage: { promptTokens: number; totalTokens: number } }> {
99105
for (let attempts = 0; attempts < MAX_RETRIES; attempts++) {
100106
try {
101-
const response = await this.embeddingsClient.embeddings.create({
107+
const params: OpenAI.Embeddings.EmbeddingCreateParams = {
102108
input: batchTexts,
103109
model: model,
104-
})
110+
}
111+
if (this.configuredDimension) {
112+
params.dimensions = this.configuredDimension
113+
}
114+
115+
const response = await this.embeddingsClient.embeddings.create(params)
105116

106117
return {
107118
embeddings: response.data.map((item) => item.embedding),

src/services/code-index/interfaces/config.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ export interface CodeIndexConfig {
1414
qdrantUrl?: string
1515
qdrantApiKey?: string
1616
searchMinScore?: number
17+
dimension?: number
1718
}
1819

1920
/**
@@ -28,4 +29,5 @@ export type PreviousConfigSnapshot = {
2829
ollamaBaseUrl?: string
2930
qdrantUrl?: string
3031
qdrantApiKey?: string
32+
dimension?: number
3133
}

0 commit comments

Comments
 (0)