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
12 changes: 6 additions & 6 deletions pnpm-lock.yaml

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

144 changes: 143 additions & 1 deletion src/api/providers/__tests__/mistral.spec.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
// Mock Mistral client - must come before other imports
const mockCreate = vi.fn()
const mockComplete = vi.fn()
vi.mock("@mistralai/mistralai", () => {
return {
Mistral: vi.fn().mockImplementation(() => ({
Expand All @@ -21,6 +22,17 @@ vi.mock("@mistralai/mistralai", () => {
}
return stream
}),
complete: mockComplete.mockImplementation(async (_options) => {
return {
choices: [
{
message: {
content: "Test response",
},
},
],
}
}),
},
})),
}
Expand All @@ -29,7 +41,7 @@ vi.mock("@mistralai/mistralai", () => {
import type { Anthropic } from "@anthropic-ai/sdk"
import { MistralHandler } from "../mistral"
import type { ApiHandlerOptions } from "../../../shared/api"
import type { ApiStreamTextChunk } from "../../transform/stream"
import type { ApiStreamTextChunk, ApiStreamReasoningChunk } from "../../transform/stream"

describe("MistralHandler", () => {
let handler: MistralHandler
Expand All @@ -44,6 +56,7 @@ describe("MistralHandler", () => {
}
handler = new MistralHandler(mockOptions)
mockCreate.mockClear()
mockComplete.mockClear()
})

describe("constructor", () => {
Expand Down Expand Up @@ -122,5 +135,134 @@ describe("MistralHandler", () => {
mockCreate.mockRejectedValueOnce(new Error("API Error"))
await expect(handler.createMessage(systemPrompt, messages).next()).rejects.toThrow("API Error")
})

it("should handle thinking content as reasoning chunks", async () => {
// Mock stream with thinking content matching new SDK structure
mockCreate.mockImplementationOnce(async (_options) => {
const stream = {
[Symbol.asyncIterator]: async function* () {
yield {
data: {
choices: [
{
delta: {
content: [
{
type: "thinking",
thinking: [{ type: "text", text: "Let me think about this..." }],
},
{ type: "text", text: "Here's the answer" },
],
},
index: 0,
},
],
},
}
},
}
return stream
})

const iterator = handler.createMessage(systemPrompt, messages)
const results: (ApiStreamTextChunk | ApiStreamReasoningChunk)[] = []

for await (const chunk of iterator) {
if ("text" in chunk) {
results.push(chunk as ApiStreamTextChunk | ApiStreamReasoningChunk)
}
}

expect(results).toHaveLength(2)
expect(results[0]).toEqual({ type: "reasoning", text: "Let me think about this..." })
expect(results[1]).toEqual({ type: "text", text: "Here's the answer" })
})

it("should handle mixed content arrays correctly", async () => {
// Mock stream with mixed content matching new SDK structure
mockCreate.mockImplementationOnce(async (_options) => {
const stream = {
[Symbol.asyncIterator]: async function* () {
yield {
data: {
choices: [
{
delta: {
content: [
{ type: "text", text: "First text" },
{
type: "thinking",
thinking: [{ type: "text", text: "Some reasoning" }],
},
{ type: "text", text: "Second text" },
],
},
index: 0,
},
],
},
}
},
}
return stream
})

const iterator = handler.createMessage(systemPrompt, messages)
const results: (ApiStreamTextChunk | ApiStreamReasoningChunk)[] = []

for await (const chunk of iterator) {
if ("text" in chunk) {
results.push(chunk as ApiStreamTextChunk | ApiStreamReasoningChunk)
}
}

expect(results).toHaveLength(3)
expect(results[0]).toEqual({ type: "text", text: "First text" })
expect(results[1]).toEqual({ type: "reasoning", text: "Some reasoning" })
expect(results[2]).toEqual({ type: "text", text: "Second text" })
})
})

describe("completePrompt", () => {
it("should complete prompt successfully", async () => {
const prompt = "Test prompt"
const result = await handler.completePrompt(prompt)

expect(mockComplete).toHaveBeenCalledWith({
model: mockOptions.apiModelId,
messages: [{ role: "user", content: prompt }],
temperature: 0,
})

expect(result).toBe("Test response")
})
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

While we test array content filtering in completePrompt, could we add an explicit test for when content is already a string (covering line 128 in mistral.ts)?


it("should filter out thinking content in completePrompt", async () => {
mockComplete.mockImplementationOnce(async (_options) => {
return {
choices: [
{
message: {
content: [
{ type: "thinking", text: "Let me think..." },
{ type: "text", text: "Answer part 1" },
{ type: "text", text: "Answer part 2" },
],
},
},
],
}
})

const prompt = "Test prompt"
const result = await handler.completePrompt(prompt)

expect(result).toBe("Answer part 1Answer part 2")
})

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we add a test for the edge case where ALL content is thinking content? This would verify the function returns an empty string correctly:

Suggested change
it("should handle all thinking content in completePrompt", async () => {
mockComplete.mockImplementationOnce(async (_options) => {
return {
choices: [
{
message: {
content: [
{ type: "thinking", text: "Let me think..." },
{ type: "thinking", text: "Still thinking..." },
],
},
},
],
}
})
const prompt = "Test prompt"
const result = await handler.completePrompt(prompt)
expect(result).toBe("")
})

it("should handle errors in completePrompt", async () => {
mockComplete.mockRejectedValueOnce(new Error("API Error"))
await expect(handler.completePrompt("Test prompt")).rejects.toThrow("Mistral completion error: API Error")
})
})
})
48 changes: 36 additions & 12 deletions src/api/providers/mistral.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,14 @@ import { ApiStream } from "../transform/stream"
import { BaseProvider } from "./base-provider"
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index"

// Type helper to handle thinking chunks from Mistral API
// The SDK includes ThinkChunk but TypeScript has trouble with the discriminated union
type ContentChunkWithThinking = {
type: string
text?: string
thinking?: Array<{ type: string; text?: string }>
}

export class MistralHandler extends BaseProvider implements SingleCompletionHandler {
protected options: ApiHandlerOptions
private client: Mistral
Expand Down Expand Up @@ -48,26 +56,38 @@ export class MistralHandler extends BaseProvider implements SingleCompletionHand
temperature,
})

