zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-08-04 23:00:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(go-models): migrate batch 5 model drivers to unified handlers (#17700)

Browse Source 
## Summary

Relate to #17284. Completes the batch 5 migration of 7 OpenAI-compatible
drivers (`vllm`, `volcengine`, `xai`, `xiaomi`, `xinference`, `xunfei`,
`zhipu-ai`) onto the unified request/response helpers
(`doRequest`/`doStreamRequest` +
`HandleNonStreamingResponse`/`HandleStreamingResponse` +
`ParserConfig`), established by `deepseek` in #17634.

This branch is rebased on the current `pr/migrate-models-batch5` and
fixes the issues in the previous state of the PR.

Co-authored-by: Haruko386 <tryeverypossible@163.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
jay77721
2026-08-03 20:19:22 +08:00
committed by GitHub
parent 1d141aff18
commit d357eea8ef
17 changed files with 478 additions and 1486 deletions
Download Patch File Download Diff File Expand all files Collapse all files

76
internal/entity/models/base_model.go
View File

@@ -30,8 +30,6 @@ import (
"sort"
"strings"
"time"
"github.com/mitchellh/mapstructure"
)
type BaseModel struct {
@@ -39,6 +37,11 @@ type BaseModel struct {
URLSuffix URLSuffix
httpClient *http.Client
AllowEmptyAPIKey bool
// authHeader, when non-nil, supplies the (name, value) pair used for
// authentication instead of the default "Authorization: Bearer <key>".
// Drivers with non-standard auth (e.g. Xiaomi's api-key header, Xunfei's
// spark_api_password bundle) set it in their constructor.
authHeader func(*APIConfig) (string, string)
}
// chatResponseParts is the provider-normalized result of a non-streaming chat
@@ -105,20 +108,6 @@ func collectModelUsage(modelUsage *common.ModelUsage, usage *TokenUsage) error {
return clickhouse.GetDriver().CollectModelUsage(modelUsage)
}
// decodeOpenAICompatibleStreamUsage extracts aggregate token usage from one
// OpenAI-compatible streaming event. A missing usage field is not an error.
func decodeOpenAICompatibleStreamUsage(event map[string]any) (*TokenUsage, bool, error) {
rawUsage, ok := event["usage"].(map[string]any)
if !ok {
return nil, false, nil
}
usage := &TokenUsage{}
if err := mapstructure.Decode(rawUsage, usage); err != nil {
return nil, false, err
}
return usage, true, nil
}
// applyStreamUsage exposes streamed token usage to the caller and records it
// for model-usage analytics when a usage event is received. Analytics failures
// are logged but do not interrupt the stream.
@@ -149,7 +138,21 @@ func (b *BaseModel) APIConfigCheck(apiConfig *APIConfig) error {
return nil
}
func newJSONPostRequest(ctx context.Context, url string, apiConfig *APIConfig, reqBody map[string]any) (*http.Request, error) {
// applyAuth sets the authentication header on req. Drivers with a custom
// authHeader hook (e.g. Xiaomi's api-key header, Xunfei's spark_api_password
// bundle) use it; the default is "Authorization: Bearer <key>".
func (b *BaseModel) applyAuth(req *http.Request, apiConfig *APIConfig) {
if b.authHeader != nil {
name, value := b.authHeader(apiConfig)
req.Header.Set(name, value)
return
}
if auth := BearerAuth(apiConfig); auth != "" {
req.Header.Set("Authorization", auth)
}
}
func (b *BaseModel) newJSONPostRequest(ctx context.Context, url string, apiConfig *APIConfig, reqBody map[string]any) (*http.Request, error) {
jsonData, err := json.Marshal(reqBody)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
@@ -161,9 +164,7 @@ func newJSONPostRequest(ctx context.Context, url string, apiConfig *APIConfig, r
}
req.Header.Set("Content-Type", "application/json")
if auth := BearerAuth(apiConfig); auth != "" {
req.Header.Set("Authorization", auth)
}
b.applyAuth(req, apiConfig)
return req, nil
}
@@ -173,7 +174,7 @@ func (b *BaseModel) doRequest(ctx context.Context, url string, apiConfig *APICon
ctx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
req, err := newJSONPostRequest(ctx, url, apiConfig, reqBody)
req, err := b.newJSONPostRequest(ctx, url, apiConfig, reqBody)
if err != nil {
return nil, err
}
@@ -196,6 +197,37 @@ func (b *BaseModel) doRequest(ctx context.Context, url string, apiConfig *APICon
return body, nil
}
// doGetRequest sends a GET request and returns the response body.
func (b *BaseModel) doGetRequest(ctx context.Context, url string, apiConfig *APIConfig, timeout time.Duration) ([]byte, error) {
ctx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
b.applyAuth(req, apiConfig)
resp, err := b.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
return body, nil
}
// mustMarshal marshals v to JSON, panicking on error.
func mustMarshal(v any) []byte {
b, err := json.Marshal(v)
@@ -210,7 +242,7 @@ func (b *BaseModel) doStreamRequest(ctx context.Context, url string, apiConfig *
ctx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
req, err := newJSONPostRequest(ctx, url, apiConfig, reqBody)
req, err := b.newJSONPostRequest(ctx, url, apiConfig, reqBody)
if err != nil {
return err
}

38
internal/entity/models/common.go
View File

@@ -18,44 +18,6 @@ package models
import "strings"
func GetThinkingAndAnswer(modelType *string, content *string) (*string, *string) {
if content == nil {
return nil, nil
}
switch NormalizeModelFamily(modelType) {
case "qwen3":
return extractThinkContent(content)
}
return nil, content
}
// NormalizeModelFamily normalizes provider-prefixed model class/name strings for shared response parsing.
func NormalizeModelFamily(modelType *string) string {
if modelType == nil {
return ""
}
family := strings.ToLower(strings.TrimSpace(*modelType))
if family == "" {
return ""
}
if slash := strings.LastIndex(family, "/"); slash >= 0 && slash < len(family)-1 {
family = family[slash+1:]
}
if family == "qwen3" || strings.HasPrefix(family, "qwen3-") || strings.HasPrefix(family, "qwen3.") {
return "qwen3"
}
if dash := strings.Index(family, "-"); dash >= 0 {
family = family[:dash]
}
return family
}
func extractThinkContent(content *string) (*string, *string) {
if content == nil {
return nil, nil

95
internal/entity/models/common_test.go
View File

@@ -10,101 +10,6 @@ import (
"time"
)
func TestNormalizeModelFamily(t *testing.T) {
tests := []struct {
name string
input *string
want string
}{
{name: "nil", input: nil, want: ""},
{name: "empty", input: modelFamilyTestString(""), want: ""},
{name: "qwen3", input: modelFamilyTestString("qwen3"), want: "qwen3"},
{name: "qwen3 hyphen variant", input: modelFamilyTestString("qwen3-8b"), want: "qwen3"},
{name: "qwen3 dot variant", input: modelFamilyTestString("qwen3.5-4b"), want: "qwen3"},
{name: "provider-prefixed qwen3", input: modelFamilyTestString("qwen/qwen3-8b"), want: "qwen3"},
{name: "case-varied qwen3", input: modelFamilyTestString("Qwen/Qwen3.5-4B"), want: "qwen3"},
{name: "provider-prefixed non-qwen", input: modelFamilyTestString("deepseek/deepseek-r1"), want: "deepseek"},
{name: "qwen plus not qwen3", input: modelFamilyTestString("qwen-plus"), want: "qwen"},
{name: "provider-prefixed qwen2.5 not qwen3", input: modelFamilyTestString("qwen/qwen2.5-32b-instruct"), want: "qwen2.5"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := NormalizeModelFamily(tt.input); got != tt.want {
t.Fatalf("NormalizeModelFamily()=%q, want %q", got, tt.want)
}
})
}
}
func TestGetThinkingAndAnswerExtractsQwenThinking(t *testing.T) {
content := "<think>\nplan</think>\nanswer"
for _, modelType := range []string{
"qwen3",
"qwen3-8b",
"qwen/qwen3",
"qwen/qwen3-8b",
"Qwen/Qwen3.5-4B",
} {
t.Run(modelType, func(t *testing.T) {
thinking, answer := GetThinkingAndAnswer(&modelType, &content)
if thinking == nil || *thinking != "plan" {
t.Fatalf("thinking=%v, want plan", thinking)
}
if answer == nil || *answer != "answer" {
t.Fatalf("answer=%v, want answer", answer)
}
})
}
}
func TestGetThinkingAndAnswerLeavesUnknownModelFamiliesUnchanged(t *testing.T) {
content := "<think>\nplan</think>\nanswer"
for _, modelType := range []string{
"deepseek",
"deepseek/deepseek-r1",
"qwen-plus",
"qwen/qwen2.5-32b-instruct",
} {
t.Run(modelType, func(t *testing.T) {
thinking, answer := GetThinkingAndAnswer(&modelType, &content)
if thinking != nil {
t.Fatalf("thinking=%v, want nil", thinking)
}
if answer != &content {
t.Fatalf("answer pointer changed")
}
})
}
}
func TestGetThinkingAndAnswerHandlesNilInputs(t *testing.T) {
thinking, answer := GetThinkingAndAnswer(nil, nil)
if thinking != nil || answer != nil {
t.Fatalf("GetThinkingAndAnswer(nil, nil)=(%v, %v), want nils", thinking, answer)
}
content := "<think>\nplan</think>\nanswer"
thinking, answer = GetThinkingAndAnswer(nil, &content)
if thinking != nil {
t.Fatalf("thinking=%v, want nil", thinking)
}
if answer != &content {
t.Fatalf("answer pointer changed")
}
modelType := "qwen3"
thinking, answer = GetThinkingAndAnswer(&modelType, nil)
if thinking != nil {
t.Fatalf("thinking=%v, want nil", thinking)
}
if answer != nil {
t.Fatalf("answer=%v, want nil", answer)
}
}
func TestBaseModelDoRequestAuthorizationHeader(t *testing.T) {
tests := []struct {
name string

26
internal/entity/models/response_handler.go
View File

@@ -27,12 +27,7 @@ import (
// HandleNonStreamingResponse processes a complete non-streaming chat
// response using the ParserConfig's ResponseParser to extract usage.
func HandleNonStreamingResponse(
body []byte,
modelUsage *common.ModelUsage,
chatConfig *ChatConfig,
cfg *ParserConfig,
) (*ChatResponse, error) {
func HandleNonStreamingResponse(body []byte, modelUsage *common.ModelUsage, chatConfig *ChatConfig, cfg *ParserConfig) (*ChatResponse, error) {
var result map[string]any
if err := json.Unmarshal(body, &result); err != nil {
return nil, fmt.Errorf("failed to parse response: %w", err)
@@ -75,13 +70,7 @@ func HandleNonStreamingResponse(
// HandleStreamingResponse processes a streaming chat response using the
// ParserConfig's StreamParser to extract usage from each event.
func HandleStreamingResponse(
body io.Reader,
modelUsage *common.ModelUsage,
chatConfig *ChatConfig,
cfg *ParserConfig,
sender func(*string, *string) error,
) error {
func HandleStreamingResponse(body io.Reader, modelUsage *common.ModelUsage, chatConfig *ChatConfig, cfg *ParserConfig, sender func(*string, *string) error) error {
if sender == nil {
return fmt.Errorf("sender is required")
}
@@ -103,6 +92,14 @@ func HandleStreamingResponse(
return fmt.Errorf("upstream stream error: %v", apiErr)
}
// Some providers emit a terminal event that carries only a root-level
// finish_reason with no choices. Check it before the choices guard so
// the stream terminates cleanly instead of being rejected as "ended
// before [DONE] or finish_reason".
if finishReason, ok := event["finish_reason"].(string); ok && finishReason != "" {
sawTerminal = true
}
choices, ok := event["choices"].([]any)
if !ok || len(choices) == 0 {
return nil
@@ -142,9 +139,6 @@ func HandleStreamingResponse(
if finishReason, ok := firstChoice["finish_reason"].(string); ok && finishReason != "" {
sawTerminal = true
}
if finishReason, ok := event["finish_reason"].(string); ok && finishReason != "" {
sawTerminal = true
}
return nil
})

81
internal/entity/models/response_handler_test.go Normal file
View File

@@ -0,0 +1,81 @@
//
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
package models
import (
"bytes"
"io"
"strings"
"testing"
)
// TestHandleStreamingResponseBareRootFinishReason verifies the shared handler
// terminates the stream when an event carries only a root-level finish_reason
// with no choices and no trailing [DONE]. Before the fix this event was
// skipped by the choices guard and the handler returned "stream ended before
// [DONE] or finish_reason".
func TestHandleStreamingResponseBareRootFinishReason(t *testing.T) {
sse := "data: {\"finish_reason\":\"stop\"}\n"
var contentChunks, reasonChunks []string
err := HandleStreamingResponse(
io.NopCloser(bytes.NewBufferString(sse)),
nil,
nil,
OpenAIParserConfig,
func(content *string, reason *string) error {
if content != nil && *content != "" {
contentChunks = append(contentChunks, *content)
}
if reason != nil && *reason != "" {
reasonChunks = append(reasonChunks, *reason)
}
return nil
},
)
if err != nil {
t.Fatalf("HandleStreamingResponse: %v", err)
}
// The terminal [DONE] marker is the only content the handler emits for a
// finish_reason-only event: no tokens arrived and no reasoning was sent.
if got := strings.Join(contentChunks, ""); got != "[DONE]" {
t.Errorf("content=%q, want only the [DONE] marker", got)
}
if len(reasonChunks) != 0 {
t.Errorf("reasoning=%q, want none for a finish_reason-only event", strings.Join(reasonChunks, ""))
}
}
// TestHandleStreamingResponseTruncatedStream verifies a stream that ends
// without [DONE], a root-level finish_reason, or a per-choice finish_reason
// is rejected as truncated. This is the counterpart of the bare
// finish_reason case: only a genuine truncation must fail.
func TestHandleStreamingResponseTruncatedStream(t *testing.T) {
sse := "data: {\"choices\":[{\"delta\":{\"content\":\"partial\"}}]}\n"
err := HandleStreamingResponse(
io.NopCloser(bytes.NewBufferString(sse)),
nil,
nil,
OpenAIParserConfig,
func(*string, *string) error { return nil },
)
if err == nil || !strings.Contains(err.Error(), "stream ended before [DONE] or finish_reason") {
t.Fatalf("expected truncation error, got %v", err)
}
}

159
internal/entity/models/vllm.go
View File

@@ -91,84 +91,12 @@ func (v *VllmModel) ChatWithMessages(ctx context.Context, modelName string, mess
}
}
jsonData, err := json.Marshal(reqBody)
body, err := v.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
return nil, err
}
ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
if auth := BearerAuth(apiConfig); auth != "" {
req.Header.Set("Authorization", auth)
}
resp, err := v.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
// Parse response
var result map[string]interface{}
if err = json.Unmarshal(body, &result); err != nil {
return nil, fmt.Errorf("failed to parse response: %w", err)
}
choices, ok := result["choices"].([]interface{})
if !ok || len(choices) == 0 {
return nil, fmt.Errorf("no choices in response")
}
firstChoice, ok := choices[0].(map[string]interface{})
if !ok {
return nil, fmt.Errorf("invalid choice format")
}
messageMap, ok := firstChoice["message"].(map[string]interface{})
if !ok {
return nil, fmt.Errorf("invalid message format")
}
content, ok := messageMap["content"].(string)
toolCalls := extractToolCalls(messageMap)
if !ok && len(toolCalls) == 0 {
return nil, fmt.Errorf("invalid content format")
}
var reasonContent string
if chatModelConfig != nil && chatModelConfig.Thinking != nil && *chatModelConfig.Thinking {
reasonContent, ok = messageMap["reasoning_content"].(string)
if !ok {
return nil, fmt.Errorf("invalid content format")
}
if reasonContent != "" && reasonContent[0] == '\n' {
reasonContent = reasonContent[1:]
}
}
chatResponse := &ChatResponse{
Answer: &content,
ReasonContent: &reasonContent,
ToolCalls: toolCalls,
}
return chatResponse, nil
return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig)
}
// ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel)
@@ -206,84 +134,11 @@ func (v *VllmModel) ChatStreamlyWithSender(ctx context.Context, modelName string
}
}
jsonData, err := json.Marshal(reqBody)
if err != nil {
return fmt.Errorf("failed to marshal request: %w", err)
}
reqBody["stream_options"] = map[string]interface{}{"include_usage": true}
ctx, cancel := context.WithTimeout(ctx, streamCallTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
if auth := BearerAuth(apiConfig); auth != "" {
req.Header.Set("Authorization", auth)
}
resp, err := v.baseModel.httpClient.Do(req)
if err != nil {
return fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
// 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))
choices, ok := event["choices"].([]interface{})
if !ok || len(choices) == 0 {
return nil
}
firstChoice, ok := choices[0].(map[string]interface{})
if !ok {
return nil
}
delta, ok := firstChoice["delta"].(map[string]interface{})
if !ok {
return nil
}
accumulateToolCallDeltas(delta, accumulatedToolCalls)
reasoningContent, ok := delta["reasoning_content"].(string)
if ok && reasoningContent != "" {
if err := sender(nil, &reasoningContent); err != nil {
return err
}
}
content, ok := delta["content"].(string)
if ok && content != "" {
if err := sender(&content, nil); err != nil {
return err
}
}
return nil
}); err != nil {
return fmt.Errorf("failed to scan response body: %w", err)
}
setSortedToolCallsResult(modelConfig, accumulatedToolCalls)
// Send [DONE] marker for OpenAI compatibility
endOfStream := "[DONE]"
if err = sender(&endOfStream, nil); err != nil {
return err
}
return nil
return v.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error {
return HandleStreamingResponse(body, modelUsage, modelConfig, OpenAIParserConfig, sender)
})
}
// Encode encodes a list of texts into embeddings

271
internal/entity/models/volcengine.go
View File

@@ -24,7 +24,6 @@ import (
"io"
"net/http"
"ragflow/internal/common"
"sort"
"strings"
)
@@ -52,37 +51,6 @@ func (v *VolcEngine) Name() string {
return "volcengine"
}
// VolcEngineChatResponse mirrors Ark's OpenAI-compatible chat response.
type VolcEngineChatResponse struct {
ID string `json:"id"`
Object string `json:"object"`
Created int64 `json:"created"`
Model string `json:"model"`
Choices []struct {
FinishReason string `json:"finish_reason"`
Index int `json:"index"`
Logprobs any `json:"logprobs"`
Message struct {
Content string `json:"content"`
ReasoningContent string `json:"reasoning_content"`
Role string `json:"role"`
ToolCalls []map[string]any `json:"tool_calls"`
} `json:"message"`
} `json:"choices"`
SystemFingerprint string `json:"system_fingerprint"`
Usage struct {
CompletionTokens int `json:"completion_tokens"`
PromptTokens int `json:"prompt_tokens"`
TotalTokens int `json:"total_tokens"`
CompletionTokensDetails struct {
ReasoningTokens int `json:"reasoning_tokens"`
} `json:"completion_tokens_details"`
PromptTokensDetails struct {
CachedTokens int `json:"cached_tokens"`
} `json:"prompt_tokens_details"`
} `json:"usage"`
}
// getAPIKey extracts the actual API key from VolcEngine's stored format.
// VolcEngine stores the api_key as a JSON string like:
//
@@ -103,6 +71,20 @@ func (v *VolcEngine) getAPIKey(apiConfig *APIConfig) string {
return key
}
// effectiveAPIConfig returns an APIConfig with the extracted API key
// suitable for use with doRequest/doStreamRequest.
func (v *VolcEngine) effectiveAPIConfig(apiConfig *APIConfig) *APIConfig {
if apiConfig == nil {
return nil
}
key := v.getAPIKey(apiConfig)
return &APIConfig{
ApiKey: &key,
Region: apiConfig.Region,
BaseURL: apiConfig.BaseURL,
}
}
// ChatWithMessages sends multiple messages with roles and returns response
func (v *VolcEngine) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) {
if err := v.baseModel.APIConfigCheck(apiConfig); err != nil {
@@ -165,75 +147,13 @@ func (v *VolcEngine) ChatWithMessages(ctx context.Context, modelName string, mes
}
jsonData, err := json.Marshal(reqBody)
effectiveConfig := v.effectiveAPIConfig(apiConfig)
body, err := v.baseModel.doRequest(ctx, url, effectiveConfig, reqBody, nonStreamCallTimeout)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
return nil, err
}
ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", v.getAPIKey(apiConfig)))
resp, err := v.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, chatConfig *ChatConfig) (chatResponseParts, error) {
var result VolcEngineChatResponse
if err := json.Unmarshal(body, &result); err != nil {
return chatResponseParts{}, fmt.Errorf("failed to unmarshal response: %w", err)
}
if len(result.Choices) == 0 {
return chatResponseParts{}, fmt.Errorf("no choices in response")
}
choice := &result.Choices[0]
content := choice.Message.Content
if content == "" && len(choice.Message.ToolCalls) == 0 {
return chatResponseParts{}, fmt.Errorf("invalid content format")
}
reasonContent := choice.Message.ReasoningContent
if chatConfig != nil && chatConfig.Thinking != nil && *chatConfig.Thinking && reasonContent == "" {
reasoning, answer := GetThinkingAndAnswer(chatConfig.ModelClass, &content)
if reasoning != nil {
reasonContent = *reasoning
content = *answer
}
}
if reasonContent != "" && reasonContent[0] == '\n' {
reasonContent = reasonContent[1:]
}
return chatResponseParts{
RequestID: result.ID,
Content: &content,
ReasonContent: &reasonContent,
ToolCalls: choice.Message.ToolCalls,
Usage: &TokenUsage{
PromptTokens: result.Usage.PromptTokens,
CompletionTokens: result.Usage.CompletionTokens,
TotalTokens: result.Usage.TotalTokens,
},
}, nil
})
return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig)
}
// ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel)
@@ -250,7 +170,7 @@ func (v *VolcEngine) ChatStreamlyWithSender(ctx context.Context, modelName strin
if err != nil {
return err
}
url := fmt.Sprintf("%s/chat/completions", resolvedBaseURL)
url := fmt.Sprintf("%s/%s", resolvedBaseURL, v.baseModel.URLSuffix.Chat)
// Build request body with streaming enabled
reqBody := buildRequestBody(modelConfig, modelName, messages, true)
@@ -309,157 +229,10 @@ func (v *VolcEngine) ChatStreamlyWithSender(ctx context.Context, modelName strin
}
jsonData, err := json.Marshal(reqBody)
if err != nil {
return fmt.Errorf("failed to marshal request: %w", err)
}
ctx, cancel := context.WithTimeout(ctx, streamCallTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", v.getAPIKey(apiConfig)))
resp, err := v.baseModel.httpClient.Do(req)
if err != nil {
return fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
accumulatedToolCalls := make(map[int]map[string]any)
sawTerminal := false
done, err := ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error {
common.Info(fmt.Sprintf("%v", event))
tokenUsage, found, usageErr := decodeOpenAICompatibleStreamUsage(event)
if usageErr != nil {
return usageErr
}
if found {
applyStreamUsage(modelConfig, modelUsage, tokenUsage)
}
choices, ok := event["choices"].([]interface{})
if !ok || len(choices) == 0 {
return nil
}
firstChoice, ok := choices[0].(map[string]interface{})
if !ok {
return nil
}
if finishReason, ok := firstChoice["finish_reason"].(string); ok && finishReason != "" {
sawTerminal = true
}
delta, ok := firstChoice["delta"].(map[string]interface{})
if !ok {
return nil
}
content, ok := delta["content"].(string)
if ok && content != "" {
if err := sender(&content, nil); err != nil {
return err
}
}
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 {
return err
}
}
return nil
effectiveConfig := v.effectiveAPIConfig(apiConfig)
return v.baseModel.doStreamRequest(ctx, url, effectiveConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error {
return HandleStreamingResponse(body, modelUsage, modelConfig, OpenAIParserConfig, sender)
})
if err != nil {
return fmt.Errorf("failed to scan response body: %w", err)
}
if !done && !sawTerminal {
return fmt.Errorf("volcengine: stream ended before [DONE] or finish_reason")
}
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 {
return err
}
return nil
}
type volcengineEmbeddingResponse struct {

11
internal/entity/models/volcengine_test.go
View File

@@ -120,6 +120,11 @@ func TestVolcEngineChatStreamSupportsMaxEffortAndUsage(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
srv := newVolcEngineServer(t, func(t *testing.T, r *http.Request, w http.ResponseWriter) {
// The streaming endpoint must honor URLSuffix.Chat, not a hardcoded
// "chat/completions" path.
if r.URL.Path != "/v1/chat/completions" {
t.Errorf("path=%s, want /v1/chat/completions", r.URL.Path)
}
var body map[string]interface{}
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
t.Fatalf("decode request: %v", err)
@@ -142,7 +147,11 @@ func TestVolcEngineChatStreamSupportsMaxEffortAndUsage(t *testing.T) {
thinking := true
effort := "max"
config := &ChatConfig{Thinking: &thinking, Effort: &effort}
if err := newVolcEngineForTest(srv.URL).ChatStreamlyWithSender(
driver := NewVolcEngine(
map[string]string{"default": srv.URL},
URLSuffix{Chat: "v1/chat/completions", Models: "models"},
)
if err := driver.ChatStreamlyWithSender(
ctx,
"doubao-seed-2-0-pro-260215",
[]Message{{Role: "user", Content: "hello"}},

185
internal/entity/models/xai.go
View File

@@ -80,81 +80,12 @@ func (x *XAIModel) ChatWithMessages(ctx context.Context, modelName string, messa
url := fmt.Sprintf("%s/%s", baseURL, x.baseModel.URLSuffix.Chat)
reqBody := buildRequestBody(chatModelConfig, modelName, messages, false)
jsonData, err := json.Marshal(reqBody)
body, err := x.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
return nil, err
}
ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := x.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
// Parse response
var result map[string]interface{}
if err = json.Unmarshal(body, &result); err != nil {
return nil, fmt.Errorf("failed to parse response: %w", err)
}
choices, ok := result["choices"].([]interface{})
if !ok || len(choices) == 0 {
return nil, fmt.Errorf("no choices in response")
}
firstChoice, ok := choices[0].(map[string]interface{})
if !ok {
return nil, fmt.Errorf("invalid choice format")
}
messageMap, ok := firstChoice["message"].(map[string]interface{})
if !ok {
return nil, fmt.Errorf("invalid message format")
}
content, ok := messageMap["content"].(string)
toolCalls := extractToolCalls(messageMap)
if !ok && len(toolCalls) == 0 {
return nil, fmt.Errorf("invalid content format")
}
// xAI reasoning models (grok-3-mini and similar) return reasoning text in
// the reasoning_content field. Pass it through when present.
var reasonContent string
if rc, ok := messageMap["reasoning_content"].(string); ok {
reasonContent = rc
if reasonContent != "" && reasonContent[0] == '\n' {
reasonContent = reasonContent[1:]
}
}
chatResponse := &ChatResponse{
Answer: &content,
ReasonContent: &reasonContent,
ToolCalls: toolCalls,
}
return chatResponse, nil
return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig)
}
// ChatStreamlyWithSender sends messages and streams the response
@@ -174,89 +105,13 @@ func (x *XAIModel) ChatStreamlyWithSender(ctx context.Context, modelName string,
baseURL = strings.TrimSuffix(baseURL, "/")
url := fmt.Sprintf("%s/%s", baseURL, x.baseModel.URLSuffix.Chat)
reqBody := buildRequestBody(chatModelConfig, modelName, messages, true)
jsonData, err := json.Marshal(reqBody)
if err != nil {
return fmt.Errorf("failed to marshal request: %w", err)
reqBody["stream_options"] = map[string]interface{}{
"include_usage": true,
}
ctx, cancel := context.WithTimeout(ctx, streamCallTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := x.baseModel.httpClient.Do(req)
if err != nil {
return fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
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 {
return nil
}
firstChoice, ok := choices[0].(map[string]interface{})
if !ok {
return nil
}
delta, ok := firstChoice["delta"].(map[string]interface{})
if !ok {
return nil
}
accumulateToolCallDeltas(delta, accumulatedToolCalls)
reasoningContent, ok := delta["reasoning_content"].(string)
if ok && reasoningContent != "" {
if err := sender(nil, &reasoningContent); err != nil {
return err
}
}
content, ok := delta["content"].(string)
if ok && content != "" {
if err := sender(&content, nil); err != nil {
return err
}
}
finishReason, ok := firstChoice["finish_reason"].(string)
if ok && finishReason != "" {
sawTerminal = true
}
return nil
return x.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error {
return HandleStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig, sender)
})
if err != nil {
return fmt.Errorf("failed to scan response body: %w", err)
}
setSortedToolCallsResult(chatModelConfig, accumulatedToolCalls)
if !done && !sawTerminal {
return fmt.Errorf("xai: stream ended before [DONE] or finish_reason")
}
// Send the [DONE] marker for OpenAI compatibility
endOfStream := "[DONE]"
if err := sender(&endOfStream, nil); err != nil {
return err
}
return nil
}
// Embed embeds a list of texts into embeddings. xAI does not expose a
@@ -282,33 +137,11 @@ func (x *XAIModel) ListModels(ctx context.Context, apiConfig *APIConfig) ([]List
}
url := fmt.Sprintf("%s/%s", strings.TrimSuffix(baseURL, "/"), modelsSuffix)
ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
body, err := x.baseModel.doGetRequest(ctx, url, apiConfig, nonStreamCallTimeout)
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
return nil, err
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := x.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
// Parse response
// Parse response
var modelList ModelList
if err = json.Unmarshal(body, &modelList); err != nil {

107
internal/entity/models/xai_test.go
View File

@@ -2,6 +2,7 @@ package models
import (
"encoding/json"
"io"
"net/http"
"net/http/httptest"
"os"
@@ -137,3 +138,109 @@ func TestXAIListModelsRequiresModelsSuffix(t *testing.T) {
t.Fatalf("expected missing models suffix error, got %v", err)
}
}
func TestXAIChatHappyPath(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
t.Errorf("method=%s, want POST", r.Method)
}
if got := r.Header.Get("Authorization"); got != "Bearer test-key" {
t.Errorf("Authorization=%q, want Bearer test-key", got)
}
var body map[string]interface{}
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
t.Errorf("decode request: %v", err)
return
}
if body["stream"] != false {
t.Errorf("stream=%v, want false", body["stream"])
}
_ = json.NewEncoder(w).Encode(map[string]interface{}{
"choices": []map[string]interface{}{{
"message": map[string]interface{}{
"content": "pong",
"reasoning_content": "\nthought",
},
}},
})
}))
defer srv.Close()
apiKey := "test-key"
resp, err := newXAIForTest(srv.URL).ChatWithMessages(
ctx,
"grok-3-mini",
[]Message{{Role: "user", Content: "ping"}},
&APIConfig{ApiKey: &apiKey},
nil,
nil,
)
if err != nil {
t.Fatalf("ChatWithMessages: %v", err)
}
if resp.Answer == nil || *resp.Answer != "pong" {
t.Errorf("Answer=%v, want pong", resp.Answer)
}
if resp.ReasonContent == nil || *resp.ReasonContent != "thought" {
t.Errorf("ReasonContent=%v, want thought", resp.ReasonContent)
}
}
func TestXAIStreamHappyPath(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if got := r.Header.Get("Accept"); got != "text/event-stream" {
t.Errorf("Accept=%q, want text/event-stream", got)
}
var body map[string]interface{}
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
t.Errorf("decode request: %v", err)
return
}
streamOptions, ok := body["stream_options"].(map[string]interface{})
if !ok || streamOptions["include_usage"] != true {
t.Errorf("stream_options=%#v, want include_usage=true", body["stream_options"])
}
w.Header().Set("Content-Type", "text/event-stream")
_, _ = io.WriteString(w, strings.Join([]string{
`data: {"choices":[{"delta":{"reasoning_content":"step "}}]}`,
`data: {"choices":[{"delta":{"content":"Hello"}}]}`,
`data: {"choices":[{"delta":{"content":" world"},"finish_reason":"stop"}]}`,
`data: [DONE]`,
``,
}, "\n"))
}))
defer srv.Close()
apiKey := "test-key"
var content, reasoning []string
err := newXAIForTest(srv.URL).ChatStreamlyWithSender(
ctx,
"grok-3-mini",
[]Message{{Role: "user", Content: "hi"}},
&APIConfig{ApiKey: &apiKey},
nil,
nil,
func(answer, reason *string) error {
if answer != nil {
content = append(content, *answer)
}
if reason != nil {
reasoning = append(reasoning, *reason)
}
return nil
},
)
if err != nil {
t.Fatalf("ChatStreamlyWithSender: %v", err)
}
if strings.Join(reasoning, "") != "step " {
t.Errorf("reasoning=%q", strings.Join(reasoning, ""))
}
if got := strings.Join(content, ""); got != "Hello world[DONE]" {
t.Errorf("content=%q, want Hello world[DONE]", got)
}
}

182
internal/entity/models/xiaomi.go
View File

@@ -41,6 +41,11 @@ func NewXiaomiModel(baseURL map[string]string, urlSuffix URLSuffix) *XiaomiModel
BaseURL: baseURL,
URLSuffix: urlSuffix,
httpClient: NewDriverHTTPClient(false),
// Xiaomi authenticates with the non-standard "api-key" header
// instead of "Authorization: Bearer".
authHeader: func(cfg *APIConfig) (string, string) {
return "api-key", *cfg.ApiKey
},
},
}
}
@@ -75,7 +80,6 @@ func (x *XiaomiModel) ChatWithMessages(ctx context.Context, modelName string, me
delete(reqBody, "max_tokens")
if chatModelConfig != nil {
if chatModelConfig.MaxTokens != nil {
reqBody["max_completion_tokens"] = *chatModelConfig.MaxTokens
}
@@ -93,94 +97,12 @@ func (x *XiaomiModel) ChatWithMessages(ctx context.Context, modelName string, me
}
}
jsonData, err := json.Marshal(reqBody)
body, err := x.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
return nil, err
}
ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("api-key", *apiConfig.ApiKey)
resp, err := x.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
// Parse response
var result map[string]interface{}
if err = json.Unmarshal(body, &result); err != nil {
return nil, fmt.Errorf("failed to parse response: %w", err)
}
choices, ok := result["choices"].([]interface{})
if !ok || len(choices) == 0 {
return nil, fmt.Errorf("no choices in response")
}
firstChoice, ok := choices[0].(map[string]interface{})
if !ok {
return nil, fmt.Errorf("invalid choice format")
}
messageMap, ok := firstChoice["message"].(map[string]interface{})
if !ok {
return nil, fmt.Errorf("invalid message format")
}
content, ok := messageMap["content"].(string)
if !ok {
return nil, fmt.Errorf("invalid content format")
}
var reasonContent string
reasonContent, _ = messageMap["reasoning_content"].(string)
if reasonContent == "" && chatModelConfig != nil && chatModelConfig.Thinking != nil && *chatModelConfig.Thinking {
// If reasoning_content not in response, try parsing from content tags
reasoning, answer := GetThinkingAndAnswer(chatModelConfig.ModelClass, &content)
if reasoning != nil {
reasonContent = *reasoning
content = *answer
}
}
// if first char of reasonContent is \n remove the '\n'
if reasonContent != "" && reasonContent[0] == '\n' {
reasonContent = reasonContent[1:]
}
var toolCalls []map[string]interface{}
if tcs, ok := messageMap["tool_calls"].([]interface{}); ok {
for _, tc := range tcs {
if tcMap, ok := tc.(map[string]interface{}); ok {
toolCalls = append(toolCalls, tcMap)
}
}
}
chatResponse := &ChatResponse{
Answer: &content,
ReasonContent: &reasonContent,
ToolCalls: toolCalls,
}
return chatResponse, nil
return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig)
}
func (x *XiaomiModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error {
@@ -204,6 +126,9 @@ func (x *XiaomiModel) ChatStreamlyWithSender(ctx context.Context, modelName stri
// Build request body with streaming enabled
reqBody := buildRequestBody(modelConfig, modelName, messages, true)
delete(reqBody, "max_tokens")
reqBody["stream_options"] = map[string]interface{}{
"include_usage": true,
}
if modelConfig != nil {
if modelConfig.MaxTokens != nil {
@@ -224,84 +149,9 @@ func (x *XiaomiModel) ChatStreamlyWithSender(ctx context.Context, modelName stri
}
jsonData, err := json.Marshal(reqBody)
if err != nil {
return fmt.Errorf("failed to marshal request: %w", err)
}
ctx, cancel := context.WithTimeout(ctx, streamCallTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("api-key", *apiConfig.ApiKey)
resp, err := x.baseModel.httpClient.Do(req)
if err != nil {
return fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
// SSE parsing: read line by line
if modelConfig != nil {
modelConfig.ToolCallsResult = nil
}
accumulatedToolCalls := make(map[int]map[string]interface{})
if _, err = ParseSSEStreamTolerant[map[string]interface{}](resp.Body, func(event map[string]interface{}) error {
choices, ok := event["choices"].([]interface{})
if !ok || len(choices) == 0 {
return nil
}
firstChoice, ok := choices[0].(map[string]interface{})
if !ok {
return nil
}
delta, ok := firstChoice["delta"].(map[string]interface{})
if !ok {
return nil
}
accumulateToolCallDeltas(delta, accumulatedToolCalls)
reasoningContent, ok := delta["reasoning_content"].(string)
if ok && reasoningContent != "" {
if err = sender(nil, &reasoningContent); err != nil {
return err
}
}
content, ok := delta["content"].(string)
if ok && content != "" {
if err = sender(&content, nil); err != nil {
return err
}
}
return nil
}); err != nil {
return fmt.Errorf("failed to scan response body: %w", err)
}
setSortedToolCallsResult(modelConfig, accumulatedToolCalls)
// Send [DONE] marker for OpenAI compatibility
endOfStream := "[DONE]"
if err = sender(&endOfStream, nil); err != nil {
return err
}
return nil
return x.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error {
return HandleStreamingResponse(body, modelUsage, modelConfig, OpenAIParserConfig, sender)
})
}
func (x *XiaomiModel) Embed(ctx context.Context, modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig, modelUsage *common.ModelUsage) ([]EmbeddingData, error) {
@@ -485,7 +335,7 @@ func (x *XiaomiModel) newXiaomiASRRequest(ctx context.Context, modelName *string
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("api-key", *apiConfig.ApiKey)
x.baseModel.applyAuth(req, apiConfig)
return req, nil
}
@@ -679,7 +529,7 @@ func (x *XiaomiModel) newXiaomiTTSRequest(ctx context.Context, modelName *string
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("api-key", *apiConfig.ApiKey)
x.baseModel.applyAuth(req, apiConfig)
return req, nil
}

11
internal/entity/models/xiaomi_test.go
View File

@@ -251,6 +251,10 @@ func TestXiaomiStreamHappyPath(t *testing.T) {
if body["stream"] != true {
t.Errorf("stream=%v want true", body["stream"])
}
streamOptions, ok := body["stream_options"].(map[string]interface{})
if !ok || streamOptions["include_usage"] != true {
t.Errorf("stream_options=%#v, want include_usage=true", body["stream_options"])
}
w.Header().Set("Content-Type", "text/event-stream")
_, _ = io.WriteString(w,
`data: {"choices":[{"delta":{"reasoning_content":"step "}}]}`+"\n\n"+
@@ -344,10 +348,11 @@ func TestXiaomiStreamRejectsMalformedFrame(t *testing.T) {
defer srv.Close()
apiKey := "test-key"
// Malformed SSE frames are silently skipped; the stream completes and sends [DONE].
// Malformed SSE frames abort the stream: Xiaomi now uses the strict
// OpenAIParserConfig shared by every OpenAI-compatible driver.
err := newXiaomiForTest(srv.URL).ChatStreamlyWithSender(ctx, "mimo-v2.5-pro", []Message{{Role: "user", Content: "x"}}, &APIConfig{ApiKey: &apiKey}, nil, nil, func(*string, *string) error { return nil })
if err != nil {
t.Errorf("expected no error on malformed frame, got %v", err)
if err == nil {
t.Error("expected error on malformed frame, got nil")
}
}

186
internal/entity/models/xinference.go
View File

@@ -28,30 +28,13 @@ import (
"ragflow/internal/common"
"strconv"
"strings"
"sync"
"time"
)
var xinferenceStreamIdleTimeout = 60 * time.Second
// XinferenceModel implements ModelDriver for Xinference chat models.
type XinferenceModel struct {
baseModel BaseModel
}
type xinferenceChatChoice struct {
Message struct {
Content string `json:"content"`
ReasoningContent string `json:"reasoning_content"`
Reasoning string `json:"reasoning"`
Thinking string `json:"thinking"`
} `json:"message"`
}
type xinferenceChatResponse struct {
Choices []xinferenceChatChoice `json:"choices"`
}
type xinferenceModelListResponse struct {
Data []ModelListItem `json:"data"`
}
@@ -87,28 +70,6 @@ func normalizeXinferenceBaseURL(base string) string {
return trimmed
}
func xinferenceReasoningFromStrings(reasoningContent string, reasoning string, thinking string) string {
switch {
case reasoningContent != "":
return reasoningContent
case reasoning != "":
return reasoning
case thinking != "":
return thinking
default:
return ""
}
}
func xinferenceReasoningFromMap(value map[string]interface{}) string {
for _, field := range []string{"reasoning_content", "reasoning", "thinking"} {
if text, ok := value[field].(string); ok && text != "" {
return text
}
}
return ""
}
// ChatWithMessages sends multiple messages with roles and returns the response.
func (x *XinferenceModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) {
if err := x.baseModel.APIConfigCheck(apiConfig); err != nil {
@@ -127,56 +88,13 @@ func (x *XinferenceModel) ChatWithMessages(ctx context.Context, modelName string
url := fmt.Sprintf("%s/%s", baseURL, x.baseModel.URLSuffix.Chat)
reqBody := buildRequestBody(chatModelConfig, modelName, messages, false)
jsonData, err := json.Marshal(reqBody)
body, err := x.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
return nil, err
}
ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
if auth := BearerAuth(apiConfig); auth != "" {
req.Header.Set("Authorization", auth)
}
resp, err := x.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
var result xinferenceChatResponse
if err = json.Unmarshal(body, &result); err != nil {
return nil, fmt.Errorf("failed to parse response: %w", err)
}
if len(result.Choices) == 0 {
return nil, fmt.Errorf("no choices in response")
}
content := result.Choices[0].Message.Content
reasonContent := xinferenceReasoningFromStrings(
result.Choices[0].Message.ReasoningContent,
result.Choices[0].Message.Reasoning,
result.Choices[0].Message.Thinking,
)
return &ChatResponse{
Answer: &content,
ReasonContent: &reasonContent,
}, nil
return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig)
}
// ChatStreamlyWithSender sends messages and streams response via sender.
@@ -203,101 +121,13 @@ func (x *XinferenceModel) ChatStreamlyWithSender(ctx context.Context, modelName
url := fmt.Sprintf("%s/%s", baseURL, x.baseModel.URLSuffix.Chat)
reqBody := buildRequestBody(chatModelConfig, modelName, messages, true)
jsonData, err := json.Marshal(reqBody)
if err != nil {
return fmt.Errorf("failed to marshal request: %w", err)
reqBody["stream_options"] = map[string]interface{}{
"include_usage": true,
}
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
if auth := BearerAuth(apiConfig); auth != "" {
req.Header.Set("Authorization", auth)
}
resp, err := x.baseModel.httpClient.Do(req)
if err != nil {
return fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
lastActive := time.Now()
var lastActiveMu sync.Mutex
done := make(chan struct{})
defer close(done)
go func() {
ticker := time.NewTicker(xinferenceStreamIdleTimeout / 4)
defer ticker.Stop()
for {
select {
case <-done:
return
case now := <-ticker.C:
lastActiveMu.Lock()
idle := now.Sub(lastActive)
lastActiveMu.Unlock()
if idle >= xinferenceStreamIdleTimeout {
cancel()
return
}
}
}
}()
sawTerminal := false
sseDone, parseErr := ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error {
lastActiveMu.Lock()
lastActive = time.Now()
lastActiveMu.Unlock()
choices, ok := event["choices"].([]interface{})
if !ok || len(choices) == 0 {
return nil
}
firstChoice, ok := choices[0].(map[string]interface{})
if !ok {
return nil
}
if delta, ok := firstChoice["delta"].(map[string]interface{}); ok {
if reasoning := xinferenceReasoningFromMap(delta); reasoning != "" {
if err := sender(nil, &reasoning); err != nil {
return err
}
}
if content, ok := delta["content"].(string); ok && content != "" {
if err := sender(&content, nil); err != nil {
return err
}
}
}
if finishReason, ok := firstChoice["finish_reason"].(string); ok && finishReason != "" {
sawTerminal = true
}
return nil
return x.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error {
return HandleStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig, sender)
})
if parseErr != nil {
if ctx.Err() != nil {
return fmt.Errorf("xinference: stream idle for more than %s, aborted", xinferenceStreamIdleTimeout)
}
return fmt.Errorf("failed to scan response body: %w", parseErr)
}
if !sseDone && !sawTerminal {
return fmt.Errorf("xinference: stream ended before [DONE] or finish_reason")
}
endOfStream := "[DONE]"
return sender(&endOfStream, nil)
}
// Index is *int so a missing JSON field is distinguishable from index 0.

42
internal/entity/models/xinference_test.go
View File

@@ -7,7 +7,6 @@ import (
"net/http/httptest"
"strings"
"testing"
"time"
)
func newXinferenceForTest(baseURL string) *XinferenceModel {
@@ -22,15 +21,6 @@ func newXinferenceForTest(baseURL string) *XinferenceModel {
)
}
func withXinferenceIdleTimeout(t *testing.T, d time.Duration) {
t.Helper()
original := xinferenceStreamIdleTimeout
xinferenceStreamIdleTimeout = d
t.Cleanup(func() {
xinferenceStreamIdleTimeout = original
})
}
func TestXinferenceName(t *testing.T) {
x := newXinferenceForTest("http://unused")
if got := x.Name(); got != "Xinference" {
@@ -229,38 +219,6 @@ func TestXinferenceStreamRejectsFalseStreamConfig(t *testing.T) {
}
}
func TestXinferenceStreamCancelsOnIdle(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
withXinferenceIdleTimeout(t, 200*time.Millisecond)
hold := make(chan struct{})
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
if f, ok := w.(http.Flusher); ok {
_, _ = io.WriteString(w, `data: {"choices":[{"delta":{"content":"hi"}}]}`+"\n")
f.Flush()
}
select {
case <-hold:
case <-r.Context().Done():
}
}))
t.Cleanup(srv.Close)
t.Cleanup(func() { close(hold) })
x := newXinferenceForTest(srv.URL)
err := x.ChatStreamlyWithSender(ctx, "qwen2.5-instruct",
[]Message{{Role: "user", Content: "x"}},
&APIConfig{}, nil,
nil,
func(*string, *string) error { return nil })
if err == nil || !strings.Contains(err.Error(), "stream idle") {
t.Errorf("expected stream-idle error, got %v", err)
}
}
func TestXinferenceListModelsAndCheckConnection(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()

202
internal/entity/models/xunfei.go
View File

@@ -17,12 +17,10 @@
package models
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"ragflow/internal/common"
"strings"
)
@@ -76,6 +74,11 @@ func NewXunFeiModel(baseURL map[string]string, urlSuffix URLSuffix) *XunFeiModel
BaseURL: baseURL,
URLSuffix: urlSuffix,
httpClient: NewDriverHTTPClient(false),
// The Spark HTTP API authenticates with the credential bundle's
// spark_api_password, not the raw stored key.
authHeader: func(cfg *APIConfig) (string, string) {
return "Authorization", "Bearer " + resolveBearerToken(cfg)
},
},
}
}
@@ -118,82 +121,12 @@ func (x *XunFeiModel) ChatWithMessages(ctx context.Context, modelName string, me
}
}
jsonData, err := json.Marshal(reqBody)
body, err := x.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout)
if err != nil {
return nil, fmt.Errorf("failed to marshal request body: %w", err)
return nil, err
}
ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", resolveBearerToken(apiConfig)))
resp, err := x.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response body: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
// Parse Response
var result map[string]interface{}
if err := json.Unmarshal(body, &result); err != nil {
return nil, fmt.Errorf("failed to unmarshal response body: %w", err)
}
choices, ok := result["choices"].([]interface{})
if !ok {
return nil, fmt.Errorf("no choices in response")
}
firstChoice, ok := choices[0].(map[string]interface{})
if !ok {
return nil, fmt.Errorf("no choices in response")
}
messageMap, ok := firstChoice["message"].(map[string]interface{})
if !ok {
return nil, fmt.Errorf("no message in response")
}
content, ok := messageMap["content"].(string)
toolCalls := extractToolCalls(messageMap)
if !ok && len(toolCalls) == 0 {
return nil, fmt.Errorf("no message in response")
}
var reasonContent string
if chatModelConfig != nil && chatModelConfig.Thinking != nil && *chatModelConfig.Thinking {
reasonContent, ok = messageMap["reasoning_content"].(string)
if !ok {
return nil, fmt.Errorf("invalid content format")
}
if reasonContent != "" && reasonContent[0] == '\n' {
reasonContent = reasonContent[1:]
}
}
chatResponse := &ChatResponse{
Answer: &content,
ReasonContent: &reasonContent,
ToolCalls: toolCalls,
}
return chatResponse, nil
return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig)
}
func (x *XunFeiModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error {
@@ -227,84 +160,13 @@ func (x *XunFeiModel) ChatStreamlyWithSender(ctx context.Context, modelName stri
}
}
jsonData, err := json.Marshal(reqBody)
if err != nil {
return fmt.Errorf("failed to marshal request: %w", err)
}
ctx, cancel := context.WithTimeout(ctx, streamCallTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", resolveBearerToken(apiConfig)))
resp, err := x.baseModel.httpClient.Do(req)
if err != nil {
return fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return fmt.Errorf("invalid status code: %d, body: %s", resp.StatusCode, string(body))
}
// 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 {
if data, marshalErr := json.Marshal(event); marshalErr == nil {
common.Info(string(data))
}
choices, ok := event["choices"].([]interface{})
if !ok || len(choices) == 0 {
return nil
}
firstChoice, ok := choices[0].(map[string]interface{})
if !ok {
return nil
}
delta, ok := firstChoice["delta"].(map[string]interface{})
if !ok {
return nil
}
accumulateToolCallDeltas(delta, accumulatedToolCalls)
reasoningContent, ok := delta["reasoning_content"].(string)
if ok && reasoningContent != "" {
if err := sender(nil, &reasoningContent); err != nil {
return err
}
}
content, ok := delta["content"].(string)
if ok && content != "" {
if err := sender(&content, nil); err != nil {
return err
}
}
return nil
}); err != nil {
return fmt.Errorf("failed to scan response body: %w", err)
}
setSortedToolCallsResult(modelConfig, accumulatedToolCalls)
// Send [DONE] marker for OpenAI compatibility
endOfStream := "[DONE]"
if err = sender(&endOfStream, nil); err != nil {
return err
}
return nil
// XunFei's OpenAI-compatible endpoint does not document support for
// stream_options.include_usage, but its streaming responses carry a
// usage object in the final chunk by default, so usage reporting works
// without it.
return x.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error {
return HandleStreamingResponse(body, modelUsage, modelConfig, OpenAIParserConfig, sender)
})
}
func (x *XunFeiModel) Embed(ctx context.Context, modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig, modelUsage *common.ModelUsage) ([]EmbeddingData, error) {
@@ -350,41 +212,11 @@ func (x *XunFeiModel) ListModels(ctx context.Context, apiConfig *APIConfig) ([]L
}
url := fmt.Sprintf("%s/%s", resolvedBaseURL, x.baseModel.URLSuffix.Models)
// Build request body
reqBody := map[string]interface{}{}
jsonData, err := json.Marshal(reqBody)
body, err := x.baseModel.doGetRequest(ctx, url, apiConfig, nonStreamCallTimeout)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
return nil, err
}
ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, "GET", url, bytes.NewBuffer(jsonData))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", resolveBearerToken(apiConfig)))
resp, err := x.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
// Parse response
// Parse response
var modelList ModelList
if err = json.Unmarshal(body, &modelList); err != nil {

120
internal/entity/models/xunfei_test.go
View File

@@ -1,6 +1,10 @@
package models
import (
"encoding/json"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
)
@@ -119,3 +123,119 @@ func TestXunFeiUnsupportedMethodsReturnNoSuchMethod(t *testing.T) {
})
}
}
func newXunFeiForTest(baseURL string) *XunFeiModel {
return NewXunFeiModel(
map[string]string{"default": baseURL},
URLSuffix{Chat: "v1/chat/completions", Models: "v1/models"},
)
}
func TestXunFeiChatUsesResolvedBearerToken(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
t.Errorf("method=%s, want POST", r.Method)
}
// The Spark credential bundle is stored as JSON; the request must
// authenticate with the extracted spark_api_password.
if got := r.Header.Get("Authorization"); got != "Bearer pwd" {
t.Errorf("Authorization=%q, want Bearer pwd", got)
}
var body map[string]interface{}
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
t.Errorf("decode request: %v", err)
return
}
if body["model"] != "lite" {
t.Errorf("model=%v, want lite (resolved Spark-Lite)", body["model"])
}
if body["stream"] != false {
t.Errorf("stream=%v, want false", body["stream"])
}
_ = json.NewEncoder(w).Encode(map[string]interface{}{
"choices": []map[string]interface{}{{
"message": map[string]interface{}{
"content": "pong",
"reasoning_content": "\nthought",
},
}},
})
}))
defer srv.Close()
bundle := `{"spark_api_password":"pwd","spark_app_id":"app","spark_api_secret":"secret","spark_api_key":"key"}`
resp, err := newXunFeiForTest(srv.URL).ChatWithMessages(
ctx,
"Spark-Lite",
[]Message{{Role: "user", Content: "ping"}},
&APIConfig{ApiKey: &bundle},
nil,
nil,
)
if err != nil {
t.Fatalf("ChatWithMessages: %v", err)
}
if resp.Answer == nil || *resp.Answer != "pong" {
t.Errorf("Answer=%v, want pong", resp.Answer)
}
if resp.ReasonContent == nil || *resp.ReasonContent != "thought" {
t.Errorf("ReasonContent=%v, want thought", resp.ReasonContent)
}
}
func TestXunFeiStreamHappyPath(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if got := r.Header.Get("Authorization"); got != "Bearer sk-plain" {
t.Errorf("Authorization=%q, want Bearer sk-plain", got)
}
w.Header().Set("Content-Type", "text/event-stream")
_, _ = io.WriteString(w, strings.Join([]string{
`data: {"choices":[{"delta":{"reasoning_content":"step "}}]}`,
`data: {"choices":[{"delta":{"content":"Hello"}}]}`,
`data: {"choices":[{"delta":{"content":" world"},"finish_reason":"stop"}]}`,
// XunFei carries usage in the final chunk without requiring
// stream_options.include_usage.
`data: {"choices":[],"usage":{"prompt_tokens":3,"completion_tokens":5,"total_tokens":8}}`,
`data: [DONE]`,
``,
}, "\n"))
}))
defer srv.Close()
apiKey := "sk-plain"
var content, reasoning []string
config := &ChatConfig{}
err := newXunFeiForTest(srv.URL).ChatStreamlyWithSender(
ctx,
"Spark-Lite",
[]Message{{Role: "user", Content: "hi"}},
&APIConfig{ApiKey: &apiKey},
config,
nil,
func(answer, reason *string) error {
if answer != nil {
content = append(content, *answer)
}
if reason != nil {
reasoning = append(reasoning, *reason)
}
return nil
},
)
if err != nil {
t.Fatalf("ChatStreamlyWithSender: %v", err)
}
if strings.Join(reasoning, "") != "step " {
t.Errorf("reasoning=%q", strings.Join(reasoning, ""))
}
if got := strings.Join(content, ""); got != "Hello world[DONE]" {
t.Errorf("content=%q, want Hello world[DONE]", got)
}
if config.UsageResult == nil || config.UsageResult.TotalTokens != 8 {
t.Errorf("UsageResult=%#v, want total tokens 8", config.UsageResult)
}
}

