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

Go: introduce content_length and max_output (#17807)

Browse Source 
Signed-off-by: Jin Hai <haijin.chn@gmail.com>
This commit is contained in:
Jin Hai
2026-08-04 19:08:31 +08:00
committed by GitHub
parent f707cca074
commit c1f960cd47
19 changed files with 254 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files

43
internal/service/model_service.go
View File

@@ -396,12 +396,13 @@ func (m *ModelProviderService) ListSupportedModels(ctx context.Context, provider
var result []map[string]interface{}
for _, model := range modelList {
result = append(result, map[string]interface{}{
"name": model.Name,
"max_dimension": model.MaxDimension,
"dimensions": model.Dimensions,
"max_tokens": model.MaxTokens,
"model_types": model.ModelTypes,
"thinking": model.Thinking,
"name": model.Name,
"max_dimension": model.MaxDimension,
"dimensions": model.Dimensions,
"content_length": model.ContentLength,
"max_output": model.MaxOutput,
"model_types": model.ModelTypes,
"thinking": model.Thinking,
})
}
return result, nil
@@ -450,8 +451,8 @@ func (m *ModelProviderService) reconcileNvidiaInstanceModels(
for _, remote := range normalized {
maxTokens := 8192
if remote.MaxTokens != nil && *remote.MaxTokens > 0 {
maxTokens = *remote.MaxTokens
if remote.MaxOutput != nil && *remote.MaxOutput > 0 {
maxTokens = *remote.MaxOutput
}
modelType := int(entity.ModelTypeFromStrings(remote.ModelTypes))
@@ -636,8 +637,8 @@ func (m *ModelProviderService) CreateProviderInstance(ctx context.Context, provi
ModelName: llm.Name,
ModelTypes: llm.ModelTypes,
MaxTokens: func() int {
if llm.MaxTokens != nil {
return *llm.MaxTokens
if llm.MaxOutput != nil {
return *llm.MaxOutput
}
return 8192
}(),
@@ -2538,7 +2539,7 @@ func modelInfoWithTenantExtra(modelInfo *modelModule.Model, modelEntity *entity.
}
if extra.MaxTokens != nil && *extra.MaxTokens > 0 {
model.MaxTokens = extra.MaxTokens
model.MaxOutput = extra.MaxTokens
}
if len(extra.ModelTypes) > 0 {
model.ModelTypes = append([]string(nil), extra.ModelTypes...)
@@ -3483,8 +3484,8 @@ func (m *ModelProviderService) GetModelConfigByID(ctx context.Context, userID st
maxTokens := 0
if mi, _ := dao.GetModelProviderManager().GetModelByName(providerEntity.ProviderName, modelEntity.ModelName); mi != nil {
if mi.MaxTokens != nil {
maxTokens = *mi.MaxTokens
if mi.MaxOutput != nil {
maxTokens = *mi.MaxOutput
}
}
maxTokens, err = maxTokensFromTenantModelExtra(modelEntity, maxTokens)
@@ -3879,10 +3880,10 @@ func (m *ModelProviderService) GetModelConfigFromProviderInstance(ctx context.Co
apiConfig := &modelModule.APIConfig{ApiKey: &apiKey, Region: &region}
maxTokens := 0
if mi, _ := dao.GetModelProviderManager().GetModelByName("Builtin", pureModelName); mi != nil {
if mi.MaxTokens == nil {
if mi.MaxOutput == nil {
maxTokens = 0
} else {
maxTokens = *mi.MaxTokens
maxTokens = *mi.MaxOutput
}
}
return builtinDriver, pureModelName, apiConfig, maxTokens, nil
@@ -3941,10 +3942,10 @@ func (m *ModelProviderService) GetModelConfigFromProviderInstance(ctx context.Co
}
maxTokens := 0
if mi, _ := dao.GetModelProviderManager().GetModelByName(providerName, pureModelName); mi != nil {
if mi.MaxTokens == nil {
if mi.MaxOutput == nil {
maxTokens = 0
} else {
maxTokens = *mi.MaxTokens
maxTokens = *mi.MaxOutput
}
}
maxTokens, driverErr = maxTokensFromTenantModelExtra(modelObj, maxTokens)
@@ -3999,8 +4000,8 @@ func (m *ModelProviderService) GetModelConfigFromProviderInstance(ctx context.Co
}
apiConfig := &modelModule.APIConfig{ApiKey: &apiKey, Region: &region, BaseURL: &baseURL}
maxTokens := 0
if llmInfo.MaxTokens != nil {
maxTokens = *llmInfo.MaxTokens
if llmInfo.MaxOutput != nil {
maxTokens = *llmInfo.MaxOutput
}
return driver, llmInfo.Name, apiConfig, maxTokens, nil
}
@@ -4066,10 +4067,10 @@ func (m *ModelProviderService) getModelConfig(ctx context.Context, tenantID, com
modelInfo, err := dao.GetModelProviderManager().GetModelByName(providerName, modelName)
maxTokens := 0
if err == nil && modelInfo != nil {
if modelInfo.MaxTokens == nil {
if modelInfo.MaxOutput == nil {
maxTokens = 0
} else {
maxTokens = *modelInfo.MaxTokens
maxTokens = *modelInfo.MaxOutput
}
}

6
internal/service/model_service_test.go
View File

@@ -282,8 +282,8 @@ func TestReconcileNvidiaInstanceModelsAddsUpdatesAndDeletes(t *testing.T) {
maxTokens := 131072
maxDimension := 2048
remote := []modelModule.ListModelResponse{
{Name: "nvidia/keep", MaxTokens: &maxTokens, ModelTypes: []string{"chat", "vision"}},
{Name: "nvidia/new-embed", MaxTokens: ptrService(8192), MaxDimension: &maxDimension, Dimensions: []int{1024, 2048}, ModelTypes: []string{"embedding"}},
{Name: "nvidia/keep", MaxOutput: &maxTokens, ModelTypes: []string{"chat", "vision"}},
{Name: "nvidia/new-embed", MaxOutput: ptrService(8192), MaxDimension: &maxDimension, Dimensions: []int{1024, 2048}, ModelTypes: []string{"embedding"}},
}
err := NewModelProviderService().reconcileNvidiaInstanceModels(context.Background(), db, provider, instance, remote)
@@ -292,7 +292,7 @@ func TestReconcileNvidiaInstanceModelsAddsUpdatesAndDeletes(t *testing.T) {
}
var got []*entity.TenantModel
if err := db.Order("model_name").Find(&got).Error; err != nil {
if err = db.Order("model_name").Find(&got).Error; err != nil {
t.Fatalf("list models: %v", err)
}
if len(got) != 2 || got[0].ModelName != "nvidia/keep" || got[1].ModelName != "nvidia/new-embed" {