diff --git a/internal/entity/models/gitee.go b/internal/entity/models/gitee.go index 57399f6886..58fa8d588b 100644 --- a/internal/entity/models/gitee.go +++ b/internal/entity/models/gitee.go @@ -25,6 +25,7 @@ import ( "mime/multipart" "net/http" "ragflow/internal/common" + "sort" "strings" "time" ) @@ -76,15 +77,20 @@ func (g *GiteeModel) ChatWithMessages(modelName string, messages []Message, apiC "role": msg.Role, "content": msg.Content, } + if msg.ToolCallID != "" { + apiMessages[i]["tool_call_id"] = msg.ToolCallID + } + if len(msg.ToolCalls) > 0 { + apiMessages[i]["tool_calls"] = msg.ToolCalls + } } common.Info(fmt.Sprintf("GiteeAPI messages: %+v", apiMessages)) // Build request body reqBody := map[string]interface{}{ - "model": modelName, - "messages": apiMessages, - "stream": false, - "temperature": 1, + "model": modelName, + "messages": apiMessages, + "stream": false, } if chatModelConfig != nil { @@ -119,6 +125,12 @@ func (g *GiteeModel) ChatWithMessages(modelName string, messages []Message, apiC } } } + if chatModelConfig.Tools != nil { + reqBody["tools"] = chatModelConfig.Tools + } + if chatModelConfig.ToolChoice != nil { + reqBody["tool_choice"] = chatModelConfig.ToolChoice + } } jsonData, err := json.Marshal(reqBody) @@ -175,10 +187,7 @@ func (g *GiteeModel) ChatWithMessages(modelName string, messages []Message, apiC return nil, fmt.Errorf("invalid message format") } - content, ok := messageMap["content"].(string) - if !ok { - return nil, fmt.Errorf("invalid content format") - } + content, _ := messageMap["content"].(string) // Handle thinking/reasoning if enabled var reasonContent string @@ -199,9 +208,24 @@ func (g *GiteeModel) ChatWithMessages(modelName string, messages []Message, apiC } } + var toolCalls []map[string]any + if tcs, ok := messageMap["tool_calls"].([]any); ok { + for _, tc := range tcs { + if toolCall, ok := tc.(map[string]any); ok { + toolCalls = append(toolCalls, toolCall) + } + } + } + chatResponse := &ChatResponse{ Answer: &content, ReasonContent: &reasonContent, + ToolCalls: toolCalls, + } + if pt, ct, tt := extractUsageFromMap(result); tt > 0 { + chatResponse.Usage = &TokenUsage{ + PromptTokens: pt, CompletionTokens: ct, TotalTokens: tt, + } } return chatResponse, nil @@ -230,50 +254,61 @@ func (g *GiteeModel) ChatStreamlyWithSender(modelName string, messages []Message "role": msg.Role, "content": msg.Content, } + if msg.ToolCallID != "" { + apiMessages[i]["tool_call_id"] = msg.ToolCallID + } + if len(msg.ToolCalls) > 0 { + apiMessages[i]["tool_calls"] = msg.ToolCalls + } } // Build request body with streaming enabled reqBody := map[string]interface{}{ - "model": modelName, - "messages": apiMessages, - "stream": true, - "temperature": 1, + "model": modelName, + "messages": apiMessages, + "stream": true, } - if chatModelConfig.Stream != nil { - reqBody["stream"] = *chatModelConfig.Stream - } + if chatModelConfig != nil { + if chatModelConfig.Stream != nil { + reqBody["stream"] = *chatModelConfig.Stream + } - if chatModelConfig.MaxTokens != nil { - reqBody["max_tokens"] = *chatModelConfig.MaxTokens - } + if chatModelConfig.MaxTokens != nil { + reqBody["max_tokens"] = *chatModelConfig.MaxTokens + } - if chatModelConfig.Temperature != nil { - reqBody["temperature"] = *chatModelConfig.Temperature - } + if chatModelConfig.Temperature != nil { + reqBody["temperature"] = *chatModelConfig.Temperature + } - if chatModelConfig.DoSample != nil { - reqBody["do_sample"] = *chatModelConfig.DoSample - } + if chatModelConfig.DoSample != nil { + reqBody["do_sample"] = *chatModelConfig.DoSample + } - if chatModelConfig.TopP != nil { - reqBody["top_p"] = *chatModelConfig.TopP - } + if chatModelConfig.TopP != nil { + reqBody["top_p"] = *chatModelConfig.TopP + } - if chatModelConfig.Stop != nil { - reqBody["stop"] = *chatModelConfig.Stop - } + if chatModelConfig.Stop != nil { + reqBody["stop"] = *chatModelConfig.Stop + } - if chatModelConfig.Thinking != nil { - if *chatModelConfig.Thinking { - reqBody["thinking"] = map[string]interface{}{ - "type": "enabled", - } - } else { - reqBody["thinking"] = map[string]interface{}{ - "type": "disabled", + if chatModelConfig.Thinking != nil { + if *chatModelConfig.Thinking { + reqBody["thinking"] = map[string]interface{}{ + "type": "enabled", + } + } else { + reqBody["thinking"] = map[string]interface{}{ + "type": "disabled", + } } } + + if chatModelConfig.Tools != nil { + reqBody["tools"] = chatModelConfig.Tools + } } jsonData, err := json.Marshal(reqBody) @@ -308,6 +343,7 @@ func (g *GiteeModel) ChatStreamlyWithSender(modelName string, messages []Message answerPhase := false sawTerminal := false + accumulatedToolCalls := make(map[int]map[string]any) done, err := ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error { choices, ok := event["choices"].([]interface{}) if !ok || len(choices) == 0 { @@ -356,6 +392,56 @@ func (g *GiteeModel) ChatStreamlyWithSender(modelName string, messages []Message } } + if tcs, ok := delta["tool_calls"].([]interface{}); ok { + for _, tc := range tcs { + tcMap, ok := tc.(map[string]interface{}) + if !ok { + continue + } + idxF, ok := tcMap["index"].(float64) + if !ok { + continue + } + idx := int(idxF) + existing, hasExisting := accumulatedToolCalls[idx] + if !hasExisting { + accumulatedToolCalls[idx] = cloneMap(tcMap) + continue + } + if id, ok := tcMap["id"].(string); ok && id != "" { + if eid, ok := existing["id"].(string); ok { + existing["id"] = eid + id + } else { + existing["id"] = id + } + } + if typ, ok := tcMap["type"].(string); ok && typ != "" { + existing["type"] = typ + } + if fn, ok := tcMap["function"].(map[string]interface{}); ok { + ef, ok := existing["function"].(map[string]interface{}) + if !ok { + ef = make(map[string]interface{}) + existing["function"] = ef + } + if name, ok := fn["name"].(string); ok && name != "" { + if en, ok := ef["name"].(string); ok { + ef["name"] = en + name + } else { + ef["name"] = name + } + } + if args, ok := fn["arguments"].(string); ok && args != "" { + if ea, ok := ef["arguments"].(string); ok { + ef["arguments"] = ea + args + } else { + ef["arguments"] = args + } + } + } + } + } + finishReason, ok := firstChoice["finish_reason"].(string) if ok && finishReason != "" { sawTerminal = true @@ -375,6 +461,19 @@ func (g *GiteeModel) ChatStreamlyWithSender(modelName string, messages []Message } } + if len(accumulatedToolCalls) > 0 && chatModelConfig != nil { + indices := make([]int, 0, len(accumulatedToolCalls)) + for idx := range accumulatedToolCalls { + indices = append(indices, idx) + } + sort.Ints(indices) + tcs := make([]map[string]interface{}, 0, len(accumulatedToolCalls)) + for _, idx := range indices { + tcs = append(tcs, accumulatedToolCalls[idx]) + } + chatModelConfig.ToolCallsResult = &tcs + } + // Send [DONE] marker for OpenAI compatibility endOfStream := "[DONE]" if err = sender(&endOfStream, nil); err != nil { diff --git a/internal/entity/models/types.go b/internal/entity/models/types.go index 086a08b2d8..6987cfd8b6 100644 --- a/internal/entity/models/types.go +++ b/internal/entity/models/types.go @@ -7,10 +7,10 @@ import ( // Message represents a chat message with role and content // -// Content is interface{} to support different formats: -// - string: plain text message (e.g., "Hello") -// - []interface{}: multimodal content array where each element is map[string]interface{} -// (e.g., [{"type": "text", "text": "..."}, {"type": "image_url", "image_url": {"url": "..."}}]) +// is interface{} to support different formats: +// - string: plain text message (e.g., "Hello") +// - []interface{}: multimodal content array where each element is map[string]interface{} +// (e.g., [{"type": "text", "text": "..."}, {"type": "image_url", "image_url": {"url": "..."}}]) type Message struct { Role string `json:"role"` Content interface{} `json:"content"` @@ -66,7 +66,7 @@ type ChatResponse struct { Usage *TokenUsage `json:"usage,omitempty"` } -// ChatUsage holds token usage split for one LLM call. Consumed by +// TokenUsage holds token usage split for one LLM call. Consumed by // LLMBundle for accurate Langfuse reporting and run aggregation. // Mirrors Python's common.token_utils.usage_from_response() split. type TokenUsage struct { diff --git a/internal/entity/models/volcengine.go b/internal/entity/models/volcengine.go index db318b7d62..74909e8f5e 100644 --- a/internal/entity/models/volcengine.go +++ b/internal/entity/models/volcengine.go @@ -24,6 +24,7 @@ import ( "io" "net/http" "ragflow/internal/common" + "sort" "strings" ) @@ -74,14 +75,19 @@ func (v *VolcEngine) ChatWithMessages(modelName string, messages []Message, apiC "role": msg.Role, "content": msg.Content, } + if msg.ToolCallID != "" { + apiMessages[i]["tool_call_id"] = msg.ToolCallID + } + if len(msg.ToolCalls) > 0 { + apiMessages[i]["tool_calls"] = msg.ToolCalls + } } // Build request body reqBody := map[string]interface{}{ - "model": modelName, - "messages": apiMessages, - "stream": false, - "temperature": 1, + "model": modelName, + "messages": apiMessages, + "stream": false, } if chatModelConfig != nil { @@ -124,6 +130,9 @@ func (v *VolcEngine) ChatWithMessages(modelName string, messages []Message, apiC case "high": thinkingFlag = "enabled" reqBody["reasoning_effort"] = "high" + case "xhigh": + thinkingFlag = "enabled" + reqBody["reasoning_effort"] = "xhigh" default: thinkingFlag = "enabled" reqBody["reasoning_effort"] = effort @@ -137,6 +146,13 @@ func (v *VolcEngine) ChatWithMessages(modelName string, messages []Message, apiC } } } + + if chatModelConfig.Tools != nil { + reqBody["tools"] = chatModelConfig.Tools + } + if chatModelConfig.ToolChoice != nil { + reqBody["tool_choice"] = chatModelConfig.ToolChoice + } } jsonData, err := json.Marshal(reqBody) @@ -191,10 +207,7 @@ func (v *VolcEngine) ChatWithMessages(modelName string, messages []Message, apiC return nil, fmt.Errorf("invalid message format") } - content, ok := messageMap["content"].(string) - if !ok { - return nil, fmt.Errorf("invalid content format") - } + content, _ := messageMap["content"].(string) var reasonContent string if chatModelConfig != nil && chatModelConfig.Thinking != nil && *chatModelConfig.Thinking { @@ -207,9 +220,24 @@ func (v *VolcEngine) ChatWithMessages(modelName string, messages []Message, apiC } } + var toolCalls []map[string]any + if tcs, ok := messageMap["tool_calls"].([]any); ok { + for _, tc := range tcs { + if toolCall, ok := tc.(map[string]any); ok { + toolCalls = append(toolCalls, toolCall) + } + } + } + chatResponse := &ChatResponse{ Answer: &content, ReasonContent: &reasonContent, + ToolCalls: toolCalls, + } + if pt, ct, tt := extractUsageFromMap(result); tt > 0 { + chatResponse.Usage = &TokenUsage{ + PromptTokens: pt, CompletionTokens: ct, TotalTokens: tt, + } } return chatResponse, nil @@ -238,84 +266,99 @@ func (v *VolcEngine) ChatStreamlyWithSender(modelName string, messages []Message "role": msg.Role, "content": msg.Content, } + if msg.ToolCallID != "" { + apiMessages[i]["tool_call_id"] = msg.ToolCallID + } + if len(msg.ToolCalls) > 0 { + apiMessages[i]["tool_calls"] = msg.ToolCalls + } } // Build request body with streaming enabled reqBody := map[string]interface{}{ - "model": modelName, - "messages": apiMessages, - "stream": true, - "temperature": 1, + "model": modelName, + "messages": apiMessages, + "stream": true, } - if modelConfig == nil { - modelConfig = &ChatConfig{} - } + if modelConfig != nil { + if modelConfig.Stream != nil { + reqBody["stream"] = *modelConfig.Stream + } - if modelConfig.Stream != nil { - reqBody["stream"] = *modelConfig.Stream - } + if modelConfig.MaxTokens != nil { + reqBody["max_tokens"] = *modelConfig.MaxTokens + } - if modelConfig.MaxTokens != nil { - reqBody["max_tokens"] = *modelConfig.MaxTokens - } + if modelConfig.Temperature != nil { + reqBody["temperature"] = *modelConfig.Temperature + } - if modelConfig.Temperature != nil { - reqBody["temperature"] = *modelConfig.Temperature - } + if modelConfig.TopP != nil { + reqBody["top_p"] = *modelConfig.TopP + } - if modelConfig.TopP != nil { - reqBody["top_p"] = *modelConfig.TopP - } + if modelConfig.DoSample != nil { + reqBody["do_sample"] = *modelConfig.DoSample + } - if modelConfig.DoSample != nil { - reqBody["do_sample"] = *modelConfig.DoSample - } + if modelConfig.Stop != nil { + reqBody["stop"] = *modelConfig.Stop + } - if modelConfig.Stop != nil { - reqBody["stop"] = *modelConfig.Stop - } - - // TODO VolcEngine has `auto` mode - if modelConfig.Thinking != nil { - if *modelConfig.Thinking { - var thinkingFlag string - effort := "medium" - if modelConfig.Effort != nil { - effort = *modelConfig.Effort - } - switch effort { - case "none", "minimal": - thinkingFlag = "disabled" - reqBody["reasoning_effort"] = "minimal" - break - case "low": - thinkingFlag = "enabled" - reqBody["reasoning_effort"] = "low" - break - case "medium": - thinkingFlag = "enabled" - reqBody["reasoning_effort"] = "medium" - break - case "auto", "default": - thinkingFlag = "enabled" - reqBody["reasoning_effort"] = "medium" - break - case "high": - thinkingFlag = "enabled" - reqBody["reasoning_effort"] = "high" - break - default: - return fmt.Errorf("invalid effort level") - } - reqBody["thinking"] = map[string]interface{}{ - "type": thinkingFlag, - } - } else { - reqBody["thinking"] = map[string]interface{}{ - "type": "disabled", + // TODO VolcEngine has `auto` mode + if modelConfig.Thinking != nil { + if *modelConfig.Thinking { + var thinkingFlag string + effort := "medium" + if modelConfig.Effort != nil { + effort = *modelConfig.Effort + } + switch effort { + case "none", "minimal": + thinkingFlag = "disabled" + reqBody["reasoning_effort"] = "minimal" + break + case "low": + thinkingFlag = "enabled" + reqBody["reasoning_effort"] = "low" + break + case "medium": + thinkingFlag = "enabled" + reqBody["reasoning_effort"] = "medium" + break + case "auto", "default": + thinkingFlag = "enabled" + reqBody["reasoning_effort"] = "medium" + break + case "high": + thinkingFlag = "enabled" + reqBody["reasoning_effort"] = "high" + break + case "xhigh": + thinkingFlag = "enabled" + reqBody["reasoning_effort"] = "xhigh" + break + default: + return fmt.Errorf("invalid effort level") + } + reqBody["thinking"] = map[string]interface{}{ + "type": thinkingFlag, + } + } else { + reqBody["thinking"] = map[string]interface{}{ + "type": "disabled", + } } } + + if modelConfig.Tools != nil { + reqBody["tools"] = modelConfig.Tools + } + + if modelConfig.ToolChoice != nil { + reqBody["tool_choice"] = modelConfig.ToolChoice + } } jsonData, err := json.Marshal(reqBody) @@ -345,6 +388,7 @@ func (v *VolcEngine) ChatStreamlyWithSender(modelName string, messages []Message return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) } + accumulatedToolCalls := make(map[int]map[string]any) if _, err := ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error { common.Info(fmt.Sprintf("%v", event)) @@ -370,6 +414,56 @@ func (v *VolcEngine) ChatStreamlyWithSender(modelName string, messages []Message } } + if tcs, ok := delta["tool_calls"].([]interface{}); ok { + for _, tc := range tcs { + tcMap, ok := tc.(map[string]interface{}) + if !ok { + continue + } + idxF, ok := tcMap["index"].(float64) + if !ok { + continue + } + idx := int(idxF) + existing, hasExisting := accumulatedToolCalls[idx] + if !hasExisting { + accumulatedToolCalls[idx] = cloneMap(tcMap) + continue + } + if id, ok := tcMap["id"].(string); ok && id != "" { + if eid, ok := existing["id"].(string); ok { + existing["id"] = eid + id + } else { + existing["id"] = id + } + } + if typ, ok := tcMap["type"].(string); ok && typ != "" { + existing["type"] = typ + } + if fn, ok := tcMap["function"].(map[string]interface{}); ok { + ef, ok := existing["function"].(map[string]interface{}) + if !ok { + ef = make(map[string]interface{}) + existing["function"] = ef + } + if name, ok := fn["name"].(string); ok && name != "" { + if en, ok := ef["name"].(string); ok { + ef["name"] = en + name + } else { + ef["name"] = name + } + } + if args, ok := fn["arguments"].(string); ok && args != "" { + if ea, ok := ef["arguments"].(string); ok { + ef["arguments"] = ea + args + } else { + ef["arguments"] = args + } + } + } + } + } + reasoningContent, ok := delta["reasoning_content"].(string) if ok && reasoningContent != "" { if err := sender(nil, &reasoningContent); err != nil { @@ -382,6 +476,19 @@ func (v *VolcEngine) ChatStreamlyWithSender(modelName string, messages []Message return fmt.Errorf("failed to scan response body: %w", err) } + if len(accumulatedToolCalls) > 0 && modelConfig != nil { + indices := make([]int, 0, len(accumulatedToolCalls)) + for idx := range accumulatedToolCalls { + indices = append(indices, idx) + } + sort.Ints(indices) + tcs := make([]map[string]interface{}, 0, len(accumulatedToolCalls)) + for _, idx := range indices { + tcs = append(tcs, accumulatedToolCalls[idx]) + } + modelConfig.ToolCallsResult = &tcs + } + // Send [DONE] marker for OpenAI compatibility endOfStream := "[DONE]" if err = sender(&endOfStream, nil); err != nil { diff --git a/internal/entity/models/zhipu-ai.go b/internal/entity/models/zhipu-ai.go index 9801fcae84..fc2d3c8057 100644 --- a/internal/entity/models/zhipu-ai.go +++ b/internal/entity/models/zhipu-ai.go @@ -29,6 +29,7 @@ import ( "path/filepath" "ragflow/internal/common" "ragflow/internal/engine/clickhouse" + "sort" "strings" "time" @@ -64,9 +65,10 @@ type ZhipuChatResponse struct { FinishReason string `json:"finish_reason"` Index int `json:"index"` Message struct { - Content string `json:"content"` - ReasoningContent string `json:"reasoning_content"` - Role string `json:"role"` + Content string `json:"content"` + ReasoningContent string `json:"reasoning_content"` + Role string `json:"role"` + ToolCalls []map[string]any `json:"tool_calls"` } `json:"message"` } `json:"choices"` Created int `json:"created"` @@ -108,14 +110,19 @@ func (z *ZhipuAIModel) ChatWithMessages(modelName string, messages []Message, ap "role": msg.Role, "content": msg.Content, } + if msg.ToolCallID != "" { + apiMessages[i]["tool_call_id"] = msg.ToolCallID + } + if len(msg.ToolCalls) > 0 { + apiMessages[i]["tool_calls"] = msg.ToolCalls + } } // Build request body reqBody := map[string]interface{}{ - "model": modelName, - "messages": apiMessages, - "stream": false, - "temperature": 1, + "model": modelName, + "messages": apiMessages, + "stream": false, } if chatModelConfig != nil { @@ -150,6 +157,13 @@ func (z *ZhipuAIModel) ChatWithMessages(modelName string, messages []Message, ap } } } + + if chatModelConfig.Tools != nil { + reqBody["tools"] = chatModelConfig.Tools + } + if chatModelConfig.ToolChoice != nil { + reqBody["tool_choice"] = chatModelConfig.ToolChoice + } } jsonData, err := json.Marshal(reqBody) @@ -200,28 +214,33 @@ func (z *ZhipuAIModel) ChatWithMessages(modelName string, messages []Message, ap reasonContent = &result.Choices[0].Message.ReasoningContent } + var toolCalls = result.Choices[0].Message.ToolCalls + usage := &TokenUsage{ PromptTokens: result.Usage.PromptTokens, CompletionTokens: result.Usage.CompletionTokens, TotalTokens: result.Usage.TotalTokens, } - modelUsage.RequestID = result.RequestId - modelUsage.InputTokens = result.Usage.PromptTokens - modelUsage.OutputTokens = result.Usage.CompletionTokens - modelUsage.TotalTokens = result.Usage.TotalTokens - modelUsage.ResponseTimeMS = time.Since(modelUsage.StartAt).Milliseconds() + if modelUsage != nil { + modelUsage.RequestID = result.RequestId + modelUsage.InputTokens = result.Usage.PromptTokens + modelUsage.OutputTokens = result.Usage.CompletionTokens + modelUsage.TotalTokens = result.Usage.TotalTokens + modelUsage.ResponseTimeMS = time.Since(modelUsage.StartAt).Milliseconds() - clickhouseDriver := clickhouse.GetDriver() - err = clickhouseDriver.CollectModelUsage(modelUsage) - if err != nil { - return nil, fmt.Errorf("failed to collect model usage: %w", err) + clickhouseDriver := clickhouse.GetDriver() + err = clickhouseDriver.CollectModelUsage(modelUsage) + if err != nil { + return nil, fmt.Errorf("failed to collect model usage: %w", err) + } } chatResponse := &ChatResponse{ Answer: content, ReasonContent: reasonContent, Usage: usage, + ToolCalls: toolCalls, } return chatResponse, nil @@ -251,14 +270,20 @@ func (z *ZhipuAIModel) ChatStreamlyWithSender(modelName string, messages []Messa "role": msg.Role, "content": msg.Content, } + if msg.ToolCallID != "" { + apiMessages[i]["tool_call_id"] = msg.ToolCallID + } + + if len(msg.ToolCalls) > 0 { + apiMessages[i]["tool_calls"] = msg.ToolCalls + } } // Build request body with streaming enabled reqBody := map[string]interface{}{ - "model": modelName, - "messages": apiMessages, - "stream": true, - "temperature": 1, + "model": modelName, + "messages": apiMessages, + "stream": true, } if chatModelConfig != nil { @@ -297,6 +322,14 @@ func (z *ZhipuAIModel) ChatStreamlyWithSender(modelName string, messages []Messa } } } + + if chatModelConfig.Tools != nil { + reqBody["tools"] = chatModelConfig.Tools + } + + if chatModelConfig.ToolChoice != nil { + reqBody["tool_choice"] = chatModelConfig.ToolChoice + } } jsonData, err := json.Marshal(reqBody) @@ -327,9 +360,31 @@ func (z *ZhipuAIModel) ChatStreamlyWithSender(modelName string, messages []Messa } // SSE parsing: read line by line + accumulatedToolCalls := make(map[int]map[string]any) if _, err = ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error { common.Info(fmt.Sprintf("%v", event)) + tokenUsageMap, ok := event["usage"].(map[string]interface{}) + if ok { + tokenUsage := TokenUsage{} + if err = mapstructure.Decode(tokenUsageMap, &tokenUsage); err != nil { + return err + } + if chatModelConfig != nil { + chatModelConfig.UsageResult = &tokenUsage + } + if modelUsage != nil { + modelUsage.InputTokens = tokenUsage.PromptTokens + modelUsage.OutputTokens = tokenUsage.CompletionTokens + modelUsage.TotalTokens = tokenUsage.TotalTokens + modelUsage.ResponseTimeMS = time.Since(modelUsage.StartAt).Milliseconds() + clickhouseDriver := clickhouse.GetDriver() + if err = clickhouseDriver.CollectModelUsage(modelUsage); err != nil { + return err + } + } + } + choices, ok := event["choices"].([]interface{}) if !ok || len(choices) == 0 { return nil @@ -352,6 +407,56 @@ func (z *ZhipuAIModel) ChatStreamlyWithSender(modelName string, messages []Messa } } + if tcs, ok := delta["tool_calls"].([]interface{}); ok { + for _, tc := range tcs { + tcMap, ok := tc.(map[string]interface{}) + if !ok { + continue + } + idxF, ok := tcMap["index"].(float64) + if !ok { + continue + } + idx := int(idxF) + existing, hasExisting := accumulatedToolCalls[idx] + if !hasExisting { + accumulatedToolCalls[idx] = cloneMap(tcMap) + continue + } + if id, ok := tcMap["id"].(string); ok && id != "" { + if eid, ok := existing["id"].(string); ok { + existing["id"] = eid + id + } else { + existing["id"] = id + } + } + if typ, ok := tcMap["type"].(string); ok && typ != "" { + existing["type"] = typ + } + if fn, ok := tcMap["function"].(map[string]interface{}); ok { + ef, ok := existing["function"].(map[string]interface{}) + if !ok { + ef = make(map[string]interface{}) + existing["function"] = ef + } + if name, ok := fn["name"].(string); ok && name != "" { + if en, ok := ef["name"].(string); ok { + ef["name"] = en + name + } else { + ef["name"] = name + } + } + if args, ok := fn["arguments"].(string); ok && args != "" { + if ea, ok := ef["arguments"].(string); ok { + ef["arguments"] = ea + args + } else { + ef["arguments"] = args + } + } + } + } + } + content, ok := delta["content"].(string) if ok && content != "" { if err = sender(&content, nil); err != nil { @@ -359,31 +464,27 @@ func (z *ZhipuAIModel) ChatStreamlyWithSender(modelName string, messages []Messa } } - tokenUsageMap, ok := event["usage"].(map[string]interface{}) - if ok { - tokenUsage := TokenUsage{} - err = mapstructure.Decode(tokenUsageMap, &tokenUsage) - modelUsage.InputTokens = tokenUsage.PromptTokens - modelUsage.OutputTokens = tokenUsage.CompletionTokens - modelUsage.TotalTokens = tokenUsage.TotalTokens - modelUsage.ResponseTimeMS = time.Since(modelUsage.StartAt).Milliseconds() - clickhouseDriver := clickhouse.GetDriver() - err = clickhouseDriver.CollectModelUsage(modelUsage) - return nil - } - return nil }); err != nil { return fmt.Errorf("failed to scan response body: %w", err) } - // Send [DONE] marker for OpenAI compatibility - endOfStream := "[DONE]" - if err = sender(&endOfStream, nil); err != nil { - return err + if len(accumulatedToolCalls) > 0 && chatModelConfig != nil { + indices := make([]int, 0, len(accumulatedToolCalls)) + for idx := range accumulatedToolCalls { + indices = append(indices, idx) + } + sort.Ints(indices) + toolCalls := make([]map[string]interface{}, 0, len(accumulatedToolCalls)) + for _, idx := range indices { + toolCalls = append(toolCalls, accumulatedToolCalls[idx]) + } + chatModelConfig.ToolCallsResult = &toolCalls } - return nil + // Send [DONE] marker for OpenAI compatibility + endOfStream := "[DONE]" + return sender(&endOfStream, nil) } type zhipuEmbeddingResponse struct { diff --git a/internal/entity/models/zhipu_ai_test.go b/internal/entity/models/zhipu_ai_test.go index 73a38c246a..586f3dc116 100644 --- a/internal/entity/models/zhipu_ai_test.go +++ b/internal/entity/models/zhipu_ai_test.go @@ -342,3 +342,99 @@ func TestZhipuAIAudioSpeechHTTPError(t *testing.T) { t.Fatalf("error=%v", err) } } + +func TestZhipuAIChatStreamlyWithSenderCollectsToolCalls(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + t.Errorf("method=%s", r.Method) + return + } + if r.URL.Path != "/chat/completions" { + t.Errorf("path=%s", r.URL.Path) + return + } + + var body map[string]interface{} + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Errorf("Decode: %v", err) + return + } + if body["stream"] != true { + t.Errorf("stream=%v", body["stream"]) + } + if body["tools"] == nil { + t.Error("tools missing from request") + } + if body["tool_choice"] != "auto" { + t.Errorf("tool_choice=%v", body["tool_choice"]) + } + + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call-1\",\"type\":\"function\",\"function\":{\"name\":\"search_knowledge\",\"arguments\":\"{\\\"query\\\":\\\"mari\"}}]}}]}\n\n") + _, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"gold\\\"}\"}}]}}]}\n\n") + _, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"tool_calls\"}]}\n\n") + _, _ = io.WriteString(w, "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":3,\"completion_tokens\":5,\"total_tokens\":8}}\n\n") + _, _ = io.WriteString(w, "data: [DONE]\n\n") + })) + defer srv.Close() + + apiKey := "test-key" + cfg := &ChatConfig{ + Tools: []map[string]interface{}{ + { + "type": "function", + "function": map[string]interface{}{ + "name": "search_knowledge", + "description": "Search the knowledge base.", + "parameters": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{"query": map[string]interface{}{"type": "string"}}, + "required": []string{"query"}, + }, + }, + }, + }, + } + toolChoice := "auto" + cfg.ToolChoice = &toolChoice + + err := NewZhipuAIModel( + map[string]string{"default": srv.URL}, + URLSuffix{Chat: "chat/completions"}, + ).ChatStreamlyWithSender( + "glm-4", + []Message{{Role: "user", Content: "what is marigold"}}, + &APIConfig{ApiKey: &apiKey}, + cfg, + nil, + func(content *string, reasoning *string) error { + if content != nil && *content != "[DONE]" { + t.Errorf("content=%q", *content) + } + if reasoning != nil { + t.Errorf("reasoning=%q", *reasoning) + } + return nil + }, + ) + if err != nil { + t.Fatalf("ChatStreamlyWithSender: %v", err) + } + if cfg.ToolCallsResult == nil || len(*cfg.ToolCallsResult) != 1 { + t.Fatalf("ToolCallsResult=%#v, want one tool call", cfg.ToolCallsResult) + } + call := (*cfg.ToolCallsResult)[0] + if call["id"] != "call-1" || call["type"] != "function" { + t.Fatalf("tool call metadata=%#v", call) + } + fn, ok := call["function"].(map[string]interface{}) + if !ok { + t.Fatalf("function=%#v", call["function"]) + } + if fn["name"] != "search_knowledge" || fn["arguments"] != `{"query":"marigold"}` { + t.Fatalf("function=%#v", fn) + } + if cfg.UsageResult == nil || cfg.UsageResult.TotalTokens != 8 { + t.Fatalf("UsageResult=%#v, want total tokens 8", cfg.UsageResult) + } +} diff --git a/internal/handler/models.go b/internal/handler/models.go index 69a2dbb227..e47877d73b 100644 --- a/internal/handler/models.go +++ b/internal/handler/models.go @@ -24,12 +24,12 @@ import ( "github.com/gin-gonic/gin" ) -// ProviderHandler provider handler +// ModelHandler ProviderHandler provider handler type ModelHandler struct { modelProviderService *service.ModelProviderService } -// NewProviderHandler create provider handler +// NewModelHandler NewProviderHandler create provider handler func NewModelHandler(modelProviderService *service.ModelProviderService) *ModelHandler { return &ModelHandler{ modelProviderService: modelProviderService,