fix: concurrent map writes due to dig.Container is not thread safe

Signed-off-by: joeyczheng <[email protected]>

Author: chengjoey

internal/application/service/model.go

@@ -20,13 +20,15 @@ var ErrModelNotFound = errors.New("model not found")
 type modelService struct {
 	repo          interfaces.ModelRepository
 	ollamaService *ollama.OllamaService
+	pooler        embedding.EmbedderPooler
 }
 
 // NewModelService creates a new model service instance
-func NewModelService(repo interfaces.ModelRepository, ollamaService *ollama.OllamaService) interfaces.ModelService {
+func NewModelService(repo interfaces.ModelRepository, ollamaService *ollama.OllamaService, pooler embedding.EmbedderPooler) interfaces.ModelService {
 	return &modelService{
 		repo:          repo,
 		ollamaService: ollamaService,
+		pooler:        pooler,
 	}
 }
 
@@ -253,7 +255,7 @@ func (s *modelService) GetEmbeddingModel(ctx context.Context, modelId string) (e
 		Dimensions:           model.Parameters.EmbeddingParameters.Dimension,
 		TruncatePromptTokens: model.Parameters.EmbeddingParameters.TruncatePromptTokens,
 		Provider:             model.Parameters.Provider,
-	})
+	}, s.pooler, s.ollamaService)
 	if err != nil {
 		logger.ErrorWithFields(ctx, err, map[string]interface{}{
 			"model_id":   model.ID,
@@ -335,7 +337,7 @@ func (s *modelService) GetChatModel(ctx context.Context, modelId string) (chat.C
 		BaseURL:   model.Parameters.BaseURL,
 		ModelName: model.Name,
 		Source:    model.Source,
-	})
+	}, s.ollamaService)
 	if err != nil {
 		logger.ErrorWithFields(ctx, err, map[string]interface{}{
 			"model_id":   model.ID,

internal/handler/initialization.go

@@ -58,6 +58,7 @@ type InitializationHandler struct {
 	knowledgeService interfaces.KnowledgeService
 	ollamaService    *ollama.OllamaService
 	docReaderClient  *client.Client
+	pooler           embedding.EmbedderPooler
 }
 
 // NewInitializationHandler 创建初始化处理器
@@ -70,6 +71,7 @@ func NewInitializationHandler(
 	knowledgeService interfaces.KnowledgeService,
 	ollamaService *ollama.OllamaService,
 	docReaderClient *client.Client,
+	pooler embedding.EmbedderPooler,
 ) *InitializationHandler {
 	return &InitializationHandler{
 		config:           config,
@@ -80,6 +82,7 @@ func NewInitializationHandler(
 		knowledgeService: knowledgeService,
 		ollamaService:    ollamaService,
 		docReaderClient:  docReaderClient,
+		pooler:           pooler,
 	}
 }
 
@@ -1544,7 +1547,7 @@ func (h *InitializationHandler) TestEmbeddingModel(c *gin.Context) {
 		Provider:             req.Provider,
 	}
 
-	emb, err := embedding.NewEmbedder(cfg)
+	emb, err := embedding.NewEmbedder(cfg, h.pooler, h.ollamaService)
 	if err != nil {
 		logger.ErrorWithFields(ctx, err, map[string]interface{}{"model": utils.SanitizeForLog(req.ModelName)})
 		c.JSON(http.StatusOK, gin.H{
@@ -1588,7 +1591,7 @@ func (h *InitializationHandler) checkRemoteModelConnection(ctx context.Context,
 	}
 
 	// 创建聊天实例
-	chatInstance, err := chat.NewChat(chatConfig)
+	chatInstance, err := chat.NewChat(chatConfig, h.ollamaService)
 	if err != nil {
 		return false, fmt.Sprintf("创建聊天实例失败: %v", err)
 	}
@@ -2087,7 +2090,7 @@ func (h *InitializationHandler) extractRelationsFromText(
 		BaseURL:   llm.BaseUrl,
 		ModelName: llm.ModelName,
 		Source:    types.ModelSource(llm.Source),
-	})
+	}, h.ollamaService)
 	if err != nil {
 		logger.Error(ctx, "初始化模型服务失败", err)
 		return nil, err
@@ -2169,7 +2172,7 @@ func (h *InitializationHandler) fabriText(ctx context.Context, tags []string, ll
 		BaseURL:   llm.BaseUrl,
 		ModelName: llm.ModelName,
 		Source:    types.ModelSource(llm.Source),
-	})
+	}, h.ollamaService)
 	if err != nil {
 		logger.Error(ctx, "初始化模型服务失败", err)
 		return "", err

internal/models/chat/chat.go

@@ -7,7 +7,6 @@ import (
 	"strings"
 
 	"github.com/Tencent/WeKnora/internal/models/utils/ollama"
-	"github.com/Tencent/WeKnora/internal/runtime"
 	"github.com/Tencent/WeKnora/internal/types"
 )
 
@@ -87,14 +86,12 @@ type ChatConfig struct {
 }
 
 // NewChat 创建聊天实例
-func NewChat(config *ChatConfig) (Chat, error) {
+func NewChat(config *ChatConfig, ollamaService *ollama.OllamaService) (Chat, error) {
 	var chat Chat
 	var err error
 	switch strings.ToLower(string(config.Source)) {
 	case string(types.ModelSourceLocal):
-		runtime.GetContainer().Invoke(func(ollamaService *ollama.OllamaService) {
-			chat, err = NewOllamaChat(config, ollamaService)
-		})
+		chat, err = NewOllamaChat(config, ollamaService)
 		if err != nil {
 			return nil, err
 		}

internal/models/embedding/embedder.go

@@ -7,7 +7,6 @@ import (
 
 	"github.com/Tencent/WeKnora/internal/models/provider"
 	"github.com/Tencent/WeKnora/internal/models/utils/ollama"
-	"github.com/Tencent/WeKnora/internal/runtime"
 	"github.com/Tencent/WeKnora/internal/types"
 )
 
@@ -51,15 +50,13 @@ type Config struct {
 }
 
 // NewEmbedder creates an embedder based on the configuration
-func NewEmbedder(config Config) (Embedder, error) {
+func NewEmbedder(config Config, pooler EmbedderPooler, ollamaService *ollama.OllamaService) (Embedder, error) {
 	var embedder Embedder
 	var err error
 	switch strings.ToLower(string(config.Source)) {
 	case string(types.ModelSourceLocal):
-		runtime.GetContainer().Invoke(func(pooler EmbedderPooler, ollamaService *ollama.OllamaService) {
-			embedder, err = NewOllamaEmbedder(config.BaseURL,
-				config.ModelName, config.TruncatePromptTokens, config.Dimensions, config.ModelID, pooler, ollamaService)
-		})
+		embedder, err = NewOllamaEmbedder(config.BaseURL,
+			config.ModelName, config.TruncatePromptTokens, config.Dimensions, config.ModelID, pooler, ollamaService)
 		return embedder, err
 	case string(types.ModelSourceRemote):
 		// Detect or use configured provider for routing
@@ -88,66 +85,56 @@ func NewEmbedder(config Config) (Embedder, error) {
 					baseURL = strings.Replace(baseURL, "/compatible-mode/v1", "", 1)
 					baseURL = strings.Replace(baseURL, "/compatible-mode", "", 1)
 				}
-				runtime.GetContainer().Invoke(func(pooler EmbedderPooler) {
-					embedder, err = NewAliyunEmbedder(config.APIKey,
-						baseURL,
-						config.ModelName,
-						config.TruncatePromptTokens,
-						config.Dimensions,
-						config.ModelID,
-						pooler)
-				})
+				embedder, err = NewAliyunEmbedder(config.APIKey,
+					baseURL,
+					config.ModelName,
+					config.TruncatePromptTokens,
+					config.Dimensions,
+					config.ModelID,
+					pooler)
 			} else {
 				baseURL := config.BaseURL
 				if baseURL == "" || !strings.Contains(baseURL, "/compatible-mode/") {
 					baseURL = "https://dashscope.aliyuncs.com/compatible-mode/v1"
 				}
-				runtime.GetContainer().Invoke(func(pooler EmbedderPooler) {
-					embedder, err = NewOpenAIEmbedder(config.APIKey,
-						baseURL,
-						config.ModelName,
-						config.TruncatePromptTokens,
-						config.Dimensions,
-						config.ModelID,
-						pooler)
-				})
-			}
-			return embedder, err
-		case provider.ProviderVolcengine:
-			// Volcengine Ark uses multimodal embedding API
-			runtime.GetContainer().Invoke(func(pooler EmbedderPooler) {
-				embedder, err = NewVolcengineEmbedder(config.APIKey,
-					config.BaseURL,
+				embedder, err = NewOpenAIEmbedder(config.APIKey,
+					baseURL,
 					config.ModelName,
 					config.TruncatePromptTokens,
 					config.Dimensions,
 					config.ModelID,
 					pooler)
-			})
+			}
+			return embedder, err
+		case provider.ProviderVolcengine:
+			// Volcengine Ark uses multimodal embedding API
+			embedder, err = NewVolcengineEmbedder(config.APIKey,
+				config.BaseURL,
+				config.ModelName,
+				config.TruncatePromptTokens,
+				config.Dimensions,
+				config.ModelID,
+				pooler)
 			return embedder, err
 		case provider.ProviderJina:
 			// Jina AI uses different API format (truncate instead of truncate_prompt_tokens)
-			runtime.GetContainer().Invoke(func(pooler EmbedderPooler) {
-				embedder, err = NewJinaEmbedder(config.APIKey,
-					config.BaseURL,
-					config.ModelName,
-					config.TruncatePromptTokens,
-					config.Dimensions,
-					config.ModelID,
-					pooler)
-			})
+			embedder, err = NewJinaEmbedder(config.APIKey,
+				config.BaseURL,
+				config.ModelName,
+				config.TruncatePromptTokens,
+				config.Dimensions,
+				config.ModelID,
+				pooler)
 			return embedder, err
 		default:
 			// Use OpenAI-compatible embedder for other providers
-			runtime.GetContainer().Invoke(func(pooler EmbedderPooler) {
-				embedder, err = NewOpenAIEmbedder(config.APIKey,
-					config.BaseURL,
-					config.ModelName,
-					config.TruncatePromptTokens,
-					config.Dimensions,
-					config.ModelID,
-					pooler)
-			})
+			embedder, err = NewOpenAIEmbedder(config.APIKey,
+				config.BaseURL,
+				config.ModelName,
+				config.TruncatePromptTokens,
+				config.Dimensions,
+				config.ModelID,
+				pooler)
 			return embedder, err
 		}
 	default: