Files
xinmi/internal/container/container.go
T
yawan f7fc992f57
Anydoc / anydoc parser engine (push) Canceled after 0s
App / Format, vet, test, and build (push) Canceled after 0s
CLI / test (macos-latest / go 1.26) (push) Canceled after 0s
CLI / test (ubuntu-latest / go 1.26) (push) Canceled after 0s
CLI / test (windows-latest / go 1.26) (push) Canceled after 0s
Build and Push Docker Image / build-ui (push) Canceled after 0s
Build and Push Docker Image / build-docreader (push) Canceled after 0s
Build and Push Docker Image / build-sandbox (push) Canceled after 0s
Build and Push Docker Image / build-app (amd64, linux/amd64, ubuntu-latest) (push) Canceled after 0s
Build and Push Docker Image / build-app (arm64, linux/arm64, ubuntu-24.04-arm) (push) Canceled after 0s
DocReader / Python tests and Go client tests (push) Canceled after 0s
Build & Publish (dsh-weknora) / Typecheck, test, pack (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (pinned) (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (latest, advisory) (push) Canceled after 0s
Frontend / Test, type-check, and build (push) Canceled after 0s
Go Lint Cache / warm lint cache (.) (push) Canceled after 0s
Go Lint Cache / warm lint cache (cli) (push) Canceled after 0s
Go Lint Cache / warm lint cache (client) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.10) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.11) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.12) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.13) (push) Canceled after 0s
Build and Push Docker Image / merge (push) Canceled after 0s
Build & Publish (dsh-weknora) / Publish to npm (push) Canceled after 0s
Build & Publish (mcp-server) / Build sdist + wheel (push) Canceled after 0s
Build & Publish (mcp-server) / Publish to PyPI (push) Canceled after 0s
WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes
2026-09-22 16:19:06 +08:00