for await (const chunk of response) {
const delta = chunk.data.choices[0]?.delta
for await (const event of response) {
const delta = event.data.choices[0]?.delta

if (delta?.content) {
let content: string = ""

if (typeof delta.content === "string") {
content = delta.content
// Handle string content as text
yield { type: "text", text: delta.content }
} else if (Array.isArray(delta.content)) {
content = delta.content.map((c) => (c.type === "text" ? c.text : "")).join("")
// Handle array of content chunks
// The SDK v1.9.18 supports ThinkChunk with type "thinking"
for (const chunk of delta.content as ContentChunkWithThinking[]) {
if (chunk.type === "thinking" && chunk.thinking) {
// Handle thinking content as reasoning chunks
// ThinkChunk has a 'thinking' property that contains an array of text/reference chunks
for (const thinkingPart of chunk.thinking) {
if (thinkingPart.type === "text" && thinkingPart.text) {
yield { type: "reasoning", text: thinkingPart.text }
}
}
} else if (chunk.type === "text" && chunk.text) {
// Handle text content normally
yield { type: "text", text: chunk.text }
}
}
}

yield { type: "text", text: content }
}

if (chunk.data.usage) {
if (event.data.usage) {
yield {
type: "usage",
inputTokens: chunk.data.usage.promptTokens || 0,
outputTokens: chunk.data.usage.completionTokens || 0,
inputTokens: event.data.usage.promptTokens || 0,
outputTokens: event.data.usage.completionTokens || 0,
}
}
}
Expand Down Expand Up @@ -97,7 +117,11 @@ export class MistralHandler extends BaseProvider implements SingleCompletionHand
const content = response.choices?.[0]?.message.content

if (Array.isArray(content)) {
return content.map((c) => (c.type === "text" ? c.text : "")).join("")
// Only return text content, filter out thinking content for non-streaming
return (content as ContentChunkWithThinking[])
.filter((c) => c.type === "text" && c.text)
.map((c) => c.text || "")
.join("")
}

return content || ""
Expand Down
2 changes: 1 addition & 1 deletion src/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -424,7 +424,7 @@
"@aws-sdk/credential-providers": "^3.848.0",
"@google/genai": "^1.0.0",
"@lmstudio/sdk": "^1.1.1",
"@mistralai/mistralai": "^1.3.6",
"@mistralai/mistralai": "^1.9.18",
"@modelcontextprotocol/sdk": "^1.9.0",
"@qdrant/js-client-rest": "^1.14.0",
"@roo-code/cloud": "^0.14.0",
Expand Down
Loading