From a879bbdbb5b566bee4cd65d8d885378af8a74765 Mon Sep 17 00:00:00 2001
From: wizardchen <wizardchen@tencent.com>
Date: Thu, 20 Nov 2025 16:35:07 +0800
Subject: [PATCH] feat: Integrate Neo4j driver into SystemHandler for enhanced
 graph database support

- Updated SystemHandler to include a Neo4j driver, allowing for improved interaction with graph databases.
- Modified NewSystemHandler to accept the Neo4j driver as a parameter.
- Adjusted getGraphDatabaseEngine method to return the appropriate engine name based on Neo4j driver availability.
---
 internal/handler/system.go                       | 16 +++++++++-------
 .../000011_remove_kb_rerank_model.down.sql       | 11 +++++++++++
 .../000011_remove_kb_rerank_model.up.sql         | 11 +++++++++++
 3 files changed, 31 insertions(+), 7 deletions(-)
 create mode 100644 migrations/versioned/000011_remove_kb_rerank_model.down.sql
 create mode 100644 migrations/versioned/000011_remove_kb_rerank_model.up.sql

diff --git a/internal/handler/system.go b/internal/handler/system.go
index 598a06e9e..b8013501a 100644
--- a/internal/handler/system.go
+++ b/internal/handler/system.go
@@ -7,17 +7,20 @@ import (
 	"github.com/Tencent/WeKnora/internal/config"
 	"github.com/Tencent/WeKnora/internal/logger"
 	"github.com/gin-gonic/gin"
+	"github.com/neo4j/neo4j-go-driver/v6/neo4j"
 )
 
 // SystemHandler handles system-related requests
 type SystemHandler struct {
-	cfg *config.Config
+	cfg         *config.Config
+	neo4jDriver neo4j.Driver
 }
 
 // NewSystemHandler creates a new system handler
-func NewSystemHandler(cfg *config.Config) *SystemHandler {
+func NewSystemHandler(cfg *config.Config, neo4jDriver neo4j.Driver) *SystemHandler {
 	return &SystemHandler{
-		cfg: cfg,
+		cfg:         cfg,
+		neo4jDriver: neo4jDriver,
 	}
 }
 
@@ -130,11 +133,10 @@ func (h *SystemHandler) getVectorStoreEngine() string {
 
 // getGraphDatabaseEngine returns the graph database engine name
 func (h *SystemHandler) getGraphDatabaseEngine() string {
-	neo4jEnable := strings.ToLower(os.Getenv("NEO4J_ENABLE"))
-	if neo4jEnable == "true" {
-		return "Neo4j"
+	if h.neo4jDriver == nil {
+		return "未启用"
 	}
-	return "未启用"
+	return "Neo4j"
 }
 
 // isMinioEnabled checks if MinIO is enabled
diff --git a/migrations/versioned/000011_remove_kb_rerank_model.down.sql b/migrations/versioned/000011_remove_kb_rerank_model.down.sql
new file mode 100644
index 000000000..d354e2315
--- /dev/null
+++ b/migrations/versioned/000011_remove_kb_rerank_model.down.sql
@@ -0,0 +1,11 @@
+-- 000011_remove_kb_rerank_model.down.sql
+-- Reintroduce rerank_model_id column to knowledge_bases table
+
+BEGIN;
+
+ALTER TABLE knowledge_bases
+    ADD COLUMN IF NOT EXISTS rerank_model_id VARCHAR(64);
+
+COMMIT;
+
+
diff --git a/migrations/versioned/000011_remove_kb_rerank_model.up.sql b/migrations/versioned/000011_remove_kb_rerank_model.up.sql
new file mode 100644
index 000000000..a22db423a
--- /dev/null
+++ b/migrations/versioned/000011_remove_kb_rerank_model.up.sql
@@ -0,0 +1,11 @@
+-- 000011_remove_kb_rerank_model.up.sql
+-- Remove rerank_model_id column from knowledge_bases table
+
+BEGIN;
+
+ALTER TABLE knowledge_bases
+    DROP COLUMN IF EXISTS rerank_model_id;
+
+COMMIT;
+
+