Skip to content
Closed
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
80 changes: 80 additions & 0 deletions src/api/providers/__tests__/mistral.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -118,6 +118,86 @@ describe("MistralHandler", () => {
expect(results[0].text).toBe("Test response")
})

it("should handle thinking type in content array", async () => {
// Override the mock to return content with thinking type
mockCreate.mockImplementationOnce(async (_options) => {
const stream = {
[Symbol.asyncIterator]: async function* () {
yield {
data: {
choices: [
{
delta: {
content: [
{ type: "thinking", text: "Let me think about this..." },
{ type: "text", text: "Here is my response" },
],
},
index: 0,
},
],
},
}
},
}
return stream
})

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

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

// Should only include the text type content, not thinking
expect(results.length).toBe(1)
expect(results[0].text).toBe("Here is my response")
})

it("should handle mixed content types gracefully", async () => {
// Override the mock to return various content types
mockCreate.mockImplementationOnce(async (_options) => {
const stream = {
[Symbol.asyncIterator]: async function* () {
yield {
data: {
choices: [
{
delta: {
content: [
{ type: "thinking", text: "Processing..." },
{ type: "text", text: "First part" },
{ type: "unknown", data: "some data" },
{ type: "text", text: " Second part" },
],
},
index: 0,
},
],
},
}
},
}
return stream
})

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

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

// Should concatenate only text type content
expect(results.length).toBe(1)
expect(results[0].text).toBe("First part Second part")
})
Copy link
Contributor Author

Choose a reason for hiding this comment

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

Great test coverage for the streaming case! However, I notice we're missing tests for the completePrompt method which also received the same fix. Would it make sense to add a test case to ensure completePrompt handles the "thinking" type correctly as well?


it("should handle errors gracefully", async () => {
mockCreate.mockRejectedValueOnce(new Error("API Error"))
await expect(handler.createMessage(systemPrompt, messages).next()).rejects.toThrow("API Error")
Expand Down
17 changes: 14 additions & 3 deletions src/api/providers/mistral.ts
Original file line number Diff line number Diff line change
Expand Up @@ -57,10 +57,17 @@ export class MistralHandler extends BaseProvider implements SingleCompletionHand
if (typeof delta.content === "string") {
content = delta.content
} else if (Array.isArray(delta.content)) {
content = delta.content.map((c) => (c.type === "text" ? c.text : "")).join("")
// Handle array content, filtering out "thinking" type and other non-text types
content = delta.content
.filter((c: any) => c.type === "text")
Copy link
Contributor Author

Choose a reason for hiding this comment

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

The use of any type here works, but could we consider defining a proper type for the content array items? Something like:

Suggested change
.filter((c: any) => c.type === "text")
.filter((c: { type: string; text?: string }) => c.type === "text")
.map((c: { type: string; text?: string }) => c.text || "")

This would improve type safety and make the code more maintainable.

.map((c: any) => c.text || "")
.join("")
}

yield { type: "text", text: content }
// Only yield if we have actual content to send
if (content) {
yield { type: "text", text: content }
}
}

if (chunk.data.usage) {
Expand Down Expand Up @@ -97,7 +104,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("")
// Handle array content, filtering out "thinking" type and other non-text types
return content
.filter((c: any) => c.type === "text")
Copy link
Contributor Author

Choose a reason for hiding this comment

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

Same type safety consideration here - could we use a proper type instead of any?

.map((c: any) => c.text || "")
.join("")
}

return content || ""
Expand Down
Loading