diff --git a/internal/service/chat_session.go b/internal/service/chat_session.go index 5a65067a05..5eaabc4553 100644 --- a/internal/service/chat_session.go +++ b/internal/service/chat_session.go @@ -24,6 +24,7 @@ import ( "math" "ragflow/internal/common" "ragflow/internal/engine" + modelModule "ragflow/internal/entity/models" "ragflow/internal/storage" "ragflow/internal/utility" "strconv" @@ -58,6 +59,10 @@ type chatPipelineRunner interface { AsyncChat(ctx context.Context, userID string, chat *entity.Chat, messages []map[string]interface{}, stream bool, kwargs map[string]interface{}) (<-chan AsyncChatResult, error) } +type chatModelConfigResolver interface { + GetChatModelConfig(tenantID, llmID string) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) +} + // chunkFeedbackApplier is the dispatch seam for chunk-level feedback // persistence. Mirrors the Python ChunkFeedbackService.apply_feedback // (api/db/services/chunk_feedback_service.py) call site at @@ -80,6 +85,7 @@ type ChatSessionService struct { chatSessionDAO chatSessionStore userTenantDAO userTenantStore pipeline chatPipelineRunner + modelProviderSvc chatModelConfigResolver chunkFeedbackApplier chunkFeedbackApplier docEngine engine.DocEngine } @@ -87,10 +93,11 @@ type ChatSessionService struct { // NewChatSessionService create chat session service func NewChatSessionService() *ChatSessionService { return &ChatSessionService{ - chatSessionDAO: dao.NewChatSessionDAO(), - userTenantDAO: dao.NewUserTenantDAO(), - pipeline: NewChatPipelineService(), - docEngine: engine.Get(), + chatSessionDAO: dao.NewChatSessionDAO(), + userTenantDAO: dao.NewUserTenantDAO(), + pipeline: NewChatPipelineService(), + modelProviderSvc: NewModelProviderService(), + docEngine: engine.Get(), } } @@ -1936,7 +1943,11 @@ func (s *ChatSessionService) initializeReference(session *entity.ChatSession) [] } func (s *ChatSessionService) checkTenantLLMAPIKey(tenantID, modelName string) (bool, error) { - _, err := NewTenantLLMService().GetAPIKeyFromInstance(tenantID, modelName) + resolver := s.modelProviderSvc + if resolver == nil { + resolver = NewModelProviderService() + } + _, _, _, _, err := resolver.GetChatModelConfig(tenantID, modelName) if err != nil { return false, err } diff --git a/internal/service/chat_session_test.go b/internal/service/chat_session_test.go index c7060461d3..795fc1d327 100644 --- a/internal/service/chat_session_test.go +++ b/internal/service/chat_session_test.go @@ -14,6 +14,7 @@ import ( "ragflow/internal/common" "ragflow/internal/engine" "ragflow/internal/entity" + modelModule "ragflow/internal/entity/models" ) // --------------------------------------------------------------------------- @@ -175,6 +176,21 @@ func makeResultChan(results ...AsyncChatResult) <-chan AsyncChatResult { return ch } +type fakeChatModelConfigResolver struct { + tenantID string + llmID string + err error +} + +func (f *fakeChatModelConfigResolver) GetChatModelConfig(tenantID, llmID string) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) { + f.tenantID = tenantID + f.llmID = llmID + if f.err != nil { + return nil, "", nil, 0, f.err + } + return nil, "resolved-model", &modelModule.APIConfig{}, 8192, nil +} + type feedbackContextKey struct{} type fakeFeedbackDocEngine struct { @@ -1020,6 +1036,65 @@ func TestChatCompletionsStreamFinalCarriesDecoratedReference(t *testing.T) { } } +func TestChatCompletionsModelIDOverrideUsesModelResolver(t *testing.T) { + store := newFakeSessionStore() + store.sessions["session-1"] = &entity.ChatSession{ + ID: "session-1", + DialogID: "dialog-1", + Message: json.RawMessage(`[]`), + Reference: json.RawMessage(`[]`), + } + store.dialogs["dialog-1"] = &entity.Chat{ + ID: "dialog-1", + TenantID: "tenant-owner", + LLMID: "old-model@default@Provider", + LLMSetting: entity.JSONMap{}, + PromptConfig: entity.JSONMap{"parameters": []interface{}{}}, + } + store.dialogExists["tenant-owner|dialog-1"] = true + + pipeline := &fakePipeline{ + resultChan: makeResultChan( + AsyncChatResult{Answer: "ok", Final: true, Reference: map[string]interface{}{"chunks": []interface{}{}}}, + ), + } + resolver := &fakeChatModelConfigResolver{} + modelID := "3d2d824e7e5d11f1a845455b140cef90" + + svc := &ChatSessionService{ + chatSessionDAO: store, + userTenantDAO: &fakeTenantStore{tenantIDs: []string{"tenant-owner"}}, + pipeline: pipeline, + modelProviderSvc: resolver, + } + + _, err := svc.ChatCompletions( + context.Background(), + "user-1", + "dialog-1", + "session-1", + []map[string]interface{}{{"role": "user", "content": "hi"}}, + "", + nil, + modelID, + nil, + nil, + false, + false, + false, + nil, + ) + if err != nil { + t.Fatalf("ChatCompletions failed: %v", err) + } + if resolver.tenantID != "tenant-owner" || resolver.llmID != modelID { + t.Fatalf("resolver got tenantID=%q llmID=%q, want tenant-owner/%s", resolver.tenantID, resolver.llmID, modelID) + } + if store.dialogs["dialog-1"].LLMID != modelID { + t.Fatalf("dialog LLMID=%q, want model id %q", store.dialogs["dialog-1"].LLMID, modelID) + } +} + func TestCompletion_EmptyMessages(t *testing.T) { svc := &ChatSessionService{ chatSessionDAO: &fakeSessionStore{},