Merge pull request #808 from sammcj/deepseek_r1_ollama

fix(r1): Ensure the DeepSeek r1 model works with Ollama
This commit is contained in:
Matt Rubens 2025-02-06 02:09:24 -05:00 committed by GitHub
commit 6037140c4f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -3,6 +3,7 @@ import OpenAI from "openai"
import { ApiHandler, SingleCompletionHandler } from "../"
import { ApiHandlerOptions, ModelInfo, openAiModelInfoSaneDefaults } from "../../shared/api"
import { convertToOpenAiMessages } from "../transform/openai-format"
import { convertToR1Format } from "../transform/r1-format"
import { ApiStream } from "../transform/stream"
export class OllamaHandler implements ApiHandler, SingleCompletionHandler {
@ -18,9 +19,11 @@ export class OllamaHandler implements ApiHandler, SingleCompletionHandler {
}
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
const modelId = this.getModel().id
const useR1Format = modelId.toLowerCase().includes('deepseek-r1')
const openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [
{ role: "system", content: systemPrompt },
...convertToOpenAiMessages(messages),
...(useR1Format ? convertToR1Format(messages) : convertToOpenAiMessages(messages)),
]
const stream = await this.client.chat.completions.create({
@ -49,9 +52,11 @@ export class OllamaHandler implements ApiHandler, SingleCompletionHandler {
async completePrompt(prompt: string): Promise<string> {
try {
const modelId = this.getModel().id
const useR1Format = modelId.toLowerCase().includes('deepseek-r1')
const response = await this.client.chat.completions.create({
model: this.getModel().id,
messages: [{ role: "user", content: prompt }],
messages: useR1Format ? convertToR1Format([{ role: "user", content: prompt }]) : [{ role: "user", content: prompt }],
temperature: 0,
stream: false,
})