Просмотр исходного кода

Merge pull request #808 from sammcj/deepseek_r1_ollama

fix(r1): Ensure the DeepSeek r1 model works with Ollama
Matt Rubens 10 месяцев назад
Родитель
Сommit
6037140c4f
1 измененных файлов с 7 добавлено и 2 удалено
  1. 7 2
      src/api/providers/ollama.ts

+ 7 - 2
src/api/providers/ollama.ts

@@ -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,
 			})