mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-04 23:00:30 +08:00
feat(go-models): migrate batch 1 model drivers to unified handlers (#17696)
## Summary Relate to #17284 Migrate 10 OpenAI-compatible drivers (`302ai`, `aliyun`, `astraflow`, `avian`, `azure_openai`, `baichuan`, `baidu`, `cometapi`, `deepinfra`, `futurmix`) to use the unified response handlers (`HandleNonStreamingResponse` / `HandleStreamingResponse`), following the same pattern established by `deepseek` in #17634. - Cut ~150 lines per driver (1507 lines removed, 144 added across 10 files). - No functional changes — pure deduplication of HTTP plumbing. - Each driver now routes through `baseModel.doRequest()` and `HandleNonStreamingResponse()`. --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
jay77721
@@ -75,7 +75,6 @@ func (a *AI302Model) ChatWithMessages(ctx context.Context, modelName string, mes
|
||||
if err := a.baseModel.APIConfigCheck(apiConfig); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
apiKey := strings.TrimSpace(*apiConfig.ApiKey)
|
||||
if strings.TrimSpace(modelName) == "" {
|
||||
return nil, fmt.Errorf("model name is required")
|
||||
}
|
||||
@@ -113,90 +112,18 @@ func (a *AI302Model) ChatWithMessages(ctx context.Context, modelName string, mes
|
||||
}
|
||||
}
|
||||
|
||||
jsonData, err := json.Marshal(reqBody)
|
||||
body, err := a.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", apiKey))
|
||||
|
||||
resp, err := a.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 first char of reasonContent is \n remove the '\n'
|
||||
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 (a *AI302Model) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error {
|
||||
if err := a.baseModel.APIConfigCheck(apiConfig); err != nil {
|
||||
return err
|
||||
}
|
||||
apiKey := strings.TrimSpace(*apiConfig.ApiKey)
|
||||
if strings.TrimSpace(modelName) == "" {
|
||||
return fmt.Errorf("model name is required")
|
||||
}
|
||||
@@ -237,77 +164,11 @@ func (a *AI302Model) ChatStreamlyWithSender(ctx context.Context, modelName strin
|
||||
}
|
||||
}
|
||||
|
||||
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", apiKey))
|
||||
req.Header.Set("Accept", "text/event-stream")
|
||||
|
||||
resp, err := a.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)
|
||||
if _, 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
|
||||
}
|
||||
}
|
||||
|
||||
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]"
|
||||
return sender(&endOfStream, nil)
|
||||
return a.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error {
|
||||
return HandleStreamingResponse(body, modelUsage, modelConfig, OpenAIParserConfig, sender)
|
||||
})
|
||||
}
|
||||
|
||||
func (a *AI302Model) Embed(ctx context.Context, modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig, modelUsage *common.ModelUsage) ([]EmbeddingData, error) {
|
||||
|
||||
@@ -52,37 +52,6 @@ func (a *AliyunModel) Name() string {
|
||||
return "Tongyi-Qianwen"
|
||||
}
|
||||
|
||||
// AliyunChatResponse mirrors DashScope's OpenAI-compatible chat response.
|
||||
type AliyunChatResponse 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"`
|
||||
}
|
||||
|
||||
func (a *AliyunModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) {
|
||||
if err := a.baseModel.APIConfigCheck(apiConfig); err != nil {
|
||||
return nil, err
|
||||
@@ -121,78 +90,12 @@ func (a *AliyunModel) ChatWithMessages(ctx context.Context, modelName string, me
|
||||
// enabled by the user, matching Python's chat_model.py behavior.
|
||||
applyQwen3ThinkingDefault(modelName, reqBody)
|
||||
|
||||
jsonData, err := json.Marshal(reqBody)
|
||||
body, err := a.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 := a.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 AliyunChatResponse
|
||||
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("no choices in response")
|
||||
}
|
||||
|
||||
choice := &result.Choices[0]
|
||||
content := choice.Message.Content
|
||||
if content == "" && len(choice.Message.ToolCalls) == 0 {
|
||||
return chatResponseParts{}, fmt.Errorf("response contains neither content nor tool calls")
|
||||
}
|
||||
reasonContent := ""
|
||||
if chatConfig != nil && chatConfig.Thinking != nil && *chatConfig.Thinking {
|
||||
reasonContent = choice.Message.ReasoningContent
|
||||
if 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)
|
||||
@@ -235,96 +138,9 @@ func (a *AliyunModel) ChatStreamlyWithSender(ctx context.Context, modelName stri
|
||||
// enabled by the user, matching Python's chat_model.py behavior.
|
||||
applyQwen3ThinkingDefault(modelName, reqBody)
|
||||
|
||||
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 := a.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]interface{})
|
||||
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(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
|
||||
}
|
||||
if finishReason, ok := firstChoice["finish_reason"].(string); ok && finishReason != "" {
|
||||
sawTerminal = true
|
||||
}
|
||||
|
||||
delta, ok := firstChoice["delta"].(map[string]interface{})
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
|
||||
if accumulateToolCallDeltas(delta, accumulatedToolCalls) {
|
||||
return nil
|
||||
}
|
||||
|
||||
content, ok := delta["content"].(string)
|
||||
if ok && content != "" {
|
||||
if err := sender(&content, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
reasoningContent, ok := delta["reasoning_content"].(string)
|
||||
if ok && reasoningContent != "" {
|
||||
if err := sender(nil, &reasoningContent); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
return a.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)
|
||||
}
|
||||
if !done && !sawTerminal {
|
||||
return fmt.Errorf("aliyun: stream ended before [DONE] or finish_reason")
|
||||
}
|
||||
|
||||
setSortedToolCallsResult(chatModelConfig, accumulatedToolCalls)
|
||||
|
||||
// Send [DONE] marker for OpenAI compatibility
|
||||
endOfStream := "[DONE]"
|
||||
return sender(&endOfStream, nil)
|
||||
}
|
||||
|
||||
// applyQwen3ThinkingDefault ensures enable_thinking=false is sent for qwen3
|
||||
|
||||
@@ -88,72 +88,12 @@ func (a *AstraflowModel) ChatWithMessages(ctx context.Context, modelName string,
|
||||
if chatModelConfig != nil {
|
||||
}
|
||||
|
||||
jsonData, err := json.Marshal(reqBody)
|
||||
body, err := a.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 := a.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 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")
|
||||
}
|
||||
|
||||
reasonContent := ""
|
||||
if r, ok := messageMap["reasoning_content"].(string); ok {
|
||||
reasonContent = r
|
||||
}
|
||||
|
||||
return &ChatResponse{
|
||||
Answer: &content,
|
||||
ReasonContent: &reasonContent,
|
||||
ToolCalls: toolCalls,
|
||||
}, nil
|
||||
return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig)
|
||||
}
|
||||
|
||||
// ChatStreamlyWithSender opens the SSE chat-completions
|
||||
@@ -180,80 +120,13 @@ func (a *AstraflowModel) ChatStreamlyWithSender(ctx context.Context, modelName s
|
||||
url := fmt.Sprintf("%s/%s", baseURL, a.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]any{
|
||||
"include_usage": true,
|
||||
}
|
||||
|
||||
// SSE is long-lived; rely on the transport's ResponseHeaderTimeout
|
||||
// to cap connection-establishment instead of a hard deadline.
|
||||
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 := a.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 {
|
||||
if apiErr, ok := event["error"]; ok {
|
||||
return fmt.Errorf("astraflow: upstream stream error: %v", apiErr)
|
||||
}
|
||||
|
||||
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 {
|
||||
accumulateToolCallDeltas(delta, accumulatedToolCalls)
|
||||
if r, ok := delta["reasoning_content"].(string); ok && r != "" {
|
||||
rr := r
|
||||
if err := sender(nil, &rr); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if c, ok := delta["content"].(string); ok && c != "" {
|
||||
cc := c
|
||||
if err := sender(&cc, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
if finish, ok := firstChoice["finish_reason"].(string); ok && finish != "" {
|
||||
sawTerminal = true
|
||||
}
|
||||
return nil
|
||||
return a.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("astraflow: stream ended before [DONE] or finish_reason")
|
||||
}
|
||||
|
||||
endOfStream := "[DONE]"
|
||||
if err := sender(&endOfStream, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (a *AstraflowModel) ListModels(ctx context.Context, apiConfig *APIConfig) ([]ListModelResponse, error) {
|
||||
|
||||
@@ -17,7 +17,6 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
@@ -54,33 +53,13 @@ func (a *AvianModel) Name() string {
|
||||
}
|
||||
|
||||
func (a *AvianModel) chatURL(apiConfig *APIConfig) (string, error) {
|
||||
|
||||
baseURL, err := a.baseModel.GetBaseURL(apiConfig)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
baseURL = strings.TrimSuffix(baseURL, "/")
|
||||
return fmt.Sprintf("%s/%s", baseURL, a.baseModel.URLSuffix.Chat), nil
|
||||
}
|
||||
|
||||
type avianChatMessage struct {
|
||||
Content string `json:"content"`
|
||||
ReasoningContent string `json:"reasoning_content"`
|
||||
Reasoning string `json:"reasoning"`
|
||||
}
|
||||
|
||||
type avianChatChoice struct {
|
||||
Message avianChatMessage `json:"message"`
|
||||
Delta avianChatMessage `json:"delta"`
|
||||
FinishReason string `json:"finish_reason"`
|
||||
}
|
||||
|
||||
type avianChatResponse struct {
|
||||
Choices []avianChatChoice `json:"choices"`
|
||||
Error interface{} `json:"error"`
|
||||
FinishReason string `json:"finish_reason"`
|
||||
}
|
||||
|
||||
func (a *AvianModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) {
|
||||
if err := a.baseModel.APIConfigCheck(apiConfig); err != nil {
|
||||
return nil, err
|
||||
@@ -97,55 +76,14 @@ func (a *AvianModel) ChatWithMessages(ctx context.Context, modelName string, mes
|
||||
return nil, err
|
||||
}
|
||||
|
||||
jsonData, err := json.Marshal(buildRequestBody(chatModelConfig, modelName, messages, false))
|
||||
reqBody := buildRequestBody(chatModelConfig, modelName, messages, false)
|
||||
|
||||
body, err := a.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, http.MethodPost, 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 := a.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 avianChatResponse
|
||||
if err = json.Unmarshal(body, &result); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse response: %w", err)
|
||||
}
|
||||
if result.Error != nil {
|
||||
return nil, fmt.Errorf("avian: upstream error: %v", result.Error)
|
||||
}
|
||||
if len(result.Choices) == 0 {
|
||||
return nil, fmt.Errorf("no choices in response")
|
||||
}
|
||||
|
||||
content := result.Choices[0].Message.Content
|
||||
reasonContent := result.Choices[0].Message.ReasoningContent
|
||||
if reasonContent == "" {
|
||||
reasonContent = result.Choices[0].Message.Reasoning
|
||||
}
|
||||
return &ChatResponse{
|
||||
Answer: &content,
|
||||
ReasonContent: &reasonContent,
|
||||
}, nil
|
||||
return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig)
|
||||
}
|
||||
|
||||
func (a *AvianModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error {
|
||||
@@ -171,79 +109,25 @@ func (a *AvianModel) ChatStreamlyWithSender(ctx context.Context, modelName strin
|
||||
return err
|
||||
}
|
||||
|
||||
jsonData, err := json.Marshal(buildRequestBody(chatModelConfig, modelName, messages, true))
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal request: %w", err)
|
||||
reqBody := buildRequestBody(chatModelConfig, modelName, messages, true)
|
||||
reqBody["stream_options"] = map[string]interface{}{
|
||||
"include_usage": true,
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(ctx, streamCallTimeout)
|
||||
defer cancel()
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, 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 := a.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]any](resp.Body, func(event map[string]any) 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
|
||||
if chatModelConfig != nil && chatModelConfig.Thinking != nil {
|
||||
if *chatModelConfig.Thinking {
|
||||
reqBody["thinking"] = map[string]interface{}{
|
||||
"type": "enabled",
|
||||
}
|
||||
} else {
|
||||
reqBody["thinking"] = map[string]interface{}{
|
||||
"type": "disabled",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
content, ok := delta["content"].(string)
|
||||
if ok && content != "" {
|
||||
if err = sender(&content, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
return a.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)
|
||||
}
|
||||
if !done && !sawTerminal {
|
||||
return fmt.Errorf("avian: stream ended before [DONE] or finish_reason")
|
||||
}
|
||||
|
||||
setSortedToolCallsResult(chatModelConfig, accumulatedToolCalls)
|
||||
|
||||
endOfStream := "[DONE]"
|
||||
return sender(&endOfStream, nil)
|
||||
}
|
||||
|
||||
type avianModelInfo struct {
|
||||
|
||||
@@ -141,43 +141,7 @@ func (a *AzureOpenAIModel) ChatWithMessages(ctx context.Context, modelName strin
|
||||
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
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
|
||||
if rc, ok := messageMap["reasoning_content"].(string); ok {
|
||||
reasonContent = rc
|
||||
if reasonContent != "" && reasonContent[0] == '\n' {
|
||||
reasonContent = reasonContent[1:]
|
||||
}
|
||||
}
|
||||
|
||||
return &ChatResponse{
|
||||
Answer: &content,
|
||||
ReasonContent: &reasonContent,
|
||||
}, nil
|
||||
return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig)
|
||||
}
|
||||
|
||||
// ChatStreamlyWithSender sends messages and streams the response via the
|
||||
@@ -241,6 +205,10 @@ func (a *AzureOpenAIModel) ChatStreamlyWithSender(ctx context.Context, modelName
|
||||
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)
|
||||
@@ -260,49 +228,7 @@ func (a *AzureOpenAIModel) ChatStreamlyWithSender(ctx context.Context, modelName
|
||||
return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
sawTerminal := false
|
||||
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
|
||||
}
|
||||
|
||||
if reasoningContent, ok := delta["reasoning_content"].(string); ok && reasoningContent != "" {
|
||||
if err := sender(nil, &reasoningContent); 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
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to scan response body: %w", err)
|
||||
}
|
||||
if !done && !sawTerminal {
|
||||
return fmt.Errorf("azure-openai: stream ended before [DONE] or finish_reason")
|
||||
}
|
||||
|
||||
endOfStream := "[DONE]"
|
||||
return sender(&endOfStream, nil)
|
||||
return HandleStreamingResponse(resp.Body, modelUsage, chatModelConfig, OpenAIParserConfig, sender)
|
||||
}
|
||||
|
||||
type azureEmbeddingResponse struct {
|
||||
|
||||
@@ -49,29 +49,6 @@ func (b *BaichuanModel) Name() string {
|
||||
return "BaiChuan"
|
||||
}
|
||||
|
||||
type BaiChuanChatResponse struct {
|
||||
ID string `json:"id"`
|
||||
Choices []struct {
|
||||
FinishReason string `json:"finish_reason"`
|
||||
Index int `json:"index"`
|
||||
Message struct {
|
||||
Content string `json:"content"`
|
||||
Role string `json:"role"`
|
||||
ToolCalls []map[string]any `json:"tool_calls"`
|
||||
} `json:"message"`
|
||||
Logprobs interface{} `json:"logprobs"`
|
||||
} `json:"choices"`
|
||||
Created int `json:"created"`
|
||||
Model string `json:"model"`
|
||||
Object string `json:"object"`
|
||||
Usage struct {
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
TotalTokens int `json:"total_tokens"`
|
||||
SearchCount int `json:"search_count"`
|
||||
} `json:"usage"`
|
||||
}
|
||||
|
||||
func (b *BaichuanModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) {
|
||||
if err := b.baseModel.APIConfigCheck(apiConfig); err != nil {
|
||||
return nil, err
|
||||
@@ -87,66 +64,12 @@ func (b *BaichuanModel) ChatWithMessages(ctx context.Context, modelName string,
|
||||
url := fmt.Sprintf("%s/%s", resolvedBaseURL, b.baseModel.URLSuffix.Chat)
|
||||
reqBody := buildRequestBody(chatModelConfig, modelName, messages, false)
|
||||
|
||||
jsonData, err := json.Marshal(reqBody)
|
||||
body, err := b.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.Add("Content-Type", "application/json")
|
||||
req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
|
||||
|
||||
resp, err := b.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("failed to send request: %d %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
// Parse response
|
||||
return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, chatConfig *ChatConfig) (chatResponseParts, error) {
|
||||
var result BaiChuanChatResponse
|
||||
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("no message in response")
|
||||
}
|
||||
|
||||
emptyReason := ""
|
||||
return chatResponseParts{
|
||||
RequestID: result.ID,
|
||||
Content: &content,
|
||||
ReasonContent: &emptyReason,
|
||||
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)
|
||||
}
|
||||
|
||||
func (b *BaichuanModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error {
|
||||
@@ -165,90 +88,11 @@ func (b *BaichuanModel) ChatStreamlyWithSender(ctx context.Context, modelName st
|
||||
url := fmt.Sprintf("%s/%s", resolvedBaseURL, b.baseModel.URLSuffix.Chat)
|
||||
|
||||
reqBody := buildRequestBody(modelConfig, modelName, messages, true)
|
||||
reqBody["stream_options"] = map[string]interface{}{"include_usage": true}
|
||||
|
||||
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 := b.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
|
||||
sawTerminal := false
|
||||
accumulatedToolCalls := make(map[int]map[string]any)
|
||||
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
|
||||
}
|
||||
|
||||
delta, ok := firstChoice["delta"].(map[string]interface{})
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
|
||||
accumulateToolCallDeltas(delta, accumulatedToolCalls)
|
||||
|
||||
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 b.baseModel.doStreamRequest(ctx, url, apiConfig, 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("baichuan: stream ended before [DONE] or finish_reason")
|
||||
}
|
||||
|
||||
setSortedToolCallsResult(modelConfig, accumulatedToolCalls)
|
||||
|
||||
// Send [DONE] marker for OpenAI compatibility
|
||||
endOfStream := "[DONE]"
|
||||
return sender(&endOfStream, nil)
|
||||
}
|
||||
|
||||
func (b *BaichuanModel) Embed(ctx context.Context, modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig, modelUsage *common.ModelUsage) ([]EmbeddingData, error) {
|
||||
|
||||
@@ -50,37 +50,6 @@ func (b *BaiduModel) Name() string {
|
||||
return "BaiduYiyan"
|
||||
}
|
||||
|
||||
type BaiduChatResponse struct {
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
Created int64 `json:"created"`
|
||||
Model string `json:"model"`
|
||||
Choices []struct {
|
||||
Index int `json:"index"`
|
||||
Message struct {
|
||||
Role string `json:"role"`
|
||||
Content *string `json:"content"`
|
||||
ToolCalls []map[string]any `json:"tool_calls"`
|
||||
ReasoningContent *string `json:"reasoning_content"`
|
||||
} `json:"message"`
|
||||
FinishReason string `json:"finish_reason"`
|
||||
} `json:"choices"`
|
||||
Usage struct {
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
PromptTokensDetails struct {
|
||||
SearchTokens int `json:"search_tokens"`
|
||||
CachedTokens int `json:"cached_tokens"`
|
||||
} `json:"prompt_tokens_details"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
TotalTokens int `json:"total_tokens"`
|
||||
} `json:"usage"`
|
||||
SearchResult struct {
|
||||
Index int `json:"index"`
|
||||
URL string `json:"url"`
|
||||
Title string `json:"title"`
|
||||
} `json:"search_result"`
|
||||
}
|
||||
|
||||
func (b *BaiduModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) {
|
||||
if err := b.baseModel.APIConfigCheck(apiConfig); err != nil {
|
||||
return nil, err
|
||||
@@ -142,82 +111,12 @@ func (b *BaiduModel) ChatWithMessages(ctx context.Context, modelName string, mes
|
||||
}
|
||||
}
|
||||
|
||||
jsonData, err := json.Marshal(reqBody)
|
||||
body, err := b.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 := b.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))
|
||||
}
|
||||
|
||||
return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, chatConfig *ChatConfig) (chatResponseParts, error) {
|
||||
var result BaiduChatResponse
|
||||
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("no choices in response")
|
||||
}
|
||||
|
||||
choice := &result.Choices[0]
|
||||
content := ""
|
||||
if choice.Message.Content != nil {
|
||||
content = *choice.Message.Content
|
||||
}
|
||||
reasonContent := ""
|
||||
if chatConfig != nil && chatConfig.Thinking != nil && *chatConfig.Thinking {
|
||||
if choice.Message.ReasoningContent != nil {
|
||||
reasonContent = *choice.Message.ReasoningContent
|
||||
}
|
||||
if reasonContent != "" && reasonContent[0] == '\n' {
|
||||
reasonContent = reasonContent[1:]
|
||||
}
|
||||
}
|
||||
|
||||
totalTokens := result.Usage.TotalTokens
|
||||
if totalTokens == 0 {
|
||||
totalTokens = result.Usage.PromptTokens + result.Usage.CompletionTokens
|
||||
}
|
||||
var usage *TokenUsage
|
||||
if totalTokens > 0 {
|
||||
usage = &TokenUsage{
|
||||
PromptTokens: result.Usage.PromptTokens,
|
||||
CompletionTokens: result.Usage.CompletionTokens,
|
||||
TotalTokens: totalTokens,
|
||||
}
|
||||
}
|
||||
|
||||
return chatResponseParts{
|
||||
RequestID: result.ID,
|
||||
Content: &content,
|
||||
ReasonContent: &reasonContent,
|
||||
ToolCalls: choice.Message.ToolCalls,
|
||||
Usage: usage,
|
||||
}, nil
|
||||
})
|
||||
return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig)
|
||||
}
|
||||
|
||||
func (b *BaiduModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error {
|
||||
@@ -237,6 +136,9 @@ func (b *BaiduModel) ChatStreamlyWithSender(ctx context.Context, modelName strin
|
||||
|
||||
// Build request body with streaming enabled
|
||||
reqBody := buildRequestBody(modelConfig, modelName, messages, true)
|
||||
reqBody["stream_options"] = map[string]interface{}{
|
||||
"include_usage": true,
|
||||
}
|
||||
if modelConfig != nil && modelConfig.Thinking != nil {
|
||||
lowerModelName := strings.ToLower(modelName)
|
||||
|
||||
@@ -278,101 +180,9 @@ func (b *BaiduModel) 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", *apiConfig.ApiKey))
|
||||
|
||||
resp, err := b.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
|
||||
sawTerminal := false
|
||||
accumulatedToolCalls := make(map[int]map[string]interface{})
|
||||
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
|
||||
}
|
||||
|
||||
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 b.baseModel.doStreamRequest(ctx, url, apiConfig, 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("baidu: stream ended before [DONE] or finish_reason")
|
||||
}
|
||||
|
||||
setSortedToolCallsResult(modelConfig, accumulatedToolCalls)
|
||||
|
||||
// Send [DONE] marker for OpenAI compatibility
|
||||
endOfStream := "[DONE]"
|
||||
if err = sender(&endOfStream, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
type baiduEmbeddingResponse struct {
|
||||
|
||||
@@ -214,6 +214,7 @@ func (b *BaseModel) doStreamRequest(ctx context.Context, url string, apiConfig *
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
req.Header.Set("Accept", "text/event-stream")
|
||||
|
||||
resp, err := b.httpClient.Do(req)
|
||||
if err != nil {
|
||||
|
||||
@@ -94,114 +94,11 @@ func (c *CometAPIModel) balanceURL(apiKey string) string {
|
||||
return parsed.String()
|
||||
}
|
||||
|
||||
func newCometAPIJSONRequest(ctx context.Context, method string, endpoint string, payload interface{}, apiKey string) (*http.Request, error) {
|
||||
jsonData, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal request: %w", err)
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, method, endpoint, bytes.NewBuffer(jsonData))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if apiKey != "" {
|
||||
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", apiKey))
|
||||
}
|
||||
return req, nil
|
||||
}
|
||||
|
||||
type cometapiHTTPResponse struct {
|
||||
StatusCode int
|
||||
Status string
|
||||
Body []byte
|
||||
}
|
||||
|
||||
func (c *CometAPIModel) doCometAPIRequest(req *http.Request) (*cometapiHTTPResponse, error) {
|
||||
resp, err := c.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)
|
||||
}
|
||||
|
||||
return &cometapiHTTPResponse{
|
||||
StatusCode: resp.StatusCode,
|
||||
Status: resp.Status,
|
||||
Body: body,
|
||||
}, nil
|
||||
}
|
||||
|
||||
type cometapiChatResponsePayload struct {
|
||||
Choices []cometapiChatChoice `json:"choices"`
|
||||
}
|
||||
|
||||
type cometapiChatChoice struct {
|
||||
Message cometapiChatMessage `json:"message"`
|
||||
Delta cometapiChatDelta `json:"delta"`
|
||||
FinishReason string `json:"finish_reason"`
|
||||
}
|
||||
|
||||
type cometapiChatMessage struct {
|
||||
Content *string `json:"content"`
|
||||
ReasoningContent string `json:"reasoning_content"`
|
||||
}
|
||||
|
||||
type cometapiChatDelta struct {
|
||||
Content string `json:"content"`
|
||||
ReasoningContent string `json:"reasoning_content"`
|
||||
}
|
||||
|
||||
func parseCometAPIChatResponse(body []byte) (*ChatResponse, error) {
|
||||
var parsed cometapiChatResponsePayload
|
||||
if err := json.Unmarshal(body, &parsed); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse response: %w", err)
|
||||
}
|
||||
if len(parsed.Choices) == 0 {
|
||||
return nil, fmt.Errorf("no choices in response")
|
||||
}
|
||||
if parsed.Choices[0].Message.Content == nil {
|
||||
return nil, fmt.Errorf("invalid content format")
|
||||
}
|
||||
|
||||
content := *parsed.Choices[0].Message.Content
|
||||
reasonContent := strings.TrimLeft(parsed.Choices[0].Message.ReasoningContent, "\n")
|
||||
return &ChatResponse{
|
||||
Answer: &content,
|
||||
ReasonContent: &reasonContent,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func parseCometAPIStreamEvent(data string) (content string, reasonContent string, terminal bool, ok bool) {
|
||||
var event cometapiChatResponsePayload
|
||||
if err := json.Unmarshal([]byte(data), &event); err != nil {
|
||||
return "", "", false, false
|
||||
}
|
||||
if len(event.Choices) == 0 {
|
||||
return "", "", false, false
|
||||
}
|
||||
choice := event.Choices[0]
|
||||
return choice.Delta.Content, choice.Delta.ReasoningContent, choice.FinishReason != "", true
|
||||
}
|
||||
|
||||
type cometapiModelCatalogResponse struct {
|
||||
Data []cometapiModelCatalogItem `json:"data"`
|
||||
}
|
||||
|
||||
type cometapiModelCatalogItem struct {
|
||||
ID string `json:"id"`
|
||||
}
|
||||
|
||||
// ChatWithMessages sends multiple messages with roles and returns the response.
|
||||
func (c *CometAPIModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) {
|
||||
if err := c.baseModel.APIConfigCheck(apiConfig); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
apiKey := *apiConfig.ApiKey
|
||||
if err := validateCometAPIModelName(modelName); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -217,22 +114,20 @@ func (c *CometAPIModel) ChatWithMessages(ctx context.Context, modelName string,
|
||||
|
||||
reqBody := buildRequestBody(chatModelConfig, modelName, messages, false)
|
||||
|
||||
ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout)
|
||||
defer cancel()
|
||||
|
||||
req, err := newCometAPIJSONRequest(ctx, "POST", url, reqBody, apiKey)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
if chatModelConfig != nil && chatModelConfig.Thinking != nil {
|
||||
if *chatModelConfig.Thinking {
|
||||
reqBody["thinking"] = map[string]interface{}{"type": "enabled"}
|
||||
} else {
|
||||
reqBody["thinking"] = map[string]interface{}{"type": "disabled"}
|
||||
}
|
||||
}
|
||||
resp, err := c.doCometAPIRequest(req)
|
||||
|
||||
body, err := c.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(resp.Body))
|
||||
}
|
||||
return parseCometAPIChatResponse(resp.Body)
|
||||
return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig)
|
||||
}
|
||||
|
||||
// ChatStreamlyWithSender sends messages and streams the response
|
||||
@@ -253,8 +148,6 @@ func (c *CometAPIModel) ChatStreamlyWithSender(ctx context.Context, modelName st
|
||||
return fmt.Errorf("messages is empty")
|
||||
}
|
||||
|
||||
apiKey := *apiConfig.ApiKey
|
||||
|
||||
url, err := c.endpointURL(cometapiRegion(apiConfig), c.baseModel.URLSuffix.Chat)
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -271,76 +164,20 @@ func (c *CometAPIModel) ChatStreamlyWithSender(ctx context.Context, modelName st
|
||||
}
|
||||
reqBody := buildRequestBody(chatModelConfig, modelName, messages, true)
|
||||
|
||||
req, err := newCometAPIJSONRequest(ctx, "POST", url, reqBody, apiKey)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
resp, err := c.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
|
||||
if chatModelConfig != nil {
|
||||
if chatModelConfig.Thinking != nil {
|
||||
if *chatModelConfig.Thinking {
|
||||
reqBody["thinking"] = map[string]interface{}{"type": "enabled"}
|
||||
} else {
|
||||
reqBody["thinking"] = map[string]interface{}{"type": "disabled"}
|
||||
}
|
||||
}
|
||||
}
|
||||
reqBody["stream_options"] = map[string]interface{}{"include_usage": true}
|
||||
|
||||
content, ok := delta["content"].(string)
|
||||
if 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 c.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)
|
||||
}
|
||||
if !done && !sawTerminal {
|
||||
return fmt.Errorf("cometapi: stream ended before [DONE] or finish_reason")
|
||||
}
|
||||
|
||||
setSortedToolCallsResult(chatModelConfig, accumulatedToolCalls)
|
||||
|
||||
endOfStream := "[DONE]"
|
||||
if err := sender(&endOfStream, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
type cometapiEmbeddingData struct {
|
||||
@@ -355,12 +192,6 @@ type cometapiEmbeddingResponse struct {
|
||||
Object string `json:"object"`
|
||||
}
|
||||
|
||||
type cometapiEmbeddingRequest struct {
|
||||
Model string `json:"model"`
|
||||
Input []string `json:"input"`
|
||||
Dimensions int `json:"dimensions,omitempty"`
|
||||
}
|
||||
|
||||
// Embed turns a list of texts into embedding vectors
|
||||
func (c *CometAPIModel) Embed(ctx context.Context, modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig, modelUsage *common.ModelUsage) ([]EmbeddingData, error) {
|
||||
if err := c.baseModel.APIConfigCheck(apiConfig); err != nil {
|
||||
@@ -371,8 +202,6 @@ func (c *CometAPIModel) Embed(ctx context.Context, modelName *string, texts []st
|
||||
return []EmbeddingData{}, nil
|
||||
}
|
||||
|
||||
apiKey := *apiConfig.ApiKey
|
||||
|
||||
if modelName == nil || strings.TrimSpace(*modelName) == "" {
|
||||
return nil, fmt.Errorf("model name is required")
|
||||
}
|
||||
@@ -382,33 +211,21 @@ func (c *CometAPIModel) Embed(ctx context.Context, modelName *string, texts []st
|
||||
return nil, err
|
||||
}
|
||||
|
||||
reqBody := cometapiEmbeddingRequest{
|
||||
Model: *modelName,
|
||||
Input: texts,
|
||||
reqBody := map[string]any{
|
||||
"model": *modelName,
|
||||
"input": texts,
|
||||
}
|
||||
if embeddingConfig != nil && embeddingConfig.Dimension > 0 {
|
||||
reqBody.Dimensions = embeddingConfig.Dimension
|
||||
reqBody["dimensions"] = embeddingConfig.Dimension
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout)
|
||||
defer cancel()
|
||||
|
||||
req, err := newCometAPIJSONRequest(ctx, "POST", url, reqBody, apiKey)
|
||||
body, err := c.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
resp, err := c.doCometAPIRequest(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("CometAPI embeddings API error: %s, body: %s", resp.Status, string(resp.Body))
|
||||
}
|
||||
|
||||
var parsed cometapiEmbeddingResponse
|
||||
if err = json.Unmarshal(resp.Body, &parsed); err != nil {
|
||||
if err = json.Unmarshal(body, &parsed); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse response: %w", err)
|
||||
}
|
||||
|
||||
@@ -450,19 +267,28 @@ func (c *CometAPIModel) ListModels(ctx context.Context, apiConfig *APIConfig) ([
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
if apiConfig != nil && apiConfig.ApiKey != nil {
|
||||
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
|
||||
}
|
||||
|
||||
resp, err := c.doCometAPIRequest(req)
|
||||
resp, err := c.baseModel.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
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(resp.Body))
|
||||
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
// Parse response
|
||||
var modelList ModelList
|
||||
if err = json.Unmarshal(resp.Body, &modelList); err != nil {
|
||||
if err = json.Unmarshal(body, &modelList); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse response: %w", err)
|
||||
}
|
||||
return ParseListModel(modelList), nil
|
||||
@@ -485,17 +311,23 @@ func (c *CometAPIModel) Balance(ctx context.Context, apiConfig *APIConfig) (map[
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
|
||||
resp, err := c.doCometAPIRequest(req)
|
||||
resp, err := c.baseModel.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
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("CometAPI quota API error: %s, body: %s", resp.Status, string(resp.Body))
|
||||
return nil, fmt.Errorf("CometAPI quota API error: %s, body: %s", resp.Status, string(body))
|
||||
}
|
||||
|
||||
var result map[string]interface{}
|
||||
if err = json.Unmarshal(resp.Body, &result); err != nil {
|
||||
if err = json.Unmarshal(body, &result); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse response: %w", err)
|
||||
}
|
||||
|
||||
|
||||
@@ -801,7 +801,7 @@ func TestCometAPIEmbedRejectsHTTPError(t *testing.T) {
|
||||
apiKey := "test-key"
|
||||
model := "text-embedding-3-small"
|
||||
_, err := m.Embed(ctx, &model, []string{"a"}, &APIConfig{ApiKey: &apiKey}, nil, nil)
|
||||
if err == nil || !strings.Contains(err.Error(), "CometAPI embeddings API error") {
|
||||
t.Errorf("expected CometAPI embeddings API error, got %v", err)
|
||||
if err == nil || !strings.Contains(err.Error(), "status 401") {
|
||||
t.Errorf("expected embeddings API error for HTTP 401, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -88,73 +88,12 @@ func (d *DeepInfraModel) ChatWithMessages(ctx context.Context, modelName string,
|
||||
}
|
||||
}
|
||||
|
||||
jsonData, err := json.Marshal(reqBody)
|
||||
body, err := d.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 := d.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)
|
||||
}
|
||||
|
||||
// Parse result
|
||||
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 || 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 rc, ok := messageMap["reasoning_content"].(string); ok {
|
||||
reasonContent = rc
|
||||
}
|
||||
|
||||
chatResponse := &ChatResponse{
|
||||
Answer: &content,
|
||||
}
|
||||
if reasonContent != "" {
|
||||
chatResponse.ReasonContent = &reasonContent
|
||||
}
|
||||
|
||||
return chatResponse, nil
|
||||
return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig)
|
||||
}
|
||||
|
||||
func (d *DeepInfraModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error {
|
||||
@@ -189,86 +128,11 @@ func (d *DeepInfraModel) ChatStreamlyWithSender(ctx context.Context, modelName s
|
||||
}
|
||||
}
|
||||
|
||||
jsonData, err := json.Marshal(reqBody)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal request: %w", err)
|
||||
}
|
||||
reqBody["stream_options"] = map[string]any{"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 := d.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 {
|
||||
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
|
||||
}
|
||||
}
|
||||
|
||||
finishReason, ok := firstChoice["finish_reason"].(string)
|
||||
if ok && finishReason != "" {
|
||||
sawTerminal = true
|
||||
}
|
||||
return nil
|
||||
return d.baseModel.doStreamRequest(ctx, url, apiConfig, 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)
|
||||
}
|
||||
setSortedToolCallsResult(modelConfig, accumulatedToolCalls)
|
||||
if !done && !sawTerminal {
|
||||
return fmt.Errorf("deepinfra: stream ended before [DONE] or finish_reason")
|
||||
}
|
||||
|
||||
// Send [DONE] marker for OpenAI compatibility
|
||||
endOfStream := "[DONE]"
|
||||
return sender(&endOfStream, nil)
|
||||
}
|
||||
|
||||
func (d *DeepInfraModel) Embed(ctx context.Context, modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig, modelUsage *common.ModelUsage) ([]EmbeddingData, error) {
|
||||
|
||||
@@ -17,12 +17,9 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"ragflow/internal/common"
|
||||
"strings"
|
||||
)
|
||||
@@ -60,54 +57,6 @@ func (f *FuturMixModel) endpointURL(region, suffix string) (string, error) {
|
||||
return fmt.Sprintf("%s/%s", baseURL, strings.TrimLeft(suffix, "/")), nil
|
||||
}
|
||||
|
||||
func futurmixRegion(apiConfig *APIConfig) string {
|
||||
if apiConfig != nil && apiConfig.Region != nil && *apiConfig.Region != "" {
|
||||
return *apiConfig.Region
|
||||
}
|
||||
return "default"
|
||||
}
|
||||
|
||||
func newFuturMixJSONRequest(ctx context.Context, method, endpoint string, payload interface{}, apiKey string) (*http.Request, error) {
|
||||
var body io.Reader
|
||||
if payload != nil {
|
||||
jsonData, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal request: %w", err)
|
||||
}
|
||||
body = bytes.NewBuffer(jsonData)
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, method, endpoint, body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if apiKey != "" {
|
||||
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", apiKey))
|
||||
}
|
||||
return req, nil
|
||||
}
|
||||
|
||||
type futurmixChatChoice struct {
|
||||
Message futurmixChatMessage `json:"message"`
|
||||
Delta futurmixChatDelta `json:"delta"`
|
||||
FinishReason string `json:"finish_reason"`
|
||||
}
|
||||
|
||||
type futurmixChatMessage struct {
|
||||
Content *string `json:"content"`
|
||||
ReasoningContent string `json:"reasoning_content"`
|
||||
}
|
||||
|
||||
type futurmixChatDelta struct {
|
||||
Content string `json:"content"`
|
||||
ReasoningContent string `json:"reasoning_content"`
|
||||
}
|
||||
|
||||
type futurmixChatResponse struct {
|
||||
Choices []futurmixChatChoice `json:"choices"`
|
||||
}
|
||||
|
||||
// ChatWithMessages sends a non-streaming chat completion
|
||||
func (f *FuturMixModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) {
|
||||
if err := f.baseModel.APIConfigCheck(apiConfig); err != nil {
|
||||
@@ -125,55 +74,24 @@ func (f *FuturMixModel) ChatWithMessages(ctx context.Context, modelName string,
|
||||
url := fmt.Sprintf("%s/%s", resolvedBaseURL, f.baseModel.URLSuffix.Chat)
|
||||
|
||||
reqBody := buildRequestBody(chatModelConfig, modelName, messages, false)
|
||||
if chatModelConfig != nil && chatModelConfig.Thinking != nil {
|
||||
if *chatModelConfig.Thinking {
|
||||
reqBody["thinking"] = map[string]any{
|
||||
"type": "enabled",
|
||||
}
|
||||
} else {
|
||||
reqBody["thinking"] = map[string]any{
|
||||
"type": "disabled",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
jsonData, err := json.Marshal(reqBody)
|
||||
body, err := f.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 := f.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("futurmix chat API error: %s, body: %s", resp.Status, string(body))
|
||||
}
|
||||
|
||||
var parsed futurmixChatResponse
|
||||
if err := json.Unmarshal(body, &parsed); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse response: %w", err)
|
||||
}
|
||||
if len(parsed.Choices) == 0 {
|
||||
return nil, fmt.Errorf("no choices in response")
|
||||
}
|
||||
if parsed.Choices[0].Message.Content == nil {
|
||||
return nil, fmt.Errorf("invalid content format")
|
||||
}
|
||||
|
||||
content := *parsed.Choices[0].Message.Content
|
||||
reasonContent := parsed.Choices[0].Message.ReasoningContent
|
||||
return &ChatResponse{
|
||||
Answer: &content,
|
||||
ReasonContent: &reasonContent,
|
||||
}, nil
|
||||
return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig)
|
||||
}
|
||||
|
||||
// ChatStreamlyWithSender sends a streaming chat completion
|
||||
@@ -192,77 +110,25 @@ func (f *FuturMixModel) ChatStreamlyWithSender(ctx context.Context, modelName st
|
||||
}
|
||||
url := fmt.Sprintf("%s/%s", resolvedBaseURL, f.baseModel.URLSuffix.Chat)
|
||||
|
||||
reqBody := buildRequestBody(chatModelConfig, modelName, messages, false)
|
||||
|
||||
jsonData, err := json.Marshal(reqBody)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal request: %w", err)
|
||||
reqBody := buildRequestBody(chatModelConfig, modelName, messages, true)
|
||||
reqBody["stream_options"] = map[string]any{
|
||||
"include_usage": true,
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout)
|
||||
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 := f.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("futurmix chat stream API error: %s, body: %s", resp.Status, string(body))
|
||||
}
|
||||
|
||||
accumulatedToolCalls := make(map[int]map[string]any)
|
||||
if _, 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
|
||||
if chatModelConfig != nil && chatModelConfig.Thinking != nil {
|
||||
if *chatModelConfig.Thinking {
|
||||
reqBody["thinking"] = map[string]any{
|
||||
"type": "enabled",
|
||||
}
|
||||
} else {
|
||||
reqBody["thinking"] = map[string]any{
|
||||
"type": "disabled",
|
||||
}
|
||||
}
|
||||
|
||||
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)
|
||||
|
||||
endOfStream := "[DONE]"
|
||||
return sender(&endOfStream, nil)
|
||||
return f.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error {
|
||||
return HandleStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig, sender)
|
||||
})
|
||||
}
|
||||
|
||||
// Embed is not exposed by the FuturMix API per the public docs.
|
||||
|
||||
@@ -191,12 +191,17 @@ func extractContentAndChoices(result map[string]any) (*string, *string, []map[st
|
||||
}
|
||||
|
||||
var reasonContent *string
|
||||
if rc, ok := messageMap["reasoning_content"].(string); ok {
|
||||
if rc, ok := messageMap["reasoning_content"].(string); ok && rc != "" {
|
||||
reason := rc
|
||||
if reason != "" && reason[0] == '\n' {
|
||||
if reason[0] == '\n' {
|
||||
reason = reason[1:]
|
||||
}
|
||||
reasonContent = &reason
|
||||
} else if rc, ok := messageMap["reasoning"].(string); ok && rc != "" {
|
||||
// Some providers (e.g. Avian) report reasoning under a top-level
|
||||
// "reasoning" field instead of "reasoning_content".
|
||||
reason := rc
|
||||
reasonContent = &reason
|
||||
} else {
|
||||
// Always return a non-nil pointer so callers can rely on it.
|
||||
empty := ""
|
||||
|
||||
Reference in New Issue
Block a user