Merge pull request #1811 from somnifex/main

refactor: 重构ollama渠道

relay/channel/ollama/adaptor.go

@@ -10,6 +10,7 @@ import (
 	relaycommon "one-api/relay/common"
 	relayconstant "one-api/relay/constant"
 	"one-api/types"
+	"strings"
 
 	"github.com/gin-gonic/gin"
 )
@@ -17,10 +18,7 @@ import (
 type Adaptor struct {
 }
 
-func (a *Adaptor) ConvertGeminiRequest(*gin.Context, *relaycommon.RelayInfo, *dto.GeminiChatRequest) (any, error) {
-	//TODO implement me
-	return nil, errors.New("not implemented")
-}
+func (a *Adaptor) ConvertGeminiRequest(*gin.Context, *relaycommon.RelayInfo, *dto.GeminiChatRequest) (any, error) { return nil, errors.New("not implemented") }
 
 func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.ClaudeRequest) (any, error) {
 	openaiAdaptor := openai.Adaptor{}
@@ -31,32 +29,21 @@ func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayIn
 	openaiRequest.(*dto.GeneralOpenAIRequest).StreamOptions = &dto.StreamOptions{
 		IncludeUsage: true,
 	}
-	return requestOpenAI2Ollama(c, openaiRequest.(*dto.GeneralOpenAIRequest))
+	// map to ollama chat request (Claude -> OpenAI -> Ollama chat)
+	return openAIChatToOllamaChat(c, openaiRequest.(*dto.GeneralOpenAIRequest))
 }
 
-func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
-	//TODO implement me
-	return nil, errors.New("not implemented")
-}
+func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) { return nil, errors.New("not implemented") }
 
-func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
-	//TODO implement me
-	return nil, errors.New("not implemented")
-}
+func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) { return nil, errors.New("not implemented") }
 
 func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
 }
 
 func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
-	if info.RelayFormat == types.RelayFormatClaude {
-		return info.ChannelBaseUrl + "/v1/chat/completions", nil
-	}
-	switch info.RelayMode {
-	case relayconstant.RelayModeEmbeddings:
-		return info.ChannelBaseUrl + "/api/embed", nil
-	default:
-		return relaycommon.GetFullRequestURL(info.ChannelBaseUrl, info.RequestURLPath, info.ChannelType), nil
-	}
+    if info.RelayMode == relayconstant.RelayModeEmbeddings { return info.ChannelBaseUrl + "/api/embed", nil }
+    if strings.Contains(info.RequestURLPath, "/v1/completions") || info.RelayMode == relayconstant.RelayModeCompletions { return info.ChannelBaseUrl + "/api/generate", nil }
+    return info.ChannelBaseUrl + "/api/chat", nil
 }
 
 func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
@@ -66,10 +53,12 @@ func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *rel
 }
 
 func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
-	if request == nil {
-		return nil, errors.New("request is nil")
+	if request == nil { return nil, errors.New("request is nil") }
+	// decide generate or chat
+	if strings.Contains(info.RequestURLPath, "/v1/completions") || info.RelayMode == relayconstant.RelayModeCompletions {
+		return openAIToGenerate(c, request)
 	}
-	return requestOpenAI2Ollama(c, request)
+	return openAIChatToOllamaChat(c, request)
 }
 
 func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
@@ -80,10 +69,7 @@ func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.Rela
 	return requestOpenAI2Embeddings(request), nil
 }
 
-func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
-	// TODO implement me
-	return nil, errors.New("not implemented")
-}
+func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) { return nil, errors.New("not implemented") }
 
 func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
 	return channel.DoApiRequest(a, c, info, requestBody)
@@ -92,15 +78,13 @@ func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, request
 func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
 	switch info.RelayMode {
 	case relayconstant.RelayModeEmbeddings:
-		usage, err = ollamaEmbeddingHandler(c, info, resp)
+		return ollamaEmbeddingHandler(c, info, resp)
 	default:
 		if info.IsStream {
-			usage, err = openai.OaiStreamHandler(c, info, resp)
-		} else {
-			usage, err = openai.OpenaiHandler(c, info, resp)
+			return ollamaStreamHandler(c, info, resp)
 		}
+		return ollamaChatHandler(c, info, resp)
 	}
