Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 50 additions & 7 deletions src/services/code-index/__tests__/service-factory.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -265,7 +265,7 @@ describe("CodeIndexServiceFactory", () => {
expect(() => factory.createEmbedder()).toThrow("serviceFactory.openAiCompatibleConfigMissing")
})

it("should create GeminiEmbedder when using Gemini provider", () => {
it("should create GeminiEmbedder with default model when no modelId specified", () => {
// Arrange
const testConfig = {
embedderProvider: "gemini",
Expand All @@ -279,7 +279,25 @@ describe("CodeIndexServiceFactory", () => {
factory.createEmbedder()

// Assert
expect(MockedGeminiEmbedder).toHaveBeenCalledWith("test-gemini-api-key")
expect(MockedGeminiEmbedder).toHaveBeenCalledWith("test-gemini-api-key", undefined)
})

it("should create GeminiEmbedder with specified modelId", () => {
// Arrange
const testConfig = {
embedderProvider: "gemini",
modelId: "text-embedding-004",
geminiOptions: {
apiKey: "test-gemini-api-key",
},
}
mockConfigManager.getConfig.mockReturnValue(testConfig as any)

// Act
factory.createEmbedder()

// Assert
expect(MockedGeminiEmbedder).toHaveBeenCalledWith("test-gemini-api-key", "text-embedding-004")
})

it("should throw error when Gemini API key is missing", () => {
Expand Down Expand Up @@ -507,26 +525,51 @@ describe("CodeIndexServiceFactory", () => {
)
})

it("should use fixed dimension 768 for Gemini provider", () => {
it("should use model-specific dimension for Gemini provider", () => {
// Arrange
const testConfig = {
embedderProvider: "gemini",
modelId: "text-embedding-004", // This is ignored by Gemini
modelId: "gemini-embedding-001",
qdrantUrl: "http://localhost:6333",
qdrantApiKey: "test-key",
}
mockConfigManager.getConfig.mockReturnValue(testConfig as any)
mockGetModelDimension.mockReturnValue(3072)

// Act
factory.createVectorStore()

// Assert
// getModelDimension should not be called for Gemini
expect(mockGetModelDimension).not.toHaveBeenCalled()
expect(mockGetModelDimension).toHaveBeenCalledWith("gemini", "gemini-embedding-001")
expect(MockedQdrantVectorStore).toHaveBeenCalledWith(
"/test/workspace",
"http://localhost:6333",
768, // Fixed dimension for Gemini
3072,
"test-key",
)
})

it("should use default model dimension for Gemini when modelId not specified", () => {
// Arrange
const testConfig = {
embedderProvider: "gemini",
qdrantUrl: "http://localhost:6333",
qdrantApiKey: "test-key",
}
mockConfigManager.getConfig.mockReturnValue(testConfig as any)
mockGetDefaultModelId.mockReturnValue("gemini-embedding-001")
mockGetModelDimension.mockReturnValue(3072)

// Act
factory.createVectorStore()

// Assert
expect(mockGetDefaultModelId).toHaveBeenCalledWith("gemini")
expect(mockGetModelDimension).toHaveBeenCalledWith("gemini", "gemini-embedding-001")
expect(MockedQdrantVectorStore).toHaveBeenCalledWith(
"/test/workspace",
"http://localhost:6333",
3072,
"test-key",
)
})
Expand Down
82 changes: 79 additions & 3 deletions src/services/code-index/embedders/__tests__/gemini.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -25,13 +25,30 @@ describe("GeminiEmbedder", () => {
})

describe("constructor", () => {
it("should create an instance with correct fixed values passed to OpenAICompatibleEmbedder", () => {
it("should create an instance with default model when no model specified", () => {
// Arrange
const apiKey = "test-gemini-api-key"

// Act
embedder = new GeminiEmbedder(apiKey)

// Assert
expect(MockedOpenAICompatibleEmbedder).toHaveBeenCalledWith(
"https://generativelanguage.googleapis.com/v1beta/openai/",
apiKey,
"gemini-embedding-001",
2048,
)
})

it("should create an instance with specified model", () => {
// Arrange
const apiKey = "test-gemini-api-key"
const modelId = "text-embedding-004"

// Act
embedder = new GeminiEmbedder(apiKey, modelId)

// Assert
expect(MockedOpenAICompatibleEmbedder).toHaveBeenCalledWith(
"https://generativelanguage.googleapis.com/v1beta/openai/",
Expand All @@ -50,7 +67,7 @@ describe("GeminiEmbedder", () => {
})

describe("embedderInfo", () => {
it("should return correct embedder info with dimension 768", () => {
it("should return correct embedder info", () => {
// Arrange
embedder = new GeminiEmbedder("test-api-key")

Expand All @@ -61,7 +78,66 @@ describe("GeminiEmbedder", () => {
expect(info).toEqual({
name: "gemini",
})
expect(GeminiEmbedder.dimension).toBe(768)
})

describe("createEmbeddings", () => {
let mockCreateEmbeddings: any

beforeEach(() => {
mockCreateEmbeddings = vitest.fn()
MockedOpenAICompatibleEmbedder.prototype.createEmbeddings = mockCreateEmbeddings
})

it("should use instance model when no model parameter provided", async () => {
// Arrange
embedder = new GeminiEmbedder("test-api-key")
const texts = ["test text 1", "test text 2"]
const mockResponse = {
embeddings: [
[0.1, 0.2],
[0.3, 0.4],
],
}
mockCreateEmbeddings.mockResolvedValue(mockResponse)

// Act
const result = await embedder.createEmbeddings(texts)

// Assert
expect(mockCreateEmbeddings).toHaveBeenCalledWith(texts, "gemini-embedding-001")
expect(result).toEqual(mockResponse)
})

it("should use provided model parameter when specified", async () => {
// Arrange
embedder = new GeminiEmbedder("test-api-key", "text-embedding-004")
const texts = ["test text 1", "test text 2"]
const mockResponse = {
embeddings: [
[0.1, 0.2],
[0.3, 0.4],
],
}
mockCreateEmbeddings.mockResolvedValue(mockResponse)

// Act
const result = await embedder.createEmbeddings(texts, "gemini-embedding-001")

// Assert
expect(mockCreateEmbeddings).toHaveBeenCalledWith(texts, "gemini-embedding-001")
expect(result).toEqual(mockResponse)
})

it("should handle errors from OpenAICompatibleEmbedder", async () => {
// Arrange
embedder = new GeminiEmbedder("test-api-key")
const texts = ["test text"]
const error = new Error("Embedding failed")
mockCreateEmbeddings.mockRejectedValue(error)

// Act & Assert
await expect(embedder.createEmbeddings(texts)).rejects.toThrow("Embedding failed")
})
})
})

Expand Down
37 changes: 17 additions & 20 deletions src/services/code-index/embedders/gemini.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,47 +7,51 @@ import { TelemetryService } from "@roo-code/telemetry"

/**
* Gemini embedder implementation that wraps the OpenAI Compatible embedder
* with fixed configuration for Google's Gemini embedding API.
* with configuration for Google's Gemini embedding API.
*
* Fixed values:
* - Base URL: https://generativelanguage.googleapis.com/v1beta/openai/
* - Model: text-embedding-004
* - Dimension: 768
* Supported models:
* - text-embedding-004 (dimension: 768)
* - gemini-embedding-001 (dimension: 2048)
*/
export class GeminiEmbedder implements IEmbedder {
private readonly openAICompatibleEmbedder: OpenAICompatibleEmbedder
private static readonly GEMINI_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/openai/"
private static readonly GEMINI_MODEL = "text-embedding-004"
private static readonly GEMINI_DIMENSION = 768
private static readonly DEFAULT_MODEL = "gemini-embedding-001"
private readonly modelId: string

/**
* Creates a new Gemini embedder
* @param apiKey The Gemini API key for authentication
* @param modelId The model ID to use (defaults to gemini-embedding-001)
*/
constructor(apiKey: string) {
constructor(apiKey: string, modelId?: string) {
if (!apiKey) {
throw new Error(t("embeddings:validation.apiKeyRequired"))
}

// Create an OpenAI Compatible embedder with Gemini's fixed configuration
// Use provided model or default
this.modelId = modelId || GeminiEmbedder.DEFAULT_MODEL

// Create an OpenAI Compatible embedder with Gemini's configuration
this.openAICompatibleEmbedder = new OpenAICompatibleEmbedder(
GeminiEmbedder.GEMINI_BASE_URL,
apiKey,
GeminiEmbedder.GEMINI_MODEL,
this.modelId,
GEMINI_MAX_ITEM_TOKENS,
)
}

/**
* Creates embeddings for the given texts using Gemini's embedding API
* @param texts Array of text strings to embed
* @param model Optional model identifier (ignored - always uses text-embedding-004)
* @param model Optional model identifier (uses constructor model if not provided)
* @returns Promise resolving to embedding response
*/
async createEmbeddings(texts: string[], model?: string): Promise<EmbeddingResponse> {
try {
// Always use the fixed Gemini model, ignoring any passed model parameter
return await this.openAICompatibleEmbedder.createEmbeddings(texts, GeminiEmbedder.GEMINI_MODEL)
// Use the provided model or fall back to the instance's model
const modelToUse = model || this.modelId
return await this.openAICompatibleEmbedder.createEmbeddings(texts, modelToUse)
} catch (error) {
TelemetryService.instance.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, {
error: error instanceof Error ? error.message : String(error),
Expand Down Expand Up @@ -85,11 +89,4 @@ export class GeminiEmbedder implements IEmbedder {
name: "gemini",
}
}

/**
* Gets the fixed dimension for Gemini embeddings
*/
static get dimension(): number {
return GeminiEmbedder.GEMINI_DIMENSION
}
}
5 changes: 1 addition & 4 deletions src/services/code-index/service-factory.ts
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ export class CodeIndexServiceFactory {
if (!config.geminiOptions?.apiKey) {
throw new Error(t("embeddings:serviceFactory.geminiConfigMissing"))
}
return new GeminiEmbedder(config.geminiOptions.apiKey)
return new GeminiEmbedder(config.geminiOptions.apiKey, config.modelId)
}

throw new Error(
Expand Down Expand Up @@ -111,9 +111,6 @@ export class CodeIndexServiceFactory {
// First check if a manual dimension is provided (works for all providers)
if (config.modelDimension && config.modelDimension > 0) {
vectorSize = config.modelDimension
} else if (provider === "gemini") {
// Gemini's text-embedding-004 has a fixed dimension of 768
vectorSize = 768
} else {
// Fall back to model-specific dimension from profiles
vectorSize = getModelDimension(provider, modelId)
Expand Down
3 changes: 2 additions & 1 deletion src/shared/embeddingModels.ts
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ export const EMBEDDING_MODEL_PROFILES: EmbeddingModelProfiles = {
},
gemini: {
"text-embedding-004": { dimension: 768 },
"gemini-embedding-001": { dimension: 3072, scoreThreshold: 0.4 },
},
}

Expand Down Expand Up @@ -134,7 +135,7 @@ export function getDefaultModelId(provider: EmbedderProvider): string {
}

case "gemini":
return "text-embedding-004"
return "gemini-embedding-001"

default:
// Fallback for unknown providers
Expand Down
Loading