1865 lines
77 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
// Package container implements dependency injection container setup
// Provides centralized configuration for services, repositories, and handlers
// This package is responsible for wiring up all dependencies and ensuring proper lifecycle management
package container
import (
"context"
"database/sql"
"errors"
"fmt"
"io"
"net/url"
"os"
"path/filepath"
"slices"
"strconv"
"strings"
"time"
sqlite_vec "github.com/asg017/sqlite-vec-go-bindings/cgo"
_ "github.com/duckdb/duckdb-go/v2"
esv7 "github.com/elastic/go-elasticsearch/v7"
"github.com/elastic/go-elasticsearch/v8"
_ "github.com/go-sql-driver/mysql" // 给 Doris (database/sql) 注册 MySQL 协议驱动
"github.com/milvus-io/milvus/client/v2/milvusclient"
"github.com/neo4j/neo4j-go-driver/v6/neo4j"
"github.com/panjf2000/ants/v2"
"github.com/qdrant/go-client/qdrant"
"github.com/redis/go-redis/v9"
"go.uber.org/dig"
"google.golang.org/grpc"
"gorm.io/driver/postgres"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
"github.com/Tencent/WeKnora/internal/accesscontrol"
"github.com/Tencent/WeKnora/internal/agent/approval"
"github.com/Tencent/WeKnora/internal/application/repository"
dorisRepo "github.com/Tencent/WeKnora/internal/application/repository/retriever/doris"
elasticsearchRepoV7 "github.com/Tencent/WeKnora/internal/application/repository/retriever/elasticsearch/v7"
elasticsearchRepoV8 "github.com/Tencent/WeKnora/internal/application/repository/retriever/elasticsearch/v8"
"github.com/Tencent/WeKnora/internal/application/repository/retriever/legacygraph"
milvusRepo "github.com/Tencent/WeKnora/internal/application/repository/retriever/milvus"
openSearchRepo "github.com/Tencent/WeKnora/internal/application/repository/retriever/opensearch"
postgresRepo "github.com/Tencent/WeKnora/internal/application/repository/retriever/postgres"
qdrantRepo "github.com/Tencent/WeKnora/internal/application/repository/retriever/qdrant"
sqliteRetrieverRepo "github.com/Tencent/WeKnora/internal/application/repository/retriever/sqlite"
tencentVectorDBRepo "github.com/Tencent/WeKnora/internal/application/repository/retriever/tencentvectordb"
weaviateRepo "github.com/Tencent/WeKnora/internal/application/repository/retriever/weaviate"
"github.com/Tencent/WeKnora/internal/application/service"
chatpipeline "github.com/Tencent/WeKnora/internal/application/service/chat_pipeline"
"github.com/Tencent/WeKnora/internal/application/service/file"
"github.com/Tencent/WeKnora/internal/application/service/memory"
"github.com/Tencent/WeKnora/internal/application/service/retriever"
"github.com/Tencent/WeKnora/internal/common"
"github.com/Tencent/WeKnora/internal/config"
"github.com/Tencent/WeKnora/internal/database"
"github.com/Tencent/WeKnora/internal/datasource"
"github.com/Tencent/WeKnora/internal/datasource/connector/feishu/core"
"github.com/Tencent/WeKnora/internal/datasource/connector/feishu/drive"
"github.com/Tencent/WeKnora/internal/datasource/connector/feishu/wiki"
gitlabConnector "github.com/Tencent/WeKnora/internal/datasource/connector/gitlab"
imaConnector "github.com/Tencent/WeKnora/internal/datasource/connector/ima"
notionConnector "github.com/Tencent/WeKnora/internal/datasource/connector/notion"
rssConnector "github.com/Tencent/WeKnora/internal/datasource/connector/rss"
yuqueConnector "github.com/Tencent/WeKnora/internal/datasource/connector/yuque"
"github.com/Tencent/WeKnora/internal/event"
"github.com/Tencent/WeKnora/internal/handler"
"github.com/Tencent/WeKnora/internal/handler/session"
workbenchHandler "github.com/Tencent/WeKnora/internal/handler/workbench"
imPkg "github.com/Tencent/WeKnora/internal/im"
"github.com/Tencent/WeKnora/internal/im/dingtalk"
"github.com/Tencent/WeKnora/internal/im/feishu"
"github.com/Tencent/WeKnora/internal/im/mattermost"
"github.com/Tencent/WeKnora/internal/im/qqbot"
"github.com/Tencent/WeKnora/internal/im/slack"
"github.com/Tencent/WeKnora/internal/im/telegram"
"github.com/Tencent/WeKnora/internal/im/wechat"
"github.com/Tencent/WeKnora/internal/im/wecom"
"github.com/Tencent/WeKnora/internal/im/yunzhijia"
"github.com/Tencent/WeKnora/internal/infrastructure/docparser"
infra_web_search "github.com/Tencent/WeKnora/internal/infrastructure/web_search"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/mcp"
"github.com/Tencent/WeKnora/internal/models/asr"
"github.com/Tencent/WeKnora/internal/models/chat"
"github.com/Tencent/WeKnora/internal/models/embedding"
"github.com/Tencent/WeKnora/internal/models/limiter"
"github.com/Tencent/WeKnora/internal/models/rerank"
"github.com/Tencent/WeKnora/internal/models/utils/ollama"
"github.com/Tencent/WeKnora/internal/models/vlm"
"github.com/Tencent/WeKnora/internal/resumable"
"github.com/Tencent/WeKnora/internal/router"
"github.com/Tencent/WeKnora/internal/storageallowlist"
"github.com/Tencent/WeKnora/internal/stream"
"github.com/Tencent/WeKnora/internal/tracing/langfuse"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
secutils "github.com/Tencent/WeKnora/internal/utils"
workbench "github.com/Tencent/WeKnora/internal/workbench"
"github.com/tencent/vectordatabase-sdk-go/tcvectordb"
"github.com/weaviate/weaviate-go-client/v5/weaviate"
"github.com/weaviate/weaviate-go-client/v5/weaviate/auth"
wgrpc "github.com/weaviate/weaviate-go-client/v5/weaviate/grpc"
)
// BuildContainer constructs the dependency injection container
// Registers all components, services, repositories and handlers needed by the application
// Creates a fully configured application container with proper dependency resolution
// Parameters:
// - container: Base dig container to add dependencies to
//
// Returns:
// - Configured container with all application dependencies registered
func BuildContainer(container *dig.Container) *dig.Container {
ctx := context.Background()
logger.Debugf(ctx, "[Container] Starting container initialization...")
// Register resource cleaner for proper cleanup of resources
must(container.Provide(NewResourceCleaner, dig.As(new(interfaces.ResourceCleaner))))
// Core infrastructure configuration
logger.Debugf(ctx, "[Container] Registering core infrastructure...")
must(container.Provide(config.LoadConfig))
must(container.Provide(initLangfuse))
must(container.Provide(initDatabase))
must(container.Provide(initFileService))
must(container.Provide(initRedisClient))
must(container.Provide(initAntsPool))
must(container.Invoke(registerLangfuseCleanup))
// Register goroutine pool cleanup handler
must(container.Invoke(registerPoolCleanup))
// Initialize retrieval engine registry for search capabilities
logger.Debugf(ctx, "[Container] Registering retrieval engine registry...")
must(container.Provide(initRetrieveEngineRegistry))
// External service clients
logger.Debugf(ctx, "[Container] Registering external service clients...")
must(container.Provide(initDocReaderClient))
must(container.Provide(docparser.NewImageResolver))
must(container.Provide(initOllamaService))
must(container.Provide(initNeo4jClient))
must(container.Provide(stream.NewStreamManager))
logger.Debugf(ctx, "[Container] Initializing DuckDB...")
must(container.Provide(NewDuckDB))
logger.Debugf(ctx, "[Container] DuckDB registered")
// Data repositories layer
logger.Debugf(ctx, "[Container] Registering repositories...")
must(container.Provide(repository.NewTenantRepository))
must(container.Provide(repository.NewTenantAPIKeyRepository))
must(container.Provide(repository.NewTenantMemberRepository))
must(container.Provide(repository.NewTenantInvitationRepository))
must(container.Provide(repository.NewAuditLogRepository))
must(container.Provide(repository.NewKnowledgeBaseRepository))
must(container.Provide(repository.NewKnowledgeRepository))
must(container.Provide(repository.NewKnowledgeSpanRepository))
must(container.Provide(repository.NewChunkRepository))
must(container.Provide(repository.NewKnowledgeTagRepository))
must(container.Provide(repository.NewSessionRepository))
must(container.Provide(repository.NewMessageRepository))
must(container.Provide(repository.NewMessageSuggestionRepository))
must(container.Provide(repository.NewModelRepository))
must(container.Provide(repository.NewUserRepository))
must(container.Provide(repository.NewAuthTokenRepository))
must(container.Provide(repository.NewSystemSettingRepository))
must(container.Provide(legacygraph.NewRetrieveRepository))
must(container.Provide(legacygraph.NewRepository))
must(container.Provide(repository.NewMCPServiceRepository))
must(container.Provide(repository.NewMCPToolApprovalRepository))
must(container.Provide(repository.NewMCPOAuthRepository))
must(container.Provide(repository.NewTenantSandboxConfigRepository))
must(container.Provide(repository.NewTenantSkillRepository))
must(container.Provide(repository.NewCustomAgentRepository))
must(container.Provide(repository.NewOrganizationRepository))
must(container.Provide(repository.NewKBShareRepository))
must(container.Provide(repository.NewAgentShareRepository))
must(container.Provide(repository.NewEmbedChannelRepository))
must(container.Provide(repository.NewTenantDisabledSharedAgentRepository))
must(container.Provide(repository.NewUserResourceFavoriteRepository))
must(container.Provide(service.NewWebSearchStateService))
must(container.Provide(repository.NewDataSourceRepository))
must(container.Provide(repository.NewSyncLogRepository))
must(container.Provide(repository.NewWikiPageRepository))
must(container.Provide(repository.NewMemoryRepository))
must(container.Provide(repository.NewDigitalHumanRepository))
must(container.Provide(repository.NewTaskPendingOpsRepository))
must(container.Provide(repository.NewTaskDeadLetterRepository))
// MCP manager for managing MCP client connections
logger.Debugf(ctx, "[Container] Registering MCP manager...")
must(container.Provide(mcp.NewMCPManager))
must(container.Provide(mcp.NewOAuthManager))
// Sandbox manager fallback is disabled; executable backends are resolved
// from named workspace configurations.
logger.Debugf(ctx, "[Container] Registering sandbox manager...")
must(container.Provide(newSandboxManager))
// Per-tenant sandbox backends: the resolver builds a manager per request
// from the tenant's own configuration, falling back to the singleton above
// for tenants that configured nothing.
must(container.Provide(service.NewTenantSandboxConfigLoader))
must(container.Provide(newTenantSandboxResolver))
// Business service layer
logger.Debugf(ctx, "[Container] Registering business services...")
must(container.Provide(service.NewTenantService))
must(container.Provide(service.NewTenantAPIKeyService))
must(container.Provide(service.NewTenantMemberService))
must(container.Provide(service.NewTenantInvitationService))
must(container.Provide(service.NewAuditLogService))
must(container.Provide(service.NewAuditLogRetentionRunner))
must(container.Provide(service.NewKnowledgeBaseService))
must(container.Provide(service.NewOrganizationService))
must(container.Provide(service.NewKBShareService)) // KBShareService must be registered before KnowledgeService and KnowledgeTagService
must(container.Provide(service.NewAgentShareService))
must(container.Provide(service.NewKnowledgeService))
must(container.Provide(service.NewSpanTracker))
must(container.Provide(service.NewChunkService))
must(container.Provide(service.NewKnowledgeTagService))
must(container.Provide(embedding.NewBatchEmbedder))
must(container.Provide(service.NewModelService))
must(container.Provide(service.NewDigitalHumanService))
must(container.Provide(service.NewResumableUploadService))
must(container.Provide(service.NewKnowledgeClassificationService))
must(container.Provide(service.NewDatasetService))
must(container.Provide(service.NewEvaluationService))
must(container.Provide(service.NewUserService))
must(container.Provide(service.NewSystemSettingService))
must(container.Provide(func(
repo repository.TenantSandboxConfigRepository,
agents interfaces.CustomAgentRepository,
skills repository.TenantSkillRepository,
files interfaces.StorageBackendResolver,
) *service.TenantSandboxConfigService {
return service.NewTenantSandboxConfigService(repo, agents, buildGlobalSandboxConfig(), skills, files)
}))
must(container.Provide(func(s *service.TenantSandboxConfigService) service.WorkspaceSandboxPolicy {
return s
}))
must(container.Provide(service.NewWeKnoraCloudService))
// Extract services - register individual extracters with names
must(container.Provide(service.NewChunkExtractService, dig.Name("chunkExtractor")))
must(container.Provide(service.NewDataTableSummaryService, dig.Name("dataTableSummary")))
must(container.Provide(service.NewImageMultimodalService, dig.Name("imageMultimodal")))
must(container.Provide(service.NewKnowledgePostProcessService, dig.Name("knowledgePostProcess")))
must(container.Provide(service.NewKnowledgeAutoTagService, dig.Name("knowledgeAutoTag")))
must(container.Provide(service.NewMessageService))
must(container.Provide(service.NewMessageSuggestionService))
must(container.Provide(service.NewMCPServiceService))
must(container.Provide(service.NewMCPToolApprovalService))
must(container.Provide(service.NewCustomAgentService))
must(container.Provide(service.NewUserResourceFavoriteService))
must(container.Provide(service.NewWikiPageService))
must(container.Provide(service.NewWikiIngestService, dig.Name("wikiIngest")))
must(container.Provide(service.NewWikiLintService))
must(container.Provide(service.NewEmbedChannelService))
// Web search service (needed by AgentService)
logger.Debugf(ctx, "[Container] Registering web search registry and providers...")
must(container.Provide(infra_web_search.NewRegistry))
must(container.Invoke(registerWebSearchProviders))
must(container.Provide(repository.NewWebSearchProviderRepository))
must(container.Provide(repository.NewVectorStoreRepository))
must(container.Provide(repository.NewStorageBackendRepository))
must(container.Provide(repository.NewResourceRepository))
must(container.Provide(repository.NewTemporaryDocumentRepository))
must(container.Provide(service.NewResourceCatalog))
// TenantStoreOwnership adapter used by the retriever factory functions
// to verify that a resolved VectorStore belongs to the caller's tenant.
must(container.Provide(retriever.NewVectorStoreRepoOwnership))
must(container.Provide(service.NewWebSearchService))
must(container.Provide(service.NewWebSearchProviderService))
must(container.Provide(NewEngineFactory))
// StoreRegistry: same instance as RetrieveEngineRegistry, exposed as StoreRegistry interface.
// NewRetrieveEngineRegistry always returns *retriever.RetrieveEngineRegistry which implements both.
must(container.Provide(func(r interfaces.RetrieveEngineRegistry) (interfaces.StoreRegistry, error) {
sr, ok := r.(*retriever.RetrieveEngineRegistry)
if !ok {
return nil, fmt.Errorf("registry does not implement StoreRegistry")
}
return sr, nil
}))
must(container.Provide(service.NewVectorStoreService))
must(container.Provide(service.NewStorageBackendServiceWithResources))
must(container.Provide(func(s *service.StorageBackendService) interfaces.StorageBackendService { return s }))
must(container.Provide(func(s *service.StorageBackendService) interfaces.StorageBackendResolver { return s }))
// Agent service layer (requires event bus, web search service)
// SessionService is passed as parameter to CreateAgentEngine method when creating AgentService
logger.Debugf(ctx, "[Container] Registering event bus and agent service...")
must(container.Provide(event.NewEventBus))
must(container.Provide(service.NewSessionSandboxPinner))
must(container.Provide(func(cfg *config.Config, s interfaces.MCPToolApprovalService, rdb *redis.Client) *approval.Gate {
return approval.NewGate(cfg, &approval.Adapter{Svc: s}, rdb)
}))
// Expose Gate as MCPApproval interface so AgentService and others can depend on the abstraction.
must(container.Provide(func(g *approval.Gate) approval.MCPApproval { return g }))
must(container.Provide(service.NewAgentService))
// Session service (depends on agent service)
// SessionService is created after AgentService and passes itself to AgentService.CreateAgentEngine when needed
logger.Debugf(ctx, "[Container] Registering memory service...")
must(container.Provide(memory.NewMemoryService))
logger.Debugf(ctx, "[Container] Registering session service...")
must(container.Provide(service.NewSessionService))
must(container.Provide(service.NewTenantSkillService))
must(container.Provide(service.NewInstructionSkillService))
// Native 私域Agent workbench resources (training maps, content tasks, flows)
// use the same tenant DB and native model/agent/session services.
must(container.Provide(func(db *gorm.DB) *workbench.Service { return workbench.NewService(db) }))
must(container.Provide(accesscontrol.NewService))
must(container.Provide(func(s *workbench.Service) *workbenchHandler.Handler { return workbenchHandler.NewHandler(s) }))
// The member-facing half of env vars is its own service because its
// authority is different in kind: it derives the identity from the context
// and touches only that identity's rows.
must(container.Provide(service.NewUserEnvService))
// ArtifactCollector drains skill-generated files from the sandbox on
// each agent turn (see spec at
// docs/superpowers/specs/2026-07-10-skill-artifact-download-design.md).
// The factory returns nil when the sandbox backend does not support
// per-session file inspection; downstream code guards on nil.
must(container.Provide(service.NewArtifactCollectorFromSandboxManager))
logger.Debugf(ctx, "[Container] Registering task enqueuer...")
redisAvailable := os.Getenv("REDIS_ADDR") != ""
if redisAvailable {
must(container.Provide(router.NewAsyncqClient, dig.As(new(interfaces.TaskEnqueuer))))
// Dedicated pools guarantee capacity for each stage. The shared pool
// additionally subscribes to core/enrichment queues to provide elastic
// borrowing while post-process and maintenance remain hard-isolated.
must(container.Provide(router.NewCoreAsynqServer, dig.Name("coreAsynqServer")))
must(container.Provide(router.NewPostProcessAsynqServer, dig.Name("postProcessAsynqServer")))
must(container.Provide(router.NewEnrichmentAsynqServer, dig.Name("enrichmentAsynqServer")))
must(container.Provide(router.NewMaintenanceAsynqServer, dig.Name("maintenanceAsynqServer")))
must(container.Provide(router.NewSharedAsynqServer, dig.Name("sharedAsynqServer")))
must(container.Provide(router.NewWikiAsynqServer, dig.Name("wikiAsynqServer")))
// Asynq inspector for cancel-by-knowledge-id (best-effort
// dequeue of pending/scheduled/retry tasks + active-task cancel).
must(container.Provide(router.NewAsynqInspector))
must(container.Provide(router.NewAsynqTaskInspector))
// Install the distributed per-model chat concurrency governor. Only
// available with Redis (the shared semaphore backend); Lite mode is
// single-process and low-volume, so it runs ungated.
must(container.Invoke(registerModelConcurrencyLimiter))
} else {
syncExec := router.NewSyncTaskExecutor()
must(container.Provide(func() interfaces.TaskEnqueuer { return syncExec }))
must(container.Provide(func() *router.SyncTaskExecutor { return syncExec }))
// Lite mode: no Redis means no asynq inspector. SyncTaskExecutor
// dispatches inline goroutines that the checkpoint-based abort
// already handles.
must(container.Provide(router.NewNoopTaskInspector))
// Even without Redis, background ingestion/enrichment can burst the
// worker pool against one provider, so install an in-process governor.
must(container.Invoke(registerLiteModelConcurrencyLimiter))
}
must(container.Provide(service.NewTemporaryDocumentService))
must(container.Invoke(startTemporaryDocumentCleanup))
// Chat pipeline components for processing chat requests
logger.Debugf(ctx, "[Container] Registering chat pipeline plugins...")
// Data source sync framework
logger.Debugf(ctx, "[Container] Registering data source sync framework...")
must(container.Provide(initConnectorRegistry))
must(container.Provide(datasource.NewScheduler))
must(container.Provide(service.NewDataSourceService))
must(container.Invoke(startDataSourceScheduler))
logger.Debugf(ctx, "[Container] Data source sync framework registered")
must(container.Invoke(startAuditLogRetention))
logger.Debugf(ctx, "[Container] Audit log retention runner registered")
must(container.Provide(service.NewHousekeepingService))
must(container.Invoke(startHousekeepingService))
logger.Debugf(ctx, "[Container] Knowledge housekeeping runner registered")
must(container.Provide(chatpipeline.NewEventManager))
must(container.Invoke(chatpipeline.NewPluginSearch))
must(container.Invoke(chatpipeline.NewPluginRerank))
must(container.Invoke(chatpipeline.NewPluginWebFetch))
must(container.Invoke(chatpipeline.NewPluginMerge))
must(container.Invoke(chatpipeline.NewPluginDataAnalysis))
must(container.Invoke(chatpipeline.NewPluginIntoChatMessage))
must(container.Invoke(chatpipeline.NewPluginChatCompletion))
must(container.Invoke(chatpipeline.NewPluginChatCompletionStream))
must(container.Invoke(chatpipeline.NewPluginFilterTopK))
must(container.Invoke(chatpipeline.NewPluginQueryUnderstand))
must(container.Invoke(chatpipeline.NewPluginLoadHistory))
must(container.Invoke(chatpipeline.NewPluginMemoryRecall))
must(container.Invoke(chatpipeline.NewPluginExtractEntity))
must(container.Invoke(chatpipeline.NewPluginSearchEntity))
must(container.Invoke(chatpipeline.NewPluginSearchParallel))
must(container.Invoke(chatpipeline.NewPluginWikiBoost))
must(container.Invoke(chatpipeline.NewPluginMemoryAffinity))
logger.Debugf(ctx, "[Container] Chat pipeline plugins registered")
// TenantSkillService is provided next to SessionService (handlers need
// it), but Invoke constructs the whole chain. SessionService needs
// *chatpipeline.EventManager, which only exists after the pipeline
// block above — starting the reaper any earlier panics.
must(container.Invoke(startTenantSkillReaper))
logger.Debugf(ctx, "[Container] Tenant skill reaper registered")
// HTTP handlers layer
logger.Debugf(ctx, "[Container] Registering HTTP handlers...")
must(container.Provide(handler.NewTenantHandler))
must(container.Provide(handler.NewTenantMemberHandler))
must(container.Provide(handler.NewTenantInvitationHandler))
must(container.Provide(handler.NewAuditLogHandler))
must(container.Provide(handler.NewKnowledgeBaseHandler))
must(container.Provide(handler.NewNativeGraphHandler))
must(container.Provide(handler.NewKnowledgeClassificationHandler))
must(container.Provide(handler.NewKnowledgeHandler))
must(container.Provide(handler.NewChunkHandler))
must(container.Provide(handler.NewFAQHandler))
must(container.Provide(handler.NewTagHandler))
must(container.Provide(session.NewHandler))
must(container.Provide(handler.NewMessageHandler))
must(container.Provide(handler.NewMessageSuggestionHandler))
must(container.Provide(handler.NewModelHandler))
must(container.Provide(handler.NewSandboxConfigHandler))
must(container.Provide(func(
s *service.TenantSkillService, streams interfaces.StreamManager,
) *handler.SandboxSkillHandler {
return handler.NewSandboxSkillHandler(s, streams)
}))
must(container.Provide(handler.NewMeEnvVarHandler))
must(container.Provide(handler.NewEvaluationHandler))
must(container.Provide(handler.NewInitializationHandler))
must(container.Provide(handler.NewAuthHandler))
must(container.Provide(handler.NewSystemHandler))
must(container.Provide(handler.NewMCPServiceHandler))
must(container.Provide(handler.NewMCPCredentialsHandler))
must(container.Provide(handler.NewMCPOAuthHandler))
must(container.Provide(handler.NewModelCredentialsHandler))
must(container.Provide(handler.NewWebSearchProviderCredentialsHandler))
must(container.Provide(handler.NewDataSourceCredentialsHandler))
must(container.Provide(handler.NewWebSearchHandler))
must(container.Provide(handler.NewWebSearchProviderHandler))
must(container.Provide(handler.NewVectorStoreHandler))
must(container.Provide(handler.NewStorageBackendHandler))
must(container.Provide(handler.NewCustomAgentHandler))
must(container.Provide(handler.NewUserResourceFavoriteHandler))
must(container.Provide(service.NewSkillService))
must(container.Provide(func(s *service.TenantSkillService, instructions *service.InstructionSkillService) *handler.SkillHandler {
h := handler.NewSkillHandler(s, s)
h.SetInstructionService(instructions)
return h
}))
must(container.Provide(handler.NewOrganizationHandler))
must(container.Provide(handler.NewMemoryHandler))
must(container.Provide(handler.NewDigitalHumanHandler))
// Data source handler
must(container.Provide(handler.NewDataSourceHandler))
// Wiki page handler
must(container.Provide(handler.NewWikiPageHandler))
// IM integration
logger.Debugf(ctx, "[Container] Registering IM integration...")
must(container.Provide(imPkg.NewService))
must(container.Invoke(registerIMService))
must(container.Provide(handler.NewIMHandler))
must(container.Provide(handler.NewEmbedChannelHandler))
must(container.Provide(handler.NewWeKnoraCloudHandler))
logger.Debugf(ctx, "[Container] HTTP handlers registered")
must(container.Invoke(func(s *workbench.Service, models interfaces.ModelService, agents interfaces.CustomAgentService, sessions interfaces.SessionService, instructions *service.InstructionSkillService, bases interfaces.KnowledgeBaseService, knowledge interfaces.KnowledgeService, shares interfaces.KBShareService, access *accesscontrol.Service, memories interfaces.MemoryService) {
s.SetDependencies(models, agents, sessions)
s.SetInstructionResolver(instructions)
s.SetKnowledgeServices(bases, knowledge, shares)
s.SetAccessSnapshot(access.Snapshot)
s.SetMemoryService(memories)
if writer, ok := memories.(interface {
RememberUsingRepositories(context.Context, types.MemoryItem, interfaces.MemoryRepository, interfaces.TenantRepository) (*types.MemoryItem, error)
}); ok {
s.SetTransactionalMemoryWriter(func(ctx context.Context, tx *gorm.DB, item types.MemoryItem) (*types.MemoryItem, error) {
return writer.RememberUsingRepositories(ctx, item, repository.NewMemoryRepository(tx), repository.NewTenantRepository(tx))
})
}
if metered, ok := models.(interface {
SetChatWrapper(func(chat.Chat) chat.Chat)
}); ok {
metered.SetChatWrapper(s.WrapChat)
}
if observed, ok := models.(interface {
SetAuxiliaryModelWrappers(func(embedding.Embedder) embedding.Embedder, func(vlm.VLM) vlm.VLM, func(asr.ASR) asr.ASR, func(rerank.Reranker) rerank.Reranker)
}); ok {
observed.SetAuxiliaryModelWrappers(s.WrapEmbedding, s.WrapVLM, s.WrapASR, s.WrapRerank)
}
s.Start(context.Background())
}))
must(container.Invoke(func(access *accesscontrol.Service, knowledge interfaces.KnowledgeService) error {
if err := access.SetKnowledgeVersionServices(knowledge); err != nil {
return err
}
access.StartKnowledgeVersionWorker(context.Background())
return nil
}))
// Wire the chat package's local image resolver so multimodal chat can read
// local:// images that live under a tenant's configured storage PathPrefix
// (which is not encoded in the local:// URL).
must(container.Invoke(registerChatLocalImageResolver))
// Router configuration
logger.Debugf(ctx, "[Container] Registering router and starting task server...")
must(container.Provide(router.NewRouter))
if redisAvailable {
must(container.Invoke(router.RunAsynqServer))
} else {
must(container.Invoke(router.RegisterSyncHandlers))
}
// Wiki operation rows are durable, while their wake-up triggers may be
// lost across a process restart (always in Lite mode, and in Redis mode if
// persistence succeeded immediately before trigger enqueue failed). Re-arm
// them only after the matching handlers are ready.
must(container.Invoke(recoverPendingWikiTasks))
must(container.Invoke(func(memories interfaces.MemoryService, cleaner interfaces.ResourceCleaner) {
if recovery, ok := memories.(interface{ StartExtractionRecovery(context.Context) func() }); ok {
stop := recovery.StartExtractionRecovery(context.Background())
cleaner.RegisterWithName("MemoryExtractionRecovery", func() error { stop(); return nil })
}
}))
must(container.Invoke(func(uploads *resumable.Service) {
uploads.Start(context.Background())
}))
must(container.Invoke(func(classification *service.KnowledgeClassificationService) {
classification.Start(context.Background())
}))
logger.Infof(ctx, "[Container] Container initialization completed successfully")
return container
}
// registerChatLocalImageResolver wires the chat package's LocalImageResolver
// hook. Stored local:// URLs are relative to the resolved storage base dir and
// do NOT encode the owning tenant's configured PathPrefix, so resolving them to
// disk bytes requires rebuilding the FileService from that tenant's storage
// config. The owning tenant is parsed from the URL's first path segment, which
// correctly handles cross-tenant shared resources (e.g. shared KB images).
func registerChatLocalImageResolver(
tenantRepo interfaces.TenantRepository,
storageResolver interfaces.StorageBackendResolver,
resourceCatalog interfaces.ResourceCatalog,
) {
chat.LocalImageResolver = func(storageURL string) ([]byte, bool) {
ctx := context.Background()
physicalPath, resource, err := resourceCatalog.ResolvePath(ctx, storageURL)
if err != nil {
return nil, false
}
tenantID := secutils.ParseTenantIDFromStoragePath(physicalPath)
if resource != nil {
tenantID = resource.TenantID
}
if tenantID == 0 {
return nil, false
}
tenant, err := tenantRepo.GetTenantByID(ctx, tenantID)
if err != nil || tenant == nil {
return nil, false
}
baseDir := strings.TrimSpace(os.Getenv("LOCAL_STORAGE_BASE_DIR"))
backendID, inner, scoped := types.ParseStorageBackendPath(physicalPath)
if resource != nil && resource.StorageBackendID != "" {
backendID = resource.StorageBackendID
}
providerPath := physicalPath
if scoped {
providerPath = inner
}
provider := types.ParseProviderScheme(providerPath)
if provider == "" {
provider = "local"
}
fileSvc, _, err := storageResolver.ResolveFileService(ctx, tenant, backendID, provider, baseDir)
if err != nil {
return nil, false
}
rc, err := fileSvc.GetFile(ctx, physicalPath)
if err != nil {
return nil, false
}
defer rc.Close()
data, err := io.ReadAll(rc)
if err != nil {
return nil, false
}
return data, true
}
}
// must is a helper function for error handling
// Panics if the error is not nil, useful for configuration steps that must succeed
// Parameters:
// - err: Error to check
func must(err error) {
if err != nil {
panic(err)
}
}
// initLangfuse initializes the Langfuse ingestion client.
// Configuration is read from LANGFUSE_* environment variables (see
// docs/langfuse.md). Returns a disabled manager if credentials are absent —
// never an error — so deployments that don't use Langfuse are unaffected.
func initLangfuse() (*langfuse.Manager, error) {
cfg := langfuse.LoadConfigFromEnv()
return langfuse.Init(cfg)
}
// defaultModelMaxConcurrency is the per-model cap on concurrent background
// (ingestion/enrichment) chat calls when WEKNORA_MODEL_MAX_CONCURRENCY /
// model.max_concurrency is unset. summary / question / graph enrichment all
// share the same model, so this bounds their combined pressure on one provider
// across every replica. Interactive chat is never gated.
const defaultModelMaxConcurrency = 32
// resolveModelMaxConcurrency reads the per-model background concurrency limit
// from system settings / env, defaulting to defaultModelMaxConcurrency when
// unset. A configured value of 0 (or negative) is honoured and disables the
// governor — that is the supported way to turn throttling off via config/env.
func resolveModelMaxConcurrency(ss interfaces.SystemSettingService) int {
if ss == nil {
return defaultModelMaxConcurrency
}
return int(ss.GetInt(context.Background(), "model.max_concurrency",
"WEKNORA_MODEL_MAX_CONCURRENCY", int64(defaultModelMaxConcurrency)))
}
// registerModelConcurrencyLimiter builds the Redis-backed per-model background
// concurrency governor (chat + vlm) and installs it. Only available with Redis
// (the shared semaphore backend); Lite mode uses registerLiteModelConcurrencyLimiter.
func registerModelConcurrencyLimiter(rdb *redis.Client, ss interfaces.SystemSettingService) {
limit := resolveModelMaxConcurrency(ss)
limiter.SetGovernor(limiter.NewRedisLimiter(rdb), limit)
if limit <= 0 {
logger.Infof(context.Background(),
"[ModelLimiter] background concurrency governor DISABLED (model.max_concurrency<=0)")
return
}
logger.Infof(context.Background(),
"[ModelLimiter] background model concurrency governed per-model, limit=%d (distributed via redis)", limit)
}
// registerLiteModelConcurrencyLimiter installs an in-process per-model governor
// for Lite mode (no Redis). Lite runs a single process, so an in-process
// semaphore is sufficient to keep a background ingestion storm from bursting
// the whole worker pool against one provider.
func registerLiteModelConcurrencyLimiter(ss interfaces.SystemSettingService) {
limit := resolveModelMaxConcurrency(ss)
limiter.SetGovernor(limiter.NewLocalLimiter(), limit)
if limit <= 0 {
logger.Infof(context.Background(),
"[ModelLimiter] background concurrency governor DISABLED (model.max_concurrency<=0)")
return
}
logger.Infof(context.Background(),
"[ModelLimiter] background model concurrency governed per-model, limit=%d (in-process, lite mode)", limit)
}
func initRedisClient() (*redis.Client, error) {
redisAddr := os.Getenv("REDIS_ADDR")
if redisAddr == "" {
logger.Infof(context.Background(), "[Redis] No REDIS_ADDR configured, Redis disabled (Lite mode)")
return nil, nil
}
db, err := strconv.Atoi(os.Getenv("REDIS_DB"))
if err != nil {
db = 0
}
client := redis.NewClient(&redis.Options{
Addr: redisAddr,
Username: os.Getenv("REDIS_USERNAME"),
Password: os.Getenv("REDIS_PASSWORD"),
DB: db,
TLSConfig: common.RedisTLSConfig(),
})
_, err = client.Ping(context.Background()).Result()
if err != nil {
return nil, fmt.Errorf("连接Redis失败: %w", err)
}
return client, nil
}
// initDatabase initializes database connection
// Creates and configures database connection based on environment configuration
// Supports multiple database backends (PostgreSQL)
// Parameters:
// - cfg: Application configuration
//
// Returns:
// - Configured database connection
// - Error if connection fails
func initDatabase(cfg *config.Config) (*gorm.DB, error) {
var dialector gorm.Dialector
var migrateDSN string
var sqliteDBPath string
switch os.Getenv("DB_DRIVER") {
case "postgres":
// DSN for GORM (key-value format)
gormDSN := fmt.Sprintf(
"host=%s port=%s user=%s password=%s dbname=%s sslmode=%s TimeZone=UTC",
os.Getenv("DB_HOST"),
os.Getenv("DB_PORT"),
os.Getenv("DB_USER"),
os.Getenv("DB_PASSWORD"),
os.Getenv("DB_NAME"),
"disable",
)
dialector = postgres.Open(gormDSN)
// DSN for golang-migrate (URL format)
// URL-encode password to handle special characters like !@#
dbPassword := os.Getenv("DB_PASSWORD")
encodedPassword := url.QueryEscape(dbPassword)
// Check if postgres is in RETRIEVE_DRIVER to determine skip_embedding
retrieveDriver := strings.Split(os.Getenv("RETRIEVE_DRIVER"), ",")
skipEmbedding := "true"
if slices.Contains(retrieveDriver, "postgres") {
skipEmbedding = "false"
}
logger.Infof(context.Background(), "Skip embedding: %s", skipEmbedding)
migrateDSN = fmt.Sprintf(
"postgres://%s:%s@%s:%s/%s?sslmode=disable&options=-c%%20app.skip_embedding=%s",
os.Getenv("DB_USER"),
encodedPassword, // Use encoded password
os.Getenv("DB_HOST"),
os.Getenv("DB_PORT"),
os.Getenv("DB_NAME"),
skipEmbedding,
)
// Debug log (don't log password)
logger.Infof(context.Background(), "DB Config: user=%s host=%s port=%s dbname=%s",
os.Getenv("DB_USER"),
os.Getenv("DB_HOST"),
os.Getenv("DB_PORT"),
os.Getenv("DB_NAME"),
)
case "sqlite":
dbPath := os.Getenv("DB_PATH")
if dbPath == "" {
dbPath = "./data/weknora.db"
}
if dir := filepath.Dir(dbPath); dir != "." && dir != "" {
if err := os.MkdirAll(dir, 0o755); err != nil {
return nil, fmt.Errorf("failed to create SQLite data directory %s: %w", dir, err)
}
}
sqlite_vec.Auto()
dsn := dbPath + "?_journal_mode=WAL&_busy_timeout=5000&_foreign_keys=on"
dialector = sqlite.Open(dsn)
sqliteDBPath = dbPath
migrateDSN = "sqlite3://" + dbPath
logger.Infof(context.Background(), "DB Config: driver=sqlite path=%s", dbPath)
default:
return nil, fmt.Errorf("unsupported database driver: %s", os.Getenv("DB_DRIVER"))
}
db, err := gorm.Open(dialector, &gorm.Config{
NowFunc: func() time.Time {
return time.Now().UTC()
},
})
if err != nil {
return nil, err
}
// Sanity check: dialect-specific code in services (notably the
// vector_stores delete guard) compares Dialector.Name() to "postgres" /
// "sqlite" string literals. A future driver swap that produces a
// different name (e.g., a wrapper dialect for managed PG) would silently
// fall back to the SQLite path, dropping the row-level X-lock. Catching
// the mismatch at startup is loud and inexpensive.
if name := db.Dialector.Name(); name != "postgres" && name != "sqlite" {
return nil, fmt.Errorf(
"unsupported gorm dialector %q; expected postgres or sqlite "+
"(see vectorStoreService.isPostgres for impact)", name)
}
if os.Getenv("DB_DRIVER") == "sqlite" {
sqlDB, err := db.DB()
if err != nil {
return nil, fmt.Errorf("failed to get underlying sql.DB: %w", err)
}
if err := sqlDB.Ping(); err != nil {
return nil, fmt.Errorf("failed to ping SQLite database: %w", err)
}
}
// Run database migrations automatically (optional, can be disabled via env var)
// To disable auto-migration, set AUTO_MIGRATE=false
// To enable auto-recovery from dirty state, set AUTO_RECOVER_DIRTY=true
if os.Getenv("AUTO_MIGRATE") != "false" {
logger.Infof(context.Background(), "Running database migrations...")
autoRecover := os.Getenv("AUTO_RECOVER_DIRTY") != "false"
migrationOpts := database.MigrationOptions{
AutoRecoverDirty: autoRecover,
SQLiteDBPath: sqliteDBPath,
}
// Run base migrations (all versioned migrations including embeddings)
// The embeddings migration will be conditionally executed based on skip_embedding parameter in DSN
if err := database.RunMigrationsWithOptions(migrateDSN, migrationOpts); err != nil {
// A partially migrated schema is not safe to serve. Workbench handlers
// rely on the latest tables/columns (for example
// workbench_jobs.generation); continuing here turns a deploy-time
// migration problem into opaque HTTP 500s during user requests.
return nil, fmt.Errorf("database migration failed; refusing to start with an incomplete schema: %w", err)
}
// Post-migration: resolve __pending_env__ storage provider markers for historical KBs.
// The SQL migration marks KBs that have documents but no provider with "__pending_env__";
// we replace that with the actual STORAGE_TYPE from the environment.
resolveStorageProviderPending(db)
migrateLegacyStorageBackends(db)
// Post-migration: declarative built-in models from config/builtin_models.yaml (optional).
if err := types.LoadBuiltinModelsConfig(context.Background(), db, config.ConfigDir()); err != nil {
logger.Warnf(context.Background(), "Load builtin models config failed: %v", err)
}
} else {
logger.Infof(context.Background(), "Auto-migration is disabled (AUTO_MIGRATE=false)")
}
// Get underlying SQL DB object
sqlDB, err := db.DB()
if err != nil {
return nil, err
}
// Configure connection pool parameters
if os.Getenv("DB_DRIVER") == "sqlite" {
// SQLite only supports one concurrent writer even in WAL mode.
// Limiting to a single open connection serialises all DB access and
// prevents "database is locked" errors from concurrent goroutines.
sqlDB.SetMaxOpenConns(1)
} else {
sqlDB.SetMaxIdleConns(10)
}
sqlDB.SetConnMaxLifetime(time.Duration(10) * time.Minute)
return db, nil
}
// resolveStorageProviderPending replaces the "__pending_env__" sentinel in
// knowledge_bases.storage_provider_config with the actual STORAGE_TYPE from the environment.
// This runs once after SQL migrations to bind historical KBs to their real storage provider.
func resolveStorageProviderPending(db *gorm.DB) {
storageType := strings.TrimSpace(os.Getenv("STORAGE_TYPE"))
if storageType == "" {
storageType = "local"
}
storageType = strings.ToLower(storageType)
result := db.Exec(
`UPDATE knowledge_bases SET storage_provider_config = ? WHERE storage_provider_config IS NOT NULL AND storage_provider_config->>'provider' = '__pending_env__'`,
fmt.Sprintf(`{"provider":"%s"}`, storageType),
)
if result.Error != nil {
logger.Warnf(context.Background(), "Failed to resolve __pending_env__ storage providers: %v", result.Error)
} else if result.RowsAffected > 0 {
logger.Infof(context.Background(), "Resolved %d knowledge bases with __pending_env__ storage provider → %s", result.RowsAffected, storageType)
}
// Sync PostgreSQL sequences with actual MAX values to prevent duplicate key
// errors. The old code assigned seq_id via SELECT MAX()+1 in application
// code, which could push values past the DB sequence counter.
syncSequences(db)
// Reset any pending tasks left over from previous aborted runs (Lite App mode)
resetPendingTasks(db)
}
// migrateLegacyStorageBackends backfills the storage_backends table from each
// workspace's legacy StorageEngineConfig (or environment defaults) and binds
// existing knowledge bases to the resulting backend.
//
// The table, columns and indexes are created by the SQL migrations
// (migrations/versioned/000068 for Postgres, migrations/sqlite/000000_init for
// SQLite); this step only handles data that cannot be expressed portably in
// SQL: environment snapshots, JSON→config mapping, AES-encrypted credentials,
// UUID generation and the per-startup refresh of env-backed aliases.
// The migration is idempotent: one legacy_alias row per tenant/provider.
func migrateLegacyStorageBackends(db *gorm.DB) {
var tenants []*types.Tenant
if err := db.Find(&tenants).Error; err != nil {
logger.Warnf(context.Background(), "Failed to load workspaces for storage backend migration: %v", err)
return
}
if len(tenants) == 0 {
return
}
// Load every alias in a single query. Probing each tenant/provider pair with
// First() makes GORM log "record not found" for every miss, which floods the
// startup log with workspaces × providers lines on fresh installs.
var aliases []*types.StorageBackend
if err := db.Where("legacy_alias = ?", true).Find(&aliases).Error; err != nil {
logger.Warnf(context.Background(), "Failed to load legacy storage aliases: %v", err)
return
}
existingAliases := make(map[uint64]map[string]*types.StorageBackend, len(aliases))
for _, alias := range aliases {
byProvider := existingAliases[alias.TenantID]
if byProvider == nil {
byProvider = make(map[string]*types.StorageBackend)
existingAliases[alias.TenantID] = byProvider
}
byProvider[alias.Provider] = alias
}
for _, tenant := range tenants {
legacy := tenant.StorageEngineConfig
defaultProvider := ""
if legacy != nil {
defaultProvider = strings.ToLower(strings.TrimSpace(legacy.DefaultProvider))
}
if defaultProvider == "" {
defaultProvider = strings.ToLower(strings.TrimSpace(os.Getenv("STORAGE_TYPE")))
}
if defaultProvider == "" {
defaultProvider = "local"
}
backendIDs := make(map[string]string)
for _, provider := range storageallowlist.Supported() {
if existing := existingAliases[tenant.ID][provider]; existing != nil {
// Environment-backed aliases are snapshots, not user-owned config.
// Refresh them at every startup so credential rotation does not
// leave the persisted resolver on stale values. If the workspace
// later gains an explicit legacy config, promote the alias to user
// source and stop automatic refreshes.
if existing.Source == types.StorageBackendSourceEnv {
desired := types.StorageBackendFromLegacy(tenant.ID, provider, legacy)
if desired == nil && provider == defaultProvider {
desired = types.StorageBackendFromEnvironment(tenant.ID)
}
if desired != nil && desired.Provider == provider {
_ = db.Model(&types.StorageBackend{}).Where("id = ?", existing.ID).Updates(map[string]interface{}{
"name": desired.Name, "config": desired.Config, "source": desired.Source, "status": desired.Status, "updated_at": time.Now(),
}).Error
}
}
backendIDs[provider] = existing.ID
continue
}
backend := types.StorageBackendFromLegacy(tenant.ID, provider, legacy)
if backend == nil && provider == defaultProvider {
backend = types.StorageBackendFromEnvironment(tenant.ID)
}
if backend == nil {
continue
}
if err := db.Create(backend).Error; err != nil {
logger.Warnf(context.Background(), "Failed to migrate %s storage for workspace %d: %v", provider, tenant.ID, err)
continue
}
backendIDs[provider] = backend.ID
}
if tenant.DefaultStorageBackendID == nil {
if id := backendIDs[defaultProvider]; id != "" {
if err := db.Model(&types.Tenant{}).Where("id = ?", tenant.ID).Update("default_storage_backend_id", id).Error; err != nil {
logger.Warnf(context.Background(), "Failed to set default storage backend for workspace %d: %v", tenant.ID, err)
}
}
}
var kbs []*types.KnowledgeBase
if err := db.Where("tenant_id = ? AND storage_backend_id IS NULL", tenant.ID).Find(&kbs).Error; err != nil {
continue
}
for _, kb := range kbs {
provider := kb.GetStorageProvider()
if provider == "" {
provider = defaultProvider
}
if id := backendIDs[provider]; id != "" {
_ = db.Model(&types.KnowledgeBase{}).Where("id = ? AND storage_backend_id IS NULL", kb.ID).Update("storage_backend_id", id).Error
}
}
}
}
// syncSequences ensures PostgreSQL sequences for auto-increment columns (seq_id)
// are at least as high as the current MAX value in each table. This is needed
// because older code assigned seq_id via application-level MAX()+1, which could
// advance values past the DB sequence counter and cause duplicate key errors.
func syncSequences(db *gorm.DB) {
if db.Dialector.Name() != "postgres" {
return
}
pairs := [][2]string{
{"chunks", "chunks_seq_id_seq"},
{"knowledge_tags", "knowledge_tags_seq_id_seq"},
}
for _, p := range pairs {
table, seq := p[0], p[1]
sql := fmt.Sprintf(
`SELECT setval('%s', GREATEST(nextval('%s'), (SELECT COALESCE(MAX(seq_id), 0) FROM %s)))`,
seq, seq, table,
)
if err := db.Exec(sql).Error; err != nil {
logger.Warnf(context.Background(), "Failed to sync sequence %s: %v", seq, err)
} else {
logger.Infof(context.Background(), "Synced sequence %s with table %s", seq, table)
}
}
}
// initFileService initializes file storage service
// Creates the appropriate file storage service based on configuration
// Supports multiple storage backends (MinIO, COS, local filesystem)
// Parameters:
// - cfg: Application configuration
//
// Returns:
// - Configured file service implementation
// - Error if initialization fails
func initFileService(cfg *config.Config, catalog interfaces.ResourceCatalog) (interfaces.FileService, error) {
inner, err := initRawFileService(cfg)
if err != nil {
return nil, err
}
return file.NewResourceCatalogFileService(inner, catalog), nil
}
func initRawFileService(_ *config.Config) (interfaces.FileService, error) {
storageType := strings.TrimSpace(os.Getenv("STORAGE_TYPE"))
if storageType == "" {
storageType = "local"
}
switch storageType {
case "minio":
if os.Getenv("MINIO_ENDPOINT") == "" ||
os.Getenv("MINIO_ACCESS_KEY_ID") == "" ||
os.Getenv("MINIO_SECRET_ACCESS_KEY") == "" ||
os.Getenv("MINIO_BUCKET_NAME") == "" {
return nil, fmt.Errorf("missing MinIO configuration")
}
return file.NewMinioFileService(
os.Getenv("MINIO_ENDPOINT"),
os.Getenv("MINIO_ACCESS_KEY_ID"),
os.Getenv("MINIO_SECRET_ACCESS_KEY"),
os.Getenv("MINIO_BUCKET_NAME"),
strings.EqualFold(os.Getenv("MINIO_USE_SSL"), "true"),
)
case "cos":
if os.Getenv("COS_BUCKET_NAME") == "" ||
os.Getenv("COS_REGION") == "" ||
os.Getenv("COS_SECRET_ID") == "" ||
os.Getenv("COS_SECRET_KEY") == "" ||
os.Getenv("COS_PATH_PREFIX") == "" {
return nil, fmt.Errorf("missing COS configuration")
}
return file.NewCosFileServiceWithTempBucket(
os.Getenv("COS_BUCKET_NAME"),
os.Getenv("COS_REGION"),
os.Getenv("COS_SECRET_ID"),
os.Getenv("COS_SECRET_KEY"),
os.Getenv("COS_PATH_PREFIX"),
os.Getenv("COS_TEMP_BUCKET_NAME"),
os.Getenv("COS_TEMP_REGION"),
)
case "tos":
if os.Getenv("TOS_ENDPOINT") == "" ||
os.Getenv("TOS_REGION") == "" ||
os.Getenv("TOS_ACCESS_KEY") == "" ||
os.Getenv("TOS_SECRET_KEY") == "" ||
os.Getenv("TOS_BUCKET_NAME") == "" {
return nil, fmt.Errorf("missing TOS configuration")
}
return file.NewTosFileServiceWithTempBucket(
os.Getenv("TOS_ENDPOINT"),
os.Getenv("TOS_REGION"),
os.Getenv("TOS_ACCESS_KEY"),
os.Getenv("TOS_SECRET_KEY"),
os.Getenv("TOS_BUCKET_NAME"),
os.Getenv("TOS_PATH_PREFIX"),
os.Getenv("TOS_TEMP_BUCKET_NAME"), // 可选:临时桶名称(桶需配置生命周期规则自动过期)
os.Getenv("TOS_TEMP_REGION"), // 可选:临时桶 region,默认与主桶相同
)
case "s3":
accessKey, secretKey := os.Getenv("S3_ACCESS_KEY"), os.Getenv("S3_SECRET_KEY")
if os.Getenv("S3_REGION") == "" ||
os.Getenv("S3_BUCKET_NAME") == "" ||
(accessKey == "") != (secretKey == "") {
return nil, fmt.Errorf("missing S3 configuration")
}
pathPrefix := os.Getenv("S3_PATH_PREFIX")
if pathPrefix == "" {
pathPrefix = "weknora/"
}
return file.NewS3FileService(
os.Getenv("S3_ENDPOINT"),
accessKey,
secretKey,
os.Getenv("S3_BUCKET_NAME"),
os.Getenv("S3_REGION"),
pathPrefix,
)
case "obs":
if os.Getenv("OBS_ENDPOINT") == "" ||
os.Getenv("OBS_ACCESS_KEY") == "" ||
os.Getenv("OBS_SECRET_KEY") == "" ||
os.Getenv("OBS_BUCKET_NAME") == "" {
return nil, fmt.Errorf("missing OBS configuration")
}
obsRegion := os.Getenv("OBS_REGION")
obsPathPrefix := os.Getenv("OBS_PATH_PREFIX")
if obsPathPrefix == "" {
obsPathPrefix = "weknora/"
}
return file.NewObsFileService(
os.Getenv("OBS_ENDPOINT"),
obsRegion,
os.Getenv("OBS_ACCESS_KEY"),
os.Getenv("OBS_SECRET_KEY"),
os.Getenv("OBS_BUCKET_NAME"),
obsPathPrefix,
)
case "oss":
if os.Getenv("OSS_ENDPOINT") == "" ||
os.Getenv("OSS_REGION") == "" ||
os.Getenv("OSS_ACCESS_KEY") == "" ||
os.Getenv("OSS_SECRET_KEY") == "" ||
os.Getenv("OSS_BUCKET_NAME") == "" {
return nil, fmt.Errorf("missing OSS configuration")
}
pathPrefix := os.Getenv("OSS_PATH_PREFIX")
if pathPrefix == "" {
pathPrefix = "weknora/"
}
return file.NewOssFileServiceWithTempBucket(
os.Getenv("OSS_ENDPOINT"),
os.Getenv("OSS_REGION"),
os.Getenv("OSS_ACCESS_KEY"),
os.Getenv("OSS_SECRET_KEY"),
os.Getenv("OSS_BUCKET_NAME"),
pathPrefix,
os.Getenv("OSS_TEMP_BUCKET_NAME"),
os.Getenv("OSS_TEMP_REGION"),
)
case "local":
baseDir := os.Getenv("LOCAL_STORAGE_BASE_DIR")
if baseDir == "" {
baseDir = "/data/files"
}
externalURL := strings.TrimSpace(os.Getenv("APP_EXTERNAL_URL"))
return file.NewLocalFileService(baseDir, externalURL), nil
case "dummy":
return file.NewDummyFileService(), nil
default:
return nil, fmt.Errorf("unsupported storage type: %s", storageType)
}
}
// initRetrieveEngineRegistry initializes the retrieval engine registry
// Sets up and configures various search engine backends based on configuration
// Supports multiple retrieval engines (PostgreSQL, ElasticsearchV7, ElasticsearchV8)
// Parameters:
// - db: Database connection
// - cfg: Application configuration
//
// Returns:
// - Configured retrieval engine registry
// - Error if initialization fails
func initRetrieveEngineRegistry(
db *gorm.DB, cfg *config.Config, auditSvc interfaces.AuditLogService,
storeRepo interfaces.VectorStoreRepository, engineFactory interfaces.EngineFactory,
) (interfaces.RetrieveEngineRegistry, error) {
// storeRepo and engineFactory let the registry rebuild a store engine that
// is absent from this process, which happens when startup skipped it after
// a construction failure or when another instance registered it.
registry := retriever.NewRetrieveEngineRegistry(storeRepo, engineFactory)
retrieveDriver := strings.Split(os.Getenv("RETRIEVE_DRIVER"), ",")
log := logger.GetLogger(context.Background())
// Audit sink for OpenSearch driver events (index created / reindex). Driver
// events fire under a tenant-scoped ctx at indexing time; the env-path
// registration ctx below has no tenant, so those emits self-skip.
auditSink := newAuditSinkAdapter(auditSvc)
if slices.Contains(retrieveDriver, "postgres") {
postgresRepo := postgresRepo.NewPostgresRetrieveEngineRepository(db)
if err := registry.Register(
retriever.NewKVHybridRetrieveEngine(postgresRepo, types.PostgresRetrieverEngineType),
); err != nil {
log.Errorf("Register postgres retrieve engine failed: %v", err)
} else {
log.Infof("Register postgres retrieve engine success")
}
}
if slices.Contains(retrieveDriver, "sqlite") {
sqliteRepo := sqliteRetrieverRepo.NewSQLiteRetrieveEngineRepository(db)
if err := registry.Register(
retriever.NewKVHybridRetrieveEngine(sqliteRepo, types.SQLiteRetrieverEngineType),
); err != nil {
log.Errorf("Register sqlite retrieve engine failed: %v", err)
} else {
log.Infof("Register sqlite retrieve engine success")
}
}
if slices.Contains(retrieveDriver, "elasticsearch_v8") {
client, err := elasticsearch.NewTypedClient(elasticsearch.Config{
Addresses: []string{os.Getenv("ELASTICSEARCH_ADDR")},
Username: os.Getenv("ELASTICSEARCH_USERNAME"),
Password: os.Getenv("ELASTICSEARCH_PASSWORD"),
})
if err != nil {
log.Errorf("Create elasticsearch_v8 client failed: %v", err)
} else {
elasticsearchRepo := elasticsearchRepoV8.NewElasticsearchEngineRepository(client, cfg, nil)
if err := registry.Register(
retriever.NewKVHybridRetrieveEngine(
elasticsearchRepo, types.ElasticsearchRetrieverEngineType,
),
); err != nil {
log.Errorf("Register elasticsearch_v8 retrieve engine failed: %v", err)
} else {
log.Infof("Register elasticsearch_v8 retrieve engine success")
}
}
}
if slices.Contains(retrieveDriver, "elasticsearch_v7") {
client, err := esv7.NewClient(esv7.Config{
Addresses: []string{os.Getenv("ELASTICSEARCH_ADDR")},
Username: os.Getenv("ELASTICSEARCH_USERNAME"),
Password: os.Getenv("ELASTICSEARCH_PASSWORD"),
})
if err != nil {
log.Errorf("Create elasticsearch_v7 client failed: %v", err)
} else {
elasticsearchRepo := elasticsearchRepoV7.NewElasticsearchEngineRepository(client, cfg, nil)
if err := registry.Register(
retriever.NewKVHybridRetrieveEngine(
elasticsearchRepo, types.ElasticsearchRetrieverEngineType,
),
); err != nil {
log.Errorf("Register elasticsearch_v7 retrieve engine failed: %v", err)
} else {
log.Infof("Register elasticsearch_v7 retrieve engine success")
}
}
}
if slices.Contains(retrieveDriver, "opensearch") {
cc := &types.ConnectionConfig{
Addr: os.Getenv("OPENSEARCH_ADDR"),
Username: os.Getenv("OPENSEARCH_USERNAME"),
Password: os.Getenv("OPENSEARCH_PASSWORD"),
InsecureSkipVerify: strings.EqualFold(os.Getenv("OPENSEARCH_INSECURE_SKIP_VERIFY"), "true"),
}
client, err := openSearchRepo.NewOpenSearchClient(cc)
if err != nil {
log.Errorf("Create opensearch client failed: %v", err)
} else if repo, err := openSearchRepo.NewRepository(
context.Background(), client, "", nil, openSearchRepo.WithAuditSink(auditSink),
); err != nil {
log.Errorf("Create opensearch repository failed: %v", err)
} else if err := registry.Register(
retriever.NewKVHybridRetrieveEngine(repo, types.OpenSearchRetrieverEngineType),
); err != nil {
log.Errorf("Register opensearch retrieve engine failed: %v", err)
} else {
log.Infof("Register opensearch retrieve engine success")
}
}
if slices.Contains(retrieveDriver, "qdrant") {
qdrantHost := os.Getenv("QDRANT_HOST")
if qdrantHost == "" {
qdrantHost = "localhost"
}
qdrantPort := 6334 // Default port
if portStr := os.Getenv("QDRANT_PORT"); portStr != "" {
if port, err := strconv.Atoi(portStr); err == nil {
qdrantPort = port
}
}
// API key for authentication (optional)
qdrantAPIKey := os.Getenv("QDRANT_API_KEY")
// TLS configuration (optional, defaults to false)
// Enable TLS unless explicitly set to "false" or "0" (case insensitive)
qdrantUseTLS := false
if useTLSStr := os.Getenv("QDRANT_USE_TLS"); useTLSStr != "" {
useTLSLower := strings.ToLower(strings.TrimSpace(useTLSStr))
qdrantUseTLS = useTLSLower != "false" && useTLSLower != "0"
}
log.Infof("Connecting to Qdrant at %s:%d (TLS: %v)", qdrantHost, qdrantPort, qdrantUseTLS)
client, err := qdrant.NewClient(&qdrant.Config{
Host: qdrantHost,
Port: qdrantPort,
APIKey: qdrantAPIKey,
UseTLS: qdrantUseTLS,
})
if err != nil {
log.Errorf("Create qdrant client failed: %v", err)
} else {
qdrantRepository := qdrantRepo.NewQdrantRetrieveEngineRepository(client, nil)
if err := registry.Register(
retriever.NewKVHybridRetrieveEngine(
qdrantRepository, types.QdrantRetrieverEngineType,
),
); err != nil {
log.Errorf("Register qdrant retrieve engine failed: %v", err)
} else {
log.Infof("Register qdrant retrieve engine success")
}
}
}
if slices.Contains(retrieveDriver, "weaviate") {
weaviateHost := os.Getenv("WEAVIATE_HOST")
if weaviateHost == "" {
// Docker compose default (service name inside network)
weaviateHost = "weaviate:8080"
}
weaviateGrpcAddress := os.Getenv("WEAVIATE_GRPC_ADDRESS")
if weaviateGrpcAddress == "" {
weaviateGrpcAddress = "weaviate:50051"
}
weaviateScheme := os.Getenv("WEAVIATE_SCHEME")
if weaviateScheme == "" {
weaviateScheme = "http"
}
var authConfig auth.Config
if strings.EqualFold(strings.TrimSpace(os.Getenv("WEAVIATE_AUTH_ENABLED")), "true") {
if apiKey := strings.TrimSpace(os.Getenv("WEAVIATE_API_KEY")); apiKey != "" {
authConfig = auth.ApiKey{Value: apiKey}
}
}
weaviateClient, err := weaviate.NewClient(weaviate.Config{
Host: weaviateHost,
GrpcConfig: &wgrpc.Config{
Host: weaviateGrpcAddress,
},
Scheme: weaviateScheme,
AuthConfig: authConfig,
})
if err != nil {
log.Errorf("Create weaviate client failed: %v", err)
} else {
weaviateRepository := weaviateRepo.NewWeaviateRetrieveEngineRepository(weaviateClient, nil)
if err := registry.Register(
retriever.NewKVHybridRetrieveEngine(
weaviateRepository, types.WeaviateRetrieverEngineType,
),
); err != nil {
log.Errorf("Register weaviate retrieve engine failed: %v", err)
} else {
log.Infof("Register weaviate retrieve engine success")
}
}
}
if slices.Contains(retrieveDriver, "milvus") {
milvusCfg := milvusclient.ClientConfig{
DialOptions: []grpc.DialOption{grpc.WithTimeout(5 * time.Second)},
}
milvusAddress := os.Getenv("MILVUS_ADDRESS")
if milvusAddress == "" {
milvusAddress = "localhost:19530"
}
milvusCfg.Address = milvusAddress
milvusUsername := os.Getenv("MILVUS_USERNAME")
if milvusUsername != "" {
milvusCfg.Username = milvusUsername
}
milvusPassword := os.Getenv("MILVUS_PASSWORD")
if milvusPassword != "" {
milvusCfg.Password = milvusPassword
}
milvusDBName := os.Getenv("MILVUS_DB_NAME")
if milvusDBName != "" {
milvusCfg.DBName = milvusDBName
}
milvusCli, err := milvusclient.New(context.Background(), &milvusCfg)
if err != nil {
log.Errorf("Create milvus client failed: %v", err)
} else {
milvusRepository := milvusRepo.NewMilvusRetrieveEngineRepository(milvusCli, nil)
if err := registry.Register(
retriever.NewKVHybridRetrieveEngine(
milvusRepository, types.MilvusRetrieverEngineType,
),
); err != nil {
log.Errorf("Register milvus retrieve engine failed: %v", err)
} else {
log.Infof("Register milvus retrieve engine success")
}
}
}
if slices.Contains(retrieveDriver, "doris") {
dorisAddr := os.Getenv("DORIS_ADDR")
if dorisAddr == "" {
// docker-compose 默认服务名 + Doris FE MySQL 端口
dorisAddr = "doris-fe:9030"
}
dorisDatabase := os.Getenv("DORIS_DATABASE")
if dorisDatabase == "" {
dorisDatabase = "weknora"
}
dorisUsername := os.Getenv("DORIS_USERNAME")
if dorisUsername == "" {
dorisUsername = "root"
}
dorisPassword := os.Getenv("DORIS_PASSWORD")
dorisHTTPPort := 8030
if portStr := os.Getenv("DORIS_HTTP_PORT"); portStr != "" {
if port, err := strconv.Atoi(portStr); err == nil {
dorisHTTPPort = port
}
}
dsn := fmt.Sprintf("%s:%s@tcp(%s)/%s?charset=utf8mb4&parseTime=true&loc=Local&interpolateParams=true",
dorisUsername, dorisPassword, dorisAddr, dorisDatabase)
dorisDB, err := sql.Open("mysql", dsn)
if err != nil {
log.Errorf("Create doris client failed: %v", err)
} else {
dorisDB.SetMaxOpenConns(20)
dorisDB.SetMaxIdleConns(5)
dorisDB.SetConnMaxLifetime(time.Hour)
httpBase := "http://" + hostFromAddr(dorisAddr) + ":" + strconv.Itoa(dorisHTTPPort)
dorisRepository := dorisRepo.NewDorisRetrieveEngineRepository(
dorisDB, httpBase, dorisUsername, dorisPassword, dorisDatabase, nil,
)
if err := registry.Register(
retriever.NewKVHybridRetrieveEngine(
dorisRepository, types.DorisRetrieverEngineType,
),
); err != nil {
log.Errorf("Register doris retrieve engine failed: %v", err)
} else {
log.Infof("Register doris retrieve engine success: %s db=%s", dorisAddr, dorisDatabase)
}
}
}
if slices.Contains(retrieveDriver, "tencent_vectordb") {
addr := os.Getenv("TENCENT_VECTORDB_ADDR")
username := os.Getenv("TENCENT_VECTORDB_USERNAME")
apiKey := os.Getenv("TENCENT_VECTORDB_API_KEY")
if addr == "" || username == "" || apiKey == "" {
log.Errorf("Missing Tencent VectorDB configuration")
} else {
client, err := tcvectordb.NewRpcClient(addr, username, apiKey, &tcvectordb.ClientOption{
ReadConsistency: tcvectordb.EventualConsistency,
Timeout: 10 * time.Second,
})
if err != nil {
log.Errorf("Create tencent_vectordb client failed: %v", err)
} else {
tencentRepository := tencentVectorDBRepo.NewTencentVectorDBRetrieveEngineRepository(
client,
os.Getenv("TENCENT_VECTORDB_DATABASE"),
nil,
)
if err := registry.Register(
retriever.NewKVHybridRetrieveEngine(
tencentRepository, types.TencentVectorDBRetrieverEngineType,
),
); err != nil {
log.Errorf("Register tencent_vectordb retrieve engine failed: %v", err)
} else {
log.Infof("Register tencent_vectordb retrieve engine success")
}
}
}
}
// ─── DB store registration (byStoreID) ───
if storeReg, ok := registry.(*retriever.RetrieveEngineRegistry); ok {
loadDBStoresIntoRegistry(storeReg, db, cfg, auditSink)
}
return registry, nil
}
// loadDBStoresIntoRegistry loads VectorStore records from DB and registers them
// in the registry's byStoreID map. Failures are logged and skipped (non-fatal).
func loadDBStoresIntoRegistry(
storeRegistry interfaces.StoreRegistry, db *gorm.DB, cfg *config.Config, auditSink openSearchRepo.AuditSink,
) {
ctx := context.Background()
log := logger.GetLogger(ctx)
var stores []types.VectorStore
// GORM soft delete automatically adds "deleted_at IS NULL" condition
if err := db.Find(&stores).Error; err != nil {
log.Warnf("Failed to load vector stores from DB: %v", err)
return
}
if len(stores) == 0 {
return
}
log.Infof("Loading %d vector store(s) from database", len(stores))
for _, store := range stores {
svc, err := createEngineServiceFromStore(ctx, store, db, cfg, auditSink)
if err != nil {
log.Errorf("Failed to create engine for store %s (%s): %v", store.ID, store.Name, err)
continue
}
storeRegistry.RegisterWithStoreID(store.ID, svc)
log.Infof("Registered DB vector store: id=%s, name=%s, engine=%s", store.ID, store.Name, store.EngineType)
}
}
// initAntsPool initializes the goroutine pool
// Creates a managed goroutine pool for concurrent task execution
// Parameters:
// - cfg: Application configuration
//
// Returns:
// - Configured goroutine pool
// - Error if initialization fails
func initAntsPool(cfg *config.Config) (*ants.Pool, error) {
// Default to 5 if not specified in config
poolSize := os.Getenv("CONCURRENCY_POOL_SIZE")
if poolSize == "" {
poolSize = "5"
}
poolSizeInt, err := strconv.Atoi(poolSize)
if err != nil {
return nil, err
}
// Set up the pool with pre-allocation for better performance
return ants.NewPool(poolSizeInt, ants.WithPreAlloc(true))
}
// registerPoolCleanup registers the goroutine pool for cleanup
// Ensures proper cleanup of the goroutine pool when application shuts down
// Parameters:
// - pool: Goroutine pool
// - cleaner: Resource cleaner
func registerPoolCleanup(pool *ants.Pool, cleaner interfaces.ResourceCleaner) {
cleaner.RegisterWithName("AntsPool", func() error {
pool.Release()
return nil
})
}
// registerLangfuseCleanup ensures buffered Langfuse events are flushed on
// shutdown. A 5-second timeout matches other external-service cleanups and
// balances data durability against a slow remote endpoint holding up exit.
func registerLangfuseCleanup(mgr *langfuse.Manager, cleaner interfaces.ResourceCleaner) {
if mgr == nil {
return
}
cleaner.RegisterWithName("Langfuse", func() error {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
return mgr.Shutdown(ctx)
})
}
// initDocReaderClient initializes the DocumentReader client (lightweight API).
func initDocReaderClient(cfg *config.Config) (interfaces.DocumentReader, error) {
addr := strings.TrimSpace(os.Getenv("DOCREADER_ADDR"))
transport := strings.TrimSpace(os.Getenv("DOCREADER_TRANSPORT"))
if transport == "" {
transport = "grpc"
}
if addr == "" {
logger.Infof(context.Background(), "[DocConverter] No DOCREADER_ADDR configured, starting disconnected")
}
transport = strings.ToLower(transport)
switch transport {
case "http", "https":
if addr != "" && !strings.HasPrefix(addr, "http://") && !strings.HasPrefix(addr, "https://") {
addr = "http://" + addr
}
return docparser.NewHTTPDocumentReader(addr)
default:
return docparser.NewGRPCDocumentReader(addr)
}
}
// initOllamaService initializes the Ollama service client
// Creates a client for interacting with Ollama API for model inference
// Parameters:
// - None
//
// Returns:
// - Configured Ollama service client
// - Error if initialization fails
func initOllamaService() (*ollama.OllamaService, error) {
// Get Ollama service from existing factory function
return ollama.GetOllamaService()
}
func initNeo4jClient() (neo4j.Driver, error) {
ctx := context.Background()
if strings.ToLower(os.Getenv("NEO4J_ENABLE")) != "true" {
logger.Debugf(ctx, "NOT SUPPORT RETRIEVE GRAPH")
return nil, nil
}
uri := os.Getenv("NEO4J_URI")
username := os.Getenv("NEO4J_USERNAME")
password := os.Getenv("NEO4J_PASSWORD")
// Retry configuration
maxRetries := 30 // Max retry attempts
retryInterval := 2 * time.Second // Wait between retries
var driver neo4j.Driver
var err error
for attempt := 1; attempt <= maxRetries; attempt++ {
driver, err = neo4j.NewDriver(uri, neo4j.BasicAuth(username, password, ""))
if err != nil {
logger.Warnf(ctx, "Failed to create Neo4j driver (attempt %d/%d): %v", attempt, maxRetries, err)
time.Sleep(retryInterval)
continue
}
err = driver.VerifyAuthentication(ctx, nil)
if err == nil {
if attempt > 1 {
logger.Infof(ctx, "Successfully connected to Neo4j after %d attempts", attempt)
}
return driver, nil
}
logger.Warnf(ctx, "Failed to verify Neo4j authentication (attempt %d/%d): %v", attempt, maxRetries, err)
driver.Close(ctx)
time.Sleep(retryInterval)
}
return nil, fmt.Errorf("failed to connect to Neo4j after %d attempts: %w", maxRetries, err)
}
func NewDuckDB() (*sql.DB, error) {
sqlDB, err := sql.Open("duckdb", ":memory:")
if err != nil {
return nil, fmt.Errorf("failed to open duckdb: %w", err)
}
// Try to install and load required extensions unless explicitly disabled.
// - spatial: used for st_read_meta() to enumerate layer (sheet) names from .xlsx/.xls
// - excel: used for read_xlsx() which gives proper type inference per sheet
//
// INSTALL hits extensions.duckdb.org (public internet). In locked-down
// runtimes with no egress, set DUCKDB_SKIP_EXTENSION_LOAD=1 to avoid a
// startup hang; xlsx/xls ingest may fail later without these extensions.
if strings.EqualFold(os.Getenv("DUCKDB_SKIP_EXTENSION_LOAD"), "true") ||
os.Getenv("DUCKDB_SKIP_EXTENSION_LOAD") == "1" {
logger.Infof(context.Background(),
"[DuckDB] Skipping spatial/excel extension install/load "+
"(DUCKDB_SKIP_EXTENSION_LOAD is set; xlsx ingest may fail without them)")
} else {
bgCtx := context.Background()
for _, ext := range []string{"spatial", "excel"} {
if _, err := sqlDB.ExecContext(bgCtx, fmt.Sprintf("INSTALL %s;", ext)); err != nil {
logger.Warnf(bgCtx, "[DuckDB] Failed to install %s extension: %v", ext, err)
}
if _, err := sqlDB.ExecContext(bgCtx, fmt.Sprintf("LOAD %s;", ext)); err != nil {
logger.Warnf(bgCtx, "[DuckDB] Failed to load %s extension: %v", ext, err)
}
}
}
return sqlDB, nil
}
// registerWebSearchProviders registers all web search provider types to the registry.
// Each provider type is registered with its factory function that accepts parameters.
// Provider instances are created on-demand when tenants configure them.
func registerWebSearchProviders(registry *infra_web_search.Registry) {
registry.Register("duckduckgo", infra_web_search.NewDuckDuckGoProvider)
registry.Register("google", infra_web_search.NewGoogleProvider)
registry.Register("bing", infra_web_search.NewBingProvider)
registry.Register("tavily", infra_web_search.NewTavilyProvider)
registry.Register("ollama", infra_web_search.NewOllamaProvider)
registry.Register("baidu", infra_web_search.NewBaiduProvider)
registry.Register("searxng", infra_web_search.NewSearxngProvider)
registry.Register("keenable", infra_web_search.NewKeenableProvider)
registry.Register("zhipu", infra_web_search.NewZhipuProvider)
registry.Register("exa", infra_web_search.NewExaProvider)
registry.Register("metaso", infra_web_search.NewMetasoProvider)
}
// registerIMService registers adapter factories, loads enabled channels, and
// wires the process-lifetime shutdown hook. Each platform's factory lives in
// its own subpackage to keep this file focused on wiring.
func registerIMService(imService *imPkg.Service, cleaner interfaces.ResourceCleaner) {
imService.RegisterAdapterFactory("wecom", wecom.NewFactory())
imService.RegisterAdapterFactory("feishu", feishu.NewFactory(feishu.RegionFeishu))
// Lark is Feishu's international cloud: same adapter, different host/tenant.
imService.RegisterAdapterFactory("lark", feishu.NewFactory(feishu.RegionLark))
imService.RegisterAdapterFactory("slack", slack.NewFactory())
imService.RegisterAdapterFactory("telegram", telegram.NewFactory())
imService.RegisterAdapterFactory("dingtalk", dingtalk.NewFactory())
imService.RegisterAdapterFactory("mattermost", mattermost.NewFactory())
imService.RegisterAdapterFactory("wechat", wechat.NewFactory())
imService.RegisterAdapterFactory("qqbot", qqbot.NewFactory())
imService.RegisterAdapterFactory("yunzhijia", yunzhijia.NewFactory())
// Load and start all enabled channels from database
if err := imService.LoadAndStartChannels(); err != nil {
logger.Warnf(context.Background(), "[IM] Failed to load channels from database: %v", err)
}
cleaner.RegisterWithName("IMService", func() error {
imService.Stop()
return nil
})
}
// initConnectorRegistry creates and populates the connector registry with all available connectors.
// Aggregates registration errors via errors.Join so a misconfigured or duplicated connector fails
// container initialization loudly instead of silently disabling the feature at runtime.
func initConnectorRegistry() (*datasource.ConnectorRegistry, error) {
registry := datasource.NewConnectorRegistry()
var errs error
if err := registry.Register(wiki.NewConnector(core.RegionFeishu)); err != nil {
errs = errors.Join(errs, fmt.Errorf("register feishu connector: %w", err))
}
// Lark is Feishu's international cloud: same connector, different host/tenant.
if err := registry.Register(wiki.NewConnector(core.RegionLark)); err != nil {
errs = errors.Join(errs, fmt.Errorf("register lark connector: %w", err))
}
// Feishu/Lark Drive (云盘) mode: different connector type so the registry
// dispatches to the Drive connector. Shares core.Client/Region/export logic
// with the wiki connector. See 飞书云盘数据源设计.md / ADR-0001.
if err := registry.Register(drive.NewDriveConnector(core.RegionFeishuDrive)); err != nil {
errs = errors.Join(errs, fmt.Errorf("register feishu_drive connector: %w", err))
}
if err := registry.Register(drive.NewDriveConnector(core.RegionLarkDrive)); err != nil {
errs = errors.Join(errs, fmt.Errorf("register lark_drive connector: %w", err))
}
if err := registry.Register(notionConnector.NewConnector()); err != nil {
errs = errors.Join(errs, fmt.Errorf("register notion connector: %w", err))
}
if err := registry.Register(yuqueConnector.NewConnector()); err != nil {
errs = errors.Join(errs, fmt.Errorf("register yuque connector: %w", err))
}
if err := registry.Register(imaConnector.NewConnector()); err != nil {
errs = errors.Join(errs, fmt.Errorf("register ima connector: %w", err))
}
if err := registry.Register(rssConnector.NewConnector()); err != nil {
errs = errors.Join(errs, fmt.Errorf("register rss connector: %w", err))
}
if err := registry.Register(gitlabConnector.NewConnector()); err != nil {
errs = errors.Join(errs, fmt.Errorf("register gitlab connector: %w", err))
}
// Future connectors will be registered here:
// if err := registry.Register(confluenceConnector.NewConnector()); err != nil { ... }
// if err := registry.Register(githubConnector.NewConnector()); err != nil { ... }
if errs != nil {
return nil, errs
}
return registry, nil
}
// startDataSourceScheduler starts the data source cron scheduler and registers cleanup.
func startDataSourceScheduler(scheduler *datasource.Scheduler, cleaner interfaces.ResourceCleaner) {
if err := scheduler.Start(context.Background()); err != nil {
logger.Warnf(context.Background(), "[Container] data source scheduler start failed: %v", err)
}
cleaner.RegisterWithName("DataSourceScheduler", func() error {
scheduler.Stop()
return nil
})
}
// startHousekeepingService starts the knowledge housekeeping cron and registers
// cleanup. This is the safety net that recovers any knowledge stuck in
// "processing" past a configurable threshold (see HousekeepingService for
// rationale). Best-effort: a startup error is logged but does NOT abort the
// container — the rest of the system stays usable.
func startHousekeepingService(svc *service.HousekeepingService, cleaner interfaces.ResourceCleaner) {
if svc == nil {
return
}
if err := svc.Start(context.Background()); err != nil {
logger.Warnf(context.Background(), "[Container] housekeeping start failed: %v", err)
}
cleaner.RegisterWithName("KnowledgeHousekeeping", func() error {
svc.Stop()
return nil
})
}
// startTenantSkillReaper starts the stuck-install / orphan-snapshot cron and
// registers cleanup. Best-effort: a startup error is logged but does NOT abort
// the container — the rest of the system stays usable.
func startTenantSkillReaper(svc *service.TenantSkillService, cleaner interfaces.ResourceCleaner) {
if svc == nil {
return
}
if err := svc.Start(context.Background()); err != nil {
logger.Warnf(context.Background(), "[Container] tenant skill reaper start failed: %v", err)
}
cleaner.RegisterWithName("TenantSkillReaper", func() error {
svc.Stop()
return nil
})
}
// startTemporaryDocumentCleanup removes expired session attachments and their
// extracted images. The durable expiry timestamp is the source of truth; the
// ticker only controls how quickly storage is reclaimed.
func startTemporaryDocumentCleanup(svc interfaces.TemporaryDocumentService, cleaner interfaces.ResourceCleaner) {
stop := make(chan struct{})
go func() {
ticker := time.NewTicker(10 * time.Minute)
defer ticker.Stop()
for {
select {
case <-ticker.C:
if err := svc.CleanupExpired(context.Background()); err != nil {
logger.Warnf(context.Background(), "[TemporaryDocument] cleanup failed: %v", err)
}
case <-stop:
return
}
}
}()
cleaner.RegisterWithName("TemporaryDocumentCleanup", func() error {
close(stop)
return nil
})
}
// startAuditLogRetention spins up the daily audit_logs purge sweep
// and registers shutdown cleanup. Mirrors the data-source-scheduler
// pattern: container init kicks the goroutine, ResourceCleaner stops
// it during graceful shutdown so a SIGTERM during a sweep doesn't
// orphan the goroutine.
//
// retention_days <= 0 is the configured way to disable retention;
// the runner short-circuits Start() on that path so we don't need
// to gate the wiring here.
func startAuditLogRetention(
runner *service.AuditLogRetentionRunner, cleaner interfaces.ResourceCleaner,
) {
runner.Start(context.Background())
cleaner.RegisterWithName("AuditLogRetentionRunner", func() error {
runner.Stop()
return nil
})
}