feat: 算力引擎 Go 后端核心（new-api）
- 统一 OpenAI 兼容 /v1 中继 + 渠道/额度/令牌/流水 - 各领域包：relay 模型网关、model 数据层、controller 管理 API
This commit is contained in:
@@ -0,0 +1,111 @@
|
||||
package ollama
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/relay/channel"
|
||||
"github.com/QuantumNous/new-api/relay/channel/openai"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
relayconstant "github.com/QuantumNous/new-api/relay/constant"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type Adaptor struct {
|
||||
}
|
||||
|
||||
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{}
|
||||
openaiRequest, err := openaiAdaptor.ConvertClaudeRequest(c, info, request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
openaiRequest.(*dto.GeneralOpenAIRequest).StreamOptions = &dto.StreamOptions{
|
||||
IncludeUsage: true,
|
||||
}
|
||||
// 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) {
|
||||
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.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 {
|
||||
channel.SetupApiRequestHeader(info, c, req)
|
||||
req.Set("Authorization", "Bearer "+info.ApiKey)
|
||||
return nil
|
||||
}
|
||||
|
||||
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")
|
||||
}
|
||||
// decide generate or chat
|
||||
if strings.Contains(info.RequestURLPath, "/v1/completions") || info.RelayMode == relayconstant.RelayModeCompletions {
|
||||
return openAIToGenerate(c, request)
|
||||
}
|
||||
return openAIChatToOllamaChat(c, request)
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
|
||||
return requestOpenAI2Embeddings(request), nil
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
|
||||
switch info.RelayMode {
|
||||
case relayconstant.RelayModeEmbeddings:
|
||||
return ollamaEmbeddingHandler(c, info, resp)
|
||||
default:
|
||||
if info.IsStream {
|
||||
return ollamaStreamHandler(c, info, resp)
|
||||
}
|
||||
return ollamaChatHandler(c, info, resp)
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Adaptor) GetModelList() []string {
|
||||
return ModelList
|
||||
}
|
||||
|
||||
func (a *Adaptor) GetChannelName() string {
|
||||
return ChannelName
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
package ollama
|
||||
|
||||
var ModelList = []string{
|
||||
"llama3-7b",
|
||||
}
|
||||
|
||||
var ChannelName = "ollama"
|
||||
@@ -0,0 +1,108 @@
|
||||
package ollama
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
)
|
||||
|
||||
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"`
|
||||
ToolCallID string `json:"tool_call_id,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 {
|
||||
ID string `json:"id,omitempty"`
|
||||
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"`
|
||||
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"`
|
||||
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"`
|
||||
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"`
|
||||
Embeddings [][]float64 `json:"embeddings"`
|
||||
PromptEvalCount int `json:"prompt_eval_count,omitempty"`
|
||||
}
|
||||
|
||||
type OllamaTagsResponse struct {
|
||||
Models []OllamaModel `json:"models"`
|
||||
}
|
||||
|
||||
type OllamaModel struct {
|
||||
Name string `json:"name"`
|
||||
Size int64 `json:"size"`
|
||||
Digest string `json:"digest,omitempty"`
|
||||
ModifiedAt string `json:"modified_at"`
|
||||
Details OllamaModelDetail `json:"details,omitempty"`
|
||||
}
|
||||
|
||||
type OllamaModelDetail struct {
|
||||
ParentModel string `json:"parent_model,omitempty"`
|
||||
Format string `json:"format,omitempty"`
|
||||
Family string `json:"family,omitempty"`
|
||||
Families []string `json:"families,omitempty"`
|
||||
ParameterSize string `json:"parameter_size,omitempty"`
|
||||
QuantizationLevel string `json:"quantization_level,omitempty"`
|
||||
}
|
||||
|
||||
type OllamaPullRequest struct {
|
||||
Name string `json:"name"`
|
||||
Stream bool `json:"stream,omitempty"`
|
||||
}
|
||||
|
||||
type OllamaPullResponse struct {
|
||||
Status string `json:"status"`
|
||||
Digest string `json:"digest,omitempty"`
|
||||
Total int64 `json:"total,omitempty"`
|
||||
Completed int64 `json:"completed,omitempty"`
|
||||
}
|
||||
|
||||
type OllamaDeleteRequest struct {
|
||||
Name string `json:"name"`
|
||||
}
|
||||
@@ -0,0 +1,575 @@
|
||||
package ollama
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relay/helper"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/samber/lo"
|
||||
)
|
||||
|
||||
func toOllamaResponseFormat(responseFormat *dto.ResponseFormat) (any, error) {
|
||||
if responseFormat == nil {
|
||||
return nil, nil
|
||||
}
|
||||
switch responseFormat.Type {
|
||||
case "json", "json_object":
|
||||
return "json", nil
|
||||
case "json_schema":
|
||||
if len(responseFormat.JsonSchema) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
var jsonSchema dto.FormatJsonSchema
|
||||
if err := common.Unmarshal(responseFormat.JsonSchema, &jsonSchema); err != nil {
|
||||
return nil, fmt.Errorf("invalid ollama response format: %w", err)
|
||||
}
|
||||
return jsonSchema.Schema, nil
|
||||
default:
|
||||
return nil, nil
|
||||
}
|
||||
}
|
||||
|
||||
func openAIChatToOllamaChat(c *gin.Context, r *dto.GeneralOpenAIRequest) (*OllamaChatRequest, error) {
|
||||
think := r.Think
|
||||
if len(think) == 0 {
|
||||
effort := r.ReasoningEffort
|
||||
if len(r.Reasoning) > 0 {
|
||||
var reasoning dto.Reasoning
|
||||
if err := common.Unmarshal(r.Reasoning, &reasoning); err != nil {
|
||||
return nil, fmt.Errorf("invalid ollama reasoning: %w", err)
|
||||
}
|
||||
effort = lo.CoalesceOrEmpty(reasoning.Effort, effort)
|
||||
}
|
||||
if effort != "" {
|
||||
var thinkValue any
|
||||
switch effort {
|
||||
case "none":
|
||||
thinkValue = false
|
||||
case "low", "medium", "high", "max":
|
||||
thinkValue = effort
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported ollama reasoning effort %q", effort)
|
||||
}
|
||||
var err error
|
||||
think, err = common.Marshal(thinkValue)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshal ollama think: %w", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
chatReq := &OllamaChatRequest{
|
||||
Model: r.Model,
|
||||
Stream: lo.FromPtrOr(r.Stream, false),
|
||||
Options: map[string]any{},
|
||||
Think: think,
|
||||
}
|
||||
format, err := toOllamaResponseFormat(r.ResponseFormat)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
chatReq.Format = format
|
||||
|
||||
// options mapping
|
||||
if r.Temperature != nil {
|
||||
chatReq.Options["temperature"] = r.Temperature
|
||||
}
|
||||
if r.TopP != nil {
|
||||
chatReq.Options["top_p"] = lo.FromPtr(r.TopP)
|
||||
}
|
||||
if r.TopK != nil {
|
||||
chatReq.Options["top_k"] = lo.FromPtr(r.TopK)
|
||||
}
|
||||
if r.FrequencyPenalty != nil {
|
||||
chatReq.Options["frequency_penalty"] = lo.FromPtr(r.FrequencyPenalty)
|
||||
}
|
||||
if r.PresencePenalty != nil {
|
||||
chatReq.Options["presence_penalty"] = lo.FromPtr(r.PresencePenalty)
|
||||
}
|
||||
if r.Seed != nil {
|
||||
chatReq.Options["seed"] = int(lo.FromPtr(r.Seed))
|
||||
}
|
||||
if mt := r.GetMaxTokens(); mt != 0 {
|
||||
chatReq.Options["num_predict"] = int(mt)
|
||||
}
|
||||
|
||||
if r.Stop != nil {
|
||||
switch v := r.Stop.(type) {
|
||||
case string:
|
||||
chatReq.Options["stop"] = []string{v}
|
||||
case []string:
|
||||
chatReq.Options["stop"] = v
|
||||
case []any:
|
||||
arr := lo.FilterMap(v, func(item any, _ int) (string, bool) {
|
||||
value, ok := item.(string)
|
||||
return value, ok
|
||||
})
|
||||
if len(arr) > 0 {
|
||||
chatReq.Options["stop"] = arr
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if len(r.Tools) > 0 {
|
||||
chatReq.Tools = lo.Map(r.Tools, func(tool dto.ToolCallRequest, _ int) OllamaTool {
|
||||
return OllamaTool{
|
||||
Type: "function",
|
||||
Function: OllamaToolFunction{
|
||||
Name: tool.Function.Name,
|
||||
Description: tool.Function.Description,
|
||||
Parameters: tool.Function.Parameters,
|
||||
},
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
chatReq.Messages = make([]OllamaChatMessage, 0, len(r.Messages))
|
||||
toolNamesByCallID := make(map[string]string)
|
||||
for _, m := range r.Messages {
|
||||
var textBuilder strings.Builder
|
||||
var images []string
|
||||
if m.IsStringContent() {
|
||||
textBuilder.WriteString(m.StringContent())
|
||||
} else {
|
||||
parts := m.ParseContent()
|
||||
for _, part := range parts {
|
||||
if part.Type == dto.ContentTypeImageURL {
|
||||
source := part.ToFileSource()
|
||||
if source != nil {
|
||||
base64Data, _, err := service.GetBase64Data(c, source, "fetch image for ollama chat")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if base64Data != "" {
|
||||
images = append(images, base64Data)
|
||||
}
|
||||
}
|
||||
} else if part.Type == dto.ContentTypeText {
|
||||
textBuilder.WriteString(part.Text)
|
||||
}
|
||||
}
|
||||
}
|
||||
cm := OllamaChatMessage{Role: m.Role, Content: textBuilder.String()}
|
||||
if len(images) > 0 {
|
||||
cm.Images = images
|
||||
}
|
||||
if m.Role == "assistant" {
|
||||
if reasoning, ok := lo.Coalesce(m.ReasoningContent, m.Reasoning); ok {
|
||||
thinking, err := common.Marshal(*reasoning)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshal ollama thinking: %w", err)
|
||||
}
|
||||
cm.Thinking = thinking
|
||||
}
|
||||
}
|
||||
if m.Role == "tool" {
|
||||
cm.ToolCallID = m.ToolCallId
|
||||
cm.ToolName = lo.CoalesceOrEmpty(lo.FromPtr(m.Name), toolNamesByCallID[m.ToolCallId])
|
||||
}
|
||||
if m.ToolCalls != nil && len(m.ToolCalls) > 0 {
|
||||
parsed := m.ParseToolCalls()
|
||||
if len(parsed) > 0 {
|
||||
calls := make([]OllamaToolCall, 0, len(parsed))
|
||||
for _, tc := range parsed {
|
||||
var args interface{}
|
||||
if tc.Function.Arguments != "" {
|
||||
_ = common.Unmarshal([]byte(tc.Function.Arguments), &args)
|
||||
}
|
||||
if args == nil {
|
||||
args = map[string]any{}
|
||||
}
|
||||
oc := OllamaToolCall{ID: tc.ID}
|
||||
oc.Function.Name = tc.Function.Name
|
||||
oc.Function.Arguments = args
|
||||
calls = append(calls, oc)
|
||||
if tc.ID != "" {
|
||||
toolNamesByCallID[tc.ID] = tc.Function.Name
|
||||
}
|
||||
}
|
||||
cm.ToolCalls = calls
|
||||
}
|
||||
}
|
||||
chatReq.Messages = append(chatReq.Messages, cm)
|
||||
}
|
||||
return chatReq, nil
|
||||
}
|
||||
|
||||
// openAIToGenerate converts OpenAI completions request to Ollama generate
|
||||
func openAIToGenerate(c *gin.Context, r *dto.GeneralOpenAIRequest) (*OllamaGenerateRequest, error) {
|
||||
gen := &OllamaGenerateRequest{
|
||||
Model: r.Model,
|
||||
Stream: lo.FromPtrOr(r.Stream, false),
|
||||
Options: map[string]any{},
|
||||
Think: r.Think,
|
||||
}
|
||||
// Prompt may be in r.Prompt (string or []any)
|
||||
if r.Prompt != nil {
|
||||
switch v := r.Prompt.(type) {
|
||||
case string:
|
||||
gen.Prompt = v
|
||||
case []any:
|
||||
var sb strings.Builder
|
||||
for _, it := range v {
|
||||
if s, ok := it.(string); ok {
|
||||
sb.WriteString(s)
|
||||
}
|
||||
}
|
||||
gen.Prompt = sb.String()
|
||||
default:
|
||||
gen.Prompt = fmt.Sprintf("%v", r.Prompt)
|
||||
}
|
||||
}
|
||||
if r.Suffix != nil {
|
||||
if s, ok := r.Suffix.(string); ok {
|
||||
gen.Suffix = s
|
||||
}
|
||||
}
|
||||
format, err := toOllamaResponseFormat(r.ResponseFormat)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
gen.Format = format
|
||||
if r.Temperature != nil {
|
||||
gen.Options["temperature"] = r.Temperature
|
||||
}
|
||||
if r.TopP != nil {
|
||||
gen.Options["top_p"] = lo.FromPtr(r.TopP)
|
||||
}
|
||||
if r.TopK != nil {
|
||||
gen.Options["top_k"] = lo.FromPtr(r.TopK)
|
||||
}
|
||||
if r.FrequencyPenalty != nil {
|
||||
gen.Options["frequency_penalty"] = lo.FromPtr(r.FrequencyPenalty)
|
||||
}
|
||||
if r.PresencePenalty != nil {
|
||||
gen.Options["presence_penalty"] = lo.FromPtr(r.PresencePenalty)
|
||||
}
|
||||
if r.Seed != nil {
|
||||
gen.Options["seed"] = int(lo.FromPtr(r.Seed))
|
||||
}
|
||||
if mt := r.GetMaxTokens(); mt != 0 {
|
||||
gen.Options["num_predict"] = int(mt)
|
||||
}
|
||||
if r.Stop != nil {
|
||||
switch v := r.Stop.(type) {
|
||||
case string:
|
||||
gen.Options["stop"] = []string{v}
|
||||
case []string:
|
||||
gen.Options["stop"] = v
|
||||
case []any:
|
||||
arr := lo.FilterMap(v, func(item any, _ int) (string, bool) {
|
||||
value, ok := item.(string)
|
||||
return value, ok
|
||||
})
|
||||
if len(arr) > 0 {
|
||||
gen.Options["stop"] = arr
|
||||
}
|
||||
}
|
||||
}
|
||||
return gen, nil
|
||||
}
|
||||
|
||||
func requestOpenAI2Embeddings(r dto.EmbeddingRequest) *OllamaEmbeddingRequest {
|
||||
opts := map[string]any{}
|
||||
if r.Temperature != nil {
|
||||
opts["temperature"] = r.Temperature
|
||||
}
|
||||
if r.TopP != nil {
|
||||
opts["top_p"] = lo.FromPtr(r.TopP)
|
||||
}
|
||||
if r.FrequencyPenalty != nil {
|
||||
opts["frequency_penalty"] = lo.FromPtr(r.FrequencyPenalty)
|
||||
}
|
||||
if r.PresencePenalty != nil {
|
||||
opts["presence_penalty"] = lo.FromPtr(r.PresencePenalty)
|
||||
}
|
||||
if r.Seed != nil {
|
||||
opts["seed"] = int(lo.FromPtr(r.Seed))
|
||||
}
|
||||
dimensions := lo.FromPtrOr(r.Dimensions, 0)
|
||||
if r.Dimensions != nil {
|
||||
opts["dimensions"] = dimensions
|
||||
}
|
||||
input := r.ParseInput()
|
||||
if len(input) == 1 {
|
||||
return &OllamaEmbeddingRequest{Model: r.Model, Input: input[0], Options: opts, Dimensions: dimensions}
|
||||
}
|
||||
return &OllamaEmbeddingRequest{Model: r.Model, Input: input, Options: opts, Dimensions: dimensions}
|
||||
}
|
||||
|
||||
func ollamaEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
|
||||
var oResp OllamaEmbeddingResponse
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
|
||||
}
|
||||
service.CloseResponseBodyGracefully(resp)
|
||||
if err = common.Unmarshal(body, &oResp); err != nil {
|
||||
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
|
||||
}
|
||||
if oResp.Error != "" {
|
||||
return nil, types.NewOpenAIError(fmt.Errorf("ollama error: %s", oResp.Error), types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
|
||||
}
|
||||
data := make([]dto.OpenAIEmbeddingResponseItem, 0, len(oResp.Embeddings))
|
||||
for i, emb := range oResp.Embeddings {
|
||||
data = append(data, dto.OpenAIEmbeddingResponseItem{Index: i, Object: "embedding", Embedding: emb})
|
||||
}
|
||||
usage := &dto.Usage{PromptTokens: oResp.PromptEvalCount, CompletionTokens: 0, TotalTokens: oResp.PromptEvalCount}
|
||||
embResp := &dto.OpenAIEmbeddingResponse{Object: "list", Data: data, Model: info.UpstreamModelName, Usage: *usage}
|
||||
out, _ := common.Marshal(embResp)
|
||||
service.IOCopyBytesGracefully(c, resp, out)
|
||||
return usage, nil
|
||||
}
|
||||
|
||||
func FetchOllamaModels(baseURL, apiKey string) ([]OllamaModel, error) {
|
||||
url := fmt.Sprintf("%s/api/tags", baseURL)
|
||||
|
||||
client := &http.Client{}
|
||||
request, err := http.NewRequest("GET", url, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("创建请求失败: %v", err)
|
||||
}
|
||||
|
||||
// Ollama 通常不需要 Bearer token，但为了兼容性保留
|
||||
if apiKey != "" {
|
||||
request.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
}
|
||||
|
||||
response, err := client.Do(request)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("请求失败: %v", err)
|
||||
}
|
||||
defer response.Body.Close()
|
||||
|
||||
if response.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(response.Body)
|
||||
return nil, fmt.Errorf("服务器返回错误 %d: %s", response.StatusCode, string(body))
|
||||
}
|
||||
|
||||
var tagsResponse OllamaTagsResponse
|
||||
body, err := io.ReadAll(response.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("读取响应失败: %v", err)
|
||||
}
|
||||
|
||||
err = common.Unmarshal(body, &tagsResponse)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("解析响应失败: %v", err)
|
||||
}
|
||||
|
||||
return tagsResponse.Models, nil
|
||||
}
|
||||
|
||||
// 拉取 Ollama 模型 (非流式)
|
||||
func PullOllamaModel(baseURL, apiKey, modelName string) error {
|
||||
url := fmt.Sprintf("%s/api/pull", baseURL)
|
||||
|
||||
pullRequest := OllamaPullRequest{
|
||||
Name: modelName,
|
||||
Stream: false, // 非流式，简化处理
|
||||
}
|
||||
|
||||
requestBody, err := common.Marshal(pullRequest)
|
||||
if err != nil {
|
||||
return fmt.Errorf("序列化请求失败: %v", err)
|
||||
}
|
||||
|
||||
client := &http.Client{
|
||||
Timeout: 30 * 60 * 1000 * time.Millisecond, // 30分钟超时，支持大模型
|
||||
}
|
||||
request, err := http.NewRequest("POST", url, strings.NewReader(string(requestBody)))
|
||||
if err != nil {
|
||||
return fmt.Errorf("创建请求失败: %v", err)
|
||||
}
|
||||
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
if apiKey != "" {
|
||||
request.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
}
|
||||
|
||||
response, err := client.Do(request)
|
||||
if err != nil {
|
||||
return fmt.Errorf("请求失败: %v", err)
|
||||
}
|
||||
defer response.Body.Close()
|
||||
|
||||
if response.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(response.Body)
|
||||
return fmt.Errorf("拉取模型失败 %d: %s", response.StatusCode, string(body))
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// 流式拉取 Ollama 模型 (支持进度回调)
|
||||
func PullOllamaModelStream(baseURL, apiKey, modelName string, progressCallback func(OllamaPullResponse)) error {
|
||||
url := fmt.Sprintf("%s/api/pull", baseURL)
|
||||
|
||||
pullRequest := OllamaPullRequest{
|
||||
Name: modelName,
|
||||
Stream: true, // 启用流式
|
||||
}
|
||||
|
||||
requestBody, err := common.Marshal(pullRequest)
|
||||
if err != nil {
|
||||
return fmt.Errorf("序列化请求失败: %v", err)
|
||||
}
|
||||
|
||||
client := &http.Client{
|
||||
Timeout: 60 * 60 * 1000 * time.Millisecond, // 1小时超时，支持超大模型
|
||||
}
|
||||
request, err := http.NewRequest("POST", url, strings.NewReader(string(requestBody)))
|
||||
if err != nil {
|
||||
return fmt.Errorf("创建请求失败: %v", err)
|
||||
}
|
||||
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
if apiKey != "" {
|
||||
request.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
}
|
||||
|
||||
response, err := client.Do(request)
|
||||
if err != nil {
|
||||
return fmt.Errorf("请求失败: %v", err)
|
||||
}
|
||||
defer response.Body.Close()
|
||||
|
||||
if response.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(response.Body)
|
||||
return fmt.Errorf("拉取模型失败 %d: %s", response.StatusCode, string(body))
|
||||
}
|
||||
|
||||
// 读取流式响应
|
||||
scanner := helper.NewStreamScanner(response.Body)
|
||||
successful := false
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
if strings.TrimSpace(line) == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
var pullResponse OllamaPullResponse
|
||||
if err := common.Unmarshal([]byte(line), &pullResponse); err != nil {
|
||||
continue // 忽略解析失败的行
|
||||
}
|
||||
|
||||
if progressCallback != nil {
|
||||
progressCallback(pullResponse)
|
||||
}
|
||||
|
||||
// 检查是否出现错误或完成
|
||||
if strings.EqualFold(pullResponse.Status, "error") {
|
||||
return fmt.Errorf("拉取模型失败: %s", strings.TrimSpace(line))
|
||||
}
|
||||
if strings.EqualFold(pullResponse.Status, "success") {
|
||||
successful = true
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if err := scanner.Err(); err != nil {
|
||||
return fmt.Errorf("读取流式响应失败: %v", err)
|
||||
}
|
||||
|
||||
if !successful {
|
||||
return fmt.Errorf("拉取模型未完成: 未收到成功状态")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// 删除 Ollama 模型
|
||||
func DeleteOllamaModel(baseURL, apiKey, modelName string) error {
|
||||
url := fmt.Sprintf("%s/api/delete", baseURL)
|
||||
|
||||
deleteRequest := OllamaDeleteRequest{
|
||||
Name: modelName,
|
||||
}
|
||||
|
||||
requestBody, err := common.Marshal(deleteRequest)
|
||||
if err != nil {
|
||||
return fmt.Errorf("序列化请求失败: %v", err)
|
||||
}
|
||||
|
||||
client := &http.Client{}
|
||||
request, err := http.NewRequest("DELETE", url, strings.NewReader(string(requestBody)))
|
||||
if err != nil {
|
||||
return fmt.Errorf("创建请求失败: %v", err)
|
||||
}
|
||||
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
if apiKey != "" {
|
||||
request.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
}
|
||||
|
||||
response, err := client.Do(request)
|
||||
if err != nil {
|
||||
return fmt.Errorf("请求失败: %v", err)
|
||||
}
|
||||
defer response.Body.Close()
|
||||
|
||||
if response.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(response.Body)
|
||||
return fmt.Errorf("删除模型失败 %d: %s", response.StatusCode, string(body))
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func FetchOllamaVersion(baseURL, apiKey string) (string, error) {
|
||||
trimmedBase := strings.TrimRight(baseURL, "/")
|
||||
if trimmedBase == "" {
|
||||
return "", fmt.Errorf("baseURL 为空")
|
||||
}
|
||||
|
||||
url := fmt.Sprintf("%s/api/version", trimmedBase)
|
||||
|
||||
client := &http.Client{Timeout: 10 * time.Second}
|
||||
request, err := http.NewRequest("GET", url, nil)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("创建请求失败: %v", err)
|
||||
}
|
||||
|
||||
if apiKey != "" {
|
||||
request.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
}
|
||||
|
||||
response, err := client.Do(request)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("请求失败: %v", err)
|
||||
}
|
||||
defer response.Body.Close()
|
||||
|
||||
body, err := io.ReadAll(response.Body)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("读取响应失败: %v", err)
|
||||
}
|
||||
|
||||
if response.StatusCode != http.StatusOK {
|
||||
return "", fmt.Errorf("查询版本失败 %d: %s", response.StatusCode, string(body))
|
||||
}
|
||||
|
||||
var versionResp struct {
|
||||
Version string `json:"version"`
|
||||
}
|
||||
|
||||
if err := common.Unmarshal(body, &versionResp); err != nil {
|
||||
return "", fmt.Errorf("解析响应失败: %v", err)
|
||||
}
|
||||
|
||||
if versionResp.Version == "" {
|
||||
return "", fmt.Errorf("未返回版本信息")
|
||||
}
|
||||
|
||||
return versionResp.Version, nil
|
||||
}
|
||||
@@ -0,0 +1,346 @@
|
||||
package ollama
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relay/helper"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
|
||||
"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 []OllamaToolCall `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 ollamaToolCallsToOpenAI(toolCalls []OllamaToolCall, startIndex int, includeIndex bool) ([]dto.ToolCallResponse, int) {
|
||||
if len(toolCalls) == 0 {
|
||||
return nil, startIndex
|
||||
}
|
||||
result := make([]dto.ToolCallResponse, 0, len(toolCalls))
|
||||
for _, tc := range toolCalls {
|
||||
var argBytes []byte
|
||||
var err error
|
||||
if tc.Function.Arguments == nil {
|
||||
argBytes = []byte("{}")
|
||||
} else {
|
||||
argBytes, err = common.Marshal(tc.Function.Arguments)
|
||||
if err != nil || len(argBytes) == 0 {
|
||||
argBytes = []byte("{}")
|
||||
}
|
||||
}
|
||||
toolCallID := tc.ID
|
||||
if toolCallID == "" {
|
||||
toolCallID = fmt.Sprintf("call_%d", startIndex)
|
||||
}
|
||||
tr := dto.ToolCallResponse{
|
||||
ID: toolCallID,
|
||||
Type: "function",
|
||||
Function: dto.FunctionResponse{
|
||||
Name: tc.Function.Name,
|
||||
Arguments: string(argBytes),
|
||||
},
|
||||
}
|
||||
if includeIndex {
|
||||
tr.SetIndex(startIndex)
|
||||
}
|
||||
startIndex++
|
||||
result = append(result, tr)
|
||||
}
|
||||
return result, startIndex
|
||||
}
|
||||
|
||||
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 := helper.NewStreamScanner(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 := common.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" {
|
||||
// Unmarshal the JSON string to get the actual content without quotes
|
||||
var thinkingContent string
|
||||
if err := common.Unmarshal(chunk.Message.Thinking, &thinkingContent); err == nil {
|
||||
delta.Choices[0].Delta.SetReasoningContent(thinkingContent)
|
||||
} else {
|
||||
// Fallback to raw string if it's not a JSON string
|
||||
delta.Choices[0].Delta.SetReasoningContent(raw)
|
||||
}
|
||||
}
|
||||
}
|
||||
// tool calls
|
||||
if chunk.Message != nil && len(chunk.Message.ToolCalls) > 0 {
|
||||
delta.Choices[0].Delta.ToolCalls, toolCallIndex = ollamaToolCallsToOpenAI(chunk.Message.ToolCalls, toolCallIndex, true)
|
||||
}
|
||||
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"
|
||||
}
|
||||
if toolCallIndex > 0 {
|
||||
finishReason = constant.FinishReasonToolCalls
|
||||
}
|
||||
// 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
|
||||
toolCallIndex int
|
||||
toolCalls []dto.ToolCallResponse
|
||||
)
|
||||
for _, ln := range lines {
|
||||
ln = strings.TrimSpace(ln)
|
||||
if ln == "" {
|
||||
continue
|
||||
}
|
||||
var ck ollamaChatStreamChunk
|
||||
if err := common.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" {
|
||||
// Unmarshal the JSON string to get the actual content without quotes
|
||||
var thinkingContent string
|
||||
if err := common.Unmarshal(ck.Message.Thinking, &thinkingContent); err == nil {
|
||||
reasoningBuilder.WriteString(thinkingContent)
|
||||
} else {
|
||||
// Fallback to raw string if it's not a JSON string
|
||||
reasoningBuilder.WriteString(raw)
|
||||
}
|
||||
}
|
||||
}
|
||||
if ck.Message != nil && ck.Message.Content != "" {
|
||||
aggContent.WriteString(ck.Message.Content)
|
||||
} else if ck.Response != "" {
|
||||
aggContent.WriteString(ck.Response)
|
||||
}
|
||||
if ck.Message != nil && len(ck.Message.ToolCalls) > 0 {
|
||||
var converted []dto.ToolCallResponse
|
||||
converted, toolCallIndex = ollamaToolCallsToOpenAI(ck.Message.ToolCalls, toolCallIndex, false)
|
||||
toolCalls = append(toolCalls, converted...)
|
||||
}
|
||||
}
|
||||
|
||||
if !parsedAny {
|
||||
var single ollamaChatStreamChunk
|
||||
if err := common.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" {
|
||||
// Unmarshal the JSON string to get the actual content without quotes
|
||||
var thinkingContent string
|
||||
if err := common.Unmarshal(single.Message.Thinking, &thinkingContent); err == nil {
|
||||
reasoningBuilder.WriteString(thinkingContent)
|
||||
} else {
|
||||
// Fallback to raw string if it's not a JSON string
|
||||
reasoningBuilder.WriteString(raw)
|
||||
}
|
||||
}
|
||||
}
|
||||
aggContent.WriteString(single.Message.Content)
|
||||
if len(single.Message.ToolCalls) > 0 {
|
||||
var converted []dto.ToolCallResponse
|
||||
converted, toolCallIndex = ollamaToolCallsToOpenAI(single.Message.ToolCalls, toolCallIndex, false)
|
||||
toolCalls = append(toolCalls, converted...)
|
||||
}
|
||||
} 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"
|
||||
}
|
||||
if len(toolCalls) > 0 {
|
||||
finishReason = constant.FinishReasonToolCalls
|
||||
}
|
||||
|
||||
msg := dto.Message{Role: "assistant", Content: contentPtr(content)}
|
||||
if len(toolCalls) > 0 {
|
||||
if rawToolCalls, err := common.Marshal(toolCalls); err == nil {
|
||||
msg.ToolCalls = rawToolCalls
|
||||
}
|
||||
}
|
||||
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
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
package ollama
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestOllamaChatHandlerNonStreamToolCalls(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
raw string
|
||||
wantID string
|
||||
}{
|
||||
{
|
||||
name: "compact json per-line parse path",
|
||||
raw: `{"model":"llama3.1","created_at":"2026-05-27T12:00:00Z","message":{"role":"assistant","content":"","tool_calls":[{"id":"call_upstream","function":{"name":"get_weather","arguments":{"city":"Paris","days":0}}}]},"done":true,"done_reason":"stop","prompt_eval_count":5,"eval_count":7}`,
|
||||
wantID: "call_upstream",
|
||||
},
|
||||
{
|
||||
name: "pretty json fallback parse path",
|
||||
raw: `{
|
||||
"model": "llama3.1",
|
||||
"created_at": "2026-05-27T12:00:00Z",
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"arguments": {
|
||||
"city": "Paris",
|
||||
"days": 0
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
"done": true,
|
||||
"done_reason": "stop",
|
||||
"prompt_eval_count": 5,
|
||||
"eval_count": 7
|
||||
}`,
|
||||
wantID: "call_0",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: make(http.Header),
|
||||
Body: io.NopCloser(strings.NewReader(tt.raw)),
|
||||
}
|
||||
|
||||
usage, apiErr := ollamaChatHandler(c, &relaycommon.RelayInfo{
|
||||
ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "fallback-model"},
|
||||
}, resp)
|
||||
require.Nil(t, apiErr)
|
||||
require.NotNil(t, usage)
|
||||
assert.Equal(t, 12, usage.TotalTokens)
|
||||
|
||||
var out dto.OpenAITextResponse
|
||||
require.NoError(t, common.Unmarshal(w.Body.Bytes(), &out))
|
||||
require.Len(t, out.Choices, 1)
|
||||
assert.Equal(t, constant.FinishReasonToolCalls, out.Choices[0].FinishReason)
|
||||
|
||||
var toolCalls []dto.ToolCallResponse
|
||||
require.NoError(t, common.Unmarshal(out.Choices[0].Message.ToolCalls, &toolCalls))
|
||||
require.Len(t, toolCalls, 1)
|
||||
assert.Equal(t, tt.wantID, toolCalls[0].ID)
|
||||
assert.Equal(t, "function", toolCalls[0].Type)
|
||||
assert.Equal(t, "get_weather", toolCalls[0].Function.Name)
|
||||
assert.Nil(t, toolCalls[0].Index)
|
||||
|
||||
var args map[string]any
|
||||
require.NoError(t, common.Unmarshal([]byte(toolCalls[0].Function.Arguments), &args))
|
||||
assert.Equal(t, "Paris", args["city"])
|
||||
assert.Equal(t, float64(0), args["days"])
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user