mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-13 08:28:22 +08:00
### Summary Closes #15381 Every provider in `internal/entity/models/` reads its streaming response with `bufio.NewScanner(resp.Body)` and iterates over `scanner.Scan()`. The default `bufio.Scanner` maximum token size is 64KB, so when an upstream sends a single SSE `data:` line larger than 64KB (long content deltas, large tool or function call argument blobs, bundled `reasoning_content`, or providers that emit a whole message in one event) `scanner.Scan()` returns `false` and `scanner.Err()` returns `bufio.ErrTooLong`. Streaming chat then ends with an error partway through the response. This change adds `scanner.Buffer(make([]byte, 64*1024), 1024*1024)` immediately after every SSE scanner that was still bare, raising the cap to 1MB. 1MB is the value already used for streaming chat in `openai.go`, `modelscope.go`, `groq.go`, `mistral.go`, `xai.go` and the other already patched providers (the 8MB cap in the repo is reserved for TTS and embedding paths), so this simply converges the remaining providers onto the established pattern. Nothing else changes: line parsing, `data:` prefix handling, `[DONE]` detection, JSON unmarshalling, error handling, and the existing `scanner.Err()` checks all stay the same. Providers covered (23 scanners across 22 files): 302ai, aliyun, baichuan, baidu, cohere, deepinfra, deepseek, gitee, huggingface, lmstudio, minimax (the chat scanner, whose TTS scanner was already bumped), moonshot, nvidia, ollama, openrouter, orcarouter, paddleocr, siliconflow, tokenhub, vllm, volcengine, xunfei, zhipu-ai. `jiekouai.go` is excluded because it is covered by the in flight #15337. A table driven regression test (`sse_scanner_buffer_test.go`) streams a single 128KB `data:` content delta followed by `data: [DONE]` through an `httptest` server and asserts that `ChatStreamlyWithSender` delivers the full content with no error across a representative subset of providers. Without the buffer fix the test fails with `bufio.Scanner: token too long`. This PR also removes three duplicate declarations of the package level `roundTripperFunc` test helper that several recently merged provider PRs each added independently, which had left the `internal/entity/models` test package unable to compile. The helper now lives in a single place and is shared. ### Type of change - [x] Bug Fix (non-breaking change which fixes an issue)
803 lines
21 KiB
Go
803 lines
21 KiB
Go
package models
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"context"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"ragflow/internal/common"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
type BaiduModel struct {
|
|
BaseURL map[string]string
|
|
URLSuffix URLSuffix
|
|
httpClient *http.Client
|
|
}
|
|
|
|
func (b *BaiduModel) NewInstance(baseURL map[string]string) ModelDriver {
|
|
return &BaiduModel{
|
|
BaseURL: baseURL,
|
|
URLSuffix: b.URLSuffix,
|
|
httpClient: &http.Client{
|
|
Timeout: 120 * time.Second,
|
|
Transport: &http.Transport{
|
|
MaxIdleConns: 10,
|
|
MaxIdleConnsPerHost: 100,
|
|
IdleConnTimeout: 90 * time.Second,
|
|
DisableCompression: false,
|
|
},
|
|
},
|
|
}
|
|
}
|
|
|
|
func NewBaiduModel(baseURL map[string]string, urlSuffix URLSuffix) *BaiduModel {
|
|
return &BaiduModel{
|
|
BaseURL: baseURL,
|
|
URLSuffix: urlSuffix,
|
|
httpClient: &http.Client{
|
|
Timeout: 120 * time.Second,
|
|
Transport: &http.Transport{
|
|
MaxConnsPerHost: 10,
|
|
MaxIdleConnsPerHost: 100,
|
|
IdleConnTimeout: 90 * time.Second,
|
|
DisableCompression: false,
|
|
},
|
|
},
|
|
}
|
|
}
|
|
|
|
func (b *BaiduModel) Name() string {
|
|
return "baidu"
|
|
}
|
|
|
|
func (b *BaiduModel) ChatWithMessages(modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig) (*ChatResponse, error) {
|
|
if apiConfig == nil || apiConfig.ApiKey == nil || *apiConfig.ApiKey == "" {
|
|
return nil, fmt.Errorf("api key is nil or empty")
|
|
}
|
|
if len(messages) == 0 {
|
|
return nil, fmt.Errorf("messages is empty")
|
|
}
|
|
|
|
var region = "default"
|
|
if apiConfig.Region != nil && *apiConfig.Region != "" {
|
|
region = *apiConfig.Region
|
|
}
|
|
|
|
url := fmt.Sprintf("%s/%s", b.BaseURL[region], b.URLSuffix.Chat)
|
|
|
|
// Convert messages to API format
|
|
apiMessages := make([]map[string]interface{}, len(messages))
|
|
for i, msg := range messages {
|
|
apiMessages[i] = map[string]interface{}{
|
|
"role": msg.Role,
|
|
"content": msg.Content,
|
|
}
|
|
}
|
|
|
|
// Build request body
|
|
reqBody := map[string]interface{}{
|
|
"model": modelName,
|
|
"messages": apiMessages,
|
|
"stream": false,
|
|
"temperature": 1,
|
|
}
|
|
|
|
if chatModelConfig != nil {
|
|
if chatModelConfig.Stream != nil {
|
|
reqBody["stream"] = *chatModelConfig.Stream
|
|
}
|
|
|
|
if chatModelConfig.MaxTokens != nil {
|
|
reqBody["max_tokens"] = *chatModelConfig.MaxTokens
|
|
}
|
|
|
|
if chatModelConfig.Temperature != nil {
|
|
reqBody["temperature"] = *chatModelConfig.Temperature
|
|
}
|
|
|
|
if chatModelConfig.TopP != nil {
|
|
reqBody["top_p"] = *chatModelConfig.TopP
|
|
}
|
|
|
|
if chatModelConfig.Stop != nil {
|
|
reqBody["stop"] = *chatModelConfig.Stop
|
|
}
|
|
|
|
if chatModelConfig.Thinking != nil {
|
|
lowerModelName := strings.ToLower(modelName)
|
|
|
|
// `enable_think` for qwen and erine
|
|
if strings.HasPrefix(lowerModelName, "qwen") || strings.HasPrefix(lowerModelName, "ernie") {
|
|
reqBody["enable_thinking"] = *chatModelConfig.Thinking
|
|
} else {
|
|
if *chatModelConfig.Thinking {
|
|
thinkingFlag := "enabled"
|
|
|
|
if strings.Contains(lowerModelName, "deepseek-v4") {
|
|
effort := "high"
|
|
if chatModelConfig.Effort != nil {
|
|
effort = *chatModelConfig.Effort
|
|
}
|
|
switch effort {
|
|
case "none", "low", "medium":
|
|
thinkingFlag = "disabled"
|
|
case "high", "default":
|
|
thinkingFlag = "enabled"
|
|
reqBody["reasoning_effort"] = "high"
|
|
case "max":
|
|
thinkingFlag = "enabled"
|
|
reqBody["reasoning_effort"] = "max"
|
|
default:
|
|
thinkingFlag = "enabled"
|
|
reqBody["reasoning_effort"] = effort
|
|
}
|
|
}
|
|
|
|
reqBody["thinking"] = map[string]interface{}{
|
|
"type": thinkingFlag,
|
|
}
|
|
} else {
|
|
reqBody["thinking"] = map[string]interface{}{
|
|
"type": "disabled",
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
jsonData, err := json.Marshal(reqBody)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to marshal request: %w", err)
|
|
}
|
|
|
|
req, err := http.NewRequest("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.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 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 chatModelConfig != nil && chatModelConfig.Thinking != nil && *chatModelConfig.Thinking {
|
|
reasonContent, ok = messageMap["reasoning_content"].(string)
|
|
if !ok {
|
|
return nil, fmt.Errorf("invalid reasoning 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,
|
|
}
|
|
|
|
return chatResponse, nil
|
|
}
|
|
|
|
func (b *BaiduModel) ChatStreamlyWithSender(modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, sender func(*string, *string) error) error {
|
|
if len(messages) == 0 {
|
|
return fmt.Errorf("messages is empty")
|
|
}
|
|
|
|
var region = "default"
|
|
if apiConfig != nil && apiConfig.Region != nil && *apiConfig.Region != "" {
|
|
region = *apiConfig.Region
|
|
}
|
|
|
|
url := fmt.Sprintf("%s/%s", strings.TrimSuffix(b.BaseURL[region], "/"), b.URLSuffix.Chat)
|
|
|
|
// Convert messages to API format
|
|
apiMessages := make([]map[string]interface{}, len(messages))
|
|
for i, msg := range messages {
|
|
apiMessages[i] = map[string]interface{}{
|
|
"role": msg.Role,
|
|
"content": msg.Content,
|
|
}
|
|
}
|
|
|
|
// Build request body with streaming enabled
|
|
reqBody := map[string]interface{}{
|
|
"model": modelName,
|
|
"messages": apiMessages,
|
|
"stream": true,
|
|
"temperature": 1,
|
|
}
|
|
|
|
if modelConfig != nil {
|
|
if modelConfig.Stream != nil {
|
|
reqBody["stream"] = *modelConfig.Stream
|
|
}
|
|
|
|
if modelConfig.MaxTokens != nil {
|
|
reqBody["max_tokens"] = *modelConfig.MaxTokens
|
|
}
|
|
|
|
if modelConfig.Temperature != nil {
|
|
reqBody["temperature"] = *modelConfig.Temperature
|
|
}
|
|
|
|
if modelConfig.DoSample != nil {
|
|
reqBody["do_sample"] = *modelConfig.DoSample
|
|
}
|
|
|
|
if modelConfig.TopP != nil {
|
|
reqBody["top_p"] = *modelConfig.TopP
|
|
}
|
|
|
|
if modelConfig.Stop != nil {
|
|
reqBody["stop"] = *modelConfig.Stop
|
|
}
|
|
|
|
if modelConfig.Thinking != nil {
|
|
lowerModelName := strings.ToLower(modelName)
|
|
|
|
// `enable_think` for qwen and erine
|
|
if strings.HasPrefix(lowerModelName, "qwen") || strings.HasPrefix(lowerModelName, "ernie") {
|
|
reqBody["enable_thinking"] = *modelConfig.Thinking
|
|
} else {
|
|
if *modelConfig.Thinking {
|
|
thinkingFlag := "enabled"
|
|
|
|
if strings.Contains(lowerModelName, "deepseek-v4") {
|
|
effort := "high"
|
|
if modelConfig.Effort != nil {
|
|
effort = *modelConfig.Effort
|
|
}
|
|
switch effort {
|
|
case "none", "low", "medium":
|
|
thinkingFlag = "disabled"
|
|
case "high", "default":
|
|
thinkingFlag = "enabled"
|
|
reqBody["reasoning_effort"] = "high"
|
|
case "max":
|
|
thinkingFlag = "enabled"
|
|
reqBody["reasoning_effort"] = "max"
|
|
default:
|
|
thinkingFlag = "enabled"
|
|
reqBody["reasoning_effort"] = effort
|
|
}
|
|
}
|
|
|
|
reqBody["thinking"] = map[string]interface{}{
|
|
"type": thinkingFlag,
|
|
}
|
|
} else {
|
|
reqBody["thinking"] = map[string]interface{}{
|
|
"type": "disabled",
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
jsonData, err := json.Marshal(reqBody)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to marshal request: %w", err)
|
|
}
|
|
|
|
req, err := http.NewRequest("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.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
|
|
scanner := bufio.NewScanner(resp.Body)
|
|
scanner.Buffer(make([]byte, 64*1024), 1024*1024)
|
|
for scanner.Scan() {
|
|
line := scanner.Text()
|
|
common.Info(line)
|
|
|
|
// SSE data line starts with "data:"
|
|
if !strings.HasPrefix(line, "data:") {
|
|
continue
|
|
}
|
|
|
|
// Extract JSON after "data:"
|
|
data := strings.TrimSpace(line[5:])
|
|
|
|
// [DONE] marks the end of stream
|
|
if data == "[DONE]" {
|
|
break
|
|
}
|
|
|
|
// Parse the JSON event
|
|
var event map[string]interface{}
|
|
if err = json.Unmarshal([]byte(data), &event); err != nil {
|
|
continue
|
|
}
|
|
|
|
choices, ok := event["choices"].([]interface{})
|
|
if !ok || len(choices) == 0 {
|
|
continue
|
|
}
|
|
|
|
firstChoice, ok := choices[0].(map[string]interface{})
|
|
if !ok {
|
|
continue
|
|
}
|
|
|
|
delta, ok := firstChoice["delta"].(map[string]interface{})
|
|
if !ok {
|
|
continue
|
|
}
|
|
|
|
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 != "" {
|
|
break
|
|
}
|
|
}
|
|
|
|
// Send [DONE] marker for OpenAI compatibility
|
|
endOfStream := "[DONE]"
|
|
if err = sender(&endOfStream, nil); err != nil {
|
|
return err
|
|
}
|
|
|
|
return scanner.Err()
|
|
}
|
|
|
|
type baiduEmbeddingResponse struct {
|
|
ID string `json:"id"`
|
|
Object string `json:"object"`
|
|
Created int64 `json:"created"`
|
|
Data []baiduEmbeddingData `json:"data"`
|
|
Model string `json:"model"`
|
|
Usage baiduUsage `json:"usage"`
|
|
}
|
|
|
|
type baiduEmbeddingData struct {
|
|
Object string `json:"object"`
|
|
Embedding []float64 `json:"embedding"`
|
|
Index *int `json:"index"`
|
|
}
|
|
|
|
type baiduUsage struct {
|
|
PromptTokens int `json:"prompt_tokens"`
|
|
TotalTokens int `json:"total_tokens"`
|
|
}
|
|
|
|
func (b *BaiduModel) Embed(modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig) ([]EmbeddingData, error) {
|
|
if apiConfig == nil || apiConfig.ApiKey == nil || *apiConfig.ApiKey == "" {
|
|
return nil, fmt.Errorf("api key is required")
|
|
}
|
|
if modelName == nil || *modelName == "" {
|
|
return nil, fmt.Errorf("model name is required")
|
|
}
|
|
if len(texts) == 0 {
|
|
return []EmbeddingData{}, nil
|
|
}
|
|
|
|
var region = "default"
|
|
if apiConfig != nil && apiConfig.Region != nil && *apiConfig.Region != "" {
|
|
region = *apiConfig.Region
|
|
}
|
|
|
|
url := fmt.Sprintf("%s/%s", b.BaseURL[region], b.URLSuffix.Embedding)
|
|
|
|
reqBody := map[string]interface{}{
|
|
"model": *modelName,
|
|
"input": texts,
|
|
}
|
|
if embeddingConfig != nil && embeddingConfig.Dimension > 0 {
|
|
reqBody["dimensions"] = embeddingConfig.Dimension
|
|
}
|
|
|
|
jsonData, err := json.Marshal(reqBody)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to marshal request: %w", err)
|
|
}
|
|
|
|
req, err := http.NewRequest("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.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("Baidu embedding API error: status %d, body: %s", resp.StatusCode, string(body))
|
|
}
|
|
|
|
var parsed baiduEmbeddingResponse
|
|
if err = json.Unmarshal(body, &parsed); err != nil {
|
|
return nil, fmt.Errorf("failed to parse response: %w", err)
|
|
}
|
|
|
|
if len(parsed.Data) != len(texts) {
|
|
return nil, fmt.Errorf("expected %d embeddings, got %d", len(texts), len(parsed.Data))
|
|
}
|
|
|
|
embeddings := make([]EmbeddingData, len(texts))
|
|
seen := make([]bool, len(texts))
|
|
for _, item := range parsed.Data {
|
|
if item.Index == nil {
|
|
return nil, fmt.Errorf("missing index field in embedding item")
|
|
}
|
|
idx := *item.Index
|
|
if idx < 0 || idx >= len(texts) {
|
|
return nil, fmt.Errorf("embedding index %d out of range", idx)
|
|
}
|
|
if seen[idx] {
|
|
return nil, fmt.Errorf("duplicate embedding index %d", idx)
|
|
}
|
|
if len(item.Embedding) == 0 {
|
|
return nil, fmt.Errorf("empty embedding at index %d", idx)
|
|
}
|
|
seen[idx] = true
|
|
embeddings[idx] = EmbeddingData{
|
|
Embedding: item.Embedding,
|
|
Index: idx,
|
|
}
|
|
}
|
|
|
|
for i, ok := range seen {
|
|
if !ok {
|
|
return nil, fmt.Errorf("missing embedding index %d", i)
|
|
}
|
|
}
|
|
|
|
return embeddings, nil
|
|
}
|
|
|
|
func (b *BaiduModel) Rerank(modelName *string, query string, documents []string, apiConfig *APIConfig, rerankConfig *RerankConfig) (*RerankResponse, error) {
|
|
if len(documents) == 0 {
|
|
return &RerankResponse{}, nil
|
|
}
|
|
|
|
var region = "default"
|
|
if apiConfig != nil && apiConfig.Region != nil && *apiConfig.Region != "" {
|
|
region = *apiConfig.Region
|
|
}
|
|
|
|
url := fmt.Sprintf("%s/%s", strings.TrimSuffix(b.BaseURL[region], "/"), b.URLSuffix.Rerank)
|
|
|
|
var topN = rerankConfig.TopN
|
|
if rerankConfig.TopN == 0 {
|
|
topN = len(documents)
|
|
}
|
|
|
|
reqBody := map[string]interface{}{
|
|
"model": *modelName,
|
|
"query": query,
|
|
"documents": documents,
|
|
"top_n": topN,
|
|
}
|
|
|
|
jsonData, err := json.Marshal(reqBody)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to marshal request: %w", err)
|
|
}
|
|
|
|
req, err := http.NewRequest("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.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("Baidu rerank API error: status %d, body: %s", resp.StatusCode, string(body))
|
|
}
|
|
|
|
var rerankResp struct {
|
|
Results []struct {
|
|
Index int `json:"index"`
|
|
RelevanceScore float64 `json:"relevance_score"`
|
|
} `json:"results"`
|
|
}
|
|
|
|
if err = json.Unmarshal(body, &rerankResp); err != nil {
|
|
return nil, fmt.Errorf("failed to decode response: %w", err)
|
|
}
|
|
|
|
var rerankResponse RerankResponse
|
|
for _, result := range rerankResp.Results {
|
|
rerankResult := RerankResult{
|
|
Index: result.Index,
|
|
RelevanceScore: result.RelevanceScore,
|
|
}
|
|
rerankResponse.Data = append(rerankResponse.Data, rerankResult)
|
|
}
|
|
|
|
return &rerankResponse, nil
|
|
}
|
|
|
|
// TranscribeAudio transcribe audio
|
|
func (b *BaiduModel) TranscribeAudio(modelName *string, file *string, apiConfig *APIConfig, asrConfig *ASRConfig) (*ASRResponse, error) {
|
|
return nil, fmt.Errorf("%s, no such method", b.Name())
|
|
}
|
|
|
|
func (z *BaiduModel) TranscribeAudioWithSender(modelName *string, file *string, apiConfig *APIConfig, asrConfig *ASRConfig, sender func(*string, *string) error) error {
|
|
return fmt.Errorf("%s, no such method", z.Name())
|
|
}
|
|
|
|
// AudioSpeech convert text to audio
|
|
func (b *BaiduModel) AudioSpeech(modelName *string, audioContent *string, apiConfig *APIConfig, ttsConfig *TTSConfig) (*TTSResponse, error) {
|
|
return nil, fmt.Errorf("%s, no such method", b.Name())
|
|
}
|
|
|
|
func (z *BaiduModel) AudioSpeechWithSender(modelName *string, audioContent *string, apiConfig *APIConfig, ttsConfig *TTSConfig, sender func(*string, *string) error) error {
|
|
return fmt.Errorf("%s, no such method", z.Name())
|
|
}
|
|
|
|
// OCRFile OCR file
|
|
type qianfanOCRResponse struct {
|
|
Id string `json:"id"`
|
|
Result struct {
|
|
LayoutParsingResults []struct {
|
|
Markdown struct {
|
|
Text string `json:"text"`
|
|
} `json:"markdown"`
|
|
} `json:"layoutParsingResults"`
|
|
} `json:"result"`
|
|
}
|
|
|
|
func (b *BaiduModel) OCRFile(modelName *string, content []byte, fileURL *string, apiConfig *APIConfig, ocrConfig *OCRConfig) (*OCRFileResponse, error) {
|
|
if (fileURL == nil || *fileURL == "") && (content == nil || len(content) == 0) {
|
|
return nil, fmt.Errorf("image url or content is required")
|
|
}
|
|
if apiConfig == nil || apiConfig.ApiKey == nil || *apiConfig.ApiKey == "" {
|
|
return nil, fmt.Errorf("api key is required")
|
|
}
|
|
if modelName == nil || *modelName == "" {
|
|
return nil, fmt.Errorf("model name is required")
|
|
}
|
|
|
|
region := "default"
|
|
if apiConfig.Region != nil && *apiConfig.Region != "" {
|
|
region = *apiConfig.Region
|
|
}
|
|
|
|
url := fmt.Sprintf("%s/%s", b.BaseURL[region], b.URLSuffix.OCR)
|
|
|
|
reqData := map[string]interface{}{
|
|
"model": *modelName,
|
|
}
|
|
|
|
if fileURL != nil && *fileURL != "" {
|
|
reqData["file"] = *fileURL
|
|
if strings.HasSuffix(strings.ToLower(*fileURL), ".pdf") {
|
|
reqData["fileType"] = 0 // PDF
|
|
} else {
|
|
reqData["fileType"] = 1 // img
|
|
}
|
|
} else if content != nil && len(content) > 0 {
|
|
reqData["file"] = base64.StdEncoding.EncodeToString(content)
|
|
|
|
mimeType := http.DetectContentType(content)
|
|
if strings.Contains(mimeType, "pdf") {
|
|
reqData["fileType"] = 0 // PDF
|
|
} else {
|
|
reqData["fileType"] = 1 // img
|
|
}
|
|
}
|
|
|
|
jsonData, err := json.Marshal(reqData)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to marshal json payload: %w", err)
|
|
}
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
|
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.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 error: %s, body: %s", resp.Status, string(body))
|
|
}
|
|
|
|
var apiResponse qianfanOCRResponse
|
|
if err = json.Unmarshal(body, &apiResponse); err != nil {
|
|
return nil, fmt.Errorf("failed to parse response json: %w", err)
|
|
}
|
|
|
|
var extractedText string
|
|
if len(apiResponse.Result.LayoutParsingResults) > 0 {
|
|
extractedText = apiResponse.Result.LayoutParsingResults[0].Markdown.Text
|
|
} else {
|
|
return nil, fmt.Errorf("no parsing results returned from API")
|
|
}
|
|
|
|
var ocrResponse = OCRFileResponse{
|
|
Text: &extractedText,
|
|
}
|
|
|
|
return &ocrResponse, nil
|
|
}
|
|
|
|
func (b *BaiduModel) ListModels(apiConfig *APIConfig) ([]string, error) {
|
|
var region = "default"
|
|
if apiConfig != nil && apiConfig.Region != nil && *apiConfig.Region != "" {
|
|
region = *apiConfig.Region
|
|
}
|
|
|
|
url := fmt.Sprintf("%s/%s", b.BaseURL[region], b.URLSuffix.Models)
|
|
|
|
reqBody := map[string]string{}
|
|
|
|
jsonData, err := json.Marshal(reqBody)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to marshal request: %w", err)
|
|
}
|
|
|
|
req, err := http.NewRequest("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", *apiConfig.ApiKey))
|
|
|
|
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 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 parse response: %w", err)
|
|
}
|
|
|
|
// convert result["data"] to []map[string]interface{}
|
|
models := make([]string, 0)
|
|
for _, model := range result["data"].([]interface{}) {
|
|
modelMap := model.(map[string]interface{})
|
|
modelName := modelMap["id"].(string)
|
|
models = append(models, modelName)
|
|
}
|
|
|
|
return models, nil
|
|
}
|
|
|
|
func (b *BaiduModel) Balance(apiConfig *APIConfig) (map[string]interface{}, error) {
|
|
return nil, fmt.Errorf("%s, no such method", b.Name())
|
|
}
|
|
|
|
func (b *BaiduModel) CheckConnection(apiConfig *APIConfig) error {
|
|
_, err := b.ListModels(apiConfig)
|
|
return err
|
|
}
|
|
|
|
func (z *BaiduModel) ParseFile(modelName *string, content []byte, url *string, apiConfig *APIConfig, parseFileConfig *ParseFileConfig) (*ParseFileResponse, error) {
|
|
return nil, fmt.Errorf("%s, no such method", z.Name())
|
|
}
|
|
|
|
func (z *BaiduModel) ListTasks(apiConfig *APIConfig) ([]ListTaskStatus, error) {
|
|
return nil, fmt.Errorf("%s, no such method", z.Name())
|
|
}
|
|
|
|
func (z *BaiduModel) ShowTask(taskID string, apiConfig *APIConfig) (*TaskResponse, error) {
|
|
return nil, fmt.Errorf("%s, no such method", z.Name())
|
|
}
|