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

fix(go-models): validate embedding request limits (#17919)

Browse Source 
## Summary

- Add embedding batch-size metadata to model responses and tenant
overrides.
- Validate embedding dimensions and batch limits across provider
verification and embedding requests.
- Expand validation tests for defaults, limits, and missing metadata.

---------

Signed-off-by: Jin Hai <haijin.chn@gmail.com>
Co-authored-by: Jin Hai <haijin.chn@gmail.com>
This commit is contained in:
Hz_
2026-08-06 16:18:12 +08:00
committed by GitHub
parent 60df86bfa2
commit ef0b293271
6 changed files with 228 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
internal/entity/models/base_model.go
View File

@@ -400,6 +400,7 @@ func ParseListModel(modelList ModelList) []ListModelResponse {
modelResponse.Name = modelName
if modelEntity != nil {
modelResponse.MaxDimension = modelEntity.MaxDimension
modelResponse.MaxBatchSize = modelEntity.MaxBatchSize
modelResponse.Dimensions = modelEntity.Dimensions
modelResponse.MaxOutput = modelEntity.MaxOutput
modelResponse.ModelTypes = modelEntity.ModelTypes

33
internal/entity/models/model.go
View File

@@ -166,7 +166,8 @@ type Model struct {
Thinking *ModelThinking `json:"thinking"`
Tools *ModelTools `json:"tools"`
Class *string `json:"class"`
MaxDimension *int `json:"max_dimension"` // used by embedding models
MaxDimension *int `json:"max_dimension"` // used by embedding models
MaxBatchSize *int `json:"max_batch_size"` // used by embedding models
Dimensions []int `json:"dimensions"`
BatchSize *int `json:"batch_size"` // max texts per Embed request; used by embedding models
Alias []string `json:"alias"`
@@ -417,6 +418,9 @@ func (pm *ProviderManager) ListAllModels() ([]map[string]interface{}, error) {
if model.MaxDimension != nil {
modelData["max_dimension"] = *model.MaxDimension
}
if model.MaxBatchSize != nil {
modelData["max_batch_size"] = *model.MaxBatchSize
}
if len(model.Dimensions) > 0 {
modelData["dimensions"] = model.Dimensions
}
@@ -516,11 +520,12 @@ func (pm *ProviderManager) ListModels(providerName string) ([]map[string]interfa
// keep the response shape stable for clients that destructure
// the object.
modelData := map[string]interface{}{
"name": model.Name,
"max_output": model.MaxOutput,
"model_types": model.ModelTypes,
"max_dimension": model.MaxDimension,
"dimensions": model.Dimensions,
"name": model.Name,
"max_output": model.MaxOutput,
"model_types": model.ModelTypes,
"max_dimension": model.MaxDimension,
"max_batch_size": model.MaxBatchSize,
"dimensions": model.Dimensions,
}
if model.BatchSize != nil {
modelData["batch_size"] = *model.BatchSize
@@ -634,9 +639,10 @@ func (pm *ProviderManager) SearchModelInfo(providerName, modelName string, filte
if matchFilter {
modelData := map[string]interface{}{
"name": model.Name,
"max_output": model.MaxOutput,
"model_types": model.ModelTypes,
"name": model.Name,
"max_output": model.MaxOutput,
"model_types": model.ModelTypes,
"max_batch_size": model.MaxBatchSize,
//"features": getFeaturesMap(model.Features),
}
@@ -694,10 +700,11 @@ func (pm *ProviderManager) SearchByType(modelType string) ModelResponse {
for _, model := range provider.Models {
if containsModelType(model.ModelTypes, modelType) {
modelData := map[string]interface{}{
"provider": provider.Name,
"name": model.Name,
"max_output": model.MaxOutput,
"model_types": model.ModelTypes,
"provider": provider.Name,
"name": model.Name,
"max_output": model.MaxOutput,
"model_types": model.ModelTypes,
"max_batch_size": model.MaxBatchSize,
//"features": getFeaturesMap(model.Features),
}
resp.Data = append(resp.Data, modelData)

1
internal/entity/models/nvidia.go
View File

@@ -694,6 +694,7 @@ func parseNvidiaModelList(modelList ModelList, provider *Provider) []ListModelRe
response.ModelTypes = append([]string(nil), preset.ModelTypes...)
response.Thinking = preset.Thinking
response.MaxDimension = preset.MaxDimension
response.MaxBatchSize = preset.MaxBatchSize
response.Dimensions = append([]int(nil), preset.Dimensions...)
} else {
maxTokens := defaultMaxTokens

3
internal/entity/models/types.go
View File

@@ -110,7 +110,8 @@ type ListModelResponse struct {
MaxOutput *int `json:"max_output"`
ModelTypes []string `json:"model_types"`
Thinking *ModelThinking `json:"thinking"`
MaxDimension *int `json:"max_dimension"` // used by embedding models
MaxDimension *int `json:"max_dimension"` // used by embedding models
MaxBatchSize *int `json:"max_batch_size"` // used by embedding models
Dimensions []int `json:"dimensions"`
}