zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-08-10 01:11:23 +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

155
internal/entity/models/stepfun.go
View File

@@ -30,28 +30,11 @@ import (
)
// StepFunModel implements ModelDriver for StepFun (阶跃星辰).
//
// StepFun exposes an OpenAI-compatible REST API at https://api.stepfun.com/v1
// (chat completions at /chat/completions, list models at /models). The wire
// shape matches OpenAI closely enough that the chat path here is a direct
// port of the OpenAI driver.
type StepFunModel struct {
BaseURL map[string]string
URLSuffix URLSuffix
httpClient *http.Client
baseModel BaseModel
}
// NewStepFunModel creates a new StepFun model instance.
//
// We clone http.DefaultTransport so we keep Go's defaults for
// ProxyFromEnvironment, DialContext (with KeepAlive), HTTP/2,
// TLSHandshakeTimeout, and ExpectContinueTimeout, and only override
// the connection-pool fields we care about.
//
// The Client itself has no Timeout. http.Client.Timeout would also
// cap the time spent reading the response body, which would cut off
// long-lived SSE streams in ChatStreamlyWithSender. Non-streaming
// callers wrap each request with context.WithTimeout instead.
func NewStepFunModel(baseURL map[string]string, urlSuffix URLSuffix) *StepFunModel {
transport := http.DefaultTransport.(*http.Transport).Clone()
transport.MaxIdleConns = 100
@@ -61,63 +44,40 @@ func NewStepFunModel(baseURL map[string]string, urlSuffix URLSuffix) *StepFunMod
transport.ResponseHeaderTimeout = 60 * time.Second
return &StepFunModel{
BaseURL: baseURL,
URLSuffix: urlSuffix,
httpClient: &http.Client{
Transport: transport,
baseModel: BaseModel{
BaseURL: baseURL,
URLSuffix: urlSuffix,
httpClient: &http.Client{
Transport: transport,
},
},
}
}
/*
RAGFlow(user)> tts with 'fnlp/MOSS-TTSD-v0.5@test@siliconflow' text 'He who desires but acts not, breeds pestilence.' play format 'wav' param '{"voice": "fnlp/MOSS-TTSD-v0.5:alex"}'
SUCCESS
RAGFlow(user)> stream tts with 'fnlp/MOSS-TTSD-v0.5@test@siliconflow' text 'He who desires but acts not, breeds pestilence.' play format 'wav' param '{"voice": "fnlp/MOSS-TTSD-v0.5:claire"}'
SUCCESS
*/
func (s *StepFunModel) NewInstance(baseURL map[string]string) ModelDriver {
return NewStepFunModel(baseURL, s.URLSuffix)
return NewStepFunModel(baseURL, s.baseModel.URLSuffix)
}
func (s *StepFunModel) Name() string {
return "stepfun"
}
// baseURLForRegion returns the base URL for the given region, or an
// error if no entry exists. This makes a misconfigured region fail
// fast with a clear message, instead of silently producing a relative
// URL that the HTTP transport then rejects.
func (s *StepFunModel) baseURLForRegion(region string) (string, error) {
base, ok := s.BaseURL[region]
if !ok || base == "" {
return "", fmt.Errorf("stepfun: no base URL configured for region %q", region)
}
return base, nil
}
// ChatWithMessages sends multiple messages with roles and returns the response.
func (s *StepFunModel) 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 := s.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 := s.baseURLForRegion(region)
baseURL, err := s.baseModel.GetBaseURL(apiConfig)
if err != nil {
return nil, err
}
url := fmt.Sprintf("%s/%s", baseURL, s.URLSuffix.Chat)
baseURL = strings.TrimSuffix(baseURL, "/")
url := fmt.Sprintf("%s/%s", baseURL, s.baseModel.URLSuffix.Chat)
apiMessages := make([]map[string]interface{}, len(messages))
for i, msg := range messages {
@@ -133,9 +93,6 @@ func (s *StepFunModel) ChatWithMessages(modelName string, messages []Message, ap
"stream": false,
}
// Note: do NOT propagate chatModelConfig.Stream into the request body
// here. ChatWithMessages parses a single JSON response, so stream must
// always be off for this code path.
if chatModelConfig != nil {
if chatModelConfig.MaxTokens != nil {
reqBody["max_tokens"] = *chatModelConfig.MaxTokens
@@ -167,7 +124,7 @@ func (s *StepFunModel) ChatWithMessages(modelName string, messages []Message, ap
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := s.httpClient.Do(req)
resp, err := s.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
@@ -214,10 +171,12 @@ func (s *StepFunModel) ChatWithMessages(modelName string, messages []Message, ap
}, nil
}
// ChatStreamlyWithSender sends messages and streams the response via the
// sender function. The StepFun SSE stream uses the same shape as OpenAI:
// "data:" lines carrying JSON events, with a final "[DONE]" line.
// ChatStreamlyWithSender sends messages and streams the response
func (s *StepFunModel) ChatStreamlyWithSender(modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, sender func(*string, *string) error) error {
if err := s.baseModel.APIConfigCheck(apiConfig); err != nil {
return err
}
if sender == nil {
return fmt.Errorf("sender is required")
}
@@ -226,20 +185,12 @@ func (s *StepFunModel) ChatStreamlyWithSender(modelName string, messages []Messa
return fmt.Errorf("messages is empty")
}
if apiConfig == nil || apiConfig.ApiKey == nil || *apiConfig.ApiKey == "" {
return fmt.Errorf("api key is required")
}
var region = "default"
if apiConfig.Region != nil && *apiConfig.Region != "" {
region = *apiConfig.Region
}
baseURL, err := s.baseURLForRegion(region)
baseURL, err := s.baseModel.GetBaseURL(apiConfig)
if err != nil {
return err
}
url := fmt.Sprintf("%s/%s", baseURL, s.URLSuffix.Chat)
baseURL = strings.TrimSuffix(baseURL, "/")
url := fmt.Sprintf("%s/%s", baseURL, s.baseModel.URLSuffix.Chat)
apiMessages := make([]map[string]interface{}, len(messages))
for i, msg := range messages {
@@ -256,10 +207,6 @@ func (s *StepFunModel) ChatStreamlyWithSender(modelName string, messages []Messa
}
if chatModelConfig != nil {
// Refuse to run if the caller explicitly asked for stream=false.
// The body of this method only knows how to read SSE, so a
// non-SSE JSON response would be parsed as if it were a stream
// and produce no chunks. Better to fail clearly.
if chatModelConfig.Stream != nil && !*chatModelConfig.Stream {
return fmt.Errorf("stream must be true in ChatStreamlyWithSender")
}
@@ -283,9 +230,6 @@ func (s *StepFunModel) ChatStreamlyWithSender(modelName string, messages []Messa
return fmt.Errorf("failed to marshal request: %w", err)
}
// SSE streams are long-lived. We rely on the transport's
// ResponseHeaderTimeout to cap the connection-establishment phase
// instead of attaching a hard deadline here.
req, err := http.NewRequestWithContext(context.Background(), "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return fmt.Errorf("failed to create request: %w", err)
@@ -294,7 +238,7 @@ func (s *StepFunModel) ChatStreamlyWithSender(modelName string, messages []Messa
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := s.httpClient.Do(req)
resp, err := s.baseModel.httpClient.Do(req)
if err != nil {
return fmt.Errorf("failed to send request: %w", err)
}
@@ -305,8 +249,6 @@ func (s *StepFunModel) ChatStreamlyWithSender(modelName string, messages []Messa
return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
// SSE parsing: bump the scanner buffer from the 64KB default to 1MB
// so we never silently truncate a long data: line.
scanner := bufio.NewScanner(resp.Body)
scanner.Buffer(make([]byte, 64*1024), 1024*1024)
sawTerminal := false
@@ -373,29 +315,23 @@ func (s *StepFunModel) ChatStreamlyWithSender(modelName string, messages []Messa
return nil
}
// Embed is left as a stub. StepFun has not advertised a public embeddings
// endpoint in the API reference linked from the umbrella issue, so any real
// implementation belongs in a follow-up only after the endpoint is verified.
// Embed is left as a stub.
func (s *StepFunModel) Embed(modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig) ([]EmbeddingData, error) {
return nil, fmt.Errorf("not implemented")
}
// ListModels returns the list of model ids visible to the API key.
func (s *StepFunModel) ListModels(apiConfig *APIConfig) ([]string, error) {
if apiConfig == nil || apiConfig.ApiKey == nil || *apiConfig.ApiKey == "" {
return nil, fmt.Errorf("api key is required")
if err := s.baseModel.APIConfigCheck(apiConfig); err != nil {
return nil, err
}
region := "default"
if apiConfig.Region != nil && *apiConfig.Region != "" {
region = *apiConfig.Region
}
baseURL, err := s.baseURLForRegion(region)
baseURL, err := s.baseModel.GetBaseURL(apiConfig)
if err != nil {
return nil, err
}
url := fmt.Sprintf("%s/%s", baseURL, s.URLSuffix.Models)
baseURL = strings.TrimSuffix(baseURL, "/")
url := fmt.Sprintf("%s/%s", baseURL, s.baseModel.URLSuffix.Models)
ctx, cancel := context.WithTimeout(context.Background(), nonStreamCallTimeout)
defer cancel()
@@ -407,7 +343,7 @@ func (s *StepFunModel) ListModels(apiConfig *APIConfig) ([]string, error) {
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := s.httpClient.Do(req)
resp, err := s.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
@@ -479,17 +415,20 @@ func (s *StepFunModel) TranscribeAudioWithSender(modelName *string, file *string
// AudioSpeech convert text to audio
func (s *StepFunModel) AudioSpeech(modelName *string, audioContent *string, apiConfig *APIConfig, ttsConfig *TTSConfig) (*TTSResponse, error) {
if err := s.baseModel.APIConfigCheck(apiConfig); err != nil {
return nil, err
}
// TODO Test it
if audioContent == nil || *audioContent == "" {
return nil, fmt.Errorf("audio content is empty")
}
var region = "default"
if apiConfig != nil && apiConfig.Region != nil && *apiConfig.Region != "" {
region = *apiConfig.Region
resolvedBaseURL, err := s.baseModel.GetBaseURL(apiConfig)
if err != nil {
return nil, err
}
url := fmt.Sprintf("%s/%s", s.BaseURL[region], s.URLSuffix.TTS)
url := fmt.Sprintf("%s/%s", resolvedBaseURL, s.baseModel.URLSuffix.TTS)
reqBody := map[string]interface{}{
"model": *modelName,
@@ -518,7 +457,7 @@ func (s *StepFunModel) AudioSpeech(modelName *string, audioContent *string, apiC
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := s.httpClient.Do(req)
resp, err := s.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
@@ -538,21 +477,19 @@ func (s *StepFunModel) AudioSpeech(modelName *string, audioContent *string, apiC
// AudioSpeechWithSender for Streaming TTS
func (s *StepFunModel) AudioSpeechWithSender(modelName *string, audioContent *string, apiConfig *APIConfig, ttsConfig *TTSConfig, sender func(*string, *string) error) error {
// TODO Test it
if apiConfig == nil || apiConfig.ApiKey == nil || *apiConfig.ApiKey == "" {
return fmt.Errorf("StepFun API key is missing")
if err := s.baseModel.APIConfigCheck(apiConfig); err != nil {
return err
}
if audioContent == nil || *audioContent == "" {
return fmt.Errorf("audio content is empty")
}
var region = "default"
if apiConfig.Region != nil && *apiConfig.Region != "" {
region = *apiConfig.Region
resolvedBaseURL, err := s.baseModel.GetBaseURL(apiConfig)
if err != nil {
return err
}
url := fmt.Sprintf("%s/%s", s.BaseURL[region], s.URLSuffix.TTS)
url := fmt.Sprintf("%s/%s", resolvedBaseURL, s.baseModel.URLSuffix.TTS)
reqBody := map[string]interface{}{
"model": *modelName,
@@ -582,7 +519,7 @@ func (s *StepFunModel) AudioSpeechWithSender(modelName *string, audioContent *st
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := s.httpClient.Do(req)
resp, err := s.baseModel.httpClient.Do(req)
if err != nil {
return fmt.Errorf("failed to send request: %w", err)
}