172
internal/entity/models/zhipu-ai.go
View File

@@ -55,32 +55,6 @@ func (z *ZhipuAIModel) Name() string {
return "zhipu"
}
type ZhipuChatResponse struct {
Choices []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"`
ToolCalls []map[string]any `json:"tool_calls"`
} `json:"message"`
} `json:"choices"`
Created int `json:"created"`
Id string `json:"id"`
Model string `json:"model"`
Object string `json:"object"`
RequestId string `json:"request_id"`
Usage struct {
CompletionTokens int `json:"completion_tokens"`
PromptTokens int `json:"prompt_tokens"`
PromptTokensDetails struct {
CachedTokens int `json:"cached_tokens"`
} `json:"prompt_tokens_details"`
TotalTokens int `json:"total_tokens"`
} `json:"usage"`
}
// ChatWithMessages sends multiple messages with roles and returns response
func (z *ZhipuAIModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) {
if err := z.baseModel.APIConfigCheck(apiConfig); err != nil {
@@ -111,68 +85,14 @@ func (z *ZhipuAIModel) ChatWithMessages(ctx context.Context, modelName string, m
}
}
}
}
jsonData, err := json.Marshal(reqBody)
body, err := z.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
return nil, err
}
ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := z.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, chatConfig *ChatConfig) (chatResponseParts, error) {
var result ZhipuChatResponse
if err := json.Unmarshal(body, &result); err != nil {
return chatResponseParts{}, fmt.Errorf("failed to parse response: %w", err)
}
if len(result.Choices) == 0 {
return chatResponseParts{}, fmt.Errorf("empty response")
}
choice := &result.Choices[0]
var reasonContent *string
if chatConfig != nil && chatConfig.Thinking != nil && *chatConfig.Thinking {
reasonContent = &choice.Message.ReasoningContent
}
return chatResponseParts{
RequestID: result.RequestId,
Content: &choice.Message.Content,
ReasonContent: reasonContent,
ToolCalls: choice.Message.ToolCalls,
Usage: &TokenUsage{
PromptTokens: result.Usage.PromptTokens,
CompletionTokens: result.Usage.CompletionTokens,
TotalTokens: result.Usage.TotalTokens,
},
}, nil
})
return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig)
}
// ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel)
@@ -192,6 +112,9 @@ func (z *ZhipuAIModel) ChatStreamlyWithSender(ctx context.Context, modelName str
url := fmt.Sprintf("%s/%s", baseURL, z.baseModel.URLSuffix.Chat)
reqBody := buildRequestBody(chatModelConfig, modelName, messages, true)
reqBody["stream_options"] = map[string]interface{}{
"include_usage": true,
}
if chatModelConfig != nil {
if chatModelConfig.Thinking != nil {
@@ -207,86 +130,9 @@ func (z *ZhipuAIModel) ChatStreamlyWithSender(ctx context.Context, modelName str
}
}
jsonData, err := json.Marshal(reqBody)
if err != nil {
return fmt.Errorf("failed to marshal request: %w", err)
}
ctx, cancel := context.WithTimeout(ctx, streamCallTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := z.baseModel.httpClient.Do(req)
if err != nil {
return fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
// 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))
tokenUsage, found, usageErr := decodeOpenAICompatibleStreamUsage(event)
if usageErr != nil {
return usageErr
}
if found {
applyStreamUsage(chatModelConfig, modelUsage, tokenUsage)
}
choices, ok := event["choices"].([]interface{})
if !ok || len(choices) == 0 {
return nil
}
firstChoice, ok := choices[0].(map[string]interface{})
if !ok {
return nil
}
delta, ok := firstChoice["delta"].(map[string]interface{})
if !ok {
return nil
}
accumulateToolCallDeltas(delta, accumulatedToolCalls)
reasoningContent, ok := delta["reasoning_content"].(string)
if ok && reasoningContent != "" {
if err = sender(nil, &reasoningContent); err != nil {
return err
}
}
content, ok := delta["content"].(string)
if ok && content != "" {
if err = sender(&content, nil); err != nil {
return err
}
}
return nil
}); err != nil {
return fmt.Errorf("failed to scan response body: %w", err)
}
setSortedToolCallsResult(chatModelConfig, accumulatedToolCalls)
// Send [DONE] marker for OpenAI compatibility
endOfStream := "[DONE]"
return sender(&endOfStream, nil)
return z.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error {
return HandleStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig, sender)
})
}
type zhipuEmbeddingResponse struct {