zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-28 11:48:10 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refactor[Go Model Provider]: refactor baseURL and modelConfig (#15627)

Browse Source 
### What problem does this PR solve?

As Title

### Type of change

- [x] Refactoring
This commit is contained in:
Haruko386
2026-06-04 17:50:22 +08:00
committed by GitHub
parent 04dc3bb19c
commit baeb0c0431
65 changed files with 2834 additions and 4410 deletions
Download Patch File Download Diff File Expand all files Collapse all files

163
internal/entity/models/astraflow.go
View File

@@ -46,17 +46,10 @@ import (
// DeepSeek-R1 applies and there's no need for an inline <think>...
// extractor like Novita's.
type AstraflowModel struct {
BaseURL map[string]string
URLSuffix URLSuffix
httpClient *http.Client
baseModel BaseModel
}
// NewAstraflowModel creates a new Astraflow model instance.
//
// Same transport convention as the other Go drivers in this package:
// clone http.DefaultTransport to keep ProxyFromEnvironment, DialContext,
// HTTP/2, and TLS defaults, and only override the connection-pool
// fields. No client-level Timeout so SSE streams aren't capped mid-flight.
func NewAstraflowModel(baseURL map[string]string, urlSuffix URLSuffix) *AstraflowModel {
transport := http.DefaultTransport.(*http.Transport).Clone()
transport.MaxIdleConns = 100
@@ -66,52 +59,39 @@ func NewAstraflowModel(baseURL map[string]string, urlSuffix URLSuffix) *Astraflo
transport.ResponseHeaderTimeout = 60 * time.Second
return &AstraflowModel{
BaseURL: baseURL,
URLSuffix: urlSuffix,
httpClient: &http.Client{
Transport: transport,
baseModel: BaseModel{
BaseURL: baseURL,
URLSuffix: urlSuffix,
httpClient: &http.Client{
Transport: transport,
},
},
}
}
func (a *AstraflowModel) NewInstance(baseURL map[string]string) ModelDriver {
return NewAstraflowModel(baseURL, a.URLSuffix)
return NewAstraflowModel(baseURL, a.baseModel.URLSuffix)
}
func (a *AstraflowModel) Name() string {
return "astraflow"
}
func (a *AstraflowModel) baseURLForRegion(region string) (string, error) {
base, ok := a.BaseURL[region]
if !ok || base == "" {
return "", fmt.Errorf("astraflow: no base URL configured for region %q", region)
}
return strings.TrimSuffix(base, "/"), nil
}
// ChatWithMessages sends a non-streaming chat request and returns the
// full response. Forwards documented OpenAI-shaped parameters when the
// caller supplies them; reasoning_content is surfaced separately so the
// visible Answer is never polluted by chain-of-thought.
// ChatWithMessages sends a non-streaming chat request
func (a *AstraflowModel) 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 required")
if err := a.baseModel.APIConfigCheck(apiConfig); err != nil {
return nil, err
}
if len(messages) == 0 {
return nil, fmt.Errorf("messages is empty")
}
region := "default"
if apiConfig.Region != nil && *apiConfig.Region != "" {
region = *apiConfig.Region
}
baseURL, err := a.baseURLForRegion(region)
baseURL, err := a.baseModel.GetBaseURL(apiConfig)
if err != nil {
return nil, err
}
url := fmt.Sprintf("%s/%s", baseURL, a.URLSuffix.Chat)
baseURL = strings.TrimSuffix(baseURL, "/")
url := fmt.Sprintf("%s/%s", baseURL, a.baseModel.URLSuffix.Chat)
apiMessages := make([]map[string]interface{}, len(messages))
for i, msg := range messages {
@@ -157,7 +137,7 @@ func (a *AstraflowModel) ChatWithMessages(modelName string, messages []Message,
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := a.httpClient.Do(req)
resp, err := a.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
@@ -196,10 +176,6 @@ func (a *AstraflowModel) ChatWithMessages(modelName string, messages []Message,
return nil, fmt.Errorf("invalid content format")
}
// Reasoning models (deepseek-r1 / kimi / glm-thinking) put
// chain-of-thought in a separate `reasoning_content` field with
// `content` already cleaned. Absent or non-string means no reasoning
// was emitted; leave it empty rather than synthesizing one.
reasonContent := ""
if r, ok := messageMap["reasoning_content"].(string); ok {
reasonContent = r
@@ -211,31 +187,25 @@ func (a *AstraflowModel) ChatWithMessages(modelName string, messages []Message,
}, nil
}
// ChatStreamlyWithSender opens the SSE chat-completions endpoint and
// forwards each delta through the supplied sender. Reasoning chunks go
// to the sender's second argument, content chunks to the first; the
// stream is terminated by either `[DONE]` or a delta with finish_reason.
// ChatStreamlyWithSender opens the SSE chat-completions
func (a *AstraflowModel) ChatStreamlyWithSender(modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, sender func(*string, *string) error) error {
if err := a.baseModel.APIConfigCheck(apiConfig); err != nil {
return err
}
if sender == nil {
return fmt.Errorf("sender is required")
}
if len(messages) == 0 {
return fmt.Errorf("messages is empty")
}
if apiConfig == nil || apiConfig.ApiKey == nil || *apiConfig.ApiKey == "" {
return fmt.Errorf("api key is required")
}
region := "default"
if apiConfig.Region != nil && *apiConfig.Region != "" {
region = *apiConfig.Region
}
baseURL, err := a.baseURLForRegion(region)
baseURL, err := a.baseModel.GetBaseURL(apiConfig)
if err != nil {
return err
}
url := fmt.Sprintf("%s/%s", baseURL, a.URLSuffix.Chat)
baseURL = strings.TrimSuffix(baseURL, "/")
url := fmt.Sprintf("%s/%s", baseURL, a.baseModel.URLSuffix.Chat)
apiMessages := make([]map[string]interface{}, len(messages))
for i, msg := range messages {
@@ -252,9 +222,6 @@ func (a *AstraflowModel) ChatStreamlyWithSender(modelName string, messages []Mes
}
if chatModelConfig != nil {
// Guard against the caller asking for stream=false on a code path
// that only knows how to read SSE. Without this, a non-SSE JSON
// body would parse as zero chunks and look like a silent timeout.
if chatModelConfig.Stream != nil && !*chatModelConfig.Stream {
return fmt.Errorf("stream must be true in ChatStreamlyWithSender")
}
@@ -286,7 +253,7 @@ func (a *AstraflowModel) ChatStreamlyWithSender(modelName string, messages []Mes
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := a.httpClient.Do(req)
resp, err := a.baseModel.httpClient.Do(req)
if err != nil {
return fmt.Errorf("failed to send request: %w", err)
}
@@ -313,15 +280,8 @@ func (a *AstraflowModel) ChatStreamlyWithSender(modelName string, messages []Mes
var event map[string]interface{}
if err = json.Unmarshal([]byte(data), &event); err != nil {
// A malformed frame usually means a truncated event or an
// upstream incident. Surface it instead of silently producing
// partial output.
return fmt.Errorf("astraflow: invalid SSE event: %w", err)
}
// Astraflow can emit a terminal `{"error": ...}` frame when the
// upstream model rejects mid-stream (rate limit, content policy).
// Surface it verbatim instead of falling through to "no choices".
if apiErr, ok := event["error"]; ok {
return fmt.Errorf("astraflow: upstream stream error: %v", apiErr)
}
@@ -334,10 +294,6 @@ func (a *AstraflowModel) ChatStreamlyWithSender(modelName string, messages []Mes
if !ok {
continue
}
// Reasoning first, content second — matches the wire ordering
// for reasoning models and lets UIs render the chain-of-thought
// before the visible token. A terminal frame may carry
// finish_reason without a delta, so don't skip when delta is absent.
if delta, ok := firstChoice["delta"].(map[string]interface{}); ok {
if r, ok := delta["reasoning_content"].(string); ok && r != "" {
rr := r
@@ -372,24 +328,17 @@ func (a *AstraflowModel) ChatStreamlyWithSender(modelName string, messages []Mes
return nil
}
// ListModels returns the model ids visible to the API key by calling
// /v1/models. Used by Add-Provider's connection check and by the UI's
// model picker.
func (a *AstraflowModel) ListModels(apiConfig *APIConfig) ([]string, error) {
if apiConfig == nil || apiConfig.ApiKey == nil || *apiConfig.ApiKey == "" {
return nil, fmt.Errorf("api key is required")
if err := a.baseModel.APIConfigCheck(apiConfig); err != nil {
return nil, err
}
region := "default"
if apiConfig.Region != nil && *apiConfig.Region != "" {
region = *apiConfig.Region
}
baseURL, err := a.baseURLForRegion(region)
baseURL, err := a.baseModel.GetBaseURL(apiConfig)
if err != nil {
return nil, err
}
url := fmt.Sprintf("%s/%s", baseURL, a.URLSuffix.Models)
baseURL = strings.TrimSuffix(baseURL, "/")
url := fmt.Sprintf("%s/%s", baseURL, a.baseModel.URLSuffix.Models)
ctx, cancel := context.WithTimeout(context.Background(), nonStreamCallTimeout)
defer cancel()
@@ -400,7 +349,7 @@ func (a *AstraflowModel) ListModels(apiConfig *APIConfig) ([]string, error) {
}
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := a.httpClient.Do(req)
resp, err := a.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
@@ -439,29 +388,25 @@ func (a *AstraflowModel) ListModels(apiConfig *APIConfig) ([]string, error) {
return models, nil
}
// CheckConnection verifies the API key by calling ListModels. The /v1/models
// endpoint is the documented lightweight way to validate credentials on
// OpenAI-compatible gateways without burning chat-completion quota.
func (a *AstraflowModel) CheckConnection(apiConfig *APIConfig) error {
_, err := a.ListModels(apiConfig)
return err
}
// Embed is reserved for a follow-up issue; the Astraflow factory tag
// includes "TEXT EMBEDDING" but this initial driver only implements
// chat, mirroring how Novita / TogetherAI / DeepInfra landed
// method-by-method.
func (a *AstraflowModel) Embed(modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig) ([]EmbeddingData, error) {
if err := a.baseModel.APIConfigCheck(apiConfig); err != nil {
return nil, err
}
if len(texts) == 0 {
return []EmbeddingData{}, nil
}
var region = "default"
if apiConfig != nil && apiConfig.Region != nil && *apiConfig.Region != "" {
region = *apiConfig.Region
resolvedBaseURL, err := a.baseModel.GetBaseURL(apiConfig)
if err != nil {
return nil, err
}
url := fmt.Sprintf("%s/%s", a.BaseURL[region], a.URLSuffix.Embedding)
url := fmt.Sprintf("%s/%s", resolvedBaseURL, a.baseModel.URLSuffix.Embedding)
reqBody := map[string]interface{}{
"model": *modelName,
@@ -481,7 +426,7 @@ func (a *AstraflowModel) Embed(modelName *string, texts []string, apiConfig *API
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := a.httpClient.Do(req)
resp, err := a.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
@@ -523,16 +468,19 @@ func (a *AstraflowModel) Embed(modelName *string, texts []string, apiConfig *API
}
func (a *AstraflowModel) Rerank(modelName *string, query string, documents []string, apiConfig *APIConfig, rerankConfig *RerankConfig) (*RerankResponse, error) {
if err := a.baseModel.APIConfigCheck(apiConfig); err != nil {
return nil, err
}
if len(documents) == 0 {
return &RerankResponse{}, nil
}
var region = "default"
if apiConfig != nil && apiConfig.Region != nil && *apiConfig.Region != "" {
region = *apiConfig.Region
resolvedBaseURL, err := a.baseModel.GetBaseURL(apiConfig)
if err != nil {
return nil, err
}
url := fmt.Sprintf("%s/%s", a.BaseURL[region], a.URLSuffix.Rerank)
url := fmt.Sprintf("%s/%s", resolvedBaseURL, a.baseModel.URLSuffix.Rerank)
var topN = rerankConfig.TopN
if rerankConfig.TopN != 0 {
@@ -559,7 +507,7 @@ func (a *AstraflowModel) Rerank(modelName *string, query string, documents []str
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := a.httpClient.Do(req)
resp, err := a.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
@@ -610,20 +558,19 @@ func (a *AstraflowModel) TranscribeAudioWithSender(modelName *string, file *stri
}
func (a *AstraflowModel) AudioSpeech(modelName *string, audioContent *string, apiConfig *APIConfig, ttsConfig *TTSConfig) (*TTSResponse, error) {
if apiConfig == nil || apiConfig.ApiKey == nil || *apiConfig.ApiKey == "" {
return nil, fmt.Errorf("Astraflow API key is missing")
if err := a.baseModel.APIConfigCheck(apiConfig); err != nil {
return nil, err
}
if audioContent == nil || *audioContent == "" {
return nil, fmt.Errorf("text content is missing")
}
var region = "default"
if apiConfig.Region != nil && *apiConfig.Region != "" {
region = *apiConfig.Region
resolvedBaseURL, err := a.baseModel.GetBaseURL(apiConfig)
if err != nil {
return nil, err
}
url := fmt.Sprintf("%s/%s", a.BaseURL[region], a.URLSuffix.TTS)
url := fmt.Sprintf("%s/%s", resolvedBaseURL, a.baseModel.URLSuffix.TTS)
reqBody := map[string]interface{}{
"model": *modelName,
@@ -652,7 +599,7 @@ func (a *AstraflowModel) AudioSpeech(modelName *string, audioContent *string, ap
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := a.httpClient.Do(req)
resp, err := a.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}