mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-09-07 08:26:51 +00:00
fix: resolve Ollama codebase indexing freeze by correcting API usage
- Fix Ollama embedder to use correct API format: "prompt" instead of "input" - Handle response structure correctly: "embedding" instead of "embeddings" - Implement sequential processing for batch embeddings since Ollama API processes one text at a time - Update validation test to use correct API format - Add comprehensive tests for createEmbeddings method Fixes #5823
This commit is contained in:
parent
fb374b3e94
commit
c587a271ff
2 changed files with 138 additions and 40 deletions
|
|
@ -103,7 +103,7 @@ describe("CodeIndexOllamaEmbedder", () => {
|
|||
status: 200,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
embeddings: [[0.1, 0.2, 0.3]],
|
||||
embedding: [0.1, 0.2, 0.3],
|
||||
}),
|
||||
} as Response),
|
||||
)
|
||||
|
|
@ -126,7 +126,7 @@ describe("CodeIndexOllamaEmbedder", () => {
|
|||
expect(secondCall[0]).toBe("http://localhost:11434/api/embed")
|
||||
expect(secondCall[1]?.method).toBe("POST")
|
||||
expect(secondCall[1]?.headers).toEqual({ "Content-Type": "application/json" })
|
||||
expect(secondCall[1]?.body).toBe(JSON.stringify({ model: "nomic-embed-text", input: ["test"] }))
|
||||
expect(secondCall[1]?.body).toBe(JSON.stringify({ model: "nomic-embed-text", prompt: "test" }))
|
||||
expect(secondCall[1]?.signal).toBeDefined() // AbortSignal for timeout
|
||||
})
|
||||
|
||||
|
|
@ -240,4 +240,94 @@ describe("CodeIndexOllamaEmbedder", () => {
|
|||
expect(result.error).toBe("Network timeout")
|
||||
})
|
||||
})
|
||||
|
||||
describe("createEmbeddings", () => {
|
||||
it("should create embeddings for multiple texts using sequential calls", async () => {
|
||||
// Mock successful responses for each individual text
|
||||
mockFetch
|
||||
.mockImplementationOnce(() =>
|
||||
Promise.resolve({
|
||||
ok: true,
|
||||
status: 200,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
embedding: [0.1, 0.2, 0.3],
|
||||
}),
|
||||
} as Response),
|
||||
)
|
||||
.mockImplementationOnce(() =>
|
||||
Promise.resolve({
|
||||
ok: true,
|
||||
status: 200,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
embedding: [0.4, 0.5, 0.6],
|
||||
}),
|
||||
} as Response),
|
||||
)
|
||||
|
||||
const result = await embedder.createEmbeddings(["text1", "text2"])
|
||||
|
||||
expect(result.embeddings).toEqual([
|
||||
[0.1, 0.2, 0.3],
|
||||
[0.4, 0.5, 0.6],
|
||||
])
|
||||
expect(mockFetch).toHaveBeenCalledTimes(2)
|
||||
|
||||
// Check first call
|
||||
const firstCall = mockFetch.mock.calls[0]
|
||||
expect(firstCall[0]).toBe("http://localhost:11434/api/embed")
|
||||
expect(firstCall[1]?.method).toBe("POST")
|
||||
expect(firstCall[1]?.headers).toEqual({ "Content-Type": "application/json" })
|
||||
expect(firstCall[1]?.body).toBe(JSON.stringify({ model: "nomic-embed-text", prompt: "text1" }))
|
||||
expect(firstCall[1]?.signal).toBeDefined()
|
||||
|
||||
// Check second call
|
||||
const secondCall = mockFetch.mock.calls[1]
|
||||
expect(secondCall[0]).toBe("http://localhost:11434/api/embed")
|
||||
expect(secondCall[1]?.method).toBe("POST")
|
||||
expect(secondCall[1]?.headers).toEqual({ "Content-Type": "application/json" })
|
||||
expect(secondCall[1]?.body).toBe(JSON.stringify({ model: "nomic-embed-text", prompt: "text2" }))
|
||||
expect(secondCall[1]?.signal).toBeDefined()
|
||||
})
|
||||
|
||||
it("should handle errors during embedding creation", async () => {
|
||||
mockFetch.mockRejectedValueOnce(new Error("ECONNREFUSED"))
|
||||
|
||||
await expect(embedder.createEmbeddings(["test"])).rejects.toThrow(
|
||||
"embeddings:ollama.serviceNotRunning",
|
||||
)
|
||||
})
|
||||
|
||||
it("should handle invalid response structure", async () => {
|
||||
mockFetch.mockImplementationOnce(() =>
|
||||
Promise.resolve({
|
||||
ok: true,
|
||||
status: 200,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
// Missing 'embedding' field
|
||||
invalid: "response",
|
||||
}),
|
||||
} as Response),
|
||||
)
|
||||
|
||||
await expect(embedder.createEmbeddings(["test"])).rejects.toThrow(
|
||||
"embeddings:ollama.invalidResponseStructure",
|
||||
)
|
||||
})
|
||||
|
||||
it("should handle HTTP error responses", async () => {
|
||||
mockFetch.mockImplementationOnce(() =>
|
||||
Promise.resolve({
|
||||
ok: false,
|
||||
status: 500,
|
||||
statusText: "Internal Server Error",
|
||||
text: () => Promise.resolve("Server error details"),
|
||||
} as Response),
|
||||
)
|
||||
|
||||
await expect(embedder.createEmbeddings(["test"])).rejects.toThrow("embeddings:ollama.requestFailed")
|
||||
})
|
||||
})
|
||||
})
|
||||
|
|
|
|||
|
|
@ -26,13 +26,14 @@ export class CodeIndexOllamaEmbedder implements IEmbedder {
|
|||
|
||||
/**
|
||||
* Creates embeddings for the given texts using the specified Ollama model.
|
||||
* Ollama's /api/embed endpoint processes one text at a time, so we make sequential calls.
|
||||
* @param texts - An array of strings to embed.
|
||||
* @param model - Optional model ID to override the default.
|
||||
* @returns A promise that resolves to an EmbeddingResponse containing the embeddings and usage data.
|
||||
*/
|
||||
async createEmbeddings(texts: string[], model?: string): Promise<EmbeddingResponse> {
|
||||
const modelToUse = model || this.defaultModelId
|
||||
const url = `${this.baseUrl}/api/embed` // Endpoint as specified
|
||||
const url = `${this.baseUrl}/api/embed`
|
||||
|
||||
// Apply model-specific query prefix if required
|
||||
const queryPrefix = getModelQueryPrefix("ollama", modelToUse)
|
||||
|
|
@ -60,48 +61,55 @@ export class CodeIndexOllamaEmbedder implements IEmbedder {
|
|||
: texts
|
||||
|
||||
try {
|
||||
// Note: Standard Ollama API uses 'prompt' for single text, not 'input' for array.
|
||||
// Implementing based on user's specific request structure.
|
||||
const embeddings: number[][] = []
|
||||
|
||||
// Add timeout to prevent indefinite hanging
|
||||
const controller = new AbortController()
|
||||
const timeoutId = setTimeout(() => controller.abort(), OLLAMA_EMBEDDING_TIMEOUT_MS)
|
||||
// Process each text individually since Ollama's /api/embed endpoint
|
||||
// expects a single 'prompt' field, not an array of inputs
|
||||
for (let i = 0; i < processedTexts.length; i++) {
|
||||
const text = processedTexts[i]
|
||||
|
||||
const response = await fetch(url, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
body: JSON.stringify({
|
||||
model: modelToUse,
|
||||
input: processedTexts, // Using 'input' as requested
|
||||
}),
|
||||
signal: controller.signal,
|
||||
})
|
||||
clearTimeout(timeoutId)
|
||||
// Add timeout to prevent indefinite hanging
|
||||
const controller = new AbortController()
|
||||
const timeoutId = setTimeout(() => controller.abort(), OLLAMA_EMBEDDING_TIMEOUT_MS)
|
||||
|
||||
if (!response.ok) {
|
||||
let errorBody = t("embeddings:ollama.couldNotReadErrorBody")
|
||||
try {
|
||||
errorBody = await response.text()
|
||||
} catch (e) {
|
||||
// Ignore error reading body
|
||||
}
|
||||
throw new Error(
|
||||
t("embeddings:ollama.requestFailed", {
|
||||
status: response.status,
|
||||
statusText: response.statusText,
|
||||
errorBody,
|
||||
const response = await fetch(url, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
body: JSON.stringify({
|
||||
model: modelToUse,
|
||||
prompt: text, // Ollama expects 'prompt', not 'input'
|
||||
}),
|
||||
)
|
||||
}
|
||||
signal: controller.signal,
|
||||
})
|
||||
clearTimeout(timeoutId)
|
||||
|
||||
const data = await response.json()
|
||||
if (!response.ok) {
|
||||
let errorBody = t("embeddings:ollama.couldNotReadErrorBody")
|
||||
try {
|
||||
errorBody = await response.text()
|
||||
} catch (e) {
|
||||
// Ignore error reading body
|
||||
}
|
||||
throw new Error(
|
||||
t("embeddings:ollama.requestFailed", {
|
||||
status: response.status,
|
||||
statusText: response.statusText,
|
||||
errorBody,
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
// Extract embeddings using 'embeddings' key as requested
|
||||
const embeddings = data.embeddings
|
||||
if (!embeddings || !Array.isArray(embeddings)) {
|
||||
throw new Error(t("embeddings:ollama.invalidResponseStructure"))
|
||||
const data = await response.json()
|
||||
|
||||
// Ollama returns 'embedding' (singular), not 'embeddings' (plural)
|
||||
const embedding = data.embedding
|
||||
if (!embedding || !Array.isArray(embedding)) {
|
||||
throw new Error(t("embeddings:ollama.invalidResponseStructure"))
|
||||
}
|
||||
|
||||
embeddings.push(embedding)
|
||||
}
|
||||
|
||||
return {
|
||||
|
|
@ -210,7 +218,7 @@ export class CodeIndexOllamaEmbedder implements IEmbedder {
|
|||
},
|
||||
body: JSON.stringify({
|
||||
model: this.defaultModelId,
|
||||
input: ["test"],
|
||||
prompt: "test", // Use 'prompt' instead of 'input' array
|
||||
}),
|
||||
signal: testController.signal,
|
||||
})
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue