zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-09-08 10:14:35 +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
internal/entity/models/vllm.go
+67 -73
View File
@@ -31,40 +31,29 @@ import (
// VllmModel implements ModelDriver for Vllm AI
type VllmModel struct {
BaseURL map[string]string
URLSuffix URLSuffix
httpClient *http.Client // Reusable HTTP client with connection pool
baseModel BaseModel
}
// NewVllmModel creates a new Vllm AI model instance
func NewVllmModel(baseURL map[string]string, urlSuffix URLSuffix) *VllmModel {
return &VllmModel{
BaseURL: baseURL,
URLSuffix: urlSuffix,
httpClient: &http.Client{
Transport: &http.Transport{
MaxIdleConns: 100,
MaxIdleConnsPerHost: 10,
IdleConnTimeout: 90 * time.Second,
DisableCompression: false,
baseModel: BaseModel{
BaseURL: baseURL,
URLSuffix: urlSuffix,
httpClient: &http.Client{
Transport: &http.Transport{
MaxIdleConns: 100,
MaxIdleConnsPerHost: 10,
IdleConnTimeout: 90 * time.Second,
DisableCompression: false,
},
},
},
}
}
func (v *VllmModel) NewInstance(baseURL map[string]string) ModelDriver {
return &VllmModel{
BaseURL: baseURL,
URLSuffix: v.URLSuffix,
httpClient: &http.Client{
Transport: &http.Transport{
MaxIdleConns: 100,
MaxIdleConnsPerHost: 10,
IdleConnTimeout: 90 * time.Second,
DisableCompression: false,
},
},
}
return NewVllmModel(baseURL, v.baseModel.URLSuffix)
}
func (v *VllmModel) Name() string {
@@ -73,21 +62,24 @@ func (v *VllmModel) Name() string {
// ChatWithMessages sends multiple messages with roles and returns response
func (v *VllmModel) ChatWithMessages(modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig) (*ChatResponse, error) {
if err := v.baseModel.APIConfigCheck(apiConfig); err != nil {
return nil, err
}
if len(messages) == 0 {
return nil, fmt.Errorf("messages is empty")
}
var region = "default"
if apiConfig != nil && apiConfig.Region != nil && *apiConfig.Region != "" {
region = *apiConfig.Region
resolvedBaseURL, err := v.baseModel.GetBaseURL(apiConfig)
if err != nil {
return nil, err
}
url := fmt.Sprintf("%s/%s", v.BaseURL[region], v.URLSuffix.Chat)
url := fmt.Sprintf("%s/%s", resolvedBaseURL, v.baseModel.URLSuffix.Chat)
// For qwen/glm models, use async chat endpoint
modelType := strings.Split(modelName, "-")[0]
if modelType == "qwen" || modelType == "glm" {
url = fmt.Sprintf("%s/%s", v.BaseURL[region], v.URLSuffix.AsyncChat)
url = fmt.Sprintf("%s/%s", resolvedBaseURL, v.baseModel.URLSuffix.AsyncChat)
}
// Convert messages to API format
@@ -157,7 +149,7 @@ func (v *VllmModel) ChatWithMessages(modelName string, messages []Message, apiCo
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := v.httpClient.Do(req)
resp, err := v.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
@@ -219,19 +211,22 @@ func (v *VllmModel) ChatWithMessages(modelName string, messages []Message, apiCo
// ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel)
func (v *VllmModel) ChatStreamlyWithSender(modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, sender func(*string, *string) error) error {
if err := v.baseModel.APIConfigCheck(apiConfig); err != nil {
return err
}
if len(messages) == 0 {
return fmt.Errorf("messages is empty")
}
var region = "default"
if apiConfig != nil && apiConfig.Region != nil && *apiConfig.Region != "" {
region = *apiConfig.Region
resolvedBaseURL, err := v.baseModel.GetBaseURL(apiConfig)
if err != nil {
return err
}
url := fmt.Sprintf("%s/%s", v.BaseURL[region], v.URLSuffix.Chat)
url := fmt.Sprintf("%s/%s", resolvedBaseURL, v.baseModel.URLSuffix.Chat)
modelType := strings.Split(modelName, "-")[0]
if modelType == "qwen" || modelType == "glm" {
url = fmt.Sprintf("%s/%s", v.BaseURL[region], v.URLSuffix.AsyncChat)
url = fmt.Sprintf("%s/%s", resolvedBaseURL, v.baseModel.URLSuffix.AsyncChat)
}
// Convert messages to API format (supporting multimodal content)
@@ -302,7 +297,7 @@ func (v *VllmModel) ChatStreamlyWithSender(modelName string, messages []Message,
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := v.httpClient.Do(req)
resp, err := v.baseModel.httpClient.Do(req)
if err != nil {
return fmt.Errorf("failed to send request: %w", err)
}
@@ -393,6 +388,10 @@ type vllmEmbeddingResponse struct {
// Embed embeds a list of texts into embeddings
func (v *VllmModel) Embed(modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig) ([]EmbeddingData, error) {
if err := v.baseModel.APIConfigCheck(apiConfig); err != nil {
return nil, err
}
if len(texts) == 0 {
return []EmbeddingData{}, nil
}
@@ -401,20 +400,19 @@ func (v *VllmModel) Embed(modelName *string, texts []string, apiConfig *APIConfi
return nil, fmt.Errorf("model name is required")
}
region := "default"
if apiConfig != nil && apiConfig.Region != nil && *apiConfig.Region != "" {
region = *apiConfig.Region
resolvedBaseURL, err := v.baseModel.GetBaseURL(apiConfig)
if err != nil {
return nil, err
}
baseURL := v.BaseURL[region]
baseURL := resolvedBaseURL
if baseURL == "" {
baseURL = v.BaseURL["default"]
baseURL = resolvedBaseURL
}
if baseURL == "" {
return nil, fmt.Errorf("missing base URL: please configure the local access address for vLLM (e.g., http://127.0.0.1:8000/v1)")
}
url := fmt.Sprintf("%s/%s", strings.TrimSuffix(baseURL, "/"), v.URLSuffix.Embedding)
url := fmt.Sprintf("%s/%s", strings.TrimSuffix(baseURL, "/"), v.baseModel.URLSuffix.Embedding)
reqBody := map[string]interface{}{
"model": *modelName,
@@ -438,11 +436,9 @@ func (v *VllmModel) Embed(modelName *string, texts []string, apiConfig *APIConfi
}
req.Header.Set("Content-Type", "application/json")
if apiConfig != nil && apiConfig.ApiKey != nil && *apiConfig.ApiKey != "" {
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
}
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := v.httpClient.Do(req)
resp, err := v.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
@@ -474,21 +470,23 @@ func (v *VllmModel) Embed(modelName *string, texts []string, apiConfig *APIConfi
}
func (v *VllmModel) ListModels(apiConfig *APIConfig) ([]string, error) {
var region = "default"
if apiConfig != nil && apiConfig.Region != nil && *apiConfig.Region != "" {
region = *apiConfig.Region
if err := v.baseModel.APIConfigCheck(apiConfig); err != nil {
return nil, err
}
baseURL := v.BaseURL[region]
resolvedBaseURL, err := v.baseModel.GetBaseURL(apiConfig)
if err != nil {
return nil, err
}
baseURL := resolvedBaseURL
if baseURL == "" {
baseURL = v.BaseURL["default"]
baseURL = resolvedBaseURL
}
if baseURL == "" {
return nil, fmt.Errorf("missing base URL: please configure the local access address for vLLM (e.g., http://127.0.0.1:8000/v1)")
}
url := fmt.Sprintf("%s/%s", baseURL, v.URLSuffix.Models)
url := fmt.Sprintf("%s/%s", baseURL, v.baseModel.URLSuffix.Models)
reqBody := map[string]interface{}{}
@@ -506,14 +504,9 @@ func (v *VllmModel) ListModels(apiConfig *APIConfig) ([]string, error) {
}
req.Header.Set("Content-Type", "application/json")
// vLLM is a local provider and the API key is optional. Only set
// the Authorization header when a non-empty key was supplied. This
// also avoids a nil-pointer dereference on apiConfig or ApiKey.
if apiConfig != nil && apiConfig.ApiKey != nil && *apiConfig.ApiKey != "" {
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
}
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := v.httpClient.Do(req)
resp, err := v.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
@@ -585,6 +578,10 @@ type vllmRerankResponse struct {
// matching the existing Embed/ListModels behaviour for this local
// driver.
func (v *VllmModel) Rerank(modelName *string, query string, documents []string, apiConfig *APIConfig, rerankConfig *RerankConfig) (*RerankResponse, error) {
if err := v.baseModel.APIConfigCheck(apiConfig); err != nil {
return nil, err
}
if len(documents) == 0 {
return &RerankResponse{}, nil
}
@@ -592,20 +589,19 @@ func (v *VllmModel) Rerank(modelName *string, query string, documents []string,
return nil, fmt.Errorf("model name is required")
}
region := "default"
if apiConfig != nil && apiConfig.Region != nil && *apiConfig.Region != "" {
region = *apiConfig.Region
resolvedBaseURL, err := v.baseModel.GetBaseURL(apiConfig)
if err != nil {
return nil, err
}
baseURL := v.BaseURL[region]
baseURL := resolvedBaseURL
if baseURL == "" {
baseURL = v.BaseURL["default"]
baseURL = resolvedBaseURL
}
if baseURL == "" {
return nil, fmt.Errorf("missing base URL: please configure the local access address for vLLM (e.g., http://127.0.0.1:8000/v1)")
}
url := fmt.Sprintf("%s/%s", strings.TrimSuffix(baseURL, "/"), v.URLSuffix.Rerank)
url := fmt.Sprintf("%s/%s", strings.TrimSuffix(baseURL, "/"), v.baseModel.URLSuffix.Rerank)
topN := len(documents)
if rerankConfig != nil && rerankConfig.TopN > 0 && rerankConfig.TopN < topN {
@@ -633,11 +629,9 @@ func (v *VllmModel) Rerank(modelName *string, query string, documents []string,
}
req.Header.Set("Content-Type", "application/json")
if apiConfig != nil && apiConfig.ApiKey != nil && *apiConfig.ApiKey != "" {
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
}
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := v.httpClient.Do(req)
resp, err := v.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}