mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-09-08 10:14:35 +08:00
feat(models): add shared HTTP client, SSE parser, and stub helpers for Go model drivers (#15821)
### What problem does this PR solve? The Go model-driver layer () has ~38,700 lines across 109 files. Roughly 74% of that is boilerplate duplicated into every driver: identical HTTP client setup, the same 65-line SSE scanner loop, and 10-11 one-line "not supported" stub methods per driver. Any fix must be manually propagated to every file. Closes #15820. This PR establishes the three shared utility files that form the foundation for incremental driver migration: --- ### Type of change - [x] New Feature (non-breaking change which adds functionality) - [x] Refactoring --------- Co-authored-by: Haruko386 <tryeverypossible@163.com>
This commit is contained in:
@@ -17,7 +17,6 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
@@ -25,7 +24,6 @@ import (
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// AstraflowModel implements ModelDriver for Astraflow (UCloud
|
||||
@@ -51,20 +49,11 @@ type AstraflowModel struct {
|
||||
|
||||
// NewAstraflowModel creates a new Astraflow model instance.
|
||||
func NewAstraflowModel(baseURL map[string]string, urlSuffix URLSuffix) *AstraflowModel {
|
||||
transport := http.DefaultTransport.(*http.Transport).Clone()
|
||||
transport.MaxIdleConns = 100
|
||||
transport.MaxIdleConnsPerHost = 10
|
||||
transport.IdleConnTimeout = 90 * time.Second
|
||||
transport.DisableCompression = false
|
||||
transport.ResponseHeaderTimeout = 60 * time.Second
|
||||
|
||||
return &AstraflowModel{
|
||||
baseModel: BaseModel{
|
||||
BaseURL: baseURL,
|
||||
URLSuffix: urlSuffix,
|
||||
httpClient: &http.Client{
|
||||
Transport: transport,
|
||||
},
|
||||
BaseURL: baseURL,
|
||||
URLSuffix: urlSuffix,
|
||||
httpClient: NewDriverHTTPClient(),
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -264,35 +253,19 @@ func (a *AstraflowModel) ChatStreamlyWithSender(modelName string, messages []Mes
|
||||
return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
scanner := bufio.NewScanner(resp.Body)
|
||||
scanner.Buffer(make([]byte, 64*1024), 1024*1024)
|
||||
sawTerminal := false
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
if !strings.HasPrefix(line, "data:") {
|
||||
continue
|
||||
}
|
||||
data := strings.TrimSpace(line[5:])
|
||||
if data == "[DONE]" {
|
||||
sawTerminal = true
|
||||
break
|
||||
}
|
||||
|
||||
var event map[string]interface{}
|
||||
if err = json.Unmarshal([]byte(data), &event); err != nil {
|
||||
return fmt.Errorf("astraflow: invalid SSE event: %w", err)
|
||||
}
|
||||
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 {
|
||||
continue
|
||||
return nil
|
||||
}
|
||||
firstChoice, ok := choices[0].(map[string]interface{})
|
||||
if !ok {
|
||||
continue
|
||||
return nil
|
||||
}
|
||||
if delta, ok := firstChoice["delta"].(map[string]interface{}); ok {
|
||||
if r, ok := delta["reasoning_content"].(string); ok && r != "" {
|
||||
@@ -310,14 +283,13 @@ func (a *AstraflowModel) ChatStreamlyWithSender(modelName string, messages []Mes
|
||||
}
|
||||
if finish, ok := firstChoice["finish_reason"].(string); ok && finish != "" {
|
||||
sawTerminal = true
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if err := scanner.Err(); err != nil {
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to scan response body: %w", err)
|
||||
}
|
||||
if !sawTerminal {
|
||||
if !done && !sawTerminal {
|
||||
return fmt.Errorf("astraflow: stream ended before [DONE] or finish_reason")
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user