-	return
 }
 
 func (a *Adaptor) GetModelList() []string {

relay/channel/ollama/dto.go

@@ -2,48 +2,69 @@ package ollama
 
 import (
 	"encoding/json"
-	"one-api/dto"
 )
 
-type OllamaRequest struct {
-	Model            string                `json:"model,omitempty"`
-	Messages         []dto.Message         `json:"messages,omitempty"`
-	Stream           bool                  `json:"stream,omitempty"`
-	Temperature      *float64              `json:"temperature,omitempty"`
-	Seed             float64               `json:"seed,omitempty"`
-	Topp             float64               `json:"top_p,omitempty"`
-	TopK             int                   `json:"top_k,omitempty"`
-	Stop             any                   `json:"stop,omitempty"`
-	MaxTokens        uint                  `json:"max_tokens,omitempty"`
-	Tools            []dto.ToolCallRequest `json:"tools,omitempty"`
-	ResponseFormat   any                   `json:"response_format,omitempty"`
-	FrequencyPenalty float64               `json:"frequency_penalty,omitempty"`
-	PresencePenalty  float64               `json:"presence_penalty,omitempty"`
-	Suffix           any                   `json:"suffix,omitempty"`
-	StreamOptions    *dto.StreamOptions    `json:"stream_options,omitempty"`
-	Prompt           any                   `json:"prompt,omitempty"`
-	Think            json.RawMessage       `json:"think,omitempty"`
-}
-
-type Options struct {
-	Seed             int      `json:"seed,omitempty"`
-	Temperature      *float64 `json:"temperature,omitempty"`
-	TopK             int      `json:"top_k,omitempty"`
-	TopP             float64  `json:"top_p,omitempty"`
-	FrequencyPenalty float64  `json:"frequency_penalty,omitempty"`
-	PresencePenalty  float64  `json:"presence_penalty,omitempty"`
-	NumPredict       int      `json:"num_predict,omitempty"`
-	NumCtx           int      `json:"num_ctx,omitempty"`
+type OllamaChatMessage struct {
+	Role      string            `json:"role"`
+	Content   string            `json:"content,omitempty"`
+	Images    []string          `json:"images,omitempty"`
+	ToolCalls []OllamaToolCall  `json:"tool_calls,omitempty"`
+	ToolName  string            `json:"tool_name,omitempty"`
+	Thinking  json.RawMessage   `json:"thinking,omitempty"`
+}
+
+type OllamaToolFunction struct {
+	Name        string      `json:"name"`
+	Description string      `json:"description,omitempty"`
+	Parameters  interface{} `json:"parameters,omitempty"`
+}
+
+type OllamaTool struct {
+	Type     string            `json:"type"`
+	Function OllamaToolFunction `json:"function"`
+}
+
+type OllamaToolCall struct {
+	Function struct {
+		Name      string      `json:"name"`
+		Arguments interface{} `json:"arguments"`
+	} `json:"function"`
+}
+
+type OllamaChatRequest struct {
+	Model     string              `json:"model"`
+	Messages  []OllamaChatMessage `json:"messages"`
+	Tools     interface{}         `json:"tools,omitempty"`
+	Format    interface{}         `json:"format,omitempty"`
+	Stream    bool                `json:"stream,omitempty"`
+	Options   map[string]any      `json:"options,omitempty"`
+	KeepAlive interface{}         `json:"keep_alive,omitempty"`
+	Think     json.RawMessage     `json:"think,omitempty"`
+}
+
+type OllamaGenerateRequest struct {
+	Model     string         `json:"model"`
+	Prompt    string         `json:"prompt,omitempty"`
+	Suffix    string         `json:"suffix,omitempty"`
+	Images    []string       `json:"images,omitempty"`
+	Format    interface{}    `json:"format,omitempty"`
+	Stream    bool           `json:"stream,omitempty"`
+	Options   map[string]any `json:"options,omitempty"`
+	KeepAlive interface{}    `json:"keep_alive,omitempty"`
+	Think     json.RawMessage `json:"think,omitempty"`
 }
 
 type OllamaEmbeddingRequest struct {
-	Model   string   `json:"model,omitempty"`
-	Input   []string `json:"input"`
-	Options *Options `json:"options,omitempty"`
+	Model     string         `json:"model"`
+	Input     interface{}    `json:"input"`
+	Options   map[string]any `json:"options,omitempty"`
+	Dimensions int            `json:"dimensions,omitempty"`
 }
 
 type OllamaEmbeddingResponse struct {
-	Error     string      `json:"error,omitempty"`
-	Model     string      `json:"model"`
-	Embedding [][]float64 `json:"embeddings,omitempty"`
+	Error           string        `json:"error,omitempty"`
+	Model           string        `json:"model"`
+	Embeddings      [][]float64   `json:"embeddings"`
+	PromptEvalCount int           `json:"prompt_eval_count,omitempty"`
 }
+

relay/channel/ollama/stream.go

+package ollama
+
+import (
+    "bufio"
+    "encoding/json"
+    "fmt"
+    "io"
+    "net/http"
+    "one-api/common"
+    "one-api/dto"
+    "one-api/logger"
+    relaycommon "one-api/relay/common"
+    "one-api/relay/helper"
+    "one-api/service"
+    "one-api/types"
+    "strings"
+    "time"
+
+    "github.com/gin-gonic/gin"
+)
+
+type ollamaChatStreamChunk struct {
+    Model            string `json:"model"`
+    CreatedAt        string `json:"created_at"`
+    // chat
+    Message *struct {
+        Role      string `json:"role"`
+        Content   string `json:"content"`
+        Thinking  json.RawMessage `json:"thinking"`
+        ToolCalls []struct {
+            Function struct {
+                Name      string      `json:"name"`
+                Arguments interface{} `json:"arguments"`
+            } `json:"function"`
+        } `json:"tool_calls"`
+    } `json:"message"`
+    // generate
+    Response string `json:"response"`
+    Done         bool    `json:"done"`
+    DoneReason   string  `json:"done_reason"`
+    TotalDuration int64  `json:"total_duration"`
+    LoadDuration  int64  `json:"load_duration"`
+    PromptEvalCount int  `json:"prompt_eval_count"`
+    EvalCount       int  `json:"eval_count"`
+    PromptEvalDuration int64 `json:"prompt_eval_duration"`
+    EvalDuration       int64 `json:"eval_duration"`
+}
+
+func toUnix(ts string) int64 {
+    if ts == "" { return time.Now().Unix() }
+    // try time.RFC3339 or with nanoseconds
+    t, err := time.Parse(time.RFC3339Nano, ts)
+    if err != nil { t2, err2 := time.Parse(time.RFC3339, ts); if err2==nil { return t2.Unix() }; return time.Now().Unix() }
+    return t.Unix()
+}
+
+func ollamaStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
+    if resp == nil || resp.Body == nil { return nil, types.NewOpenAIError(fmt.Errorf("empty response"), types.ErrorCodeBadResponse, http.StatusBadRequest) }
+    defer service.CloseResponseBodyGracefully(resp)
+
+    helper.SetEventStreamHeaders(c)
+    scanner := bufio.NewScanner(resp.Body)
+    usage := &dto.Usage{}
+    var model = info.UpstreamModelName
+    var responseId = common.GetUUID()
+    var created = time.Now().Unix()
+    var toolCallIndex int
+    start := helper.GenerateStartEmptyResponse(responseId, created, model, nil)
+    if data, err := common.Marshal(start); err == nil { _ = helper.StringData(c, string(data)) }
+
+    for scanner.Scan() {
+        line := scanner.Text()
+        line = strings.TrimSpace(line)
+        if line == "" { continue }
+        var chunk ollamaChatStreamChunk
+        if err := json.Unmarshal([]byte(line), &chunk); err != nil {
+            logger.LogError(c, "ollama stream json decode error: "+err.Error()+" line="+line)
+            return usage, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
+        }
+        if chunk.Model != "" { model = chunk.Model }
+        created = toUnix(chunk.CreatedAt)
+
+        if !chunk.Done {
+            // delta content
+            var content string
+            if chunk.Message != nil { content = chunk.Message.Content } else { content = chunk.Response }
+            delta := dto.ChatCompletionsStreamResponse{
+                Id:      responseId,
+                Object:  "chat.completion.chunk",
+                Created: created,
+                Model:   model,
+                Choices: []dto.ChatCompletionsStreamResponseChoice{ {
+                    Index: 0,
+                    Delta: dto.ChatCompletionsStreamResponseChoiceDelta{ Role: "assistant" },
+                } },
+            }
+            if content != "" { delta.Choices[0].Delta.SetContentString(content) }
+            if chunk.Message != nil && len(chunk.Message.Thinking) > 0 {
+                raw := strings.TrimSpace(string(chunk.Message.Thinking))
+                if raw != "" && raw != "null" { delta.Choices[0].Delta.SetReasoningContent(raw) }
+            }
+            // tool calls
+            if chunk.Message != nil && len(chunk.Message.ToolCalls) > 0 {
+                delta.Choices[0].Delta.ToolCalls = make([]dto.ToolCallResponse,0,len(chunk.Message.ToolCalls))
+                for _, tc := range chunk.Message.ToolCalls {
+                    // arguments -> string
+                    argBytes, _ := json.Marshal(tc.Function.Arguments)
+                    toolId := fmt.Sprintf("call_%d", toolCallIndex)
+                    tr := dto.ToolCallResponse{ID:toolId, Type:"function", Function: dto.FunctionResponse{Name: tc.Function.Name, Arguments: string(argBytes)}}
+                    tr.SetIndex(toolCallIndex)
+                    toolCallIndex++
+                    delta.Choices[0].Delta.ToolCalls = append(delta.Choices[0].Delta.ToolCalls, tr)
+                }
+            }
+            if data, err := common.Marshal(delta); err == nil { _ = helper.StringData(c, string(data)) }
+            continue
+        }
+        // done frame
+        // finalize once and break loop
+        usage.PromptTokens = chunk.PromptEvalCount
+        usage.CompletionTokens = chunk.EvalCount
+        usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
+    finishReason := chunk.DoneReason
+    if finishReason == "" { finishReason = "stop" }
+        // emit stop delta
+        if stop := helper.GenerateStopResponse(responseId, created, model, finishReason); stop != nil {
+            if data, err := common.Marshal(stop); err == nil { _ = helper.StringData(c, string(data)) }
+        }
+        // emit usage frame
+        if final := helper.GenerateFinalUsageResponse(responseId, created, model, *usage); final != nil {
+            if data, err := common.Marshal(final); err == nil { _ = helper.StringData(c, string(data)) }
+        }
+        // send [DONE]
+        helper.Done(c)
+        break
+    }
+    if err := scanner.Err(); err != nil && err != io.EOF { logger.LogError(c, "ollama stream scan error: "+err.Error()) }
+    return usage, nil
+}
+
+// non-stream handler for chat/generate
+func ollamaChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
+    body, err := io.ReadAll(resp.Body)
+    if err != nil { return nil, types.NewOpenAIError(err, types.ErrorCodeReadResponseBodyFailed, http.StatusInternalServerError) }
+    service.CloseResponseBodyGracefully(resp)
+    raw := string(body)
+    if common.DebugEnabled { println("ollama non-stream raw resp:", raw) }
+
+    lines := strings.Split(raw, "\n")
+    var (
+        aggContent strings.Builder
+        reasoningBuilder strings.Builder
+        lastChunk ollamaChatStreamChunk
+        parsedAny bool
+    )
+    for _, ln := range lines {
+        ln = strings.TrimSpace(ln)
+        if ln == "" { continue }
+        var ck ollamaChatStreamChunk
+        if err := json.Unmarshal([]byte(ln), &ck); err != nil {
+            if len(lines) == 1 { return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) }
+            continue
+        }
+        parsedAny = true
+        lastChunk = ck
+        if ck.Message != nil && len(ck.Message.Thinking) > 0 {
+            raw := strings.TrimSpace(string(ck.Message.Thinking))
+            if raw != "" && raw != "null" { reasoningBuilder.WriteString(raw) }
+        }
+        if ck.Message != nil && ck.Message.Content != "" { aggContent.WriteString(ck.Message.Content) } else if ck.Response != "" { aggContent.WriteString(ck.Response) }
+    }
+
+    if !parsedAny {
+        var single ollamaChatStreamChunk
+        if err := json.Unmarshal(body, &single); err != nil { return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) }
+        lastChunk = single
+        if single.Message != nil {
+            if len(single.Message.Thinking) > 0 { raw := strings.TrimSpace(string(single.Message.Thinking)); if raw != "" && raw != "null" { reasoningBuilder.WriteString(raw) } }
+            aggContent.WriteString(single.Message.Content)
+        } else { aggContent.WriteString(single.Response) }
+    }
+
+    model := lastChunk.Model
+    if model == "" { model = info.UpstreamModelName }
+    created := toUnix(lastChunk.CreatedAt)
+    usage := &dto.Usage{PromptTokens: lastChunk.PromptEvalCount, CompletionTokens: lastChunk.EvalCount, TotalTokens: lastChunk.PromptEvalCount + lastChunk.EvalCount}
+    content := aggContent.String()
+    finishReason := lastChunk.DoneReason
+    if finishReason == "" { finishReason = "stop" }
+
+    msg := dto.Message{Role: "assistant", Content: contentPtr(content)}
+    if rc := reasoningBuilder.String(); rc != "" { msg.ReasoningContent = rc }
+    full := dto.OpenAITextResponse{
+        Id:      common.GetUUID(),
+        Model:   model,
+        Object:  "chat.completion",
+        Created: created,
+        Choices: []dto.OpenAITextResponseChoice{ {
+            Index: 0,
+            Message: msg,
+            FinishReason: finishReason,
+        } },
+        Usage: *usage,
+    }
+    out, _ := common.Marshal(full)
+    service.IOCopyBytesGracefully(c, resp, out)
+    return usage, nil
+}
+
+func contentPtr(s string) *string { if s=="" { return nil }; return &s }