mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-11 06:05:45 +08:00
### What problem does this PR solve? Now user can use 'think mode' to chat with LLM ### Type of change - [x] New Feature (non-breaking change which adds functionality) --------- Signed-off-by: Jin Hai <haijin.chn@gmail.com>
666 lines
20 KiB
Go
666 lines
20 KiB
Go
//
|
|
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
//
|
|
|
|
package service
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"ragflow/internal/common"
|
|
"ragflow/internal/dao"
|
|
"ragflow/internal/entity"
|
|
"strings"
|
|
"time"
|
|
|
|
model "ragflow/internal/entity/models"
|
|
"ragflow/internal/service/models"
|
|
)
|
|
|
|
// ModelProvider provides model instances based on tenant and model type
|
|
type ModelProvider interface {
|
|
// GetEmbeddingModel returns an embedding model for the given tenant
|
|
GetEmbeddingModel(ctx context.Context, tenantID string, modelName string) (entity.EmbeddingModel, error)
|
|
// GetChatModel returns a chat model for the given tenant
|
|
GetChatModel(ctx context.Context, tenantID string, modelName string) (entity.ChatModel, error)
|
|
// GetRerankModel returns a rerank model for the given tenant
|
|
GetRerankModel(ctx context.Context, tenantID string, modelName string) (entity.RerankModel, error)
|
|
}
|
|
|
|
// ModelProviderImpl implements ModelProvider
|
|
type ModelProviderImpl struct {
|
|
httpClient *http.Client
|
|
}
|
|
|
|
// NewModelProvider creates a new ModelProvider
|
|
func NewModelProvider() *ModelProviderImpl {
|
|
return &ModelProviderImpl{
|
|
httpClient: &http.Client{
|
|
Timeout: 30 * time.Second,
|
|
},
|
|
}
|
|
}
|
|
|
|
// parseModelName parses a composite model name in format "model_name@provider"
|
|
// Returns modelName and provider separately
|
|
func parseModelName(compositeName string) (modelName, provider string, err error) {
|
|
parts := strings.Split(compositeName, "@")
|
|
if len(parts) == 2 {
|
|
return parts[0], parts[1], nil
|
|
} else if len(parts) == 1 {
|
|
return parts[0], "", fmt.Errorf("provider name missing in model name: %s", compositeName)
|
|
} else {
|
|
return "", "", fmt.Errorf("invalid model name format: %s", compositeName)
|
|
}
|
|
}
|
|
|
|
// GetEmbeddingModel returns an embedding model for the given tenant
|
|
func (p *ModelProviderImpl) GetEmbeddingModel(ctx context.Context, tenantID string, compositeModelName string) (entity.EmbeddingModel, error) {
|
|
// Parse composite model name to extract model name and provider
|
|
modelName, provider, err := parseModelName(compositeModelName)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// Get API key and configuration
|
|
embeddingModel, err := dao.NewTenantLLMDAO().GetByTenantFactoryAndModelName(tenantID, provider, modelName)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
apiKey := embeddingModel.APIKey
|
|
if apiKey == nil || *apiKey == "" {
|
|
return nil, fmt.Errorf("no API key found for tenant %s and model %s", tenantID, compositeModelName)
|
|
}
|
|
// Always get API base from model provider configuration
|
|
providerDAO := dao.NewModelProviderDAO()
|
|
providerConfig := providerDAO.GetProviderByName(provider)
|
|
if providerConfig == nil || providerConfig.DefaultURL == "" {
|
|
return nil, fmt.Errorf("no API base found for provider %s", provider)
|
|
}
|
|
apiBase := fmt.Sprintf("%sembeddings/", providerConfig.DefaultURL)
|
|
|
|
return models.CreateEmbeddingModel(provider, *apiKey, apiBase, modelName, p.httpClient)
|
|
}
|
|
|
|
// GetChatModel returns a chat model for the given tenant
|
|
func (p *ModelProviderImpl) GetChatModel(ctx context.Context, tenantID string, compositeModelName string) (entity.ChatModel, error) {
|
|
// Parse composite model name to extract model name and provider
|
|
_, _, err := parseModelName(compositeModelName)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
// TODO: implement chat model creation
|
|
return nil, fmt.Errorf("chat model not implemented yet for model: %s", compositeModelName)
|
|
}
|
|
|
|
// GetRerankModel returns a rerank model for the given tenant
|
|
func (p *ModelProviderImpl) GetRerankModel(ctx context.Context, tenantID string, compositeModelName string) (entity.RerankModel, error) {
|
|
// Parse composite model name to extract model name and provider
|
|
_, _, err := parseModelName(compositeModelName)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
// TODO: implement rerank model creation
|
|
return nil, fmt.Errorf("rerank model not implemented yet for model: %s", compositeModelName)
|
|
}
|
|
|
|
func NewModelProviderService() *ModelProviderService {
|
|
return &ModelProviderService{
|
|
modelProviderDAO: dao.NewTenantModelProviderDAO(),
|
|
modelInstanceDAO: dao.NewTenantModelInstanceDAO(),
|
|
modelDAO: dao.NewTenantModelDAO(),
|
|
modelGroupDAO: dao.NewTenantModelGroupDAO(),
|
|
modelGroupMappingDAO: dao.NewTenantModelGroupMappingDAO(),
|
|
userTenantDAO: dao.NewUserTenantDAO(),
|
|
}
|
|
}
|
|
|
|
type ModelProviderService struct {
|
|
modelProviderDAO *dao.TenantModelProviderDAO
|
|
modelInstanceDAO *dao.TenantModelInstanceDAO
|
|
modelDAO *dao.TenantModelDAO
|
|
modelGroupDAO *dao.TenantModelGroupDAO
|
|
modelGroupMappingDAO *dao.TenantModelGroupMappingDAO
|
|
userTenantDAO *dao.UserTenantDAO
|
|
}
|
|
|
|
func (m *ModelProviderService) AddModelProvider(providerName, userID string) (common.ErrorCode, error) {
|
|
|
|
_, err := dao.GetModelProviderManager().GetProviderByName(providerName)
|
|
if err != nil {
|
|
return common.CodeNotFound, err
|
|
}
|
|
|
|
tenants, err := m.userTenantDAO.GetByUserIDAndRole(userID, "owner")
|
|
if err != nil {
|
|
return common.CodeServerError, err
|
|
}
|
|
|
|
if len(tenants) == 0 {
|
|
return common.CodeNotFound, errors.New("user has no tenants")
|
|
}
|
|
|
|
tenantID := tenants[0].TenantID
|
|
|
|
providerID, err := generateUUID1Hex()
|
|
if err != nil {
|
|
return common.CodeServerError, errors.New("fail to get UUID")
|
|
}
|
|
|
|
now := time.Now().Unix()
|
|
nowDate := time.Now().Truncate(time.Second)
|
|
tenantModelProvider := &entity.TenantModelProvider{
|
|
ID: providerID,
|
|
ProviderName: providerName,
|
|
TenantID: tenantID,
|
|
}
|
|
tenantModelProvider.CreateTime = &now
|
|
tenantModelProvider.UpdateTime = &now
|
|
tenantModelProvider.CreateDate = &nowDate
|
|
tenantModelProvider.UpdateDate = &nowDate
|
|
err = m.modelProviderDAO.Create(tenantModelProvider)
|
|
if err != nil {
|
|
return common.CodeServerError, errors.New("fail to create model provider")
|
|
}
|
|
return common.CodeSuccess, nil
|
|
}
|
|
|
|
func (m *ModelProviderService) ListProvidersOfTenant(userID string) ([]map[string]interface{}, common.ErrorCode, error) {
|
|
|
|
tenants, err := m.userTenantDAO.GetByUserIDAndRole(userID, "owner")
|
|
if err != nil {
|
|
return nil, common.CodeServerError, err
|
|
}
|
|
|
|
if len(tenants) == 0 {
|
|
return nil, common.CodeNotFound, errors.New("user has no tenants")
|
|
}
|
|
|
|
tenantID := tenants[0].TenantID
|
|
|
|
providerNames, err := m.modelProviderDAO.ListByID(tenantID)
|
|
if err != nil {
|
|
return nil, common.CodeServerError, err
|
|
}
|
|
|
|
var result []map[string]interface{}
|
|
for _, providerName := range providerNames {
|
|
provider, err := dao.GetModelProviderManager().GetProviderByName(providerName)
|
|
if err != nil {
|
|
return nil, common.CodeServerError, err
|
|
}
|
|
result = append(result, provider)
|
|
}
|
|
|
|
return result, common.CodeSuccess, nil
|
|
}
|
|
|
|
func (m *ModelProviderService) DeleteModelProvider(providerName, userID string) (common.ErrorCode, error) {
|
|
tenants, err := m.userTenantDAO.GetByUserIDAndRole(userID, "owner")
|
|
if err != nil {
|
|
return common.CodeServerError, err
|
|
}
|
|
if len(tenants) == 0 {
|
|
return common.CodeNotFound, errors.New("user has no tenants")
|
|
}
|
|
tenantID := tenants[0].TenantID
|
|
|
|
_, err = m.modelProviderDAO.DeleteByTenantIDAndProviderName(tenantID, providerName)
|
|
if err != nil {
|
|
return common.CodeServerError, err
|
|
}
|
|
|
|
return common.CodeSuccess, nil
|
|
}
|
|
|
|
func (m *ModelProviderService) CreateProviderInstance(providerName, instanceName, apiKey, userID string) (common.ErrorCode, error) {
|
|
// Get tenant ID from user
|
|
tenants, err := m.userTenantDAO.GetByUserIDAndRole(userID, "owner")
|
|
if err != nil {
|
|
return common.CodeServerError, err
|
|
}
|
|
|
|
if len(tenants) == 0 {
|
|
return common.CodeNotFound, errors.New("user has no tenants")
|
|
}
|
|
|
|
tenantID := tenants[0].TenantID
|
|
|
|
// Check if provider exists
|
|
provider, err := m.modelProviderDAO.GetByTenantIDAndProviderName(tenantID, providerName)
|
|
if err != nil {
|
|
return common.CodeServerError, err
|
|
}
|
|
|
|
instanceID, err := generateUUID1Hex()
|
|
if err != nil {
|
|
return common.CodeServerError, errors.New("fail to get UUID")
|
|
}
|
|
|
|
now := time.Now().Unix()
|
|
nowDate := time.Now().Truncate(time.Second)
|
|
tenantModelProvider := &entity.TenantModelInstance{
|
|
ID: instanceID,
|
|
InstanceName: instanceName,
|
|
ProviderID: provider.ID,
|
|
APIKey: apiKey,
|
|
Status: "active",
|
|
}
|
|
tenantModelProvider.CreateTime = &now
|
|
tenantModelProvider.UpdateTime = &now
|
|
tenantModelProvider.CreateDate = &nowDate
|
|
tenantModelProvider.UpdateDate = &nowDate
|
|
err = m.modelInstanceDAO.Create(tenantModelProvider)
|
|
|
|
if err != nil {
|
|
return common.CodeServerError, errors.New("fail to create model provider")
|
|
}
|
|
return common.CodeSuccess, nil
|
|
}
|
|
|
|
func (m *ModelProviderService) ListProviderInstances(providerName, userID string) ([]map[string]interface{}, common.ErrorCode, error) {
|
|
|
|
// Get tenant ID from user
|
|
tenants, err := m.userTenantDAO.GetByUserIDAndRole(userID, "owner")
|
|
if err != nil {
|
|
return nil, common.CodeServerError, err
|
|
}
|
|
|
|
if len(tenants) == 0 {
|
|
return nil, common.CodeNotFound, errors.New("user has no tenants")
|
|
}
|
|
|
|
tenantID := tenants[0].TenantID
|
|
|
|
// Check if provider exists
|
|
provider, err := m.modelProviderDAO.GetByTenantIDAndProviderName(tenantID, providerName)
|
|
if err != nil {
|
|
return nil, common.CodeServerError, err
|
|
}
|
|
|
|
// Check if provider exists
|
|
instances, err := m.modelInstanceDAO.GetAllInstancesByProviderID(provider.ID)
|
|
if err != nil {
|
|
return nil, common.CodeServerError, err
|
|
}
|
|
|
|
var result []map[string]interface{}
|
|
for _, instance := range instances {
|
|
result = append(result, map[string]interface{}{
|
|
"id": instance.ID,
|
|
"instanceName": instance.InstanceName,
|
|
"providerID": instance.ProviderID,
|
|
"apiKey": instance.APIKey,
|
|
"status": instance.Status,
|
|
})
|
|
}
|
|
|
|
return result, common.CodeSuccess, nil
|
|
}
|
|
|
|
func (m *ModelProviderService) ShowProviderInstance(providerName, instanceName, userID string) (map[string]interface{}, common.ErrorCode, error) {
|
|
|
|
// Get tenant ID from user
|
|
tenants, err := m.userTenantDAO.GetByUserIDAndRole(userID, "owner")
|
|
if err != nil {
|
|
return nil, common.CodeServerError, err
|
|
}
|
|
|
|
if len(tenants) == 0 {
|
|
return nil, common.CodeNotFound, errors.New("user has no tenants")
|
|
}
|
|
|
|
tenantID := tenants[0].TenantID
|
|
|
|
// Check if provider exists
|
|
provider, err := m.modelProviderDAO.GetByTenantIDAndProviderName(tenantID, providerName)
|
|
if err != nil {
|
|
return nil, common.CodeServerError, err
|
|
}
|
|
|
|
instance, err := m.modelInstanceDAO.GetByProviderIDAndInstanceName(provider.ID, instanceName)
|
|
if err != nil {
|
|
return nil, common.CodeServerError, err
|
|
}
|
|
|
|
result := map[string]interface{}{
|
|
"id": instance.ID,
|
|
"instanceName": instance.InstanceName,
|
|
"providerID": instance.ProviderID,
|
|
"status": instance.Status,
|
|
}
|
|
|
|
return result, common.CodeSuccess, nil
|
|
}
|
|
|
|
func (m *ModelProviderService) AlterProviderInstance(providerName, instanceName, newInstanceName, apiKey, userID string) (common.ErrorCode, error) {
|
|
return common.CodeSuccess, nil
|
|
}
|
|
func (m *ModelProviderService) DropProviderInstance(providerName, instanceName, userID string) (common.ErrorCode, error) {
|
|
|
|
// Get tenant ID from user
|
|
tenants, err := m.userTenantDAO.GetByUserIDAndRole(userID, "owner")
|
|
if err != nil {
|
|
return common.CodeServerError, err
|
|
}
|
|
|
|
if len(tenants) == 0 {
|
|
return common.CodeNotFound, errors.New("user has no tenants")
|
|
}
|
|
|
|
tenantID := tenants[0].TenantID
|
|
|
|
// Check if provider exists
|
|
provider, err := m.modelProviderDAO.GetByTenantIDAndProviderName(tenantID, providerName)
|
|
if err != nil {
|
|
return common.CodeServerError, err
|
|
}
|
|
|
|
count, err := m.modelInstanceDAO.DeleteByProviderIDAndInstanceName(provider.ID, instanceName)
|
|
if err != nil {
|
|
return common.CodeServerError, err
|
|
}
|
|
|
|
if count == 0 {
|
|
return common.CodeNotFound, errors.New("provider instance not found")
|
|
}
|
|
|
|
return common.CodeSuccess, nil
|
|
}
|
|
|
|
func (m *ModelProviderService) ListInstanceModels(providerName, instanceName, userID string) ([]map[string]interface{}, error) {
|
|
// Get tenant ID from user
|
|
tenants, err := m.userTenantDAO.GetByUserIDAndRole(userID, "owner")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if len(tenants) == 0 {
|
|
return nil, errors.New("user has no tenants")
|
|
}
|
|
|
|
tenantID := tenants[0].TenantID
|
|
|
|
// Check if provider exists
|
|
provider, err := m.modelProviderDAO.GetByTenantIDAndProviderName(tenantID, providerName)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// Get instance
|
|
instance, err := m.modelInstanceDAO.GetByProviderIDAndInstanceName(provider.ID, instanceName)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// Get all models for this instance
|
|
disabledModels, err := m.modelDAO.GetModelsByInstanceID(instance.ID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// insert models name into a set
|
|
modelNames := make(map[string]bool)
|
|
for _, model := range disabledModels {
|
|
modelNames[model.ModelName] = true
|
|
}
|
|
|
|
allModels, err := dao.GetModelProviderManager().ListModels(providerName)
|
|
|
|
for _, model := range allModels {
|
|
// convert model["name"] to string
|
|
modelName := model["name"].(string)
|
|
if modelNames[modelName] {
|
|
model["status"] = "disabled"
|
|
} else {
|
|
model["status"] = "enabled"
|
|
}
|
|
|
|
}
|
|
|
|
return allModels, nil
|
|
}
|
|
|
|
func (m *ModelProviderService) UpdateModelStatus(providerName, instanceName, modelName, userID, status string) (common.ErrorCode, error) {
|
|
|
|
// Get tenant ID from user
|
|
tenants, err := m.userTenantDAO.GetByUserIDAndRole(userID, "owner")
|
|
if err != nil {
|
|
return common.CodeServerError, err
|
|
}
|
|
|
|
if len(tenants) == 0 {
|
|
return common.CodeNotFound, errors.New("user has no tenants")
|
|
}
|
|
|
|
tenantID := tenants[0].TenantID
|
|
|
|
// Check if provider exists
|
|
provider, err := m.modelProviderDAO.GetByTenantIDAndProviderName(tenantID, providerName)
|
|
if err != nil {
|
|
return common.CodeServerError, err
|
|
}
|
|
|
|
instance, err := m.modelInstanceDAO.GetByProviderIDAndInstanceName(provider.ID, instanceName)
|
|
if err != nil {
|
|
return common.CodeServerError, err
|
|
}
|
|
|
|
model, err := m.modelDAO.GetModelByProviderIDAndInstanceIDAndModelName(provider.ID, instance.ID, modelName)
|
|
if err != nil {
|
|
var modelID string
|
|
modelID, err = generateUUID1Hex()
|
|
if err != nil {
|
|
return common.CodeServerError, errors.New("fail to get UUID")
|
|
}
|
|
// Get model info from provider
|
|
model = &entity.TenantModel{
|
|
ID: modelID,
|
|
ModelName: modelName,
|
|
ModelType: model.ModelType,
|
|
ProviderID: provider.ID,
|
|
InstanceID: instance.ID,
|
|
Status: status,
|
|
}
|
|
err = m.modelDAO.Create(model)
|
|
if err != nil {
|
|
return common.CodeServerError, errors.New("fail to create model")
|
|
}
|
|
return common.CodeSuccess, nil
|
|
}
|
|
|
|
count, err := m.modelDAO.DeleteByModelID(model.ID)
|
|
if err != nil {
|
|
return common.CodeServerError, err
|
|
}
|
|
if count == 0 {
|
|
return common.CodeNotFound, errors.New("model not found")
|
|
}
|
|
|
|
return common.CodeSuccess, nil
|
|
}
|
|
|
|
func (m *ModelProviderService) ChatToModel(providerName, instanceName, modelName, userID, message string) (*string, common.ErrorCode, error) {
|
|
|
|
// Get tenant ID from user
|
|
tenants, err := m.userTenantDAO.GetByUserIDAndRole(userID, "owner")
|
|
if err != nil {
|
|
return nil, common.CodeServerError, err
|
|
}
|
|
|
|
if len(tenants) == 0 {
|
|
return nil, common.CodeNotFound, errors.New("user has no tenants")
|
|
}
|
|
|
|
tenantID := tenants[0].TenantID
|
|
|
|
// Check if provider exists
|
|
provider, err := m.modelProviderDAO.GetByTenantIDAndProviderName(tenantID, providerName)
|
|
if err != nil {
|
|
return nil, common.CodeServerError, err
|
|
}
|
|
|
|
instance, err := m.modelInstanceDAO.GetByProviderIDAndInstanceName(provider.ID, instanceName)
|
|
if err != nil {
|
|
return nil, common.CodeServerError, err
|
|
}
|
|
|
|
_, err = m.modelDAO.GetModelByProviderIDAndInstanceIDAndModelName(provider.ID, instance.ID, modelName)
|
|
if err != nil {
|
|
providerInfo := dao.GetModelProviderManager().FindProvider(providerName)
|
|
if providerInfo == nil {
|
|
return nil, common.CodeNotFound, errors.New("provider not found")
|
|
}
|
|
|
|
_, err = dao.GetModelProviderManager().GetModelByName(providerName, modelName)
|
|
if err != nil {
|
|
return nil, common.CodeNotFound, errors.New(fmt.Sprintf("provider %s model %s not found", providerName, modelName))
|
|
}
|
|
|
|
var response string
|
|
response, err = providerInfo.ModelDriver.Chat(&modelName, &instance.APIKey, &message, nil)
|
|
if err != nil {
|
|
return nil, common.CodeServerError, err
|
|
}
|
|
|
|
return &response, common.CodeSuccess, nil
|
|
}
|
|
|
|
return nil, common.CodeServerError, errors.New("model is disabled")
|
|
}
|
|
|
|
// ChatToModelStream
|
|
func (m *ModelProviderService) ChatToModelStream(providerName, instanceName, modelName, userID, message string) (<-chan string, <-chan error, common.ErrorCode, error) {
|
|
streamChan := make(chan string)
|
|
errChan := make(chan error, 1)
|
|
|
|
// Get tenant ID from user
|
|
tenants, err := m.userTenantDAO.GetByUserIDAndRole(userID, "owner")
|
|
if err != nil {
|
|
close(streamChan)
|
|
close(errChan)
|
|
return streamChan, errChan, common.CodeServerError, err
|
|
}
|
|
|
|
if len(tenants) == 0 {
|
|
close(streamChan)
|
|
close(errChan)
|
|
return streamChan, errChan, common.CodeNotFound, errors.New("user has no tenants")
|
|
}
|
|
|
|
tenantID := tenants[0].TenantID
|
|
|
|
// Check if provider exists
|
|
provider, err := m.modelProviderDAO.GetByTenantIDAndProviderName(tenantID, providerName)
|
|
if err != nil {
|
|
close(streamChan)
|
|
close(errChan)
|
|
return streamChan, errChan, common.CodeServerError, err
|
|
}
|
|
|
|
instance, err := m.modelInstanceDAO.GetByProviderIDAndInstanceName(provider.ID, instanceName)
|
|
if err != nil {
|
|
close(streamChan)
|
|
close(errChan)
|
|
return streamChan, errChan, common.CodeServerError, err
|
|
}
|
|
|
|
_, err = m.modelDAO.GetModelByProviderIDAndInstanceIDAndModelName(provider.ID, instance.ID, modelName)
|
|
if err != nil {
|
|
providerInfo := dao.GetModelProviderManager().FindProvider(providerName)
|
|
if providerInfo == nil {
|
|
close(streamChan)
|
|
close(errChan)
|
|
return streamChan, errChan, common.CodeNotFound, errors.New("provider not found")
|
|
}
|
|
|
|
_, err = dao.GetModelProviderManager().GetModelByName(providerName, modelName)
|
|
if err != nil {
|
|
close(streamChan)
|
|
close(errChan)
|
|
return streamChan, errChan, common.CodeNotFound, errors.New(fmt.Sprintf("provider %s model %s not found", providerName, modelName))
|
|
}
|
|
|
|
// Async call stream interface using channel for better performance
|
|
go func() {
|
|
defer close(streamChan)
|
|
defer close(errChan)
|
|
|
|
err := providerInfo.ModelDriver.ChatStreamlyWithChannel(&modelName, &instance.APIKey, &message, nil, streamChan)
|
|
if err != nil {
|
|
errChan <- err
|
|
}
|
|
}()
|
|
|
|
return streamChan, errChan, common.CodeSuccess, nil
|
|
}
|
|
|
|
close(streamChan)
|
|
close(errChan)
|
|
return streamChan, errChan, common.CodeServerError, errors.New("model is disabled")
|
|
}
|
|
|
|
// ChatToModelStreamWithSender streams chat response directly via sender function (best performance, no channel)
|
|
func (m *ModelProviderService) ChatToModelStreamWithSender(providerName, instanceName, modelName, userID, message string, modelConfig *model.ChatConfig, sender func(*string, *string) error) common.ErrorCode {
|
|
// Get tenant ID from user
|
|
tenants, err := m.userTenantDAO.GetByUserIDAndRole(userID, "owner")
|
|
if err != nil {
|
|
return common.CodeServerError
|
|
}
|
|
|
|
if len(tenants) == 0 {
|
|
return common.CodeNotFound
|
|
}
|
|
|
|
tenantID := tenants[0].TenantID
|
|
|
|
// Check if provider exists
|
|
provider, err := m.modelProviderDAO.GetByTenantIDAndProviderName(tenantID, providerName)
|
|
if err != nil {
|
|
return common.CodeServerError
|
|
}
|
|
|
|
instance, err := m.modelInstanceDAO.GetByProviderIDAndInstanceName(provider.ID, instanceName)
|
|
if err != nil {
|
|
return common.CodeServerError
|
|
}
|
|
|
|
_, err = m.modelDAO.GetModelByProviderIDAndInstanceIDAndModelName(provider.ID, instance.ID, modelName)
|
|
if err != nil {
|
|
providerInfo := dao.GetModelProviderManager().FindProvider(providerName)
|
|
if providerInfo == nil {
|
|
return common.CodeNotFound
|
|
}
|
|
|
|
_, err = dao.GetModelProviderManager().GetModelByName(providerName, modelName)
|
|
if err != nil {
|
|
return common.CodeNotFound
|
|
}
|
|
|
|
// Direct call with sender function
|
|
err := providerInfo.ModelDriver.ChatStreamlyWithSender(&modelName, &instance.APIKey, &message, modelConfig, sender)
|
|
if err != nil {
|
|
return common.CodeServerError
|
|
}
|
|
|
|
return common.CodeSuccess
|
|
}
|
|
|
|
return common.CodeServerError
|
|
}
|