✨ feat: support Ollama (#168)

providers/ollama/base.go

@@ -0,0 +1,79 @@
+package ollama
+
+import (
+	"encoding/json"
+	"net/http"
+	"one-api/common/requester"
+	"one-api/model"
+	"one-api/types"
+
+	"one-api/providers/base"
+)
+
+type OllamaProviderFactory struct{}
+
+type OllamaProvider struct {
+	base.BaseProvider
+}
+
+// 创建 OllamaProvider
+func (f OllamaProviderFactory) Create(channel *model.Channel) base.ProviderInterface {
+	config := getOllamaConfig()
+
+	return &OllamaProvider{
+		BaseProvider: base.BaseProvider{
+			Config:    config,
+			Channel:   channel,
+			Requester: requester.NewHTTPRequester(*channel.Proxy, RequestErrorHandle),
+		},
+	}
+}
+
+func getOllamaConfig() base.ProviderConfig {
+	return base.ProviderConfig{
+		BaseURL:         "",
+		ChatCompletions: "/api/chat",
+		Embeddings:      "/api/embeddings",
+	}
+}
+
+// 请求错误处理
+func RequestErrorHandle(resp *http.Response) *types.OpenAIError {
+	errorResponse := &OllamaError{}
+	err := json.NewDecoder(resp.Body).Decode(errorResponse)
+	if err != nil {
+		return nil
+	}
+
+	return errorHandle(errorResponse)
+}
+
+// 错误处理
+func errorHandle(OllamaError *OllamaError) *types.OpenAIError {
+	if OllamaError.Error == "" {
+		return nil
+	}
+	return &types.OpenAIError{
+		Message: OllamaError.Error,
+		Type:    "Ollama Error",
+	}
+}
+
+// 获取请求头
+func (p *OllamaProvider) GetRequestHeaders() (headers map[string]string) {
+	headers = make(map[string]string)
+	p.CommonRequestHeaders(headers)
+
+	otherHeaders := p.Channel.Plugin.Data()["headers"]
+
+	for key, value := range otherHeaders {
+		headerValue, isString := value.(string)
+		if !isString || headerValue == "" {
+			continue
+		}
+
+		headers[key] = headerValue
+	}
+
+	return headers
+}

providers/ollama/chat.go

@@ -0,0 +1,207 @@
+package ollama
+
+import (
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"one-api/common"
+	"one-api/common/image"
+	"one-api/common/requester"
+	"one-api/types"
+	"strings"
+)
+
+type ollamaStreamHandler struct {
+	Usage   *types.Usage
+	Request *types.ChatCompletionRequest
+}
+
+func (p *OllamaProvider) CreateChatCompletion(request *types.ChatCompletionRequest) (*types.ChatCompletionResponse, *types.OpenAIErrorWithStatusCode) {
+	req, errWithCode := p.getChatRequest(request)
+	if errWithCode != nil {
+		return nil, errWithCode
+	}
+	defer req.Body.Close()
+
+	response := &ChatResponse{}
+	// 发送请求
+	_, errWithCode = p.Requester.SendRequest(req, response, false)
+	if errWithCode != nil {
+		return nil, errWithCode
+	}
+
+	return p.convertToChatOpenai(response, request)
+}
+
+func (p *OllamaProvider) CreateChatCompletionStream(request *types.ChatCompletionRequest) (requester.StreamReaderInterface[string], *types.OpenAIErrorWithStatusCode) {
+	req, errWithCode := p.getChatRequest(request)
+	if errWithCode != nil {
+		return nil, errWithCode
+	}
+	defer req.Body.Close()
+
+	// 发送请求
+	resp, errWithCode := p.Requester.SendRequestRaw(req)
+	if errWithCode != nil {
+		return nil, errWithCode
+	}
+
+	chatHandler := &ollamaStreamHandler{
+		Usage:   p.Usage,
+		Request: request,
+	}
+
+	return requester.RequestStream(p.Requester, resp, chatHandler.handlerStream)
+}
+
+func (p *OllamaProvider) getChatRequest(request *types.ChatCompletionRequest) (*http.Request, *types.OpenAIErrorWithStatusCode) {
+	url, errWithCode := p.GetSupportedAPIUri(common.RelayModeChatCompletions)
+	if errWithCode != nil {
+		return nil, errWithCode
+	}
+
+	// 获取请求地址
+	fullRequestURL := p.GetFullRequestURL(url, request.Model)
+
+	// 获取请求头
+	headers := p.GetRequestHeaders()
+
+	ollamaRequest, errWithCode := convertFromChatOpenai(request)
+	if errWithCode != nil {
+		return nil, errWithCode
+	}
+
+	// 创建请求
+	req, err := p.Requester.NewRequest(http.MethodPost, fullRequestURL, p.Requester.WithBody(ollamaRequest), p.Requester.WithHeader(headers))
+	if err != nil {
+		return nil, common.ErrorWrapper(err, "new_request_failed", http.StatusInternalServerError)
+	}
+
+	return req, nil
+}
+
+func (p *OllamaProvider) convertToChatOpenai(response *ChatResponse, request *types.ChatCompletionRequest) (openaiResponse *types.ChatCompletionResponse, errWithCode *types.OpenAIErrorWithStatusCode) {
+	err := errorHandle(&response.OllamaError)
+	if err != nil {
+		errWithCode = &types.OpenAIErrorWithStatusCode{
+			OpenAIError: *err,
+			StatusCode:  http.StatusBadRequest,
+		}
+		return
+	}
+
+	choices := types.ChatCompletionChoice{
+		Index: 0,
+		Message: types.ChatCompletionMessage{
+			Role:    response.Message.Role,
+			Content: response.Message.Content,
+		},
+		FinishReason: types.FinishReasonStop,
+	}
+
+	openaiResponse = &types.ChatCompletionResponse{
+		ID:      fmt.Sprintf("chatcmpl-%s", common.GetUUID()),
+		Object:  "chat.completion",
+		Created: common.GetTimestamp(),
+		Model:   request.Model,
+		Choices: []types.ChatCompletionChoice{choices},
+		Usage: &types.Usage{
+			PromptTokens:     response.PromptEvalCount,
+			CompletionTokens: response.EvalCount,
+			TotalTokens:      response.PromptEvalCount + response.EvalCount,
+		},
+	}
+
+	*p.Usage = *openaiResponse.Usage
+
+	return openaiResponse, nil
+}
+
+func convertFromChatOpenai(request *types.ChatCompletionRequest) (*ChatRequest, *types.OpenAIErrorWithStatusCode) {
+	ollamaRequest := &ChatRequest{
+		Model:    request.Model,
+		Stream:   request.Stream,
+		Messages: make([]Message, 0, len(request.Messages)),
+		Options: Option{
+			Temperature: request.Temperature,
+			TopP:        request.TopP,
+			Seed:        request.Seed,
+		},
+	}
+
+	for _, message := range request.Messages {
+		ollamaMessage := Message{
+			Role:    message.Role,
+			Content: "",
+		}
+
+		openaiMessagePart := message.ParseContent()
+		for _, openaiPart := range openaiMessagePart {
+			if openaiPart.Type == types.ContentTypeText {
+				ollamaMessage.Content += openaiPart.Text
+			} else if openaiPart.Type == types.ContentTypeImageURL {
+				_, data, err := image.GetImageFromUrl(openaiPart.ImageURL.URL)
+				if err != nil {
+					return nil, common.ErrorWrapper(err, "image_url_invalid", http.StatusBadRequest)
+				}
+				ollamaMessage.Images = append(ollamaMessage.Images, data)
+			}
+		}
+		ollamaRequest.Messages = append(ollamaRequest.Messages, ollamaMessage)
+	}
+
+	return ollamaRequest, nil
+}
+
+// 转换为OpenAI聊天流式请求体
+func (h *ollamaStreamHandler) handlerStream(rawLine *[]byte, dataChan chan string, errChan chan error) {
+	if !strings.HasPrefix(string(*rawLine), "{") {
+		*rawLine = nil
+		return
+	}
+
+	var chatResponse ChatResponse
+	err := json.Unmarshal(*rawLine, &chatResponse)
+	if err != nil {
+		errChan <- common.ErrorToOpenAIError(err)
+		return
+	}
+
+	errWithCode := errorHandle(&chatResponse.OllamaError)
+	if errWithCode != nil {
+		errChan <- errWithCode
+		return
+	}
+
+	choice := types.ChatCompletionStreamChoice{
+		Index: 0,
+	}
+
+	if chatResponse.Message.Content != "" {
+		choice.Delta = types.ChatCompletionStreamChoiceDelta{
+			Role:    types.ChatMessageRoleAssistant,
+			Content: chatResponse.Message.Content,
+		}
+	}
+
+	if chatResponse.Done {
+		choice.FinishReason = types.FinishReasonStop
+	}
+
+	if chatResponse.EvalCount > 0 {
+		h.Usage.PromptTokens = chatResponse.PromptEvalCount
+		h.Usage.CompletionTokens = chatResponse.EvalCount
+		h.Usage.TotalTokens = h.Usage.PromptTokens + chatResponse.EvalCount
+	}
+
+	chatCompletion := types.ChatCompletionStreamResponse{
+		ID:      fmt.Sprintf("chatcmpl-%s", common.GetUUID()),
+		Object:  "chat.completion.chunk",
+		Created: common.GetTimestamp(),
+		Model:   h.Request.Model,
+		Choices: []types.ChatCompletionStreamChoice{choice},
+	}
+
+	responseBody, _ := json.Marshal(chatCompletion)
+	dataChan <- string(responseBody)
+}

providers/ollama/embeddings.go

@@ -0,0 +1,70 @@
+package ollama
+
+import (
+	"net/http"
+	"one-api/common"
+	"one-api/types"
+)
+
+func (p *OllamaProvider) CreateEmbeddings(request *types.EmbeddingRequest) (*types.EmbeddingResponse, *types.OpenAIErrorWithStatusCode) {
+	url, errWithCode := p.GetSupportedAPIUri(common.RelayModeEmbeddings)
+	if errWithCode != nil {
+		return nil, errWithCode
+	}
+	// 获取请求地址
+	fullRequestURL := p.GetFullRequestURL(url, request.Model)
+	if fullRequestURL == "" {
+		return nil, common.ErrorWrapper(nil, "invalid_ollama_config", http.StatusInternalServerError)
+	}
+
+	// 获取请求头
+	headers := p.GetRequestHeaders()
+
+	ollamaRequest := &EmbeddingRequest{
+		Model:  request.Model,
+		Prompt: request.ParseInputString(),
+	}
+
+	// 创建请求
+	req, err := p.Requester.NewRequest(http.MethodPost, fullRequestURL, p.Requester.WithBody(ollamaRequest), p.Requester.WithHeader(headers))
+	if err != nil {
+		return nil, common.ErrorWrapper(err, "new_request_failed", http.StatusInternalServerError)
+	}
+	defer req.Body.Close()
+
+	ollamaResponse := &EmbeddingResponse{}
+
+	// 发送请求
+	var res *http.Response
+	res, errWithCode = p.Requester.SendRequest(req, ollamaResponse, false)
+	if errWithCode != nil {
+		return nil, errWithCode
+	}
+
+	errWithOP := errorHandle(&ollamaResponse.OllamaError)
+	if errWithOP != nil {
+		return nil, &types.OpenAIErrorWithStatusCode{
+			OpenAIError: *errWithOP,
+			StatusCode:  res.StatusCode,
+		}
+	}
+
+	response := &types.EmbeddingResponse{
+		Object: "list",
+		Model:  request.Model,
+		Data: []types.Embedding{{
+			Object:    "embedding",
+			Index:     0,
+			Embedding: ollamaResponse.Embedding,
+		}},
+		Usage: &types.Usage{
+			TotalTokens:      0,
+			CompletionTokens: 0,
+			PromptTokens:     0,
+		},
+	}
+
+	*p.Usage = *response.Usage
+
+	return response, nil
+}

providers/ollama/type.go

@@ -0,0 +1,48 @@
+package ollama
+
+import "time"
+
+type OllamaError struct {
+	Error string `json:"error,omitempty"`
+}
+
+type ChatRequest struct {
+	Model    string    `json:"model"`
+	Messages []Message `json:"messages,omitempty"`
+	Stream   bool      `json:"stream"`
+	Format   string    `json:"format,omitempty"`
+	Options  Option    `json:"options,omitempty"`
+}
+
+type Option struct {
+	Temperature float64 `json:"temperature,omitempty"`
+	Seed        *int    `json:"seed,omitempty"`
+	TopP        float64 `json:"top_p,omitempty"`
+	TopK        int     `json:"top_k,omitempty"`
+}
+
+type ChatResponse struct {
+	OllamaError
+	Model           string    `json:"model"`
+	CreatedAt       time.Time `json:"created_at"`
+	Message         Message   `json:"message,omitempty"`
+	Done            bool      `json:"done"`
+	EvalCount       int       `json:"eval_count,omitempty"`
+	PromptEvalCount int       `json:"prompt_eval_count,omitempty"`
+}
+
+type Message struct {
+	Role    string   `json:"role,omitempty"`
+	Content string   `json:"content,omitempty"`
+	Images  []string `json:"images,omitempty"`
+}
+
+type EmbeddingRequest struct {
+	Model  string `json:"model"`
+	Prompt string `json:"prompt"`
+}
+
+type EmbeddingResponse struct {
+	OllamaError
+	Embedding []float64 `json:"embedding,omitempty"`
+}

providers/providers.go

@@ -20,6 +20,7 @@ import (
 	"one-api/providers/midjourney"
 	"one-api/providers/minimax"
 	"one-api/providers/mistral"
+	"one-api/providers/ollama"
 	"one-api/providers/openai"
 	"one-api/providers/palm"
 	"one-api/providers/stabilityAI"
@@ -62,6 +63,7 @@ func init() {
 	providerFactories[common.ChannelTypeCohere] = cohere.CohereProviderFactory{}
 	providerFactories[common.ChannelTypeStabilityAI] = stabilityAI.StabilityAIProviderFactory{}
 	providerFactories[common.ChannelTypeCoze] = coze.CozeProviderFactory{}
+	providerFactories[common.ChannelTypeOllama] = ollama.OllamaProviderFactory{}
 
 }