scriptoriumadmin/Roo-Code
1
0
Fork 0
mirror of https://github.com/RooVetGit/Roo-Code.git synced 2026-08-28 05:27:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

update IBM watsonx embeddeer for code indeix

Browse source
This commit is contained in:
Prasang Prajapati 2025-09-17 11:05:15 -04:00
parent 99bdf52263
commit e196ea1bc8
7 changed files with 217 additions and 140 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
src/api/providers/fetchers/modelCache.ts
View file

@ -83,7 +83,7 @@ export const getModels = async (options: GetModelsOptions): Promise<ModelRecord>
models = await getIOIntelligenceModels(options.apiKey)
break
case "watsonx":
models = await getWatsonxModels(options.apiKey, false)
models = await getWatsonxModels(options.apiKey)
break
default: {
// Ensures router is exhaustively checked if RouterName is a strict union

94
src/api/providers/fetchers/watsonx.ts
View file

@ -16,8 +16,6 @@ import { WatsonXAI } from "@ibm-cloud/watsonx-ai"
*/
export async function getWatsonxModels(
apiKey: string,
embedded: boolean,
projectId?: string,
baseUrl?: string,
platform: "ibmCloud" | "cloudPak" = "ibmCloud",
username?: string,
@ -63,13 +61,7 @@ export async function getWatsonxModels(
let knownModels: Record<string, ModelInfo> = {}
try {
let response
if (embedded) {
response = await service.listFoundationModelSpecs({ filters: "function_embedding" })
} else {
response = await service.listFoundationModelSpecs({ filters: "!function_embedding" })
}
const response = await service.listFoundationModelSpecs({ filters: "!function_embedding" })
if (response && response.result) {
const result = response.result as any
const modelsList = result.resources
@ -113,3 +105,87 @@ export async function getWatsonxModels(
return {}
}
}
/**
* Fetches available embedded watsonx models
*
* @param apiKey - The watsonx API key (for IBM Cloud or Cloud Pak with API key auth)
* @param baseUrl - Optional base URL for the watsonx API
* @param platform - Optional platform type (ibmCloud or cloudPak)
* @param username - Optional username for Cloud Pak for Data
* @param password - Optional password for Cloud Pak for Data (when using password auth)
* @returns A promise resolving to an object with model IDs as keys and model info as values
*/
export async function getEmbeddedWatsonxModels(
apiKey: string,
baseUrl?: string,
platform: "ibmCloud" | "cloudPak" = "ibmCloud",
username?: string,
password?: string,
): Promise<Record<string, { dimension: number }>> {
try {
let options: any = {
version: "2024-05-31",
}
if (platform === "ibmCloud" || !platform) {
if (apiKey) {
options.authenticator = new IamAuthenticator({
apikey: apiKey,
})
} else {
return {}
}
} else if (platform === "cloudPak") {
if (!baseUrl) {
throw new Error("Base URL is required for IBM Cloud Pak for Data")
}
if (username) {
if (password) {
options.authenticator = new CloudPakForDataAuthenticator({
url: baseUrl,
username: username,
password: password,
})
} else if (apiKey) {
options.authenticator = new CloudPakForDataAuthenticator({
url: baseUrl,
username: username,
apikey: apiKey,
})
}
}
}
const service = WatsonXAI.newInstance(options)
let knownModels: Record<string, { dimension: number }> = {}
try {
const response = await service.listFoundationModelSpecs({ filters: "function_embedding" })
if (response && response.result) {
const result = response.result as any
const modelsList = result.models || result.resources || result.foundation_models || []
if (Array.isArray(modelsList)) {
for (const model of modelsList) {
const modelId = model.id || model.name || model.model_id
if (modelId.startsWith("ibm")) {
const dimension = model.model_limits.embedding_dimension
knownModels[modelId] = { dimension }
}
}
}
}
} catch (error) {
console.warn("Error fetching embedded models from IBM watsonx API:", error)
return {}
}
return knownModels
} catch (apiError) {
console.error("Error fetching embedded IBM watsonx models:", apiError)
return {}
}
}

8
src/core/webview/webviewMessageHandler.ts
View file

@ -55,7 +55,7 @@ const ALLOWED_VSCODE_SETTINGS = new Set(["terminal.integrated.inheritEnv"])
import { MarketplaceManager, MarketplaceItemType } from "../../services/marketplace"
import { setPendingTodoList } from "../tools/updateTodoListTool"
import { getWatsonxModels } from "../../api/providers/fetchers/watsonx"
import { getEmbeddedWatsonxModels, getWatsonxModels } from "../../api/providers/fetchers/watsonx"
export const webviewMessageHandler = async (
provider: ClineProvider,
@ -758,8 +758,6 @@ export const webviewMessageHandler = async (
const watsonxModels = await getWatsonxModels(
apiKey,
false,
projectId,
effectiveBaseUrl,
platform,
username,
@ -816,10 +814,8 @@ export const webviewMessageHandler = async (
effectiveBaseUrl = regionToUrl[region] || "https://us-south.ml.cloud.ibm.com"
}
const watsonxModels = await getWatsonxModels(
const watsonxModels = await getEmbeddedWatsonxModels(
apiKey,
true,
projectId,
effectiveBaseUrl,
platform as "ibmCloud" | "cloudPak",
username,

238
src/services/code-index/embedders/watsonx.ts
View file

@ -1,16 +1,12 @@
import { IEmbedder, EmbeddingResponse, EmbedderInfo } from "../interfaces/embedder"
import { MAX_ITEM_TOKENS } from "../constants"
import { t } from "../../../i18n"
import { TelemetryEventName } from "@roo-code/types"
import { TelemetryService } from "@roo-code/telemetry"
import { WatsonXAI } from "@ibm-cloud/watsonx-ai"
import { IamAuthenticator, CloudPakForDataAuthenticator } from "ibm-cloud-sdk-core"
/**
* IBM watsonx embedder implementation using the native IBM Cloud watsonx.ai package.
*
* Supported models:
* - ibm/slate-125m-english-rtrvr-v2 (dimension: 768)
*/
export class WatsonxEmbedder implements IEmbedder {
private readonly watsonxClient: WatsonXAI
@ -91,101 +87,135 @@ export class WatsonxEmbedder implements IEmbedder {
}
/**
* Creates embeddings for the given texts using watsonx's embedding API
* @param texts Array of text strings to embed
* @param model Optional model identifier (uses constructor model if not provided)
* @returns Promise resolving to embedding response
* Gets the expected dimension for a given model ID
* @param modelId The model ID to get the dimension for
* @returns The expected dimension for the model, or 768 if unknown
*/
private getExpectedDimension(modelId: string): number {
// Known dimensions for watsonx models
const knownDimensions: Record<string, number> = {
"ibm/slate-125m-english-rtrvr-v2": 768,
"ibm/slate-125m-english-rtrvr": 768,
"ibm/slate-30m-english-rtrvr-v2": 384,
"ibm/slate-30m-english-rtrvr": 384,
"ibm/granite-embedding-107m-multilingual": 384,
"ibm/granite-embedding-278M-multilingual": 768,
}
return knownDimensions[modelId] || 768
}
async createEmbeddings(texts: string[], model?: string): Promise<EmbeddingResponse> {
const MAX_RETRIES = 3
const INITIAL_DELAY_MS = 1000
try {
const modelToUse = model || this.modelId
const embeddings: number[][] = []
let promptTokens = 0
let totalTokens = 0
for (const text of texts) {
if (!text.trim()) {
embeddings.push([])
continue
}
const estimatedTokens = Math.ceil(text.length / 4)
if (estimatedTokens > MAX_ITEM_TOKENS) {
console.warn(
t("embeddings:textExceedsTokenLimit", {
index: texts.indexOf(text),
itemTokens: estimatedTokens,
maxTokens: MAX_ITEM_TOKENS,
}),
)
embeddings.push([])
continue
}
let lastError
for (let attempt = 0; attempt < MAX_RETRIES; attempt++) {
try {
const response = await this.watsonxClient.embedText({
modelId: modelToUse,
inputs: [text],
projectId: this.projectId,
parameters: {
truncate_input_tokens: MAX_ITEM_TOKENS,
return_options: {
input_text: true,
const MAX_CONCURRENT_REQUESTS = 1
const REQUEST_DELAY_MS = 500
const modelToUse = model || this.modelId
const embeddings: number[][] = []
let promptTokens = 0
let totalTokens = 0
for (let i = 0; i < texts.length; i += MAX_CONCURRENT_REQUESTS) {
const batch = texts.slice(i, i + MAX_CONCURRENT_REQUESTS)
const batchResults = await Promise.all(
batch.map(async (text, batchIndex) => {
const textIndex = i + batchIndex
if (!text.trim()) {
return { index: textIndex, embedding: [], tokens: 0 }
}
const estimatedTokens = Math.ceil(text.length / 4)
if (estimatedTokens > MAX_ITEM_TOKENS) {
console.warn(
t("embeddings:textExceedsTokenLimit", {
index: textIndex,
itemTokens: estimatedTokens,
maxTokens: MAX_ITEM_TOKENS,
}),
)
return { index: textIndex, embedding: [], tokens: 0 }
}
let lastError
for (let attempt = 0; attempt < MAX_RETRIES; attempt++) {
try {
await delay(1000)
const response = await this.watsonxClient.embedText({
modelId: modelToUse,
inputs: [text],
projectId: this.projectId,
parameters: {
truncate_input_tokens: MAX_ITEM_TOKENS,
return_options: {
input_text: true,
},
},
},
})
})
if (response.result && response.result.results && response.result.results.length > 0) {
let embedding = response.result.results[0].embedding
if (!embedding || embedding.length === 0) {
console.error(`Empty embedding returned for text at index ${textIndex}`)
const expectedDimension = this.getExpectedDimension(modelToUse)
if (expectedDimension > 0) {
embedding = new Array(expectedDimension).fill(0.0001)
} else {
throw new Error(`Cannot determine expected dimension for model ${modelToUse}`)
}
}
if (!embedding || embedding.length === 0) {
throw new Error("Failed to create valid embedding")
}
if (response.result && response.result.results && response.result.results.length > 0) {
embeddings.push(response.result.results[0].embedding)
if (response.result.input_token_count) {
promptTokens += response.result.input_token_count
totalTokens += response.result.input_token_count
const tokens = response.result.input_token_count || 0
return { index: textIndex, embedding, tokens }
} else {
console.warn(`No embedding results for text at index ${textIndex}`)
const expectedDimension = this.getExpectedDimension(modelToUse)
if (expectedDimension > 0) {
console.log(`Creating fallback embedding with dimension ${expectedDimension}`)
const fallbackEmbedding = new Array(expectedDimension).fill(0.0001)
return { index: textIndex, embedding: fallbackEmbedding, tokens: 0 }
} else {
return { index: textIndex, embedding: [], tokens: 0 }
}
}
break
} else {
embeddings.push([])
break
}
} catch (error) {
lastError = error
} catch (error) {
lastError = error
if (attempt < MAX_RETRIES - 1) {
const delayMs = INITIAL_DELAY_MS * Math.pow(2, attempt)
console.warn(
`IBM watsonx API call failed, retrying in ${delayMs}ms (attempt ${attempt + 1}/${MAX_RETRIES})`,
)
await new Promise((resolve) => setTimeout(resolve, delayMs))
if (attempt < MAX_RETRIES - 1) {
const delayMs = INITIAL_DELAY_MS * Math.pow(2, attempt)
console.warn(
`IBM watsonx API call failed, retrying in ${delayMs}ms (attempt ${attempt + 1}/${MAX_RETRIES})`,
)
}
}
}
}
if (lastError && embeddings.length < texts.indexOf(text) + 1) {
console.error(
`Failed to embed text at index ${textIndex} after ${MAX_RETRIES} attempts:`,
lastError,
)
return { index: textIndex, embedding: [], tokens: 0 }
}),
)
if (i + MAX_CONCURRENT_REQUESTS < texts.length) {
await new Promise((resolve) => setTimeout(resolve, REQUEST_DELAY_MS * 2))
}
// Process batch results
for (const result of batchResults) {
while (embeddings.length <= result.index) {
embeddings.push([])
console.error(`Failed to embed text after ${MAX_RETRIES} attempts:`, lastError)
}
}
return {
embeddings,
usage: {
promptTokens,
totalTokens,
},
embeddings[result.index] = result.embedding
promptTokens += result.tokens
totalTokens += result.tokens
}
} catch (error) {
TelemetryService.instance.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, {
error: error instanceof Error ? error.message : String(error),
stack: error instanceof Error ? error.stack : undefined,
location: "WatsonxEmbedder:createEmbeddings",
})
throw error
}
return {
embeddings,
usage: {
promptTokens,
totalTokens,
},
}
}
@ -196,9 +226,6 @@ export class WatsonxEmbedder implements IEmbedder {
async validateConfiguration(): Promise<{ valid: boolean; error?: string }> {
try {
const testText = "test"
console.log("Testing IBM watsonx.ai configuration with model:", this.modelId)
const response = await this.watsonxClient.embedText({
modelId: this.modelId,
inputs: [testText],
@ -218,17 +245,9 @@ export class WatsonxEmbedder implements IEmbedder {
error: "embeddings:validation.invalidResponse",
}
}
console.log("IBM watsonx configuration validated successfully")
return { valid: true }
} catch (error) {
console.error("IBM watsonx validation error:", error)
TelemetryService.instance.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, {
error: error instanceof Error ? error.message : String(error),
stack: error instanceof Error ? error.stack : undefined,
location: "WatsonxEmbedder:validateConfiguration",
})
let errorMessage = "embeddings:validation.unknownError"
let errorDetails = ""
@ -246,7 +265,6 @@ export class WatsonxEmbedder implements IEmbedder {
errorMessage = "embeddings:validation.invalidModelId"
}
}
return {
valid: false,
error: `${errorMessage} (${errorDetails})`,
@ -265,16 +283,8 @@ export class WatsonxEmbedder implements IEmbedder {
const knownModels: Record<string, { dimension: number }> = {
"ibm/slate-125m-english-rtrvr-v2": { dimension: 768 },
}
try {
const response = await this.watsonxClient.listFoundationModelSpecs()
console.log(
"IBM watsonx API response structure:",
Object.keys(response || {}).join(", "),
Object.keys(response?.result || {}).join(", "),
)
const response = await this.watsonxClient.listFoundationModelSpecs({ filters: "function_embedding" })
if (response && response.result) {
const result = response.result as any
@ -283,24 +293,14 @@ export class WatsonxEmbedder implements IEmbedder {
if (Array.isArray(modelsList)) {
for (const model of modelsList) {
const modelId = model.id || model.name || model.model_id
const modelInfo = JSON.stringify(model).toLowerCase()
if (
modelId &&
(modelInfo.includes("embed") ||
modelInfo.includes("rtrvr") ||
modelInfo.includes("retriev"))
) {
const dimension = model.dimension || model.vector_size || model.embedding_size || 1536
knownModels[modelId] = { dimension }
}
const dimension = model.model_limits.embedding_dimension
knownModels[modelId] = { dimension }
}
}
}
} catch (apiError) {
console.warn("Error fetching models from IBM watsonx API:", apiError)
}
console.log(`Found ${Object.keys(knownModels).length} IBM watsonx embedding models`)
return knownModels
} catch (error) {
console.error("Error in getAvailableModels:", error)
@ -319,3 +319,7 @@ export class WatsonxEmbedder implements IEmbedder {
}
}
}
function delay(ms: number): Promise<void> {
return new Promise((resolve) => setTimeout(resolve, ms))
}

2
src/shared/ExtensionMessage.ts
View file

@ -152,7 +152,7 @@ export interface ExtensionMessage {
lmStudioModels?: ModelRecord
vsCodeLmModels?: { vendor?: string; family?: string; version?: string; id?: string }[]
watsonxModels?: Record<string, ModelInfo>
embeddedWatsonxModels?: Record<string, ModelInfo>
embeddedWatsonxModels?: Record<string, { dimension: number }>
huggingFaceModels?: Array<{
id: string
object: string

5
src/shared/embeddingModels.ts
View file

@ -54,7 +54,12 @@ export const EMBEDDING_MODEL_PROFILES: EmbeddingModelProfiles = {
"codestral-embed-2505": { dimension: 1536, scoreThreshold: 0.4 },
},
watsonx: {
"ibm/granite-embedding-107m-multilingual": { dimension: 384, scoreThreshold: 0.4 },
"ibm/granite-embedding-278M-multilingual": { dimension: 768, scoreThreshold: 0.4 },
"ibm/slate-125m-english-rtrvr-v2": { dimension: 768, scoreThreshold: 0.4 },
"ibm/slate-125m-english-rtrvr": { dimension: 768, scoreThreshold: 0.4 },
"ibm/slate-30m-english-rtrvr-v2": { dimension: 384, scoreThreshold: 0.4 },
"ibm/slate-30m-english-rtrvr": { dimension: 384, scoreThreshold: 0.4 },
},
}

8
webview-ui/src/components/chat/CodeIndexPopover.tsx
View file

@ -325,7 +325,7 @@ export const CodeIndexPopover: React.FC<CodeIndexPopoverProps> = ({
} else if (event.data.type === "embeddedWatsonxModels" && event.data.embeddedWatsonxModels) {
try {
console.log("Received IBM Embeded watsonx models:", event.data.embeddedWatsonxModels)
const embeddedWatsonxModels: Record<string, { dimension: number }> = {}
let embeddedWatsonxModels: Record<string, { dimension: number }> = {}
if (
!event.data.embeddedWatsonxModels ||
Object.keys(event.data.embeddedWatsonxModels).length === 0
@ -333,11 +333,7 @@ export const CodeIndexPopover: React.FC<CodeIndexPopoverProps> = ({
console.warn("No models received from server, adding default model")
embeddedWatsonxModels["ibm/slate-125m-english-rtrvr-v2"] = { dimension: 768 }
} else {
Object.keys(event.data.embeddedWatsonxModels).forEach((modelId) => {
embeddedWatsonxModels[modelId] = {
dimension: 768,
}
})
embeddedWatsonxModels = event.data.embeddedWatsonxModels
}
if (codebaseIndexModels) {
codebaseIndexModels.watsonx = { ...embeddedWatsonxModels }