WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes
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

This commit is contained in:
yawan
2026-09-22 16:19:06 +08:00
commit f7fc992f57
3929 changed files with 1006607 additions and 0 deletions
@@ -0,0 +1,228 @@
package doris
import (
"context"
"database/sql"
"fmt"
"os"
"sort"
"strings"
"github.com/Tencent/WeKnora/internal/logger"
)
const envDorisCompatMode = "DORIS_COMPAT_MODE"
type dorisCompatMode string
const (
dorisCompatModeAuto dorisCompatMode = "auto"
dorisCompatModeLegacy dorisCompatMode = "legacy"
dorisCompatModeInnerProductDuplicate dorisCompatMode = "inner_product_duplicate"
)
type dorisCompatProbe struct {
innerProductApproximate bool
cosineDistanceApproximate bool
}
func resolveConfiguredDorisCompatMode() (mode dorisCompatMode, invalid string) {
raw := strings.TrimSpace(os.Getenv(envDorisCompatMode))
if raw == "" {
return dorisCompatModeAuto, ""
}
switch strings.ToLower(raw) {
case string(dorisCompatModeAuto):
return dorisCompatModeAuto, ""
case string(dorisCompatModeLegacy):
return dorisCompatModeLegacy, ""
case string(dorisCompatModeInnerProductDuplicate), "inner-product-duplicate", "inner_product", "inner-product":
return dorisCompatModeInnerProductDuplicate, ""
default:
return dorisCompatModeAuto, raw
}
}
func (r *dorisRepository) resolveCompatMode(ctx context.Context) (dorisCompatMode, error) {
r.compatResolveOnce.Do(func() {
log := logger.GetLogger(ctx)
requested := r.compatModeRequested
if requested == "" {
r.compatModeResolved = dorisCompatModeInnerProductDuplicate
return
}
existingMode, exampleTable, foundExisting, err := r.detectExistingCompatMode(ctx)
if err != nil {
r.compatResolveErr = err
return
}
if foundExisting {
if requested != dorisCompatModeAuto && requested != existingMode {
r.compatResolveErr = fmt.Errorf(
"Doris compat mode %q does not match existing embedding tables (detected %q from %s). %s is not interchangeable after %s_* tables are created. Recreate the existing %s_* tables before switching modes, or set %s=%s",
requested,
existingMode,
exampleTable,
envDorisCompatMode,
r.tableBaseName,
r.tableBaseName,
envDorisCompatMode,
existingMode,
)
log.Errorf("[Doris] %v", r.compatResolveErr)
return
}
r.compatModeResolved = existingMode
log.Warnf(
"[Doris] Using compat mode %s from existing embedding tables (%s). %s is not interchangeable after %s_* tables are created; recreate those tables before switching modes.",
existingMode,
exampleTable,
envDorisCompatMode,
r.tableBaseName,
)
return
}
if requested == dorisCompatModeAuto {
probe, probeErr := r.probeCompatMode(ctx)
if probeErr != nil {
r.compatResolveErr = probeErr
log.Errorf("[Doris] %v", probeErr)
return
}
log.Infof(
"[Doris] Compat probe result: inner_product_approximate=%t, cosine_distance_approximate=%t",
probe.innerProductApproximate,
probe.cosineDistanceApproximate,
)
switch {
case probe.innerProductApproximate:
r.compatModeResolved = dorisCompatModeInnerProductDuplicate
case probe.cosineDistanceApproximate:
r.compatModeResolved = dorisCompatModeLegacy
default:
r.compatResolveErr = fmt.Errorf(
"Doris compatibility auto-detection could not find a supported vector function. Set %s=%s or %s=%s explicitly after verifying your Doris build. %s is not interchangeable after %s_* tables are created",
envDorisCompatMode,
dorisCompatModeInnerProductDuplicate,
envDorisCompatMode,
dorisCompatModeLegacy,
envDorisCompatMode,
r.tableBaseName,
)
log.Errorf("[Doris] %v", r.compatResolveErr)
return
}
log.Warnf(
"[Doris] Auto-selected compat mode %s for new embedding tables. %s is not interchangeable after %s_* tables are created; recreate those tables before switching modes.",
r.compatModeResolved,
envDorisCompatMode,
r.tableBaseName,
)
return
}
r.compatModeResolved = requested
log.Warnf(
"[Doris] Using configured compat mode %s for new embedding tables. %s is not interchangeable after %s_* tables are created; recreate those tables before switching modes.",
r.compatModeResolved,
envDorisCompatMode,
r.tableBaseName,
)
})
return r.compatModeResolved, r.compatResolveErr
}
func (r *dorisRepository) detectExistingCompatMode(ctx context.Context) (dorisCompatMode, string, bool, error) {
tables, err := r.listEmbeddingTables(ctx)
if err != nil {
return "", "", false, fmt.Errorf("list Doris embedding tables: %w", err)
}
if len(tables) == 0 {
return "", "", false, nil
}
sort.Strings(tables)
var detected dorisCompatMode
for _, table := range tables {
ddl, err := r.showCreateTable(ctx, table)
if err != nil {
return "", "", false, fmt.Errorf("show create table %s: %w", table, err)
}
mode, err := compatModeFromDDL(ddl)
if err != nil {
return "", "", false, fmt.Errorf("detect compat mode from %s: %w", table, err)
}
if detected == "" {
detected = mode
continue
}
if detected != mode {
return "", "", false, fmt.Errorf(
"existing Doris embedding tables use mixed compat modes (%s and %s). %s is not interchangeable after %s_* tables are created; recreate the existing %s_* tables with a single mode",
detected,
mode,
envDorisCompatMode,
r.tableBaseName,
r.tableBaseName,
)
}
}
return detected, tables[0], true, nil
}
func (r *dorisRepository) showCreateTable(ctx context.Context, table string) (string, error) {
row := r.db.QueryRowContext(ctx, fmt.Sprintf("SHOW CREATE TABLE `%s`", table))
var tableName, ddl string
if err := row.Scan(&tableName, &ddl); err != nil {
return "", err
}
return ddl, nil
}
func compatModeFromDDL(ddl string) (dorisCompatMode, error) {
normalized := strings.ToLower(ddl)
switch {
case strings.Contains(normalized, "duplicate key("):
return dorisCompatModeInnerProductDuplicate, nil
case strings.Contains(normalized, "unique key("):
return dorisCompatModeLegacy, nil
default:
return "", fmt.Errorf("unsupported table definition; expected UNIQUE KEY or DUPLICATE KEY")
}
}
func (r *dorisRepository) probeCompatMode(ctx context.Context) (dorisCompatProbe, error) {
probe := dorisCompatProbe{}
probe.innerProductApproximate = r.vectorFunctionSupported(ctx, "inner_product_approximate([1.0],[1.0])")
probe.cosineDistanceApproximate = r.vectorFunctionSupported(ctx, "cosine_distance_approximate([1.0],[1.0])")
return probe, nil
}
func (r *dorisRepository) vectorFunctionSupported(ctx context.Context, expr string) bool {
var value sql.NullFloat64
if err := r.db.QueryRowContext(ctx, "SELECT "+expr).Scan(&value); err != nil {
return false
}
return true
}
func (m dorisCompatMode) normalizeEmbeddings() bool {
return m != dorisCompatModeLegacy
}
func (m dorisCompatMode) usesReplaceWrite() bool {
return m != dorisCompatModeLegacy
}
func (m dorisCompatMode) usesRewriteChunkUpdates() bool {
return m != dorisCompatModeLegacy
}
@@ -0,0 +1,252 @@
package doris
import (
"math"
"strconv"
"strings"
"github.com/Tencent/WeKnora/internal/types"
)
// 字段名常量。Doris 是 SQL 库,字段名要在 SELECT/WHERE/INSERT 多处复用,
// 用常量统一防止笔误。
const (
fieldID = "id"
fieldContent = "content"
fieldSourceID = "source_id"
fieldSourceType = "source_type"
fieldChunkID = "chunk_id"
fieldKnowledgeID = "knowledge_id"
fieldKnowledgeBaseID = "knowledge_base_id"
fieldTagID = "tag_id"
fieldIsEnabled = "is_enabled"
fieldEmbedding = "embedding"
)
// columns 是 INSERT / SELECT 时使用的标准列序。
var columns = []string{
fieldID, fieldContent, fieldSourceID, fieldSourceType,
fieldChunkID, fieldKnowledgeID, fieldKnowledgeBaseID, fieldTagID,
fieldIsEnabled, fieldEmbedding,
}
// columnsForRetrieve 是 Retrieve 时 SELECT 的列序,
// 不包含 embedding(向量本身查询结果中无需返回,省带宽)。
var columnsForRetrieve = []string{
fieldID, fieldContent, fieldSourceID, fieldSourceType,
fieldChunkID, fieldKnowledgeID, fieldKnowledgeBaseID, fieldTagID,
fieldIsEnabled,
}
// columnsForCopy 是 CopyIndices 中分页 SELECT 时使用的列序,
// 比 columnsForRetrieve 多 embedding,因为复制目的是搬运向量本身。
var columnsForCopy = []string{
fieldID, fieldContent, fieldSourceID, fieldSourceType,
fieldChunkID, fieldKnowledgeID, fieldKnowledgeBaseID, fieldTagID,
fieldIsEnabled, fieldEmbedding,
}
// whereCond 表示一个 WHERE 子条件:clause 是参数化 SQL 片段(带 ? 占位),
// args 是对应顺序的参数值。所有用户输入字段(IDs)必须通过 args 传入,
// 严禁拼到 clause 字符串里。
type whereCond struct {
clause string
args []any
}
// whereBuilder 用于把 RetrieveParams 中的过滤条件翻译成 SQL WHERE 子句。
//
// 每个 add* 方法对应一种 IN / NOT IN / = 算子;最终 build() 用 AND 拼接。
type whereBuilder struct {
conds []whereCond
}
// addEqual 追加一个 field = ? 条件。
func (w *whereBuilder) addEqual(field string, value any) {
w.conds = append(w.conds, whereCond{
clause: field + " = ?",
args: []any{value},
})
}
// addIn 追加一个 field IN (?, ?, ...) 条件。values 为空时不追加任何东西。
func (w *whereBuilder) addIn(field string, values []string) {
if len(values) == 0 {
return
}
placeholders := make([]string, len(values))
args := make([]any, len(values))
for i, v := range values {
placeholders[i] = "?"
args[i] = v
}
w.conds = append(w.conds, whereCond{
clause: field + " IN (" + strings.Join(placeholders, ", ") + ")",
args: args,
})
}
// addNotIn 追加一个 field NOT IN (?, ?, ...) 条件。
func (w *whereBuilder) addNotIn(field string, values []string) {
if len(values) == 0 {
return
}
placeholders := make([]string, len(values))
args := make([]any, len(values))
for i, v := range values {
placeholders[i] = "?"
args[i] = v
}
w.conds = append(w.conds, whereCond{
clause: field + " NOT IN (" + strings.Join(placeholders, ", ") + ")",
args: args,
})
}
// build 返回 WHERE 子句(不含 "WHERE " 前缀)和参数数组。
// 没有任何条件时返回 ("1 = 1", nil),方便调用方无脑拼接。
func (w *whereBuilder) build() (string, []any) {
if len(w.conds) == 0 {
return "1 = 1", nil
}
parts := make([]string, len(w.conds))
var args []any
for i, c := range w.conds {
parts[i] = c.clause
args = append(args, c.args...)
}
return strings.Join(parts, " AND "), args
}
// buildBaseFilter 将 RetrieveParams 中的过滤条件翻译为 whereBuilder。
// 默认追加 is_enabled = TRUE,与 Qdrant/Milvus/Weaviate 保持一致:
// 关闭的 chunk 不参与检索。
func buildBaseFilter(params types.RetrieveParams) *whereBuilder {
w := &whereBuilder{}
w.addEqual(fieldIsEnabled, true)
if len(params.KnowledgeBaseIDs) > 0 {
w.addIn(fieldKnowledgeBaseID, params.KnowledgeBaseIDs)
}
if len(params.KnowledgeIDs) > 0 {
w.addIn(fieldKnowledgeID, params.KnowledgeIDs)
}
if len(params.TagIDs) > 0 {
w.addIn(fieldTagID, params.TagIDs)
}
if len(params.ExcludeKnowledgeIDs) > 0 {
w.addNotIn(fieldKnowledgeID, params.ExcludeKnowledgeIDs)
}
if len(params.ExcludeChunkIDs) > 0 {
w.addNotIn(fieldChunkID, params.ExcludeChunkIDs)
}
return w
}
// parseEmbeddingLiteral 解析 Doris ARRAY<FLOAT> 通过 MySQL 协议返回的
// 字面量字符串(形如 "[1,2,3]")为 []float32。
//
// CopyIndices 路径需要从源行读出向量本身再写回目标行;此处的解析容错优先:
// 不带 [] 也接受、空数组返回 nil。
func parseEmbeddingLiteral(raw []byte) ([]float32, error) {
s := strings.TrimSpace(string(raw))
if s == "" {
return nil, nil
}
s = strings.TrimPrefix(s, "[")
s = strings.TrimSuffix(s, "]")
if s == "" {
return nil, nil
}
parts := strings.Split(s, ",")
out := make([]float32, 0, len(parts))
for _, p := range parts {
p = strings.TrimSpace(p)
if p == "" {
continue
}
f, err := strconv.ParseFloat(p, 32)
if err != nil {
return nil, err
}
out = append(out, float32(f))
}
return out, nil
}
// validateEmbedding 校验向量元素均为有限值。
//
// strconv.FormatFloat 对 NaN/±Inf 会输出 "NaN"/"+Inf"/"-Inf",
// 这些字面量会让 Doris 拼出来的 SQL 报语法错误(或在某些版本下产生未定义结果)。
// 上游(嵌入模型)正常情况下不会输出非有限值,但 GPU OOM、上游 bug、
// 测试桩都可能触发;这里 fail-fast 比悄悄写脏数据安全。
func validateEmbedding(vec []float32) error {
for i, v := range vec {
if f := float64(v); math.IsNaN(f) || math.IsInf(f, 0) {
return errInvalidEmbedding{index: i, value: v}
}
}
return nil
}
// normalizeEmbedding returns a unit-length copy of vec so Doris inner-product
// ANN search can preserve cosine-style similarity semantics.
func normalizeEmbedding(vec []float32) []float32 {
if len(vec) == 0 {
return nil
}
var sumSquares float64
for _, value := range vec {
f := float64(value)
sumSquares += f * f
}
if sumSquares == 0 {
return append([]float32(nil), vec...)
}
norm := float32(math.Sqrt(sumSquares))
normalized := make([]float32, len(vec))
for i, value := range vec {
normalized[i] = value / norm
}
return normalized
}
// errInvalidEmbedding 描述哪个下标含非有限值;用结构体而非 fmt.Errorf
// 是为了让上层在日志里能拿到下标做问题定位。
type errInvalidEmbedding struct {
index int
value float32
}
func (e errInvalidEmbedding) Error() string {
return "doris: embedding[" + strconv.Itoa(e.index) +
"] is not finite: " + strconv.FormatFloat(float64(e.value), 'g', -1, 32)
}
// embeddingLiteral 把 []float32 转为 Doris ARRAY<FLOAT> 字面量字符串:
// "[1.23,4.56,...]"。
//
// 为何不用占位符:go-sql-driver/mysql 不支持 ARRAY 类型的参数绑定,
// Doris 端也只接受字面量形式。这里用 strconv.FormatFloat('g' + bitSize=32)
// 而不是 fmt.Sprintf("%f", v),原因有二:
// 1. fmt 在某些 locale 下会用千分位分隔符,破坏 SQL 语法;
// 2. 'g' 比 'f' 短且不会丢精度。
//
// 注入风险:[]float32 元素是嵌入模型输出的有限位数浮点数,序列化后只可能
// 包含 [0-9eE+-.\s] 字符,不会逃逸出字面量上下文。
func embeddingLiteral(vec []float32) string {
if len(vec) == 0 {
return "[]"
}
var sb strings.Builder
sb.Grow(len(vec) * 12)
sb.WriteByte('[')
for i, v := range vec {
if i > 0 {
sb.WriteByte(',')
}
sb.WriteString(strconv.FormatFloat(float64(v), 'g', -1, 32))
}
sb.WriteByte(']')
return sb.String()
}
@@ -0,0 +1,699 @@
package doris
import (
"context"
"database/sql"
"fmt"
"os"
"strings"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
"github.com/google/uuid"
)
const (
defaultTableBaseName = "weknora_embeddings"
envDorisTablePrefix = "DORIS_TABLE_PREFIX"
)
// NewDorisRetrieveEngineRepository 创建 Doris 检索引擎仓储。
//
// 参数:
// - db:MySQL 协议的 *sql.DB 实例。调用方负责 SetMaxOpenConns 等参数。
// - feHTTPBase:Stream Load 用的 FE HTTP 基地址(含 scheme),例如 "http://doris-fe:8030"。
// - username/password:MySQL 与 Stream Load 共用的凭据。
// - database:目标数据库名(既用于 MySQL DSN,也用于 Stream Load URL 路径)。
// - indexCfg:可空。为 nil 时退化为环境变量 + 默认值(env 路径)。
func NewDorisRetrieveEngineRepository(
db *sql.DB,
feHTTPBase, username, password, database string,
indexCfg *types.IndexConfig,
) interfaces.RetrieveEngineRepository {
log := logger.GetLogger(context.Background())
log.Info("[Doris] Initializing Doris retriever engine repository")
tableBaseName := types.ResolveCollectionName(indexCfg, envDorisTablePrefix, defaultTableBaseName)
compatMode, invalidCompatMode := resolveConfiguredDorisCompatMode()
if invalidCompatMode != "" {
log.Warnf("[Doris] Invalid %s=%q, defaulting to %s", envDorisCompatMode, invalidCompatMode, dorisCompatModeAuto)
}
repo := &dorisRepository{
db: db,
httpClient: newDorisStreamLoadHTTPClient(),
feHTTPBase: strings.TrimRight(feHTTPBase, "/"),
username: username,
password: password,
database: database,
tableBaseName: tableBaseName,
bucketsNum: indexCfg.GetBucketsNum(0),
replicationNum: indexCfg.GetReplicationNum(0),
compatModeRequested: compatMode,
}
log.Infof("[Doris] Repository initialized: db=%s, base=%s, fe_http=%s, compat_mode=%s",
database, tableBaseName, repo.feHTTPBase, repo.compatModeRequested)
if os.Getenv(envDorisCompatMode) == "" {
log.Infof("[Doris] %s not set, defaulting to %s and probing on first use", envDorisCompatMode, dorisCompatModeAuto)
}
return repo
}
func (r *dorisRepository) EngineType() types.RetrieverEngineType {
return types.DorisRetrieverEngineType
}
func (r *dorisRepository) Support() []types.RetrieverType {
return []types.RetrieverType{types.KeywordsRetrieverType, types.VectorRetrieverType}
}
// EstimateStorageSize 估算给定 IndexInfo 列表的存储字节数。
//
// 参考 Qdrant 的算法:payload 字段长度 + 向量字节 + HNSW 邻居 + 元数据。
func (r *dorisRepository) EstimateStorageSize(_ context.Context,
indexInfoList []*types.IndexInfo, params map[string]any,
) int64 {
var total int64
for _, info := range indexInfoList {
emb := toDorisVectorEmbedding(info, params, dorisCompatModeInnerProductDuplicate)
total += calculateStorageSize(emb)
}
return total
}
// Save 写入单条记录到对应维度的表。空向量在 BatchSave 内部统一拒绝。
func (r *dorisRepository) Save(ctx context.Context,
info *types.IndexInfo, additionalParams map[string]any,
) error {
return r.BatchSave(ctx, []*types.IndexInfo{info}, additionalParams)
}
// BatchSave 把同一批 IndexInfo 按维度分组;Doris ANN 兼容表使用
// DUPLICATE KEY,因此这里显式按 id 做 delete + insert,保持原先的替换语义。
func (r *dorisRepository) BatchSave(ctx context.Context,
indexInfoList []*types.IndexInfo, additionalParams map[string]any,
) error {
log := logger.GetLogger(ctx)
if len(indexInfoList) == 0 {
return nil
}
compatMode, err := r.resolveCompatMode(ctx)
if err != nil {
return err
}
groups := make(map[int][]*DorisVectorEmbedding)
for _, info := range indexInfoList {
emb := toDorisVectorEmbedding(info, additionalParams, compatMode)
if len(emb.Embedding) == 0 {
log.Warnf("[Doris] Skipping empty embedding for chunk %s", info.ChunkID)
continue
}
if err := validateEmbedding(emb.Embedding); err != nil {
return fmt.Errorf("invalid embedding for chunk %s: %w", info.ChunkID, err)
}
// 给一个稳定的主键。SourceID 是上层最有意义的"行身份",
// 但同 chunk 多 question 的场景下 SourceID 已经唯一,所以直接用它。
if emb.ID == "" {
emb.ID = emb.SourceID
}
if emb.ID == "" {
emb.ID = uuid.New().String()
}
dim := len(emb.Embedding)
groups[dim] = append(groups[dim], emb)
}
for dim, rows := range groups {
if err := r.ensureTable(ctx, dim); err != nil {
return err
}
if compatMode.usesReplaceWrite() {
err = r.replaceRows(ctx, r.getTableName(dim), rows)
} else {
err = r.insertRows(ctx, r.getTableName(dim), rows)
}
if err != nil {
return fmt.Errorf("batch save dim=%d: %w", dim, err)
}
log.Infof("[Doris] Saved %d rows to %s", len(rows), r.getTableName(dim))
}
return nil
}
// insertRows 按列序拼一条多 VALUES 的 INSERT。embedding 列由于
// go-sql-driver/mysql 不支持 ARRAY 占位符,必须以字面量形式拼到 SQL 文本中。
func (r *dorisRepository) insertRows(ctx context.Context,
table string, rows []*DorisVectorEmbedding,
) error {
if len(rows) == 0 {
return nil
}
// 9 个普通占位符 + 1 个 embedding 字面量。
const perRowPlaceholders = "(?, ?, ?, ?, ?, ?, ?, ?, ?, %s)"
parts := make([]string, len(rows))
args := make([]any, 0, len(rows)*9)
for i, e := range rows {
parts[i] = fmt.Sprintf(perRowPlaceholders, embeddingLiteral(e.Embedding))
args = append(args,
e.ID, e.Content, e.SourceID, e.SourceType,
e.ChunkID, e.KnowledgeID, e.KnowledgeBaseID, e.TagID,
e.IsEnabled,
)
}
stmt := fmt.Sprintf("INSERT INTO `%s` (%s) VALUES %s",
table,
strings.Join(columns, ", "),
strings.Join(parts, ", "),
)
_, err := r.db.ExecContext(ctx, stmt, args...)
return err
}
// replaceRows 在 DUPLICATE KEY 表上显式模拟“按 id 覆盖”的语义。
func (r *dorisRepository) replaceRows(ctx context.Context,
table string, rows []*DorisVectorEmbedding,
) error {
rows = dedupeRowsByID(rows)
if len(rows) == 0 {
return nil
}
ids := make([]string, 0, len(rows))
for _, row := range rows {
ids = append(ids, row.ID)
}
if err := r.deleteRowsByID(ctx, table, ids); err != nil {
return err
}
return r.insertRows(ctx, table, rows)
}
func (r *dorisRepository) deleteRowsByID(ctx context.Context,
table string, ids []string,
) error {
if len(ids) == 0 {
return nil
}
placeholders := make([]string, len(ids))
args := make([]any, len(ids))
for i, id := range ids {
placeholders[i] = "?"
args[i] = id
}
stmt := fmt.Sprintf("DELETE FROM `%s` WHERE %s IN (%s)",
table, fieldID, strings.Join(placeholders, ", "))
_, err := r.db.ExecContext(ctx, stmt, args...)
return err
}
func dedupeRowsByID(rows []*DorisVectorEmbedding) []*DorisVectorEmbedding {
if len(rows) < 2 {
return rows
}
out := make([]*DorisVectorEmbedding, 0, len(rows))
positions := make(map[string]int, len(rows))
for _, row := range rows {
if idx, ok := positions[row.ID]; ok {
out[idx] = row
continue
}
positions[row.ID] = len(out)
out = append(out, row)
}
return out
}
// DeleteByChunkIDList 用 chunk_id 列删除。dimension 用于定位具体表。
func (r *dorisRepository) DeleteByChunkIDList(ctx context.Context,
chunkIDList []string, dimension int, _ string,
) error {
return r.deleteByField(ctx, fieldChunkID, chunkIDList, dimension)
}
// DeleteByKnowledgeIDList 用 knowledge_id 列删除。
func (r *dorisRepository) DeleteByKnowledgeIDList(ctx context.Context,
knowledgeIDList []string, dimension int, _ string,
) error {
return r.deleteByField(ctx, fieldKnowledgeID, knowledgeIDList, dimension)
}
// DeleteBySourceIDList 用 source_id 列删除。
func (r *dorisRepository) DeleteBySourceIDList(ctx context.Context,
sourceIDList []string, dimension int, _ string,
) error {
return r.deleteByField(ctx, fieldSourceID, sourceIDList, dimension)
}
// deleteByField 是三个 Delete* 方法的统一实现:
// DELETE FROM <table> WHERE <field> IN (?, ?, ...)。
func (r *dorisRepository) deleteByField(ctx context.Context,
field string, ids []string, dimension int,
) error {
log := logger.GetLogger(ctx)
if len(ids) == 0 {
return nil
}
table := r.getTableName(dimension)
placeholders := make([]string, len(ids))
args := make([]any, len(ids))
for i, v := range ids {
placeholders[i] = "?"
args[i] = v
}
stmt := fmt.Sprintf("DELETE FROM `%s` WHERE %s IN (%s)",
table, field, strings.Join(placeholders, ", "))
if _, err := r.db.ExecContext(ctx, stmt, args...); err != nil {
log.Errorf("[Doris] Delete by %s failed: %v", field, err)
return fmt.Errorf("delete by %s: %w", field, err)
}
log.Infof("[Doris] Deleted %d rows from %s by %s", len(ids), table, field)
return nil
}
// Retrieve 根据 RetrieverType 分发到向量检索或关键词检索。
func (r *dorisRepository) Retrieve(ctx context.Context,
params types.RetrieveParams,
) ([]*types.RetrieveResult, error) {
switch params.RetrieverType {
case types.VectorRetrieverType:
return r.VectorRetrieve(ctx, params)
case types.KeywordsRetrieverType:
return r.KeywordsRetrieve(ctx, params)
}
return nil, fmt.Errorf("invalid retriever type: %v", params.RetrieverType)
}
// VectorRetrieve 对查询向量先做单位化,再调用 inner_product_approximate 做
// ANN 搜索;对单位向量而言,inner product 与 cosine similarity 等价,
// 因此 score 仍然保持“越大越相似”的语义。
func (r *dorisRepository) VectorRetrieve(ctx context.Context,
params types.RetrieveParams,
) ([]*types.RetrieveResult, error) {
log := logger.GetLogger(ctx)
if err := validateEmbedding(params.Embedding); err != nil {
return nil, fmt.Errorf("invalid query embedding: %w", err)
}
compatMode, err := r.resolveCompatMode(ctx)
if err != nil {
return nil, err
}
queryEmbedding := append([]float32(nil), params.Embedding...)
if compatMode.normalizeEmbeddings() {
queryEmbedding = normalizeEmbedding(queryEmbedding)
}
dim := len(params.Embedding)
table := r.getTableName(dim)
exists, err := r.tableExists(ctx, table)
if err != nil {
return nil, fmt.Errorf("check table %s: %w", table, err)
}
if !exists {
log.Warnf("[Doris] Table %s does not exist, returning empty results", table)
return buildRetrieveResult(nil, types.VectorRetrieverType), nil
}
wb := buildBaseFilter(params)
whereClause, whereArgs := wb.build()
// embedding 必须用字面量,Doris 不支持 LIMIT/OFFSET 使用占位符,必须内联为字面量。
// 使用 HAVING 是因为 score 是 SELECT 列别名,WHERE 阶段还看不到。
scoreExpr := fmt.Sprintf("inner_product_approximate(`%s`, %s)", fieldEmbedding, embeddingLiteral(queryEmbedding))
if compatMode == dorisCompatModeLegacy {
scoreExpr = fmt.Sprintf("(1 - cosine_distance_approximate(`%s`, %s))", fieldEmbedding, embeddingLiteral(queryEmbedding))
}
stmt := fmt.Sprintf(
"SELECT %s, %s AS score "+
"FROM `%s` WHERE %s "+
"HAVING score >= ? "+
"ORDER BY score DESC LIMIT %d",
strings.Join(columnsForRetrieve, ", "),
scoreExpr,
table,
whereClause,
params.TopK,
)
args := append(whereArgs, params.Threshold)
rows, err := r.db.QueryContext(ctx, stmt, args...)
if err != nil {
return nil, r.wrapVectorRetrieveError(table, compatMode, err)
}
defer rows.Close()
results, err := scanRetrieveRows(rows, types.MatchTypeEmbedding)
if err != nil {
return nil, err
}
log.Infof("[Doris] Vector retrieval found %d results in %s", len(results), table)
return buildRetrieveResult(results, types.VectorRetrieverType), nil
}
// KeywordsRetrieve 用 Doris 倒排索引 + MATCH_ANY 做关键词匹配。
//
// 不需要 jieba 客户端分词:CREATE TABLE 时 idx_content 已经声明了 chinese parser。
// 不同维度的表跨表合并取 topK,与 Milvus/Weaviate 现状一致。
func (r *dorisRepository) KeywordsRetrieve(ctx context.Context,
params types.RetrieveParams,
) ([]*types.RetrieveResult, error) {
log := logger.GetLogger(ctx)
query := strings.TrimSpace(params.Query)
if query == "" {
return buildRetrieveResult(nil, types.KeywordsRetrieverType), nil
}
tables, err := r.listEmbeddingTables(ctx)
if err != nil {
return nil, fmt.Errorf("list tables: %w", err)
}
if len(tables) == 0 {
return buildRetrieveResult(nil, types.KeywordsRetrieverType), nil
}
wb := buildBaseFilter(params)
whereClause, whereArgs := wb.build()
var all []*types.IndexWithScore
for _, table := range tables {
stmt := fmt.Sprintf(
"SELECT %s FROM `%s` WHERE %s AND %s MATCH_ANY ? LIMIT %d",
strings.Join(columnsForRetrieve, ", "),
table, whereClause, fieldContent,
params.TopK,
)
args := append(append([]any{}, whereArgs...), query)
rows, err := r.db.QueryContext(ctx, stmt, args...)
if err != nil {
log.Warnf("[Doris] Keyword retrieve in %s failed: %v", table, err)
continue
}
// score 在 KeywordsRetrieve 中固定 1.0,与 Qdrant 行为一致。
batch, scanErr := scanRetrieveRows(rows, types.MatchTypeKeywords)
_ = rows.Close()
if scanErr != nil {
return nil, scanErr
}
all = append(all, batch...)
}
if len(all) > params.TopK {
all = all[:params.TopK]
}
log.Infof("[Doris] Keywords retrieval found %d results across %d tables", len(all), len(tables))
return buildRetrieveResult(all, types.KeywordsRetrieverType), nil
}
// CopyIndices 把源知识库的 chunk 复制到目标知识库,避免重新生成 embedding。
//
// 与 Qdrant 的实现完全镜像:
// - 分页扫描源表
// - 按 sourceToTargetChunkIDMap 把 chunk_id 翻译过去
// - 处理 source_id 翻译规则(普通 chunk / 生成型问题 / 其他)
// - 把目标行写回同一个表
func (r *dorisRepository) CopyIndices(ctx context.Context,
sourceKnowledgeBaseID string,
sourceToTargetKBIDMap map[string]string,
sourceToTargetChunkIDMap map[string]string,
targetKnowledgeBaseID string,
dimension int,
_ string,
) error {
log := logger.GetLogger(ctx)
if len(sourceToTargetChunkIDMap) == 0 {
return nil
}
if err := r.ensureTable(ctx, dimension); err != nil {
return err
}
table := r.getTableName(dimension)
const pageSize = 64
offset := 0
totalCopied := 0
for {
stmt := fmt.Sprintf(
"SELECT %s FROM `%s` WHERE %s = ? ORDER BY %s LIMIT %d OFFSET %d",
strings.Join(columnsForCopy, ", "),
table, fieldKnowledgeBaseID, fieldID,
pageSize, offset,
)
rows, err := r.db.QueryContext(ctx, stmt, sourceKnowledgeBaseID)
if err != nil {
return fmt.Errorf("copy indices scan: %w", err)
}
batch, err := scanCopyRows(rows)
_ = rows.Close()
if err != nil {
return err
}
if len(batch) == 0 {
break
}
var targets []*DorisVectorEmbedding
for _, src := range batch {
targetChunkID, ok := sourceToTargetChunkIDMap[src.ChunkID]
if !ok {
log.Warnf("[Doris] Source chunk %s not in target mapping", src.ChunkID)
continue
}
targetKnowledgeID, ok := sourceToTargetKBIDMap[src.KnowledgeID]
if !ok {
log.Warnf("[Doris] Source knowledge %s not in target mapping", src.KnowledgeID)
continue
}
targetSourceID := translateSourceID(src.SourceID, src.ChunkID, targetChunkID)
targets = append(targets, &DorisVectorEmbedding{
ID: uuid.New().String(),
Content: src.Content,
SourceID: targetSourceID,
SourceType: src.SourceType,
ChunkID: targetChunkID,
KnowledgeID: targetKnowledgeID,
KnowledgeBaseID: targetKnowledgeBaseID,
TagID: src.TagID,
IsEnabled: src.IsEnabled,
Embedding: src.Embedding,
})
}
if len(targets) > 0 {
if err := r.insertRows(ctx, table, targets); err != nil {
return fmt.Errorf("copy indices insert: %w", err)
}
totalCopied += len(targets)
}
if len(batch) < pageSize {
break
}
offset += pageSize
}
log.Infof("[Doris] CopyIndices done, dim=%d, copied=%d", dimension, totalCopied)
return nil
}
// BatchUpdateChunkEnabledStatus / BatchUpdateChunkTagID 实际实现位于 streamload.go,
// 会按 compat mode 选择 partial update 或 rewrite rows。
// ---------------------------------------------------------------------------
// 私有辅助
// ---------------------------------------------------------------------------
// toDorisVectorEmbedding 把 IndexInfo + 上层传入的 embedding 映射 转换为
// Doris 行模型。Embedding 通过 additionalParams[fieldEmbedding] 中的
// map[string][]float32 按 SourceID 取出,与 Qdrant/Milvus 完全一致。
// inner_product_duplicate 模式会先单位化,legacy 模式保留原始向量。
func toDorisVectorEmbedding(
info *types.IndexInfo,
additionalParams map[string]any,
compatMode dorisCompatMode,
) *DorisVectorEmbedding {
emb := &DorisVectorEmbedding{
ID: info.ID,
Content: info.Content,
SourceID: info.SourceID,
SourceType: int(info.SourceType),
ChunkID: info.ChunkID,
KnowledgeID: info.KnowledgeID,
KnowledgeBaseID: info.KnowledgeBaseID,
TagID: info.TagID,
IsEnabled: info.IsEnabled,
}
if additionalParams != nil {
if v, ok := additionalParams[fieldEmbedding]; ok {
if m, ok := v.(map[string][]float32); ok {
emb.Embedding = append([]float32(nil), m[info.SourceID]...)
if compatMode.normalizeEmbeddings() {
emb.Embedding = normalizeEmbedding(emb.Embedding)
}
}
}
}
return emb
}
func (r *dorisRepository) wrapVectorRetrieveError(table string, compatMode dorisCompatMode, err error) error {
if compatMode == dorisCompatModeLegacy {
return fmt.Errorf(
"vector retrieve %s in Doris compat mode %s: %w. If your Doris build does not support cosine_distance_approximate or ANN on UNIQUE KEY tables, set %s=%s before creating embedding tables. %s is not interchangeable after %s_* tables are created",
table,
compatMode,
err,
envDorisCompatMode,
dorisCompatModeInnerProductDuplicate,
envDorisCompatMode,
r.tableBaseName,
)
}
return fmt.Errorf("vector retrieve %s in Doris compat mode %s: %w", table, compatMode, err)
}
// translateSourceID 把源 SourceID 翻译到目标 SourceID,与 Qdrant 实现完全镜像:
// - 普通 chunk:SourceID == ChunkID -> 使用 targetChunkID
// - 生成型问题:SourceID == "<chunkID>-<questionID>" -> "<targetChunkID>-<questionID>"
// - 其他场景:生成新的 UUID(保持唯一性)
func translateSourceID(originalSourceID, sourceChunkID, targetChunkID string) string {
switch {
case originalSourceID == sourceChunkID:
return targetChunkID
case strings.HasPrefix(originalSourceID, sourceChunkID+"-"):
questionID := strings.TrimPrefix(originalSourceID, sourceChunkID+"-")
return fmt.Sprintf("%s-%s", targetChunkID, questionID)
default:
return uuid.New().String()
}
}
// scanRetrieveRows 把 Retrieve 阶段的 rows 反序列化为 IndexWithScore 列表。
//
// 分两路:
// - 列数 == columnsForRetrieve+1:第 N+1 列是 score(向量检索路径)
// - 列数 == columnsForRetrieve:score 统一赋 1.0(关键词检索路径)
func scanRetrieveRows(rows *sql.Rows, matchType types.MatchType) ([]*types.IndexWithScore, error) {
cols, err := rows.Columns()
if err != nil {
return nil, err
}
withScore := len(cols) == len(columnsForRetrieve)+1
var out []*types.IndexWithScore
for rows.Next() {
var (
id, content, sourceID, chunkID string
knowledgeID, knowledgeBaseID, tagID string
sourceType int
isEnabled bool
score float64
err error
)
if withScore {
err = rows.Scan(&id, &content, &sourceID, &sourceType,
&chunkID, &knowledgeID, &knowledgeBaseID, &tagID, &isEnabled, &score)
} else {
err = rows.Scan(&id, &content, &sourceID, &sourceType,
&chunkID, &knowledgeID, &knowledgeBaseID, &tagID, &isEnabled)
score = 1.0
}
if err != nil {
return nil, fmt.Errorf("scan row: %w", err)
}
out = append(out, &types.IndexWithScore{
ID: id,
Content: content,
SourceID: sourceID,
SourceType: types.SourceType(sourceType),
ChunkID: chunkID,
KnowledgeID: knowledgeID,
KnowledgeBaseID: knowledgeBaseID,
TagID: tagID,
Score: score,
MatchType: matchType,
})
}
return out, rows.Err()
}
// scanCopyRows 反序列化 CopyIndices 的分页查询结果。
//
// 与 scanRetrieveRows 不同,这里需要 embedding 字段(去复制原始向量)。
// Doris 的 ARRAY<FLOAT> 通过 mysql 协议返回的是字符串字面量 "[1,2,3]"。
func scanCopyRows(rows *sql.Rows) ([]*DorisVectorEmbedding, error) {
var out []*DorisVectorEmbedding
for rows.Next() {
var (
id, content, sourceID, chunkID string
knowledgeID, knowledgeBaseID, tagID string
sourceType int
isEnabled bool
embeddingRaw sql.RawBytes
)
if err := rows.Scan(&id, &content, &sourceID, &sourceType,
&chunkID, &knowledgeID, &knowledgeBaseID, &tagID, &isEnabled, &embeddingRaw); err != nil {
return nil, fmt.Errorf("scan copy row: %w", err)
}
vec, err := parseEmbeddingLiteral(embeddingRaw)
if err != nil {
return nil, fmt.Errorf("parse embedding: %w", err)
}
out = append(out, &DorisVectorEmbedding{
ID: id,
Content: content,
SourceID: sourceID,
SourceType: sourceType,
ChunkID: chunkID,
KnowledgeID: knowledgeID,
KnowledgeBaseID: knowledgeBaseID,
TagID: tagID,
IsEnabled: isEnabled,
Embedding: vec,
})
}
return out, rows.Err()
}
// buildRetrieveResult 把 IndexWithScore 列表包装成 RetrieveResult。
func buildRetrieveResult(results []*types.IndexWithScore, retrieverType types.RetrieverType) []*types.RetrieveResult {
return []*types.RetrieveResult{{
Results: results,
RetrieverEngineType: types.DorisRetrieverEngineType,
RetrieverType: retrieverType,
Error: nil,
}}
}
// calculateStorageSize 估算单行的存储成本。
//
// 与 Qdrant 一致:payload 字符串字节 + 向量 (dim*4) + HNSW M*2*8 + 元数据 24。
func calculateStorageSize(emb *DorisVectorEmbedding) int64 {
var payload int64
payload += int64(len(emb.Content))
payload += int64(len(emb.SourceID))
payload += int64(len(emb.ChunkID))
payload += int64(len(emb.KnowledgeID))
payload += int64(len(emb.KnowledgeBaseID))
payload += int64(len(emb.TagID))
payload += 8 // source_type int
var vec int64
var hnsw int64
if len(emb.Embedding) > 0 {
vec = int64(len(emb.Embedding)) * 4
const hnswM = 32
hnsw = hnswM * 2 * 8
}
const metaBytes int64 = 24
return payload + vec + hnsw + metaBytes
}
@@ -0,0 +1,762 @@
package doris
import (
"context"
"database/sql"
"encoding/json"
"errors"
"io"
"net/http"
"net/http/httptest"
"regexp"
"strings"
"testing"
"time"
"github.com/DATA-DOG/go-sqlmock"
"github.com/Tencent/WeKnora/internal/types"
secutils "github.com/Tencent/WeKnora/internal/utils"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// newTestRepo 构造一个共享 sqlmock 的 dorisRepository,
// 默认绕过 ensureTable(initializedTables 已置位),便于专注测试 SQL 形态。
//
// 返回的 cleanup 用 defer 调用即可。
func newTestRepo(t *testing.T) (*dorisRepository, sqlmock.Sqlmock, *httptest.Server, func()) {
t.Helper()
// httptest binds to loopback, which production SSRF policy correctly blocks.
// Make that one test host explicit and restore the process-global policy when
// the test completes.
secutils.SetSSRFWhitelistFromRaw("127.0.0.1")
t.Cleanup(secutils.ResetSSRFWhitelistForTest)
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
require.NoError(t, err)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
// 默认行为:成功,回 1 行。
body, _ := io.ReadAll(r.Body)
_ = body
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(map[string]any{
"Status": "Success",
"NumberTotalRows": 1,
"NumberLoadedRows": 1,
"Label": "test",
})
}))
repo := &dorisRepository{
db: db,
httpClient: srv.Client(),
feHTTPBase: srv.URL,
username: "u",
password: "p",
database: "weknora",
tableBaseName: "weknora_embeddings",
}
cleanup := func() {
_ = db.Close()
srv.Close()
}
return repo, mock, srv, cleanup
}
func primeCompatMode(repo *dorisRepository, mode dorisCompatMode, err error) {
repo.compatModeResolved = mode
repo.compatResolveErr = err
repo.compatResolveOnce.Do(func() {})
}
// ---------------------------------------------------------------------------
// query.go:whereBuilder / embeddingLiteral / parseEmbeddingLiteral
// ---------------------------------------------------------------------------
func TestEmbeddingLiteralRoundTrip(t *testing.T) {
t.Run("empty vector returns []", func(t *testing.T) {
assert.Equal(t, "[]", embeddingLiteral(nil))
assert.Equal(t, "[]", embeddingLiteral([]float32{}))
})
t.Run("does not contain locale-sensitive separators", func(t *testing.T) {
s := embeddingLiteral([]float32{1.5, -2.25, 0.001})
// 必须只包含数字 / 点 / 负号 / e / 逗号 / 方括号
assert.Regexp(t, regexp.MustCompile(`^\[[\-+0-9eE.,]+\]$`), s)
})
t.Run("round trip", func(t *testing.T) {
orig := []float32{1.5, -2.25, 0.0625}
s := embeddingLiteral(orig)
got, err := parseEmbeddingLiteral([]byte(s))
require.NoError(t, err)
assert.Equal(t, orig, got)
})
t.Run("parse handles whitespace and missing brackets", func(t *testing.T) {
v, err := parseEmbeddingLiteral([]byte(" 1.0 , 2.0 ,3.0 "))
require.NoError(t, err)
assert.Equal(t, []float32{1.0, 2.0, 3.0}, v)
})
}
func TestWhereBuilder(t *testing.T) {
t.Run("empty builder returns 1 = 1", func(t *testing.T) {
w := &whereBuilder{}
clause, args := w.build()
assert.Equal(t, "1 = 1", clause)
assert.Nil(t, args)
})
t.Run("equal + IN + NOT IN", func(t *testing.T) {
w := &whereBuilder{}
w.addEqual("is_enabled", true)
w.addIn("knowledge_base_id", []string{"kb1", "kb2"})
w.addNotIn("chunk_id", []string{"x"})
clause, args := w.build()
assert.Contains(t, clause, "is_enabled = ?")
assert.Contains(t, clause, "knowledge_base_id IN (?, ?)")
assert.Contains(t, clause, "chunk_id NOT IN (?)")
assert.Equal(t, []any{true, "kb1", "kb2", "x"}, args)
})
t.Run("buildBaseFilter applies all params", func(t *testing.T) {
w := buildBaseFilter(types.RetrieveParams{
KnowledgeBaseIDs: []string{"kb1"},
KnowledgeIDs: []string{"k1", "k2"},
TagIDs: []string{"t1"},
ExcludeKnowledgeIDs: []string{"k9"},
ExcludeChunkIDs: []string{"c9"},
})
clause, _ := w.build()
assert.Contains(t, clause, "is_enabled = ?")
assert.Contains(t, clause, "knowledge_base_id IN (?)")
assert.Contains(t, clause, "knowledge_id IN (?, ?)")
assert.Contains(t, clause, "tag_id IN (?)")
assert.Contains(t, clause, "knowledge_id NOT IN (?)")
assert.Contains(t, clause, "chunk_id NOT IN (?)")
})
}
// ---------------------------------------------------------------------------
// streamload.go:chunkRows / partialUpdateRows
// ---------------------------------------------------------------------------
func TestChunkRows(t *testing.T) {
rows := []map[string]any{
{"id": "a", "is_enabled": true},
{"id": "b", "is_enabled": false},
{"id": "c", "is_enabled": true},
}
t.Run("single batch when fits", func(t *testing.T) {
batches := chunkRows(rows, 4096)
require.Len(t, batches, 1)
assert.Len(t, batches[0], 3)
})
t.Run("splits when exceeding maxBytes", func(t *testing.T) {
// 给一个非常小的上限,强制每行单独成段
batches := chunkRows(rows, 16)
assert.GreaterOrEqual(t, len(batches), 2)
var total int
for _, b := range batches {
total += len(b)
}
assert.Equal(t, len(rows), total)
})
}
func TestPartialUpdateRows_HappyPath(t *testing.T) {
repo, _, _, cleanup := newTestRepo(t)
defer cleanup()
// 自定义 server 验证请求形态。
var captured *http.Request
var capturedBody []byte
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
captured = r
capturedBody, _ = io.ReadAll(r.Body)
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(map[string]any{
"Status": "Success",
"NumberTotalRows": 2,
"NumberLoadedRows": 2,
"Label": "ok",
})
}))
defer srv.Close()
repo.feHTTPBase = srv.URL
rows := []map[string]any{
{"id": "id1", "is_enabled": true},
{"id": "id2", "is_enabled": false},
}
require.NoError(t, repo.partialUpdateRows(context.Background(),
"weknora_embeddings_768", []string{"id", "is_enabled"}, rows))
require.NotNil(t, captured)
assert.Equal(t, http.MethodPut, captured.Method)
assert.Equal(t, "/api/weknora/weknora_embeddings_768/_stream_load", captured.URL.Path)
assert.Equal(t, "true", captured.Header.Get("partial_columns"))
assert.Equal(t, "true", captured.Header.Get("strip_outer_array"))
assert.Equal(t, "json", captured.Header.Get("format"))
assert.Equal(t, "id,is_enabled", captured.Header.Get("columns"))
assert.Equal(t, "APPEND", captured.Header.Get("merge_type"))
assert.True(t, strings.HasPrefix(captured.Header.Get("Authorization"), "Basic "))
var got []map[string]any
require.NoError(t, json.Unmarshal(capturedBody, &got))
assert.Len(t, got, 2)
assert.Equal(t, "id1", got[0]["id"])
}
func TestPartialUpdateRows_FailureSurfaced(t *testing.T) {
repo, _, _, cleanup := newTestRepo(t)
defer cleanup()
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(map[string]any{
"Status": "Fail",
"Message": "label exists",
"ErrorURL": "http://...",
})
}))
defer srv.Close()
repo.feHTTPBase = srv.URL
err := repo.partialUpdateRows(context.Background(),
"t", []string{"id", "is_enabled"},
[]map[string]any{{"id": "x", "is_enabled": true}},
)
require.Error(t, err)
assert.Contains(t, err.Error(), "stream load failed")
}
func TestStreamLoadOnce_BlocksUnsafeTarget(t *testing.T) {
repo, _, _, cleanup := newTestRepo(t)
defer cleanup()
// Override the helper's loopback allowance: link-local metadata endpoints
// must be rejected before the HTTP client is called.
secutils.SetSSRFWhitelistFromRaw("")
repo.feHTTPBase = "http://169.254.169.254"
err := repo.streamLoadOnce(context.Background(), "t", []string{"id"},
[]map[string]any{{"id": "x"}})
require.Error(t, err)
assert.Contains(t, err.Error(), "blocked by SSRF validation")
}
func TestDorisStreamLoadHTTPClient_BlocksUnsafeRedirect(t *testing.T) {
secutils.SetSSRFWhitelistFromRaw("127.0.0.1")
t.Cleanup(secutils.ResetSSRFWhitelistForTest)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
http.Redirect(w, r, "http://169.254.169.254/latest/meta-data", http.StatusTemporaryRedirect)
}))
defer server.Close()
req, err := http.NewRequest(http.MethodPut, server.URL, strings.NewReader("[]"))
require.NoError(t, err)
_, err = newDorisStreamLoadHTTPClient().Do(req)
require.Error(t, err)
assert.ErrorIs(t, err, secutils.ErrSSRFRedirectBlocked)
}
func TestDorisStreamLoadHTTPClient_ForwardsAuthorizationToTrustedRedirect(t *testing.T) {
secutils.SetSSRFWhitelistFromRaw("127.0.0.1")
t.Cleanup(secutils.ResetSSRFWhitelistForTest)
var gotAuthorization string
be := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
gotAuthorization = r.Header.Get(headerAuthorization)
w.WriteHeader(http.StatusOK)
}))
defer be.Close()
fe := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
http.Redirect(w, r, be.URL, http.StatusTemporaryRedirect)
}))
defer fe.Close()
req, err := http.NewRequest(http.MethodPut, fe.URL, strings.NewReader("[]"))
require.NoError(t, err)
req.Header.Set(headerAuthorization, "Basic trusted-doris-credential")
resp, err := newDorisStreamLoadHTTPClient().Do(req)
require.NoError(t, err)
defer resp.Body.Close()
assert.Equal(t, "Basic trusted-doris-credential", gotAuthorization)
}
// ---------------------------------------------------------------------------
// repository.go:SQL 形态
// ---------------------------------------------------------------------------
func TestDeleteByChunkIDList_SQLShape(t *testing.T) {
repo, mock, _, cleanup := newTestRepo(t)
defer cleanup()
mock.ExpectExec(`DELETE FROM .*weknora_embeddings_768.* WHERE chunk_id IN \(\?, \?\)`).
WithArgs("c1", "c2").
WillReturnResult(sqlmock.NewResult(0, 2))
err := repo.DeleteByChunkIDList(context.Background(), []string{"c1", "c2"}, 768, "")
require.NoError(t, err)
assert.NoError(t, mock.ExpectationsWereMet())
}
func TestDeleteByKnowledgeIDList_NoOpOnEmpty(t *testing.T) {
repo, mock, _, cleanup := newTestRepo(t)
defer cleanup()
require.NoError(t, repo.DeleteByKnowledgeIDList(context.Background(), nil, 768, ""))
// 空列表不应触发任何 query
assert.NoError(t, mock.ExpectationsWereMet())
}
func TestVectorRetrieve_SQLShape(t *testing.T) {
repo, mock, _, cleanup := newTestRepo(t)
defer cleanup()
mock.ExpectQuery(`SELECT COUNT\(1\) FROM information_schema.tables`).
WithArgs("weknora", "weknora_embeddings_3").
WillReturnRows(sqlmock.NewRows([]string{"c"}).AddRow(1))
// Doris does not support placeholders for LIMIT, so TopK is inlined as a
// literal and is not a bound argument.
mock.ExpectQuery(`SELECT id, content, .*inner_product_approximate.*HAVING score >= \? ORDER BY score DESC LIMIT \d+`).
WithArgs(true, 0.5).
WillReturnRows(
sqlmock.NewRows([]string{
"id", "content", "source_id", "source_type",
"chunk_id", "knowledge_id", "knowledge_base_id", "tag_id",
"is_enabled", "score",
}).AddRow("id1", "hello", "src", 0, "c1", "k1", "kb1", "t1", true, 0.95),
)
results, err := repo.VectorRetrieve(context.Background(), types.RetrieveParams{
Embedding: []float32{1, 2, 3},
TopK: 5,
Threshold: 0.5,
RetrieverType: types.VectorRetrieverType,
})
require.NoError(t, err)
require.Len(t, results, 1)
require.Len(t, results[0].Results, 1)
assert.Equal(t, "id1", results[0].Results[0].ID)
assert.InDelta(t, 0.95, results[0].Results[0].Score, 1e-9)
assert.NoError(t, mock.ExpectationsWereMet())
}
func TestVectorRetrieve_SQLShape_LegacyMode(t *testing.T) {
repo, mock, _, cleanup := newTestRepo(t)
defer cleanup()
primeCompatMode(repo, dorisCompatModeLegacy, nil)
mock.ExpectQuery(`SELECT COUNT\(1\) FROM information_schema.tables`).
WithArgs("weknora", "weknora_embeddings_3").
WillReturnRows(sqlmock.NewRows([]string{"c"}).AddRow(1))
mock.ExpectQuery(`SELECT id, content, .*cosine_distance_approximate.*HAVING score >= \? ORDER BY score DESC LIMIT \d+`).
WithArgs(true, 0.5).
WillReturnRows(
sqlmock.NewRows([]string{
"id", "content", "source_id", "source_type",
"chunk_id", "knowledge_id", "knowledge_base_id", "tag_id",
"is_enabled", "score",
}).AddRow("id1", "hello", "src", 0, "c1", "k1", "kb1", "t1", true, 0.8),
)
results, err := repo.VectorRetrieve(context.Background(), types.RetrieveParams{
Embedding: []float32{1, 2, 3},
TopK: 5,
Threshold: 0.5,
RetrieverType: types.VectorRetrieverType,
})
require.NoError(t, err)
require.Len(t, results, 1)
require.Len(t, results[0].Results, 1)
assert.Equal(t, "id1", results[0].Results[0].ID)
assert.NoError(t, mock.ExpectationsWereMet())
}
func TestNormalizeEmbedding(t *testing.T) {
t.Run("returns unit vector copy", func(t *testing.T) {
got := normalizeEmbedding([]float32{3, 4})
require.Len(t, got, 2)
assert.InDelta(t, 0.6, got[0], 1e-6)
assert.InDelta(t, 0.8, got[1], 1e-6)
})
t.Run("keeps zero vector finite", func(t *testing.T) {
got := normalizeEmbedding([]float32{0, 0})
require.Len(t, got, 2)
assert.Equal(t, []float32{0, 0}, got)
})
}
func TestKeywordsRetrieve_SQLShape(t *testing.T) {
repo, mock, _, cleanup := newTestRepo(t)
defer cleanup()
mock.ExpectQuery(`SELECT TABLE_NAME FROM information_schema.tables`).
WithArgs("weknora", "weknora_embeddings\\_%").
WillReturnRows(sqlmock.NewRows([]string{"TABLE_NAME"}).
AddRow("weknora_embeddings_768"))
// LIMIT is inlined (see VectorRetrieve), so only is_enabled and the
// MATCH_ANY query string are bound arguments.
mock.ExpectQuery(`MATCH_ANY \?`).
WithArgs(true, "你好").
WillReturnRows(
sqlmock.NewRows([]string{
"id", "content", "source_id", "source_type",
"chunk_id", "knowledge_id", "knowledge_base_id", "tag_id",
"is_enabled",
}).AddRow("id1", "你好世界", "src", 0, "c1", "k1", "kb1", "", true),
)
results, err := repo.KeywordsRetrieve(context.Background(), types.RetrieveParams{
Query: "你好",
TopK: 3,
RetrieverType: types.KeywordsRetrieverType,
})
require.NoError(t, err)
require.Len(t, results, 1)
require.Len(t, results[0].Results, 1)
assert.Equal(t, "id1", results[0].Results[0].ID)
assert.Equal(t, 1.0, results[0].Results[0].Score)
assert.NoError(t, mock.ExpectationsWereMet())
}
func TestBatchUpdateChunkEnabledStatus_RewritesRows(t *testing.T) {
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
require.NoError(t, err)
defer db.Close()
repo := &dorisRepository{
db: db,
database: "weknora",
tableBaseName: "weknora_embeddings",
}
primeCompatMode(repo, dorisCompatModeInnerProductDuplicate, nil)
mock.ExpectQuery(`SELECT TABLE_NAME FROM information_schema.tables`).
WithArgs("weknora", "weknora_embeddings\\_%").
WillReturnRows(sqlmock.NewRows([]string{"TABLE_NAME"}).AddRow("weknora_embeddings_768"))
mock.ExpectQuery(`SELECT id, content, source_id, source_type, chunk_id, knowledge_id, knowledge_base_id, tag_id, is_enabled, embedding FROM .*weknora_embeddings_768.* WHERE chunk_id IN`).
WithArgs("c1").
WillReturnRows(sqlmock.NewRows([]string{
"id", "content", "source_id", "source_type",
"chunk_id", "knowledge_id", "knowledge_base_id", "tag_id",
"is_enabled", "embedding",
}).AddRow("row-1", "hello", "src1", 0, "c1", "k1", "kb1", "t1", true, "[1,2,3]"))
mock.ExpectExec(`DELETE FROM .*weknora_embeddings_768.* WHERE id IN \(\?\)`).
WithArgs("row-1").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectExec(`INSERT INTO .*weknora_embeddings_768.*VALUES \(\?, \?, \?, \?, \?, \?, \?, \?, \?, \[`).
WithArgs("row-1", "hello", "src1", 0, "c1", "k1", "kb1", "t1", false).
WillReturnResult(sqlmock.NewResult(0, 1))
require.NoError(t, repo.BatchUpdateChunkEnabledStatus(
context.Background(), map[string]bool{"c1": false}))
assert.NoError(t, mock.ExpectationsWereMet())
}
func TestBatchUpdateChunkEnabledStatus_LegacyModeUsesPartialUpdate(t *testing.T) {
repo, mock, _, cleanup := newTestRepo(t)
defer cleanup()
primeCompatMode(repo, dorisCompatModeLegacy, nil)
mock.ExpectQuery(`SELECT TABLE_NAME FROM information_schema.tables`).
WithArgs("weknora", "weknora_embeddings\\_%").
WillReturnRows(sqlmock.NewRows([]string{"TABLE_NAME"}).AddRow("weknora_embeddings_768"))
mock.ExpectQuery(`SELECT id, chunk_id FROM .*weknora_embeddings_768.* WHERE chunk_id IN`).
WithArgs("c1").
WillReturnRows(sqlmock.NewRows([]string{"id", "chunk_id"}).AddRow("row-1", "c1"))
require.NoError(t, repo.BatchUpdateChunkEnabledStatus(
context.Background(), map[string]bool{"c1": false}))
assert.NoError(t, mock.ExpectationsWereMet())
}
func TestEnsureTable_DDLShape(t *testing.T) {
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
require.NoError(t, err)
defer db.Close()
repo := &dorisRepository{
db: db,
database: "weknora",
tableBaseName: "weknora_embeddings",
bucketsNum: 5,
replicationNum: 2,
}
primeCompatMode(repo, dorisCompatModeInnerProductDuplicate, nil)
// 表不存在
mock.ExpectQuery(`SELECT COUNT\(1\) FROM information_schema.tables`).
WithArgs("weknora", "weknora_embeddings_768").
WillReturnRows(sqlmock.NewRows([]string{"c"}).AddRow(0))
// CREATE TABLE 应包含关键属性和 Doris 支持的 inner_product ANN metric
mock.ExpectExec(`CREATE TABLE IF NOT EXISTS .*weknora_embeddings_768.*metric_type"="inner_product".*DUPLICATE KEY\(id\).*BUCKETS 5.*replication_num.*=.*2`).
WillReturnResult(sqlmock.NewResult(0, 0))
// SHOW INDEX 一次即返回 ANN 已 FINISHED
mock.ExpectQuery(`SHOW INDEX FROM .*weknora_embeddings_768.*`).
WillReturnRows(
sqlmock.NewRows([]string{"Table", "Key_name", "State"}).
AddRow("weknora_embeddings_768", "idx_emb", "FINISHED"),
)
require.NoError(t, repo.ensureTable(context.Background(), 768))
// waitANNReady 在后台 goroutine 里执行,轮询 ExpectationsWereMet 直到 SHOW INDEX 也被消费。
require.Eventually(t, func() bool {
return mock.ExpectationsWereMet() == nil
}, 2*time.Second, 10*time.Millisecond, "expectations should be met after async ANN poll")
}
func TestEnsureTable_DDLShape_LegacyMode(t *testing.T) {
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
require.NoError(t, err)
defer db.Close()
repo := &dorisRepository{
db: db,
database: "weknora",
tableBaseName: "weknora_embeddings",
bucketsNum: 5,
replicationNum: 2,
}
primeCompatMode(repo, dorisCompatModeLegacy, nil)
mock.ExpectQuery(`SELECT COUNT\(1\) FROM information_schema.tables`).
WithArgs("weknora", "weknora_embeddings_768").
WillReturnRows(sqlmock.NewRows([]string{"c"}).AddRow(0))
mock.ExpectExec(`CREATE TABLE IF NOT EXISTS .*weknora_embeddings_768.*metric_type"="cosine_distance".*UNIQUE KEY\(id\).*enable_unique_key_merge_on_write.*true`).
WillReturnResult(sqlmock.NewResult(0, 0))
mock.ExpectQuery(`SHOW INDEX FROM .*weknora_embeddings_768.*`).
WillReturnRows(
sqlmock.NewRows([]string{"Table", "Key_name", "State"}).
AddRow("weknora_embeddings_768", "idx_emb", "FINISHED"),
)
require.NoError(t, repo.ensureTable(context.Background(), 768))
require.Eventually(t, func() bool {
return mock.ExpectationsWereMet() == nil
}, 2*time.Second, 10*time.Millisecond, "expectations should be met after async ANN poll")
}
func TestBatchSave_SQLShape(t *testing.T) {
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
require.NoError(t, err)
defer db.Close()
repo := &dorisRepository{
db: db,
database: "weknora",
tableBaseName: "weknora_embeddings",
}
primeCompatMode(repo, dorisCompatModeInnerProductDuplicate, nil)
repo.initializedTables.Store(3, true) // 跳过 ensureTable
mock.ExpectExec(`DELETE FROM .*weknora_embeddings_3.* WHERE id IN \(\?\)`).
WithArgs("src1").
WillReturnResult(sqlmock.NewResult(0, 0))
mock.ExpectExec(`INSERT INTO .*weknora_embeddings_3.*VALUES \(\?, \?, \?, \?, \?, \?, \?, \?, \?, \[`).
WithArgs(
"src1", "hello", "src1", 0,
"c1", "k1", "kb1", "",
true,
).
WillReturnResult(sqlmock.NewResult(0, 1))
err = repo.BatchSave(context.Background(),
[]*types.IndexInfo{{
Content: "hello",
SourceID: "src1",
ChunkID: "c1",
KnowledgeID: "k1",
KnowledgeBaseID: "kb1",
IsEnabled: true,
}},
map[string]any{
fieldEmbedding: map[string][]float32{
"src1": {1, 2, 3},
},
},
)
require.NoError(t, err)
assert.NoError(t, mock.ExpectationsWereMet())
}
func TestBatchSave_SQLShape_LegacyMode(t *testing.T) {
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
require.NoError(t, err)
defer db.Close()
repo := &dorisRepository{
db: db,
database: "weknora",
tableBaseName: "weknora_embeddings",
}
primeCompatMode(repo, dorisCompatModeLegacy, nil)
repo.initializedTables.Store(3, true)
mock.ExpectExec(`INSERT INTO .*weknora_embeddings_3.*VALUES \(\?, \?, \?, \?, \?, \?, \?, \?, \?, \[`).
WithArgs(
"src1", "hello", "src1", 0,
"c1", "k1", "kb1", "",
true,
).
WillReturnResult(sqlmock.NewResult(0, 1))
err = repo.BatchSave(context.Background(),
[]*types.IndexInfo{{
Content: "hello",
SourceID: "src1",
ChunkID: "c1",
KnowledgeID: "k1",
KnowledgeBaseID: "kb1",
IsEnabled: true,
}},
map[string]any{
fieldEmbedding: map[string][]float32{
"src1": {1, 2, 3},
},
},
)
require.NoError(t, err)
assert.NoError(t, mock.ExpectationsWereMet())
}
func TestResolveCompatMode_RejectsExistingModeSwitch(t *testing.T) {
repo, mock, _, cleanup := newTestRepo(t)
defer cleanup()
repo.compatModeRequested = dorisCompatModeLegacy
mock.ExpectQuery(`SELECT TABLE_NAME FROM information_schema.tables`).
WithArgs("weknora", "weknora_embeddings\\_%").
WillReturnRows(sqlmock.NewRows([]string{"TABLE_NAME"}).AddRow("weknora_embeddings_768"))
mock.ExpectQuery(`SHOW CREATE TABLE .*weknora_embeddings_768.*`).
WillReturnRows(sqlmock.NewRows([]string{"Table", "Create Table"}).AddRow(
"weknora_embeddings_768",
"CREATE TABLE `weknora_embeddings_768` (id VARCHAR(64)) ENGINE=OLAP DUPLICATE KEY(id) DISTRIBUTED BY HASH(id) BUCKETS 10 PROPERTIES(\"replication_num\"=\"1\")",
))
_, err := repo.resolveCompatMode(context.Background())
require.Error(t, err)
assert.Contains(t, err.Error(), "not interchangeable after weknora_embeddings_*")
assert.Contains(t, err.Error(), string(dorisCompatModeInnerProductDuplicate))
assert.NoError(t, mock.ExpectationsWereMet())
}
func TestResolveCompatMode_AutoPrefersInnerProductDuplicate(t *testing.T) {
repo, mock, _, cleanup := newTestRepo(t)
defer cleanup()
repo.compatModeRequested = dorisCompatModeAuto
mock.ExpectQuery(`SELECT TABLE_NAME FROM information_schema.tables`).
WithArgs("weknora", "weknora_embeddings\\_%").
WillReturnRows(sqlmock.NewRows([]string{"TABLE_NAME"}))
mock.ExpectQuery(`SELECT inner_product_approximate\(\[1.0\],\[1.0\]\)`).
WillReturnRows(sqlmock.NewRows([]string{"v"}).AddRow(1.0))
mock.ExpectQuery(`SELECT cosine_distance_approximate\(\[1.0\],\[1.0\]\)`).
WillReturnError(errors.New("unsupported"))
mode, err := repo.resolveCompatMode(context.Background())
require.NoError(t, err)
assert.Equal(t, dorisCompatModeInnerProductDuplicate, mode)
assert.NoError(t, mock.ExpectationsWereMet())
}
// ---------------------------------------------------------------------------
// Engine wiring
// ---------------------------------------------------------------------------
func TestEngineTypeAndSupport(t *testing.T) {
repo, _, _, cleanup := newTestRepo(t)
defer cleanup()
assert.Equal(t, types.DorisRetrieverEngineType, repo.EngineType())
supports := repo.Support()
assert.Contains(t, supports, types.KeywordsRetrieverType)
assert.Contains(t, supports, types.VectorRetrieverType)
}
func TestRetrieve_DispatchesByType(t *testing.T) {
repo, mock, _, cleanup := newTestRepo(t)
defer cleanup()
// invalid retriever type -> error,不会触发任何 SQL
_, err := repo.Retrieve(context.Background(), types.RetrieveParams{
RetrieverType: "unknown",
})
require.Error(t, err)
assert.NoError(t, mock.ExpectationsWereMet())
}
// ---------------------------------------------------------------------------
// regression / utility
// ---------------------------------------------------------------------------
func TestTranslateSourceID(t *testing.T) {
t.Run("plain chunk uses target chunk", func(t *testing.T) {
got := translateSourceID("c1", "c1", "tc1")
assert.Equal(t, "tc1", got)
})
t.Run("generated question preserves question id", func(t *testing.T) {
got := translateSourceID("c1-q9", "c1", "tc1")
assert.Equal(t, "tc1-q9", got)
})
t.Run("unrecognized source falls back to fresh uuid", func(t *testing.T) {
got := translateSourceID("totally-other", "c1", "tc1")
assert.NotEqual(t, "totally-other", got)
assert.Len(t, got, 36) // UUID 长度
})
}
func TestEstimateStorageSize(t *testing.T) {
repo, _, _, cleanup := newTestRepo(t)
defer cleanup()
out := repo.EstimateStorageSize(context.Background(),
[]*types.IndexInfo{{Content: "hello", ChunkID: "c1", KnowledgeID: "k1", KnowledgeBaseID: "kb1"}},
map[string]any{
fieldEmbedding: map[string][]float32{
"": {1, 2, 3},
},
},
)
assert.Greater(t, out, int64(0))
}
// 保证 *sql.Rows 错误不会被吞掉。
func TestScanRetrieveRows_Error(t *testing.T) {
db, mock, err := sqlmock.New()
require.NoError(t, err)
defer db.Close()
mock.ExpectQuery("SELECT 1").
WillReturnError(errors.New("boom"))
rows, err := db.Query("SELECT 1")
if err != nil {
// query 直接失败也算预期
assert.Equal(t, "boom", err.Error())
return
}
_, scanErr := scanRetrieveRows(rows, types.MatchTypeEmbedding)
if scanErr != nil {
assert.Error(t, scanErr)
}
}
// silence unused import warning when sql isn't directly used
var _ = sql.ErrNoRows
@@ -0,0 +1,310 @@
package doris
import (
"context"
"fmt"
"strings"
"time"
"github.com/Tencent/WeKnora/internal/logger"
)
// 默认的桶数 / 副本数。Doris 在 PROPERTIES 不指定时会用集群默认值,
// 这里给一个对单机/小集群更友好的保守值。
const (
defaultBucketsNum = 10
defaultReplicationNum = 1
// ANN 索引就绪轮询的最大等待时间。索引未就绪不阻塞写入路径,
// 只阻塞 ensureTable 自身(首次建表场景),所以 30s 是可接受的。
annReadyTimeout = 30 * time.Second
annReadyPoll = 1 * time.Second
)
// getTableName 返回某个维度对应的物理表名:<base>_<dim>。
//
// 与 Qdrant/Milvus/Weaviate 的 collection 命名约定一致,
// 这样不同 embedding 模型(不同维度)的数据互不冲突。
func (r *dorisRepository) getTableName(dimension int) string {
return fmt.Sprintf("%s_%d", r.tableBaseName, dimension)
}
// ensureTable 保证目标维度对应的表已经存在;
// 不存在则用 CREATE TABLE IF NOT EXISTS 创建,并在创建后轮询 ANN 索引就绪。
//
// 该方法在每次 Save / BatchSave 之前调用,结果缓存在 initializedTables 中,
// 同一进程内同一 dimension 只会真正打一次 SHOW TABLES + DDL。
func (r *dorisRepository) ensureTable(ctx context.Context, dimension int) error {
if _, ok := r.initializedTables.Load(dimension); ok {
return nil
}
compatMode, err := r.resolveCompatMode(ctx)
if err != nil {
return err
}
log := logger.GetLogger(ctx)
tableName := r.getTableName(dimension)
exists, err := r.tableExists(ctx, tableName)
if err != nil {
log.Errorf("[Doris] Failed to check table existence: %v", err)
return fmt.Errorf("check table existence: %w", err)
}
if !exists {
log.Infof("[Doris] Creating table %s with dimension %d in compat mode %s", tableName, dimension, compatMode)
if err := r.createTable(ctx, tableName, dimension, compatMode); err != nil {
log.Errorf("[Doris] Failed to create table: %v", err)
return fmt.Errorf("create table: %w", err)
}
// ANN 索引在 Doris 端异步构建。这里在后台 goroutine 里轮询就绪,
// 写入路径不阻塞——索引未就绪期间检索会退化为 brute-force(结果对、速度慢),
// 比让首批写入卡 30s 更可接受。
go func(tn string) {
// 用独立 context(带 timeout),避免请求级 ctx 取消把后台轮询也带走。
bgCtx, cancel := context.WithTimeout(context.Background(), annReadyTimeout)
defer cancel()
if err := r.waitANNReady(bgCtx, tn); err != nil {
logger.GetLogger(bgCtx).Warnf(
"[Doris] ANN index for %s not ready within %s: %v "+
"(queries may fall back to brute force temporarily)",
tn, annReadyTimeout, err)
return
}
logger.GetLogger(bgCtx).Infof("[Doris] ANN index for %s ready", tn)
}(tableName)
}
r.initializedTables.Store(dimension, true)
return nil
}
// tableExists 通过 information_schema 判断表是否存在。
//
// 不直接用 SHOW TABLES 是因为 Doris 4.1 对 SHOW TABLES LIKE 大小写敏感,
// 而 information_schema 与 MySQL 兼容性更好。
func (r *dorisRepository) tableExists(ctx context.Context, tableName string) (bool, error) {
const q = `SELECT COUNT(1) FROM information_schema.tables
WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ?`
var n int
if err := r.db.QueryRowContext(ctx, q, r.database, tableName).Scan(&n); err != nil {
return false, err
}
return n > 0, nil
}
// createTable 发出 CREATE TABLE DDL。Doris DDL 是同步的(除 ANN 索引构建外),
// 返回成功即代表表已可写。
func (r *dorisRepository) createTable(ctx context.Context, tableName string, dimension int, compatMode dorisCompatMode) error {
buckets := r.bucketsNum
if buckets <= 0 {
buckets = defaultBucketsNum
}
replication := r.replicationNum
if replication <= 0 {
replication = defaultReplicationNum
}
ddl := buildCreateTableDDL(tableName, dimension, buckets, replication, compatMode)
_, err := r.db.ExecContext(ctx, ddl)
if err != nil && compatMode == dorisCompatModeLegacy {
return fmt.Errorf(
"legacy Doris table creation failed: %w. If your Doris build rejects ANN indexes on UNIQUE KEY tables, set %s=%s before creating embedding tables. %s is not interchangeable after %s_* tables are created",
err,
envDorisCompatMode,
dorisCompatModeInnerProductDuplicate,
envDorisCompatMode,
r.tableBaseName,
)
}
return err
}
// buildCreateTableDDL 根据维度生成 CREATE TABLE DDL。
//
// 关键点:
// - DUPLICATE KEY(id):兼容当前 Doris/SelectDB 对 ANN 索引的表模型要求。
// WeKnora 在 Go 端用 delete + insert 保持按 id 替换的写入语义。
// - INVERTED 索引覆盖所有过滤字段 + 中文分词的 content 全文索引。
// - ANN 索引使用 HNSW + inner_product;Doris 写入/查询前会对向量单位化,
// 因此整体仍保持与其他向量库一致的 cosine 相似度语义。
//
// 注意:DDL 中 dimension / buckets / replication 三个数值字段是 Go 端格式化拼接的,
// 不存在 SQL 注入风险(来源都是受控的 IndexConfig int)。
func buildCreateTableDDL(tableName string, dimension, buckets, replication int, compatMode dorisCompatMode) string {
metricType := "inner_product"
keyMode := "DUPLICATE KEY(id)"
properties := fmt.Sprintf("\t\"replication_num\"=\"%d\"", replication)
if compatMode == dorisCompatModeLegacy {
metricType = "cosine_distance"
keyMode = "UNIQUE KEY(id)"
properties = fmt.Sprintf("\t\"replication_num\"=\"%d\",\n\t\"enable_unique_key_merge_on_write\"=\"true\"", replication)
}
const tpl = `CREATE TABLE IF NOT EXISTS ` + "`%s`" + ` (
id VARCHAR(64) NOT NULL,
chunk_id VARCHAR(64),
knowledge_id VARCHAR(64),
knowledge_base_id VARCHAR(64),
source_id VARCHAR(255),
source_type INT,
tag_id VARCHAR(64),
is_enabled BOOLEAN,
content TEXT,
embedding ARRAY<FLOAT> NOT NULL,
INDEX idx_chunk (chunk_id) USING INVERTED,
INDEX idx_kb (knowledge_base_id) USING INVERTED,
INDEX idx_kid (knowledge_id) USING INVERTED,
INDEX idx_src (source_id) USING INVERTED,
INDEX idx_tag (tag_id) USING INVERTED,
INDEX idx_enabled (is_enabled) USING INVERTED,
INDEX idx_content (content) USING INVERTED PROPERTIES("parser"="chinese","support_phrase"="true"),
INDEX idx_emb (embedding) USING ANN PROPERTIES(
"index_type"="hnsw",
"metric_type"="%s",
"dim"="%d",
"max_degree"="32",
"ef_construction"="200"
)
) ENGINE=OLAP
%s
DISTRIBUTED BY HASH(id) BUCKETS %d
PROPERTIES(
%s
);`
return fmt.Sprintf(tpl, tableName, metricType, dimension, keyMode, buckets, properties)
}
// waitANNReady 轮询 SHOW INDEX,等待 ANN 索引进入 FINISHED 状态。
//
// Doris 的 ANN 索引在建表后会异步构建,期间查询会退化为 brute-force(结果对,速度慢)。
// 此处仅做"尽力而为"的等待:到点未就绪只记 warning,不阻塞写入。
func (r *dorisRepository) waitANNReady(ctx context.Context, tableName string) error {
deadline := time.Now().Add(annReadyTimeout)
for {
ready, err := r.annIndexReady(ctx, tableName)
if err != nil {
return err
}
if ready {
return nil
}
if time.Now().After(deadline) {
return fmt.Errorf("ann index not ready within %s", annReadyTimeout)
}
select {
case <-ctx.Done():
return ctx.Err()
case <-time.After(annReadyPoll):
}
}
}
// annIndexReady 检查 ANN 索引的 State 是否为 FINISHED。
//
// SHOW INDEX FROM <table> 在 Doris 上返回多列;不同小版本列序略有差异,
// 这里以列名匹配(来自 information_schema.statistics + 自定义 view 不可行,
// 直接用 SHOW INDEX 然后扫描即可)。
//
// 兼容策略:如果 SHOW INDEX 返回中找不到 idx_emb 行(极旧版本),视为已就绪,
// 避免因为不同 Doris 版本的输出差异把启动卡死。
func (r *dorisRepository) annIndexReady(ctx context.Context, tableName string) (bool, error) {
rows, err := r.db.QueryContext(ctx,
fmt.Sprintf("SHOW INDEX FROM `%s`", tableName))
if err != nil {
return false, err
}
defer rows.Close()
cols, err := rows.Columns()
if err != nil {
return false, err
}
keyNameIdx, stateIdx := -1, -1
for i, c := range cols {
switch strings.ToLower(c) {
case "key_name":
keyNameIdx = i
case "state", "index_state":
stateIdx = i
}
}
for rows.Next() {
// 使用 sql.RawBytes 接收以兼容不同列类型。
raw := make([]any, len(cols))
ptrs := make([]any, len(cols))
for i := range raw {
ptrs[i] = &raw[i]
}
if err := rows.Scan(ptrs...); err != nil {
return false, err
}
var keyName, state string
if keyNameIdx >= 0 {
keyName = bytesToString(raw[keyNameIdx])
}
if stateIdx >= 0 {
state = bytesToString(raw[stateIdx])
}
if keyName != "idx_emb" {
continue
}
if stateIdx < 0 {
// 旧版本不暴露 state 列,乐观认为已就绪。
return true, nil
}
if !strings.EqualFold(state, "FINISHED") &&
!strings.EqualFold(state, "NORMAL") {
return false, nil
}
}
if err := rows.Err(); err != nil {
return false, err
}
// 走到这里有两种情况:
// 1. 找到了 idx_emb 行,且 state 已是 FINISHED/NORMAL(或 stateIdx<0 的旧版本);
// 2. 没找到 idx_emb 行(极旧 Doris 不暴露该索引名);
// 都视为已就绪,不阻塞。未就绪的分支已在循环内提前 return false。
return true, nil
}
// listEmbeddingTables 返回当前 database 下所有 <base>_% 命名的表,
// 用于关键词检索 / 跨维度 BatchUpdate。
func (r *dorisRepository) listEmbeddingTables(ctx context.Context) ([]string, error) {
const q = `SELECT TABLE_NAME FROM information_schema.tables
WHERE TABLE_SCHEMA = ? AND TABLE_NAME LIKE ?`
rows, err := r.db.QueryContext(ctx, q, r.database, r.tableBaseName+"\\_%")
if err != nil {
return nil, err
}
defer rows.Close()
var names []string
for rows.Next() {
var n string
if err := rows.Scan(&n); err != nil {
return nil, err
}
names = append(names, n)
}
return names, rows.Err()
}
// bytesToString 把 SHOW INDEX 返回的 raw any(通常是 []byte 或 string)
// 安全转成字符串。
func bytesToString(v any) string {
switch s := v.(type) {
case []byte:
return string(s)
case string:
return s
case nil:
return ""
default:
return fmt.Sprintf("%v", s)
}
}
@@ -0,0 +1,508 @@
package doris
import (
"bytes"
"context"
"encoding/base64"
"encoding/json"
"fmt"
"io"
"net/http"
"net/url"
"strings"
"github.com/Tencent/WeKnora/internal/logger"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
// Stream Load 相关常量。
const (
// 单批 Stream Load 的 JSON body 上限(保守值,远小于 Doris 默认 streaming_load_max_mb=10240)。
// 主要目的是控制单次 HTTP 请求的尾延迟,超过则自动拆批。
streamLoadMaxBatchBytes = 1 << 20 // 1 MiB
// HTTP 头 Authorization 用 Basic auth;Stream Load 也支持 token,
// 这里走最常用的用户名/密码方案与 MySQL 协议保持一致。
headerAuthorization = "Authorization"
headerExpect = "Expect"
headerContentType = "Content-Type"
)
// streamLoadResponse 是 Doris FE/BE 返回的 Stream Load 结果体。
//
// 关键字段:Status 应为 "Success" 或 "Publish Timeout"(后者表示数据已写入但发布事务超时,
// 仍视为成功)。其它状态都视为失败。
type streamLoadResponse struct {
TxnId int64 `json:"TxnId"`
Label string `json:"Label"`
Status string `json:"Status"`
Message string `json:"Message"`
NumberTotalRows int64 `json:"NumberTotalRows"`
NumberLoadedRows int64 `json:"NumberLoadedRows"`
NumberFilteredRows int64 `json:"NumberFilteredRows"`
NumberUnselectedRows int64 `json:"NumberUnselectedRows"`
LoadBytes int64 `json:"LoadBytes"`
LoadTimeMs int64 `json:"LoadTimeMs"`
BeginTxnTimeMs int64 `json:"BeginTxnTimeMs"`
StreamLoadPutTimeMs int64 `json:"StreamLoadPutTimeMs"`
ReadDataTimeMs int64 `json:"ReadDataTimeMs"`
WriteDataTimeMs int64 `json:"WriteDataTimeMs"`
CommitAndPublishTimeMs int64 `json:"CommitAndPublishTimeMs"`
ErrorURL string `json:"ErrorURL"`
}
// streamLoadURL 拼装某张表的 Stream Load HTTP 端点。
func (r *dorisRepository) streamLoadURL(table string) string {
return fmt.Sprintf("%s/api/%s/%s/_stream_load",
r.feHTTPBase, url.PathEscape(r.database), url.PathEscape(table))
}
// newDorisStreamLoadHTTPClient protects both the initial FE request and every
// FE -> BE redirect at connection time. Doris requires Basic auth to survive
// the 307 redirect, so redirect targets on another host must be explicitly
// trusted through the SSRF whitelist before credentials are forwarded.
func newDorisStreamLoadHTTPClient() *http.Client {
cfg := secutils.DefaultSSRFSafeHTTPClientConfig()
// Stream Load calls carry their own context deadline. Preserve the previous
// client behaviour instead of imposing the generic 30-second HTTP timeout.
cfg.Timeout = 0
client := secutils.NewSSRFSafeHTTPClient(cfg)
ssrfCheckRedirect := client.CheckRedirect
client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
// Keep the shared client's redirect limit, scheme/target validation and
// sensitive-header stripping as the authoritative SSRF behaviour.
if err := ssrfCheckRedirect(req, via); err != nil {
return err
}
if len(via) == 0 {
return nil
}
sourceHost := via[len(via)-1].URL.Hostname()
targetHost := req.URL.Hostname()
if !strings.EqualFold(sourceHost, targetHost) && !secutils.IsSSRFWhitelisted(targetHost) {
return fmt.Errorf(
"stream load redirect blocked: target host %q is not trusted to receive credentials",
targetHost,
)
}
// The shared client deliberately strips Authorization on cross-host
// redirects. Doris is the exceptional case where an explicitly trusted
// BE needs the same Basic credentials after the FE's 307 response.
req.Header.Set(headerAuthorization, via[0].Header.Get(headerAuthorization))
return nil
}
return client
}
// partialUpdateRows 把若干行通过 Stream Load 的 partial update 模式写回目标表。
//
// columns 是参与本次 partial update 的列(必须包含 UNIQUE KEY 列,即 "id")。
// rows 中每一项是一个与 columns 等长的字段值数组。
//
// 实现要点:
// 1. 用 JSON 数组的 body 形式,header 加 strip_outer_array=true。
// 2. 设置 partial_columns=true、merge_type=APPEND,触发 Doris 的 partial update 模式
// (Doris 4.1 + UNIQUE KEY MoW 表的标准玩法)。
// 3. 按 streamLoadMaxBatchBytes 自动拆批,避免单次过大。
// 4. 处理 307:Doris 的 FE 会 redirect 到 BE,net/http 默认会跟随;
// 这里需要保证 GetBody 可重发(已通过 bytes.NewReader 构造 Body 满足)。
func (r *dorisRepository) partialUpdateRows(ctx context.Context,
table string, columns []string, rows []map[string]any,
) error {
if len(rows) == 0 {
return nil
}
for _, batch := range chunkRows(rows, streamLoadMaxBatchBytes) {
if err := r.streamLoadOnce(ctx, table, columns, batch); err != nil {
return err
}
}
return nil
}
// streamLoadOnce 发出一次 Stream Load HTTP 请求。
func (r *dorisRepository) streamLoadOnce(ctx context.Context,
table string, columns []string, rows []map[string]any,
) error {
log := logger.GetLogger(ctx)
if len(rows) == 0 {
return nil
}
body, err := json.Marshal(rows)
if err != nil {
return fmt.Errorf("marshal stream load body: %w", err)
}
url := r.streamLoadURL(table)
// Validate at the final outbound boundary as well as when user-supplied
// vector-store configuration is created. This covers stored/env configs and
// keeps the tainted value from reaching http.Client.Do unchecked.
if err := secutils.ValidateURLForSSRF(url); err != nil {
return fmt.Errorf("stream load URL blocked by SSRF validation: %w", err)
}
req, err := http.NewRequestWithContext(ctx, http.MethodPut, url, bytes.NewReader(body))
if err != nil {
return fmt.Errorf("build stream load request: %w", err)
}
// GetBody 让 redirect 时可以重新读 body(FE -> BE 的 307 需要重发 PUT body)。
bodyCopy := append([]byte(nil), body...)
req.GetBody = func() (io.ReadCloser, error) {
return io.NopCloser(bytes.NewReader(bodyCopy)), nil
}
req.ContentLength = int64(len(body))
auth := base64.StdEncoding.EncodeToString([]byte(r.username + ":" + r.password))
req.Header.Set(headerAuthorization, "Basic "+auth)
req.Header.Set(headerExpect, "100-continue")
req.Header.Set(headerContentType, "application/json")
req.Header.Set("format", "json")
req.Header.Set("strip_outer_array", "true")
req.Header.Set("partial_columns", "true")
req.Header.Set("columns", strings.Join(columns, ","))
req.Header.Set("merge_type", "APPEND")
resp, err := r.httpClient.Do(req)
if err != nil {
return fmt.Errorf("stream load HTTP: %w", err)
}
defer resp.Body.Close()
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return fmt.Errorf("read stream load response: %w", err)
}
if resp.StatusCode/100 != 2 {
return fmt.Errorf("stream load HTTP %d: %s", resp.StatusCode, string(respBody))
}
var result streamLoadResponse
if err := json.Unmarshal(respBody, &result); err != nil {
return fmt.Errorf("decode stream load response: %w (raw=%s)", err, string(respBody))
}
switch result.Status {
case "Success", "Publish Timeout":
log.Infof("[Doris] Stream load %s OK: rows=%d, loaded=%d, label=%s",
table, result.NumberTotalRows, result.NumberLoadedRows, result.Label)
return nil
default:
return fmt.Errorf("stream load failed: status=%s msg=%s err_url=%s",
result.Status, result.Message, result.ErrorURL)
}
}
// chunkRows 把行按累积 JSON 体大小切分,每段不超过 maxBytes。
//
// 注意:JSON 序列化的实际开销大约是 marshal 后的字节数,而单行 marshal
// 加上逗号 + 数组括号约等于本估算。这里用粗略估计避免每段都 marshal。
func chunkRows(rows []map[string]any, maxBytes int) [][]map[string]any {
if len(rows) == 0 {
return nil
}
var (
out [][]map[string]any
curr []map[string]any
size int
header = 2 // "[" + "]"
)
for _, row := range rows {
raw, err := json.Marshal(row)
if err != nil {
// marshal 失败时把这一行单独成段,由上层 streamLoadOnce 再次 marshal 报错。
if len(curr) > 0 {
out = append(out, curr)
}
out = append(out, []map[string]any{row})
curr = nil
size = 0
continue
}
// 加上逗号位(除第一行之外)。
need := len(raw)
if len(curr) > 0 {
need++
}
if size+need+header > maxBytes && len(curr) > 0 {
out = append(out, curr)
curr = nil
size = 0
}
curr = append(curr, row)
size += need
}
if len(curr) > 0 {
out = append(out, curr)
}
return out
}
// ---------------------------------------------------------------------------
// 业务面方法:BatchUpdateChunkEnabledStatus / BatchUpdateChunkTagID
// ---------------------------------------------------------------------------
// BatchUpdateChunkEnabledStatus 批量更新 chunk 的 is_enabled 字段。
// legacy 模式走 Stream Load partial update;inner_product_duplicate 模式改为读整行后 replaceRows 写回。
func (r *dorisRepository) BatchUpdateChunkEnabledStatus(ctx context.Context,
chunkStatusMap map[string]bool,
) error {
if len(chunkStatusMap) == 0 {
return nil
}
compatMode, err := r.resolveCompatMode(ctx)
if err != nil {
return err
}
if !compatMode.usesRewriteChunkUpdates() {
return r.batchUpdateChunkEnabledStatusLegacy(ctx, chunkStatusMap)
}
chunkIDs := make([]string, 0, len(chunkStatusMap))
for id := range chunkStatusMap {
chunkIDs = append(chunkIDs, id)
}
return r.rewriteChunkRows(ctx, chunkIDs, func(row *DorisVectorEmbedding) bool {
enabled, ok := chunkStatusMap[row.ChunkID]
if !ok || row.IsEnabled == enabled {
return false
}
row.IsEnabled = enabled
return true
}, "rewrite is_enabled")
}
// BatchUpdateChunkTagID 批量更新 chunk 的 tag_id 字段。逻辑与 EnabledStatus 一致。
func (r *dorisRepository) BatchUpdateChunkTagID(ctx context.Context,
chunkTagMap map[string]string,
) error {
if len(chunkTagMap) == 0 {
return nil
}
compatMode, err := r.resolveCompatMode(ctx)
if err != nil {
return err
}
if !compatMode.usesRewriteChunkUpdates() {
return r.batchUpdateChunkTagIDLegacy(ctx, chunkTagMap)
}
chunkIDs := make([]string, 0, len(chunkTagMap))
for id := range chunkTagMap {
chunkIDs = append(chunkIDs, id)
}
return r.rewriteChunkRows(ctx, chunkIDs, func(row *DorisVectorEmbedding) bool {
tagID, ok := chunkTagMap[row.ChunkID]
if !ok || row.TagID == tagID {
return false
}
row.TagID = tagID
return true
}, "rewrite tag_id")
}
func (r *dorisRepository) rewriteChunkRows(ctx context.Context,
chunkIDs []string,
mutate func(*DorisVectorEmbedding) bool,
action string,
) error {
if len(chunkIDs) == 0 {
return nil
}
tables, err := r.listEmbeddingTables(ctx)
if err != nil {
return fmt.Errorf("list tables: %w", err)
}
for _, table := range tables {
rows, err := r.loadRowsByChunkIDs(ctx, table, chunkIDs)
if err != nil {
return fmt.Errorf("load chunk rows from %s: %w", table, err)
}
updated := make([]*DorisVectorEmbedding, 0, len(rows))
for _, row := range rows {
if !mutate(row) {
continue
}
updated = append(updated, row)
}
if len(updated) == 0 {
continue
}
if err := r.replaceRows(ctx, table, updated); err != nil {
return fmt.Errorf("%s in %s: %w", action, table, err)
}
}
return nil
}
func (r *dorisRepository) batchUpdateChunkEnabledStatusLegacy(ctx context.Context,
chunkStatusMap map[string]bool,
) error {
chunkIDs := make([]string, 0, len(chunkStatusMap))
for id := range chunkStatusMap {
chunkIDs = append(chunkIDs, id)
}
mapping, err := r.lookupChunkRowKeys(ctx, chunkIDs)
if err != nil {
return err
}
byTable := make(map[string][]map[string]any)
for chunkID, locations := range mapping {
enabled, ok := chunkStatusMap[chunkID]
if !ok {
continue
}
for _, loc := range locations {
byTable[loc.table] = append(byTable[loc.table], map[string]any{
fieldID: loc.id,
fieldIsEnabled: enabled,
})
}
}
for table, rows := range byTable {
if err := r.partialUpdateRows(ctx, table, []string{fieldID, fieldIsEnabled}, rows); err != nil {
return fmt.Errorf("partial update is_enabled in %s: %w", table, err)
}
}
return nil
}
func (r *dorisRepository) batchUpdateChunkTagIDLegacy(ctx context.Context,
chunkTagMap map[string]string,
) error {
chunkIDs := make([]string, 0, len(chunkTagMap))
for id := range chunkTagMap {
chunkIDs = append(chunkIDs, id)
}
mapping, err := r.lookupChunkRowKeys(ctx, chunkIDs)
if err != nil {
return err
}
byTable := make(map[string][]map[string]any)
for chunkID, locations := range mapping {
tagID, ok := chunkTagMap[chunkID]
if !ok {
continue
}
for _, loc := range locations {
byTable[loc.table] = append(byTable[loc.table], map[string]any{
fieldID: loc.id,
fieldTagID: tagID,
})
}
}
for table, rows := range byTable {
if err := r.partialUpdateRows(ctx, table, []string{fieldID, fieldTagID}, rows); err != nil {
return fmt.Errorf("partial update tag_id in %s: %w", table, err)
}
}
return nil
}
func (r *dorisRepository) loadRowsByChunkIDs(ctx context.Context,
table string, chunkIDs []string,
) ([]*DorisVectorEmbedding, error) {
if len(chunkIDs) == 0 {
return nil, nil
}
placeholders := make([]string, len(chunkIDs))
args := make([]any, len(chunkIDs))
for i, v := range chunkIDs {
placeholders[i] = "?"
args[i] = v
}
stmt := fmt.Sprintf(
"SELECT %s FROM `%s` WHERE %s IN (%s)",
strings.Join(columnsForCopy, ", "),
table,
fieldChunkID,
strings.Join(placeholders, ", "),
)
rows, err := r.db.QueryContext(ctx, stmt, args...)
if err != nil {
return nil, err
}
batch, err := scanCopyRows(rows)
_ = rows.Close()
if err != nil {
return nil, err
}
return batch, nil
}
// rowLocation 表示某行在哪个表里、主键 id 是什么。
type rowLocation struct {
table string
id string
}
// lookupChunkRowKeys 查询给定的 chunkIDs 在所有 <base>_<dim> 表中的物理位置:
// - key:chunk_id
// - value:[(table, id), ...],因为同一 chunk 可能在多个维度的表里都有副本。
//
// 跨表查询使用 listEmbeddingTables 列出的所有匹配表;每张表执行一次
// SELECT id, chunk_id FROM <table> WHERE chunk_id IN (?, ?, ...)。
func (r *dorisRepository) lookupChunkRowKeys(ctx context.Context,
chunkIDs []string,
) (map[string][]rowLocation, error) {
if len(chunkIDs) == 0 {
return nil, nil
}
tables, err := r.listEmbeddingTables(ctx)
if err != nil {
return nil, fmt.Errorf("list tables: %w", err)
}
if len(tables) == 0 {
return nil, nil
}
placeholders := make([]string, len(chunkIDs))
args := make([]any, len(chunkIDs))
for i, v := range chunkIDs {
placeholders[i] = "?"
args[i] = v
}
out := make(map[string][]rowLocation)
for _, table := range tables {
stmt := fmt.Sprintf(
"SELECT %s, %s FROM `%s` WHERE %s IN (%s)",
fieldID, fieldChunkID, table, fieldChunkID, strings.Join(placeholders, ", "),
)
rows, err := r.db.QueryContext(ctx, stmt, args...)
if err != nil {
return nil, fmt.Errorf("lookup chunk row keys in %s: %w", table, err)
}
for rows.Next() {
var id, chunkID string
if err := rows.Scan(&id, &chunkID); err != nil {
_ = rows.Close()
return nil, fmt.Errorf("scan row keys: %w", err)
}
out[chunkID] = append(out[chunkID], rowLocation{table: table, id: id})
}
if err := rows.Err(); err != nil {
_ = rows.Close()
return nil, err
}
_ = rows.Close()
}
return out, nil
}
@@ -0,0 +1,70 @@
package doris
import (
"database/sql"
"net/http"
"sync"
)
// dorisRepository 是 Apache Doris 4.1 的检索引擎仓储实现。
//
// 通信通道:
// - 读写主链路:MySQL 协议(database/sql + go-sql-driver/mysql),FE 默认 9030 端口。
// - Stream Load:HTTP(FE 默认 8030 端口),legacy 模式用于 partial update。
//
// 表结构按维度分表:<tableBaseName>_<dim>。
// 兼容模式由 DORIS_COMPAT_MODE 决定:
// - legacy:UNIQUE KEY(id) + cosine_distance ANN + Stream Load partial update
// - inner_product_duplicate:DUPLICATE KEY(id) + normalized inner product + delete/insert rewrite
//
// 该设置在 embedding 表创建后不可直接互换;切换模式前需要重建这些表。
//
// 与 Qdrant/Milvus/Weaviate 一样,initializedTables 缓存"已确保存在"的维度,
// 避免每次写入都打 SHOW TABLES。
type dorisRepository struct {
db *sql.DB
httpClient *http.Client
// fe HTTP base,例如 "http://doris-fe:8030"。Stream Load 路径
// 由 streamLoadURL(table) 拼接:<feHTTPBase>/api/<database>/<table>/_stream_load。
feHTTPBase string
username string
password string
database string
tableBaseName string
bucketsNum int // 0 -> default 10
replicationNum int // 0 -> default 1
compatModeRequested dorisCompatMode
compatModeResolved dorisCompatMode
compatResolveOnce sync.Once
compatResolveErr error
// 已经确保过 ensureTable 的维度集合:dim -> true。
initializedTables sync.Map
}
// DorisVectorEmbedding 是落到 Doris 表里的一行的领域模型。
//
// 字段顺序与 schema.go 中的 INSERT 列序保持一致,
// 调整时需要同时更新 createInsert 与 columns。
type DorisVectorEmbedding struct {
ID string
Content string
SourceID string
SourceType int
ChunkID string
KnowledgeID string
KnowledgeBaseID string
TagID string
IsEnabled bool
Embedding []float32
}
// DorisVectorEmbeddingWithScore 是检索结果的领域模型,
// Score 在向量检索时按当前 compat mode 计算,在关键词检索时统一赋 1.0。
type DorisVectorEmbeddingWithScore struct {
DorisVectorEmbedding
Score float64
}
@@ -0,0 +1,78 @@
package elasticsearch
import (
"maps"
"slices"
"github.com/Tencent/WeKnora/internal/types"
)
// VectorEmbedding defines the Elasticsearch document structure for vector embeddings
type VectorEmbedding struct {
Content string `json:"content" gorm:"column:content;not null"` // Text content of the chunk
SourceID string `json:"source_id" gorm:"column:source_id;not null"` // ID of the source document
SourceType int `json:"source_type" gorm:"column:source_type;not null"` // Type of the source document
ChunkID string `json:"chunk_id" gorm:"column:chunk_id"` // Unique ID of the text chunk
KnowledgeID string `json:"knowledge_id" gorm:"column:knowledge_id"` // ID of the knowledge item
KnowledgeBaseID string `json:"knowledge_base_id" gorm:"column:knowledge_base_id"` // ID of the knowledge base
TagID string `json:"tag_id" gorm:"column:tag_id"` // Tag ID for categorization
Embedding []float32 `json:"embedding" gorm:"column:embedding;not null"` // Vector embedding of the content
IsEnabled bool `json:"is_enabled"` // Whether the chunk is enabled
IsRecommended bool `json:"is_recommended"` // Whether the chunk is recommended
}
// VectorEmbeddingWithScore extends VectorEmbedding with similarity score
type VectorEmbeddingWithScore struct {
VectorEmbedding
Score float64 // Similarity score from vector search
}
// ToDBVectorEmbedding converts IndexInfo to Elasticsearch document format
func ToDBVectorEmbedding(embedding *types.IndexInfo, additionalParams map[string]interface{}) *VectorEmbedding {
vector := &VectorEmbedding{
Content: embedding.Content,
SourceID: embedding.SourceID,
SourceType: int(embedding.SourceType),
ChunkID: embedding.ChunkID,
KnowledgeID: embedding.KnowledgeID,
KnowledgeBaseID: embedding.KnowledgeBaseID,
TagID: embedding.TagID,
IsEnabled: embedding.IsEnabled,
IsRecommended: embedding.IsRecommended,
}
// Add embedding data if available in additionalParams
if additionalParams != nil && slices.Contains(slices.Collect(maps.Keys(additionalParams)), "embedding") {
if embeddingMap, ok := additionalParams["embedding"].(map[string][]float32); ok {
vector.Embedding = embeddingMap[embedding.SourceID]
}
}
// Get is_enabled from additionalParams if available
if additionalParams != nil {
if chunkEnabledMap, ok := additionalParams["chunk_enabled"].(map[string]bool); ok {
if enabled, exists := chunkEnabledMap[embedding.ChunkID]; exists {
vector.IsEnabled = enabled
}
}
}
return vector
}
// FromDBVectorEmbeddingWithScore converts Elasticsearch document to IndexWithScore domain model
func FromDBVectorEmbeddingWithScore(id string,
embedding *VectorEmbeddingWithScore,
matchType types.MatchType,
) *types.IndexWithScore {
return &types.IndexWithScore{
ID: id,
SourceID: embedding.SourceID,
SourceType: types.SourceType(embedding.SourceType),
ChunkID: embedding.ChunkID,
KnowledgeID: embedding.KnowledgeID,
KnowledgeBaseID: embedding.KnowledgeBaseID,
TagID: embedding.TagID,
Content: embedding.Content,
Score: embedding.Score,
MatchType: matchType,
IsEnabled: embedding.IsEnabled,
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,820 @@
package v8
import (
"context"
"encoding/json"
"fmt"
"strings"
elasticsearchRetriever "github.com/Tencent/WeKnora/internal/application/repository/retriever/elasticsearch"
"github.com/Tencent/WeKnora/internal/config"
"github.com/Tencent/WeKnora/internal/logger"
typesLocal "github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
"github.com/elastic/go-elasticsearch/v8"
"github.com/elastic/go-elasticsearch/v8/typedapi/core/search"
"github.com/elastic/go-elasticsearch/v8/typedapi/types"
"github.com/elastic/go-elasticsearch/v8/typedapi/types/enums/scriptlanguage"
"github.com/google/uuid"
)
// elasticsearchRepository implements the RetrieveEngineRepository interface for Elasticsearch v8
type elasticsearchRepository struct {
client *elasticsearch.TypedClient // Elasticsearch client instance
index string // Name of the Elasticsearch index to use
useKeywordSuffix bool // Whether to append .keyword suffix to ID field names in queries
numberOfShards int // Shard count for index creation (0 = ES default)
numberOfReplicas int // Replica count for index creation (-1 = unset, use ES default)
}
// NewElasticsearchEngineRepository creates and initializes a new Elasticsearch v8 repository.
// indexCfg is optional — pass nil to use env var / default values (env path).
func NewElasticsearchEngineRepository(client *elasticsearch.TypedClient,
config *config.Config,
indexCfg *typesLocal.IndexConfig,
) interfaces.RetrieveEngineRepository {
log := logger.GetLogger(context.Background())
log.Info("[Elasticsearch] Initializing Elasticsearch v8 retriever engine repository")
indexName := typesLocal.ResolveIndexName(indexCfg, "ELASTICSEARCH_INDEX", "xwrag_default")
// Create repository instance and ensure index exists
res := &elasticsearchRepository{
client: client,
index: indexName,
numberOfShards: indexCfg.GetNumberOfShards(0),
numberOfReplicas: indexCfg.GetNumberOfReplicas(-1),
}
if err := res.createIndexIfNotExists(context.Background()); err != nil {
log.Errorf("[Elasticsearch] Failed to create index: %v", err)
} else {
log.Info("[Elasticsearch] Successfully initialized repository")
}
res.detectFieldTypes(context.Background())
return res
}
// idField returns the query field name for an ID field, appending ".keyword"
// suffix when the index uses text-type mappings with keyword sub-fields.
func (e *elasticsearchRepository) idField(name string) string {
if e.useKeywordSuffix {
return name + ".keyword"
}
return name
}
// detectFieldTypes inspects the index mapping to determine whether ID fields
// are mapped as "keyword" (no suffix needed) or "text" (needs ".keyword" suffix).
func (e *elasticsearchRepository) detectFieldTypes(ctx context.Context) {
log := logger.GetLogger(ctx)
mappingResp, err := e.client.Indices.GetMapping().Index(e.index).Do(ctx)
if err != nil {
log.Warnf("[Elasticsearch] Failed to get index mapping, defaulting to .keyword suffix: %v", err)
e.useKeywordSuffix = true
return
}
indexMapping, ok := mappingResp[e.index]
if !ok {
log.Warnf("[Elasticsearch] Index %s not found in mapping response, defaulting to .keyword suffix", e.index)
e.useKeywordSuffix = true
return
}
if prop, ok := indexMapping.Mappings.Properties["chunk_id"]; ok {
propBytes, err := json.Marshal(prop)
if err == nil {
var propInfo struct {
Type string `json:"type"`
}
if json.Unmarshal(propBytes, &propInfo) == nil && propInfo.Type == "keyword" {
e.useKeywordSuffix = false
log.Infof("[Elasticsearch] Detected keyword type for ID fields, querying without .keyword suffix")
return
}
}
e.useKeywordSuffix = true
log.Infof("[Elasticsearch] ID fields are not keyword type, querying with .keyword suffix")
return
}
log.Infof("[Elasticsearch] No mapping detected for chunk_id (empty index?), defaulting to .keyword suffix")
e.useKeywordSuffix = true
}
// EngineType returns the type of retriever engine (Elasticsearch)
func (e *elasticsearchRepository) EngineType() typesLocal.RetrieverEngineType {
return typesLocal.ElasticsearchRetrieverEngineType
}
// Support returns the retrieval types supported by this repository (Keywords and Vector)
func (e *elasticsearchRepository) Support() []typesLocal.RetrieverType {
return []typesLocal.RetrieverType{typesLocal.KeywordsRetrieverType, typesLocal.VectorRetrieverType}
}
// calculateStorageSize estimates the storage size in bytes for a single index document
func (e *elasticsearchRepository) calculateStorageSize(embedding *elasticsearchRetriever.VectorEmbedding) int64 {
// 1. Content text size
contentSizeBytes := int64(len(embedding.Content))
// 2. Vector embedding size
var vectorSizeBytes int64 = 0
if embedding.Embedding != nil {
// 4 bytes per dimension (full precision float)
vectorSizeBytes = int64(len(embedding.Embedding) * 4)
}
// 3. Metadata size (IDs, timestamps, and other fixed overhead)
metadataSizeBytes := int64(250) // Approximately 250 bytes of metadata
// 4. Index overhead (Elasticsearch index expansion factor ~1.5)
indexOverheadBytes := (contentSizeBytes + vectorSizeBytes) * 5 / 10
// Total size in bytes
totalSizeBytes := contentSizeBytes + vectorSizeBytes + metadataSizeBytes + indexOverheadBytes
return totalSizeBytes
}
// EstimateStorageSize calculates the estimated storage size for a list of indices
// Returns the total size in bytes
func (e *elasticsearchRepository) EstimateStorageSize(ctx context.Context,
indexInfoList []*typesLocal.IndexInfo, params map[string]any,
) int64 {
var totalStorageSize int64 = 0
for _, embedding := range indexInfoList {
embeddingDB := elasticsearchRetriever.ToDBVectorEmbedding(embedding, params)
totalStorageSize += e.calculateStorageSize(embeddingDB)
}
logger.GetLogger(ctx).Infof(
"[Elasticsearch] Storage size for %d indices: %d bytes", len(indexInfoList), totalStorageSize,
)
return totalStorageSize
}
// Save stores a single index document in Elasticsearch
// Returns an error if the operation fails
func (e *elasticsearchRepository) Save(ctx context.Context,
embedding *typesLocal.IndexInfo,
additionalParams map[string]any,
) error {
log := logger.GetLogger(ctx)
log.Debugf("[Elasticsearch] Saving index for chunk ID: %s", embedding.ChunkID)
// Convert to database format
embeddingDB := elasticsearchRetriever.ToDBVectorEmbedding(embedding, additionalParams)
if len(embeddingDB.Embedding) == 0 {
err := fmt.Errorf("empty embedding vector for chunk ID: %s", embedding.ChunkID)
log.Errorf("[Elasticsearch] %v", err)
return err
}
// Index the document
resp, err := e.client.Index(e.index).Request(embeddingDB).Do(ctx)
if err != nil {
log.Errorf("[Elasticsearch] Failed to save index: %v", err)
return err
}
log.Infof("[Elasticsearch] Successfully saved index for chunk ID: %s, document ID: %s", embedding.ChunkID, resp.Id_)
return nil
}
// BatchSave stores multiple index documents in Elasticsearch using bulk API
// Returns an error if the operation fails
func (e *elasticsearchRepository) BatchSave(ctx context.Context,
embeddingList []*typesLocal.IndexInfo, additionalParams map[string]any,
) error {
log := logger.GetLogger(ctx)
if len(embeddingList) == 0 {
log.Warn("[Elasticsearch] Empty list provided to BatchSave, skipping")
return nil
}
log.Infof("[Elasticsearch] Batch saving %d indices", len(embeddingList))
indexRequest := e.client.Bulk().Index(e.index)
// Add each document to the bulk request
for _, embedding := range embeddingList {
embeddingDB := elasticsearchRetriever.ToDBVectorEmbedding(embedding, additionalParams)
err := indexRequest.CreateOp(types.CreateOperation{Index_: &e.index}, embeddingDB)
if err != nil {
log.Errorf("[Elasticsearch] Failed to create bulk operation: %v", err)
return fmt.Errorf("failed to create op: %w", err)
}
log.Debugf("[Elasticsearch] Added chunk ID %s to bulk request", embedding.ChunkID)
}
// Execute the bulk request
_, err := indexRequest.Do(ctx)
if err != nil {
log.Errorf("[Elasticsearch] Failed to execute bulk operation: %v", err)
return fmt.Errorf("failed to do bulk: %w", err)
}
log.Infof("[Elasticsearch] Successfully batch saved %d indices", len(embeddingList))
return nil
}
// DeleteByChunkIDList removes documents from the index based on chunk IDs
// Returns an error if the delete operation fails
func (e *elasticsearchRepository) DeleteByChunkIDList(ctx context.Context, chunkIDList []string, dimension int, knowledgeType string) error {
log := logger.GetLogger(ctx)
if len(chunkIDList) == 0 {
log.Warn("[Elasticsearch] Empty chunk ID list provided for deletion, skipping")
return nil
}
log.Infof("[Elasticsearch] Deleting indices by chunk IDs, count: %d", len(chunkIDList))
// Use DeleteByQuery to delete all documents matching the chunk IDs
_, err := e.client.DeleteByQuery(e.index).Query(&types.Query{
Terms: &types.TermsQuery{TermsQuery: map[string]types.TermsQueryField{e.idField("chunk_id"): chunkIDList}},
}).Do(ctx)
if err != nil {
log.Errorf("[Elasticsearch] Failed to delete by chunk IDs: %v", err)
return fmt.Errorf("failed to delete by query: %w", err)
}
log.Infof("[Elasticsearch] Successfully deleted documents by chunk IDs")
return nil
}
// DeleteBySourceIDList removes documents from the index based on source IDs
// Returns an error if the delete operation fails
func (e *elasticsearchRepository) DeleteBySourceIDList(ctx context.Context, sourceIDList []string, dimension int, knowledgeType string) error {
log := logger.GetLogger(ctx)
if len(sourceIDList) == 0 {
log.Warn("[Elasticsearch] Empty source ID list provided for deletion, skipping")
return nil
}
log.Infof("[Elasticsearch] Deleting indices by source IDs, count: %d", len(sourceIDList))
// Use DeleteByQuery to delete all documents matching the source IDs
_, err := e.client.DeleteByQuery(e.index).Query(&types.Query{
Terms: &types.TermsQuery{TermsQuery: map[string]types.TermsQueryField{e.idField("source_id"): sourceIDList}},
}).Do(ctx)
if err != nil {
log.Errorf("[Elasticsearch] Failed to delete by source IDs: %v", err)
return fmt.Errorf("failed to delete by query: %w", err)
}
log.Infof("[Elasticsearch] Successfully deleted documents by source IDs")
return nil
}
// DeleteByKnowledgeIDList removes documents from the index based on knowledge IDs
// Returns an error if the delete operation fails
func (e *elasticsearchRepository) DeleteByKnowledgeIDList(ctx context.Context,
knowledgeIDList []string, dimension int, knowledgeType string,
) error {
log := logger.GetLogger(ctx)
if len(knowledgeIDList) == 0 {
log.Warn("[Elasticsearch] Empty knowledge ID list provided for deletion, skipping")
return nil
}
log.Infof("[Elasticsearch] Deleting indices by knowledge IDs, count: %d", len(knowledgeIDList))
// Use DeleteByQuery to delete all documents matching the knowledge IDs
_, err := e.client.DeleteByQuery(e.index).Query(&types.Query{
Terms: &types.TermsQuery{TermsQuery: map[string]types.TermsQueryField{e.idField("knowledge_id"): knowledgeIDList}},
}).Do(ctx)
if err != nil {
log.Errorf("[Elasticsearch] Failed to delete by knowledge IDs: %v", err)
return fmt.Errorf("failed to delete by query: %w", err)
}
log.Infof("[Elasticsearch] Successfully deleted documents by knowledge IDs")
return nil
}
// getBaseConds creates the base query conditions for retrieval operations
// Returns a slice of Query objects with must and must_not conditions
// KnowledgeBaseIDs and KnowledgeIDs use AND logic (search specific documents within knowledge bases)
func (e *elasticsearchRepository) getBaseConds(params typesLocal.RetrieveParams) []types.Query {
must := []types.Query{}
// KnowledgeBaseIDs and KnowledgeIDs use AND logic
// - If only KnowledgeBaseIDs: search entire knowledge bases
// - If only KnowledgeIDs: search specific documents
// - If both: search specific documents within the knowledge bases (AND)
if len(params.KnowledgeBaseIDs) > 0 {
must = append(must, types.Query{Terms: &types.TermsQuery{
TermsQuery: map[string]types.TermsQueryField{
e.idField("knowledge_base_id"): params.KnowledgeBaseIDs,
},
}})
}
if len(params.KnowledgeIDs) > 0 {
must = append(must, types.Query{Terms: &types.TermsQuery{
TermsQuery: map[string]types.TermsQueryField{
e.idField("knowledge_id"): params.KnowledgeIDs,
},
}})
}
// Filter by tag IDs if specified
if len(params.TagIDs) > 0 {
must = append(must, types.Query{Terms: &types.TermsQuery{
TermsQuery: map[string]types.TermsQueryField{
e.idField("tag_id"): params.TagIDs,
},
}})
}
mustNot := make([]types.Query, 0)
// Exclude disabled chunks (is_enabled = false)
// Note: Historical data without is_enabled field will be included (not matching must_not)
mustNot = append(mustNot, types.Query{Term: map[string]types.TermQuery{
"is_enabled": {Value: false},
}})
if len(params.ExcludeKnowledgeIDs) > 0 {
mustNot = append(mustNot, types.Query{Terms: &types.TermsQuery{
TermsQuery: map[string]types.TermsQueryField{e.idField("knowledge_id"): params.ExcludeKnowledgeIDs},
}})
}
if len(params.ExcludeChunkIDs) > 0 {
mustNot = append(mustNot, types.Query{Terms: &types.TermsQuery{
TermsQuery: map[string]types.TermsQueryField{e.idField("chunk_id"): params.ExcludeChunkIDs},
}})
}
return []types.Query{{Bool: &types.BoolQuery{Must: must, MustNot: mustNot}}}
}
// createIndexIfNotExists checks if the specified index exists and creates it if not
// Returns an error if the operation fails
func (e *elasticsearchRepository) createIndexIfNotExists(ctx context.Context) error {
log := logger.GetLogger(ctx)
log.Debugf("[Elasticsearch] Checking if index exists: %s", e.index)
// Check if index exists
exists, err := e.client.Indices.Exists(e.index).Do(ctx)
if err != nil {
log.Errorf("[Elasticsearch] Failed to check if index exists: %v", err)
return err
}
if exists {
log.Debugf("[Elasticsearch] Index already exists: %s", e.index)
return nil
}
// Create index if it doesn't exist, with optional shards/replicas settings
log.Infof("[Elasticsearch] Creating index: %s", e.index)
createReq := e.client.Indices.Create(e.index)
if e.numberOfShards > 0 || e.numberOfReplicas >= 0 {
settings := &types.IndexSettings{}
if e.numberOfShards > 0 {
shards := fmt.Sprintf("%d", e.numberOfShards)
settings.NumberOfShards = &shards
}
if e.numberOfReplicas >= 0 {
replicas := fmt.Sprintf("%d", e.numberOfReplicas)
settings.NumberOfReplicas = &replicas
}
createReq = createReq.Settings(settings)
}
_, err = createReq.Do(ctx)
if err != nil {
log.Errorf("[Elasticsearch] Failed to create index: %v", err)
return err
}
log.Infof("[Elasticsearch] Index created successfully: %s", e.index)
return nil
}
// Retrieve dispatches the retrieval operation to the appropriate method based on retriever type
// Returns a slice of RetrieveResult and an error if the operation fails
func (e *elasticsearchRepository) Retrieve(ctx context.Context,
params typesLocal.RetrieveParams,
) ([]*typesLocal.RetrieveResult, error) {
log := logger.GetLogger(ctx)
log.Debugf("[Elasticsearch] Processing retrieval request of type: %s", params.RetrieverType)
// Route to appropriate retrieval method
switch params.RetrieverType {
case typesLocal.VectorRetrieverType:
return e.VectorRetrieve(ctx, params)
case typesLocal.KeywordsRetrieverType:
return e.KeywordsRetrieve(ctx, params)
}
err := fmt.Errorf("invalid retriever type: %v", params.RetrieverType)
log.Errorf("[Elasticsearch] %v", err)
return nil, err
}
// VectorRetrieve performs vector similarity search using cosine similarity
// Returns a slice of RetrieveResult containing matching documents
func (e *elasticsearchRepository) VectorRetrieve(ctx context.Context,
params typesLocal.RetrieveParams,
) ([]*typesLocal.RetrieveResult, error) {
log := logger.GetLogger(ctx)
log.Infof("[Elasticsearch] Vector retrieval: dim=%d, topK=%d, threshold=%.4f",
len(params.Embedding), params.TopK, params.Threshold)
filter := e.getBaseConds(params)
// Build script scoring query with cosine similarity
queryVectorJSON, err := json.Marshal(params.Embedding)
if err != nil {
log.Errorf("[Elasticsearch] Failed to marshal query vector: %v", err)
return nil, fmt.Errorf("failed to marshal query embedding: %w", err)
}
scoreSource := "cosineSimilarity(params.query_vector, 'embedding')"
minScore := float32(params.Threshold)
scriptScore := &types.ScriptScoreQuery{
Query: types.Query{Bool: &types.BoolQuery{Filter: filter}},
Script: types.Script{
Source: &scoreSource,
Params: map[string]json.RawMessage{
"query_vector": json.RawMessage(queryVectorJSON),
},
},
MinScore: &minScore,
}
// Exclude embedding field from source to reduce response size
sourceFilter := &types.SourceFilter{
Excludes: []string{"embedding"},
}
log.Debugf("[Elasticsearch] Executing vector search in index: %s", e.index)
// Execute search with minimum score threshold
response, err := e.client.Search().Index(e.index).Request(&search.Request{
Query: &types.Query{ScriptScore: scriptScore},
Size: &params.TopK,
Source_: sourceFilter,
}).Do(ctx)
if err != nil {
log.Errorf("[Elasticsearch] Vector search failed: %v", err)
return nil, err
}
// Process search results
var results []*typesLocal.IndexWithScore
for _, hit := range response.Hits.Hits {
var embedding *elasticsearchRetriever.VectorEmbeddingWithScore
if err := json.Unmarshal(hit.Source_, &embedding); err != nil {
log.Errorf("[Elasticsearch] Failed to unmarshal search result: %v", err)
return nil, err
}
embedding.Score = float64(*hit.Score_)
results = append(results,
elasticsearchRetriever.FromDBVectorEmbeddingWithScore(*hit.Id_, embedding, typesLocal.MatchTypeEmbedding))
}
if len(results) == 0 {
log.Warnf("[Elasticsearch] No vector matches found that meet threshold %.4f", params.Threshold)
} else {
log.Infof("[Elasticsearch] Vector retrieval found %d results", len(results))
log.Debugf("[Elasticsearch] Top result score: %.4f", results[0].Score)
}
return []*typesLocal.RetrieveResult{
{
Results: results,
RetrieverEngineType: typesLocal.ElasticsearchRetrieverEngineType,
RetrieverType: typesLocal.VectorRetrieverType,
Error: nil,
},
}, nil
}
// KeywordsRetrieve performs keyword-based search in document content
// Returns a slice of RetrieveResult containing matching documents
func (e *elasticsearchRepository) KeywordsRetrieve(ctx context.Context,
params typesLocal.RetrieveParams,
) ([]*typesLocal.RetrieveResult, error) {
log := logger.GetLogger(ctx)
log.Infof("[Elasticsearch] Performing keywords retrieval with query: %s, topK: %d", params.Query, params.TopK)
filter := e.getBaseConds(params)
// Build must conditions for content matching
must := []types.Query{
{Match: map[string]types.MatchQuery{"content": {Query: params.Query}}},
}
// Exclude embedding field from source to reduce response size
sourceFilter := &types.SourceFilter{
Excludes: []string{"embedding"},
}
log.Debugf("[Elasticsearch] Executing keyword search in index: %s", e.index)
response, err := e.client.Search().Index(e.index).Request(&search.Request{
Query: &types.Query{Bool: &types.BoolQuery{Filter: filter, Must: must}},
Size: &params.TopK,
Source_: sourceFilter,
}).Do(ctx)
if err != nil {
log.Errorf("[Elasticsearch] Keywords search failed: %v", err)
return nil, err
}
// Process search results
var results []*typesLocal.IndexWithScore
for _, hit := range response.Hits.Hits {
var embedding *elasticsearchRetriever.VectorEmbeddingWithScore
if err := json.Unmarshal(hit.Source_, &embedding); err != nil {
log.Errorf("[Elasticsearch] Failed to unmarshal search result: %v", err)
return nil, err
}
embedding.Score = float64(*hit.Score_)
results = append(results,
elasticsearchRetriever.FromDBVectorEmbeddingWithScore(*hit.Id_, embedding, typesLocal.MatchTypeKeywords),
)
}
if len(results) == 0 {
log.Warnf("[Elasticsearch] No keyword matches found for query: %s", params.Query)
} else {
log.Infof("[Elasticsearch] Keywords retrieval found %d results", len(results))
log.Debugf("[Elasticsearch] Top result score: %.4f", results[0].Score)
}
return []*typesLocal.RetrieveResult{
{
Results: results,
RetrieverEngineType: typesLocal.ElasticsearchRetrieverEngineType,
RetrieverType: typesLocal.KeywordsRetrieverType,
Error: nil,
},
}, nil
}
// CopyIndices 复制索引数据
func (e *elasticsearchRepository) CopyIndices(ctx context.Context,
sourceKnowledgeBaseID string,
sourceToTargetKBIDMap map[string]string,
sourceToTargetChunkIDMap map[string]string,
targetKnowledgeBaseID string,
dimension int,
knowledgeType string,
) error {
log := logger.GetLogger(ctx)
log.Infof(
"[Elasticsearch] Copying indices from source knowledge base %s to target knowledge base %s, count: %d",
sourceKnowledgeBaseID, targetKnowledgeBaseID, len(sourceToTargetChunkIDMap),
)
if len(sourceToTargetChunkIDMap) == 0 {
log.Warn("[Elasticsearch] Empty mapping, skipping copy")
return nil
}
// Build query parameters
params := typesLocal.RetrieveParams{
KnowledgeBaseIDs: []string{sourceKnowledgeBaseID},
}
// Build base query conditions
filter := e.getBaseConds(params)
// Set batch processing parameters
batchSize := 500
from := 0
totalCopied := 0
for {
// Execute pagination query
searchResponse, err := e.client.Search().Index(e.index).
Query(&types.Query{Bool: &types.BoolQuery{Filter: filter}}).
From(from).
Size(batchSize).
Do(ctx)
if err != nil {
log.Errorf("[Elasticsearch] Failed to query source index data: %v", err)
return err
}
hitsCount := len(searchResponse.Hits.Hits)
if hitsCount == 0 {
break
}
log.Infof("[Elasticsearch] Found %d source index data, batch start position: %d", hitsCount, from)
// Prepare index list for BatchSave
var indexInfoList []*typesLocal.IndexInfo
// Collect all embedding vector data for additionalParams
embeddingMap := make(map[string][]float32)
for _, hit := range searchResponse.Hits.Hits {
// Parse source document
var sourceDoc elasticsearchRetriever.VectorEmbedding
if err := json.Unmarshal(hit.Source_, &sourceDoc); err != nil {
log.Errorf("[Elasticsearch] Failed to parse source index data: %v", err)
continue
}
// Get mapped target chunk ID and knowledge ID
targetChunkID, ok := sourceToTargetChunkIDMap[sourceDoc.ChunkID]
if !ok {
log.Warnf("[Elasticsearch] Source chunk %s not found in target mapping, skipping", sourceDoc.ChunkID)
continue
}
targetKnowledgeID, ok := sourceToTargetKBIDMap[sourceDoc.KnowledgeID]
if !ok {
log.Warnf(
"[Elasticsearch] Source knowledge %s not found in target mapping, skipping",
sourceDoc.KnowledgeID,
)
continue
}
// Save embedding vector to embeddingMap
if len(sourceDoc.Embedding) > 0 {
embeddingMap[targetChunkID] = sourceDoc.Embedding
}
// Handle SourceID transformation for generated questions
// Generated questions have SourceID format: {chunkID}-{questionID}
// Regular chunks have SourceID == ChunkID
var targetSourceID string
if sourceDoc.SourceID == sourceDoc.ChunkID {
// Regular chunk, use targetChunkID as SourceID
targetSourceID = targetChunkID
} else if strings.HasPrefix(sourceDoc.SourceID, sourceDoc.ChunkID+"-") {
// This is a generated question, preserve the questionID part
questionID := strings.TrimPrefix(sourceDoc.SourceID, sourceDoc.ChunkID+"-")
targetSourceID = fmt.Sprintf("%s-%s", targetChunkID, questionID)
} else {
// For other complex scenarios, generate new unique SourceID
targetSourceID = uuid.New().String()
}
// Create new index information
indexInfo := &typesLocal.IndexInfo{
Content: sourceDoc.Content,
SourceID: targetSourceID,
SourceType: typesLocal.SourceType(sourceDoc.SourceType),
ChunkID: targetChunkID,
KnowledgeID: targetKnowledgeID,
KnowledgeBaseID: targetKnowledgeBaseID,
}
indexInfoList = append(indexInfoList, indexInfo)
totalCopied++
}
// Use BatchSave function to save index
if len(indexInfoList) > 0 {
// Add embedding vector to additional parameters
additionalParams := map[string]any{
"embedding": embeddingMap,
}
if err := e.BatchSave(ctx, indexInfoList, additionalParams); err != nil {
log.Errorf("[Elasticsearch] Failed to batch save index: %v", err)
return err
}
log.Infof("[Elasticsearch] Successfully copied batch data, batch size: %d, total copied: %d",
len(indexInfoList), totalCopied)
}
// Move to next batch
from += hitsCount
// If the number of returned records is less than the request size, it means the last page has been reached
if hitsCount < batchSize {
break
}
}
log.Infof("[Elasticsearch] Index copy completed, total copied: %d", totalCopied)
return nil
}
// BatchUpdateChunkEnabledStatus updates the enabled status of chunks in batch
func (e *elasticsearchRepository) BatchUpdateChunkEnabledStatus(
ctx context.Context,
chunkStatusMap map[string]bool,
) error {
log := logger.GetLogger(ctx)
if len(chunkStatusMap) == 0 {
log.Warnf("[Elasticsearch] Chunk status map is empty, skipping update")
return nil
}
log.Infof("[Elasticsearch] Batch updating chunk enabled status, count: %d", len(chunkStatusMap))
// Group chunks by enabled status for batch updates
enabledChunkIDs := make([]string, 0)
disabledChunkIDs := make([]string, 0)
for chunkID, enabled := range chunkStatusMap {
if enabled {
enabledChunkIDs = append(enabledChunkIDs, chunkID)
} else {
disabledChunkIDs = append(disabledChunkIDs, chunkID)
}
}
// Batch update enabled chunks using update_by_query
if len(enabledChunkIDs) > 0 {
query := types.NewQuery()
query.Bool = &types.BoolQuery{
Must: []types.Query{
{Terms: &types.TermsQuery{
TermsQuery: map[string]types.TermsQueryField{
e.idField("chunk_id"): enabledChunkIDs,
},
}},
},
}
source := "ctx._source.is_enabled = true"
lang := scriptlanguage.Painless
script := types.Script{
Source: &source,
Lang: &lang,
}
_, err := e.client.UpdateByQuery(e.index).Query(query).Script(&script).Do(ctx)
if err != nil {
log.Errorf("[Elasticsearch] Failed to update enabled chunks: %v", err)
return err
}
log.Infof("[Elasticsearch] Updated %d chunks to enabled", len(enabledChunkIDs))
}
// Batch update disabled chunks using update_by_query
if len(disabledChunkIDs) > 0 {
query := types.NewQuery()
query.Bool = &types.BoolQuery{
Must: []types.Query{
{Terms: &types.TermsQuery{
TermsQuery: map[string]types.TermsQueryField{
e.idField("chunk_id"): disabledChunkIDs,
},
}},
},
}
source := "ctx._source.is_enabled = false"
lang := scriptlanguage.Painless
script := types.Script{
Source: &source,
Lang: &lang,
}
_, err := e.client.UpdateByQuery(e.index).Query(query).Script(&script).Do(ctx)
if err != nil {
log.Errorf("[Elasticsearch] Failed to update disabled chunks: %v", err)
return err
}
log.Infof("[Elasticsearch] Updated %d chunks to disabled", len(disabledChunkIDs))
}
log.Infof("[Elasticsearch] Successfully batch updated chunk enabled status")
return nil
}
// BatchUpdateChunkTagID updates the tag ID of chunks in batch
func (e *elasticsearchRepository) BatchUpdateChunkTagID(
ctx context.Context,
chunkTagMap map[string]string,
) error {
log := logger.GetLogger(ctx)
if len(chunkTagMap) == 0 {
log.Warnf("[Elasticsearch] Chunk tag map is empty, skipping update")
return nil
}
log.Infof("[Elasticsearch] Batch updating chunk tag ID, count: %d", len(chunkTagMap))
// Group chunks by tag ID for batch updates
tagGroups := make(map[string][]string)
for chunkID, tagID := range chunkTagMap {
tagGroups[tagID] = append(tagGroups[tagID], chunkID)
}
// Batch update chunks for each tag ID using update_by_query
for tagID, chunkIDs := range tagGroups {
query := types.NewQuery()
query.Bool = &types.BoolQuery{
Must: []types.Query{
{Terms: &types.TermsQuery{
TermsQuery: map[string]types.TermsQueryField{
e.idField("chunk_id"): chunkIDs,
},
}},
},
}
source := "ctx._source.tag_id = params.tag_id"
lang := scriptlanguage.Painless
script := types.Script{
Source: &source,
Lang: &lang,
Params: map[string]json.RawMessage{
"tag_id": json.RawMessage(`"` + tagID + `"`),
},
}
_, err := e.client.UpdateByQuery(e.index).Query(query).Script(&script).Do(ctx)
if err != nil {
log.Errorf("[Elasticsearch] Failed to update chunks with tag_id %s: %v", tagID, err)
return err
}
log.Infof("[Elasticsearch] Updated %d chunks to tag_id=%s", len(chunkIDs), tagID)
}
log.Infof("[Elasticsearch] Successfully batch updated chunk tag ID")
return nil
}
@@ -0,0 +1,455 @@
package legacygraph
import (
"context"
"encoding/json"
"fmt"
"slices"
"strconv"
"strings"
neo4jrepo "github.com/Tencent/WeKnora/internal/application/repository/retriever/neo4j"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
"github.com/neo4j/neo4j-go-driver/v6/neo4j"
"gorm.io/gorm"
)
// Repository preserves source records while exposing verified native chunks to RAG.
type Repository struct {
db *gorm.DB
fallback interfaces.RetrieveGraphRepository
}
func NewRepository(db *gorm.DB, driver neo4j.Driver) *Repository {
r := &Repository{db: db}
if driver != nil {
r.fallback = neo4jrepo.NewNeo4jRepository(driver)
}
return r
}
func NewRetrieveRepository(db *gorm.DB, driver neo4j.Driver) interfaces.RetrieveGraphRepository {
return NewRepository(db, driver)
}
func (r *Repository) AddGraph(ctx context.Context, ns types.NameSpace, graphs []*types.GraphData) error {
return r.addNativeGraphs(ctx, ns, graphs, nil)
}
func (r *Repository) DelGraph(ctx context.Context, ns []types.NameSpace) error {
return r.deleteNativeGraphs(ctx, ns)
}
type Record struct {
ID int64 `json:"id" gorm:"column:legacy_id"`
Payload json.RawMessage `json:"payload" gorm:"column:payload"`
}
type Evidence struct {
Record
NativeKnowledgeID string `json:"native_knowledge_id"`
NativeChunkID string `json:"native_chunk_id"`
MappingState string `json:"mapping_state"`
}
type Overview struct {
Counts map[string]int64 `json:"counts"`
Sources int64 `json:"sources"`
MappedSources int64 `json:"mapped_sources"`
Nodes []Record `json:"nodes"`
Relations []Record `json:"relations"`
Evidence []Evidence `json:"evidence"`
NextCursor int64 `json:"next_cursor"`
}
func (r *Repository) scope(ctx context.Context, kb string) (*gorm.DB, error) {
tenant, ok := types.TenantIDFromContext(ctx)
if !ok || tenant == 0 || kb == "" {
return nil, fmt.Errorf("graph tenant and knowledge base are required")
}
return r.db.WithContext(ctx).Table("native_graph_records").Where("tenant_id = ? AND knowledge_base_id = ?", tenant, kb), nil
}
func (r *Repository) Browse(ctx context.Context, kb, query string, cursor, entity int64) (*Overview, error) {
scope, err := r.scope(ctx, kb)
if err != nil {
return nil, err
}
out := &Overview{Counts: map[string]int64{}, Nodes: []Record{}, Relations: []Record{}, Evidence: []Evidence{}}
var counts []struct {
SourceTable string
N int64
}
if err = scope.Session(&gorm.Session{}).Select("source_table, count(*) AS n").Group("source_table").Scan(&counts).Error; err != nil {
return nil, err
}
for _, c := range counts {
out.Counts[c.SourceTable] = c.N
}
tenant := types.MustTenantIDFromContext(ctx)
var coverage struct {
Sources int64
Mapped int64
}
// Overview coverage records migration-time matches; each returned link is revalidated below.
err = r.db.WithContext(ctx).Raw(`SELECT count(*) AS sources, count(native_chunk_id) AS mapped FROM native_graph_sources WHERE tenant_id=? AND knowledge_base_id=?`, tenant, kb).Scan(&coverage).Error
if err != nil {
return nil, err
}
out.Sources = coverage.Sources
out.MappedSources = coverage.Mapped
q := scope.Session(&gorm.Session{}).Select("legacy_id, payload").Where("source_table = ? AND legacy_id > ?", "knowledgeentity", cursor)
if entity > 0 {
q = q.Where("legacy_id = ?", entity)
} else if strings.TrimSpace(query) != "" {
q = q.Where("strpos(lower(payload->>'canonical_name'), lower(?)) > 0 OR strpos(lower(?), lower(payload->>'canonical_name')) > 0", query, query)
}
if err = q.Order("legacy_id").Limit(51).Find(&out.Nodes).Error; err != nil {
return nil, err
}
if len(out.Nodes) > 50 {
out.NextCursor = out.Nodes[49].ID
out.Nodes = out.Nodes[:50]
}
if err = r.projectNativeEntityAttributes(ctx, types.NameSpace{KnowledgeBase: kb}, out.Nodes, nil); err != nil {
return nil, err
}
if entity <= 0 {
return out, nil
}
err = r.db.WithContext(ctx).Raw(`SELECT r.legacy_id, r.payload || jsonb_build_object('from_name',a.payload->>'canonical_name','to_name',b.payload->>'canonical_name') AS payload FROM native_graph_records r LEFT JOIN native_graph_records a ON a.tenant_id=r.tenant_id AND a.knowledge_base_id=r.knowledge_base_id AND a.source_table='knowledgeentity' AND a.legacy_id=(r.payload->>'from_entity_id')::bigint LEFT JOIN native_graph_records b ON b.tenant_id=r.tenant_id AND b.knowledge_base_id=r.knowledge_base_id AND b.source_table='knowledgeentity' AND b.legacy_id=(r.payload->>'to_entity_id')::bigint WHERE r.tenant_id=? AND r.knowledge_base_id=? AND r.source_table='knowledgerelation' AND (r.payload->>'from_entity_id'=? OR r.payload->>'to_entity_id'=?) ORDER BY r.legacy_id LIMIT 100`, tenant, kb, strconv.FormatInt(entity, 10), strconv.FormatInt(entity, 10)).Scan(&out.Relations).Error
if err != nil {
return nil, err
}
err = r.db.WithContext(ctx).Raw(`SELECT r.legacy_id,r.payload, COALESCE(c.knowledge_id,'') AS native_knowledge_id, COALESCE(c.id,'') AS native_chunk_id, CASE WHEN c.id IS NOT NULL THEN s.mapping_state WHEN s.native_chunk_id IS NOT NULL THEN 'stale' ELSE COALESCE(s.mapping_state,'unmapped') END AS mapping_state FROM native_graph_records r LEFT JOIN native_graph_sources s ON s.tenant_id=r.tenant_id AND s.knowledge_base_id=r.knowledge_base_id AND s.source_chunk_id=r.payload->>'source_chunk_id' LEFT JOIN chunks c ON c.id=s.native_chunk_id AND c.tenant_id=s.tenant_id AND c.knowledge_base_id=s.knowledge_base_id AND c.deleted_at IS NULL AND c.is_enabled=true AND encode(sha256(convert_to(c.content,'UTF8')),'hex')=s.native_content_sha256 AND EXISTS (SELECT 1 FROM knowledges k WHERE k.id=c.knowledge_id AND k.tenant_id=c.tenant_id AND k.knowledge_base_id=c.knowledge_base_id AND k.deleted_at IS NULL) AND EXISTS (SELECT 1 FROM knowledge_bases kb WHERE kb.id=c.knowledge_base_id AND kb.tenant_id=c.tenant_id AND kb.deleted_at IS NULL) WHERE r.tenant_id=? AND r.knowledge_base_id=? AND r.source_table='knowledgegraphevidence' AND r.payload->>'subject_type'='entity' AND r.payload->>'subject_id'=? ORDER BY r.legacy_id LIMIT 100`, tenant, kb, strconv.FormatInt(entity, 10)).Scan(&out.Evidence).Error
return out, err
}
func (r *Repository) SearchNode(ctx context.Context, ns types.NameSpace, terms []string) (*types.GraphData, error) {
out := &types.GraphData{Node: []*types.GraphNode{}, Relation: []*types.GraphRelation{}}
scope, err := r.scope(ctx, ns.KnowledgeBase)
if err != nil {
return nil, err
}
terms = boundedTerms(terms)
if len(terms) == 0 {
return out, nil
}
clauses := make([]string, 0, len(terms))
args := make([]any, 0, len(terms)*2)
for _, term := range terms {
clauses = append(clauses, "(strpos(lower(payload->>'canonical_name'),lower(?))>0 OR strpos(lower(?),lower(payload->>'canonical_name'))>0)")
args = append(args, term, term)
}
var nodes []Record
err = scope.Session(&gorm.Session{}).Select("legacy_id,payload").Where("source_table='knowledgeentity' AND payload->>'status'='active'").Where(strings.Join(clauses, " OR "), args...).Order("legacy_id").Limit(16).Find(&nodes).Error
if err != nil {
return nil, err
}
tenant := types.MustTenantIDFromContext(ctx)
seedIDs := []string{}
seenIDs := map[int64]bool{}
for _, node := range nodes {
seedIDs = append(seedIDs, strconv.FormatInt(node.ID, 10))
seenIDs[node.ID] = true
}
var relations []Record
if len(seedIDs) > 0 {
q, _ := r.scope(ctx, ns.KnowledgeBase)
err = q.Select("legacy_id,payload").Where("source_table='knowledgerelation' AND payload->>'status'='active' AND (payload->>'from_entity_id' IN ? OR payload->>'to_entity_id' IN ?)", seedIDs, seedIDs).Order("legacy_id").Limit(32).Find(&relations).Error
if err != nil {
return nil, err
}
neighborIDs := []int64{}
for _, rel := range relations {
var edge struct {
From int64 `json:"from_entity_id"`
To int64 `json:"to_entity_id"`
}
if err = json.Unmarshal(rel.Payload, &edge); err != nil {
return nil, err
}
for _, id := range []int64{edge.From, edge.To} {
if !seenIDs[id] && len(neighborIDs) < 16 {
seenIDs[id] = true
neighborIDs = append(neighborIDs, id)
}
}
}
if len(neighborIDs) > 0 {
var neighbors []Record
q, _ = r.scope(ctx, ns.KnowledgeBase)
if err = q.Select("legacy_id,payload").Where("source_table='knowledgeentity' AND payload->>'status'='active' AND legacy_id IN ?", neighborIDs).Order("legacy_id").Find(&neighbors).Error; err != nil {
return nil, err
}
nodes = append(nodes, neighbors...)
}
}
chunksByEntity := map[int64][]string{}
chunksByRelation := map[int64][]string{}
if len(nodes) > 0 {
type subject struct {
ID int64 `json:"subject_id"`
Type string `json:"subject_type"`
}
ids := make([]subject, 0, len(nodes)+len(relations))
for _, node := range nodes {
ids = append(ids, subject{node.ID, "entity"})
}
for _, relation := range relations {
ids = append(ids, subject{relation.ID, "relation"})
}
encodedIDs, err := json.Marshal(ids)
if err != nil {
return nil, err
}
var matches []struct {
SubjectID int64
SubjectType string
ChunkID string
}
// One bounded database request, with the evidence index and four-chunk cap per entity.
query := `SELECT seed.subject_id,seed.subject_type,matched.chunk_id FROM jsonb_to_recordset(?::jsonb) AS seed(subject_id bigint,subject_type text) CROSS JOIN LATERAL (
SELECT DISTINCT c.id AS chunk_id FROM native_graph_records e
JOIN native_graph_sources s ON s.tenant_id=e.tenant_id AND s.knowledge_base_id=e.knowledge_base_id AND s.source_chunk_id=e.payload->>'source_chunk_id'
JOIN chunks c ON c.id=s.native_chunk_id AND c.tenant_id=s.tenant_id AND c.knowledge_base_id=s.knowledge_base_id AND c.deleted_at IS NULL AND c.is_enabled=true AND encode(sha256(convert_to(c.content,'UTF8')),'hex')=s.native_content_sha256
JOIN knowledges k ON k.id=c.knowledge_id AND k.tenant_id=c.tenant_id AND k.knowledge_base_id=c.knowledge_base_id AND k.deleted_at IS NULL
JOIN knowledge_bases kb ON kb.id=c.knowledge_base_id AND kb.tenant_id=c.tenant_id AND kb.deleted_at IS NULL
WHERE e.tenant_id=? AND e.knowledge_base_id=? AND e.source_table='knowledgegraphevidence' AND e.payload->>'subject_type'=seed.subject_type AND e.payload->>'subject_id'=seed.subject_id::text AND e.payload->>'status'='active'`
args := []any{string(encodedIDs), tenant, ns.KnowledgeBase}
if ns.Knowledge != "" {
query += " AND c.knowledge_id=?"
args = append(args, ns.Knowledge)
}
query += " ORDER BY c.id LIMIT 4) matched"
if err = r.db.WithContext(ctx).Raw(query, args...).Scan(&matches).Error; err != nil {
return nil, err
}
for _, match := range matches {
if match.SubjectType == "entity" {
chunksByEntity[match.SubjectID] = append(chunksByEntity[match.SubjectID], match.ChunkID)
} else {
chunksByRelation[match.SubjectID] = append(chunksByRelation[match.SubjectID], match.ChunkID)
}
}
}
if err = r.projectNativeEntityAttributes(ctx, ns, nodes, chunksByEntity); err != nil {
return nil, err
}
names := map[int64]string{}
for _, node := range nodes {
var entity struct {
Name string `json:"canonical_name"`
Type string `json:"entity_type"`
Attributes []string `json:"attributes"`
}
if err = json.Unmarshal(node.Payload, &entity); err != nil {
return nil, err
}
chunks := chunksByEntity[node.ID]
if len(chunks) > 0 {
names[node.ID] = entity.Name
attributes := []string{entity.Type}
identity := fmt.Sprintf("legacy_entity:%d", node.ID)
for _, attribute := range entity.Attributes {
if attribute != entity.Type && attribute != identity {
attributes = append(attributes, attribute)
}
}
attributes = append(attributes, identity)
out.Node = append(out.Node, &types.GraphNode{Name: entity.Name, Chunks: chunks, Attributes: attributes})
}
}
for _, rel := range relations {
var edge struct {
From int64 `json:"from_entity_id"`
To int64 `json:"to_entity_id"`
Type string `json:"relation_type"`
Origin string `json:"origin"`
}
if err = json.Unmarshal(rel.Payload, &edge); err != nil {
return nil, err
}
if edge.Origin == nativeOrigin && len(chunksByRelation[rel.ID]) == 0 {
continue
}
if names[edge.From] != "" && names[edge.To] != "" {
out.Relation = append(out.Relation, &types.GraphRelation{Node1: names[edge.From], Node2: names[edge.To], Type: edge.Type})
}
}
if r.fallback != nil {
fallback, err := r.fallback.SearchNode(ctx, ns, terms)
if err != nil {
return nil, err
}
if fallback != nil {
out.Node = append(out.Node, fallback.Node...)
out.Relation = append(out.Relation, fallback.Relation...)
}
}
return out, nil
}
// Attribute ownership is checked independently of the four returned chunks:
// an entity may still have a bare reference after its attribute source expires.
func (r *Repository) projectNativeEntityAttributes(ctx context.Context, ns types.NameSpace, nodes []Record, chunks map[int64][]string) error {
type subject struct {
ID int64 `json:"subject_id"`
Knowledge string `json:"knowledge_id"`
Type string `json:"entity_type"`
Attributes json.RawMessage `json:"attributes"`
}
seeds := []subject{}
payloads := map[int64]map[string]json.RawMessage{}
for _, node := range nodes {
var entity struct {
Origin string `json:"origin"`
Knowledge string `json:"native_knowledge_id"`
Type string `json:"entity_type"`
Attributes json.RawMessage `json:"attributes"`
}
if err := json.Unmarshal(node.Payload, &entity); err != nil {
return err
}
if entity.Origin != nativeOrigin {
continue
}
var payload map[string]json.RawMessage
if err := json.Unmarshal(node.Payload, &payload); err != nil {
return err
}
payloads[node.ID] = payload
seeds = append(seeds, subject{node.ID, entity.Knowledge, entity.Type, entity.Attributes})
}
if len(seeds) == 0 {
return nil
}
encoded, err := json.Marshal(seeds)
if err != nil {
return err
}
var matches []struct {
SubjectID int64
Snapshots json.RawMessage
Chunks json.RawMessage
AttributeChunk string
}
query := `SELECT seed.subject_id,jsonb_agg(DISTINCT e.payload->'entity_attributes') AS snapshots,
to_jsonb((array_agg(DISTINCT c.id ORDER BY c.id))[1:4]) AS chunks,
COALESCE(min(c.id) FILTER (WHERE e.payload->'entity_attributes'->>0 IS NOT NULL
AND (e.payload->'entity_attributes'=seed.attributes OR e.payload->'entity_attributes'->>0=seed.entity_type)),'') AS attribute_chunk
FROM jsonb_to_recordset(?::jsonb) AS seed(subject_id bigint,knowledge_id text,entity_type text,attributes jsonb)
JOIN native_graph_records e ON e.payload->>'subject_id'=seed.subject_id::text AND e.payload->>'native_knowledge_id'=seed.knowledge_id
JOIN native_graph_sources s ON s.tenant_id=e.tenant_id AND s.knowledge_base_id=e.knowledge_base_id AND s.source_chunk_id=e.payload->>'source_chunk_id'
JOIN chunks c ON c.id=s.native_chunk_id AND c.tenant_id=s.tenant_id AND c.knowledge_base_id=s.knowledge_base_id
AND s.native_knowledge_id=c.knowledge_id AND c.deleted_at IS NULL AND c.is_enabled=true
AND encode(sha256(convert_to(c.content,'UTF8')),'hex')=s.native_content_sha256
JOIN knowledges k ON k.id=c.knowledge_id AND k.tenant_id=c.tenant_id AND k.knowledge_base_id=c.knowledge_base_id AND k.deleted_at IS NULL
JOIN knowledge_bases kb ON kb.id=c.knowledge_base_id AND kb.tenant_id=c.tenant_id AND kb.deleted_at IS NULL
WHERE e.tenant_id=? AND e.knowledge_base_id=? AND e.source_table='knowledgegraphevidence' AND e.payload->>'subject_type'='entity'
AND e.payload->>'origin'=? AND e.payload->>'status'='active' AND s.mapping_state='native_extraction'
AND e.payload->>'source_content_hash'=s.native_content_sha256 AND e.source_snapshot_sha256=s.native_content_sha256
AND s.source_content_hash=s.native_content_sha256`
args := []any{string(encoded), types.MustTenantIDFromContext(ctx), ns.KnowledgeBase, nativeOrigin}
if ns.Knowledge != "" {
query += " AND c.knowledge_id=?"
args = append(args, ns.Knowledge)
}
query += " GROUP BY seed.subject_id"
if err = r.db.WithContext(ctx).Raw(query, args...).Scan(&matches).Error; err != nil {
return err
}
supported := map[int64][]string{}
attributeChunks := map[int64]string{}
verifiedChunks := map[int64][]string{}
for _, match := range matches {
supported[match.SubjectID] = compatibleNativeAttributeSnapshots(match.Snapshots)
attributeChunks[match.SubjectID] = match.AttributeChunk
var current []string
if err = json.Unmarshal(match.Chunks, &current); err != nil {
return err
}
verifiedChunks[match.SubjectID] = current
}
for i, node := range nodes {
payload, native := payloads[node.ID]
if !native {
continue
}
if chunks != nil {
chunks[node.ID] = verifiedChunks[node.ID]
}
attributes := []string{}
entityType := "concept"
if snapshot := supported[node.ID]; len(snapshot) > 0 {
var stored []string
if err := json.Unmarshal(payload["attributes"], &stored); err == nil && slices.Equal(stored, snapshot) {
attributes = stored
entityType = snapshot[0]
}
var storedType string
if err := json.Unmarshal(payload["entity_type"], &storedType); err == nil && storedType == snapshot[0] {
entityType = storedType
}
}
if chunk := attributeChunks[node.ID]; chunks != nil && chunk != "" && (len(attributes) > 0 || entityType != "concept") {
verified := []string{chunk}
for _, prior := range chunks[node.ID] {
if prior != chunk && len(verified) < 4 {
verified = append(verified, prior)
}
}
chunks[node.ID] = verified
}
payload["attributes"], _ = json.Marshal(attributes)
payload["entity_type"], _ = json.Marshal(entityType)
if nodes[i].Payload, err = json.Marshal(payload); err != nil {
return err
}
}
return nil
}
func compatibleNativeAttributeSnapshots(raw json.RawMessage) []string {
var snapshots []json.RawMessage
if err := json.Unmarshal(raw, &snapshots); err != nil {
return nil
}
var attributes []string
for _, snapshot := range snapshots {
var values []*string
if err := json.Unmarshal(snapshot, &values); err != nil {
return nil
}
if len(values) == 0 {
continue
}
current := make([]string, len(values))
for i, value := range values {
if value == nil {
return nil
}
current[i] = *value
}
if len(attributes) > 0 && !slices.Equal(attributes, current) {
return nil
}
attributes = current
}
return attributes
}
func boundedTerms(terms []string) []string {
out := []string{}
seen := map[string]bool{}
for _, term := range terms {
term = strings.TrimSpace(term)
if term == "" || seen[term] {
continue
}
if len([]rune(term)) > 200 {
term = string([]rune(term)[:200])
}
seen[term] = true
out = append(out, term)
if len(out) == 8 {
break
}
}
return out
}
@@ -0,0 +1,109 @@
package legacygraph
import (
"context"
"fmt"
"reflect"
"testing"
"github.com/Tencent/WeKnora/internal/types"
)
func TestSearchNodePreservesStoredAttributes(t *testing.T) {
for _, origin := range []string{"native", "legacy"} {
for _, stored := range []bool{true, false} {
t.Run(fmt.Sprintf("%s/attributes=%t", origin, stored), func(t *testing.T) {
tx, r, ctx := nativeWriteFixture(t)
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
graph := testGraph("chunk")
graph.Node[0].Attributes = []string{"concept", "Not a medicine", "No treatment, nursing or health-care effect", "Evidence keeps spaces\nand line breaks", "concept"}
if err := r.AddGraph(ctx, ns, []*types.GraphData{graph}); err != nil {
t.Fatal(err)
}
if origin == "legacy" {
if err := tx.Exec(`UPDATE native_graph_records SET payload=payload-'origin' WHERE source_table='knowledgeentity'`).Error; err != nil {
t.Fatal(err)
}
if err := tx.Exec(`UPDATE native_graph_sources SET mapping_state='exact_content_hash'`).Error; err != nil {
t.Fatal(err)
}
}
if !stored {
if err := tx.Exec(`UPDATE native_graph_records SET payload=payload-'attributes' WHERE source_table='knowledgeentity'`).Error; err != nil {
t.Fatal(err)
}
}
id, _ := nativeID("doc", "entity", "trust")
want := []string{"concept"}
if stored {
want = append(want, graph.Node[0].Attributes[1:4]...)
}
want = append(want, fmt.Sprintf("legacy_entity:%d", id))
// The same record must preserve attributes as either a seed or a neighbor.
for _, term := range []string{"Trust", "Learning"} {
found, err := r.SearchNode(ctx, ns, []string{term})
if err != nil {
t.Fatal(err)
}
var got *types.GraphNode
for _, node := range found.Node {
if node.Name == "Trust" {
got = node
}
}
if got == nil || !reflect.DeepEqual(got.Attributes, want) || !reflect.DeepEqual(got.Chunks, []string{"chunk"}) {
t.Fatalf("%s lost source attributes or compatibility metadata: got=%#v want=%#v", term, got, want)
}
}
})
}
}
}
func TestSearchNodeAttributesRespectSourceAvailabilityAndOwnership(t *testing.T) {
tx, r, ctx := nativeWriteFixture(t)
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
graph := testGraph("chunk")
graph.Node[0].Attributes = []string{"concept", "Not a medicine", "No treatment, nursing or health-care effect"}
if err := r.AddGraph(ctx, ns, []*types.GraphData{graph}); err != nil {
t.Fatal(err)
}
for _, scope := range []struct {
name string
tenant uint64
ns types.NameSpace
}{
{"other tenant", 8, ns},
{"other knowledge base", 7, types.NameSpace{KnowledgeBase: "other-kb"}},
{"other document", 7, types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc2"}},
} {
t.Run(scope.name, func(t *testing.T) {
otherCtx := context.WithValue(ctx, types.TenantIDContextKey, scope.tenant)
got, err := r.SearchNode(otherCtx, scope.ns, []string{"Trust"})
if err != nil || len(got.Node) != 0 || len(got.Relation) != 0 {
t.Fatalf("out-of-scope attributes returned: %#v %v", got, err)
}
})
}
for _, change := range []struct{ name, sql string }{
{"edited source", `UPDATE chunks SET content=content || ' changed' WHERE id='chunk'`},
{"disabled source", `UPDATE chunks SET is_enabled=false WHERE id='chunk'`},
{"deleted source", `UPDATE chunks SET deleted_at=now() WHERE id='chunk'`},
{"deleted document", `UPDATE knowledges SET deleted_at=now() WHERE id='doc' AND tenant_id=7`},
{"deleted knowledge base", `UPDATE knowledge_bases SET deleted_at=now() WHERE id='kb' AND tenant_id=7`},
} {
t.Run(change.name, func(t *testing.T) {
if err := tx.SavePoint("attribute_source_check").Error; err != nil {
t.Fatal(err)
}
defer tx.RollbackTo("attribute_source_check")
if err := tx.Exec(change.sql).Error; err != nil {
t.Fatal(err)
}
got, err := r.SearchNode(ctx, ns, []string{"Trust"})
if err != nil || len(got.Node) != 0 || len(got.Relation) != 0 {
t.Fatalf("unavailable source attributes returned: %#v %v", got, err)
}
})
}
}
@@ -0,0 +1,450 @@
package legacygraph
import (
"context"
"encoding/json"
"fmt"
"os"
"slices"
"strings"
"testing"
"time"
"github.com/Tencent/WeKnora/internal/types"
"gorm.io/driver/postgres"
"gorm.io/gorm"
)
func TestMissingGraphScope(t *testing.T) {
r := NewRepository(nil, nil)
if _, err := r.Browse(context.Background(), "kb", "", 0, 0); err == nil {
t.Fatal("missing tenant accepted")
}
ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(10059))
if _, err := r.SearchNode(ctx, types.NameSpace{}, []string{"term"}); err == nil {
t.Fatal("missing KB accepted")
}
}
func TestBoundedTerms(t *testing.T) {
terms := boundedTerms([]string{"", " ", "a", "a", strings.Repeat("中", 240), "b", "c", "d", "e", "f", "g", "h"})
if len(terms) != 8 || terms[0] != "a" || len([]rune(terms[1])) != 200 {
t.Fatalf("unbounded terms: %#v", terms)
}
}
func TestSearchNodeKeepsGenericSourcesIsolated(t *testing.T) {
tx, r, ctx := nativeWriteFixture(t)
if err := tx.Exec(`INSERT INTO chunks VALUES('sibling',7,'kb','doc','Evidence about trust and learning.',NULL,true)`).Error; err != nil {
t.Fatal(err)
}
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
for chunk, target := range map[string]string{"chunk": "First target", "sibling": "Second target"} {
graph := &types.GraphData{
Node: []*types.GraphNode{
{Name: "本品", Chunks: []string{chunk}},
{Name: target, Chunks: []string{chunk}},
},
Relation: []*types.GraphRelation{{Node1: "本品", Node2: target, Type: "关联"}},
}
if err := r.AddGraph(ctx, ns, []*types.GraphData{graph}); err != nil {
t.Fatal(err)
}
}
for chunk, target := range map[string]string{"chunk": "First target", "sibling": "Second target"} {
graph, err := r.SearchNode(ctx, ns, []string{target})
if err != nil || len(graph.Node) != 2 || len(graph.Relation) != 1 {
t.Fatalf("unexpected single-product result: %#v %v", graph, err)
}
for _, node := range graph.Node {
if len(node.Chunks) != 1 || node.Chunks[0] != chunk {
t.Fatalf("same-name neighbor expanded another product: %#v", node)
}
}
if edge := graph.Relation[0]; edge.Node1 != "本品" || edge.Node2 != target {
t.Fatalf("unexpected relation projection: %#v", edge)
}
}
graph, err := r.SearchNode(ctx, ns, []string{"本品"})
if err != nil || len(graph.Node) != 4 || len(graph.Relation) != 2 {
t.Fatalf("same-name nodes were collapsed: %#v %v", graph, err)
}
sources := map[string]bool{}
identities := map[string]bool{}
for _, node := range graph.Node {
if node.Name != "本品" {
continue
}
if len(node.Chunks) != 1 || len(node.Attributes) != 2 {
t.Fatalf("generic source identity was lost: %#v", node)
}
sources[node.Chunks[0]] = true
identities[node.Attributes[1]] = true
}
if len(sources) != 2 || !sources["chunk"] || !sources["sibling"] || len(identities) != 2 {
t.Fatalf("generic products share retrieval identity: sources=%v identities=%v", sources, identities)
}
}
func assertReaderTrustAttributes(t *testing.T, r *Repository, ctx context.Context, want []string) {
t.Helper()
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
id, _ := nativeID("doc", "entity", "trust")
wantType := "concept"
if len(want) > 0 {
wantType = want[0]
}
wantGraph := append([]string{wantType}, want...)
if len(want) > 0 {
wantGraph = wantGraph[1:]
}
wantGraph = append(wantGraph, fmt.Sprintf("legacy_entity:%d", id))
for _, term := range []string{"Trust", "Learning"} {
graph, err := r.SearchNode(ctx, ns, []string{term})
if err != nil {
t.Fatal(err)
}
var found *types.GraphNode
for _, node := range graph.Node {
if node.Name == "Trust" {
found = node
}
}
if found == nil || !slices.Equal(found.Attributes, wantGraph) {
t.Fatalf("%s returned unsupported native attributes: got=%#v want=%v", term, found, wantGraph)
}
}
for _, entityID := range []int64{0, id} {
page, err := r.Browse(ctx, "kb", "", 0, entityID)
if err != nil {
t.Fatal(err)
}
found := false
for _, node := range page.Nodes {
if node.ID != id {
continue
}
found = true
var payload struct {
Attributes []string `json:"attributes"`
Type string `json:"entity_type"`
}
if err := json.Unmarshal(node.Payload, &payload); err != nil {
t.Fatal(err)
}
if !slices.Equal(payload.Attributes, want) || payload.Type != wantType {
t.Fatalf("browse returned unsupported native attributes: %#v; want %v / %s", payload, want, wantType)
}
}
if !found {
t.Fatal("browse lost the entity's surviving bare reference")
}
}
}
func TestNativeReadersRevalidateSharedAttributeSources(t *testing.T) {
for name, invalidate := range map[string]string{
"changed source": `UPDATE chunks SET content=content || ' changed' WHERE id='chunk' AND tenant_id=7`,
"disabled source": `UPDATE chunks SET is_enabled=false WHERE id='chunk' AND tenant_id=7`,
"deleted source": `UPDATE chunks SET deleted_at=now() WHERE id='chunk' AND tenant_id=7`,
"moved source document": `UPDATE chunks SET knowledge_id='doc2' WHERE id='chunk' AND tenant_id=7`,
"wrong evidence hash": `UPDATE native_graph_records SET payload=jsonb_set(payload,'{source_content_hash}',to_jsonb(repeat('a',64))) WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
"wrong row hash": `UPDATE native_graph_records SET source_snapshot_sha256=repeat('a',64) WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
"inactive evidence": `UPDATE native_graph_records SET payload=jsonb_set(payload,'{status}','"inactive"'::jsonb) WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
"wrong origin": `UPDATE native_graph_records SET payload=jsonb_set(payload,'{origin}','"foreign"'::jsonb) WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
"wrong tenant": `UPDATE native_graph_records SET tenant_id=8 WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
"wrong knowledge base": `UPDATE native_graph_records SET knowledge_base_id='other-kb' WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
"wrong document": `UPDATE native_graph_records SET payload=jsonb_set(payload,'{native_knowledge_id}','"doc2"'::jsonb) WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
"wrong mapping document": `UPDATE native_graph_sources SET native_knowledge_id='doc2' WHERE source_chunk_id='native:chunk'`,
"wrong mapping hash": `UPDATE native_graph_sources SET source_content_hash=repeat('a',64) WHERE source_chunk_id='native:chunk'`,
"inactive mapping": `UPDATE native_graph_sources SET mapping_state='stale' WHERE source_chunk_id='native:chunk'`,
"missing old snapshot": `UPDATE native_graph_records SET payload=payload-'entity_attributes' WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
"null snapshot": `UPDATE native_graph_records SET payload=jsonb_set(payload,'{entity_attributes}','null'::jsonb) WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk' AND payload->>'subject_type'='entity'`,
"null attribute": `UPDATE native_graph_records SET payload=jsonb_set(payload,'{entity_attributes}','[null]'::jsonb) WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk' AND payload->>'subject_type'='entity'`,
"conflicting snapshot": `UPDATE native_graph_records SET payload=jsonb_set(payload,'{entity_attributes}','["contradiction"]'::jsonb) WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:sibling' AND payload->>'subject_type'='entity'`,
} {
t.Run(name, func(t *testing.T) {
tx, r, ctx := nativeWriteFixture(t)
if err := tx.Exec(`INSERT INTO chunks VALUES('sibling',7,'kb','doc','Another source about trust and learning.',NULL,true)`).Error; err != nil {
t.Fatal(err)
}
owner, bare := testGraph("chunk"), testGraph("sibling")
owner.Node[0].Attributes = []string{"restriction", "Cannot be taken orally"}
bare.Node[0].Attributes = nil
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
if err := r.AddGraph(ctx, ns, []*types.GraphData{owner, bare}); err != nil {
t.Fatal(err)
}
assertReaderTrustAttributes(t, r, ctx, owner.Node[0].Attributes)
if err := tx.Exec(invalidate).Error; err != nil {
t.Fatal(err)
}
before := nativeGraphWriteState(t, tx)
assertReaderTrustAttributes(t, r, ctx, nil)
if after := nativeGraphWriteState(t, tx); after != before {
t.Fatal("reader projection changed stored graph records or mappings")
}
})
}
}
func TestNativeReadersKeepVerifiedAttributesBeyondChunkCap(t *testing.T) {
tx, r, ctx := nativeWriteFixture(t)
if err := tx.Exec(`INSERT INTO chunks SELECT id,7,'kb','doc','Another source about trust and learning.',NULL,true FROM unnest(ARRAY['aa','ab','ac','ad']) id`).Error; err != nil {
t.Fatal(err)
}
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
owner := testGraph("chunk")
owner.Node[0].Attributes = []string{"restriction", "Cannot be taken orally"}
graphs := []*types.GraphData{owner}
for _, id := range []string{"aa", "ab", "ac", "ad"} {
bare := testGraph(id)
bare.Node[0].Attributes = nil
graphs = append(graphs, bare)
}
if err := r.AddGraph(ctx, ns, graphs); err != nil {
t.Fatal(err)
}
// An old bare source cannot erase attributes proved by an explicit source.
if err := tx.Exec(`UPDATE native_graph_records SET payload=payload-'entity_attributes' WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:aa'`).Error; err != nil {
t.Fatal(err)
}
assertReaderTrustAttributes(t, r, ctx, owner.Node[0].Attributes)
graph, err := r.SearchNode(ctx, ns, []string{"Trust"})
if err != nil {
t.Fatal(err)
}
for _, node := range graph.Node {
if len(node.Chunks) != 4 || (node.Name == "Trust" && !slices.Contains(node.Chunks, "chunk")) {
t.Fatalf("attribute owner was omitted from the bounded references: %#v", node)
}
}
if err := tx.Exec(`UPDATE chunks SET is_enabled=false WHERE id='chunk' AND tenant_id=7`).Error; err != nil {
t.Fatal(err)
}
assertReaderTrustAttributes(t, r, ctx, nil)
}
func TestNativeReadersRespectCopiedSourceMappings(t *testing.T) {
tx, r, ctx := nativeWriteFixture(t)
graph := testGraph("chunk")
graph.Node[0].Attributes = []string{"restriction", "Cannot be taken orally"}
if err := r.AddGraph(ctx, types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}, []*types.GraphData{graph}); err != nil {
t.Fatal(err)
}
for _, sql := range []string{
`INSERT INTO knowledges VALUES('copy-doc',7,'other-kb','copy.md',NULL)`,
`INSERT INTO chunks VALUES('copy-chunk',7,'other-kb','copy-doc','Evidence about trust and learning.',NULL,true)`,
`INSERT INTO native_graph_records SELECT tenant_id,'other-kb',source_table,legacy_id,graph_version_id,payload,source_snapshot_sha256,imported_at FROM native_graph_records WHERE tenant_id=7 AND knowledge_base_id='kb'`,
`INSERT INTO native_graph_sources SELECT tenant_id,'other-kb',source_chunk_id,source_file,source_content_hash,excerpt,'copy-doc','copy-chunk',native_content_sha256,mapping_state FROM native_graph_sources WHERE tenant_id=7 AND knowledge_base_id='kb'`,
`UPDATE chunks SET is_enabled=false WHERE id='chunk' AND tenant_id=7`,
} {
if err := tx.Exec(sql).Error; err != nil {
t.Fatal(err)
}
}
ns := types.NameSpace{KnowledgeBase: "other-kb", Knowledge: "copy-doc"}
id, _ := nativeID("doc", "entity", "trust")
found, err := r.SearchNode(ctx, ns, []string{"Trust"})
if err != nil || len(found.Node) != 2 {
t.Fatalf("copied native graph was not retrieved: %#v %v", found, err)
}
for _, node := range found.Node {
if !slices.Equal(node.Chunks, []string{"copy-chunk"}) {
t.Fatalf("copied graph retained an original chunk reference: %#v", node)
}
if node.Name == "Trust" && !slices.Equal(node.Attributes, []string{"restriction", "Cannot be taken orally", fmt.Sprintf("legacy_entity:%d", id)}) {
t.Fatalf("copied source snapshot lost its attributes: %#v", node)
}
}
wrong, err := r.SearchNode(ctx, types.NameSpace{KnowledgeBase: "other-kb", Knowledge: "doc"}, []string{"Trust"})
if err != nil || len(wrong.Node) != 0 {
t.Fatalf("copied graph leaked through the original document scope: %#v %v", wrong, err)
}
page, err := r.Browse(ctx, "other-kb", "", 0, id)
if err != nil || len(page.Nodes) != 1 {
t.Fatalf("copied entity was not browsable: %#v %v", page, err)
}
var payload struct {
Attributes []string `json:"attributes"`
}
if err := json.Unmarshal(page.Nodes[0].Payload, &payload); err != nil || !slices.Equal(payload.Attributes, graph.Node[0].Attributes) {
t.Fatalf("copied browse attributes were not revalidated: %#v %v", payload, err)
}
}
// Fixture mutations are confined to a rolled-back transaction.
func TestAcceptancePreservedGraph(t *testing.T) {
dsn := os.Getenv("GRAPH_TEST_DSN")
if dsn == "" {
t.Skip("GRAPH_TEST_DSN is required for acceptance integration")
}
db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
if err != nil {
t.Fatal(err)
}
tx := db.Begin()
defer tx.Rollback()
if tx.Error != nil {
t.Fatal(tx.Error)
}
ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(10059))
r := NewRepository(tx, nil)
var fixture struct {
KB string
Entity int64
Name string
Chunk string
Knowledge string
}
err = tx.Raw(`SELECT r.knowledge_base_id AS kb,r.legacy_id AS entity,r.payload->>'canonical_name' AS name,s.native_chunk_id AS chunk,s.native_knowledge_id AS knowledge FROM native_graph_records r JOIN native_graph_records e ON e.tenant_id=r.tenant_id AND e.knowledge_base_id=r.knowledge_base_id AND e.source_table='knowledgegraphevidence' AND e.payload->>'subject_type'='entity' AND e.payload->>'subject_id'=r.legacy_id::text JOIN native_graph_sources s ON s.tenant_id=e.tenant_id AND s.knowledge_base_id=e.knowledge_base_id AND s.source_chunk_id=e.payload->>'source_chunk_id' WHERE r.tenant_id=10059 AND r.source_table='knowledgeentity' AND r.knowledge_base_id='16c3e59c-4bb6-41cf-ae0a-6d0941b3d032' AND s.native_chunk_id IS NOT NULL ORDER BY r.legacy_id LIMIT 1`).Scan(&fixture).Error
if err != nil || fixture.Chunk == "" {
t.Fatalf("missing mapped fixture: %v", err)
}
overview, err := r.Browse(ctx, fixture.KB, "", 0, fixture.Entity)
if err != nil {
t.Fatal(err)
}
var preservedEntities, preservedMappings int64
if err := tx.Raw(`SELECT count(*) FROM native_graph_records WHERE tenant_id=10059 AND knowledge_base_id=? AND source_table='knowledgeentity' AND payload->>'origin' IS DISTINCT FROM 'native_extraction_v1'`, fixture.KB).Scan(&preservedEntities).Error; err != nil {
t.Fatal(err)
}
if err := tx.Raw(`SELECT count(*) FROM native_graph_sources WHERE tenant_id=10059 AND knowledge_base_id=? AND mapping_state<>'native_extraction' AND native_chunk_id IS NOT NULL`, fixture.KB).Scan(&preservedMappings).Error; err != nil {
t.Fatal(err)
}
if len(overview.Nodes) != 1 || len(overview.Evidence) == 0 || preservedEntities != 2 || preservedMappings != 3 || overview.Counts["knowledgeentity"] < preservedEntities || overview.MappedSources < preservedMappings {
t.Fatalf("preserved graph changed: entities=%d mappings=%d counts=%v", preservedEntities, preservedMappings, overview.Counts)
}
found := false
for _, e := range overview.Evidence {
if e.NativeChunkID == fixture.Chunk {
found = true
}
}
if !found {
t.Fatal("mapped evidence was not returned")
}
graph, err := r.SearchNode(ctx, types.NameSpace{KnowledgeBase: fixture.KB, Knowledge: fixture.Knowledge}, []string{fixture.Name})
if err != nil {
t.Fatal(err)
}
if len(graph.Node) == 0 {
t.Fatal("native graph retrieval produced no nodes")
}
for _, n := range graph.Node {
for _, c := range n.Chunks {
if strings.HasPrefix(c, "chunk-") {
t.Fatal("legacy chunk ID leaked into native retrieval")
}
}
}
other := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(999999999))
denied, err := r.Browse(other, fixture.KB, "", 0, fixture.Entity)
if err != nil || len(denied.Nodes) != 0 || len(denied.Evidence) != 0 {
t.Fatal("cross-tenant graph leak", err)
}
wrongFile, err := r.SearchNode(ctx, types.NameSpace{KnowledgeBase: fixture.KB, Knowledge: "other-file"}, []string{fixture.Name})
if err != nil || len(wrongFile.Node) != 0 {
t.Fatal("file scope was ignored", err)
}
for _, change := range []struct {
name, sql, id string
}{
{"edited chunk", "UPDATE chunks SET content=content || ' stale-check' WHERE id=? AND tenant_id=10059", fixture.Chunk},
{"disabled chunk", "UPDATE chunks SET is_enabled=false WHERE id=? AND tenant_id=10059", fixture.Chunk},
{"deleted chunk", "UPDATE chunks SET deleted_at=now() WHERE id=? AND tenant_id=10059", fixture.Chunk},
{"deleted document", "UPDATE knowledges SET deleted_at=now() WHERE id=? AND tenant_id=10059", fixture.Knowledge},
{"deleted knowledge base", "UPDATE knowledge_bases SET deleted_at=now() WHERE id=? AND tenant_id=10059", fixture.KB},
} {
t.Run(change.name, func(t *testing.T) {
if err := tx.SavePoint("resource_check").Error; err != nil {
t.Fatal(err)
}
defer tx.RollbackTo("resource_check")
if err := tx.Exec(change.sql, change.id).Error; err != nil {
t.Fatal(err)
}
stale, err := r.Browse(ctx, fixture.KB, "", 0, fixture.Entity)
if err != nil {
t.Fatal(err)
}
for _, e := range stale.Evidence {
if e.NativeChunkID == fixture.Chunk {
t.Fatal("unavailable chunk remained linked")
}
}
staleGraph, err := r.SearchNode(ctx, types.NameSpace{KnowledgeBase: fixture.KB}, []string{fixture.Name})
if err != nil {
t.Fatal(err)
}
for _, n := range staleGraph.Node {
for _, c := range n.Chunks {
if c == fixture.Chunk {
t.Fatal("unavailable chunk leaked into retrieval")
}
}
}
})
}
}
func TestAcceptanceLargestGraph(t *testing.T) {
dsn := os.Getenv("GRAPH_TEST_DSN")
if dsn == "" {
t.Skip("GRAPH_TEST_DSN is required for acceptance integration")
}
db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
if err != nil {
t.Fatal(err)
}
ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(10059))
kb := "bb86b44a-576b-4bf8-a6bf-8f1baca3a845"
r := NewRepository(db, nil)
started := time.Now()
overview, err := r.Browse(ctx, kb, "", 0, 0)
if err != nil {
t.Fatal(err)
}
var preserved struct{ Sources, Mapped int64 }
if err := db.Raw(`SELECT count(*) AS sources,count(native_chunk_id) AS mapped FROM native_graph_sources WHERE tenant_id=10059 AND knowledge_base_id=? AND mapping_state<>'native_extraction'`, kb).Scan(&preserved).Error; err != nil {
t.Fatal(err)
}
if len(overview.Nodes) != 50 || preserved.Sources != 66204 || preserved.Mapped != 7000 || overview.Sources < preserved.Sources || overview.MappedSources < preserved.Mapped {
t.Fatalf("unexpected largest graph overview: preserved=%+v sources=%d mapped=%d", preserved, overview.Sources, overview.MappedSources)
}
t.Logf("largest KB overview: %s", time.Since(started))
started = time.Now()
graph, err := r.SearchNode(ctx, types.NameSpace{KnowledgeBase: kb}, []string{"销售", "客户", "信任", "心法"})
if err != nil {
t.Fatal(err)
}
if len(graph.Node) == 0 || len(graph.Node) > 32 || len(graph.Relation) > 32 {
t.Fatalf("invalid bounded retrieval: %d nodes, %d relations", len(graph.Node), len(graph.Relation))
}
for _, node := range graph.Node {
if len(node.Chunks) == 0 || len(node.Chunks) > 4 {
t.Fatalf("invalid chunk cap: %#v", node)
}
}
t.Logf("largest KB graph retrieval: %s; %d nodes, %d relations", time.Since(started), len(graph.Node), len(graph.Relation))
var relationFixture struct {
KB, From, To, Relation string
}
err = db.Raw(`SELECT r.knowledge_base_id AS kb,a.payload->>'canonical_name' AS "from",b.payload->>'canonical_name' AS "to",r.payload->>'relation_type' AS relation FROM native_graph_records r JOIN native_graph_records a ON a.tenant_id=r.tenant_id AND a.knowledge_base_id=r.knowledge_base_id AND a.source_table='knowledgeentity' AND a.legacy_id=(r.payload->>'from_entity_id')::bigint JOIN native_graph_records b ON b.tenant_id=r.tenant_id AND b.knowledge_base_id=r.knowledge_base_id AND b.source_table='knowledgeentity' AND b.legacy_id=(r.payload->>'to_entity_id')::bigint WHERE r.tenant_id=10059 AND r.source_table='knowledgerelation' AND r.legacy_id=1`).Scan(&relationFixture).Error
if err != nil || relationFixture.KB == "" {
t.Fatalf("missing relation fixture: %v", err)
}
expanded, err := r.SearchNode(ctx, types.NameSpace{KnowledgeBase: relationFixture.KB}, []string{relationFixture.From})
if err != nil {
t.Fatal(err)
}
found := false
for _, relation := range expanded.Relation {
if relation.Node1 == relationFixture.From && relation.Node2 == relationFixture.To && relation.Type == relationFixture.Relation {
found = true
}
}
if !found {
t.Fatal("verified relation neighbor was not expanded")
}
}
@@ -0,0 +1,522 @@
package legacygraph
import (
"context"
"crypto/sha256"
"encoding/binary"
"encoding/hex"
"encoding/json"
"fmt"
"slices"
"sort"
"strings"
"unicode"
"github.com/Tencent/WeKnora/internal/types"
"gorm.io/gorm"
)
const nativeOrigin = "native_extraction_v1"
type sourceChunk struct {
ID, KnowledgeID, KnowledgeBaseID, Content, FileName string
}
func contentHash(content string) string {
h := sha256.Sum256([]byte(content))
return hex.EncodeToString(h[:])
}
// IDs stay exactly representable in the browser, outside historical IDs. The
// full identity is also checked on conflict, so a hash collision fails closed.
func nativeID(parts ...string) (int64, string) {
identity, _ := json.Marshal(parts)
h := sha256.Sum256(identity)
return int64(binary.BigEndian.Uint64(h[:8])&((1<<52)-1)) | 1<<52, hex.EncodeToString(h[:])
}
func normalizedName(name string) string { return strings.ToLower(strings.TrimSpace(name)) }
func sourceScopedEntity(name string) bool {
name = strings.Map(func(r rune) rune {
if unicode.IsSpace(r) {
return -1
}
return r
}, normalizedName(name))
switch name {
case "本品", "该产品", "本产品", "产品", "本商品", "此产品", "本公司", "该公司", "本文", "本方法", "该方法", "thisproduct", "theproduct":
return true
default:
// The extractor uses these names to separate unresolved local products.
return strings.HasSuffix(name, "未名产品") || strings.HasPrefix(name, "本品(")
}
}
func nativeEntityIdentity(document string, node *types.GraphNode) (int64, string, error) {
key := normalizedName(node.Name)
parts := []string{document, "entity", key}
if sourceScopedEntity(key) {
// An unnamed referent is not a document-wide entity. Multiple sources
// require the extractor to supply an unambiguous, specifically named node.
if len(node.Chunks) == 0 || node.Chunks[0] == "" {
return 0, "", fmt.Errorf("generic graph entity requires one source chunk")
}
for _, chunk := range node.Chunks {
if chunk != node.Chunks[0] {
return 0, "", fmt.Errorf("generic graph entity cannot merge distinct source chunks")
}
}
parts = append(parts, "source", node.Chunks[0])
}
id, identity := nativeID(parts...)
return id, identity, nil
}
// HasChunkGraph accepts only current checkpoints with extracted relations.
// Legacy empty results and excerpt mappings require coverage review.
func (r *Repository) HasChunkGraph(ctx context.Context, chunk *types.Chunk) (bool, error) {
tenant, ok := types.TenantIDFromContext(ctx)
if !ok || tenant == 0 || chunk == nil || chunk.TenantID != tenant || r.db == nil {
return false, fmt.Errorf("graph checkpoint requires source tenant and chunk")
}
var count int64
err := r.db.WithContext(ctx).Table("native_graph_records").Where("tenant_id=? AND knowledge_base_id=? AND source_table='knowledgegraphchunkstate' AND payload->>'origin'=? AND payload->>'native_chunk_id'=? AND payload->>'native_knowledge_id'=? AND payload->>'source_content_hash'=? AND payload->>'extraction_status'='complete' AND CAST(COALESCE(payload->>'relation_count','0') AS INTEGER)>0", tenant, chunk.KnowledgeBaseID, nativeOrigin, chunk.ID, chunk.KnowledgeID, contentHash(chunk.Content)).Count(&count).Error
return count > 0, err
}
// AddChunkGraph rejects unreviewed empty coverage and results from changed sources.
func (r *Repository) AddChunkGraph(ctx context.Context, ns types.NameSpace, chunk *types.Chunk, graph *types.GraphData) error {
if chunk == nil || chunk.ID == "" || graph == nil {
return fmt.Errorf("graph extraction requires its source chunk and model result")
}
if len(graph.Relation) == 0 {
return fmt.Errorf("graph extraction without relations requires coverage review")
}
return r.addNativeGraphs(ctx, ns, []*types.GraphData{graph}, map[string]string{chunk.ID: contentHash(chunk.Content)})
}
func (r *Repository) addNativeGraphs(ctx context.Context, ns types.NameSpace, graphs []*types.GraphData, expected map[string]string) error {
tenant, ok := types.TenantIDFromContext(ctx)
if !ok || tenant == 0 || ns.KnowledgeBase == "" || ns.Knowledge == "" || r.db == nil {
return fmt.Errorf("graph writes require tenant, knowledge base, document and database")
}
ids := map[string]bool{}
for id := range expected {
ids[id] = true
}
for _, graph := range graphs {
if graph == nil {
return fmt.Errorf("nil graph extraction")
}
for _, node := range graph.Node {
if node == nil || normalizedName(node.Name) == "" || len(node.Chunks) == 0 {
return fmt.Errorf("extracted graph node requires name and source chunks")
}
for _, id := range node.Chunks {
if id == "" || (expected != nil && expected[id] == "") {
return fmt.Errorf("graph node refers to a chunk outside extraction scope")
}
ids[id] = true
}
}
}
if len(ids) == 0 {
return fmt.Errorf("graph extraction has no source chunks")
}
chunkIDs := make([]string, 0, len(ids))
for id := range ids {
chunkIDs = append(chunkIDs, id)
}
sort.Strings(chunkIDs)
return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
if err := lockGraph(tx, tenant, ns.KnowledgeBase); err != nil {
return err
}
var chunks []sourceChunk
if err := tx.Raw(`SELECT c.id,c.knowledge_id,c.knowledge_base_id,c.content,k.file_name FROM chunks c
JOIN knowledges k ON k.id=c.knowledge_id AND k.tenant_id=c.tenant_id AND k.knowledge_base_id=c.knowledge_base_id AND k.deleted_at IS NULL
JOIN knowledge_bases kb ON kb.id=c.knowledge_base_id AND kb.tenant_id=c.tenant_id AND kb.deleted_at IS NULL
WHERE c.tenant_id=? AND c.knowledge_base_id=? AND c.knowledge_id=? AND c.id IN ? AND c.deleted_at IS NULL AND c.is_enabled=true
ORDER BY c.id FOR SHARE OF c,k,kb`, tenant, ns.KnowledgeBase, ns.Knowledge, chunkIDs).Scan(&chunks).Error; err != nil {
return err
}
if len(chunks) != len(chunkIDs) {
return fmt.Errorf("graph source chunks are unavailable or outside document/tenant scope")
}
byID := map[string]sourceChunk{}
for _, chunk := range chunks {
if expected != nil && expected[chunk.ID] != contentHash(chunk.Content) {
return fmt.Errorf("graph source chunk %s changed during extraction", chunk.ID)
}
byID[chunk.ID] = chunk
}
// Replace only this extraction's evidence. Other documents, chunks and
// archived migration records are untouched, including on retries.
sources := make([]string, 0, len(chunks))
for _, chunk := range chunks {
source := "native:" + chunk.ID
sources = append(sources, source)
excerpt := []rune(chunk.Content)
if len(excerpt) > 2000 {
excerpt = excerpt[:2000]
}
if err := tx.Exec(`INSERT INTO native_graph_sources(tenant_id,knowledge_base_id,source_chunk_id,source_file,source_content_hash,excerpt,native_knowledge_id,native_chunk_id,native_content_sha256,mapping_state)
VALUES(?,?,?,?,?,?,?,?,?,'native_extraction') ON CONFLICT(tenant_id,knowledge_base_id,source_chunk_id) DO UPDATE SET source_file=excluded.source_file,source_content_hash=excluded.source_content_hash,excerpt=excluded.excerpt,native_knowledge_id=excluded.native_knowledge_id,native_chunk_id=excluded.native_chunk_id,native_content_sha256=excluded.native_content_sha256,mapping_state=excluded.mapping_state`, tenant, ns.KnowledgeBase, source, chunk.FileName, contentHash(chunk.Content), string(excerpt), ns.Knowledge, chunk.ID, contentHash(chunk.Content)).Error; err != nil {
return err
}
}
type pendingEntity struct {
identity string
payload map[string]any
}
entities := map[int64]pendingEntity{}
// Removed nodes also lose their source's attributes, even when a bare
// reference from another chunk keeps the entity alive.
var replacedEntities []Record
if err := tx.Raw(`SELECT r.legacy_id,r.payload FROM native_graph_records r
WHERE r.tenant_id=? AND r.knowledge_base_id=? AND r.source_table='knowledgeentity'
AND r.payload->>'origin'=? AND r.payload->>'native_knowledge_id'=?
AND EXISTS (SELECT 1 FROM native_graph_records e WHERE e.tenant_id=r.tenant_id AND e.knowledge_base_id=r.knowledge_base_id
AND e.source_table='knowledgegraphevidence' AND e.payload->>'origin'=?
AND e.payload->>'subject_type'='entity' AND e.payload->>'subject_id'=r.legacy_id::text AND e.payload->>'source_chunk_id' IN ?)`, tenant, ns.KnowledgeBase, nativeOrigin, ns.Knowledge, nativeOrigin, sources).Scan(&replacedEntities).Error; err != nil {
return err
}
for _, record := range replacedEntities {
var payload map[string]any
if err := json.Unmarshal(record.Payload, &payload); err != nil {
return err
}
identity, _ := payload["native_identity"].(string)
entities[record.ID] = pendingEntity{identity, payload}
}
if err := tx.Exec(`DELETE FROM native_graph_records WHERE tenant_id=? AND knowledge_base_id=? AND source_table='knowledgegraphevidence' AND payload->>'origin'=? AND payload->>'source_chunk_id' IN ?`, tenant, ns.KnowledgeBase, nativeOrigin, sources).Error; err != nil {
return err
}
counts := map[string]map[string]int{}
entitySources := map[int64]map[string][]string{}
for _, id := range chunkIDs {
counts[id] = map[string]int{"entity": 0, "relation": 0}
}
for _, graph := range graphs {
nodes := map[string]*types.GraphNode{}
entityIDs := map[string]int64{}
for _, node := range graph.Node {
key := normalizedName(node.Name)
id, identity, err := nativeEntityIdentity(ns.Knowledge, node)
if err != nil {
return err
}
if priorID, exists := entityIDs[key]; exists && priorID != id {
return fmt.Errorf("generic graph entity is ambiguous within extraction")
}
name := strings.TrimSpace(node.Name)
payload := map[string]any{"canonical_name": name, "entity_type": "concept", "attributes": node.Attributes}
if len(node.Attributes) > 0 {
payload["entity_type"] = node.Attributes[0]
}
if prior, exists := entities[id]; exists && prior.identity != identity {
return fmt.Errorf("native graph identity collision; preserved graph was not modified")
}
if _, exists := entitySources[id]; !exists {
entities[id] = pendingEntity{identity, payload}
entitySources[id] = map[string][]string{}
} else if name < entities[id].payload["canonical_name"].(string) {
entities[id].payload["canonical_name"] = name
}
if prior := nodes[key]; prior != nil {
prior.Chunks = append(prior.Chunks, node.Chunks...)
} else {
copyNode := *node
copyNode.Chunks = append([]string(nil), node.Chunks...)
nodes[key] = &copyNode
}
entityIDs[key] = id
seen := map[string]bool{}
for _, chunkID := range node.Chunks {
if seen[chunkID] {
continue
}
seen[chunkID] = true
if attributes, exists := entitySources[id][chunkID]; exists {
if !slices.Equal(attributes, node.Attributes) {
return fmt.Errorf("duplicate native graph entity source attributes conflict; extraction requires review")
}
continue
}
entitySources[id][chunkID] = node.Attributes
if err := putEvidence(tx, tenant, ns, "entity", id, byID[chunkID], "", node.Attributes); err != nil {
return err
}
counts[chunkID]["entity"]++
}
}
for _, rel := range graph.Relation {
if rel == nil || strings.TrimSpace(rel.Type) == "" {
return fmt.Errorf("invalid extracted relation")
}
from, to := normalizedName(rel.Node1), normalizedName(rel.Node2)
if nodes[from] == nil || nodes[to] == nil {
return fmt.Errorf("relation references an entity outside extraction")
}
shared := []string{}
fromChunks := map[string]bool{}
for _, id := range nodes[from].Chunks {
fromChunks[id] = true
}
for _, id := range nodes[to].Chunks {
if fromChunks[id] {
shared = append(shared, id)
delete(fromChunks, id)
}
}
if len(shared) == 0 {
return fmt.Errorf("relation has no shared source evidence")
}
if rel.Evidence != "" {
grounded := shared[:0]
for _, chunkID := range shared {
if strings.Contains(byID[chunkID].Content, rel.Evidence) {
grounded = append(grounded, chunkID)
}
}
shared = grounded
if len(shared) == 0 {
return fmt.Errorf("relation quotation is absent from current source evidence")
}
}
identityParts := []string{ns.Knowledge, "relation", from, to, normalizedName(rel.Type)}
if sourceScopedEntity(from) || sourceScopedEntity(to) {
identityParts = append(identityParts, "scoped_endpoints", fmt.Sprint(entityIDs[from]), fmt.Sprint(entityIDs[to]))
}
id, identity := nativeID(identityParts...)
if err := putNativeRecord(tx, tenant, ns, "knowledgerelation", id, identity, map[string]any{"from_entity_id": entityIDs[from], "to_entity_id": entityIDs[to], "relation_type": strings.TrimSpace(rel.Type), "evidence": rel.Evidence}, ""); err != nil {
return err
}
for _, chunkID := range shared {
if err := putEvidence(tx, tenant, ns, "relation", id, byID[chunkID], rel.Evidence); err != nil {
return err
}
counts[chunkID]["relation"]++
}
}
}
// Reconcile after every source snapshot is present, so batch ordering
// cannot decide which source owns the entity's retained attributes.
entityIDs := make([]int64, 0, len(entities))
for id := range entities {
entityIDs = append(entityIDs, id)
}
slices.Sort(entityIDs)
for _, id := range entityIDs {
entity := entities[id]
if err := putNativeRecord(tx, tenant, ns, "knowledgeentity", id, entity.identity, entity.payload, ""); err != nil {
return err
}
}
for _, chunk := range chunks {
id, identity := nativeID(ns.Knowledge, "chunkstate", chunk.ID)
status := "complete"
if counts[chunk.ID]["relation"] == 0 {
status = "needs_review"
}
if err := putNativeRecord(tx, tenant, ns, "knowledgegraphchunkstate", id, identity, map[string]any{"source_chunk_id": "native:" + chunk.ID, "native_chunk_id": chunk.ID, "source_content_hash": contentHash(chunk.Content), "extraction_status": status, "entity_count": counts[chunk.ID]["entity"], "relation_count": counts[chunk.ID]["relation"]}, contentHash(chunk.Content)); err != nil {
return err
}
}
return pruneNativeOrphans(tx, tenant, ns)
})
}
func lockGraph(tx *gorm.DB, tenant uint64, kb string) error {
return tx.Exec("SELECT pg_advisory_xact_lock(hashtextextended(?,0))", fmt.Sprintf("native-graph:%d:%s", tenant, kb)).Error
}
func putNativeRecord(tx *gorm.DB, tenant uint64, ns types.NameSpace, table string, id int64, identity string, payload map[string]any, hash string) error {
if table == "knowledgeentity" {
if err := reconcileNativeEntityAttributes(tx, tenant, ns, id, identity, payload); err != nil {
return err
}
}
payload["origin"], payload["native_identity"], payload["native_knowledge_id"], payload["status"] = nativeOrigin, identity, ns.Knowledge, "active"
data, err := json.Marshal(payload)
if err != nil {
return err
}
if hash == "" {
hash = contentHash(string(data))
}
result := tx.Exec(`INSERT INTO native_graph_records(tenant_id,knowledge_base_id,source_table,legacy_id,graph_version_id,payload,source_snapshot_sha256)
VALUES(?,?,?,?,0,?::jsonb,?) ON CONFLICT(tenant_id,knowledge_base_id,source_table,legacy_id) DO UPDATE SET payload=excluded.payload,source_snapshot_sha256=excluded.source_snapshot_sha256
WHERE native_graph_records.payload->>'native_identity'=excluded.payload->>'native_identity' AND native_graph_records.payload->>'origin'=?`, tenant, ns.KnowledgeBase, table, id, string(data), hash, nativeOrigin)
if result.Error != nil {
return result.Error
}
if result.RowsAffected != 1 {
return fmt.Errorf("native graph identity collision; preserved graph was not modified")
}
return nil
}
func reconcileNativeEntityAttributes(tx *gorm.DB, tenant uint64, ns types.NameSpace, id int64, identity string, payload map[string]any) error {
var evidence []struct {
Payload json.RawMessage
}
// Evidence for the replaced sources has already been rewritten with raw
// node attributes. Only still-current sources may preserve shared facts.
if err := tx.Raw(`SELECT e.payload FROM native_graph_records e
JOIN native_graph_sources s ON s.tenant_id=e.tenant_id AND s.knowledge_base_id=e.knowledge_base_id AND s.source_chunk_id=e.payload->>'source_chunk_id'
JOIN chunks c ON c.id=s.native_chunk_id AND c.tenant_id=s.tenant_id AND c.knowledge_base_id=s.knowledge_base_id
AND c.deleted_at IS NULL AND c.is_enabled=true AND encode(sha256(convert_to(c.content,'UTF8')),'hex')=s.native_content_sha256
JOIN knowledges k ON k.id=c.knowledge_id AND k.tenant_id=c.tenant_id AND k.knowledge_base_id=c.knowledge_base_id AND k.deleted_at IS NULL
JOIN knowledge_bases kb ON kb.id=c.knowledge_base_id AND kb.tenant_id=c.tenant_id AND kb.deleted_at IS NULL
WHERE e.tenant_id=? AND e.knowledge_base_id=? AND e.source_table='knowledgegraphevidence'
AND e.payload->>'subject_type'='entity' AND e.payload->>'subject_id'=?
AND e.payload->>'origin'=? AND e.payload->>'status'='active' AND e.payload->>'native_knowledge_id'=?
AND c.knowledge_id=? AND s.native_knowledge_id=c.knowledge_id AND s.mapping_state='native_extraction'
AND e.payload->>'source_content_hash'=s.native_content_sha256 AND s.source_content_hash=s.native_content_sha256
AND e.source_snapshot_sha256=s.native_content_sha256
ORDER BY c.id,e.legacy_id FOR SHARE OF c,k,kb`, tenant, ns.KnowledgeBase, fmt.Sprint(id), nativeOrigin, ns.Knowledge, ns.Knowledge).Scan(&evidence).Error; err != nil {
return err
}
attributes := []string{}
unknown := false
for _, row := range evidence {
var source struct {
Attributes json.RawMessage `json:"entity_attributes"`
}
if err := json.Unmarshal(row.Payload, &source); err != nil {
return err
}
if len(source.Attributes) == 0 {
unknown = true
continue
}
var snapshot []*string
if err := json.Unmarshal(source.Attributes, &snapshot); err != nil {
return fmt.Errorf("invalid native graph entity attribute evidence: %w", err)
}
if snapshot == nil {
unknown = true
continue
}
values := make([]string, len(snapshot))
for i, value := range snapshot {
if value == nil {
return fmt.Errorf("invalid native graph entity attribute evidence: null attribute")
}
values[i] = *value
}
if len(values) == 0 {
continue
}
if len(attributes) > 0 && !slices.Equal(attributes, values) {
return fmt.Errorf("shared native graph entity attributes conflict; extraction requires review")
}
attributes = values
}
entityType := "concept"
if len(attributes) > 0 {
entityType = attributes[0]
}
if unknown {
var existing Record
result := tx.Raw(`SELECT payload FROM native_graph_records WHERE tenant_id=? AND knowledge_base_id=?
AND source_table='knowledgeentity' AND legacy_id=? AND payload->>'origin'=? AND payload->>'native_identity'=?`, tenant, ns.KnowledgeBase, id, nativeOrigin, identity).Scan(&existing)
if result.Error != nil {
return result.Error
}
var prior struct {
Attributes []string `json:"attributes"`
Type string `json:"entity_type"`
}
if result.RowsAffected != 1 {
return fmt.Errorf("shared native graph entity attribute provenance is unavailable; extraction requires review")
}
if err := json.Unmarshal(existing.Payload, &prior); err != nil {
return err
}
// Legacy evidence cannot prove which surviving source owns a fact.
// Refuse changing its aggregate until those sources are re-extracted.
if !slices.Equal(prior.Attributes, attributes) || prior.Type != entityType {
return fmt.Errorf("shared native graph entity attribute provenance is unavailable; extraction requires review")
}
attributes = prior.Attributes
}
payload["attributes"], payload["entity_type"] = attributes, entityType
return nil
}
func putEvidence(tx *gorm.DB, tenant uint64, ns types.NameSpace, subject string, subjectID int64, chunk sourceChunk, quotation string, attributes ...[]string) error {
id, identity := nativeID(ns.Knowledge, "evidence", subject, fmt.Sprint(subjectID), chunk.ID)
excerpt := []rune(chunk.Content)
if len(excerpt) > 2000 {
excerpt = excerpt[:2000]
}
if quotation != "" {
excerpt = []rune(quotation)
}
payload := map[string]any{"subject_type": subject, "subject_id": subjectID, "source_chunk_id": "native:" + chunk.ID, "source_file": chunk.FileName, "source_content_hash": contentHash(chunk.Content), "excerpt": string(excerpt)}
if subject == "entity" {
if len(attributes) != 1 {
return fmt.Errorf("native graph entity evidence requires source attributes")
}
payload["entity_attributes"] = append([]string{}, attributes[0]...)
}
return putNativeRecord(tx, tenant, ns, "knowledgegraphevidence", id, identity, payload, contentHash(chunk.Content))
}
func pruneNativeOrphans(tx *gorm.DB, tenant uint64, ns types.NameSpace) error {
query := `DELETE FROM native_graph_records r WHERE r.tenant_id=? AND r.knowledge_base_id=? AND r.payload->>'origin'=? AND r.source_table IN ('knowledgeentity','knowledgerelation')`
args := []any{tenant, ns.KnowledgeBase, nativeOrigin}
if ns.Knowledge != "" {
query += " AND r.payload->>'native_knowledge_id'=?"
args = append(args, ns.Knowledge)
}
query += ` AND NOT EXISTS (SELECT 1 FROM native_graph_records e WHERE e.tenant_id=r.tenant_id AND e.knowledge_base_id=r.knowledge_base_id AND e.source_table='knowledgegraphevidence' AND e.payload->>'subject_id'=r.legacy_id::text AND e.payload->>'subject_type'=CASE r.source_table WHEN 'knowledgeentity' THEN 'entity' ELSE 'relation' END)`
return tx.Exec(query, args...).Error
}
func (r *Repository) deleteNativeGraphs(ctx context.Context, namespaces []types.NameSpace) error {
tenant, ok := types.TenantIDFromContext(ctx)
if !ok || tenant == 0 || r.db == nil {
return fmt.Errorf("graph deletion requires tenant and database")
}
for _, ns := range namespaces {
if ns.KnowledgeBase == "" {
return fmt.Errorf("graph deletion requires knowledge base scope")
}
}
for _, ns := range namespaces {
if err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
if err := lockGraph(tx, tenant, ns.KnowledgeBase); err != nil {
return err
}
query := tx.Table("native_graph_records").Where("tenant_id=? AND knowledge_base_id=? AND payload->>'origin'=?", tenant, ns.KnowledgeBase, nativeOrigin)
if ns.Knowledge != "" {
// Version copies preserve original payloads but remap source
// pointers. Scope deletions through those current pointers too.
query = query.Where("source_table IN ('knowledgegraphevidence','knowledgegraphchunkstate') AND (payload->>'native_knowledge_id'=? OR payload->>'source_chunk_id' IN (SELECT source_chunk_id FROM native_graph_sources WHERE tenant_id=? AND knowledge_base_id=? AND native_knowledge_id=?))", ns.Knowledge, tenant, ns.KnowledgeBase, ns.Knowledge)
}
if err := query.Delete(&Record{}).Error; err != nil {
return err
}
if ns.Knowledge != "" {
if err := pruneNativeOrphans(tx, tenant, types.NameSpace{KnowledgeBase: ns.KnowledgeBase}); err != nil {
return err
}
}
query = tx.Table("native_graph_sources").Where("tenant_id=? AND knowledge_base_id=? AND source_chunk_id LIKE 'native:%'", tenant, ns.KnowledgeBase)
if ns.Knowledge != "" {
query = query.Where("native_knowledge_id=?", ns.Knowledge)
}
return query.Delete(&Record{}).Error
}); err != nil {
return err
}
}
if r.fallback != nil {
return r.fallback.DelGraph(ctx, namespaces)
}
return nil
}
@@ -0,0 +1,657 @@
package legacygraph
import (
"context"
"encoding/json"
"fmt"
"os"
"slices"
"strings"
"testing"
"github.com/Tencent/WeKnora/internal/types"
"gorm.io/driver/postgres"
"gorm.io/gorm"
)
func TestNativeIdentity(t *testing.T) {
a, identity := nativeID("doc", "entity", "concept")
b, _ := nativeID("other-doc", "entity", "concept")
if a < 1<<52 || a > (1<<53)-1 || a == b || len(identity) != 64 {
t.Fatal("invalid stable document-scoped browser-safe ID")
}
c, again := nativeID("doc", "entity", "concept")
if a != c || identity != again {
t.Fatal("unstable identity")
}
}
func TestNativeGenericEntityIdentityIsSourceScoped(t *testing.T) {
for _, name := range []string{"本品", "本商品", "前段未名产品", "后段未名产品", "本品(前段)", "前段 未名产品", "本 品 (前段)", "This Product"} {
t.Run(name, func(t *testing.T) {
first, _, err := nativeEntityIdentity("doc", &types.GraphNode{Name: name, Chunks: []string{"first"}})
if err != nil {
t.Fatal(err)
}
second, _, err := nativeEntityIdentity("doc", &types.GraphNode{Name: name, Chunks: []string{"second"}})
if err != nil || first == second {
t.Fatal("unnamed products from separate chunks were merged", err)
}
for _, chunks := range [][]string{nil, {""}, {"first", "second"}} {
if _, _, err := nativeEntityIdentity("doc", &types.GraphNode{Name: name, Chunks: chunks}); err == nil {
t.Fatalf("missing or ambiguous source accepted: %v", chunks)
}
}
})
}
for _, name := range []string{"Named product", "修复凝胶", "本品专用护理套装", "未名产品研究中心"} {
for _, chunk := range []string{"first", "second"} {
id, _, err := nativeEntityIdentity("doc", &types.GraphNode{Name: name, Chunks: []string{chunk}})
expected, _ := nativeID("doc", "entity", normalizedName(name))
if err != nil || id != expected {
t.Fatal("specific product lost its stable document identity", err)
}
}
}
}
func TestNativeGenericProductsDoNotShareRelations(t *testing.T) {
for _, name := range []string{"本品", "本商品", "前段未名产品", "后段未名产品", "本品(前段)"} {
t.Run(name, func(t *testing.T) {
tx, r, ctx := nativeWriteFixture(t)
if err := tx.Exec(`INSERT INTO chunks VALUES('sibling',7,'kb','doc','Evidence about trust and learning.',NULL,true)`).Error; err != nil {
t.Fatal(err)
}
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
for _, chunkID := range []string{"chunk", "sibling"} {
graph := testGraph(chunkID)
graph.Node[0].Name = name
graph.Relation[0].Node1 = name
if err := r.AddGraph(ctx, ns, []*types.GraphData{graph}); err != nil {
t.Fatal(err)
}
}
overview, err := r.Browse(ctx, "kb", "", 0, 0)
if err != nil || overview.Counts["knowledgeentity"] != 3 || overview.Counts["knowledgerelation"] != 2 {
t.Fatalf("separate unnamed products conflated: %#v %v", overview, err)
}
for _, chunkID := range []string{"chunk", "sibling"} {
id, _, _ := nativeEntityIdentity("doc", &types.GraphNode{Name: name, Chunks: []string{chunkID}})
detail, err := r.Browse(ctx, "kb", "", 0, id)
if err != nil || len(detail.Relations) != 1 || len(detail.Evidence) != 1 || detail.Evidence[0].NativeChunkID != chunkID {
t.Fatalf("unnamed product points to another source: %#v %v", detail, err)
}
}
if err := r.AddGraph(ctx, ns, []*types.GraphData{{Node: []*types.GraphNode{{Name: name, Chunks: []string{"chunk", "sibling"}}}}}); err == nil {
t.Fatal("ambiguous multi-source write accepted")
}
after, err := r.Browse(ctx, "kb", "", 0, 0)
if err != nil || after.Counts["knowledgerelation"] != 2 {
t.Fatal("failed ambiguous write damaged committed graph", err)
}
})
}
}
func nativeWriteFixture(t *testing.T) (*gorm.DB, *Repository, context.Context) {
t.Helper()
dsn := os.Getenv("GRAPH_TEST_DSN")
if dsn == "" {
t.Skip("GRAPH_TEST_DSN required for temporary-table integration")
}
db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
if err != nil {
t.Fatal(err)
}
sqlDB, err := db.DB()
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { sqlDB.Close() })
tx := db.Begin()
if tx.Error != nil {
t.Fatal(tx.Error)
}
t.Cleanup(func() { tx.Rollback() })
// Temporary tables shadow live names on this transaction's connection;
// tests never mutate application records or permanent schemas.
for _, sql := range []string{
`CREATE TEMP TABLE knowledge_bases(id text,tenant_id bigint,deleted_at timestamptz) ON COMMIT DROP`,
`CREATE TEMP TABLE knowledges(id text,tenant_id bigint,knowledge_base_id text,file_name text,deleted_at timestamptz) ON COMMIT DROP`,
`CREATE TEMP TABLE chunks(id text,tenant_id bigint,knowledge_base_id text,knowledge_id text,content text,deleted_at timestamptz,is_enabled boolean) ON COMMIT DROP`,
`CREATE TEMP TABLE native_graph_records(tenant_id bigint,knowledge_base_id text,source_table text,legacy_id bigint,graph_version_id bigint,payload jsonb,source_snapshot_sha256 text,imported_at timestamptz DEFAULT now(),PRIMARY KEY(tenant_id,knowledge_base_id,source_table,legacy_id)) ON COMMIT DROP`,
`CREATE TEMP TABLE native_graph_sources(tenant_id bigint,knowledge_base_id text,source_chunk_id text,source_file text,source_content_hash text,excerpt text,native_knowledge_id text,native_chunk_id text,native_content_sha256 text,mapping_state text,PRIMARY KEY(tenant_id,knowledge_base_id,source_chunk_id)) ON COMMIT DROP`,
`INSERT INTO knowledge_bases VALUES('kb',7,NULL),('kb',8,NULL),('other-kb',7,NULL)`,
`INSERT INTO knowledges VALUES('doc',7,'kb','source.md',NULL),('doc2',7,'kb','other.md',NULL),('doc',8,'kb','private.md',NULL)`,
`INSERT INTO chunks VALUES('chunk',7,'kb','doc','Evidence about trust and learning.',NULL,true),('chunk2',7,'kb','doc2','Evidence about trust and learning.',NULL,true),('foreign',8,'kb','doc','Private source.',NULL,true)`,
} {
if err := tx.Exec(sql).Error; err != nil {
t.Fatal(err)
}
}
ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(7))
return tx, NewRepository(tx, nil), ctx
}
func testGraph(chunk string) *types.GraphData {
return &types.GraphData{Node: []*types.GraphNode{{Name: "Trust", Chunks: []string{chunk}, Attributes: []string{"concept", "Earned through reliable evidence"}}, {Name: "Learning", Chunks: []string{chunk}, Attributes: []string{"process"}}}, Relation: []*types.GraphRelation{{Node1: "Trust", Node2: "Learning", Type: "supports"}}}
}
func TestNativeGraphWriteReplayBrowseAndIsolation(t *testing.T) {
tx, r, ctx := nativeWriteFixture(t)
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
chunk := &types.Chunk{ID: "chunk", TenantID: 7, KnowledgeBaseID: "kb", KnowledgeID: "doc", Content: "Evidence about trust and learning."}
for i := 0; i < 2; i++ {
if err := r.AddChunkGraph(ctx, ns, chunk, testGraph("chunk")); err != nil {
t.Fatal(err)
}
}
overview, err := r.Browse(ctx, "kb", "", 0, 0)
if err != nil {
t.Fatal(err)
}
if overview.Counts["knowledgeentity"] != 2 || overview.Counts["knowledgerelation"] != 1 || overview.Counts["knowledgegraphevidence"] != 3 || overview.Counts["knowledgegraphchunkstate"] != 1 || overview.MappedSources != 1 {
t.Fatalf("not durable/idempotent: %#v", overview)
}
id, _ := nativeID("doc", "entity", "trust")
detail, err := r.Browse(ctx, "kb", "", 0, id)
if err != nil || len(detail.Relations) != 1 || len(detail.Evidence) != 1 || detail.Evidence[0].NativeChunkID != "chunk" {
t.Fatalf("detail missing actual evidence: %#v %v", detail, err)
}
found, err := r.SearchNode(ctx, ns, []string{"Trust"})
if err != nil || len(found.Node) != 2 || len(found.Relation) != 1 {
t.Fatalf("native RAG cannot find persisted graph: %#v %v", found, err)
}
complete, err := r.HasChunkGraph(ctx, chunk)
if err != nil || !complete {
t.Fatalf("missing checkpoint: %v", err)
}
if err := r.AddGraph(ctx, ns, []*types.GraphData{testGraph("foreign")}); err == nil {
t.Fatal("cross-tenant source accepted")
}
if err := r.AddGraph(ctx, ns, []*types.GraphData{testGraph("chunk2")}); err == nil {
t.Fatal("cross-document source accepted")
}
changed := *chunk
changed.Content = "old model input"
if err := r.AddChunkGraph(ctx, ns, &changed, testGraph("chunk")); err == nil {
t.Fatal("stale model output accepted")
}
bad := testGraph("chunk")
bad.Relation[0].Node2 = "invented"
if err := r.AddChunkGraph(ctx, ns, chunk, bad); err == nil {
t.Fatal("ungrounded relation accepted")
}
after, err := r.Browse(ctx, "kb", "", 0, id)
if err != nil || len(after.Evidence) != 1 {
t.Fatal("failed extraction destroyed committed evidence", err)
}
if err := r.AddGraph(ctx, types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc2"}, []*types.GraphData{testGraph("chunk2")}); err != nil {
t.Fatal(err)
}
if err := tx.Exec(`INSERT INTO native_graph_records VALUES(7,'kb','knowledgeentity',1,1,'{"canonical_name":"Archived","status":"active"}',repeat('a',64),now())`).Error; err != nil {
t.Fatal(err)
}
if err := r.DelGraph(ctx, []types.NameSpace{ns}); err != nil {
t.Fatal(err)
}
left, err := r.Browse(ctx, "kb", "", 0, 0)
if err != nil {
t.Fatal(err)
}
if left.Counts["knowledgeentity"] != 3 || left.Sources != 1 {
t.Fatalf("deletion damaged other document/archive: %#v", left)
}
wrong := context.WithValue(ctx, types.TenantIDContextKey, uint64(8))
if err := r.DelGraph(wrong, []types.NameSpace{{KnowledgeBase: "kb"}}); err != nil {
t.Fatal(err)
}
left, err = r.Browse(ctx, "kb", "", 0, 0)
if err != nil || left.Counts["knowledgeentity"] != 3 {
t.Fatal("cross-tenant deletion", err)
}
}
func TestNativeEmptyResultAndReextraction(t *testing.T) {
tx, r, ctx := nativeWriteFixture(t)
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
chunk := &types.Chunk{ID: "chunk", TenantID: 7, KnowledgeBaseID: "kb", KnowledgeID: "doc", Content: "Evidence about trust and learning."}
if err := r.AddChunkGraph(ctx, ns, chunk, testGraph("chunk")); err != nil {
t.Fatal(err)
}
if err := r.AddChunkGraph(ctx, ns, chunk, &types.GraphData{}); err == nil {
t.Fatal("empty extraction accepted without review")
}
overview, err := r.Browse(ctx, "kb", "", 0, 0)
if err != nil || overview.Counts["knowledgerelation"] == 0 || overview.Counts["knowledgegraphchunkstate"] != 1 {
t.Fatalf("rejected result replaced prior graph: %#v %v", overview, err)
}
complete, err := r.HasChunkGraph(ctx, chunk)
if err != nil || !complete {
t.Fatal("prior verified graph was lost", err)
}
if err := tx.Exec(`UPDATE native_graph_records SET payload=jsonb_set(payload,'{relation_count}','0'::jsonb) WHERE source_table='knowledgegraphchunkstate'`).Error; err != nil {
t.Fatal(err)
}
complete, err = r.HasChunkGraph(ctx, chunk)
if err != nil || complete {
t.Fatal("legacy empty checkpoint was accepted", err)
}
if err := r.AddGraph(ctx, ns, []*types.GraphData{{}}); err == nil {
t.Fatal("empty extraction without source accepted")
}
}
func TestNativeRelationQuotationIsPersistedAndChecked(t *testing.T) {
tx, r, ctx := nativeWriteFixture(t)
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
graph := testGraph("chunk")
graph.Relation[0].Evidence = "trust and learning"
if err := r.AddGraph(ctx, ns, []*types.GraphData{graph}); err != nil {
t.Fatal(err)
}
var quote, excerpt string
if err := tx.Raw(`SELECT payload->>'evidence' FROM native_graph_records WHERE source_table='knowledgerelation'`).Scan(&quote).Error; err != nil {
t.Fatal(err)
}
if err := tx.Raw(`SELECT payload->>'excerpt' FROM native_graph_records WHERE source_table='knowledgegraphevidence' AND payload->>'subject_type'='relation'`).Scan(&excerpt).Error; err != nil {
t.Fatal(err)
}
if quote != graph.Relation[0].Evidence || excerpt != quote {
t.Fatalf("relation quotation not retained: %q %q", quote, excerpt)
}
graph.Relation[0].Evidence = "invented quotation"
if err := r.AddGraph(ctx, ns, []*types.GraphData{graph}); err == nil {
t.Fatal("relation quotation absent from source accepted")
}
if err := tx.Raw(`SELECT payload->>'evidence' FROM native_graph_records WHERE source_table='knowledgerelation'`).Scan(&quote).Error; err != nil || quote != "trust and learning" {
t.Fatal("invalid quotation damaged prior graph", err)
}
}
func TestNativeRelationRequiresCurrentSourceEvidence(t *testing.T) {
tx, r, ctx := nativeWriteFixture(t)
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
if err := tx.Exec(`INSERT INTO chunks VALUES('sibling',7,'kb','doc','Another source discussing trust and learning.',NULL,true)`).Error; err != nil {
t.Fatal(err)
}
if err := r.AddGraph(ctx, ns, []*types.GraphData{testGraph("chunk")}); err != nil {
t.Fatal(err)
}
graph := testGraph("sibling")
graph.Relation = nil
if err := r.AddGraph(ctx, ns, []*types.GraphData{graph}); err != nil {
t.Fatal(err)
}
if err := tx.Exec(`UPDATE chunks SET content='The original relation is no longer supported.' WHERE id='chunk' AND tenant_id=7`).Error; err != nil {
t.Fatal(err)
}
found, err := r.SearchNode(ctx, ns, []string{"Trust"})
if err != nil || len(found.Node) != 2 || len(found.Relation) != 0 {
t.Fatalf("stale relation remained available through separately supported entities: %#v %v", found, err)
}
}
func TestNativeCopiedVersionDeletionUsesRemappedSources(t *testing.T) {
tx, r, ctx := nativeWriteFixture(t)
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
if err := r.AddGraph(ctx, ns, []*types.GraphData{testGraph("chunk")}); err != nil {
t.Fatal(err)
}
for _, sql := range []string{
`INSERT INTO knowledges VALUES('copy-doc',7,'other-kb','copy.md',NULL)`,
`INSERT INTO chunks VALUES('copy-chunk',7,'other-kb','copy-doc','Evidence about trust and learning.',NULL,true)`,
`INSERT INTO native_graph_records SELECT tenant_id,'other-kb',source_table,legacy_id,graph_version_id,payload,source_snapshot_sha256,imported_at FROM native_graph_records WHERE tenant_id=7 AND knowledge_base_id='kb'`,
`INSERT INTO native_graph_sources SELECT tenant_id,'other-kb',source_chunk_id,source_file,source_content_hash,excerpt,'copy-doc','copy-chunk',native_content_sha256,mapping_state FROM native_graph_sources WHERE tenant_id=7 AND knowledge_base_id='kb'`,
} {
if err := tx.Exec(sql).Error; err != nil {
t.Fatal(err)
}
}
if err := r.DelGraph(ctx, []types.NameSpace{{KnowledgeBase: "other-kb", Knowledge: "copy-doc"}}); err != nil {
t.Fatal(err)
}
copy, err := r.Browse(ctx, "other-kb", "", 0, 0)
if err != nil || len(copy.Counts) != 0 || copy.Sources != 0 {
t.Fatalf("copied graph was not deleted via its current document: %#v %v", copy, err)
}
original, err := r.Browse(ctx, "kb", "", 0, 0)
if err != nil || original.Counts["knowledgeentity"] != 2 || original.Sources != 1 {
t.Fatal("copy deletion changed original graph", err)
}
}
func nativeGraphWriteState(t *testing.T, tx *gorm.DB) string {
t.Helper()
var state string
if err := tx.Raw(`SELECT jsonb_build_object(
'records',COALESCE((SELECT jsonb_agg(to_jsonb(r) ORDER BY tenant_id,knowledge_base_id,source_table,legacy_id) FROM native_graph_records r),'[]'::jsonb),
'sources',COALESCE((SELECT jsonb_agg(to_jsonb(s) ORDER BY tenant_id,knowledge_base_id,source_chunk_id) FROM native_graph_sources s),'[]'::jsonb))::text`).Scan(&state).Error; err != nil {
t.Fatal(err)
}
return state
}
func assertNativeTrustAttributes(t *testing.T, tx *gorm.DB, want []string) string {
t.Helper()
id, _ := nativeID("doc", "entity", "trust")
var record struct {
Payload json.RawMessage
Hash string
State string
}
if err := tx.Raw(`SELECT payload,source_snapshot_sha256 AS hash,to_jsonb(r)::text AS state FROM native_graph_records r
WHERE tenant_id=7 AND knowledge_base_id='kb' AND source_table='knowledgeentity' AND legacy_id=?`, id).Scan(&record).Error; err != nil {
t.Fatal(err)
}
var entity struct {
Attributes []string `json:"attributes"`
Type string `json:"entity_type"`
}
if err := json.Unmarshal(record.Payload, &entity); err != nil {
t.Fatal(err)
}
wantType := "concept"
if len(want) > 0 {
wantType = want[0]
}
if !slices.Equal(entity.Attributes, want) || entity.Type != wantType {
t.Fatalf("unexpected aggregate attributes/type: %#v; want %v / %s", entity, want, wantType)
}
var payload map[string]any
if err := json.Unmarshal(record.Payload, &payload); err != nil {
t.Fatal(err)
}
encoded, err := json.Marshal(payload)
if err != nil || record.Hash != contentHash(string(encoded)) {
t.Fatalf("entity payload hash is stale: %s %v", record.Hash, err)
}
return record.State
}
func TestNativeSharedEntityAttributesFollowTheirSources(t *testing.T) {
for _, siblingHasAttributes := range []bool{false, true} {
t.Run(fmt.Sprintf("sibling_attributes_%t", siblingHasAttributes), func(t *testing.T) {
tx, r, ctx := nativeWriteFixture(t)
if err := tx.Exec(`INSERT INTO chunks VALUES('sibling',7,'kb','doc','Another source about trust and learning.',NULL,true)`).Error; err != nil {
t.Fatal(err)
}
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
original := testGraph("chunk")
if err := r.AddGraph(ctx, ns, []*types.GraphData{original}); err != nil {
t.Fatal(err)
}
before := assertNativeTrustAttributes(t, tx, original.Node[0].Attributes)
sibling := testGraph("sibling")
if !siblingHasAttributes {
sibling.Node[0].Attributes = nil
}
if err := r.AddGraph(ctx, ns, []*types.GraphData{sibling}); err != nil {
t.Fatal(err)
}
if after := assertNativeTrustAttributes(t, tx, original.Node[0].Attributes); after != before {
t.Fatal("a second source changed the existing entity row")
}
var sourceAttributes string
id, _ := nativeID("doc", "entity", "trust")
if err := tx.Raw(`SELECT payload->'entity_attributes' FROM native_graph_records WHERE source_table='knowledgegraphevidence'
AND payload->>'subject_type'='entity' AND payload->>'subject_id'=? AND payload->>'source_chunk_id'='native:sibling'`, fmt.Sprint(id)).Scan(&sourceAttributes).Error; err != nil {
t.Fatal(err)
}
if !siblingHasAttributes && sourceAttributes != "[]" {
t.Fatalf("bare source inherited aggregate attributes: %s", sourceAttributes)
}
original.Node[0].Attributes = nil
if err := r.AddGraph(ctx, ns, []*types.GraphData{original}); err != nil {
t.Fatal(err)
}
assertNativeTrustAttributes(t, tx, sibling.Node[0].Attributes)
sibling.Node[0].Attributes = []string{}
if err := r.AddGraph(ctx, ns, []*types.GraphData{sibling}); err != nil {
t.Fatal(err)
}
assertNativeTrustAttributes(t, tx, nil)
})
}
}
func TestNativeSharedEntityAttributeConflictRollsBack(t *testing.T) {
tx, r, ctx := nativeWriteFixture(t)
if err := tx.Exec(`INSERT INTO chunks VALUES('sibling',7,'kb','doc','Another source about trust and learning.',NULL,true)`).Error; err != nil {
t.Fatal(err)
}
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
if err := r.AddGraph(ctx, ns, []*types.GraphData{testGraph("chunk")}); err != nil {
t.Fatal(err)
}
before := nativeGraphWriteState(t, tx)
conflicting := testGraph("sibling")
conflicting.Node[0].Attributes = []string{"concept", "Never earned through reliable evidence"}
if err := r.AddGraph(ctx, ns, []*types.GraphData{conflicting}); err == nil || !strings.Contains(err.Error(), "attributes conflict") {
t.Fatalf("conflicting source did not require review: %v", err)
}
if after := nativeGraphWriteState(t, tx); after != before {
t.Fatal("failed conflict changed records, evidence, mapping or checkpoint")
}
}
func TestNativeSharedEntityAttributesAreRemovedWithTheirNode(t *testing.T) {
for _, siblingHasAttributes := range []bool{false, true} {
t.Run(fmt.Sprintf("sibling_attributes_%t", siblingHasAttributes), func(t *testing.T) {
tx, r, ctx := nativeWriteFixture(t)
if err := tx.Exec(`INSERT INTO chunks VALUES('sibling',7,'kb','doc','Another source about trust and learning.',NULL,true)`).Error; err != nil {
t.Fatal(err)
}
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
original, sibling := testGraph("chunk"), testGraph("sibling")
if !siblingHasAttributes {
sibling.Node[0].Attributes = nil
}
if err := r.AddGraph(ctx, ns, []*types.GraphData{original, sibling}); err != nil {
t.Fatal(err)
}
original.Node[0].Name = "Practice"
original.Relation[0].Node1 = "Practice"
chunk := &types.Chunk{ID: "chunk", TenantID: 7, KnowledgeBaseID: "kb", KnowledgeID: "doc", Content: "Evidence about trust and learning."}
if err := r.AddChunkGraph(ctx, ns, chunk, original); err != nil {
t.Fatal(err)
}
assertNativeTrustAttributes(t, tx, sibling.Node[0].Attributes)
sibling.Node[0].Name = "Practice"
sibling.Node[0].Attributes = original.Node[0].Attributes
sibling.Relation[0].Node1 = "Practice"
if err := r.AddGraph(ctx, ns, []*types.GraphData{sibling}); err != nil {
t.Fatal(err)
}
id, _ := nativeID("doc", "entity", "trust")
var count int64
if err := tx.Table("native_graph_records").Where("source_table='knowledgeentity' AND legacy_id=?", id).Count(&count).Error; err != nil || count != 0 {
t.Fatalf("entity without surviving sources was not pruned: %d %v", count, err)
}
})
}
}
func TestNativeLegacyEntityAttributeProvenanceRequiresReview(t *testing.T) {
tx, r, ctx := nativeWriteFixture(t)
if err := tx.Exec(`INSERT INTO chunks VALUES('sibling',7,'kb','doc','Another source about trust and learning.',NULL,true)`).Error; err != nil {
t.Fatal(err)
}
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
if err := r.AddGraph(ctx, ns, []*types.GraphData{testGraph("chunk")}); err != nil {
t.Fatal(err)
}
id, _ := nativeID("doc", "entity", "trust")
if err := tx.Exec(`UPDATE native_graph_records SET payload=payload-'entity_attributes' WHERE source_table='knowledgegraphevidence'
AND payload->>'subject_type'='entity' AND payload->>'subject_id'=?`, fmt.Sprint(id)).Error; err != nil {
t.Fatal(err)
}
before := nativeGraphWriteState(t, tx)
bare := testGraph("sibling")
bare.Node[0].Attributes = nil
if err := r.AddGraph(ctx, ns, []*types.GraphData{bare}); err == nil || !strings.Contains(err.Error(), "provenance is unavailable") {
t.Fatalf("legacy unknown provenance treated as known empty: %v", err)
}
if after := nativeGraphWriteState(t, tx); after != before {
t.Fatal("unknown provenance failure changed prior graph")
}
// A matching explicit source can coexist with unknown older evidence.
if err := r.AddGraph(ctx, ns, []*types.GraphData{testGraph("sibling")}); err != nil {
t.Fatal(err)
}
original := testGraph("chunk")
original.Node[0].Attributes = nil
if err := r.AddGraph(ctx, ns, []*types.GraphData{original}); err != nil {
t.Fatal(err)
}
assertNativeTrustAttributes(t, tx, testGraph("sibling").Node[0].Attributes)
if err := r.AddGraph(ctx, ns, []*types.GraphData{bare}); err != nil {
t.Fatal(err)
}
assertNativeTrustAttributes(t, tx, nil)
}
func TestNativeNullEntityAttributeEvidenceRequiresReview(t *testing.T) {
for _, snapshot := range []string{"null", "[null]"} {
t.Run(snapshot, func(t *testing.T) {
tx, r, ctx := nativeWriteFixture(t)
if err := tx.Exec(`INSERT INTO chunks VALUES('sibling',7,'kb','doc','Another source about trust and learning.',NULL,true)`).Error; err != nil {
t.Fatal(err)
}
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
if err := r.AddGraph(ctx, ns, []*types.GraphData{testGraph("chunk")}); err != nil {
t.Fatal(err)
}
id, _ := nativeID("doc", "entity", "trust")
if err := tx.Exec(`UPDATE native_graph_records SET payload=jsonb_set(payload,'{entity_attributes}',?::jsonb)
WHERE source_table='knowledgegraphevidence' AND payload->>'subject_type'='entity' AND payload->>'subject_id'=?`, snapshot, fmt.Sprint(id)).Error; err != nil {
t.Fatal(err)
}
before := nativeGraphWriteState(t, tx)
bare := testGraph("sibling")
bare.Node[0].Attributes = nil
if err := r.AddGraph(ctx, ns, []*types.GraphData{bare}); err == nil {
t.Fatal("null attribute evidence was accepted as known empty")
}
if after := nativeGraphWriteState(t, tx); after != before {
t.Fatal("null attribute evidence failure changed prior graph")
}
})
}
}
func TestNativeSharedEntityIgnoresInvalidAttributeEvidence(t *testing.T) {
for name, invalidate := range map[string]string{
"changed source": `UPDATE chunks SET content='Changed source.' WHERE id='chunk' AND tenant_id=7`,
"disabled source": `UPDATE chunks SET is_enabled=false WHERE id='chunk' AND tenant_id=7`,
"deleted source": `UPDATE chunks SET deleted_at=now() WHERE id='chunk' AND tenant_id=7`,
"wrong evidence hash": `UPDATE native_graph_records SET payload=jsonb_set(payload,'{source_content_hash}',to_jsonb(repeat('a',64))) WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
"wrong row hash": `UPDATE native_graph_records SET source_snapshot_sha256=repeat('a',64) WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
"inactive evidence": `UPDATE native_graph_records SET payload=jsonb_set(payload,'{status}','"inactive"'::jsonb) WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
"wrong origin": `UPDATE native_graph_records SET payload=jsonb_set(payload,'{origin}','"foreign"'::jsonb) WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
"wrong tenant": `UPDATE native_graph_records SET tenant_id=8 WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
"wrong knowledge base": `UPDATE native_graph_records SET knowledge_base_id='other-kb' WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
"wrong document": `UPDATE native_graph_records SET payload=jsonb_set(payload,'{native_knowledge_id}','"doc2"'::jsonb) WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
"wrong mapping document": `UPDATE native_graph_sources SET native_knowledge_id='doc2' WHERE source_chunk_id='native:chunk'`,
"inactive mapping": `UPDATE native_graph_sources SET mapping_state='stale' WHERE source_chunk_id='native:chunk'`,
} {
t.Run(name, func(t *testing.T) {
tx, r, ctx := nativeWriteFixture(t)
if err := tx.Exec(`INSERT INTO chunks VALUES('sibling',7,'kb','doc','Another source about trust and learning.',NULL,true)`).Error; err != nil {
t.Fatal(err)
}
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
if err := r.AddGraph(ctx, ns, []*types.GraphData{testGraph("chunk")}); err != nil {
t.Fatal(err)
}
if err := tx.Exec(invalidate).Error; err != nil {
t.Fatal(err)
}
bare := testGraph("sibling")
bare.Node[0].Attributes = nil
if err := r.AddGraph(ctx, ns, []*types.GraphData{bare}); err != nil {
t.Fatal(err)
}
assertNativeTrustAttributes(t, tx, nil)
})
}
}
func TestNativeEntityAttributeBatchOrderAndDuplicates(t *testing.T) {
for _, reverse := range []bool{false, true} {
t.Run(fmt.Sprintf("reverse_%t", reverse), func(t *testing.T) {
tx, r, ctx := nativeWriteFixture(t)
if err := tx.Exec(`INSERT INTO chunks VALUES('sibling',7,'kb','doc','Another source about trust and learning.',NULL,true)`).Error; err != nil {
t.Fatal(err)
}
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
first, second := testGraph("chunk"), testGraph("sibling")
first.Node = append(first.Node, &types.GraphNode{Name: " trust ", Chunks: []string{"chunk", "chunk"}, Attributes: first.Node[0].Attributes})
second.Node[0].Name = "TRUST"
second.Node[0].Attributes = nil
graphs := []*types.GraphData{first, second}
if reverse {
slices.Reverse(graphs)
}
if err := r.AddGraph(ctx, ns, graphs); err != nil {
t.Fatal(err)
}
entityBefore := assertNativeTrustAttributes(t, tx, first.Node[0].Attributes)
slices.Reverse(graphs)
if err := r.AddGraph(ctx, ns, graphs); err != nil {
t.Fatal(err)
}
if entityAfter := assertNativeTrustAttributes(t, tx, first.Node[0].Attributes); entityAfter != entityBefore {
t.Fatal("batch order changed the entity payload or hash")
}
var count int
if err := tx.Raw(`SELECT (payload->>'entity_count')::int FROM native_graph_records WHERE source_table='knowledgegraphchunkstate' AND payload->>'native_chunk_id'='chunk'`).Scan(&count).Error; err != nil || count != 2 {
t.Fatalf("duplicate entity evidence inflated checkpoint: %d %v", count, err)
}
before := nativeGraphWriteState(t, tx)
first.Node[2].Attributes = []string{"conflicting duplicate"}
if err := r.AddGraph(ctx, ns, graphs); err == nil || !strings.Contains(err.Error(), "attributes conflict") {
t.Fatalf("duplicate source attributes silently overwritten: %v", err)
}
if after := nativeGraphWriteState(t, tx); after != before {
t.Fatal("failed duplicate source write changed the committed graph")
}
})
}
}
func TestNativeLegacyEntityAttributeBatchOrder(t *testing.T) {
for _, reverse := range []bool{false, true} {
t.Run(fmt.Sprintf("reverse_%t", reverse), func(t *testing.T) {
tx, r, ctx := nativeWriteFixture(t)
if err := tx.Exec(`INSERT INTO chunks VALUES('bare',7,'kb','doc','A bare reference to trust.',NULL,true),('explicit',7,'kb','doc','An explicit source about trust.',NULL,true)`).Error; err != nil {
t.Fatal(err)
}
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
original := testGraph("chunk")
if err := r.AddGraph(ctx, ns, []*types.GraphData{original}); err != nil {
t.Fatal(err)
}
if err := tx.Exec(`UPDATE native_graph_records SET payload=payload-'entity_attributes' WHERE source_table='knowledgegraphevidence' AND payload->>'subject_type'='entity'`).Error; err != nil {
t.Fatal(err)
}
bare, explicit := testGraph("bare"), testGraph("explicit")
bare.Node[0].Attributes = nil
graphs := []*types.GraphData{bare, explicit}
if reverse {
slices.Reverse(graphs)
}
if err := r.AddGraph(ctx, ns, graphs); err != nil {
t.Fatal(err)
}
assertNativeTrustAttributes(t, tx, original.Node[0].Attributes)
original.Node[0].Attributes = nil
explicit.Node[0].Attributes = nil
graphs = append(graphs, original)
if err := r.AddGraph(ctx, ns, graphs); err != nil {
t.Fatal(err)
}
assertNativeTrustAttributes(t, tx, nil)
})
}
}
@@ -0,0 +1,304 @@
package milvus
import (
"encoding/json"
"fmt"
"maps"
"reflect"
"strings"
"time"
)
const (
// operatorAnd is the "and" operator.
operatorAnd = "and"
// operatorOr is the "or" operator.
operatorOr = "or"
// operatorEqual is the "equal" operator.
operatorEqual = "eq"
// operatorNotEqual is the "not equal" operator.
operatorNotEqual = "ne"
// operatorGreaterThan is the "greater than" operator.
operatorGreaterThan = "gt"
// operatorGreaterThanOrEqual is the "greater than or equal" operator.
operatorGreaterThanOrEqual = "gte"
// operatorLessThan is the "less than" operator.
operatorLessThan = "lt"
// operatorLessThanOrEqual is the "less than or equal" operator.
operatorLessThanOrEqual = "lte"
// operatorIn is the "in" operator.
operatorIn = "in"
// operatorNotIn is the "not in" operator.
operatorNotIn = "not in"
// operatorLike is the "contains" operator.
operatorLike = "like"
// operatorNotLike is the "not contains" operator.
operatorNotLike = "not like"
// operatorBetween is the "between" operator.
operatorBetween = "between"
)
var comparisonOperators = map[string]string{
operatorEqual: "==",
operatorNotEqual: "!=",
operatorGreaterThan: ">",
operatorGreaterThanOrEqual: ">=",
operatorLessThan: "<",
operatorLessThanOrEqual: "<=",
operatorLike: "like",
operatorNotLike: "not like",
}
type convertResult struct {
exprStr string
params map[string]any
}
type filter struct{}
func (c *filter) Convert(cond *universalFilterCondition) (*convertResult, error) {
var counter int
return c.convertCondition(cond, &counter)
}
func (c *filter) convertComparisonCondition(
cond *universalFilterCondition,
counter *int,
) (*convertResult, error) {
condField := cond.Field
if condField == "" || cond.Value == nil {
return nil, fmt.Errorf("milvus filter condition is nil")
}
operator, ok := comparisonOperators[cond.Operator]
if !ok {
return nil, fmt.Errorf("unsupported comparison operator: %s", cond.Operator)
}
paramName := c.convertParamName(cond.Field, counter)
return &convertResult{
exprStr: fmt.Sprintf("%s %s {%s}", condField, operator, paramName),
params: map[string]any{paramName: cond.Value},
}, nil
}
func (c *filter) convertLogicalCondition(
cond *universalFilterCondition,
counter *int,
) (*convertResult, error) {
if cond.Value == nil {
return nil, fmt.Errorf("milvus filter condition is nil")
}
conds, ok := cond.Value.([]*universalFilterCondition)
if !ok {
return nil, fmt.Errorf("invalid logical condition value type")
}
var condResult *convertResult
for _, childCond := range conds {
childRes, err := c.convertCondition(childCond, counter)
if err != nil {
return nil, err
}
if childRes == nil || childRes.exprStr == "" {
continue
}
if condResult == nil {
condResult = childRes
continue
}
condResult.exprStr = fmt.Sprintf(
"(%s) %s (%s)",
condResult.exprStr,
strings.ToLower(cond.Operator),
childRes.exprStr,
)
maps.Copy(condResult.params, childRes.params)
}
if condResult == nil {
return nil, fmt.Errorf("empty logical condition")
}
return condResult, nil
}
func (c *filter) convertCondition(
cond *universalFilterCondition,
counter *int,
) (*convertResult, error) {
if cond == nil {
return nil, fmt.Errorf("milvus filter condition is nil")
}
switch cond.Operator {
case operatorEqual, operatorNotEqual, operatorGreaterThan,
operatorGreaterThanOrEqual, operatorLessThan,
operatorLessThanOrEqual, operatorLike, operatorNotLike:
return c.convertComparisonCondition(cond, counter)
case operatorAnd, operatorOr:
return c.convertLogicalCondition(cond, counter)
case operatorIn, operatorNotIn:
return c.convertInCondition(cond, counter)
case operatorBetween:
return c.convertBetweenCondition(cond, counter)
default:
return nil, fmt.Errorf("unsupported operator: %v", cond.Operator)
}
}
func (c *filter) convertInCondition(
cond *universalFilterCondition,
counter *int,
) (*convertResult, error) {
condField := cond.Field
if condField == "" || cond.Value == nil {
return nil, fmt.Errorf("milvus filter condition is nil")
}
s := reflect.ValueOf(cond.Value)
if s.Kind() != reflect.Slice || s.Len() <= 0 {
return nil, fmt.Errorf("in operator value must be a slice with at least one value: %v", cond.Value)
}
paramName := c.convertParamName(cond.Field, counter)
return &convertResult{
exprStr: fmt.Sprintf("%s %s {%s}", condField, strings.ToLower(cond.Operator), paramName),
params: map[string]any{paramName: cond.Value},
}, nil
}
func (c *filter) convertBetweenCondition(
cond *universalFilterCondition,
counter *int,
) (*convertResult, error) {
condField := cond.Field
if condField == "" || cond.Value == nil {
return nil, fmt.Errorf("milvus filter condition is nil")
}
value := reflect.ValueOf(cond.Value)
if value.Kind() != reflect.Slice || value.Len() != 2 {
return nil, fmt.Errorf("between operator value must be a slice with two elements: %v", cond.Value)
}
paramBase := c.convertParamName(cond.Field, counter)
paramName1 := fmt.Sprintf("%s_%d", paramBase, 0)
paramName2 := fmt.Sprintf("%s_%d", paramBase, 1)
return &convertResult{
exprStr: fmt.Sprintf("%s >= {%s} and %s <= {%s}", condField, paramName1, condField, paramName2),
params: map[string]any{
paramName1: value.Index(0).Interface(),
paramName2: value.Index(1).Interface(),
},
}, nil
}
func formatValue(value any) string {
switch v := value.(type) {
case string:
return fmt.Sprintf("\"%s\"", escapeDoubleQuotes(v))
case int, int8, int16, int32, int64:
return fmt.Sprintf("%d", v)
case uint, uint8, uint16, uint32, uint64:
return fmt.Sprintf("%d", v)
case float32, float64:
return fmt.Sprintf("%v", v)
case bool:
if v {
return "true"
}
return "false"
case time.Time:
return fmt.Sprintf("%d", v.Unix())
default:
return fmt.Sprintf("\"%v\"", value)
}
}
// escapeDoubleQuotes escapes double quotes in a string for use in Milvus expressions.
func escapeDoubleQuotes(s string) string {
return strings.ReplaceAll(s, "\"", "\\\"")
}
// convertParamName converts field name to a valid Milvus template parameter name.
// Milvus template parameters don't support '.' character, so we replace it with '_'.
func (c *filter) convertParamName(field string, counter *int) string {
*counter++
return fmt.Sprintf("%s_%d", strings.ReplaceAll(field, ".", "_"), *counter)
}
type universalFilterCondition struct {
Field string `json:"field,omitempty" jsonschema:"description=The metadata field to filter on (required for comparison operators)"`
Operator string `json:"operator" jsonschema:"description=The operator to use,enum=eq,enum=ne,enum=gt,enum=gte,enum=lt,enum=lte,enum=in,enum=not in,enum=like,enum=not like,enum=between,enum=and,enum=or"`
Value any `json:"value,omitempty" jsonschema:"description=The value to compare against (for comparison operators) or array of sub-conditions (for logical operators and/or)"`
}
func (c *universalFilterCondition) UnmarshalJSON(data []byte) error {
type Alias struct {
Field string `json:"field,omitempty"`
Operator string `json:"operator"`
Value any `json:"value,omitempty"`
}
var aux Alias
if err := json.Unmarshal(data, &aux); err != nil {
return err
}
c.Field = aux.Field
c.Operator = strings.ToLower(aux.Operator)
// Handle logical operators (and/or) - Value should be []*UniversalFilterCondition
if c.Operator == operatorAnd || c.Operator == operatorOr {
// Value can be an array of conditions
valueSlice, ok := aux.Value.([]any)
if !ok {
return fmt.Errorf("logical operator %s requires an array of conditions", c.Operator)
}
conditions := make([]*universalFilterCondition, 0, len(valueSlice))
for i, v := range valueSlice {
condBytes, err := json.Marshal(v)
if err != nil {
return fmt.Errorf("failed to marshal condition at index %d: %w", i, err)
}
var cond universalFilterCondition
if err := json.Unmarshal(condBytes, &cond); err != nil {
return fmt.Errorf("failed to unmarshal condition at index %d: %w", i, err)
}
conditions = append(conditions, &cond)
}
c.Value = conditions
} else {
c.Value = aux.Value
}
return nil
}
// MarshalJSON implements custom JSON marshaling for UniversalFilterCondition.
func (c *universalFilterCondition) MarshalJSON() ([]byte, error) {
type Alias struct {
Field string `json:"field,omitempty"`
Operator string `json:"operator"`
Value any `json:"value,omitempty"`
}
return json.Marshal(&Alias{
Field: c.Field,
Operator: c.Operator,
Value: c.Value,
})
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,243 @@
package milvus
import (
"context"
"errors"
"testing"
"github.com/milvus-io/milvus/client/v2/column"
"github.com/milvus-io/milvus/client/v2/entity"
client "github.com/milvus-io/milvus/client/v2/milvusclient"
"github.com/stretchr/testify/require"
"github.com/Tencent/WeKnora/internal/types"
)
func TestContentFieldUsesChineseAnalyzer(t *testing.T) {
field := newContentField().ProtoMessage()
require.Equal(t, fieldContent, field.GetName())
require.Equal(t, int32(entity.FieldTypeVarChar), int32(field.GetDataType()))
params := make(map[string]string)
for _, param := range field.GetTypeParams() {
params[param.GetKey()] = param.GetValue()
}
require.Equal(t, "65535", params["max_length"])
require.Equal(t, "true", params["enable_analyzer"])
require.Equal(t, "true", params["enable_match"])
require.JSONEq(t, `{"type":"chinese"}`, params["analyzer_params"])
}
func TestUpdateChunkEnabledStatusInCollectionSkipsEmptyChunkIDs(t *testing.T) {
repo := &milvusRepository{}
require.NoError(t, repo.updateChunkEnabledStatusInCollection(
context.Background(),
"weknora_embeddings_1024",
nil,
false,
))
require.NoError(t, repo.updateChunkEnabledStatusInCollection(
context.Background(),
"weknora_embeddings_1024",
[]string{},
true,
))
}
func TestUpdateChunkEnabledStatusInCollectionsPropagatesFailure(t *testing.T) {
wantErr := errors.New("upsert failed")
err := updateChunkEnabledStatusInCollections(
context.Background(),
[]string{"other_collection", "weknora_embeddings_1024"},
"weknora_embeddings",
nil,
[]string{"chunk-1"},
func(_ context.Context, collection string, _ []string, enabled bool) error {
if collection == "weknora_embeddings_1024" && !enabled {
return wantErr
}
return nil
},
)
require.ErrorIs(t, err, wantErr)
}
func TestMatchesCollection(t *testing.T) {
for _, name := range []string{"weknora_embeddings_1", "weknora_embeddings_1024"} {
require.True(t, matchesCollection(name, "weknora_embeddings"), name)
}
for _, name := range []string{
"weknora_embeddings", "weknora_embeddings_", "weknora_embeddings_0",
"weknora_embeddings_-1", "weknora_embeddings_+1", "weknora_embeddings_01024",
"weknora_embeddings_1024_backup", "weknora_embeddings_backup_1024",
"weknora_embeddings_other_tenant_1024", "weknora_embeddings2_1024",
"weknora_embeddings_9999999999999999999999999", "other_1024", "_1024",
} {
require.False(t, matchesCollection(name, "weknora_embeddings"), name)
}
require.False(t, matchesCollection("_1024", ""))
}
func TestUpdateChunkEnabledStatusDoesNotCrossNamespaces(t *testing.T) {
var touched []string
err := updateChunkEnabledStatusInCollections(
context.Background(),
[]string{"weknora_embeddings_1024_backup", "weknora_embeddings_other_1024", "weknora_embeddings_1024"},
"weknora_embeddings", nil, []string{"chunk-1"},
func(_ context.Context, collection string, _ []string, _ bool) error {
touched = append(touched, collection)
return nil
},
)
require.NoError(t, err)
require.Equal(t, []string{"weknora_embeddings_1024", "weknora_embeddings_1024"}, touched)
}
func keywordResult(chunkID string, score float32) []client.ResultSet {
return []client.ResultSet{{
ResultCount: 1,
Scores: []float32{score},
Fields: client.DataSet{
column.NewColumnVarChar(fieldID, []string{"index-" + chunkID}),
column.NewColumnVarChar(fieldChunkID, []string{chunkID}),
column.NewColumnVarChar(fieldContent, []string{"test content"}),
},
}}
}
func TestKeywordsRetrieveReturnsAllCollectionFailures(t *testing.T) {
repo := &milvusRepository{collectionBaseName: "weknora_embeddings"}
missingSparseField := errors.New("field content_sparse does not exist")
conversionFailure := errors.New("cannot decode search result")
var searched []string
results, err := repo.keywordsRetrieveInCollections(context.Background(), types.RetrieveParams{
Query: "test query", TopK: 10,
}, []string{"weknora_embeddings_384", "weknora_embeddings_1024", "weknora_embeddings_backup_1024"},
func(_ context.Context, option client.SearchOption) ([]client.ResultSet, error) {
request, err := option.Request()
require.NoError(t, err)
searched = append(searched, request.CollectionName)
if request.CollectionName == "weknora_embeddings_384" {
return nil, missingSparseField
}
return []client.ResultSet{{Err: conversionFailure}}, nil
})
require.Nil(t, results)
require.ErrorIs(t, err, missingSparseField)
require.ErrorIs(t, err, conversionFailure)
require.ErrorContains(t, err, "all 2 matching Milvus collections")
require.ErrorContains(t, err, "weknora_embeddings_384")
require.ErrorContains(t, err, "weknora_embeddings_1024")
require.Equal(t, []string{"weknora_embeddings_384", "weknora_embeddings_1024"}, searched)
}
func TestKeywordsRetrievePreservesBM25AndRanksAcrossCollections(t *testing.T) {
repo := &milvusRepository{collectionBaseName: "weknora_embeddings"}
results, err := repo.keywordsRetrieveInCollections(context.Background(), types.RetrieveParams{
Query: "test query", TopK: 1, KnowledgeBaseIDs: []string{"kb-allowed"},
TagIDs: []string{"tag-allowed"}, ExcludeChunkIDs: []string{"chunk-excluded"},
}, []string{"weknora_embeddings_384", "weknora_embeddings_1024", "weknora_embeddings_other_1536"},
func(_ context.Context, option client.SearchOption) ([]client.ResultSet, error) {
request, err := option.Request()
require.NoError(t, err)
require.Contains(t, request.Dsl, fieldKnowledgeBaseID)
require.Contains(t, request.Dsl, fieldTagID)
require.Contains(t, request.Dsl, fieldChunkID)
require.Contains(t, request.Dsl, fieldIsEnabled)
require.NotEmpty(t, request.ExprTemplateValues)
for _, param := range request.SearchParams {
if param.Key == "anns_field" {
require.Equal(t, fieldContentSparse, param.Value)
}
}
switch request.CollectionName {
case "weknora_embeddings_384":
return keywordResult("low-score", 2.5), nil
case "weknora_embeddings_1024":
return keywordResult("high-score", 17.25), nil
default:
t.Fatalf("searched unrelated collection %s", request.CollectionName)
return nil, nil
}
})
require.NoError(t, err)
require.Len(t, results, 1)
require.Equal(t, types.KeywordsRetrieverType, results[0].RetrieverType)
require.Len(t, results[0].Results, 1)
require.Equal(t, "high-score", results[0].Results[0].ChunkID)
require.Equal(t, 17.25, results[0].Results[0].Score)
require.Equal(t, types.MatchTypeKeywords, results[0].Results[0].MatchType)
}
func TestKeywordsRetrieveKeepsSuccessfulCollections(t *testing.T) {
for _, empty := range []bool{false, true} {
t.Run(map[bool]string{false: "results", true: "empty"}[empty], func(t *testing.T) {
repo := &milvusRepository{collectionBaseName: "weknora_embeddings"}
results, err := repo.keywordsRetrieveInCollections(context.Background(), types.RetrieveParams{
Query: "query", TopK: 5,
}, []string{"weknora_embeddings_384", "weknora_embeddings_1024"},
func(_ context.Context, option client.SearchOption) ([]client.ResultSet, error) {
request, err := option.Request()
require.NoError(t, err)
if request.CollectionName == "weknora_embeddings_384" {
return nil, errors.New("missing sparse field")
}
if empty {
return nil, nil
}
return keywordResult("available", 8), nil
})
require.NoError(t, err)
require.Len(t, results, 1)
if empty {
require.Empty(t, results[0].Results)
} else {
require.Len(t, results[0].Results, 1)
require.Equal(t, "available", results[0].Results[0].ChunkID)
}
})
}
}
func TestKeywordsRetrieveRejectsMissingScores(t *testing.T) {
repo := &milvusRepository{collectionBaseName: "weknora_embeddings"}
results, err := repo.keywordsRetrieveInCollections(context.Background(), types.RetrieveParams{
Query: "query", TopK: 5,
}, []string{"weknora_embeddings_384"},
func(context.Context, client.SearchOption) ([]client.ResultSet, error) {
result := keywordResult("missing-score", 0)
result[0].Scores = nil
return result, nil
})
require.Nil(t, results)
require.ErrorContains(t, err, "expected 1 BM25 scores, received 0")
}
func TestKeywordsRetrieveNoMatchingCollections(t *testing.T) {
repo := &milvusRepository{collectionBaseName: "weknora_embeddings"}
results, err := repo.keywordsRetrieveInCollections(context.Background(), types.RetrieveParams{
Query: "query", TopK: 5,
}, []string{"weknora_embeddings_backup_384"},
func(context.Context, client.SearchOption) ([]client.ResultSet, error) {
t.Fatal("must not search unrelated collection")
return nil, nil
})
require.NoError(t, err)
require.Len(t, results, 1)
require.Empty(t, results[0].Results)
}
func TestKeywordsRetrieveHonorsCancellation(t *testing.T) {
repo := &milvusRepository{collectionBaseName: "weknora_embeddings"}
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
results, err := repo.keywordsRetrieveInCollections(ctx, types.RetrieveParams{
Query: "query", TopK: 5,
}, []string{"weknora_embeddings_384", "weknora_embeddings_1024"},
func(context.Context, client.SearchOption) ([]client.ResultSet, error) {
cancel()
return keywordResult("cancelled", 8), nil
})
require.Nil(t, results)
require.ErrorIs(t, err, context.Canceled)
}
@@ -0,0 +1,37 @@
package milvus
import (
"sync"
"github.com/milvus-io/milvus/client/v2/entity"
client "github.com/milvus-io/milvus/client/v2/milvusclient"
)
type milvusRepository struct {
filter
client *client.Client
collectionBaseName string
metricType entity.MetricType
shardsNum int // 0 = use Milvus default (1)
replicaNumber int // 0 = use Milvus default (1); set at LoadCollection time
// Cache for initialized collections (dimension -> true)
initializedCollections sync.Map
}
type MilvusVectorEmbedding struct {
ID string `json:"id"`
Content string `json:"content"`
SourceID string `json:"source_id"`
SourceType int `json:"source_type"`
ChunkID string `json:"chunk_id"`
KnowledgeID string `json:"knowledge_id"`
KnowledgeBaseID string `json:"knowledge_base_id"`
TagID string `json:"tag_id"`
Embedding []float32 `json:"embedding"`
IsEnabled bool `json:"is_enabled"`
}
type MilvusVectorEmbeddingWithScore struct {
MilvusVectorEmbedding
Score float64
}
@@ -0,0 +1,29 @@
package milvus
import (
"reflect"
"testing"
"github.com/milvus-io/milvus/client/v2/column"
client "github.com/milvus-io/milvus/client/v2/milvusclient"
)
func TestCopyResultPreservesNativeFloatVector(t *testing.T) {
want := []float32{0.1, 0.25, -0.5}
for _, field := range []column.Column{
column.NewColumnFloatVector(fieldEmbedding, 3, [][]float32{want}),
column.NewColumnDoubleArray(fieldEmbedding, [][]float64{{0.1, 0.25, -0.5}}),
} {
result, _, err := convertResultSet([]client.ResultSet{{ResultCount: 1, Fields: client.DataSet{column.NewColumnVarChar(fieldID, []string{"source"}), field}}})
if err != nil || len(result) != 1 || !reflect.DeepEqual(result[0].Embedding, want) {
t.Fatalf("copy dropped vector from %T: %#v, %v", field, result, err)
}
}
}
func TestCopyResultRejectsUnknownVectorColumn(t *testing.T) {
_, _, err := convertResultSet([]client.ResultSet{{ResultCount: 1, Fields: client.DataSet{column.NewColumnVarChar(fieldEmbedding, []string{"invalid"})}}})
if err == nil {
t.Fatal("invalid vector was silently converted into a zero-dimensional embedding")
}
}
@@ -0,0 +1,236 @@
package neo4j
import (
"context"
"fmt"
"strings"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
"github.com/neo4j/neo4j-go-driver/v6/neo4j"
)
// Neo4jRepository is a repository for Neo4j
type Neo4jRepository struct {
driver neo4j.Driver
nodePrefix string
}
// NewNeo4jRepository creates a new Neo4j repository
func NewNeo4jRepository(driver neo4j.Driver) interfaces.RetrieveGraphRepository {
return &Neo4jRepository{driver: driver, nodePrefix: "ENTITY"}
}
// _remove_hyphen removes hyphens from a string
func _remove_hyphen(s string) string {
return strings.ReplaceAll(s, "-", "_")
}
// Labels returns the labels for a namespace
func (n *Neo4jRepository) Labels(namespace types.NameSpace) []string {
res := make([]string, 0)
for _, label := range namespace.Labels() {
res = append(res, n.nodePrefix+_remove_hyphen(label))
}
return res
}
// Label returns the label for a namespace
func (n *Neo4jRepository) Label(namespace types.NameSpace) string {
labels := n.Labels(namespace)
return strings.Join(labels, ":")
}
// AddGraph adds a graph to the Neo4j repository
func (n *Neo4jRepository) AddGraph(ctx context.Context, namespace types.NameSpace, graphs []*types.GraphData) error {
if n.driver == nil {
logger.Warnf(ctx, "NOT SUPPORT RETRIEVE GRAPH")
return nil
}
for _, graph := range graphs {
if err := n.addGraph(ctx, namespace, graph); err != nil {
return err
}
}
return nil
}
// addGraph adds a graph to the Neo4j repository
func (n *Neo4jRepository) addGraph(ctx context.Context, namespace types.NameSpace, graph *types.GraphData) error {
session := n.driver.NewSession(ctx, neo4j.SessionConfig{AccessMode: neo4j.AccessModeWrite})
defer session.Close(ctx)
_, err := session.ExecuteWrite(ctx, func(tx neo4j.ManagedTransaction) (interface{}, error) {
// Node import query
node_import_query := `
UNWIND $data AS row
CALL apoc.merge.node(row.labels, {name: row.name, kg: row.knowledge_id}, row.props, {}) YIELD node
SET node.chunks = apoc.coll.union(node.chunks, row.chunks)
RETURN distinct 'done' AS result
`
nodeData := []map[string]interface{}{}
for _, node := range graph.Node {
nodeData = append(nodeData, map[string]interface{}{
"name": node.Name,
"knowledge_id": namespace.Knowledge,
"props": map[string][]string{"attributes": node.Attributes},
"chunks": node.Chunks,
"labels": n.Labels(namespace),
})
}
if _, err := tx.Run(ctx, node_import_query, map[string]interface{}{"data": nodeData}); err != nil {
return nil, fmt.Errorf("failed to create nodes: %v", err)
}
// Relationship import query
rel_import_query := `
UNWIND $data AS row
CALL apoc.merge.node(row.source_labels, {name: row.source, kg: row.knowledge_id}, {}, {}) YIELD node as source
CALL apoc.merge.node(row.target_labels, {name: row.target, kg: row.knowledge_id}, {}, {}) YIELD node as target
CALL apoc.merge.relationship(source, row.type, {}, row.attributes, target) YIELD rel
RETURN distinct 'done'
`
relData := []map[string]interface{}{}
for _, rel := range graph.Relation {
relData = append(relData, map[string]interface{}{
"source": rel.Node1,
"target": rel.Node2,
"knowledge_id": namespace.Knowledge,
"type": rel.Type,
"source_labels": n.Labels(namespace),
"target_labels": n.Labels(namespace),
})
}
if _, err := tx.Run(ctx, rel_import_query, map[string]interface{}{"data": relData}); err != nil {
return nil, fmt.Errorf("failed to create relationships: %v", err)
}
return nil, nil
})
if err != nil {
logger.Errorf(ctx, "failed to add graph: %v", err)
return err
}
return nil
}
// DelGraph deletes a graph from the Neo4j repository
func (n *Neo4jRepository) DelGraph(ctx context.Context, namespaces []types.NameSpace) error {
if n.driver == nil {
logger.Warnf(ctx, "NOT SUPPORT RETRIEVE GRAPH")
return nil
}
session := n.driver.NewSession(ctx, neo4j.SessionConfig{AccessMode: neo4j.AccessModeWrite})
defer session.Close(ctx)
result, err := session.ExecuteWrite(ctx, func(tx neo4j.ManagedTransaction) (interface{}, error) {
for _, namespace := range namespaces {
labelExpr := n.Label(namespace)
deleteRelsQuery := `
CALL apoc.periodic.iterate(
"MATCH (n:` + labelExpr + ` {kg: $knowledge_id})-[r]-(m:` + labelExpr + ` {kg: $knowledge_id}) RETURN r",
"DELETE r",
{batchSize: 1000, parallel: true, params: {knowledge_id: $knowledge_id}}
) YIELD batches, total
RETURN total
`
if _, err := tx.Run(ctx, deleteRelsQuery, map[string]interface{}{"knowledge_id": namespace.Knowledge}); err != nil {
return nil, fmt.Errorf("failed to delete relationships: %v", err)
}
deleteNodesQuery := `
CALL apoc.periodic.iterate(
"MATCH (n:` + labelExpr + ` {kg: $knowledge_id}) RETURN n",
"DELETE n",
{batchSize: 1000, parallel: true, params: {knowledge_id: $knowledge_id}}
) YIELD batches, total
RETURN total
`
if _, err := tx.Run(ctx, deleteNodesQuery, map[string]interface{}{"knowledge_id": namespace.Knowledge}); err != nil {
return nil, fmt.Errorf("failed to delete nodes: %v", err)
}
}
return nil, nil
})
if err != nil {
return err
}
logger.Infof(ctx, "delete graph result: %v", result)
return nil
}
// SearchNode searches for nodes in the Neo4j repository
func (n *Neo4jRepository) SearchNode(
ctx context.Context,
namespace types.NameSpace,
nodes []string,
) (*types.GraphData, error) {
if n.driver == nil {
logger.Warnf(ctx, "NOT SUPPORT RETRIEVE GRAPH")
return nil, nil
}
session := n.driver.NewSession(ctx, neo4j.SessionConfig{AccessMode: neo4j.AccessModeRead})
defer session.Close(ctx)
result, err := session.ExecuteRead(ctx, func(tx neo4j.ManagedTransaction) (interface{}, error) {
labelExpr := n.Label(namespace)
query := `
MATCH (n:` + labelExpr + `)-[r]-(m:` + labelExpr + `)
WHERE ANY(nodeText IN $nodes WHERE n.name CONTAINS nodeText)
RETURN n, r, m
`
params := map[string]interface{}{"nodes": nodes}
result, err := tx.Run(ctx, query, params)
if err != nil {
return nil, fmt.Errorf("failed to run query: %v", err)
}
graphData := &types.GraphData{}
nodeSeen := make(map[string]bool)
for result.Next(ctx) {
record := result.Record()
node, _ := record.Get("n")
rel, _ := record.Get("r")
targetNode, _ := record.Get("m")
nodeData := node.(neo4j.Node)
targetNodeData := targetNode.(neo4j.Node)
// Convert node to types.Node
for _, n := range []neo4j.Node{nodeData, targetNodeData} {
nameStr := n.Props["name"].(string)
if _, ok := nodeSeen[nameStr]; !ok {
nodeSeen[nameStr] = true
graphData.Node = append(graphData.Node, &types.GraphNode{
Name: nameStr,
Chunks: listI2listS(n.Props["chunks"].([]interface{})),
Attributes: listI2listS(n.Props["attributes"].([]interface{})),
})
}
}
// Convert relationship to types.Relation
relData := rel.(neo4j.Relationship)
graphData.Relation = append(graphData.Relation, &types.GraphRelation{
Node1: nodeData.Props["name"].(string),
Node2: targetNodeData.Props["name"].(string),
Type: relData.Type,
})
}
return graphData, nil
})
if err != nil {
logger.Errorf(ctx, "search node failed: %v", err)
return nil, err
}
return result.(*types.GraphData), nil
}
func listI2listS(list []any) []string {
result := make([]string, len(list))
for i, v := range list {
result[i] = fmt.Sprintf("%v", v)
}
return result
}
@@ -0,0 +1,50 @@
package opensearch
import "context"
// AuditSink receives audit events emitted from within the driver at the
// exact moment they occur (index provisioned, reindex executed). The driver
// owns this abstraction so it imports no service package — the dependency
// arrow stays one-way (service implements AuditSink; the driver only invokes
// it). A nil sink is a no-op, so tests and the env-path (no audit service)
// need no special casing.
type AuditSink interface {
// EmitIndexCreated fires once when the driver provisions a new k-NN
// index. alias is the per-dimension alias (or the keyword-only index
// name); dim is the embedding dimension (0 for the keyword-only index).
EmitIndexCreated(ctx context.Context, alias string, dim int)
// EmitReindexExecuted fires when CopyIndices finishes copying docs from
// one index to another.
EmitReindexExecuted(ctx context.Context, srcAlias, dstAlias string, docs int64)
}
// nopSink is the null-object used when no sink is configured, so emit call
// sites never need a nil check.
type nopSink struct{}
func (nopSink) EmitIndexCreated(context.Context, string, int) {}
func (nopSink) EmitReindexExecuted(context.Context, string, string, int64) {}
var _ AuditSink = nopSink{}
// Option configures a Repository at construction time.
type Option func(*Repository)
// WithAuditSink injects an audit sink. A nil sink is ignored (the Repository
// keeps its default no-op behavior).
func WithAuditSink(s AuditSink) Option {
return func(r *Repository) {
if s != nil {
r.sink = s
}
}
}
// auditSink returns the configured sink, or a no-op if none was set (e.g. a
// Repository built directly in tests, or constructed without WithAuditSink).
func (r *Repository) auditSink() AuditSink {
if r.sink == nil {
return nopSink{}
}
return r.sink
}
@@ -0,0 +1,153 @@
package opensearch
import (
"context"
"strings"
"sync"
"testing"
"github.com/Tencent/WeKnora/internal/types"
)
// spySink records audit events for assertions.
type spySink struct {
mu sync.Mutex
indexCreated []indexCreatedEvent
reindex []reindexEvent
}
type indexCreatedEvent struct {
alias string
dim int
}
type reindexEvent struct {
src, dst string
docs int64
}
func (s *spySink) EmitIndexCreated(_ context.Context, alias string, dim int) {
s.mu.Lock()
defer s.mu.Unlock()
s.indexCreated = append(s.indexCreated, indexCreatedEvent{alias, dim})
}
func (s *spySink) EmitReindexExecuted(_ context.Context, src, dst string, docs int64) {
s.mu.Lock()
defer s.mu.Unlock()
s.reindex = append(s.reindex, reindexEvent{src, dst, docs})
}
// TestBuildInternalCfg_ReadsHNSWFields verifies the HNSW IndexConfig fields
// flow into internalCfg, falling back to defaults when unset.
func TestBuildInternalCfg_ReadsHNSWFields(t *testing.T) {
t.Run("reads set fields", func(t *testing.T) {
cfg, err := buildInternalCfg(&types.IndexConfig{
HNSWM: 24,
HNSWEFConstruction: 200,
HNSWEFSearch: 128,
KNNEngine: "faiss",
NumberOfShards: 2,
})
if err != nil {
t.Fatal(err)
}
if cfg.hnswM != 24 || cfg.hnswEFConstruction != 200 || cfg.efSearch != 128 || cfg.knnEngine != "faiss" {
t.Errorf("HNSW not wired: %+v", cfg)
}
if cfg.shards != 2 {
t.Errorf("shards: want 2, got %d", cfg.shards)
}
})
t.Run("defaults when unset", func(t *testing.T) {
cfg, err := buildInternalCfg(&types.IndexConfig{})
if err != nil {
t.Fatal(err)
}
if cfg.hnswM != 16 || cfg.hnswEFConstruction != 100 || cfg.efSearch != 100 || cfg.knnEngine != "lucene" {
t.Errorf("defaults not applied: %+v", cfg)
}
})
t.Run("nil config is all defaults", func(t *testing.T) {
cfg, err := buildInternalCfg(nil)
if err != nil {
t.Fatal(err)
}
if cfg.knnEngine != "lucene" || cfg.hnswM != 16 {
t.Errorf("nil defaults wrong: %+v", cfg)
}
})
}
// TestBuildIndexMapping_ReflectsHNSWConfig verifies the end-to-end path
// IndexConfig → buildInternalCfg → buildIndexMapping carries the operator's
// HNSW values into the cluster mapping JSON (regression guard for the wire-
// through, since defaults coincide with common values).
func TestBuildIndexMapping_ReflectsHNSWConfig(t *testing.T) {
cfg, err := buildInternalCfg(&types.IndexConfig{
HNSWM: 24,
HNSWEFConstruction: 200,
HNSWEFSearch: 128,
KNNEngine: "faiss",
})
if err != nil {
t.Fatal(err)
}
body, err := buildIndexMapping(cfg, 768)
if err != nil {
t.Fatal(err)
}
s := string(body)
for _, want := range []string{`"m":24`, `"ef_construction":200`, `"engine":"faiss"`, `"knn.algo_param.ef_search":128`} {
if !strings.Contains(s, want) {
t.Errorf("mapping JSON missing %q\n%s", want, s)
}
}
}
// TestWithAuditSink_SetsAndNilSafe verifies the functional option.
func TestWithAuditSink_SetsAndNilSafe(t *testing.T) {
spy := &spySink{}
r := &Repository{}
WithAuditSink(spy)(r)
if r.sink != spy {
t.Fatal("WithAuditSink did not set the sink")
}
// nil must not clobber an already-set sink
WithAuditSink(nil)(r)
if r.sink != spy {
t.Fatal("WithAuditSink(nil) clobbered the sink")
}
}
// TestAuditSink_NopByDefault verifies a Repository with no sink does not panic
// when the audit accessor is used (nopSink fallback).
func TestAuditSink_NopByDefault(t *testing.T) {
var _ AuditSink = nopSink{} // compile-time assertion
r := &Repository{} // sink left nil
// Must not panic.
r.auditSink().EmitIndexCreated(context.Background(), "weknora_768", 768)
r.auditSink().EmitReindexExecuted(context.Background(), "a", "b", 3)
}
// TestAuditSink_EmitIndexCreated_OnEnsureReady verifies createIndexAndAlias
// emits exactly one index-created event with the per-dim alias when a new
// index is provisioned.
func TestAuditSink_EmitIndexCreated_OnEnsureReady(t *testing.T) {
repo, ts := newTestRepo(t, (&indexLifecycleHandler{}).ServeHTTP)
defer ts.Close()
spy := &spySink{}
repo.sink = spy
if err := repo.ensureReady(context.Background(), 768); err != nil {
t.Fatalf("ensureReady: %v", err)
}
if len(spy.indexCreated) != 1 {
t.Fatalf("want 1 index_created event, got %d", len(spy.indexCreated))
}
got := spy.indexCreated[0]
if got.alias != "weknora_test_768" || got.dim != 768 {
t.Errorf("event mismatch: %+v", got)
}
}
@@ -0,0 +1,107 @@
package opensearch
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"sort"
osapi "github.com/opensearch-project/opensearch-go/v4/opensearchapi"
)
// BatchUpdateChunkEnabledStatus flips is_enabled for the given chunks. The
// status map is grouped by value into one _update_by_query per distinct value
// (mirrors the Qdrant grouping pattern), so the request body carries the
// chunk ids via a terms filter and the new value via bound script params —
// never per-chunk string interpolation.
//
// Targets the cross-dim <base>_* pattern: a chunk's embedding dimension is
// not known here, and the same chunk_id is unique across the store's dim
// indices + the keyword-only index.
func (r *Repository) BatchUpdateChunkEnabledStatus(ctx context.Context, chunkStatusMap map[string]bool) error {
if len(chunkStatusMap) == 0 {
return nil
}
groups := map[bool][]string{}
for id, v := range chunkStatusMap {
groups[v] = append(groups[v], id)
}
// Deterministic order (false then true) for predictable behavior/tests.
for _, v := range []bool{false, true} {
ids := groups[v]
if len(ids) == 0 {
continue
}
sort.Strings(ids)
if err := r.updateByQueryScript(ctx, ids,
"ctx._source.is_enabled = params.v", map[string]any{"v": v}); err != nil {
return err
}
}
return nil
}
// BatchUpdateChunkTagID sets tag_id for the given chunks, grouped by tag.
func (r *Repository) BatchUpdateChunkTagID(ctx context.Context, chunkTagMap map[string]string) error {
if len(chunkTagMap) == 0 {
return nil
}
groups := map[string][]string{}
for id, tag := range chunkTagMap {
groups[tag] = append(groups[tag], id)
}
tags := make([]string, 0, len(groups))
for tag := range groups {
tags = append(tags, tag)
}
sort.Strings(tags)
for _, tag := range tags {
ids := groups[tag]
sort.Strings(ids)
if err := r.updateByQueryScript(ctx, ids,
"ctx._source.tag_id = params.v", map[string]any{"v": tag}); err != nil {
return err
}
}
return nil
}
// updateByQueryScript runs an _update_by_query over the cross-dim <base>_*
// pattern, matching the given chunk ids via a terms filter and applying a
// constant Painless source with caller values flowing only through bound
// params (Painless-injection-safe).
func (r *Repository) updateByQueryScript(
ctx context.Context, chunkIDs []string, source string, params map[string]any,
) error {
body, err := json.Marshal(map[string]any{
"query": map[string]any{
"terms": map[string]any{"chunk_id": chunkIDs},
},
"script": map[string]any{
"lang": "painless",
"source": source,
"params": params,
},
})
if err != nil {
return fmt.Errorf("opensearch: marshal update_by_query body: %w", err)
}
// Q2: UpdateByQueryParams.Refresh is *bool — the wire value "wait_for" is
// not expressible via the typed SDK, so we force an immediate refresh.
refresh := true
resp, err := r.client.UpdateByQuery(ctx, osapi.UpdateByQueryReq{
Indices: []string{r.baseIndex + "_*"},
Body: bytes.NewReader(body),
Params: osapi.UpdateByQueryParams{Refresh: &refresh},
})
if err != nil {
return wrapTransport(err)
}
if resp == nil {
return nil
}
defer drainAndClose(resp.Inspect().Response.Body)
return inspectByQueryResponse(io.LimitReader(resp.Inspect().Response.Body, 16<<20))
}
@@ -0,0 +1,51 @@
package opensearch
import (
"context"
"encoding/json"
"fmt"
"io"
"strings"
"github.com/Tencent/WeKnora/internal/logger"
)
// inspectByQueryResponse parses an _update_by_query / _delete_by_query
// response and surfaces partial failures without leaking cluster-side reason
// strings (which may embed document content). Mirrors inspectBulkResponse:
// the full reason goes to the debug log only; the returned error carries the
// bounded id + type. A non-zero version_conflicts count with no hard failures
// is logged as a warning but not treated as an error.
func inspectByQueryResponse(body io.Reader) error {
var r struct {
VersionConflicts int `json:"version_conflicts"`
Failures []struct {
ID string `json:"id"`
Cause struct {
Type string `json:"type"`
Reason string `json:"reason"`
} `json:"cause"`
} `json:"failures"`
}
if err := json.NewDecoder(body).Decode(&r); err != nil {
return fmt.Errorf("opensearch: parse by-query response: %w", ErrTransport)
}
log := logger.GetLogger(context.Background())
if len(r.Failures) == 0 {
if r.VersionConflicts > 0 {
log.Warnf("[OpenSearch] by-query had %d version conflicts (proceeded)", r.VersionConflicts)
}
return nil
}
var msgs []string
for _, f := range r.Failures {
// Full reason → debug log only (may contain document content).
log.Debugf("[OpenSearch] by-query failure: id=%s type=%s reason=%s",
f.ID, f.Cause.Type, f.Cause.Reason)
if len(msgs) < 5 {
msgs = append(msgs, fmt.Sprintf("[%s] %s", f.ID, f.Cause.Type))
}
}
return fmt.Errorf("opensearch: by-query partial failure (%d failed, first 5: %s): %w",
len(r.Failures), strings.Join(msgs, "; "), ErrTransport)
}
@@ -0,0 +1,58 @@
package opensearch
import "github.com/Tencent/WeKnora/internal/types"
// internalCfg is the driver-internal, immutable view of IndexConfig.
// Defaults are chosen so the env-path (no IndexConfig) and DB-store path
// (with IndexConfig) produce identical mappings.
type internalCfg struct {
shards int
replicas int
knnEngine string // "lucene" | "faiss"
hnswM int
hnswEFConstruction int
efSearch int
}
// buildInternalCfg projects IndexConfig to the driver-internal view,
// substituting defaults for unset (zero / empty) fields. Validation of value
// ranges (e.g. hnsw_m / ef_construction caps) is a service-layer concern
// handled elsewhere (validateOpenSearchIndexConfig at CreateStore); this
// function applies defaults only and never rejects. The env-path bypasses
// service validation entirely, so the defaults below are its safety net.
//
// The OpenSearch-specific HNSW fields (knn_engine, hnsw_m,
// hnsw_ef_construction, hnsw_ef_search) are read here. They are omitempty on
// IndexConfig, so they do not affect other drivers' serialized config.
func buildInternalCfg(c *types.IndexConfig) (internalCfg, error) {
cfg := internalCfg{
shards: 4, // matches the keyword-index default upstream
replicas: 1, // assumes >= 2-node cluster
knnEngine: "lucene", // OS default; Faiss preferred only at >= 10M docs
hnswM: 16, // OS official default
hnswEFConstruction: 100, // OS official default
efSearch: 100, // OS default
}
if c == nil {
return cfg, nil
}
if c.NumberOfShards > 0 {
cfg.shards = c.NumberOfShards
}
if c.NumberOfReplicas > 0 {
cfg.replicas = c.NumberOfReplicas
}
if c.KNNEngine != "" {
cfg.knnEngine = c.KNNEngine
}
if c.HNSWM > 0 {
cfg.hnswM = c.HNSWM
}
if c.HNSWEFConstruction > 0 {
cfg.hnswEFConstruction = c.HNSWEFConstruction
}
if c.HNSWEFSearch > 0 {
cfg.efSearch = c.HNSWEFSearch
}
return cfg, nil
}
@@ -0,0 +1,194 @@
package opensearch
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"strings"
"github.com/google/uuid"
osapi "github.com/opensearch-project/opensearch-go/v4/opensearchapi"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
)
// copyBatchSize is the pagination size for the source scan. Kept under the
// BatchSave per-call cap so each copied page is a single bulk request.
const copyBatchSize = 500
// copySourceDoc is the full _source read during CopyIndices — it includes the
// embedding vector and is_recommended, which the retrieve-path hit struct
// omits because retrieval does not need them.
type copySourceDoc struct {
Content string `json:"content"`
SourceID string `json:"source_id"`
SourceType int `json:"source_type"`
ChunkID string `json:"chunk_id"`
KnowledgeID string `json:"knowledge_id"`
KnowledgeBaseID string `json:"knowledge_base_id"`
TagID string `json:"tag_id"`
IsEnabled bool `json:"is_enabled"`
IsRecommended bool `json:"is_recommended"`
Embedding []float32 `json:"embedding"`
}
// transformSourceID mirrors the sibling drivers' source_id remap:
// - regular chunk (source_id == chunk_id) → target chunk id
// - generated question (source_id == "<chunk>-<q>") → "<targetChunk>-<q>"
// - anything else → a fresh uuid
func transformSourceID(sourceID, chunkID, targetChunkID string) string {
switch {
case sourceID == chunkID:
return targetChunkID
case strings.HasPrefix(sourceID, chunkID+"-"):
return targetChunkID + "-" + strings.TrimPrefix(sourceID, chunkID+"-")
default:
return uuid.New().String()
}
}
// CopyIndices copies all docs of one knowledge base into another (within the
// same store) by scanning the source and re-saving via BatchSave — mirroring
// the Elasticsearch / Qdrant drivers (search→BatchSave), which yields the
// source_id transformation and dim/keyword routing for free. Runs
// synchronously and paginates; the large-batch background-task path is a
// later change.
//
// NOTE: from/size pagination is bounded by the index's max_result_window
// (default 10000). Copies larger than that require the scroll-based async
// path (a later change).
func (r *Repository) CopyIndices(
ctx context.Context,
sourceKnowledgeBaseID string,
sourceToTargetKBIDMap map[string]string, // keyed by source knowledge_id (mirrors sibling drivers)
sourceToTargetChunkIDMap map[string]string,
targetKnowledgeBaseID string,
dimension int,
knowledgeType string,
) error {
log := logger.GetLogger(ctx)
if len(sourceToTargetChunkIDMap) == 0 {
log.Warn("[OpenSearch] CopyIndices: empty chunk mapping, skipping")
return nil
}
if dimension <= 0 {
return fmt.Errorf("opensearch: CopyIndices requires dim > 0, got %d: %w",
dimension, ErrDimensionMismatch)
}
if err := r.ensureReady(ctx, dimension); err != nil {
return err
}
alias := r.indexAlias(dimension)
var total int64
for from := 0; ; from += copyBatchSize {
docs, err := r.copyScanBatch(ctx, alias, sourceKnowledgeBaseID, from, copyBatchSize)
if err != nil {
return err
}
if len(docs) == 0 {
break
}
infos := make([]*types.IndexInfo, 0, len(docs))
embMap := make(map[string][]float32, len(docs))
enabledMap := make(map[string]bool, len(docs))
for i := range docs {
d := &docs[i]
targetChunkID, ok := sourceToTargetChunkIDMap[d.ChunkID]
if !ok {
log.Warnf("[OpenSearch] CopyIndices: source chunk %s not mapped, skipping", d.ChunkID)
continue
}
targetKnowledgeID, ok := sourceToTargetKBIDMap[d.KnowledgeID]
if !ok {
log.Warnf("[OpenSearch] CopyIndices: source knowledge %s not mapped, skipping", d.KnowledgeID)
continue
}
targetSourceID := transformSourceID(d.SourceID, d.ChunkID, targetChunkID)
if len(d.Embedding) > 0 {
// BatchSave looks up embeddings by SourceID (lookupEmbedding),
// so key by the target source id — not the chunk id, which is
// the Elasticsearch driver's convention.
embMap[targetSourceID] = d.Embedding
}
enabledMap[targetChunkID] = d.IsEnabled
infos = append(infos, &types.IndexInfo{
Content: d.Content,
SourceID: targetSourceID,
SourceType: types.SourceType(d.SourceType),
ChunkID: targetChunkID,
KnowledgeID: targetKnowledgeID,
KnowledgeBaseID: targetKnowledgeBaseID,
KnowledgeType: knowledgeType,
TagID: d.TagID,
IsEnabled: d.IsEnabled,
IsRecommended: d.IsRecommended,
})
}
if len(infos) > 0 {
params := map[string]any{
"embedding": embMap,
"chunk_enabled": enabledMap,
}
if err := r.BatchSave(ctx, infos, params); err != nil {
return fmt.Errorf("opensearch: CopyIndices batch save: %w", err)
}
total += int64(len(infos))
}
if len(docs) < copyBatchSize {
break
}
}
log.Infof("[OpenSearch] CopyIndices: copied %d docs (KB %s → %s, dim=%d)",
total, sourceKnowledgeBaseID, targetKnowledgeBaseID, dimension)
r.auditSink().EmitReindexExecuted(ctx, alias, alias, total)
return nil
}
// copyScanBatch reads one page of docs belonging to sourceKB from the per-dim
// index, decoding the full _source (including the embedding vector).
func (r *Repository) copyScanBatch(
ctx context.Context, index, sourceKB string, from, size int,
) ([]copySourceDoc, error) {
body, err := json.Marshal(map[string]any{
"from": from,
"size": size,
"query": map[string]any{
"bool": map[string]any{
"filter": []any{
map[string]any{"term": map[string]any{"knowledge_base_id": sourceKB}},
},
},
},
})
if err != nil {
return nil, fmt.Errorf("opensearch: marshal copy scan body: %w", err)
}
req := osapi.SearchReq{Indices: []string{index}, Body: bytes.NewReader(body)}
resp, err := r.client.Search(ctx, &req)
if err != nil {
if isNotFound(err) {
return nil, fmt.Errorf("opensearch: index %s missing: %w", index, ErrIndexNotFound)
}
return nil, wrapTransport(err)
}
defer drainAndClose(resp.Inspect().Response.Body)
var parsed struct {
Hits struct {
Hits []struct {
Source copySourceDoc `json:"_source"`
} `json:"hits"`
} `json:"hits"`
}
if err := json.NewDecoder(io.LimitReader(resp.Inspect().Response.Body, 64<<20)).Decode(&parsed); err != nil {
return nil, fmt.Errorf("opensearch: parse copy scan response: %w", ErrTransport)
}
out := make([]copySourceDoc, len(parsed.Hits.Hits))
for i, h := range parsed.Hits.Hits {
out[i] = h.Source
}
return out, nil
}
@@ -0,0 +1,193 @@
package opensearch
import (
"context"
"io"
"net/http"
"strings"
"sync"
"testing"
)
// TestTransformSourceID covers the generated-question / regular / fallback
// branches that CopyIndices uses to remap source_id.
func TestTransformSourceID(t *testing.T) {
t.Run("regular chunk uses target chunk id", func(t *testing.T) {
if got := transformSourceID("chunk1", "chunk1", "tgt1"); got != "tgt1" {
t.Errorf("want tgt1, got %s", got)
}
})
t.Run("generated question preserves question id", func(t *testing.T) {
if got := transformSourceID("chunk1-q7", "chunk1", "tgt1"); got != "tgt1-q7" {
t.Errorf("want tgt1-q7, got %s", got)
}
})
t.Run("unrelated source id gets fresh uuid", func(t *testing.T) {
got := transformSourceID("totally-different", "chunk1", "tgt1")
if got == "totally-different" || got == "tgt1" || len(got) != 36 {
t.Errorf("want fresh uuid, got %q", got)
}
})
}
// TestCopyIndices_EmptyMapping_NoOp verifies an empty chunk map short-circuits
// before any HTTP call.
func TestCopyIndices_EmptyMapping_NoOp(t *testing.T) {
repo, ts := newTestRepo(t, func(w http.ResponseWriter, r *http.Request) {
t.Errorf("unexpected HTTP call: %s %s", r.Method, r.URL.Path)
})
defer ts.Close()
err := repo.CopyIndices(context.Background(), "kbSrc", map[string]string{}, map[string]string{}, "kbDst", 768, "manual")
if err != nil {
t.Fatalf("want nil, got %v", err)
}
}
// TestCopyIndices_ScanThenBatchSave verifies the search→BatchSave path:
// remaps IDs, keys the embedding by the *target source id* (OpenSearch
// BatchSave's lookup key), and emits one reindex audit event.
func TestCopyIndices_ScanThenBatchSave(t *testing.T) {
var (
mu sync.Mutex
bulkBody string
searchCnt int
)
handler := func(w http.ResponseWriter, r *http.Request) {
switch {
case r.Method == http.MethodHead:
w.WriteHeader(http.StatusOK) // alias exists → ensureReady no-op
case strings.Contains(r.URL.Path, "_search"):
mu.Lock()
searchCnt++
first := searchCnt == 1
mu.Unlock()
if first {
_, _ = w.Write([]byte(`{"hits":{"hits":[
{"_source":{"content":"c","source_id":"srcChunk","source_type":1,"chunk_id":"srcChunk","knowledge_id":"srcKnow","knowledge_base_id":"kbSrc","tag_id":"t","is_enabled":true,"is_recommended":false,"embedding":[0.1,0.2,0.3]}}
]}}`))
} else {
_, _ = w.Write([]byte(`{"hits":{"hits":[]}}`))
}
case strings.HasSuffix(r.URL.Path, "/_bulk"):
b, _ := io.ReadAll(r.Body)
mu.Lock()
bulkBody = string(b)
mu.Unlock()
_, _ = w.Write([]byte(`{"errors":false,"items":[]}`))
default:
_, _ = w.Write([]byte(`{}`))
}
}
repo, ts := newTestRepo(t, handler)
defer ts.Close()
spy := &spySink{}
repo.sink = spy
err := repo.CopyIndices(context.Background(), "kbSrc",
map[string]string{"srcKnow": "tgtKnow"}, // knowledge_id remap (sourceToTargetKBIDMap is keyed by knowledge_id, mirroring ES)
map[string]string{"srcChunk": "tgtChunk"},
"kbDst", 768, "manual")
if err != nil {
t.Fatalf("CopyIndices: %v", err)
}
mu.Lock()
defer mu.Unlock()
if bulkBody == "" {
t.Fatal("no bulk request captured")
}
// Target IDs present, source IDs gone from the written doc.
for _, want := range []string{`"chunk_id":"tgtChunk"`, `"knowledge_id":"tgtKnow"`, `"knowledge_base_id":"kbDst"`, `"source_id":"tgtChunk"`} {
if !strings.Contains(bulkBody, want) {
t.Errorf("bulk body missing %q\n%s", want, bulkBody)
}
}
if strings.Contains(bulkBody, `"knowledge_base_id":"kbSrc"`) {
t.Errorf("bulk body leaked source KB id\n%s", bulkBody)
}
// Embedding written (keyed by target source id internally).
if !strings.Contains(bulkBody, "0.1") {
t.Errorf("embedding not written\n%s", bulkBody)
}
if len(spy.reindex) != 1 || spy.reindex[0].docs != 1 {
t.Errorf("want 1 reindex event with docs=1, got %+v", spy.reindex)
}
}
// TestBatchUpdateChunkEnabledStatus_GroupedUpdateByQuery verifies the status
// map is grouped by value into one _update_by_query per distinct value, each
// passing chunk ids via terms + the new value via bound script params.
func TestBatchUpdateChunkEnabledStatus_GroupedUpdateByQuery(t *testing.T) {
var (
mu sync.Mutex
bodies []string
)
handler := func(w http.ResponseWriter, r *http.Request) {
if strings.Contains(r.URL.Path, "_update_by_query") {
b, _ := io.ReadAll(r.Body)
mu.Lock()
bodies = append(bodies, string(b))
mu.Unlock()
_, _ = w.Write([]byte(`{"updated":1,"version_conflicts":0,"failures":[]}`))
return
}
_, _ = w.Write([]byte(`{}`))
}
repo, ts := newTestRepo(t, handler)
defer ts.Close()
err := repo.BatchUpdateChunkEnabledStatus(context.Background(), map[string]bool{
"c1": true, "c2": false, "c3": true,
})
if err != nil {
t.Fatalf("BatchUpdateChunkEnabledStatus: %v", err)
}
mu.Lock()
defer mu.Unlock()
if len(bodies) != 2 {
t.Fatalf("want 2 grouped update_by_query calls (true/false), got %d: %v", len(bodies), bodies)
}
joined := strings.Join(bodies, "\n")
for _, want := range []string{"c1", "c2", "c3", "is_enabled", "params"} {
if !strings.Contains(joined, want) {
t.Errorf("update_by_query bodies missing %q\n%s", want, joined)
}
}
}
func TestBatchUpdateChunkEnabledStatus_Empty_NoOp(t *testing.T) {
repo, ts := newTestRepo(t, func(w http.ResponseWriter, r *http.Request) {
t.Errorf("unexpected HTTP call: %s %s", r.Method, r.URL.Path)
})
defer ts.Close()
if err := repo.BatchUpdateChunkEnabledStatus(context.Background(), map[string]bool{}); err != nil {
t.Fatalf("want nil, got %v", err)
}
}
// TestInspectByQueryResponse covers the success path and the failure path,
// asserting cluster-side reason text is NOT surfaced in the returned error.
func TestInspectByQueryResponse(t *testing.T) {
t.Run("clean response", func(t *testing.T) {
body := strings.NewReader(`{"updated":5,"version_conflicts":0,"failures":[]}`)
if err := inspectByQueryResponse(body); err != nil {
t.Fatalf("want nil, got %v", err)
}
})
t.Run("failures do not leak reason", func(t *testing.T) {
body := strings.NewReader(`{"updated":1,"version_conflicts":0,"failures":[
{"id":"c9","cause":{"type":"version_conflict_engine_exception","reason":"SECRET document body leaked here"}}
]}`)
err := inspectByQueryResponse(body)
if err == nil {
t.Fatal("want error for failures, got nil")
}
if strings.Contains(err.Error(), "SECRET") {
t.Errorf("error leaked cluster reason: %v", err)
}
if !strings.Contains(err.Error(), "version_conflict_engine_exception") {
t.Errorf("error should surface bounded type: %v", err)
}
})
}
@@ -0,0 +1,405 @@
package opensearch
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"strings"
osapi "github.com/opensearch-project/opensearch-go/v4/opensearchapi"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
)
// additionalParams contract (verified against
// elasticsearch/structs.go:ToDBVectorEmbedding and
// keywords_vector_hybrid_indexer.go:concurrentBatchSave on
// upstream/main):
//
// additionalParams["embedding"] map[string][]float32 // keyed by IndexInfo.SourceID
// additionalParams["chunk_enabled"] map[string]bool // keyed by IndexInfo.ChunkID
//
// Both keys may be absent. Missing "embedding" → keyword-only path
// (vector field omitted from the OpenSearch document, routed to the
// dim-less keywords index). Missing "chunk_enabled" → IndexInfo's own
// IsEnabled field is used as-is.
// Save indexes a single chunk. Idempotent: _id is set to chunk_id so
// re-indexing the same chunk overwrites instead of creating a duplicate.
func (r *Repository) Save(
ctx context.Context,
info *types.IndexInfo,
params map[string]any,
) error {
emb := lookupEmbedding(params, info.SourceID)
enabled := lookupChunkEnabled(params, info.ChunkID, info.IsEnabled)
dim := len(emb)
var targetIndex string
if dim > 0 {
if err := r.ensureReady(ctx, dim); err != nil {
return err
}
targetIndex = r.indexAlias(dim)
} else {
// Keyword-only path: no embedding supplied. Route to the
// dim-less <base>_keywords index.
if err := r.ensureKeywordsIndex(ctx); err != nil {
return err
}
targetIndex = r.keywordsIndex()
}
doc, err := json.Marshal(toDoc(info, emb, enabled))
if err != nil {
return fmt.Errorf("opensearch: marshal doc: %w", err)
}
req := osapi.IndexReq{
Index: targetIndex,
DocumentID: info.ChunkID, // _id = chunk_id (idempotent)
Body: bytes.NewReader(doc),
}
resp, err := r.client.Index(ctx, req)
if err != nil {
return wrapTransport(err)
}
if resp != nil {
drainAndClose(resp.Inspect().Response.Body)
}
return nil
}
// BatchSave NDJSON-bulk-indexes chunks. Inspects per-item errors and
// returns a typed error if ANY item fails.
//
// Dim-aware pre-marshal size cap: estimateBulkBodyBytes predicts NDJSON
// size before allocating; over-cap callers get ErrBatchTooLarge instead
// of a 50MB allocation that's then rejected. The service layer is
// responsible for dim-aware sub-batching above this.
func (r *Repository) BatchSave(
ctx context.Context,
infos []*types.IndexInfo,
params map[string]any,
) error {
if len(infos) == 0 {
return nil
}
emb, dim, err := extractBatchEmbeddings(params, infos)
if err != nil {
return err
}
// Pre-marshal size + count check.
if estimated := estimateBulkBodyBytes(len(infos), dim); estimated > 10*1024*1024 {
return fmt.Errorf(
"opensearch: estimated bulk body %dB exceeds 10MB cap (n=%d, dim=%d): %w",
estimated, len(infos), dim, ErrBatchTooLarge,
)
}
if len(infos) > 1000 {
return fmt.Errorf("opensearch: bulk n=%d exceeds 1000-doc cap: %w",
len(infos), ErrBatchTooLarge)
}
// Route by dim: dim==0 means the entire batch is keyword-only
// (concurrentBatchSaveNoEmbedding path), use the dim-less index.
var alias string
if dim == 0 {
if err := r.ensureKeywordsIndex(ctx); err != nil {
return err
}
alias = r.keywordsIndex()
} else {
if err := r.ensureReady(ctx, dim); err != nil {
return err
}
alias = r.indexAlias(dim)
}
body, err := buildBulkBody(alias, infos, emb, params)
if err != nil {
return err
}
req := osapi.BulkReq{Body: bytes.NewReader(body)}
resp, err := r.client.Bulk(ctx, req)
if err != nil {
return wrapTransport(err)
}
if resp != nil {
defer drainAndClose(resp.Inspect().Response.Body)
}
// Cap response body at 64MB (bulks of 1000 items are ~1-2MB success,
// can balloon with full per-item errors).
return inspectBulkResponse(io.LimitReader(resp.Inspect().Response.Body, 64<<20))
}
// estimateBulkBodyBytes returns a rough NDJSON size in bytes. Slight
// overestimate so the cap fires before real allocation. Per-item:
// - action header `{"index":{"_index":"...","_id":"..."}}\n` ≈ 100B
// - doc body: dim*5 (float ASCII) + ~700B metadata + content ≈ 1024B
func estimateBulkBodyBytes(n, dim int) int {
return n * (100 + dim*5 + 1024)
}
// buildBulkBody constructs the NDJSON request body. emb[i] may be nil/
// empty — toDoc handles the keyword-only case by omitting the
// "embedding" field. chunk_enabled lookup also happens per-doc.
func buildBulkBody(alias string, infos []*types.IndexInfo, emb [][]float32, params map[string]any) ([]byte, error) {
var buf bytes.Buffer
for i, info := range infos {
action := map[string]any{
"index": map[string]any{
"_index": alias,
"_id": info.ChunkID,
},
}
actionJSON, err := json.Marshal(action)
if err != nil {
return nil, fmt.Errorf("opensearch: marshal bulk action %d: %w", i, err)
}
buf.Write(actionJSON)
buf.WriteByte('\n')
enabled := lookupChunkEnabled(params, info.ChunkID, info.IsEnabled)
docJSON, err := json.Marshal(toDoc(info, emb[i], enabled))
if err != nil {
return nil, fmt.Errorf("opensearch: marshal doc %d: %w", i, err)
}
buf.Write(docJSON)
buf.WriteByte('\n')
}
return buf.Bytes(), nil
}
// inspectBulkResponse parses OpenSearch bulk response, returns a typed
// error if any item has a per-item .error. Surfaces up to 5 distinct
// error messages + total count for operator diagnosis.
//
// Does NOT include det.Error.Reason in the wrapped public error
// (Reason can contain document content fragments). Only ID + bounded
// enum Type is included. Full per-item Reason goes to log.Debug.
func inspectBulkResponse(body io.Reader) error {
var r struct {
Errors bool `json:"errors"`
Items []map[string]struct {
ID string `json:"_id"`
Status int `json:"status"`
Error *struct {
Type string `json:"type"`
Reason string `json:"reason"`
} `json:"error,omitempty"`
} `json:"items"`
}
if err := json.NewDecoder(body).Decode(&r); err != nil {
return fmt.Errorf("opensearch: parse bulk response: %w", ErrTransport)
}
if !r.Errors {
return nil
}
log := logger.GetLogger(context.Background())
var msgs []string
total := 0
for _, item := range r.Items {
for op, det := range item { // {"index": {...}}
if det.Error == nil {
continue
}
total++
// Full reason → debug log only (may contain doc content).
log.Debugf("[OpenSearch] bulk item err: op=%s id=%s type=%s reason=%s",
op, det.ID, det.Error.Type, det.Error.Reason)
if len(msgs) < 5 {
// Public message: ID (operator-meaningful) + Type (bounded enum).
msgs = append(msgs, fmt.Sprintf("[%s %s] %s", op, det.ID, det.Error.Type))
}
}
}
if total == 0 {
return nil // shouldn't happen if r.Errors=true, but defensive
}
return fmt.Errorf("opensearch: bulk partial failure (%d items failed, first 5: %s): %w",
total, strings.Join(msgs, "; "), ErrTransport)
}
// toDoc projects IndexInfo + embedding to the OpenSearch document
// shape. source_type is `int` (not stringified) — mapping declared as
// `integer` field for stable cross-enum-extension behavior. Omits the
// "embedding" field entirely when emb is empty (keyword-only doc, no
// k-NN graph entry).
func toDoc(info *types.IndexInfo, emb []float32, enabled bool) map[string]any {
doc := map[string]any{
"chunk_id": info.ChunkID,
"knowledge_id": info.KnowledgeID,
"knowledge_base_id": info.KnowledgeBaseID,
"source_id": info.SourceID,
"source_type": int(info.SourceType),
"tag_id": info.TagID,
"content": info.Content,
"is_enabled": enabled,
"is_recommended": info.IsRecommended,
}
if len(emb) > 0 {
doc["embedding"] = emb
}
return doc
}
// lookupEmbedding returns the embedding vector for info.SourceID.
// Empty slice (not nil) indicates "no embedding provided" — caller
// MUST check len(emb) before calling ensureReady.
//
// Defense in depth: tolerates the wrong shape rather than crash. The
// service layer is responsible for getting params right; we degrade
// to keyword-only for this doc and log a warning instead of panicking
// downstream.
func lookupEmbedding(params map[string]any, sourceID string) []float32 {
if params == nil {
return nil
}
raw, ok := params["embedding"]
if !ok {
return nil
}
embMap, ok := raw.(map[string][]float32)
if !ok {
logger.GetLogger(context.Background()).Warnf(
"[OpenSearch] additionalParams[\"embedding\"] is %T, want map[string][]float32", raw)
return nil
}
return embMap[sourceID] // nil/empty if SourceID not in map
}
// lookupChunkEnabled overlays the chunk_enabled map on top of
// info.IsEnabled. Returns the overlay value if present, else the
// default (info.IsEnabled).
func lookupChunkEnabled(params map[string]any, chunkID string, def bool) bool {
if params == nil {
return def
}
raw, ok := params["chunk_enabled"]
if !ok {
return def
}
enMap, ok := raw.(map[string]bool)
if !ok {
return def
}
if v, present := enMap[chunkID]; present {
return v
}
return def
}
// extractBatchEmbeddings projects the params["embedding"] map onto the
// info list, returning per-info embeddings (some may be empty) and the
// dimension inferred from the first non-empty entry. Returns dim=0 if
// no info has an embedding — caller routes to the keywords-only path.
//
// Defense in depth: all non-empty embeddings MUST share the same
// dimension. Mixed-dim → ErrDimensionMismatch.
func extractBatchEmbeddings(params map[string]any, infos []*types.IndexInfo) ([][]float32, int, error) {
out := make([][]float32, len(infos))
dim := 0
for i, info := range infos {
emb := lookupEmbedding(params, info.SourceID)
out[i] = emb
if len(emb) > 0 {
if dim == 0 {
dim = len(emb)
} else if len(emb) != dim {
return nil, 0, fmt.Errorf(
"opensearch: embedding[%s] dim=%d != first non-empty dim=%d: %w",
info.SourceID, len(emb), dim, ErrDimensionMismatch,
)
}
}
}
return out, dim, nil
}
// DeleteByChunkIDList runs _delete_by_query with a terms filter on
// chunk_id. Sync only — large batches (>1000) are rejected with
// ErrBatchTooLarge; the async task path that lifts the cap arrives in
// a later change.
//
// Refresh policy: opensearch-go v4's typed Params.Refresh is *bool, so
// the wire-level "wait_for" value documented in the OpenSearch REST
// API is not directly expressible. Refresh:&true forces immediate
// segment flush (read-your-writes guaranteed, but expensive at scale).
// A follow-up can revisit via a custom raw request if telemetry shows
// the cost is material at production batch frequencies.
func (r *Repository) DeleteByChunkIDList(
ctx context.Context, chunkIDs []string, dim int, _ string,
) error {
return r.deleteByList(ctx, chunkIDs, dim, "chunk_id")
}
func (r *Repository) DeleteBySourceIDList(
ctx context.Context, sourceIDs []string, dim int, _ string,
) error {
return r.deleteByList(ctx, sourceIDs, dim, "source_id")
}
func (r *Repository) DeleteByKnowledgeIDList(
ctx context.Context, knowledgeIDs []string, dim int, _ string,
) error {
return r.deleteByList(ctx, knowledgeIDs, dim, "knowledge_id")
}
// deleteByList factors the common cap / empty / ensureReady / dispatch
// logic out of the three DeleteBy* methods. dim==0 routes to the
// dim-less keywords index.
func (r *Repository) deleteByList(
ctx context.Context, ids []string, dim int, field string,
) error {
if len(ids) == 0 {
return nil
}
if len(ids) > 1000 {
return fmt.Errorf("opensearch: %s-delete batch %d > 1000 cap: %w",
field, len(ids), ErrBatchTooLarge)
}
var index string
if dim == 0 {
if err := r.ensureKeywordsIndex(ctx); err != nil {
return err
}
index = r.keywordsIndex()
} else {
if err := r.ensureReady(ctx, dim); err != nil {
return err
}
index = r.indexAlias(dim)
}
return r.deleteByTerms(ctx, index, field, ids)
}
func (r *Repository) deleteByTerms(
ctx context.Context, index, field string, values []string,
) error {
body, err := json.Marshal(map[string]any{
"query": map[string]any{
"terms": map[string]any{field: values},
},
})
if err != nil {
return fmt.Errorf("opensearch: marshal delete body: %w", err)
}
refresh := true
req := osapi.DocumentDeleteByQueryReq{
Indices: []string{index},
Body: bytes.NewReader(body),
Params: osapi.DocumentDeleteByQueryParams{Refresh: &refresh},
}
resp, err := r.client.Document.DeleteByQuery(ctx, req)
if err != nil {
return wrapTransport(err)
}
if resp != nil {
drainAndClose(resp.Inspect().Response.Body)
}
return nil
}
@@ -0,0 +1,62 @@
// Package opensearch implements the RetrieveEngineRepository interface
// for OpenSearch k-NN native vector search.
package opensearch
import "errors"
// Sentinel errors returned by Repository. The service-layer factory wraps
// these into typed AppError values (2200/2201) — the repository itself
// never imports internal/errors. The boundary is intentional (directional
// dependency).
var (
// ErrIndexNotFound — alias / underlying index missing. Search and
// delete-by-query operations return this when the per-dim alias has
// not been created yet (no Save has been issued for that dim).
ErrIndexNotFound = errors.New("opensearch: index not found")
// ErrDimensionMismatch — embedding dimension violates the per-dim
// invariant (e.g. dim <= 0, dim > 16000, or embeddings within a
// single batch disagree).
ErrDimensionMismatch = errors.New("opensearch: embedding dimension mismatch")
// ErrAuth — cluster returned 401 / 403. Distinguished from ErrTransport
// so the service layer can map to a clean 4xx instead of 503.
ErrAuth = errors.New("opensearch: authentication failed")
// ErrTransport — network / 5xx / opaque cluster error. Classified as
// transient: ensureReady does NOT persist this in initErr, so the next
// caller will retry.
ErrTransport = errors.New("opensearch: transport error")
// ErrVersionUnsupported — cluster is not OpenSearch, is OS 1.x, or is
// OS 2.0~2.3 (pre-Lucene-HNSW-GA). probeVersion enforces.
ErrVersionUnsupported = errors.New("opensearch: cluster version unsupported")
// ErrConfigInvalid — IndexConfig / storeID / sanitizeIndexName guard
// failed, or the k-NN plugin is missing on one or more cluster nodes.
ErrConfigInvalid = errors.New("opensearch: invalid index config")
// ErrFeatureNotEnabled — stubs.go returns this from methods whose real
// implementation has not landed yet (CopyIndices / BatchUpdateChunk* /
// swapToVersion, plus the read/write methods that a follow-up commit
// will replace with production code).
ErrFeatureNotEnabled = errors.New("opensearch: feature not enabled in this build")
// ErrBatchTooLarge — Save / Delete batch exceeded the driver's sync
// cap. Distinct from ErrFeatureNotEnabled so the service layer can
// chunk + retry rather than treat the failure as "waiting on a future
// implementation."
ErrBatchTooLarge = errors.New("opensearch: batch size exceeds driver cap")
// ErrCircuitBreaker — OpenSearch k-NN circuit breaker returned 429
// (knn_circuit_breaker_exception). Classified as transient so the
// caller can retry after the operator scales the cluster.
ErrCircuitBreaker = errors.New("opensearch: knn circuit breaker open")
)
// isTransientErr classifies sentinel errors. Transient errors do not get
// persisted in ensureReady's initErr cache — they can be retried by the
// next caller after the underlying cause clears.
func isTransientErr(err error) bool {
return errors.Is(err, ErrTransport) || errors.Is(err, ErrCircuitBreaker)
}
@@ -0,0 +1,23 @@
package opensearch
import (
"context"
"github.com/Tencent/WeKnora/internal/types"
)
// TestConnection verifies an OpenSearch cluster is reachable, runs a
// supported version, and has the k-NN plugin installed on every node. It is
// the connectivity probe used by the VectorStore service's CreateStore
// health-check (the driver's unexported probes are reused here). Returns a
// wrapped sentinel error on failure; nil on success.
func TestConnection(ctx context.Context, cfg *types.ConnectionConfig) error {
client, err := NewOpenSearchClient(cfg)
if err != nil {
return err
}
if err := probeVersion(ctx, client); err != nil {
return err
}
return probeKNNPlugin(ctx, client)
}
@@ -0,0 +1,50 @@
package opensearch
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
osapi "github.com/opensearch-project/opensearch-go/v4/opensearchapi"
"github.com/Tencent/WeKnora/internal/types"
)
// clusterHandler serves both GET / (version info) and /_cat/plugins, so the
// full TestConnection probe (version + k-NN plugin) can run end to end.
func clusterHandler(distribution, number string, plugins []osapi.CatPluginResp) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
if r.URL.Path == "/_cat/plugins" {
_ = json.NewEncoder(w).Encode(plugins)
return
}
_, _ = w.Write([]byte(`{"version":{"distribution":"` + distribution + `","number":"` + number + `"}}`))
}
}
func TestTestConnection_Success(t *testing.T) {
ts := httptest.NewServer(clusterHandler("opensearch", "3.3.2", []osapi.CatPluginResp{
{Name: "node-1", Component: "opensearch-knn"},
}))
defer ts.Close()
if err := TestConnection(context.Background(), &types.ConnectionConfig{Addr: ts.URL}); err != nil {
t.Errorf("healthy cluster: want nil, got %v", err)
}
}
func TestTestConnection_RejectsElasticsearch(t *testing.T) {
ts := httptest.NewServer(clusterHandler("elasticsearch", "8.10.4", nil))
defer ts.Close()
if err := TestConnection(context.Background(), &types.ConnectionConfig{Addr: ts.URL}); err == nil {
t.Error("elasticsearch cluster should be rejected")
}
}
func TestTestConnection_EmptyAddr(t *testing.T) {
if err := TestConnection(context.Background(), &types.ConnectionConfig{}); err == nil {
t.Error("empty addr should be rejected")
}
}
@@ -0,0 +1,420 @@
package opensearch
import (
"bytes"
"context"
"encoding/json"
"fmt"
osapi "github.com/opensearch-project/opensearch-go/v4/opensearchapi"
"github.com/Tencent/WeKnora/internal/logger"
)
// createIndexAndAlias creates "<alias>_v1" with the full k-NN mapping
// and then aliases "<alias>" → "<alias>_v1". Idempotent: returns nil if
// the alias already exists (so concurrent NewRepository calls in
// different processes don't race the PUT).
//
// Orphan cleanup: if aliasPut fails after a successful indicesCreate,
// the function attempts best-effort DELETE of the orphan _v1 index.
// Cleanup failure is logged but does not mask the original aliasPut
// error.
//
// Mapping-drift detection: if indicesCreate returns
// resource_already_exists_exception (cross-process race), the existing
// mapping's structural fingerprint is checked against the expected
// fingerprint. Drift returns ErrConfigInvalid wrapping a clear "manual
// reindex required" message instead of silently overlaying a foreign
// mapping.
func (r *Repository) createIndexAndAlias(ctx context.Context, dim int) error {
log := logger.GetLogger(ctx)
alias := r.indexAlias(dim)
realIndex := alias + "_v1"
exists, err := r.aliasExists(ctx, alias)
if err != nil {
return fmt.Errorf("alias check %s: %w", alias, err)
}
if exists {
return nil // another writer already set it up
}
body, err := buildIndexMapping(r.cfg, dim)
if err != nil {
return fmt.Errorf("build mapping for dim=%d: %w", dim, err)
}
indexCreated := false
if err := r.indicesCreate(ctx, realIndex, body); err != nil {
if isAlreadyExistsError(err) {
// Cross-process race or stale orphan. Verify the existing
// mapping matches what we expected; drift → ErrConfigInvalid.
if driftErr := r.verifyMappingMatches(ctx, realIndex, body); driftErr != nil {
return fmt.Errorf(
"mapping drift on %s: existing index has incompatible mapping; "+
"manual reindex required: %w",
realIndex, driftErr,
)
}
// Mapping matches → safe to proceed with alias setup.
} else {
return fmt.Errorf("create index %s: %w", realIndex, err)
}
} else {
indexCreated = true
}
if err := r.aliasPut(ctx, realIndex, alias); err != nil {
// Best-effort orphan cleanup so future retries find a clean slate.
if indexCreated {
if delErr := r.indicesDelete(ctx, realIndex); delErr != nil {
log.Warnf("[OpenSearch] orphan cleanup failed for %s: %v "+
"(operator must DELETE manually)", realIndex, delErr)
} else {
log.Infof("[OpenSearch] cleaned up orphan %s after aliasPut failure", realIndex)
}
}
return fmt.Errorf("put alias %s → %s: %w", alias, realIndex, err)
}
if indexCreated {
// Emit only when we actually provisioned the index (not when a
// concurrent writer / existing index short-circuited above).
r.auditSink().EmitIndexCreated(ctx, alias, dim)
}
return nil
}
// buildIndexMapping returns the full mapping JSON for a new OpenSearch
// k-NN index. content uses the standard analyzer (BM25 default); all
// *_id fields are explicit keyword sub-fields so DeleteByQuery / terms
// filters hit them without the ES v8 "useKeywordSuffix" auto-detection
// dance.
//
// match_type is intentionally NOT in the mapping — it's a per-result
// classification derived from RetrieverType at parse time. source_type
// is stored as `integer` (not stringified) so the wire format remains
// stable across SourceType enum extension.
//
// Pure function — no IO, no logging. Easy to unit-test (byte-exact diff).
func buildIndexMapping(cfg internalCfg, dim int) ([]byte, error) {
type mapping struct {
Settings struct {
Index struct {
Knn bool `json:"knn"`
NumberOfShards int `json:"number_of_shards"`
NumberOfReplicas int `json:"number_of_replicas"`
RefreshInterval string `json:"refresh_interval"`
KnnAlgoParamEFSearch int `json:"knn.algo_param.ef_search"`
} `json:"index"`
} `json:"settings"`
Mappings struct {
Properties map[string]any `json:"properties"`
} `json:"mappings"`
}
var m mapping
m.Settings.Index.Knn = true
m.Settings.Index.NumberOfShards = cfg.shards
m.Settings.Index.NumberOfReplicas = cfg.replicas
m.Settings.Index.RefreshInterval = "1s"
m.Settings.Index.KnnAlgoParamEFSearch = cfg.efSearch
m.Mappings.Properties = map[string]any{
"embedding": map[string]any{
"type": "knn_vector",
"dimension": dim,
"method": map[string]any{
"name": "hnsw",
"space_type": "cosinesimil",
"engine": cfg.knnEngine,
"parameters": map[string]any{
"m": cfg.hnswM,
"ef_construction": cfg.hnswEFConstruction,
},
},
},
"content": map[string]any{"type": "text", "analyzer": "standard"},
"chunk_id": map[string]any{"type": "keyword"},
"knowledge_id": map[string]any{"type": "keyword"},
"knowledge_base_id": map[string]any{"type": "keyword"},
"tag_id": map[string]any{"type": "keyword"},
"source_id": map[string]any{"type": "keyword"},
"source_type": map[string]any{"type": "integer"},
"is_enabled": map[string]any{"type": "boolean"},
"is_recommended": map[string]any{"type": "boolean"},
}
out, err := json.Marshal(&m)
if err != nil {
// Cannot happen for our internal struct, but propagate for symmetry
// with the other JSON builders in the package.
return nil, fmt.Errorf("marshal mapping: %w", err)
}
return out, nil
}
// buildKeywordsMapping returns the mapping for the dim-less keyword-only
// index. Same as buildIndexMapping but with the embedding field omitted
// — used by the no-embedding save path (the service layer dispatches
// there when the knowledge base has vector indexing disabled).
func buildKeywordsMapping(cfg internalCfg) ([]byte, error) {
type mapping struct {
Settings struct {
Index struct {
NumberOfShards int `json:"number_of_shards"`
NumberOfReplicas int `json:"number_of_replicas"`
RefreshInterval string `json:"refresh_interval"`
} `json:"index"`
} `json:"settings"`
Mappings struct {
Properties map[string]any `json:"properties"`
} `json:"mappings"`
}
var m mapping
m.Settings.Index.NumberOfShards = cfg.shards
m.Settings.Index.NumberOfReplicas = cfg.replicas
m.Settings.Index.RefreshInterval = "1s"
m.Mappings.Properties = map[string]any{
"content": map[string]any{"type": "text", "analyzer": "standard"},
"chunk_id": map[string]any{"type": "keyword"},
"knowledge_id": map[string]any{"type": "keyword"},
"knowledge_base_id": map[string]any{"type": "keyword"},
"tag_id": map[string]any{"type": "keyword"},
"source_id": map[string]any{"type": "keyword"},
"source_type": map[string]any{"type": "integer"},
"is_enabled": map[string]any{"type": "boolean"},
"is_recommended": map[string]any{"type": "boolean"},
}
out, err := json.Marshal(&m)
if err != nil {
return nil, fmt.Errorf("marshal keywords mapping: %w", err)
}
return out, nil
}
// ensureKeywordsIndex creates the dim-less keyword-only index used by
// the no-embedding save path. mutex + flag pattern so transient
// failures can be retried by the next caller (sync.Once cannot be
// reset).
func (r *Repository) ensureKeywordsIndex(ctx context.Context) error {
r.keywordsMu.Lock()
defer r.keywordsMu.Unlock()
if r.keywordsReady {
return nil
}
if r.keywordsErr != nil && !isTransientErr(r.keywordsErr) {
// Permanent failure — don't retry, surface the same error.
return r.keywordsErr
}
name := r.keywordsIndex()
exists, err := r.aliasExists(ctx, name) // returns true if index exists too
if err != nil {
r.keywordsErr = err
return err
}
if exists {
r.keywordsReady = true
r.keywordsErr = nil
return nil
}
body, err := buildKeywordsMapping(r.cfg)
if err != nil {
r.keywordsErr = err
return err
}
created := false
if err := r.indicesCreate(ctx, name, body); err != nil {
if !isAlreadyExistsError(err) {
r.keywordsErr = err
return err
}
// resource_already_exists_exception — race with concurrent process,
// treat as success.
} else {
created = true
}
r.keywordsReady = true
r.keywordsErr = nil
if created {
// dim=0 marks the dim-less keyword-only index.
r.auditSink().EmitIndexCreated(ctx, name, 0)
}
return nil
}
// indicesCreate is the low-level wrapper around Indices.Create. Body is
// the JSON mapping produced by buildIndexMapping.
func (r *Repository) indicesCreate(ctx context.Context, name string, body []byte) error {
req := osapi.IndicesCreateReq{Index: name, Body: bytes.NewReader(body)}
resp, err := r.client.Indices.Create(ctx, req)
if err != nil {
return wrapTransport(err)
}
if resp != nil {
drainAndClose(resp.Inspect().Response.Body)
}
return nil
}
// indicesDelete is used by createIndexAndAlias's orphan cleanup path.
// Caller decides what to do with the error.
func (r *Repository) indicesDelete(ctx context.Context, name string) error {
req := osapi.IndicesDeleteReq{Indices: []string{name}}
resp, err := r.client.Indices.Delete(ctx, req)
if err != nil {
return wrapTransport(err)
}
if resp != nil {
drainAndClose(resp.Inspect().Response.Body)
}
return nil
}
// aliasExists uses HEAD /_alias/<name> (Indices.Alias.Exists) — returns
// true on 200, false on 404, error on 5xx / transport.
//
// SDK quirk (opensearch-go v4.6.0): the Exists method passes
// dataPointer=nil to its internal do(), which means non-2xx responses
// come back as a plain *errors.errorString ("status: 404 Not Found")
// rather than as a *opensearch.StructError. We therefore inspect
// resp.StatusCode directly (resp is always returned even when err is
// non-nil) and only fall back to wrapTransport for the "no response at
// all" case (true network failure).
func (r *Repository) aliasExists(ctx context.Context, alias string) (bool, error) {
req := osapi.AliasExistsReq{Alias: []string{alias}}
resp, err := r.client.Indices.Alias.Exists(ctx, req)
if resp != nil {
statusCode := resp.StatusCode
drainAndClose(resp.Body)
switch statusCode {
case 200:
return true, nil
case 404:
return false, nil
}
// Unexpected status (5xx) — fall through to the err path so
// wrapTransport can classify (and so the test handler's
// errBody is logged for diagnosis).
}
if err != nil {
return false, wrapTransport(err)
}
return false, nil
}
// aliasPut creates an alias pointing at index. Single add operation
// (rolling-reindex atomic swap is deferred to a follow-up PR — see
// stubs.go swapToVersion).
func (r *Repository) aliasPut(ctx context.Context, index, alias string) error {
req := osapi.AliasPutReq{
Indices: []string{index},
Alias: alias,
}
resp, err := r.client.Indices.Alias.Put(ctx, req)
if err != nil {
return wrapTransport(err)
}
if resp != nil {
drainAndClose(resp.Inspect().Response.Body)
}
return nil
}
// verifyMappingMatches GETs the existing index's mapping and compares
// the embedding field's structural fingerprint (dimension, HNSW
// parameters, engine) against the expected. A full field-by-field diff
// would be ideal but expensive; we check only the parts that matter for
// query correctness.
func (r *Repository) verifyMappingMatches(ctx context.Context, index string, expected []byte) error {
req := osapi.MappingGetReq{Indices: []string{index}}
resp, err := r.client.Indices.Mapping.Get(ctx, &req)
if err != nil {
return fmt.Errorf("get mapping %s: %w", index, wrapTransport(err))
}
if resp != nil {
drainAndClose(resp.Inspect().Response.Body)
}
// Both fingerprints are derived from the mapping's embedding-field
// HNSW parameters: dimension, m, ef_construction, engine, space_type.
// We parse the expected (our own JSON) and the actual (cluster
// response) and compare just those bits.
expectedFP, err := extractEmbeddingFingerprint(expected)
if err != nil {
return fmt.Errorf("parse expected fingerprint: %w", err)
}
actualFP, err := extractMappingFingerprintFromResp(resp, index)
if err != nil {
return fmt.Errorf("parse actual fingerprint: %w", err)
}
if expectedFP != actualFP {
return fmt.Errorf("mapping fingerprint mismatch: expected %+v, got %+v: %w",
expectedFP, actualFP, ErrConfigInvalid)
}
return nil
}
// embFingerprint is the minimal set of HNSW parameters that uniquely
// determine query semantics for a k-NN index. If any of these differ
// between create-time and an existing index, retrieval is incorrect.
type embFingerprint struct {
Dimension int
M int
EFConstruction int
Engine string
SpaceType string
}
// extractEmbeddingFingerprint parses our own mapping JSON (produced by
// buildIndexMapping) to pull out the HNSW parameters.
func extractEmbeddingFingerprint(body []byte) (embFingerprint, error) {
var m struct {
Mappings struct {
Properties map[string]any `json:"properties"`
} `json:"mappings"`
}
if err := json.Unmarshal(body, &m); err != nil {
return embFingerprint{}, err
}
return extractFingerprintFromProps(m.Mappings.Properties)
}
// extractMappingFingerprintFromResp parses the cluster's GET-mapping
// response and pulls out the embedding-field HNSW fingerprint for the
// given index.
func extractMappingFingerprintFromResp(resp *osapi.MappingGetResp, index string) (embFingerprint, error) {
indexMapping, ok := resp.Indices[index]
if !ok {
return embFingerprint{}, fmt.Errorf("index %s not in mapping response", index)
}
// MappingGetResp.Indices[index].Mappings is a json.RawMessage in
// the SDK — we decode the inner properties shape ourselves.
var m struct {
Properties map[string]any `json:"properties"`
}
if err := json.Unmarshal(indexMapping.Mappings, &m); err != nil {
return embFingerprint{}, err
}
return extractFingerprintFromProps(m.Properties)
}
func extractFingerprintFromProps(props map[string]any) (embFingerprint, error) {
emb, ok := props["embedding"].(map[string]any)
if !ok {
return embFingerprint{}, fmt.Errorf("embedding field missing or wrong type")
}
dim, _ := emb["dimension"].(float64) // json.Unmarshal produces float64 for numbers
method, _ := emb["method"].(map[string]any)
engine, _ := method["engine"].(string)
spaceType, _ := method["space_type"].(string)
params, _ := method["parameters"].(map[string]any)
m, _ := params["m"].(float64)
efc, _ := params["ef_construction"].(float64)
return embFingerprint{
Dimension: int(dim),
M: int(m),
EFConstruction: int(efc),
Engine: engine,
SpaceType: spaceType,
}, nil
}
@@ -0,0 +1,154 @@
package opensearch
import (
"encoding/json"
"fmt"
"github.com/Tencent/WeKnora/internal/types"
)
// retrieveFilters is the typed shape of OpenSearch query filter
// clauses. It is the ONLY way callers (Retrieve) can constrain a
// query — there is no JSON-fragment escape hatch. This prevents the
// "map[string]interface{} injection" failure mode where a hostile
// chunk_id like {"$or":[...]} would otherwise be inlined as JSON DSL.
//
// All slice fields are interpreted as IN-set membership (OS terms
// query). Empty slice = no constraint on that dimension.
//
// IncludeDisabled controls the implicit is_enabled = true clause:
// - false (default) — appends is_enabled term, hiding disabled chunks
// - true — omits the clause; admin / audit callers can
// observe disabled chunks. Service layer pins false for user-facing
// APIs.
type retrieveFilters struct {
KBIDs []string
KnowledgeIDs []string
TagIDs []string
ExcludeChunkIDs []string
ExcludeKnowledgeIDs []string
IncludeDisabled bool
}
// fromParams projects RetrieveParams to retrieveFilters, applying the
// multi-KB input from the service-layer dispatcher as terms.
// ExcludeKnowledgeIDs is included so the OpenSearch driver matches the
// behaviour of the Qdrant / pgvector drivers.
func fromParams(p types.RetrieveParams) *retrieveFilters {
return &retrieveFilters{
KBIDs: p.KnowledgeBaseIDs,
KnowledgeIDs: p.KnowledgeIDs,
TagIDs: p.TagIDs,
ExcludeChunkIDs: p.ExcludeChunkIDs,
ExcludeKnowledgeIDs: p.ExcludeKnowledgeIDs,
// IncludeDisabled stays false — set explicitly by admin callers
// only. Driver receives this from a typed field, not from
// AdditionalParams, so the contract is checked at compile time.
}
}
// toBoolMust emits the OpenSearch bool.must clause list for this
// filter. Returns a freshly allocated slice each call so callers can
// safely append additional clauses (e.g. the keyword match clause in
// buildKeywordQuery) without aliasing.
func (f *retrieveFilters) toBoolMust() []map[string]any {
var must []map[string]any
if len(f.KBIDs) > 0 {
must = append(must, map[string]any{
"terms": map[string]any{"knowledge_base_id": f.KBIDs},
})
}
if len(f.KnowledgeIDs) > 0 {
must = append(must, map[string]any{
"terms": map[string]any{"knowledge_id": f.KnowledgeIDs},
})
}
if len(f.TagIDs) > 0 {
must = append(must, map[string]any{
"terms": map[string]any{"tag_id": f.TagIDs},
})
}
if len(f.ExcludeChunkIDs) > 0 {
must = append(must, map[string]any{
"bool": map[string]any{
"must_not": map[string]any{
"terms": map[string]any{"chunk_id": f.ExcludeChunkIDs},
},
},
})
}
if len(f.ExcludeKnowledgeIDs) > 0 {
must = append(must, map[string]any{
"bool": map[string]any{
"must_not": map[string]any{
"terms": map[string]any{"knowledge_id": f.ExcludeKnowledgeIDs},
},
},
})
}
if !f.IncludeDisabled {
must = append(must, map[string]any{
"term": map[string]any{"is_enabled": true},
})
}
return must
}
// buildKNNQuery composes a pure k-NN ANN query with caller filters. No
// hybrid pipeline — native hybrid is out of scope for this PR; fusion
// stays in the service-layer RRF.
//
// Applies `min_score` when threshold > 0. The OpenSearch k-NN plugin's
// SpaceType.COSINESIMIL.scoreTranslation produces (1 + cos) / 2 ∈ [0, 1]
// at the client, so the caller's Threshold value passes through to
// min_score directly without further scaling.
//
// Signature returns ([]byte, error) so callers can propagate
// json.Marshal failures (reachable for NaN / +/-Inf in embedding).
func buildKNNQuery(embedding []float32, topK int, threshold float64, f *retrieveFilters) ([]byte, error) {
knn := map[string]any{
"embedding": map[string]any{
"vector": embedding,
"k": topK,
"filter": map[string]any{
"bool": map[string]any{"must": f.toBoolMust()},
},
},
}
body := map[string]any{
"size": topK,
"query": map[string]any{"knn": knn},
}
if threshold > 0 {
body["min_score"] = threshold
}
out, err := json.Marshal(body)
if err != nil {
return nil, fmt.Errorf("marshal knn query: %w", err)
}
return out, nil
}
// buildKeywordQuery composes a BM25 match query against the content
// field with caller filters. Score range is OpenSearch BM25 native (no
// upper bound) — the normalizer keyword-passthrough rule treats keyword
// retrievers as rank-based for RRF, so the score scale does not need
// normalization here. min_score still applied if requested.
func buildKeywordQuery(queryText string, topK int, threshold float64, f *retrieveFilters) ([]byte, error) {
must := f.toBoolMust()
must = append(must, map[string]any{
"match": map[string]any{"content": queryText},
})
body := map[string]any{
"size": topK,
"query": map[string]any{"bool": map[string]any{"must": must}},
}
if threshold > 0 {
body["min_score"] = threshold
}
out, err := json.Marshal(body)
if err != nil {
return nil, fmt.Errorf("marshal keyword query: %w", err)
}
return out, nil
}
@@ -0,0 +1,456 @@
package opensearch
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"regexp"
"strconv"
"strings"
"sync"
"github.com/opensearch-project/opensearch-go/v4"
osapi "github.com/opensearch-project/opensearch-go/v4/opensearchapi"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
)
// EngineAwareNormalizer applies the documented per-engine cosine-score
// formula. (Repository implements interfaces.RetrieveEngineRepository for
// OpenSearch k-NN. Safe for concurrent use after construction: all fields
// except the once / initErr maps are read-only after NewRepository returns;
// per-dimension index initialization is guarded by sync.Once.
//
// Lifecycle: NewRepository validates connectivity + cluster version +
// k-NN plugin, but does NOT create any index. Index creation happens
// lazily on first Save / BatchSave / Retrieve once the embedding
// dimension is known (see ensureReady — per-dimension index naming).
//
// Concurrency: A single Repository instance is shared across N goroutines
// concurrently retrieving from the same store (the multi-store fan-out
// dispatcher in the service layer). The struct is therefore designed to
// be allocation-stable after construction; the only mutable state is the
// sync.Once + initErr maps, each guarded by its own mutex.
//
// CRITICAL INVARIANT: initErr[dim] = err must be written INSIDE the
// once.Do closure — moving it outside opens a TOCTOU where a concurrent
// caller reads zero before the first caller persists.
type Repository struct {
client *osapi.Client
baseIndex string // pre-suffix base name, e.g. "weknora_abcdef012345"
cfg internalCfg // immutable after construction (value, not pointer)
// Per-dimension lazy initialization. ensureReady(ctx, dim) inserts an
// entry the first time dim is seen, runs the create-index work inside
// sync.Once, and persists any permanent error so subsequent callers
// see the same failure without re-attempting the PUT. Transient
// errors (ErrTransport, ErrCircuitBreaker) are NOT persisted — the
// next caller retries.
onceMu sync.Mutex
once map[int]*sync.Once
initMu sync.Mutex
initErr map[int]error
// Lazy init state for the dim-less keyword-only index used by the
// no-embedding path (concurrentBatchSaveNoEmbedding in the service
// layer). Mutex + flag pattern (rather than sync.Once) so transient
// failures can be retried by the next caller — sync.Once does not
// support reset.
keywordsMu sync.Mutex
keywordsReady bool
keywordsErr error
// sink receives audit events (index created / reindex executed). nil
// means no auditing; use r.auditSink() to get a non-nil sink. Set via
// WithAuditSink at construction.
sink AuditSink
}
// Compile-time interface satisfaction (Go best practice — keeps the build
// red if the interface drifts and our implementation lags).
var _ interfaces.RetrieveEngineRepository = (*Repository)(nil)
// NewRepository builds a new OpenSearch k-NN repository and verifies the
// backing cluster is reachable + version-compatible + has the k-NN
// plugin installed on every cluster node. It does NOT create any index —
// callers (Save / Retrieve) trigger lazy per-dimension creation via
// ensureReady on first use.
//
// storeID is the VectorStore.ID owning this repository instance. It is
// folded into the base index name so multiple OpenSearch VectorStores
// pointing at the same cluster do not collide:
// - storeID == "" — env-path; no prefix, base index = "weknora".
// - len(storeID) >= 16 — DB-store path; prefix = storeID[:12]
// (48-bit collision space).
// - 1..15 chars — rejected with ErrConfigInvalid.
//
// indexCfg is optional — pass nil to use env var (OPENSEARCH_INDEX) or
// default ("weknora") values.
//
// Optional behavior is configured via functional options (e.g.
// WithAuditSink). Passing no options keeps audit emission as a no-op, so the
// env-path and tests need no extra wiring.
//
// Returns a typed sentinel error wrapped with %w; callers translate to
// AppError at the engine-factory boundary.
func NewRepository(
ctx context.Context,
client *osapi.Client,
storeID string,
indexCfg *types.IndexConfig,
opts ...Option,
) (interfaces.RetrieveEngineRepository, error) {
log := logger.GetLogger(ctx)
if storeID != "" && len(storeID) < 16 {
return nil, fmt.Errorf(
"opensearch: storeID must be empty or >=16 chars, got %d: %w",
len(storeID), ErrConfigInvalid,
)
}
base := types.ResolveIndexName(indexCfg, "OPENSEARCH_INDEX", "weknora")
if storeID != "" {
// 12 hex chars → 48-bit collision space (UUID is 32 chars; using
// half gives strong-enough collision avoidance for the realistic
// "two OpenSearch VectorStores on the same cluster" case).
base = fmt.Sprintf("%s_%s", base, storeID[:12])
}
base, err := sanitizeIndexName(base)
if err != nil {
return nil, fmt.Errorf("opensearch: invalid index base name: %w", err)
}
icfg, err := buildInternalCfg(indexCfg)
if err != nil {
return nil, fmt.Errorf("opensearch: invalid index config: %w", err)
}
if err := probeVersion(ctx, client); err != nil {
return nil, err // already wraps ErrVersionUnsupported / ErrTransport
}
if err := probeKNNPlugin(ctx, client); err != nil {
return nil, err // already wraps ErrConfigInvalid / ErrTransport
}
r := &Repository{
client: client,
baseIndex: base,
cfg: icfg,
once: make(map[int]*sync.Once),
initErr: make(map[int]error),
}
for _, opt := range opts {
opt(r)
}
log.Infof("[OpenSearch] repository ready (baseIndex=%s, knn_engine=%s, hnsw_m=%d)",
base, icfg.knnEngine, icfg.hnswM)
return r, nil
}
// ensureReady creates the per-dimension index (alias-backed) the first
// time a given embedding dimension is seen. Concurrent callers for the
// same dim block on the same sync.Once.
//
// Transient vs permanent error handling:
// - PERMANENT failures (ErrConfigInvalid, ErrAuth, ErrVersionUnsupported,
// ErrDimensionMismatch) are persisted in initErr[dim] and re-served to
// all future callers without retry — caller must restart or fix config.
// - TRANSIENT failures (ErrTransport, ErrCircuitBreaker) are NOT
// persisted; the sync.Once is reset for that dim so the next caller
// can retry. This prevents a single cluster blip from permanently
// locking out a dim.
//
// Caller MUST pass a real ctx; nil is rejected. dim must be in (0, 16000]
// (OpenSearch knn_vector hard limit).
//
// CRITICAL INVARIANT: initErr[dim] = err is written INSIDE the once.Do
// closure. Moving it outside opens a TOCTOU where a concurrent caller
// reads zero before the first caller persists.
func (r *Repository) ensureReady(ctx context.Context, dim int) error {
if ctx == nil {
return fmt.Errorf("opensearch: ensureReady requires non-nil ctx")
}
if dim <= 0 || dim > 16000 {
return fmt.Errorf("opensearch: dim %d out of range (1..16000): %w",
dim, ErrDimensionMismatch)
}
r.onceMu.Lock()
once, ok := r.once[dim]
if !ok {
once = &sync.Once{}
r.once[dim] = once
}
r.onceMu.Unlock()
once.Do(func() {
err := r.createIndexAndAlias(ctx, dim)
if err != nil && isTransientErr(err) {
// Do NOT persist transient errors. Reset the sync.Once so
// the next caller retries (requires once map mutation).
r.onceMu.Lock()
delete(r.once, dim)
r.onceMu.Unlock()
// Fall through — caller still sees this attempt's err.
}
r.initMu.Lock()
if err == nil || !isTransientErr(err) {
r.initErr[dim] = err // nil on success, persist permanent failures
}
r.initMu.Unlock()
})
r.initMu.Lock()
err := r.initErr[dim]
r.initMu.Unlock()
return err
}
// indexAlias returns the alias name for a given dimension. All reads /
// writes target the alias; the alias points at the actual <alias>_v1
// index, enabling future rolling-reindex swap workflows (the swap path
// itself is deferred to a follow-up PR; see stubs.go swapToVersion).
func (r *Repository) indexAlias(dim int) string {
return fmt.Sprintf("%s_%d", r.baseIndex, dim)
}
// keywordsIndex returns the name of the dim-less keyword-only index used
// by the no-embedding save path (the service layer dispatches there when
// the knowledge base has vector indexing disabled). The mapping is the
// same as the per-dim index minus the knn_vector field.
func (r *Repository) keywordsIndex() string {
return r.baseIndex + "_keywords"
}
// EngineType returns the engine constant for OpenSearch. Used by the
// fan-out normalizer to route to the [0, 1] passthrough group.
func (r *Repository) EngineType() types.RetrieverEngineType {
return types.OpenSearchRetrieverEngineType
}
// Support reports the retrieval modes this driver implements. OpenSearch
// k-NN supports both ANN (vector) and BM25 (keyword) within a single
// document; the keyword-only index keeps BM25 functional for callers
// that omit the embedding map.
func (r *Repository) Support() []types.RetrieverType {
return []types.RetrieverType{types.KeywordsRetrieverType, types.VectorRetrieverType}
}
// probeVersion sends GET / and validates the cluster is OpenSearch in
// a supported version range:
// - Reject: ES (any), OS 1.x, OS 2.0~2.3 (Lucene HNSW preview).
// - Warn-but-accept: OS 2.4~2.10 (Lucene HNSW GA, pre-LTS).
// - Clean-accept: OS 2.11+, OS 3.x (primary tested: 3.3.2).
func probeVersion(ctx context.Context, client *osapi.Client) error {
resp, err := client.Info(ctx, nil)
if err != nil {
return fmt.Errorf("opensearch: cluster info: %w", wrapTransport(err))
}
distribution := resp.Version.Distribution
number := resp.Version.Number
if distribution != "opensearch" {
return fmt.Errorf("opensearch: unsupported distribution %q: %w",
distribution, ErrVersionUnsupported)
}
maj, min := parseMajorMinor(number)
switch {
case maj == 1:
return fmt.Errorf("opensearch: 1.x EOL: %w", ErrVersionUnsupported)
case maj == 2 && min >= 0 && min <= 3:
return fmt.Errorf("opensearch: %s lacks Lucene HNSW GA (need 2.4+): %w",
number, ErrVersionUnsupported)
case maj == 2 && min >= 4 && min <= 10:
logger.GetLogger(ctx).Warnf(
"[OpenSearch] using pre-2.11 cluster %s; recommend 2.11+ LTS", number)
return nil
case maj == 2 || maj == 3:
return nil
default:
return fmt.Errorf("opensearch: unsupported version %s: %w",
number, ErrVersionUnsupported)
}
}
// parseMajorMinor extracts (major, minor) from a semver string such as
// "3.3.2", "2.10.0-SNAPSHOT", "2.10". Returns (0, 0) for unparseable
// input. Robust to pre-release suffixes (strips after '-') and missing
// patch component (e.g. "2.10").
func parseMajorMinor(num string) (major, minor int) {
base := strings.SplitN(num, "-", 2)[0]
parts := strings.Split(base, ".")
if len(parts) < 2 {
return 0, 0
}
m1, err1 := strconv.Atoi(parts[0])
m2, err2 := strconv.Atoi(parts[1])
if err1 != nil || err2 != nil {
return 0, 0
}
return m1, m2
}
// probeKNNPlugin verifies the OpenSearch cluster has the opensearch-knn
// plugin installed on EVERY node. Without it on at least one node,
// knn_vector queries that route a shard to that node fail with opaque
// shard-failure errors — we surface this as a clear "plugin missing on
// N nodes" message at registration.
//
// Implementation: _cat/plugins returns one row per (node, plugin) pair.
// We group by node name and verify every distinct node has the
// opensearch-knn component.
func probeKNNPlugin(ctx context.Context, client *osapi.Client) error {
// CatPluginsReq's Params struct has no Format field — the SDK
// negotiates JSON via Accept header and returns a typed
// CatPluginsResp with .Plugins []CatPluginResp already parsed.
resp, err := client.Cat.Plugins(ctx, &osapi.CatPluginsReq{})
if err != nil {
return fmt.Errorf("opensearch: cat plugins: %w", wrapTransport(err))
}
// Group by node, verify every distinct node has opensearch-knn.
nodes := make(map[string]bool) // nodeName -> hasKNN
for _, p := range resp.Plugins {
if p.Name == "" {
continue
}
if _, seen := nodes[p.Name]; !seen {
nodes[p.Name] = false
}
if p.Component == "opensearch-knn" {
nodes[p.Name] = true
}
}
if len(nodes) == 0 {
return fmt.Errorf("opensearch: _cat/plugins returned no rows: %w", ErrConfigInvalid)
}
var missing []string
for node, hasKNN := range nodes {
if !hasKNN {
missing = append(missing, node)
}
}
if len(missing) > 0 {
return fmt.Errorf(
"opensearch: opensearch-knn plugin missing on %d/%d nodes (%v): %w",
len(missing), len(nodes), missing, ErrConfigInvalid,
)
}
return nil
}
// wrapTransport classifies a raw transport error into one of the
// driver's sentinel errors based on HTTP status ONLY. The cluster-side
// error message (err.Error()) is intentionally NOT included in the
// wrapped public error — it may contain internal index names, shard
// IDs, node hostnames, or document content snippets. The full SDK
// error is routed to log.Debug for operator forensics.
//
// Classification:
// - 401 / 403 → ErrAuth
// - 404 → caller decides (use isNotFound separately for index ops)
// - 429 + knn_circuit_breaker_exception → ErrCircuitBreaker
// - 5xx / network / timeout → ErrTransport
//
// This is the ONE place in the package that maps wire errors to
// sentinels; downstream functions just propagate.
func wrapTransport(err error) error {
if err == nil {
return nil
}
// Route the verbose SDK message to debug log (operator-visible, not
// user-visible). Background ctx is intentional — wrapTransport may
// be invoked from any caller including the constructor.
logger.GetLogger(context.Background()).Debugf("[OpenSearch] transport err: %v", err)
var se *opensearch.StructError
if errors.As(err, &se) {
switch se.Status {
case http.StatusUnauthorized, http.StatusForbidden:
return fmt.Errorf("authentication failed: %w", ErrAuth)
case http.StatusTooManyRequests:
if se.Err.Type == "knn_circuit_breaker_exception" {
return fmt.Errorf("circuit breaker open: %w", ErrCircuitBreaker)
}
}
}
return fmt.Errorf("transport error: %w", ErrTransport)
}
// isNotFound returns true if err carries an HTTP 404. Used by index /
// search code paths that want to distinguish "missing alias" from
// "transport problem".
func isNotFound(err error) bool {
var se *opensearch.StructError
return errors.As(err, &se) && se.Status == http.StatusNotFound
}
// isAlreadyExistsError returns true if err is a 400 with OpenSearch's
// resource_already_exists_exception type. Used by createIndexAndAlias
// to handle the cross-process race where two replicas try to create
// the same per-dim index simultaneously.
func isAlreadyExistsError(err error) bool {
if err == nil {
return false
}
var se *opensearch.StructError
if !errors.As(err, &se) {
return false
}
return se.Status == http.StatusBadRequest &&
se.Err.Type == "resource_already_exists_exception"
}
// sanitizeRe enforces the OpenSearch index-name contract:
// - lowercase ASCII letters, digits, underscore, hyphen only
// - must start with [a-z0-9]
// - length 1..255
var sanitizeRe = regexp.MustCompile(`^[a-z0-9][a-z0-9_-]{0,254}$`)
// sanitizeIndexName rejects names that violate the OpenSearch index-name
// rules. Any non-conforming input returns ErrConfigInvalid — the driver
// does NOT silently rewrite invalid input so operators get a clear
// error at registration.
func sanitizeIndexName(name string) (string, error) {
if name == "" {
return "", fmt.Errorf("empty index name: %w", ErrConfigInvalid)
}
// Explicit reject of OpenSearch-reserved wildcards / separators /
// control characters. Some of these would already be caught by the
// regexp below, but a dedicated check keeps the error message
// specific.
if strings.ContainsAny(name, "*?,\n\r\t/\\") {
return "", fmt.Errorf("invalid char in %q: %w", name, ErrConfigInvalid)
}
if !sanitizeRe.MatchString(name) {
return "", fmt.Errorf("name %q must match %s: %w",
name, sanitizeRe.String(), ErrConfigInvalid)
}
return name, nil
}
// drainAndClose closes the response body after draining unread bytes —
// required for HTTP keep-alive connection reuse. opensearch-go does NOT
// drain the body for us; failing to drain on every response will starve
// the per-host connection pool under load.
func drainAndClose(body io.ReadCloser) {
if body == nil {
return
}
_, _ = io.Copy(io.Discard, body)
_ = body.Close()
}
// limitedDecode wraps json.NewDecoder with an io.LimitReader so a
// pathological or malicious cluster response cannot exhaust process
// memory. Callers pass the documented per-endpoint cap (e.g. 64MB for
// bulk, 16MB for search, 1MB for plugin probe).
func limitedDecode(body io.Reader, maxBytes int64, v any) error {
return json.NewDecoder(io.LimitReader(body, maxBytes)).Decode(v)
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,205 @@
package opensearch
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
osapi "github.com/opensearch-project/opensearch-go/v4/opensearchapi"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
)
// Retrieve dispatches to vectorRetrieve or keywordsRetrieve based on
// params.RetrieverType. Returns exactly one *RetrieveResult — the
// service-layer multi-store dispatcher depends on this one-bundle-per-
// driver protocol. Score in result is RAW [0, 1] passthrough (the
// OpenSearch k-NN plugin's SpaceType.COSINESIMIL.scoreTranslation has
// already mapped to (1 + cos) / 2 before the score reaches us).
//
// Dim resolution order:
// 1. params.AdditionalParams["dim"] if present and positive
// 2. len(params.Embedding) if positive
// 3. cross-dim multi-index search for keyword-only paths via
// "<base>_*" wildcard pattern
//
// Vector retrieve always requires a positive dim (it's intrinsic to k-NN).
func (r *Repository) Retrieve(
ctx context.Context,
params types.RetrieveParams,
) ([]*types.RetrieveResult, error) {
dim, multiIndex := resolveDim(params)
switch params.RetrieverType {
case types.VectorRetrieverType:
if dim == 0 {
return nil, fmt.Errorf(
"opensearch: vector retrieve requires embedding or AdditionalParams[\"dim\"]: %w",
ErrDimensionMismatch)
}
if err := r.ensureReady(ctx, dim); err != nil {
return nil, err
}
return r.vectorRetrieve(ctx, params, dim)
case types.KeywordsRetrieverType:
if multiIndex {
// BM25-only path without dim — search across "<base>_*"
// wildcard pattern. Slower than a single-index search but
// correct for admin / keyword-only callers.
return r.keywordsRetrieve(ctx, params, r.baseIndex+"_*")
}
if err := r.ensureReady(ctx, dim); err != nil {
return nil, err
}
return r.keywordsRetrieve(ctx, params, r.indexAlias(dim))
default:
return nil, fmt.Errorf("opensearch: unsupported retriever type %q", params.RetrieverType)
}
}
// resolveDim returns (dim, multiIndex). multiIndex=true means dim is
// unknown and the caller should use a wildcard alias pattern (keyword
// path only — vector retrieve cannot be dim-less).
func resolveDim(p types.RetrieveParams) (int, bool) {
if p.AdditionalParams != nil {
if v, ok := p.AdditionalParams["dim"].(int); ok && v > 0 {
return v, false
}
}
if len(p.Embedding) > 0 {
return len(p.Embedding), false
}
return 0, true // unknown — multi-index fallback for keyword
}
func (r *Repository) vectorRetrieve(
ctx context.Context, p types.RetrieveParams, dim int,
) ([]*types.RetrieveResult, error) {
body, err := buildKNNQuery(p.Embedding, effectiveTopK(ctx, p), p.Threshold, fromParams(p))
if err != nil {
return nil, err
}
hits, err := r.search(ctx, r.indexAlias(dim), body)
if err != nil {
return nil, err
}
return wrapResults(ctx, hits, types.VectorRetrieverType, types.MatchTypeEmbedding), nil
}
func (r *Repository) keywordsRetrieve(
ctx context.Context, p types.RetrieveParams, indexPattern string,
) ([]*types.RetrieveResult, error) {
body, err := buildKeywordQuery(p.Query, effectiveTopK(ctx, p), p.Threshold, fromParams(p))
if err != nil {
return nil, err
}
hits, err := r.search(ctx, indexPattern, body)
if err != nil {
return nil, err
}
return wrapResults(ctx, hits, types.KeywordsRetrieverType, types.MatchTypeKeywords), nil
}
// effectiveTopK returns the caller's TopK or 10 if unset. The 10 is a
// defensive default — every real caller sets TopK, so a TopK=10 in
// production logs indicates a caller bug. Caps at 10000 to stay below
// OpenSearch's index.max_result_window default.
func effectiveTopK(ctx context.Context, p types.RetrieveParams) int {
if p.TopK <= 0 {
logger.GetLogger(ctx).Warnf("[OpenSearch] Retrieve called with TopK<=0; defaulting to 10 (caller bug?)")
return 10
}
if p.TopK > 10000 {
return 10000
}
return p.TopK
}
func (r *Repository) search(ctx context.Context, indexPattern string, body []byte) ([]hit, error) {
req := osapi.SearchReq{Indices: []string{indexPattern}, Body: bytes.NewReader(body)}
resp, err := r.client.Search(ctx, &req)
if err != nil {
if isNotFound(err) {
return nil, fmt.Errorf("opensearch: index %s missing: %w", indexPattern, ErrIndexNotFound)
}
return nil, wrapTransport(err)
}
defer func() {
if resp != nil {
drainAndClose(resp.Inspect().Response.Body)
}
}()
// Cap response body at 16MB (search results bounded by max_result_window * doc_size).
return parseSearchHits(io.LimitReader(resp.Inspect().Response.Body, 16<<20))
}
// hit is the structural shape we pull from each OpenSearch search hit.
// Field-by-field decode (vs map[string]any) keeps the JSON shape
// pinned at compile time.
type hit struct {
ID string `json:"_id"` // equals chunk_id per the indexing invariant
Score float64 `json:"_score"`
Source struct {
Content string `json:"content"`
ChunkID string `json:"chunk_id"`
KnowledgeID string `json:"knowledge_id"`
KnowledgeBaseID string `json:"knowledge_base_id"`
SourceID string `json:"source_id"`
SourceType int `json:"source_type"` // integer, not stringified
TagID string `json:"tag_id"`
IsEnabled bool `json:"is_enabled"`
} `json:"_source"`
}
func parseSearchHits(body io.Reader) ([]hit, error) {
var resp struct {
Hits struct {
Hits []hit `json:"hits"`
} `json:"hits"`
}
if err := json.NewDecoder(body).Decode(&resp); err != nil {
return nil, fmt.Errorf("opensearch: parse search response: %w", ErrTransport)
}
return resp.Hits.Hits, nil
}
// wrapResults converts OpenSearch hits to the upstream IndexWithScore
// shape. Always returns exactly one *RetrieveResult — the service-layer
// multi-store dispatcher depends on this "one engine result per driver"
// contract.
//
// Defensive: warns if hit._id != _source.chunk_id (we set _id=chunk_id
// on every write, so a mismatch means external tooling has been writing
// into the index) but does not fail the query — the result is still
// usable for ranking, just suspicious.
func wrapResults(ctx context.Context, hits []hit, rt types.RetrieverType, mt types.MatchType) []*types.RetrieveResult {
log := logger.GetLogger(ctx)
out := make([]*types.IndexWithScore, 0, len(hits))
for _, h := range hits {
if h.ID != h.Source.ChunkID {
log.Warnf("[OpenSearch] hit._id=%q != _source.chunk_id=%q (D12 invariant violation)",
h.ID, h.Source.ChunkID)
}
out = append(out, &types.IndexWithScore{
ID: h.ID,
ChunkID: h.Source.ChunkID,
KnowledgeID: h.Source.KnowledgeID,
KnowledgeBaseID: h.Source.KnowledgeBaseID,
SourceID: h.Source.SourceID,
SourceType: types.SourceType(h.Source.SourceType),
TagID: h.Source.TagID,
Content: h.Source.Content,
Score: h.Score,
MatchType: mt,
IsEnabled: h.Source.IsEnabled,
})
}
return []*types.RetrieveResult{{
Results: out,
RetrieverEngineType: types.OpenSearchRetrieverEngineType,
RetrieverType: rt,
}}
}
@@ -0,0 +1,16 @@
package opensearch
import (
"os"
"testing"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
func TestMain(m *testing.M) {
// End-to-end repository tests intentionally use loopback httptest servers.
secutils.SetSSRFWhitelistFromRaw("127.0.0.1,::1,localhost,opensearch.example.com")
code := m.Run()
secutils.SetSSRFWhitelistFromRaw("")
os.Exit(code)
}
@@ -0,0 +1,52 @@
package opensearch
import (
"context"
"github.com/Tencent/WeKnora/internal/types"
)
// This file holds the remaining stubs for methods whose real implementation
// has not landed yet — the rolling-reindex swap (swapToVersion) and the
// precise storage estimate. CopyIndices (copy.go), BatchUpdateChunkEnabledStatus
// and BatchUpdateChunkTagID (bulk_update.go) are now implemented. Each
// remaining stub returns ErrFeatureNotEnabled (or, for EstimateStorageSize, a
// conservative lower-bound) so any accidental invocation surfaces loudly.
// EstimateStorageSize: the real implementation that reads cluster
// `_stats` for the per-dim alias arrives in a later change. For now we
// return a conservative lower-bound estimate using the HNSW memory
// formula so the upstream KB-delete guard fails-closed (treats non-empty
// KBs as "may free non-trivial storage, force confirmation") rather than
// failing open.
//
// Formula: N * (1024 content bytes + 4*dimGuess float + 128 HNSW M=16 overhead)
// dimGuess = 768 (common embedding size; the real implementation reads
// the actual dim from the cluster).
func (r *Repository) EstimateStorageSize(
_ context.Context,
indexInfoList []*types.IndexInfo,
_ map[string]any,
) int64 {
if len(indexInfoList) == 0 {
return 0
}
const (
contentBytes = 1024 // average chunk content
embDimGuess = 768 // common embedding size
hnswOverhead = 128 // M=16 → 8*M = 128 bytes/vector
)
return int64(len(indexInfoList)) * int64(contentBytes+4*embDimGuess+hnswOverhead)
}
// swapToVersion is a stub for the future rolling-reindex swap path.
// Calling it is illegal at this stage — alias is fixed at "_v1" and only
// a follow-up change ships the swap orchestration. The surface exists so
// the future API contract is reviewable now.
//
// Unexported because it is not part of the public
// RetrieveEngineRepository interface — exposing it would require an
// interface widening in internal/types/interfaces.
func (r *Repository) swapToVersion(_ context.Context, _ int, _ int) error {
return ErrFeatureNotEnabled
}
@@ -0,0 +1,80 @@
package opensearch
import (
"crypto/tls"
"fmt"
"net/http"
"time"
"github.com/opensearch-project/opensearch-go/v4"
osapi "github.com/opensearch-project/opensearch-go/v4/opensearchapi"
"github.com/Tencent/WeKnora/internal/types"
secutils "github.com/Tencent/WeKnora/internal/utils"
)
// NewOpenSearchClient builds a TLS-hardened, pool-tuned *osapi.Client for
// the OpenSearch driver. The caller wires it into the registry from the
// env path (container) and the DB-store path (engine factory); the
// Repository constructor itself receives the pre-built client.
//
// TLS posture:
// - MinVersion: TLS 1.2 (TLS 1.3 negotiated when both ends support).
// - InsecureSkipVerify: opt-in via ConnectionConfig.InsecureSkipVerify.
// Default false. Operators set true ONLY for self-signed dev clusters.
// - CipherSuites: explicit forward-secrecy-only list. TLS 1.2's default
// cipher list includes TLS_RSA_* (no forward secrecy); we drop them.
// TLS 1.3 cipher selection is automatic and safe.
//
// Transport tuning:
// - MaxIdleConnsPerHost: 32 (Go default 2 is too low when callers
// dispatch concurrent queries across multiple knowledge bases)
// - IdleConnTimeout: 90s (typical LB keep-alive)
// - ResponseHeaderTimeout: 30s (per-request safety net)
// - ExpectContinueTimeout: 1s
func NewOpenSearchClient(cfg *types.ConnectionConfig) (*osapi.Client, error) {
if cfg == nil || cfg.Addr == "" {
return nil, fmt.Errorf("opensearch: ConnectionConfig.Addr required: %w", ErrConfigInvalid)
}
if err := secutils.ValidateURLForSSRF(cfg.Addr); err != nil {
return nil, fmt.Errorf("opensearch: address failed SSRF validation: %w", err)
}
transport := buildHTTPTransport(cfg.InsecureSkipVerify)
return osapi.NewClient(osapi.Config{
Client: opensearch.Config{
Addresses: []string{cfg.Addr},
Username: cfg.Username,
Password: cfg.Password,
Transport: &secutils.SSRFValidatingRoundTripper{Base: transport},
},
})
}
// buildHTTPTransport constructs the TLS-hardened, pool-tuned
// *http.Transport used by every OpenSearch client this driver creates.
// Extracted from newOpenSearchClient so unit tests can pin the TLS
// posture directly (the opensearch-go SDK wraps the transport in an
// internal interface, making reflection-based introspection of the
// final client fragile across SDK versions).
func buildHTTPTransport(insecureSkipVerify bool) *http.Transport {
return &http.Transport{
Proxy: http.ProxyFromEnvironment,
DialContext: secutils.SSRFSafeDialContext,
TLSClientConfig: &tls.Config{
MinVersion: tls.VersionTLS12,
InsecureSkipVerify: insecureSkipVerify, //nolint:gosec — operator opt-in flag
CipherSuites: []uint16{
tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
tls.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384,
tls.TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305,
tls.TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305,
},
},
MaxIdleConnsPerHost: 32,
IdleConnTimeout: 90 * time.Second,
ResponseHeaderTimeout: 30 * time.Second,
ExpectContinueTimeout: 1 * time.Second,
}
}
@@ -0,0 +1,137 @@
package opensearch
import (
"crypto/tls"
"errors"
"testing"
"github.com/Tencent/WeKnora/internal/types"
)
// TestNewOpenSearchClient_RejectsEmptyAddr verifies the constructor
// refuses an empty Addr — operators who forget to set OPENSEARCH_ADDR
// get a clear ErrConfigInvalid at registration instead of a cryptic
// transport error later.
func TestNewOpenSearchClient_RejectsEmptyAddr(t *testing.T) {
t.Parallel()
_, err := NewOpenSearchClient(&types.ConnectionConfig{Addr: ""})
if !errors.Is(err, ErrConfigInvalid) {
t.Fatalf("empty addr: want ErrConfigInvalid, got %v", err)
}
_, err = NewOpenSearchClient(nil)
if !errors.Is(err, ErrConfigInvalid) {
t.Fatalf("nil cfg: want ErrConfigInvalid, got %v", err)
}
}
// TestNewOpenSearchClient_Succeeds_OnValidAddr is a smoke test for the
// happy path — a non-empty Addr must produce a non-nil *osapi.Client.
// We don't probe the cluster (no network), just verify the constructor
// returns successfully.
func TestNewOpenSearchClient_Succeeds_OnValidAddr(t *testing.T) {
t.Parallel()
client, err := NewOpenSearchClient(&types.ConnectionConfig{
Addr: "https://opensearch.example.com:9200",
Username: "admin",
Password: "secret", // not a real password — wire-format only
})
if err != nil {
t.Fatalf("NewOpenSearchClient: %v", err)
}
if client == nil {
t.Fatal("client must be non-nil on success")
}
}
// TestBuildHTTPTransport_DefaultsToSecureTLS pins the default
// InsecureSkipVerify=false. Operators who don't opt in MUST get a
// strict-verify HTTPS client.
func TestBuildHTTPTransport_DefaultsToSecureTLS(t *testing.T) {
t.Parallel()
tr := buildHTTPTransport(false)
if tr.TLSClientConfig == nil {
t.Fatal("TLSClientConfig must be set")
}
if tr.TLSClientConfig.InsecureSkipVerify {
t.Fatal("default InsecureSkipVerify must be false (secure)")
}
}
// TestBuildHTTPTransport_InsecureSkipVerifyOptIn confirms the flag
// flows through to tls.Config. Self-signed dev clusters can opt in;
// production deployments MUST leave it false.
func TestBuildHTTPTransport_InsecureSkipVerifyOptIn(t *testing.T) {
t.Parallel()
tr := buildHTTPTransport(true)
if tr.TLSClientConfig == nil {
t.Fatal("TLSClientConfig must be set")
}
if !tr.TLSClientConfig.InsecureSkipVerify {
t.Fatal("InsecureSkipVerify=true must propagate to tls.Config")
}
}
// TestBuildHTTPTransport_TLS12MinVersion pins the TLS 1.2 hard floor.
// TLS 1.3 negotiates automatically when both ends support it; 1.0 / 1.1
// must be rejected at handshake time.
func TestBuildHTTPTransport_TLS12MinVersion(t *testing.T) {
t.Parallel()
tr := buildHTTPTransport(false)
if tr.TLSClientConfig.MinVersion != tls.VersionTLS12 {
t.Fatalf("MinVersion: want tls.VersionTLS12 (0x%x), got 0x%x",
tls.VersionTLS12, tr.TLSClientConfig.MinVersion)
}
}
// TestBuildHTTPTransport_CipherSuitesExcludeNonECDHE pins the explicit
// cipher list to forward-secrecy-only ECDHE suites. TLS 1.2's default
// list includes TLS_RSA_* (no forward secrecy); we drop them. TLS 1.3
// cipher selection is automatic so this only matters for the 1.2 path.
func TestBuildHTTPTransport_CipherSuitesExcludeNonECDHE(t *testing.T) {
t.Parallel()
tr := buildHTTPTransport(false)
cs := tr.TLSClientConfig.CipherSuites
if len(cs) == 0 {
t.Fatal("CipherSuites must be explicitly pinned (got empty)")
}
allowed := map[uint16]bool{
tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256: true,
tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256: true,
tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384: true,
tls.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384: true,
tls.TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305: true,
tls.TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305: true,
}
for _, suite := range cs {
if !allowed[suite] {
t.Fatalf("cipher suite 0x%x is not in the forward-secrecy whitelist", suite)
}
}
// Spot-check that a known-weak non-ECDHE cipher is NOT in the list.
for _, suite := range cs {
if suite == tls.TLS_RSA_WITH_AES_128_GCM_SHA256 {
t.Fatal("TLS_RSA_WITH_AES_128_GCM_SHA256 must NOT be present (no forward secrecy)")
}
}
}
// TestBuildHTTPTransport_PoolTuning pins the connection-pool settings.
// The Go default MaxIdleConnsPerHost is 2, which is too low for concurrent
// queries across multiple knowledge bases.
func TestBuildHTTPTransport_PoolTuning(t *testing.T) {
t.Parallel()
tr := buildHTTPTransport(false)
if tr.MaxIdleConnsPerHost < 32 {
t.Errorf("MaxIdleConnsPerHost: want >=32 for concurrent KB queries, got %d",
tr.MaxIdleConnsPerHost)
}
if tr.IdleConnTimeout == 0 {
t.Error("IdleConnTimeout must be set (Go default of 0 leaks idle conns)")
}
if tr.ResponseHeaderTimeout == 0 {
t.Error("ResponseHeaderTimeout must be set (Go default of 0 means unbounded)")
}
if tr.ExpectContinueTimeout == 0 {
t.Error("ExpectContinueTimeout must be set (avoids 100-continue stalls)")
}
}
@@ -0,0 +1,699 @@
package postgres
import (
"context"
"errors"
"fmt"
"strings"
"github.com/Tencent/WeKnora/internal/common"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
"github.com/google/uuid"
"github.com/pgvector/pgvector-go"
"gorm.io/gorm"
"gorm.io/gorm/clause"
)
// pgRepository implements PostgreSQL-based retrieval operations
type pgRepository struct {
db *gorm.DB // Database connection
}
// NewPostgresRetrieveEngineRepository creates a new PostgreSQL retriever repository
func NewPostgresRetrieveEngineRepository(db *gorm.DB) interfaces.RetrieveEngineRepository {
logger.GetLogger(context.Background()).Info("[Postgres] Initializing PostgreSQL retriever engine repository")
return &pgRepository{db: db}
}
// EngineType returns the retriever engine type (PostgreSQL)
func (r *pgRepository) EngineType() types.RetrieverEngineType {
return types.PostgresRetrieverEngineType
}
// Support returns supported retriever types (keywords and vector)
func (r *pgRepository) Support() []types.RetrieverType {
return []types.RetrieverType{types.KeywordsRetrieverType, types.VectorRetrieverType}
}
// calculateIndexStorageSize calculates storage size for a single index entry
func (g *pgRepository) calculateIndexStorageSize(embeddingDB *pgVector) int64 {
// 1. Text content size
contentSizeBytes := int64(len(embeddingDB.Content))
// 2. Vector storage size (2 bytes per dimension for half-precision float)
var vectorSizeBytes int64 = 0
if embeddingDB.Dimension > 0 {
vectorSizeBytes = int64(embeddingDB.Dimension * 2)
}
// 3. Metadata size (fixed overhead for IDs, timestamps etc.)
metadataSizeBytes := int64(200)
// 4. Index overhead (HNSW index is ~2x vector size)
indexOverheadBytes := vectorSizeBytes * 2
// Total size in bytes
totalSizeBytes := contentSizeBytes + vectorSizeBytes + metadataSizeBytes + indexOverheadBytes
return totalSizeBytes
}
// EstimateStorageSize estimates total storage size for multiple indices
func (g *pgRepository) EstimateStorageSize(
ctx context.Context, indexInfoList []*types.IndexInfo, additionalParams map[string]any,
) int64 {
var totalStorageSize int64 = 0
for _, indexInfo := range indexInfoList {
embeddingDB := toDBVectorEmbedding(indexInfo, additionalParams)
totalStorageSize += g.calculateIndexStorageSize(embeddingDB)
}
logger.GetLogger(ctx).Infof(
"[Postgres] Estimated storage size for %d indices: %d bytes",
len(indexInfoList), totalStorageSize,
)
return totalStorageSize
}
// Save stores a single index entry
func (g *pgRepository) Save(ctx context.Context, indexInfo *types.IndexInfo, additionalParams map[string]any) error {
logger.GetLogger(ctx).Debugf("[Postgres] Saving index for source ID: %s", indexInfo.SourceID)
embeddingDB := toDBVectorEmbedding(indexInfo, additionalParams)
err := g.db.WithContext(ctx).Create(embeddingDB).Error
if err != nil {
logger.GetLogger(ctx).Errorf("[Postgres] Failed to save index: %v", err)
return err
}
logger.GetLogger(ctx).Infof("[Postgres] Successfully saved index for source ID: %s", indexInfo.SourceID)
return nil
}
// BatchSave stores multiple index entries in batch
func (g *pgRepository) BatchSave(
ctx context.Context, indexInfoList []*types.IndexInfo, additionalParams map[string]any,
) error {
logger.GetLogger(ctx).Infof("[Postgres] Batch saving %d indices", len(indexInfoList))
indexInfoDBList := make([]*pgVector, len(indexInfoList))
for i := range indexInfoList {
indexInfoDBList[i] = toDBVectorEmbedding(indexInfoList[i], additionalParams)
}
err := g.db.WithContext(ctx).Clauses(clause.OnConflict{DoNothing: true}).Create(indexInfoDBList).Error
if err != nil {
logger.GetLogger(ctx).Errorf("[Postgres] Batch save failed: %v", err)
return err
}
logger.GetLogger(ctx).Infof("[Postgres] Successfully batch saved %d indices", len(indexInfoList))
return nil
}
// DeleteByChunkIDList deletes indices by chunk IDs
func (g *pgRepository) DeleteByChunkIDList(ctx context.Context, chunkIDList []string, dimension int, knowledgeType string) error {
logger.GetLogger(ctx).Infof("[Postgres] Deleting indices by chunk IDs, count: %d", len(chunkIDList))
result := g.db.WithContext(ctx).Where("chunk_id IN ?", chunkIDList).Delete(&pgVector{})
if result.Error != nil {
logger.GetLogger(ctx).Errorf("[Postgres] Failed to delete indices by chunk IDs: %v", result.Error)
return result.Error
}
logger.GetLogger(ctx).Infof("[Postgres] Successfully deleted %d indices by chunk IDs", result.RowsAffected)
return nil
}
// DeleteBySourceIDList deletes indices by source IDs
func (g *pgRepository) DeleteBySourceIDList(ctx context.Context, sourceIDList []string, dimension int, knowledgeType string) error {
if len(sourceIDList) == 0 {
return nil
}
logger.GetLogger(ctx).Infof("[Postgres] Deleting indices by source IDs, count: %d", len(sourceIDList))
result := g.db.WithContext(ctx).Where("source_id IN ?", sourceIDList).Delete(&pgVector{})
if result.Error != nil {
logger.GetLogger(ctx).Errorf("[Postgres] Failed to delete indices by source IDs: %v", result.Error)
return result.Error
}
logger.GetLogger(ctx).Infof("[Postgres] Successfully deleted %d indices by source IDs", result.RowsAffected)
return nil
}
// DeleteByKnowledgeIDList deletes indices by knowledge IDs
func (g *pgRepository) DeleteByKnowledgeIDList(ctx context.Context, knowledgeIDList []string, dimension int, knowledgeType string) error {
logger.GetLogger(ctx).Infof("[Postgres] Deleting indices by knowledge IDs, count: %d", len(knowledgeIDList))
result := g.db.WithContext(ctx).Where("knowledge_id IN ?", knowledgeIDList).Delete(&pgVector{})
if result.Error != nil {
logger.GetLogger(ctx).Errorf("[Postgres] Failed to delete indices by knowledge IDs: %v", result.Error)
return result.Error
}
logger.GetLogger(ctx).Infof("[Postgres] Successfully deleted %d indices by knowledge IDs", result.RowsAffected)
return nil
}
// Retrieve handles retrieval requests and routes to appropriate method
func (g *pgRepository) Retrieve(ctx context.Context, params types.RetrieveParams) ([]*types.RetrieveResult, error) {
logger.GetLogger(ctx).Debugf("[Postgres] Processing retrieval request of type: %s", params.RetrieverType)
switch params.RetrieverType {
case types.KeywordsRetrieverType:
return g.KeywordsRetrieve(ctx, params)
case types.VectorRetrieverType:
return g.VectorRetrieve(ctx, params)
}
err := errors.New("invalid retriever type")
logger.GetLogger(ctx).Errorf("[Postgres] %v: %s", err, params.RetrieverType)
return nil, err
}
// KeywordsRetrieve performs keyword-based search using PostgreSQL full-text search
func (g *pgRepository) KeywordsRetrieve(ctx context.Context,
params types.RetrieveParams,
) ([]*types.RetrieveResult, error) {
logger.GetLogger(ctx).Infof("[Postgres] Keywords retrieval: query=%s, topK=%d", params.Query, params.TopK)
conds := make([]clause.Expression, 0)
// KnowledgeBaseIDs and KnowledgeIDs use AND logic
// - If only KnowledgeBaseIDs: search entire knowledge bases
// - If only KnowledgeIDs: search specific documents
// - If both: search specific documents within the knowledge bases (AND)
if len(params.KnowledgeBaseIDs) > 0 {
logger.GetLogger(ctx).Debugf("[Postgres] Filtering by knowledge base IDs: %v", params.KnowledgeBaseIDs)
conds = append(conds, clause.IN{
Column: "knowledge_base_id",
Values: common.ToInterfaceSlice(params.KnowledgeBaseIDs),
})
}
if len(params.KnowledgeIDs) > 0 {
logger.GetLogger(ctx).Debugf("[Postgres] Filtering by knowledge IDs: %v", params.KnowledgeIDs)
conds = append(conds, clause.IN{
Column: "knowledge_id",
Values: common.ToInterfaceSlice(params.KnowledgeIDs),
})
}
// Filter by tag IDs if specified
if len(params.TagIDs) > 0 {
logger.GetLogger(ctx).Debugf("[Postgres] Filtering by tag IDs: %v", params.TagIDs)
conds = append(conds, clause.IN{
Column: "tag_id",
Values: common.ToInterfaceSlice(params.TagIDs),
})
}
// Use ParadeDB's ||| operator for matching any token
conds = append(conds, clause.Expr{
SQL: "content ||| ?",
Vars: []interface{}{params.Query},
})
// Filter by is_enabled = true or NULL (NULL means enabled for historical data)
conds = append(conds, clause.Expr{
SQL: "(is_enabled IS NULL OR is_enabled = ?)",
Vars: []interface{}{true},
})
conds = append(conds, clause.OrderBy{Columns: []clause.OrderByColumn{
{Column: clause.Column{Name: "score"}, Desc: true},
}})
var embeddingDBList []pgVectorWithScore
err := g.db.WithContext(ctx).Clauses(conds...).Debug().
Select([]string{
"paradedb.score(id) as score",
"id",
"content",
"source_id",
"source_type",
"chunk_id",
"knowledge_id",
"knowledge_base_id",
"tag_id",
}).
Limit(int(params.TopK)).
Find(&embeddingDBList).Error
if err == gorm.ErrRecordNotFound {
logger.GetLogger(ctx).Warnf("[Postgres] No records found for keywords query: %s", params.Query)
return nil, nil
}
if err != nil {
logger.GetLogger(ctx).Errorf("[Postgres] Keywords retrieval failed: %v", err)
return nil, err
}
logger.GetLogger(ctx).Infof("[Postgres] Keywords retrieval found %d results", len(embeddingDBList))
results := make([]*types.IndexWithScore, len(embeddingDBList))
const maxKeywordResultLog = 8
for i := range embeddingDBList {
results[i] = fromDBVectorEmbeddingWithScore(&embeddingDBList[i], types.MatchTypeKeywords)
if i < maxKeywordResultLog {
logger.GetLogger(ctx).Debugf("[Postgres] Keywords result %d: chunk=%s, score=%f",
i, results[i].ChunkID, results[i].Score)
}
}
if len(results) > maxKeywordResultLog {
logger.GetLogger(ctx).Debugf(
"[Postgres] Keywords result summary: total=%d logged=%d truncated=%d",
len(results), maxKeywordResultLog, len(results)-maxKeywordResultLog,
)
}
return []*types.RetrieveResult{
{
Results: results,
RetrieverEngineType: types.PostgresRetrieverEngineType,
RetrieverType: types.KeywordsRetrieverType,
Error: nil,
},
}, nil
}
// VectorRetrieve performs vector similarity search using pgvector
// Optimized to use HNSW index efficiently and avoid recalculating vector distance
func (g *pgRepository) VectorRetrieve(ctx context.Context,
params types.RetrieveParams,
) ([]*types.RetrieveResult, error) {
logger.GetLogger(ctx).Infof("[Postgres] Vector retrieval: dim=%d, topK=%d, threshold=%.4f",
len(params.Embedding), params.TopK, params.Threshold)
dimension := len(params.Embedding)
queryVector := pgvector.NewHalfVector(params.Embedding)
// Build WHERE conditions for filtering
whereParts := make([]string, 0)
allVars := make([]interface{}, 0)
// Add query vector first (used in ORDER BY for HNSW index)
allVars = append(allVars, queryVector)
// Dimension filter (required for HNSW index WHERE clause)
whereParts = append(whereParts, fmt.Sprintf("dimension = $%d", len(allVars)+1))
allVars = append(allVars, dimension)
// KnowledgeBaseIDs and KnowledgeIDs use AND logic
// - If only KnowledgeBaseIDs: search entire knowledge bases
// - If only KnowledgeIDs: search specific documents
// - If both: search specific documents within the knowledge bases (AND)
if len(params.KnowledgeBaseIDs) > 0 {
logger.GetLogger(ctx).Debugf(
"[Postgres] Filtering vector search by knowledge base IDs: %v",
params.KnowledgeBaseIDs,
)
placeholders := make([]string, len(params.KnowledgeBaseIDs))
paramStart := len(allVars) + 1
for i := range params.KnowledgeBaseIDs {
placeholders[i] = fmt.Sprintf("$%d", paramStart+i)
allVars = append(allVars, params.KnowledgeBaseIDs[i])
}
whereParts = append(whereParts, fmt.Sprintf("knowledge_base_id IN (%s)",
strings.Join(placeholders, ", ")))
}
if len(params.KnowledgeIDs) > 0 {
logger.GetLogger(ctx).Debugf(
"[Postgres] Filtering vector search by knowledge IDs: %v",
params.KnowledgeIDs,
)
placeholders := make([]string, len(params.KnowledgeIDs))
paramStart := len(allVars) + 1
for i := range params.KnowledgeIDs {
placeholders[i] = fmt.Sprintf("$%d", paramStart+i)
allVars = append(allVars, params.KnowledgeIDs[i])
}
whereParts = append(whereParts, fmt.Sprintf("knowledge_id IN (%s)",
strings.Join(placeholders, ", ")))
}
// Filter by tag IDs if specified
if len(params.TagIDs) > 0 {
logger.GetLogger(ctx).Debugf(
"[Postgres] Filtering vector search by tag IDs: %v",
params.TagIDs,
)
placeholders := make([]string, len(params.TagIDs))
paramStart := len(allVars) + 1
for i := range params.TagIDs {
placeholders[i] = fmt.Sprintf("$%d", paramStart+i)
allVars = append(allVars, params.TagIDs[i])
}
whereParts = append(whereParts, fmt.Sprintf("tag_id IN (%s)",
strings.Join(placeholders, ", ")))
}
// is_enabled filter
whereParts = append(whereParts, fmt.Sprintf("(is_enabled IS NULL OR is_enabled = $%d)", len(allVars)+1))
allVars = append(allVars, true)
// Build WHERE clause string
whereClause := ""
if len(whereParts) > 0 {
whereClause = "WHERE " + strings.Join(whereParts, " AND ")
}
// Expand TopK to get more candidates before threshold filtering.
//
// HNSW requires `ef_search >= LIMIT`, and a very large LIMIT (e.g. 1000)
// forces HNSW to walk a near-exhaustive portion of the graph, often making
// it slower than a sequential scan and pushing the planner to pick Seq Scan
// even when an index exists. 200 is a good sweet spot: it gives enough
// headroom for threshold/filter post-processing without ballooning ef_search.
expandedTopK := params.TopK * 2
if expandedTopK < 100 {
expandedTopK = 100 // Minimum 100 candidates
}
if expandedTopK > 200 {
expandedTopK = 200 // Maximum 200 candidates (keeps HNSW efficient)
}
if expandedTopK < params.TopK {
expandedTopK = params.TopK // Ensure subquery limit is at least final limit
}
// Optimized query: Use subquery to calculate distance once.
//
// IMPORTANT: The HNSW index in this project is built on the EXPRESSION
// (embedding::halfvec(<dim>)) halfvec_cosine_ops
// because the `embedding` column itself is `halfvec` without a fixed dimension
// (so the table can store multiple embedding sizes such as 798 / 3584 / ...).
//
// pgvector requires the ORDER BY expression to match the indexed expression
// EXACTLY, otherwise the planner falls back to a sequential scan. The
// `embedding::halfvec(%d)` cast on both sides of `<=>` is therefore NOT
// redundant — it is the only way to make the HNSW index get used at all.
// See: pgvector issues #702, #835 and ParadeDB "indexing-expressions" docs.
subqueryLimitParam := len(allVars) + 1
thresholdParam := len(allVars) + 2
finalLimitParam := len(allVars) + 3
querySQL := fmt.Sprintf(`
SELECT
id, content, source_id, source_type, chunk_id, knowledge_id, knowledge_base_id, tag_id,
(1 - distance) as score
FROM (
SELECT
id, content, source_id, source_type, chunk_id, knowledge_id, knowledge_base_id, tag_id,
embedding::halfvec(%[1]d) <=> $1::halfvec(%[1]d) as distance
FROM embeddings
%[2]s
ORDER BY embedding::halfvec(%[1]d) <=> $1::halfvec(%[1]d)
LIMIT $%[3]d
) AS candidates
WHERE distance <= $%[4]d
ORDER BY distance ASC
LIMIT $%[5]d
`, dimension, whereClause, subqueryLimitParam, thresholdParam, finalLimitParam)
allVars = append(allVars, expandedTopK) // LIMIT in subquery
allVars = append(allVars, 1-params.Threshold) // Distance threshold
allVars = append(allVars, params.TopK) // Final LIMIT
// HNSW's `ef_search` defaults to 40, which is much smaller than our
// `expandedTopK` budget (up to 1000). Without raising it, HNSW would only
// return ~40 candidates per scan and our outer threshold/filter step would
// silently lose recall. `SET LOCAL` requires a transaction so we wrap the
// query in one. We never write inside this transaction, so the cost is
// negligible.
efSearch := expandedTopK
if efSearch < 40 {
efSearch = 40
}
var embeddingDBList []pgVectorWithScore
err := g.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
if err := tx.Exec(fmt.Sprintf("SET LOCAL hnsw.ef_search = %d", efSearch)).Error; err != nil {
// Treat as non-fatal: pgvector should always expose this GUC, but if
// for any reason it does not we still want the query to run (just
// with default recall). We must rollback first because a failed
// statement aborts the transaction in PostgreSQL.
logger.GetLogger(ctx).Warnf("[Postgres] Failed to set hnsw.ef_search=%d: %v", efSearch, err)
return err
}
// pgvector >= 0.8 supports iterative scan, which keeps pulling more
// candidates from HNSW until the post-filter (knowledge_base_id /
// knowledge_id / tag_id / is_enabled) yields enough rows. Without it,
// HNSW returns at most ef_search candidates and the outer filter may
// silently lose recall when the filter is selective.
// Best-effort: ignore failure on older pgvector versions.
if err := tx.Exec("SET LOCAL hnsw.iterative_scan = strict_order").Error; err != nil {
logger.GetLogger(ctx).Debugf("[Postgres] hnsw.iterative_scan not available: %v", err)
// abort transaction and let the fallback path below handle it.
return err
}
return tx.Raw(querySQL, allVars...).Scan(&embeddingDBList).Error
})
// Fallback: if the transaction failed because of an unsupported GUC (e.g.
// older pgvector that doesn't have hnsw.ef_search or hnsw.iterative_scan),
// retry the query without the SETs so we still return results.
if err != nil && len(embeddingDBList) == 0 &&
(strings.Contains(err.Error(), "hnsw.ef_search") ||
strings.Contains(err.Error(), "hnsw.iterative_scan")) {
logger.GetLogger(ctx).Warnf("[Postgres] Retrying vector query without HNSW GUC overrides: %v", err)
err = g.db.WithContext(ctx).Raw(querySQL, allVars...).Scan(&embeddingDBList).Error
}
if err == gorm.ErrRecordNotFound {
logger.GetLogger(ctx).Warnf("[Postgres] No vector matches found that meet threshold %.4f", params.Threshold)
return nil, nil
}
if err != nil {
logger.GetLogger(ctx).Errorf("[Postgres] Vector retrieval failed: %v", err)
return nil, err
}
// Apply final TopK limit (in case we got more results than needed)
if len(embeddingDBList) > int(params.TopK) {
embeddingDBList = embeddingDBList[:params.TopK]
}
logger.GetLogger(ctx).Infof("[Postgres] Vector retrieval found %d results", len(embeddingDBList))
results := make([]*types.IndexWithScore, len(embeddingDBList))
const maxVectorResultLog = 8
for i := range embeddingDBList {
results[i] = fromDBVectorEmbeddingWithScore(&embeddingDBList[i], types.MatchTypeEmbedding)
if i < maxVectorResultLog {
logger.GetLogger(ctx).Debugf("[Postgres] Vector search result %d: chunk_id %s, score %.4f",
i, results[i].ChunkID, results[i].Score)
}
}
if len(results) > maxVectorResultLog {
logger.GetLogger(ctx).Debugf(
"[Postgres] Vector search result summary: total=%d logged=%d truncated=%d",
len(results), maxVectorResultLog, len(results)-maxVectorResultLog,
)
}
return []*types.RetrieveResult{
{
Results: results,
RetrieverEngineType: types.PostgresRetrieverEngineType,
RetrieverType: types.VectorRetrieverType,
Error: nil,
},
}, nil
}
// CopyIndices copies index data
func (g *pgRepository) CopyIndices(ctx context.Context,
sourceKnowledgeBaseID string,
sourceToTargetKBIDMap map[string]string,
sourceToTargetChunkIDMap map[string]string,
targetKnowledgeBaseID string,
dimension int,
knowledgeType string,
) error {
logger.GetLogger(ctx).Infof(
"[Postgres] Copying indices, source knowledge base: %s, target knowledge base: %s, mapping count: %d",
sourceKnowledgeBaseID, targetKnowledgeBaseID, len(sourceToTargetChunkIDMap),
)
if len(sourceToTargetChunkIDMap) == 0 {
logger.GetLogger(ctx).Warnf("[Postgres] Mapping is empty, no need to copy")
return nil
}
// Batch processing parameters
batchSize := 500 // Number of records to process per batch
offset := 0 // Offset for pagination
totalCopied := 0 // Total number of copied records
for {
// Paginated query for source data
var sourceVectors []*pgVector
if err := g.db.WithContext(ctx).
Where("knowledge_base_id = ?", sourceKnowledgeBaseID).
Limit(batchSize).
Offset(offset).
Find(&sourceVectors).Error; err != nil {
logger.GetLogger(ctx).Errorf("[Postgres] Failed to query source index data: %v", err)
return err
}
// If no more data, exit the loop
if len(sourceVectors) == 0 {
if offset == 0 {
logger.GetLogger(ctx).Warnf("[Postgres] No source index data found")
}
break
}
batchCount := len(sourceVectors)
logger.GetLogger(ctx).Infof(
"[Postgres] Found %d source index data, batch start position: %d",
batchCount, offset,
)
// Create target vector index
targetVectors := make([]*pgVector, 0, batchCount)
for _, sourceVector := range sourceVectors {
// Get the mapped target chunk ID
targetChunkID, ok := sourceToTargetChunkIDMap[sourceVector.ChunkID]
if !ok {
logger.GetLogger(ctx).Warnf(
"[Postgres] Source chunk %s not found in target chunk mapping, skipping",
sourceVector.ChunkID,
)
continue
}
// Get the mapped target knowledge ID
targetKnowledgeID, ok := sourceToTargetKBIDMap[sourceVector.KnowledgeID]
if !ok {
logger.GetLogger(ctx).Warnf(
"[Postgres] Source knowledge %s not found in target knowledge mapping, skipping",
sourceVector.KnowledgeID,
)
continue
}
// Handle SourceID transformation for generated questions
// Generated questions have SourceID format: {chunkID}-{questionID}
// Regular chunks have SourceID == ChunkID
var targetSourceID string
if sourceVector.SourceID == sourceVector.ChunkID {
// Regular chunk, use targetChunkID as SourceID
targetSourceID = targetChunkID
} else if strings.HasPrefix(sourceVector.SourceID, sourceVector.ChunkID+"-") {
// This is a generated question, preserve the questionID part
questionID := strings.TrimPrefix(sourceVector.SourceID, sourceVector.ChunkID+"-")
targetSourceID = fmt.Sprintf("%s-%s", targetChunkID, questionID)
} else {
// For other complex scenarios, generate new unique SourceID
targetSourceID = uuid.New().String()
}
// Create new vector index, copy the content and vector of the source index
targetVector := &pgVector{
Content: sourceVector.Content,
SourceID: targetSourceID, // Handle SourceID transformation properly
SourceType: sourceVector.SourceType,
ChunkID: targetChunkID, // Update to target chunk ID
KnowledgeID: targetKnowledgeID, // Update to target knowledge ID
KnowledgeBaseID: targetKnowledgeBaseID, // Update to target knowledge base ID
Dimension: sourceVector.Dimension,
Embedding: sourceVector.Embedding, // Copy the vector embedding directly, avoid recalculation
}
targetVectors = append(targetVectors, targetVector)
}
// Batch insert target vector index
if len(targetVectors) > 0 {
if err := g.db.WithContext(ctx).
Clauses(clause.OnConflict{DoNothing: true}).Create(targetVectors).Error; err != nil {
logger.GetLogger(ctx).Errorf("[Postgres] Failed to batch create target index: %v", err)
return err
}
totalCopied += len(targetVectors)
logger.GetLogger(ctx).Infof(
"[Postgres] Successfully copied batch data, batch size: %d, total copied: %d",
len(targetVectors),
totalCopied,
)
}
// Move to the next batch
offset += batchCount
// If the number of returned records is less than the requested size, it means the last page has been reached
if batchCount < batchSize {
break
}
}
logger.GetLogger(ctx).Infof("[Postgres] Index copying completed, total copied: %d", totalCopied)
return nil
}
// BatchUpdateChunkEnabledStatus updates the enabled status of chunks in batch
func (g *pgRepository) BatchUpdateChunkEnabledStatus(ctx context.Context, chunkStatusMap map[string]bool) error {
if len(chunkStatusMap) == 0 {
logger.GetLogger(ctx).Warnf("[Postgres] Chunk status map is empty, skipping update")
return nil
}
logger.GetLogger(ctx).Infof("[Postgres] Batch updating chunk enabled status, count: %d", len(chunkStatusMap))
// Group chunks by enabled status for batch updates
enabledChunkIDs := make([]string, 0)
disabledChunkIDs := make([]string, 0)
for chunkID, enabled := range chunkStatusMap {
if enabled {
enabledChunkIDs = append(enabledChunkIDs, chunkID)
} else {
disabledChunkIDs = append(disabledChunkIDs, chunkID)
}
}
// Batch update enabled chunks
if len(enabledChunkIDs) > 0 {
result := g.db.WithContext(ctx).Model(&pgVector{}).
Where("chunk_id IN ?", enabledChunkIDs).
Update("is_enabled", true)
if result.Error != nil {
logger.GetLogger(ctx).Errorf("[Postgres] Failed to update enabled chunks: %v", result.Error)
return result.Error
}
logger.GetLogger(ctx).
Infof("[Postgres] Updated %d chunks to enabled, rows affected: %d", len(enabledChunkIDs), result.RowsAffected)
}
// Batch update disabled chunks
if len(disabledChunkIDs) > 0 {
result := g.db.WithContext(ctx).Model(&pgVector{}).
Where("chunk_id IN ?", disabledChunkIDs).
Update("is_enabled", false)
if result.Error != nil {
logger.GetLogger(ctx).Errorf("[Postgres] Failed to update disabled chunks: %v", result.Error)
return result.Error
}
logger.GetLogger(ctx).
Infof("[Postgres] Updated %d chunks to disabled, rows affected: %d", len(disabledChunkIDs), result.RowsAffected)
}
logger.GetLogger(ctx).Infof("[Postgres] Successfully batch updated chunk enabled status")
return nil
}
// BatchUpdateChunkTagID updates the tag ID of chunks in batch
func (g *pgRepository) BatchUpdateChunkTagID(ctx context.Context, chunkTagMap map[string]string) error {
if len(chunkTagMap) == 0 {
logger.GetLogger(ctx).Warnf("[Postgres] Chunk tag map is empty, skipping update")
return nil
}
logger.GetLogger(ctx).Infof("[Postgres] Batch updating chunk tag ID, count: %d", len(chunkTagMap))
// Group chunks by tag ID for batch updates
tagGroups := make(map[string][]string)
for chunkID, tagID := range chunkTagMap {
tagGroups[tagID] = append(tagGroups[tagID], chunkID)
}
// Batch update chunks for each tag ID
for tagID, chunkIDs := range tagGroups {
result := g.db.WithContext(ctx).Model(&pgVector{}).
Where("chunk_id IN ?", chunkIDs).
Update("tag_id", tagID)
if result.Error != nil {
logger.GetLogger(ctx).Errorf("[Postgres] Failed to update chunks with tag_id %s: %v", tagID, result.Error)
return result.Error
}
logger.GetLogger(ctx).
Infof("[Postgres] Updated %d chunks to tag_id=%s, rows affected: %d", len(chunkIDs), tagID, result.RowsAffected)
}
logger.GetLogger(ctx).Infof("[Postgres] Successfully batch updated chunk tag ID")
return nil
}
@@ -0,0 +1,103 @@
package postgres
import (
"maps"
"slices"
"strconv"
"time"
"github.com/Tencent/WeKnora/internal/common"
"github.com/Tencent/WeKnora/internal/types"
"github.com/pgvector/pgvector-go"
)
// pgVector defines the database model for vector embeddings storage
type pgVector struct {
ID uint `json:"id" gorm:"primarykey"`
CreatedAt time.Time `json:"created_at" gorm:"column:created_at"`
UpdatedAt time.Time `json:"updated_at" gorm:"column:updated_at"`
SourceID string `json:"source_id" gorm:"column:source_id;not null"`
SourceType int `json:"source_type" gorm:"column:source_type;not null"`
ChunkID string `json:"chunk_id" gorm:"column:chunk_id"`
KnowledgeID string `json:"knowledge_id" gorm:"column:knowledge_id"`
KnowledgeBaseID string `json:"knowledge_base_id" gorm:"column:knowledge_base_id"`
TagID string `json:"tag_id" gorm:"column:tag_id;index"`
Content string `json:"content" gorm:"column:content;not null"`
Dimension int `json:"dimension" gorm:"column:dimension;not null"`
Embedding pgvector.HalfVector `json:"embedding" gorm:"column:embedding;not null"`
IsEnabled bool `json:"is_enabled" gorm:"column:is_enabled;default:true;index"`
}
// pgVectorWithScore extends pgVector with similarity score field
type pgVectorWithScore struct {
ID uint `json:"id" gorm:"primarykey"`
CreatedAt time.Time `json:"created_at" gorm:"column:created_at"`
UpdatedAt time.Time `json:"updated_at" gorm:"column:updated_at"`
SourceID string `json:"source_id" gorm:"column:source_id;not null"`
SourceType int `json:"source_type" gorm:"column:source_type;not null"`
ChunkID string `json:"chunk_id" gorm:"column:chunk_id"`
KnowledgeID string `json:"knowledge_id" gorm:"column:knowledge_id"`
KnowledgeBaseID string `json:"knowledge_base_id" gorm:"column:knowledge_base_id"`
TagID string `json:"tag_id" gorm:"column:tag_id;index"`
Content string `json:"content" gorm:"column:content;not null"`
Dimension int `json:"dimension" gorm:"column:dimension;not null"`
Embedding pgvector.HalfVector `json:"embedding" gorm:"column:embedding;not null"`
IsEnabled bool `json:"is_enabled" gorm:"column:is_enabled;default:true;index"`
Score float64 `json:"score" gorm:"column:score"`
}
// TableName specifies the database table name for pgVectorWithScore
func (pgVectorWithScore) TableName() string {
return "embeddings"
}
// TableName specifies the database table name for pgVector
func (pgVector) TableName() string {
return "embeddings"
}
// toDBVectorEmbedding converts IndexInfo to pgVector database model
func toDBVectorEmbedding(indexInfo *types.IndexInfo, additionalParams map[string]any) *pgVector {
pgVector := &pgVector{
SourceID: indexInfo.SourceID,
SourceType: int(indexInfo.SourceType),
ChunkID: indexInfo.ChunkID,
KnowledgeID: indexInfo.KnowledgeID,
KnowledgeBaseID: indexInfo.KnowledgeBaseID,
TagID: indexInfo.TagID,
Content: common.CleanInvalidUTF8(indexInfo.Content),
IsEnabled: indexInfo.IsEnabled,
}
// Add embedding data if available in additionalParams
if additionalParams != nil && slices.Contains(slices.Collect(maps.Keys(additionalParams)), "embedding") {
if embeddingMap, ok := additionalParams["embedding"].(map[string][]float32); ok {
pgVector.Embedding = pgvector.NewHalfVector(embeddingMap[indexInfo.SourceID])
pgVector.Dimension = len(pgVector.Embedding.Slice())
}
}
// Get is_enabled from additionalParams if available
if additionalParams != nil {
if chunkEnabledMap, ok := additionalParams["chunk_enabled"].(map[string]bool); ok {
if enabled, exists := chunkEnabledMap[indexInfo.ChunkID]; exists {
pgVector.IsEnabled = enabled
}
}
}
return pgVector
}
// fromDBVectorEmbeddingWithScore converts pgVectorWithScore to IndexWithScore domain model
func fromDBVectorEmbeddingWithScore(embedding *pgVectorWithScore, matchType types.MatchType) *types.IndexWithScore {
return &types.IndexWithScore{
ID: strconv.FormatInt(int64(embedding.ID), 10),
SourceID: embedding.SourceID,
SourceType: types.SourceType(embedding.SourceType),
ChunkID: embedding.ChunkID,
KnowledgeID: embedding.KnowledgeID,
KnowledgeBaseID: embedding.KnowledgeBaseID,
TagID: embedding.TagID,
Content: embedding.Content,
Score: embedding.Score,
MatchType: matchType,
}
}
@@ -0,0 +1,995 @@
package qdrant
import (
"context"
"fmt"
"strings"
"unicode/utf8"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
"github.com/google/uuid"
"github.com/qdrant/go-client/qdrant"
)
const (
envQdrantCollection = "QDRANT_COLLECTION"
defaultCollectionName = "weknora_embeddings"
fieldContent = "content"
fieldSourceID = "source_id"
fieldSourceType = "source_type"
fieldChunkID = "chunk_id"
fieldKnowledgeID = "knowledge_id"
fieldKnowledgeBaseID = "knowledge_base_id"
fieldTagID = "tag_id"
fieldEmbedding = "embedding"
fieldIsEnabled = "is_enabled"
)
// NewQdrantRetrieveEngineRepository creates and initializes a new Qdrant repository.
// indexCfg is optional — pass nil to use env var / default values (env path).
func NewQdrantRetrieveEngineRepository(client *qdrant.Client, indexCfg *types.IndexConfig) interfaces.RetrieveEngineRepository {
log := logger.GetLogger(context.Background())
log.Info("[Qdrant] Initializing Qdrant retriever engine repository")
collectionBaseName := types.ResolveCollectionName(indexCfg, envQdrantCollection, defaultCollectionName)
res := &qdrantRepository{
client: client,
collectionBaseName: collectionBaseName,
shardNumber: indexCfg.GetShardNumber(0),
replicationFactor: indexCfg.GetReplicationFactor(0),
}
log.Info("[Qdrant] Successfully initialized repository")
return res
}
// getCollectionName returns the collection name for a specific dimension
func (q *qdrantRepository) getCollectionName(dimension int) string {
return fmt.Sprintf("%s_%d", q.collectionBaseName, dimension)
}
// ensureCollection ensures the collection exists for the given dimension
func (q *qdrantRepository) ensureCollection(ctx context.Context, dimension int) error {
collectionName := q.getCollectionName(dimension)
// Check cache first
if _, ok := q.initializedCollections.Load(dimension); ok {
return nil
}
log := logger.GetLogger(ctx)
// Check if collection exists
exists, err := q.client.CollectionExists(ctx, collectionName)
if err != nil {
log.Errorf("[Qdrant] Failed to check collection existence: %v", err)
return fmt.Errorf("failed to check collection existence: %w", err)
}
if !exists {
log.Infof("[Qdrant] Creating collection %s with dimension %d", collectionName, dimension)
err = q.client.CreateCollection(ctx, &qdrant.CreateCollection{
CollectionName: collectionName,
VectorsConfig: qdrant.NewVectorsConfig(&qdrant.VectorParams{
Size: uint64(dimension),
Distance: qdrant.Distance_Cosine,
}),
ShardNumber: types.OptionalUint32(q.shardNumber),
ReplicationFactor: types.OptionalUint32(q.replicationFactor),
})
if err != nil {
log.Errorf("[Qdrant] Failed to create collection: %v", err)
return fmt.Errorf("failed to create collection: %w", err)
}
// Create payload indexes for filtering
indexFields := []string{fieldChunkID, fieldKnowledgeID, fieldKnowledgeBaseID, fieldSourceID}
for _, field := range indexFields {
_, err = q.client.CreateFieldIndex(ctx, &qdrant.CreateFieldIndexCollection{
CollectionName: collectionName,
FieldName: field,
FieldType: qdrant.FieldType_FieldTypeKeyword.Enum(),
})
if err != nil {
log.Warnf("[Qdrant] Failed to create index for field %s: %v", field, err)
}
}
// Create bool index for is_enabled
_, err = q.client.CreateFieldIndex(ctx, &qdrant.CreateFieldIndexCollection{
CollectionName: collectionName,
FieldName: fieldIsEnabled,
FieldType: qdrant.FieldType_FieldTypeBool.Enum(),
})
if err != nil {
log.Warnf("[Qdrant] Failed to create index for field %s: %v", fieldIsEnabled, err)
}
// Create text index for content (for keyword search) with multilingual tokenizer
// This supports Chinese, Japanese, Korean and other languages
lowercase := true
_, err = q.client.CreateFieldIndex(ctx, &qdrant.CreateFieldIndexCollection{
CollectionName: collectionName,
FieldName: fieldContent,
FieldType: qdrant.FieldType_FieldTypeText.Enum(),
FieldIndexParams: &qdrant.PayloadIndexParams{
IndexParams: &qdrant.PayloadIndexParams_TextIndexParams{
TextIndexParams: &qdrant.TextIndexParams{
Tokenizer: qdrant.TokenizerType_Multilingual,
Lowercase: &lowercase,
},
},
},
})
if err != nil {
log.Warnf("[Qdrant] Failed to create text index for content: %v", err)
}
log.Infof("[Qdrant] Successfully created collection %s", collectionName)
}
// Mark as initialized
q.initializedCollections.Store(dimension, true)
return nil
}
func (q *qdrantRepository) EngineType() types.RetrieverEngineType {
return types.QdrantRetrieverEngineType
}
func (q *qdrantRepository) Support() []types.RetrieverType {
return []types.RetrieverType{types.KeywordsRetrieverType, types.VectorRetrieverType}
}
// EstimateStorageSize calculates the estimated storage size for a list of indices
func (q *qdrantRepository) EstimateStorageSize(ctx context.Context,
indexInfoList []*types.IndexInfo, params map[string]any,
) int64 {
var totalStorageSize int64
for _, embedding := range indexInfoList {
embeddingDB := toQdrantVectorEmbedding(embedding, params)
totalStorageSize += q.calculateStorageSize(embeddingDB)
}
logger.GetLogger(ctx).Infof(
"[Qdrant] Storage size for %d indices: %d bytes", len(indexInfoList), totalStorageSize,
)
return totalStorageSize
}
// Save stores a single point in Qdrant
func (q *qdrantRepository) Save(ctx context.Context,
embedding *types.IndexInfo,
additionalParams map[string]any,
) error {
log := logger.GetLogger(ctx)
log.Debugf("[Qdrant] Saving index for chunk ID: %s", embedding.ChunkID)
embeddingDB := toQdrantVectorEmbedding(embedding, additionalParams)
if len(embeddingDB.Embedding) == 0 {
err := fmt.Errorf("empty embedding vector for chunk ID: %s", embedding.ChunkID)
log.Errorf("[Qdrant] %v", err)
return err
}
dimension := len(embeddingDB.Embedding)
if err := q.ensureCollection(ctx, dimension); err != nil {
return err
}
collectionName := q.getCollectionName(dimension)
pointID := uuid.New().String()
point := &qdrant.PointStruct{
Id: qdrant.NewID(pointID),
Vectors: qdrant.NewVectors(embeddingDB.Embedding...),
Payload: createPayload(embeddingDB),
}
_, err := q.client.Upsert(ctx, &qdrant.UpsertPoints{
CollectionName: collectionName,
Points: []*qdrant.PointStruct{point},
})
if err != nil {
log.Errorf("[Qdrant] Failed to save index: %v", err)
return fmt.Errorf("failed to save index for chunk ID %s: %w", embedding.ChunkID, err)
}
log.Infof("[Qdrant] Successfully saved index for chunk ID: %s, point ID: %s", embedding.ChunkID, pointID)
return nil
}
// BatchSave stores multiple points in Qdrant using batch upsert
func (q *qdrantRepository) BatchSave(ctx context.Context,
embeddingList []*types.IndexInfo, additionalParams map[string]any,
) error {
log := logger.GetLogger(ctx)
if len(embeddingList) == 0 {
log.Warn("[Qdrant] Empty list provided to BatchSave, skipping")
return nil
}
log.Infof("[Qdrant] Batch saving %d indices", len(embeddingList))
// Group points by dimension
pointsByDimension := make(map[int][]*qdrant.PointStruct)
for _, embedding := range embeddingList {
embeddingDB := toQdrantVectorEmbedding(embedding, additionalParams)
if len(embeddingDB.Embedding) == 0 {
log.Warnf("[Qdrant] Skipping empty embedding for chunk ID: %s", embedding.ChunkID)
continue
}
dimension := len(embeddingDB.Embedding)
point := &qdrant.PointStruct{
Id: qdrant.NewID(uuid.New().String()),
Vectors: qdrant.NewVectors(embeddingDB.Embedding...),
Payload: createPayload(embeddingDB),
}
pointsByDimension[dimension] = append(pointsByDimension[dimension], point)
log.Debugf("[Qdrant] Added chunk ID %s to batch request (dimension: %d)", embedding.ChunkID, dimension)
}
if len(pointsByDimension) == 0 {
log.Warn("[Qdrant] No valid points to save after filtering")
return nil
}
// Save points to each dimension-specific collection
totalSaved := 0
const batchSize = 100
for dimension, points := range pointsByDimension {
if err := q.ensureCollection(ctx, dimension); err != nil {
return err
}
collectionName := q.getCollectionName(dimension)
for i := 0; i < len(points); i += batchSize {
end := i + batchSize
if end > len(points) {
end = len(points)
}
batch := points[i:end]
_, err := q.client.Upsert(ctx, &qdrant.UpsertPoints{
CollectionName: collectionName,
Points: batch,
})
if err != nil {
return fmt.Errorf("failed to upsert batch: %w", err)
}
}
totalSaved += len(points)
log.Infof("[Qdrant] Saved %d points to collection %s", len(points), collectionName)
}
log.Infof("[Qdrant] Successfully batch saved %d indices", totalSaved)
return nil
}
// DeleteByChunkIDList removes points from the collection based on chunk IDs
func (q *qdrantRepository) DeleteByChunkIDList(ctx context.Context, chunkIDList []string, dimension int, knowledgeType string) error {
log := logger.GetLogger(ctx)
if len(chunkIDList) == 0 {
log.Warn("[Qdrant] Empty chunk ID list provided for deletion, skipping")
return nil
}
collectionName := q.getCollectionName(dimension)
log.Infof("[Qdrant] Deleting indices by chunk IDs from %s, count: %d", collectionName, len(chunkIDList))
_, err := q.client.Delete(ctx, &qdrant.DeletePoints{
CollectionName: collectionName,
Points: qdrant.NewPointsSelectorFilter(&qdrant.Filter{
Must: []*qdrant.Condition{
qdrant.NewMatchKeywords(fieldChunkID, chunkIDList...),
},
}),
})
if err != nil {
log.Errorf("[Qdrant] Failed to delete by chunk IDs: %v", err)
return fmt.Errorf("failed to delete by chunk IDs: %w", err)
}
log.Infof("[Qdrant] Successfully deleted documents by chunk IDs")
return nil
}
// DeleteByKnowledgeIDList removes points from the collection based on knowledge IDs
func (q *qdrantRepository) DeleteByKnowledgeIDList(ctx context.Context,
knowledgeIDList []string, dimension int, knowledgeType string,
) error {
log := logger.GetLogger(ctx)
if len(knowledgeIDList) == 0 {
log.Warn("[Qdrant] Empty knowledge ID list provided for deletion, skipping")
return nil
}
collectionName := q.getCollectionName(dimension)
log.Infof("[Qdrant] Deleting indices by knowledge IDs from %s, count: %d", collectionName, len(knowledgeIDList))
_, err := q.client.Delete(ctx, &qdrant.DeletePoints{
CollectionName: collectionName,
Points: qdrant.NewPointsSelectorFilter(&qdrant.Filter{
Must: []*qdrant.Condition{
qdrant.NewMatchKeywords(fieldKnowledgeID, knowledgeIDList...),
},
}),
})
if err != nil {
log.Errorf("[Qdrant] Failed to delete by knowledge IDs: %v", err)
return fmt.Errorf("failed to delete by knowledge IDs: %w", err)
}
log.Infof("[Qdrant] Successfully deleted documents by knowledge IDs")
return nil
}
// DeleteBySourceIDList removes points from the collection based on source IDs
func (q *qdrantRepository) DeleteBySourceIDList(ctx context.Context,
sourceIDList []string, dimension int, knowledgeType string,
) error {
log := logger.GetLogger(ctx)
if len(sourceIDList) == 0 {
log.Warn("[Qdrant] Empty source ID list provided for deletion, skipping")
return nil
}
collectionName := q.getCollectionName(dimension)
log.Infof("[Qdrant] Deleting indices by source IDs from %s, count: %d", collectionName, len(sourceIDList))
_, err := q.client.Delete(ctx, &qdrant.DeletePoints{
CollectionName: collectionName,
Points: qdrant.NewPointsSelectorFilter(&qdrant.Filter{
Must: []*qdrant.Condition{
qdrant.NewMatchKeywords(fieldSourceID, sourceIDList...),
},
}),
})
if err != nil {
log.Errorf("[Qdrant] Failed to delete by source IDs: %v", err)
return fmt.Errorf("failed to delete by source IDs: %w", err)
}
log.Infof("[Qdrant] Successfully deleted documents by source IDs")
return nil
}
// BatchUpdateChunkEnabledStatus updates the enabled status of chunks in batch
// This method operates on all collections since dimension is not provided
func (q *qdrantRepository) BatchUpdateChunkEnabledStatus(ctx context.Context, chunkStatusMap map[string]bool) error {
log := logger.GetLogger(ctx)
if len(chunkStatusMap) == 0 {
log.Warn("[Qdrant] Empty chunk status map provided, skipping")
return nil
}
log.Infof("[Qdrant] Batch updating chunk enabled status, count: %d", len(chunkStatusMap))
// Get all collections that match our base name pattern
collections, err := q.client.ListCollections(ctx)
if err != nil {
log.Errorf("[Qdrant] Failed to list collections: %v", err)
return fmt.Errorf("failed to list collections: %w", err)
}
// Group chunks by enabled status for batch updates
enabledChunkIDs := make([]string, 0)
disabledChunkIDs := make([]string, 0)
for chunkID, enabled := range chunkStatusMap {
if enabled {
enabledChunkIDs = append(enabledChunkIDs, chunkID)
} else {
disabledChunkIDs = append(disabledChunkIDs, chunkID)
}
}
// Update in all matching collections
for _, collectionName := range collections {
// Only process collections that start with our base name
if len(collectionName) <= len(q.collectionBaseName) ||
collectionName[:len(q.collectionBaseName)] != q.collectionBaseName {
continue
}
// Update enabled chunks
if len(enabledChunkIDs) > 0 {
_, err := q.client.SetPayload(ctx, &qdrant.SetPayloadPoints{
CollectionName: collectionName,
Payload: qdrant.NewValueMap(map[string]any{fieldIsEnabled: true}),
PointsSelector: qdrant.NewPointsSelectorFilter(&qdrant.Filter{
Must: []*qdrant.Condition{
qdrant.NewMatchKeywords(fieldChunkID, enabledChunkIDs...),
},
}),
})
if err != nil {
log.Warnf("[Qdrant] Failed to update enabled chunks in %s: %v", collectionName, err)
}
}
// Update disabled chunks
if len(disabledChunkIDs) > 0 {
_, err := q.client.SetPayload(ctx, &qdrant.SetPayloadPoints{
CollectionName: collectionName,
Payload: qdrant.NewValueMap(map[string]any{fieldIsEnabled: false}),
PointsSelector: qdrant.NewPointsSelectorFilter(&qdrant.Filter{
Must: []*qdrant.Condition{
qdrant.NewMatchKeywords(fieldChunkID, disabledChunkIDs...),
},
}),
})
if err != nil {
log.Warnf("[Qdrant] Failed to update disabled chunks in %s: %v", collectionName, err)
}
}
}
log.Infof("[Qdrant] Batch update chunk enabled status completed")
return nil
}
// BatchUpdateChunkTagID updates the tag ID of chunks in batch
func (q *qdrantRepository) BatchUpdateChunkTagID(ctx context.Context, chunkTagMap map[string]string) error {
log := logger.GetLogger(ctx)
if len(chunkTagMap) == 0 {
log.Warn("[Qdrant] Empty chunk tag map provided, skipping")
return nil
}
log.Infof("[Qdrant] Batch updating chunk tag ID, count: %d", len(chunkTagMap))
// Get all collections that match our base name pattern
collections, err := q.client.ListCollections(ctx)
if err != nil {
log.Errorf("[Qdrant] Failed to list collections: %v", err)
return fmt.Errorf("failed to list collections: %w", err)
}
// Group chunks by tag ID for batch updates
tagGroups := make(map[string][]string)
for chunkID, tagID := range chunkTagMap {
tagGroups[tagID] = append(tagGroups[tagID], chunkID)
}
// Update in all matching collections
for _, collectionName := range collections {
// Only process collections that start with our base name
if len(collectionName) <= len(q.collectionBaseName) ||
collectionName[:len(q.collectionBaseName)] != q.collectionBaseName {
continue
}
// Update chunks for each tag ID
for tagID, chunkIDs := range tagGroups {
_, err := q.client.SetPayload(ctx, &qdrant.SetPayloadPoints{
CollectionName: collectionName,
Payload: qdrant.NewValueMap(map[string]any{fieldTagID: tagID}),
PointsSelector: qdrant.NewPointsSelectorFilter(&qdrant.Filter{
Must: []*qdrant.Condition{
qdrant.NewMatchKeywords(fieldChunkID, chunkIDs...),
},
}),
})
if err != nil {
log.Warnf("[Qdrant] Failed to update chunks with tag_id %s in %s: %v", tagID, collectionName, err)
}
}
}
log.Infof("[Qdrant] Batch update chunk tag ID completed")
return nil
}
func (q *qdrantRepository) getBaseFilter(params types.RetrieveParams) *qdrant.Filter {
must := make([]*qdrant.Condition, 0)
mustNot := make([]*qdrant.Condition, 0)
// Only retrieve enabled chunks
must = append(must, qdrant.NewMatchBool(fieldIsEnabled, true))
// KnowledgeBaseIDs and KnowledgeIDs use AND logic
// - If only KnowledgeBaseIDs: search entire knowledge bases
// - If only KnowledgeIDs: search specific documents
// - If both: search specific documents within the knowledge bases (AND)
if len(params.KnowledgeBaseIDs) > 0 {
must = append(must, qdrant.NewMatchKeywords(fieldKnowledgeBaseID, params.KnowledgeBaseIDs...))
}
if len(params.KnowledgeIDs) > 0 {
must = append(must, qdrant.NewMatchKeywords(fieldKnowledgeID, params.KnowledgeIDs...))
}
// Filter by tag IDs if specified
if len(params.TagIDs) > 0 {
must = append(must, qdrant.NewMatchKeywords(fieldTagID, params.TagIDs...))
}
if len(params.ExcludeKnowledgeIDs) > 0 {
mustNot = append(mustNot, qdrant.NewMatchKeywords(fieldKnowledgeID, params.ExcludeKnowledgeIDs...))
}
if len(params.ExcludeChunkIDs) > 0 {
mustNot = append(mustNot, qdrant.NewMatchKeywords(fieldChunkID, params.ExcludeChunkIDs...))
}
filter := &qdrant.Filter{
Must: must,
MustNot: mustNot,
}
return filter
}
// Retrieve dispatches the retrieval operation to the appropriate method based on retriever type
func (q *qdrantRepository) Retrieve(ctx context.Context,
params types.RetrieveParams,
) ([]*types.RetrieveResult, error) {
log := logger.GetLogger(ctx)
log.Debugf("[Qdrant] Processing retrieval request of type: %s", params.RetrieverType)
switch params.RetrieverType {
case types.VectorRetrieverType:
return q.VectorRetrieve(ctx, params)
case types.KeywordsRetrieverType:
return q.KeywordsRetrieve(ctx, params)
}
err := fmt.Errorf("invalid retriever type: %v", params.RetrieverType)
log.Errorf("[Qdrant] %v", err)
return nil, err
}
// VectorRetrieve performs vector similarity search
func (q *qdrantRepository) VectorRetrieve(ctx context.Context,
params types.RetrieveParams,
) ([]*types.RetrieveResult, error) {
log := logger.GetLogger(ctx)
dimension := len(params.Embedding)
log.Infof("[Qdrant] Vector retrieval: dim=%d, topK=%d, threshold=%.4f",
dimension, params.TopK, params.Threshold)
// Get collection name based on embedding dimension
collectionName := q.getCollectionName(dimension)
// Check if collection exists
exists, err := q.client.CollectionExists(ctx, collectionName)
if err != nil {
log.Errorf("[Qdrant] Failed to check collection existence: %v", err)
return nil, fmt.Errorf("failed to check collection: %w", err)
}
if !exists {
log.Warnf("[Qdrant] Collection %s does not exist, returning empty results", collectionName)
return buildRetrieveResult(nil, types.VectorRetrieverType), nil
}
filter := q.getBaseFilter(params)
limit := uint64(params.TopK)
scoreThreshold := float32(params.Threshold)
searchResult, err := q.client.Query(ctx, &qdrant.QueryPoints{
CollectionName: collectionName,
Query: qdrant.NewQuery(params.Embedding...),
Filter: filter,
Limit: &limit,
ScoreThreshold: &scoreThreshold,
WithPayload: qdrant.NewWithPayload(true),
})
if err != nil {
log.Errorf("[Qdrant] Vector search failed: %v", err)
return nil, fmt.Errorf("%s: %w", collectionName, err)
}
var results []*types.IndexWithScore
for _, point := range searchResult {
payload := point.Payload
embedding := &QdrantVectorEmbeddingWithScore{
QdrantVectorEmbedding: QdrantVectorEmbedding{
Content: payload[fieldContent].GetStringValue(),
SourceID: payload[fieldSourceID].GetStringValue(),
SourceType: int(payload[fieldSourceType].GetIntegerValue()),
ChunkID: payload[fieldChunkID].GetStringValue(),
KnowledgeID: payload[fieldKnowledgeID].GetStringValue(),
KnowledgeBaseID: payload[fieldKnowledgeBaseID].GetStringValue(),
TagID: payload[fieldTagID].GetStringValue(),
},
Score: float64(point.Score),
}
pointID := point.Id.GetUuid()
results = append(results, fromQdrantVectorEmbedding(pointID, embedding, types.MatchTypeEmbedding))
}
if len(results) == 0 {
log.Warnf("[Qdrant] No vector matches found that meet threshold %.4f", params.Threshold)
} else {
log.Infof("[Qdrant] Vector retrieval found %d results", len(results))
log.Debugf("[Qdrant] Top result score: %.4f", results[0].Score)
}
return buildRetrieveResult(results, types.VectorRetrieverType), nil
}
// KeywordsRetrieve performs keyword-based search in document content
// This searches across all collections since keyword search doesn't depend on dimension
func (q *qdrantRepository) KeywordsRetrieve(ctx context.Context,
params types.RetrieveParams,
) ([]*types.RetrieveResult, error) {
log := logger.GetLogger(ctx)
log.Infof("[Qdrant] Performing keywords retrieval with query: %s, topK: %d", params.Query, params.TopK)
// Get all collections that match our base name pattern
collections, err := q.client.ListCollections(ctx)
if err != nil {
log.Errorf("[Qdrant] Failed to list collections: %v", err)
return nil, fmt.Errorf("failed to list collections: %w", err)
}
var allResults []*types.IndexWithScore
limit := uint32(params.TopK)
log.Debugf("[Qdrant] Found %d collections, base name: %s", len(collections), q.collectionBaseName)
// Tokenize query for OR-based search (better for Chinese and multi-word queries)
queryTokens := tokenizeQuery(params.Query)
log.Debugf("[Qdrant] Tokenized query into %d tokens: %v", len(queryTokens), queryTokens)
// Search in all matching collections
for _, collectionName := range collections {
log.Debugf("[Qdrant] Checking collection: %s", collectionName)
// Only process collections that start with our base name
if len(collectionName) <= len(q.collectionBaseName) ||
collectionName[:len(q.collectionBaseName)] != q.collectionBaseName {
log.Debugf("[Qdrant] Skipping collection %s (doesn't match base name %s)", collectionName, q.collectionBaseName)
continue
}
filter := q.getBaseFilter(params)
// Build should conditions for each token (OR logic)
// This allows matching documents that contain any of the query tokens
if len(queryTokens) > 0 {
shouldConditions := make([]*qdrant.Condition, 0, len(queryTokens))
for _, token := range queryTokens {
shouldConditions = append(shouldConditions, qdrant.NewMatchText(fieldContent, token))
}
filter.Should = shouldConditions
} else {
// Fallback to original query if tokenization fails
filter.Must = append(filter.Must, qdrant.NewMatchText(fieldContent, params.Query))
}
log.Debugf("[Qdrant] Searching in collection %s with %d should conditions", collectionName, len(filter.Should))
scrollResult, err := q.client.Scroll(ctx, &qdrant.ScrollPoints{
CollectionName: collectionName,
Filter: filter,
Limit: &limit,
WithPayload: qdrant.NewWithPayload(true),
})
if err != nil {
log.Warnf("[Qdrant] Keywords search failed in %s: %v", collectionName, err)
continue
}
log.Debugf("[Qdrant] Found %d results in collection %s", len(scrollResult), collectionName)
for _, point := range scrollResult {
payload := point.Payload
embedding := &QdrantVectorEmbeddingWithScore{
QdrantVectorEmbedding: QdrantVectorEmbedding{
Content: payload[fieldContent].GetStringValue(),
SourceID: payload[fieldSourceID].GetStringValue(),
SourceType: int(payload[fieldSourceType].GetIntegerValue()),
ChunkID: payload[fieldChunkID].GetStringValue(),
KnowledgeID: payload[fieldKnowledgeID].GetStringValue(),
KnowledgeBaseID: payload[fieldKnowledgeBaseID].GetStringValue(),
TagID: payload[fieldTagID].GetStringValue(),
},
Score: 1.0,
}
pointID := point.Id.GetUuid()
allResults = append(allResults, fromQdrantVectorEmbedding(pointID, embedding, types.MatchTypeKeywords))
}
}
// Limit results to topK
if len(allResults) > params.TopK {
allResults = allResults[:params.TopK]
}
if len(allResults) == 0 {
log.Warnf("[Qdrant] No keyword matches found for query: %s", params.Query)
} else {
log.Infof("[Qdrant] Keywords retrieval found %d results", len(allResults))
}
return buildRetrieveResult(allResults, types.KeywordsRetrieverType), nil
}
// CopyIndices copies index data from source knowledge base to target knowledge base
func (q *qdrantRepository) CopyIndices(ctx context.Context,
sourceKnowledgeBaseID string,
sourceToTargetKBIDMap map[string]string,
sourceToTargetChunkIDMap map[string]string,
targetKnowledgeBaseID string,
dimension int,
knowledgeType string,
) error {
log := logger.GetLogger(ctx)
log.Infof(
"[Qdrant] Copying indices from source knowledge base %s to target knowledge base %s, count: %d, dimension: %d",
sourceKnowledgeBaseID, targetKnowledgeBaseID, len(sourceToTargetChunkIDMap), dimension,
)
if len(sourceToTargetChunkIDMap) == 0 {
log.Warn("[Qdrant] Empty mapping, skipping copy")
return nil
}
collectionName := q.getCollectionName(dimension)
// Ensure target collection exists
if err := q.ensureCollection(ctx, dimension); err != nil {
return err
}
batchSize := uint32(64)
var offset *qdrant.PointId = nil
totalCopied := 0
for {
scrollResult, err := q.client.Scroll(ctx, &qdrant.ScrollPoints{
CollectionName: collectionName,
Filter: &qdrant.Filter{
Must: []*qdrant.Condition{
qdrant.NewMatch(fieldKnowledgeBaseID, sourceKnowledgeBaseID),
},
},
Limit: &batchSize,
Offset: offset,
WithPayload: qdrant.NewWithPayload(true),
WithVectors: qdrant.NewWithVectors(true),
})
if err != nil {
log.Errorf("[Qdrant] Failed to query source points: %v", err)
return err
}
pointsCount := len(scrollResult)
if pointsCount == 0 {
break
}
log.Infof("[Qdrant] Found %d source points in batch", pointsCount)
targetPoints := make([]*qdrant.PointStruct, 0, pointsCount)
for _, sourcePoint := range scrollResult {
payload := sourcePoint.Payload
sourceChunkID := payload[fieldChunkID].GetStringValue()
sourceKnowledgeID := payload[fieldKnowledgeID].GetStringValue()
originalSourceID := payload[fieldSourceID].GetStringValue()
targetChunkID, ok := sourceToTargetChunkIDMap[sourceChunkID]
if !ok {
log.Warnf("[Qdrant] Source chunk %s not found in target mapping, skipping", sourceChunkID)
continue
}
targetKnowledgeID, ok := sourceToTargetKBIDMap[sourceKnowledgeID]
if !ok {
log.Warnf("[Qdrant] Source knowledge %s not found in target mapping, skipping", sourceKnowledgeID)
continue
}
// Handle SourceID transformation for generated questions
// Generated questions have SourceID format: {chunkID}-{questionID}
// Regular chunks have SourceID == ChunkID
var targetSourceID string
if originalSourceID == sourceChunkID {
// Regular chunk, use targetChunkID as SourceID
targetSourceID = targetChunkID
} else if strings.HasPrefix(originalSourceID, sourceChunkID+"-") {
// This is a generated question, preserve the questionID part
questionID := strings.TrimPrefix(originalSourceID, sourceChunkID+"-")
targetSourceID = fmt.Sprintf("%s-%s", targetChunkID, questionID)
} else {
// For other complex scenarios, generate new unique SourceID
targetSourceID = uuid.New().String()
}
isEnabled := true
if v, ok := payload[fieldIsEnabled]; ok {
isEnabled = v.GetBoolValue()
}
newPayload := qdrant.NewValueMap(map[string]any{
fieldContent: payload[fieldContent].GetStringValue(),
fieldSourceID: targetSourceID,
fieldSourceType: payload[fieldSourceType].GetIntegerValue(),
fieldChunkID: targetChunkID,
fieldKnowledgeID: targetKnowledgeID,
fieldKnowledgeBaseID: targetKnowledgeBaseID,
fieldTagID: payload[fieldTagID].GetStringValue(),
fieldIsEnabled: isEnabled,
})
var vectors *qdrant.Vectors
if vectorOutput := sourcePoint.Vectors.GetVector(); vectorOutput != nil {
if denseVector := vectorOutput.GetDenseVector(); denseVector != nil {
vectors = qdrant.NewVectors(denseVector.Data...)
}
}
if vectors == nil {
log.Warnf("[Qdrant] No vectors found for source point with chunk %s, skipping", sourceChunkID)
continue
}
newPoint := &qdrant.PointStruct{
Id: qdrant.NewID(uuid.New().String()),
Vectors: vectors,
Payload: newPayload,
}
targetPoints = append(targetPoints, newPoint)
}
if len(targetPoints) > 0 {
_, err := q.client.Upsert(ctx, &qdrant.UpsertPoints{
CollectionName: collectionName,
Points: targetPoints,
})
if err != nil {
log.Errorf("[Qdrant] Failed to batch upsert target points: %v", err)
return fmt.Errorf("failed to batch upsert target points during copy: %w", err)
}
totalCopied += len(targetPoints)
log.Infof("[Qdrant] Successfully copied batch, batch size: %d, total copied: %d",
len(targetPoints), totalCopied)
}
if pointsCount > 0 {
offset = scrollResult[pointsCount-1].Id
}
if pointsCount < int(batchSize) {
break
}
}
log.Infof("[Qdrant] Index copy completed, total copied: %d", totalCopied)
return nil
}
func createPayload(embedding *QdrantVectorEmbedding) map[string]*qdrant.Value {
payload := map[string]any{
fieldContent: embedding.Content,
fieldSourceID: embedding.SourceID,
fieldSourceType: int64(embedding.SourceType),
fieldChunkID: embedding.ChunkID,
fieldKnowledgeID: embedding.KnowledgeID,
fieldKnowledgeBaseID: embedding.KnowledgeBaseID,
fieldTagID: embedding.TagID,
fieldIsEnabled: embedding.IsEnabled,
}
return qdrant.NewValueMap(payload)
}
func buildRetrieveResult(results []*types.IndexWithScore, retrieverType types.RetrieverType) []*types.RetrieveResult {
return []*types.RetrieveResult{
{
Results: results,
RetrieverEngineType: types.QdrantRetrieverEngineType,
RetrieverType: retrieverType,
Error: nil,
},
}
}
// Ref: https://github.com/qdrant/qdrant-sizing-calculator
func (q *qdrantRepository) calculateStorageSize(embedding *QdrantVectorEmbedding) int64 {
// Payload fields
payloadSizeBytes := int64(0)
payloadSizeBytes += int64(len(embedding.Content)) // content string
payloadSizeBytes += int64(len(embedding.SourceID)) // source_id string
payloadSizeBytes += int64(len(embedding.ChunkID)) // chunk_id string
payloadSizeBytes += int64(len(embedding.KnowledgeID)) // knowledge_id string
payloadSizeBytes += int64(len(embedding.KnowledgeBaseID)) // knowledge_base_id string
payloadSizeBytes += 8 // source_type int64
// Vector storage and index
var vectorSizeBytes int64 = 0
var hnswIndexBytes int64 = 0
if embedding.Embedding != nil {
dimensions := int64(len(embedding.Embedding))
vectorSizeBytes = dimensions * 4
// HNSW graph links per vector: M×2 neighbors in layer 0, ~8 bytes per link
// (4 bytes PointOffsetType + multi-layer amortization).
// Graph link count depends on M, NOT on vector dimensions.
const hnswM = 16
hnswIndexBytes = hnswM * 2 * 8
}
// ID tracker metadata: 24 bytes per vector
// (forward refs + backward refs + version tracking = 8 + 8 + 8)
const idTrackerBytes int64 = 24
totalSizeBytes := payloadSizeBytes + vectorSizeBytes + hnswIndexBytes + idTrackerBytes
return totalSizeBytes
}
// toQdrantVectorEmbedding converts IndexInfo to Qdrant payload format
func toQdrantVectorEmbedding(embedding *types.IndexInfo, additionalParams map[string]interface{}) *QdrantVectorEmbedding {
vector := &QdrantVectorEmbedding{
Content: embedding.Content,
SourceID: embedding.SourceID,
SourceType: int(embedding.SourceType),
ChunkID: embedding.ChunkID,
KnowledgeID: embedding.KnowledgeID,
KnowledgeBaseID: embedding.KnowledgeBaseID,
TagID: embedding.TagID,
IsEnabled: embedding.IsEnabled,
}
if additionalParams != nil {
if val, exists := additionalParams[fieldEmbedding]; exists {
if embeddingMap, ok := val.(map[string][]float32); ok {
vector.Embedding = embeddingMap[embedding.SourceID]
}
}
}
return vector
}
// fromQdrantVectorEmbedding converts Qdrant point to IndexWithScore domain model
func fromQdrantVectorEmbedding(id string,
embedding *QdrantVectorEmbeddingWithScore,
matchType types.MatchType,
) *types.IndexWithScore {
return &types.IndexWithScore{
ID: id,
SourceID: embedding.SourceID,
SourceType: types.SourceType(embedding.SourceType),
ChunkID: embedding.ChunkID,
KnowledgeID: embedding.KnowledgeID,
KnowledgeBaseID: embedding.KnowledgeBaseID,
TagID: embedding.TagID,
Content: embedding.Content,
Score: embedding.Score,
MatchType: matchType,
}
}
// tokenizeQuery splits a query string into tokens for OR-based full-text search.
// It uses jieba for professional Chinese word segmentation.
func tokenizeQuery(query string) []string {
query = strings.TrimSpace(query)
if query == "" {
return nil
}
// Use jieba for segmentation (search mode for better recall)
words := types.Jieba.CutForSearch(query, true)
// Filter and deduplicate
seen := make(map[string]bool)
result := make([]string, 0, len(words))
for _, word := range words {
word = strings.TrimSpace(strings.ToLower(word))
// Skip empty, single-char, and already seen words
if utf8.RuneCountInString(word) < 2 || seen[word] {
continue
}
seen[word] = true
result = append(result, word)
}
return result
}
@@ -0,0 +1,33 @@
package qdrant
import (
"sync"
"github.com/qdrant/go-client/qdrant"
)
type qdrantRepository struct {
client *qdrant.Client
collectionBaseName string
shardNumber int // 0 = use Qdrant server default
replicationFactor int // 0 = use Qdrant server default
// Cache for initialized collections (dimension -> true)
initializedCollections sync.Map
}
type QdrantVectorEmbedding struct {
Content string `json:"content"`
SourceID string `json:"source_id"`
SourceType int `json:"source_type"`
ChunkID string `json:"chunk_id"`
KnowledgeID string `json:"knowledge_id"`
KnowledgeBaseID string `json:"knowledge_base_id"`
TagID string `json:"tag_id"`
Embedding []float32 `json:"embedding"`
IsEnabled bool `json:"is_enabled"`
}
type QdrantVectorEmbeddingWithScore struct {
QdrantVectorEmbedding
Score float64
}
@@ -0,0 +1,661 @@
package sqlite
import (
"context"
"fmt"
"strings"
"time"
"unicode"
"github.com/Tencent/WeKnora/internal/common"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
sqlite_vec "github.com/asg017/sqlite-vec-go-bindings/cgo"
"github.com/google/uuid"
"gorm.io/gorm"
"gorm.io/gorm/clause"
)
// sqliteEmbedding stores metadata alongside the vec0 virtual table rows
type sqliteEmbedding struct {
ID uint `gorm:"primarykey;autoIncrement"`
CreatedAt time.Time `gorm:"column:created_at"`
UpdatedAt time.Time `gorm:"column:updated_at"`
SourceID string `gorm:"column:source_id;not null;uniqueIndex:idx_sqlite_emb_source"`
SourceType int `gorm:"column:source_type;not null;uniqueIndex:idx_sqlite_emb_source"`
ChunkID string `gorm:"column:chunk_id;index"`
KnowledgeID string `gorm:"column:knowledge_id;index"`
KnowledgeBaseID string `gorm:"column:knowledge_base_id;index"`
TagID string `gorm:"column:tag_id;index"`
Content string `gorm:"column:content;not null"`
Dimension int `gorm:"column:dimension;not null"`
IsEnabled *bool `gorm:"column:is_enabled;default:true;index"`
}
func (sqliteEmbedding) TableName() string { return "lite_embeddings" }
type sqliteRepository struct {
db *gorm.DB
vecTables map[int]bool // tracks which vec0 tables have been created (keyed by dimension)
}
func NewSQLiteRetrieveEngineRepository(db *gorm.DB) interfaces.RetrieveEngineRepository {
logger.GetLogger(context.Background()).Info("[SQLite] Initializing SQLite retriever engine repository with sqlite-vec")
if err := db.AutoMigrate(&sqliteEmbedding{}); err != nil {
logger.GetLogger(context.Background()).Errorf("[SQLite] Failed to auto-migrate lite_embeddings: %v", err)
}
initFTS5(db)
repo := &sqliteRepository{
db: db,
vecTables: make(map[int]bool),
}
repo.ensureExistingVecTables()
return repo
}
func initFTS5(db *gorm.DB) {
var sqlStr string
db.Raw("SELECT sql FROM sqlite_master WHERE type='table' AND name='lite_embeddings_fts'").Scan(&sqlStr)
// If the table exists but is not contentless (i.e. uses content='lite_embeddings'), recreate it
if sqlStr != "" && strings.Contains(sqlStr, "content='lite_embeddings'") {
logger.GetLogger(context.Background()).Infof("[SQLite] Migrating FTS5 table to contentless table with manual bigram tokenization")
db.Exec("DROP TABLE IF EXISTS lite_embeddings_fts")
sqlStr = "" // Trigger recreate
}
if sqlStr == "" {
createSQL := `CREATE VIRTUAL TABLE IF NOT EXISTS lite_embeddings_fts USING fts5(
content, source_id, chunk_id, knowledge_id, knowledge_base_id,
content='',
contentless_delete=1,
tokenize='unicode61'
)`
if err := db.Exec(createSQL).Error; err != nil {
logger.GetLogger(context.Background()).Warnf("[SQLite] Failed to create FTS5 table: %v", err)
return
}
logger.GetLogger(context.Background()).Infof("[SQLite] Populating contentless FTS5 table from lite_embeddings with bigrams")
// To populate, we need to read all rows and insert them via Go to apply bigram tokenization
type Row struct {
ID uint
Content string
SourceID string
ChunkID string
KnowledgeID string
KnowledgeBaseID string
}
var rows []Row
db.Raw("SELECT id, content, source_id, chunk_id, knowledge_id, knowledge_base_id FROM lite_embeddings").Scan(&rows)
for _, r := range rows {
db.Exec(`INSERT INTO lite_embeddings_fts(rowid, content, source_id, chunk_id, knowledge_id, knowledge_base_id)
VALUES(?, ?, ?, ?, ?, ?)`,
r.ID, tokenizeCJKBigram(r.Content), r.SourceID, r.ChunkID, r.KnowledgeID, r.KnowledgeBaseID)
}
}
}
func vecTableName(dim int) string {
return fmt.Sprintf("vec_embeddings_%d", dim)
}
func (r *sqliteRepository) ensureVecTable(dim int) {
if dim <= 0 || r.vecTables[dim] {
return
}
tbl := vecTableName(dim)
createSQL := fmt.Sprintf(
`CREATE VIRTUAL TABLE IF NOT EXISTS %s USING vec0(embedding float[%d] distance_metric=cosine)`,
tbl, dim,
)
if err := r.db.Exec(createSQL).Error; err != nil {
if strings.Contains(err.Error(), "already exists") {
r.vecTables[dim] = true
return
}
logger.GetLogger(context.Background()).Errorf("[SQLite] Failed to create vec0 table for dim %d: %v", dim, err)
return
}
r.vecTables[dim] = true
}
func (r *sqliteRepository) ensureExistingVecTables() {
var dims []int
r.db.Model(&sqliteEmbedding{}).Distinct("dimension").Where("dimension > 0").Pluck("dimension", &dims)
for _, dim := range dims {
r.ensureVecTable(dim)
}
}
func (r *sqliteRepository) EngineType() types.RetrieverEngineType {
return types.SQLiteRetrieverEngineType
}
func (r *sqliteRepository) Support() []types.RetrieverType {
return []types.RetrieverType{types.KeywordsRetrieverType, types.VectorRetrieverType}
}
func (r *sqliteRepository) Save(ctx context.Context, indexInfo *types.IndexInfo, params map[string]any) error {
row := toSQLiteEmbedding(indexInfo)
emb := extractEmbedding(params, indexInfo.SourceID)
if len(emb) > 0 {
row.Dimension = len(emb)
}
if err := r.db.WithContext(ctx).Clauses(clause.OnConflict{DoNothing: true}).Create(row).Error; err != nil {
return err
}
r.syncFTS5Insert(ctx, row)
if len(emb) > 0 && row.ID > 0 {
r.insertVec(ctx, row.ID, row.Dimension, emb)
}
return nil
}
func (r *sqliteRepository) BatchSave(ctx context.Context, indexInfoList []*types.IndexInfo, params map[string]any) error {
if len(indexInfoList) == 0 {
return nil
}
rows := make([]*sqliteEmbedding, len(indexInfoList))
embs := make([][]float32, len(indexInfoList))
for i, info := range indexInfoList {
rows[i] = toSQLiteEmbedding(info)
emb := extractEmbedding(params, info.SourceID)
embs[i] = emb
if len(emb) > 0 {
rows[i].Dimension = len(emb)
}
}
if err := r.db.WithContext(ctx).Clauses(clause.OnConflict{DoNothing: true}).Create(rows).Error; err != nil {
return err
}
for i, row := range rows {
r.syncFTS5Insert(ctx, row)
if len(embs[i]) > 0 && row.ID > 0 {
r.insertVec(ctx, row.ID, row.Dimension, embs[i])
}
}
return nil
}
func (r *sqliteRepository) EstimateStorageSize(_ context.Context, indexInfoList []*types.IndexInfo, _ map[string]any) int64 {
var total int64
for _, info := range indexInfoList {
total += int64(len(info.Content)) + 200
}
return total
}
func (r *sqliteRepository) DeleteByChunkIDList(ctx context.Context, chunkIDList []string, _ int, _ string) error {
var rows []sqliteEmbedding
r.db.WithContext(ctx).Where("chunk_id IN ?", chunkIDList).Find(&rows)
r.deleteRowsAndVecs(ctx, rows)
return r.db.WithContext(ctx).Where("chunk_id IN ?", chunkIDList).Delete(&sqliteEmbedding{}).Error
}
func (r *sqliteRepository) DeleteBySourceIDList(ctx context.Context, sourceIDList []string, _ int, _ string) error {
var rows []sqliteEmbedding
r.db.WithContext(ctx).Where("source_id IN ?", sourceIDList).Find(&rows)
r.deleteRowsAndVecs(ctx, rows)
return r.db.WithContext(ctx).Where("source_id IN ?", sourceIDList).Delete(&sqliteEmbedding{}).Error
}
func (r *sqliteRepository) DeleteByKnowledgeIDList(ctx context.Context, knowledgeIDList []string, _ int, _ string) error {
var rows []sqliteEmbedding
r.db.WithContext(ctx).Where("knowledge_id IN ?", knowledgeIDList).Find(&rows)
r.deleteRowsAndVecs(ctx, rows)
return r.db.WithContext(ctx).Where("knowledge_id IN ?", knowledgeIDList).Delete(&sqliteEmbedding{}).Error
}
func (r *sqliteRepository) CopyIndices(ctx context.Context,
_ string,
sourceToTargetKBIDMap map[string]string,
sourceToTargetChunkIDMap map[string]string,
targetKnowledgeBaseID string,
_ int, _ string,
) error {
for sourceChunkID, targetChunkID := range sourceToTargetChunkIDMap {
var src sqliteEmbedding
if err := r.db.WithContext(ctx).Where("chunk_id = ?", sourceChunkID).First(&src).Error; err != nil {
continue
}
newRow := sqliteEmbedding{
SourceID: uuid.New().String(),
SourceType: src.SourceType,
ChunkID: targetChunkID,
KnowledgeID: sourceToTargetKBIDMap[src.KnowledgeID],
KnowledgeBaseID: targetKnowledgeBaseID,
TagID: src.TagID,
Content: src.Content,
Dimension: src.Dimension,
IsEnabled: src.IsEnabled,
}
if err := r.db.WithContext(ctx).Create(&newRow).Error; err != nil {
logger.GetLogger(ctx).Warnf("[SQLite] CopyIndices: failed to copy chunk %s: %v", sourceChunkID, err)
continue
}
r.syncFTS5Insert(ctx, &newRow)
if src.Dimension > 0 && newRow.ID > 0 {
r.copyVec(ctx, src.ID, newRow.ID, src.Dimension)
}
}
return nil
}
func (r *sqliteRepository) BatchUpdateChunkEnabledStatus(ctx context.Context, chunkStatusMap map[string]bool) error {
for chunkID, enabled := range chunkStatusMap {
r.db.WithContext(ctx).Model(&sqliteEmbedding{}).Where("chunk_id = ?", chunkID).Update("is_enabled", enabled)
}
return nil
}
func (r *sqliteRepository) BatchUpdateChunkTagID(ctx context.Context, chunkTagMap map[string]string) error {
for chunkID, tagID := range chunkTagMap {
r.db.WithContext(ctx).Model(&sqliteEmbedding{}).Where("chunk_id = ?", chunkID).Update("tag_id", tagID)
}
return nil
}
// --- Retrieve ---
func (r *sqliteRepository) Retrieve(ctx context.Context, params types.RetrieveParams) ([]*types.RetrieveResult, error) {
var results []*types.RetrieveResult
if params.RetrieverType == types.KeywordsRetrieverType || params.RetrieverType == "" {
res, err := r.keywordsRetrieve(ctx, params)
if err != nil {
return nil, err
}
results = append(results, res...)
}
if params.RetrieverType == types.VectorRetrieverType || params.RetrieverType == "" {
res, err := r.vectorRetrieve(ctx, params)
if err != nil {
return nil, err
}
results = append(results, res...)
}
return results, nil
}
// --- Keywords retrieval via FTS5 ---
func (r *sqliteRepository) keywordsRetrieve(ctx context.Context, params types.RetrieveParams) ([]*types.RetrieveResult, error) {
if params.Query == "" {
return nil, nil
}
ftsQuery := sanitizeFTS5Query(params.Query)
sql := `
SELECT e.id, e.source_id, e.source_type, e.chunk_id,
e.knowledge_id, e.knowledge_base_id, e.tag_id,
e.content,
(bm25(lite_embeddings_fts) * -1000000.0) AS score
FROM lite_embeddings_fts
JOIN lite_embeddings e ON e.id = lite_embeddings_fts.rowid
WHERE lite_embeddings_fts MATCH ?
AND (e.is_enabled IS NULL OR e.is_enabled = 1)
`
args := []interface{}{ftsQuery}
for _, wp := range buildFilterWhere(params, "e") {
sql += " AND " + wp.clause
args = append(args, wp.args...)
}
sql += " ORDER BY score DESC LIMIT ?"
args = append(args, params.TopK)
type ftsResult struct {
ID uint
SourceID string
SourceType int
ChunkID string
KnowledgeID string
KnowledgeBaseID string
TagID string
Content string
Score float64
}
var rows []ftsResult
if err := r.db.WithContext(ctx).Raw(sql, args...).Scan(&rows).Error; err != nil {
return nil, fmt.Errorf("FTS5 query failed: %w", err)
}
logger.GetLogger(ctx).Infof("[SQLite] keywordsRetrieve: query=%q, ftsQuery=%q, matched=%d rows", params.Query, ftsQuery, len(rows))
items := make([]*types.IndexWithScore, len(rows))
for i, row := range rows {
// bm25 is originally negative and very small, we multiplied it by -1000000.0 in SQL
// to make it positive and human-readable.
score := row.Score
logger.GetLogger(ctx).Infof("[SQLite] keywordsRetrieve: #%d chunk_id=%s, score=%.4f, content_preview=%.60s",
i+1, row.ChunkID, score, row.Content)
items[i] = &types.IndexWithScore{
ID: fmt.Sprintf("%d", row.ID),
SourceID: row.SourceID,
SourceType: types.SourceType(row.SourceType),
ChunkID: row.ChunkID,
KnowledgeID: row.KnowledgeID,
KnowledgeBaseID: row.KnowledgeBaseID,
TagID: row.TagID,
Content: row.Content,
Score: score,
MatchType: types.MatchTypeKeywords,
}
}
return []*types.RetrieveResult{{
Results: items,
RetrieverEngineType: types.SQLiteRetrieverEngineType,
RetrieverType: types.KeywordsRetrieverType,
}}, nil
}
func (r *sqliteRepository) vectorRetrieve(ctx context.Context, params types.RetrieveParams) ([]*types.RetrieveResult, error) {
if len(params.Embedding) == 0 {
return nil, nil
}
dim := len(params.Embedding)
r.ensureVecTable(dim)
queryBlob, err := sqlite_vec.SerializeFloat32(params.Embedding)
if err != nil {
return nil, fmt.Errorf("serialize query vector failed: %w", err)
}
tbl := vecTableName(dim)
// ⚠️ sqlite-vec 要求必须有 k = ?
vecSQL := fmt.Sprintf(`
SELECT v.rowid, v.distance,
e.source_id, e.source_type, e.chunk_id,
e.knowledge_id, e.knowledge_base_id,
e.tag_id, e.content
FROM %s v
JOIN lite_embeddings e ON e.id = v.rowid
WHERE v.embedding MATCH ?
AND k = ?
AND v.rowid IN (
SELECT filtered.id
FROM lite_embeddings filtered
WHERE (filtered.is_enabled IS NULL OR filtered.is_enabled = 1)
`, tbl)
args := []interface{}{
queryBlob,
params.TopK, // 这里就是 k
}
// 追加过滤条件
for _, wp := range buildFilterWhere(params, "filtered") {
vecSQL += " AND " + wp.clause
args = append(args, wp.args...)
}
// ⚠️ 这里仍然建议加 ORDER BY,虽然 vec0 已经按距离返回
vecSQL += ") ORDER BY v.distance ASC"
type row struct {
Rowid uint
Distance float64
SourceID string
SourceType int
ChunkID string
KnowledgeID string
KnowledgeBaseID string
TagID string
Content string
}
var rows []row
if err := r.db.WithContext(ctx).
Raw(vecSQL, args...).
Scan(&rows).Error; err != nil {
return nil, fmt.Errorf("sqlite-vec query failed: %w", err)
}
logger.GetLogger(ctx).Infof("[SQLite] vectorRetrieve: query_dim=%d, threshold=%.4f, matched=%d rows", dim, params.Threshold, len(rows))
items := make([]*types.IndexWithScore, 0, len(rows))
for i, v := range rows {
// cosine distance = 1 - cosine_similarity
score := 1 - v.Distance
if params.Threshold > 0 && score < params.Threshold {
continue
}
logger.GetLogger(ctx).Infof("[SQLite] vectorRetrieve: #%d chunk_id=%s, distance=%.4f, score=%.4f, content_preview=%.60s",
i+1, v.ChunkID, v.Distance, score, v.Content)
items = append(items, &types.IndexWithScore{
ID: fmt.Sprintf("%d", v.Rowid),
SourceID: v.SourceID,
SourceType: types.SourceType(v.SourceType),
ChunkID: v.ChunkID,
KnowledgeID: v.KnowledgeID,
KnowledgeBaseID: v.KnowledgeBaseID,
TagID: v.TagID,
Content: v.Content,
Score: score,
MatchType: types.MatchTypeEmbedding,
})
}
return []*types.RetrieveResult{{
Results: items,
RetrieverEngineType: types.SQLiteRetrieverEngineType,
RetrieverType: types.VectorRetrieverType,
}}, nil
}
// --- Internal helpers ---
func toSQLiteEmbedding(info *types.IndexInfo) *sqliteEmbedding {
enabled := info.IsEnabled
return &sqliteEmbedding{
SourceID: info.SourceID,
SourceType: int(info.SourceType),
ChunkID: info.ChunkID,
KnowledgeID: info.KnowledgeID,
KnowledgeBaseID: info.KnowledgeBaseID,
TagID: info.TagID,
Content: common.CleanInvalidUTF8(info.Content),
Dimension: 0,
IsEnabled: &enabled,
}
}
func extractEmbedding(params map[string]any, sourceID string) []float32 {
if params == nil {
return nil
}
embMap, ok := params["embedding"].(map[string][]float32)
if !ok {
return nil
}
return embMap[sourceID]
}
func (r *sqliteRepository) insertVec(_ context.Context, rowID uint, dim int, emb []float32) {
r.ensureVecTable(dim)
blob, err := sqlite_vec.SerializeFloat32(emb)
if err != nil {
return
}
sql := fmt.Sprintf("INSERT INTO %s(rowid, embedding) VALUES (?, ?)", vecTableName(dim))
r.db.Exec(sql, rowID, blob)
}
func (r *sqliteRepository) deleteRowsAndVecs(_ context.Context, rows []sqliteEmbedding) {
dimIDs := make(map[int][]uint)
for _, row := range rows {
if row.Dimension > 0 {
dimIDs[row.Dimension] = append(dimIDs[row.Dimension], row.ID)
}
}
for dim, ids := range dimIDs {
if !r.vecTables[dim] {
continue
}
tbl := vecTableName(dim)
for _, id := range ids {
r.db.Exec(fmt.Sprintf("DELETE FROM %s WHERE rowid = ?", tbl), id)
}
}
for _, row := range rows {
r.db.Exec("DELETE FROM lite_embeddings_fts WHERE rowid = ?", row.ID)
}
}
func (r *sqliteRepository) copyVec(_ context.Context, srcID, dstID uint, dim int) {
if !r.vecTables[dim] {
return
}
tbl := vecTableName(dim)
r.db.Exec(fmt.Sprintf(
"INSERT INTO %s(rowid, embedding) SELECT ?, embedding FROM %s WHERE rowid = ?",
tbl, tbl,
), dstID, srcID)
}
func (r *sqliteRepository) syncFTS5Insert(_ context.Context, row *sqliteEmbedding) {
if row.ID == 0 {
return
}
tokenizedContent := tokenizeCJKBigram(row.Content)
sql := `INSERT INTO lite_embeddings_fts(rowid, content, source_id, chunk_id, knowledge_id, knowledge_base_id) VALUES(?, ?, ?, ?, ?, ?)`
r.db.Exec(sql, row.ID, tokenizedContent, row.SourceID, row.ChunkID, row.KnowledgeID, row.KnowledgeBaseID)
}
type whereClause struct {
clause string
args []interface{}
}
func buildFilterWhere(params types.RetrieveParams, tableAlias string) []whereClause {
var parts []whereClause
if len(params.KnowledgeBaseIDs) > 0 {
parts = append(parts, whereClause{
clause: tableAlias + ".knowledge_base_id IN (" + placeholders(len(params.KnowledgeBaseIDs)) + ")",
args: toInterfaceSlice(params.KnowledgeBaseIDs),
})
}
if len(params.KnowledgeIDs) > 0 {
parts = append(parts, whereClause{
clause: tableAlias + ".knowledge_id IN (" + placeholders(len(params.KnowledgeIDs)) + ")",
args: toInterfaceSlice(params.KnowledgeIDs),
})
}
if len(params.TagIDs) > 0 {
parts = append(parts, whereClause{
clause: tableAlias + ".tag_id IN (" + placeholders(len(params.TagIDs)) + ")",
args: toInterfaceSlice(params.TagIDs),
})
}
return parts
}
func placeholders(n int) string {
p := make([]string, n)
for i := range p {
p[i] = "?"
}
return strings.Join(p, ",")
}
func toInterfaceSlice(ss []string) []interface{} {
out := make([]interface{}, len(ss))
for i, s := range ss {
out[i] = s
}
return out
}
// tokenizeCJKBigram splits continuous CJK character sequences into overlapping bigrams.
// Non-CJK words are kept intact. This approach maximizes recall for CJK search.
func tokenizeCJKBigram(text string) string {
var parts []string
var currentCJK []rune
var currentNonCJK strings.Builder
flushCJK := func() {
if len(currentCJK) == 0 {
return
}
if len(currentCJK) == 1 {
parts = append(parts, string(currentCJK[0]))
} else {
for i := 0; i < len(currentCJK)-1; i++ {
parts = append(parts, string(currentCJK[i])+string(currentCJK[i+1]))
}
}
currentCJK = currentCJK[:0]
}
flushNonCJK := func() {
if currentNonCJK.Len() > 0 {
parts = append(parts, currentNonCJK.String())
currentNonCJK.Reset()
}
}
for _, r := range text {
// unicode.Han covers Chinese characters
if unicode.Is(unicode.Han, r) {
flushNonCJK()
currentCJK = append(currentCJK, r)
} else if unicode.IsSpace(r) || unicode.IsPunct(r) || unicode.IsSymbol(r) {
// Delimiters
flushCJK()
flushNonCJK()
} else {
// Alphanumeric or other languages
flushCJK()
currentNonCJK.WriteRune(r)
}
}
flushCJK()
flushNonCJK()
return strings.Join(parts, " ")
}
// sanitizeFTS5Query builds an FTS5 query from user input by applying bigram tokenization.
func sanitizeFTS5Query(q string) string {
q = strings.TrimSpace(q)
if q == "" {
return q
}
// Tokenize input query with bigrams (same as during indexing)
tokenized := tokenizeCJKBigram(q)
fields := strings.Fields(tokenized)
var parts []string
for _, f := range fields {
if f != "" {
parts = append(parts, `"`+f+`"`)
}
}
if len(parts) == 0 {
return ""
}
// Use OR because we want fuzzy match across multiple bigrams/words
return strings.Join(parts, " OR ")
}
@@ -0,0 +1,188 @@
package sqlite
import (
"context"
"fmt"
"strings"
"sync"
"testing"
"github.com/Tencent/WeKnora/internal/types"
sqlite_vec "github.com/asg017/sqlite-vec-go-bindings/cgo"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
gormsqlite "gorm.io/driver/sqlite"
"gorm.io/gorm"
)
var registerSQLiteVec sync.Once
func newSQLiteRetrieverTestRepository(t *testing.T) *sqliteRepository {
t.Helper()
registerSQLiteVec.Do(sqlite_vec.Auto)
dbName := strings.NewReplacer("/", "_", " ", "_").Replace(t.Name())
db, err := gorm.Open(gormsqlite.Open(fmt.Sprintf("file:%s?mode=memory&cache=shared", dbName)), &gorm.Config{})
require.NoError(t, err)
sqlDB, err := db.DB()
require.NoError(t, err)
t.Cleanup(func() {
require.NoError(t, sqlDB.Close())
})
repository, ok := NewSQLiteRetrieveEngineRepository(db).(*sqliteRepository)
require.True(t, ok)
return repository
}
func saveSQLiteTestVector(t *testing.T, repository *sqliteRepository, info *types.IndexInfo, embedding []float32) {
t.Helper()
require.NoError(t, repository.Save(context.Background(), info, map[string]any{
"embedding": map[string][]float32{info.SourceID: embedding},
}))
}
func sqliteTestIndex(chunkID, knowledgeBaseID, knowledgeID, tagID string, enabled bool) *types.IndexInfo {
return &types.IndexInfo{
Content: chunkID,
SourceID: "source-" + chunkID,
SourceType: types.ChunkSourceType,
ChunkID: chunkID,
KnowledgeID: knowledgeID,
KnowledgeBaseID: knowledgeBaseID,
TagID: tagID,
IsEnabled: enabled,
}
}
func TestVectorRetrieveFiltersBeforeTopK(t *testing.T) {
testCases := []struct {
name string
blocker *types.IndexInfo
configure func(*types.RetrieveParams)
}{
{
name: "knowledge base",
blocker: sqliteTestIndex("blocker", "kb-other", "knowledge-target", "tag-target", true),
configure: func(params *types.RetrieveParams) {
params.KnowledgeBaseIDs = []string{"kb-target"}
},
},
{
name: "knowledge",
blocker: sqliteTestIndex("blocker", "kb-target", "knowledge-other", "tag-target", true),
configure: func(params *types.RetrieveParams) {
params.KnowledgeIDs = []string{"knowledge-target"}
},
},
{
name: "tag",
blocker: sqliteTestIndex("blocker", "kb-target", "knowledge-target", "tag-other", true),
configure: func(params *types.RetrieveParams) {
params.TagIDs = []string{"tag-target"}
},
},
{
name: "enabled status",
blocker: sqliteTestIndex("blocker", "kb-target", "knowledge-target", "tag-target", false),
configure: func(_ *types.RetrieveParams) {},
},
}
for _, testCase := range testCases {
t.Run(testCase.name, func(t *testing.T) {
repository := newSQLiteRetrieverTestRepository(t)
saveSQLiteTestVector(t, repository, testCase.blocker, []float32{1, 0})
saveSQLiteTestVector(t, repository,
sqliteTestIndex("target", "kb-target", "knowledge-target", "tag-target", true),
[]float32{0.8, 0.6},
)
params := types.RetrieveParams{
Embedding: []float32{1, 0},
TopK: 1,
RetrieverType: types.VectorRetrieverType,
}
testCase.configure(&params)
results, err := repository.vectorRetrieve(context.Background(), params)
require.NoError(t, err)
require.Len(t, results, 1)
require.Len(t, results[0].Results, 1)
assert.Equal(t, "target", results[0].Results[0].ChunkID)
})
}
}
func TestVectorRetrieveZeroThresholdDoesNotFilter(t *testing.T) {
repository := newSQLiteRetrieverTestRepository(t)
saveSQLiteTestVector(t, repository,
sqliteTestIndex("anti-correlated", "kb-target", "knowledge-target", "tag-target", true),
[]float32{-1, 0},
)
results, err := repository.vectorRetrieve(context.Background(), types.RetrieveParams{
Embedding: []float32{1, 0},
KnowledgeBaseIDs: []string{"kb-target"},
TopK: 1,
Threshold: 0,
RetrieverType: types.VectorRetrieverType,
})
require.NoError(t, err)
require.Len(t, results, 1)
require.Len(t, results[0].Results, 1)
assert.Equal(t, "anti-correlated", results[0].Results[0].ChunkID)
assert.Less(t, results[0].Results[0].Score, 0.0)
}
func TestVectorRetrieveAppliesSimilarityThreshold(t *testing.T) {
repository := newSQLiteRetrieverTestRepository(t)
saveSQLiteTestVector(t, repository,
sqliteTestIndex("below-threshold", "kb-target", "knowledge-target", "tag-target", true),
[]float32{0.5, 0.8660254},
)
results, err := repository.vectorRetrieve(context.Background(), types.RetrieveParams{
Embedding: []float32{1, 0},
KnowledgeBaseIDs: []string{"kb-target"},
TopK: 1,
Threshold: 0.75,
RetrieverType: types.VectorRetrieverType,
})
require.NoError(t, err)
require.Len(t, results, 1)
assert.Empty(t, results[0].Results)
}
func TestRetrieveReturnsVectorQueryError(t *testing.T) {
repository := newSQLiteRetrieverTestRepository(t)
saveSQLiteTestVector(t, repository,
sqliteTestIndex("chunk", "kb-target", "knowledge-target", "tag-target", true),
[]float32{1, 0},
)
require.NoError(t, repository.db.Exec("DROP TABLE "+vecTableName(2)).Error)
results, err := repository.Retrieve(context.Background(), types.RetrieveParams{
Embedding: []float32{1, 0},
TopK: 1,
RetrieverType: types.VectorRetrieverType,
})
require.Error(t, err)
assert.Nil(t, results)
assert.Contains(t, err.Error(), "sqlite-vec query failed")
}
func TestRetrieveReturnsKeywordQueryError(t *testing.T) {
repository := newSQLiteRetrieverTestRepository(t)
require.NoError(t, repository.db.Exec("DROP TABLE IF EXISTS lite_embeddings_fts").Error)
results, err := repository.Retrieve(context.Background(), types.RetrieveParams{
Query: "missing fts table",
TopK: 1,
RetrieverType: types.KeywordsRetrieverType,
})
require.Error(t, err)
assert.Nil(t, results)
assert.Contains(t, err.Error(), "FTS5 query failed")
}
@@ -0,0 +1,782 @@
package tencentvectordb
import (
"context"
"crypto/sha256"
"encoding/hex"
"fmt"
"maps"
"os"
"slices"
"sort"
"strconv"
"strings"
"unicode/utf8"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
"github.com/tencent/vectordatabase-sdk-go/tcvdbtext/encoder"
"github.com/tencent/vectordatabase-sdk-go/tcvectordb"
)
const copyIndicesQueryPageSize int64 = 500
// NewTencentVectorDBRetrieveEngineRepository creates a Tencent VectorDB-backed retrieve repository.
func NewTencentVectorDBRetrieveEngineRepository(
client *tcvectordb.RpcClient,
databaseName string,
indexCfg *types.IndexConfig,
) interfaces.RetrieveEngineRepository {
if databaseName == "" {
databaseName = os.Getenv(envTencentVectorDBDatabase)
}
if databaseName == "" {
databaseName = defaultDatabaseName
}
collectionBaseName := types.ResolveCollectionName(indexCfg, envTencentVectorDBCollection, defaultCollectionName)
return &repository{
client: client,
databaseName: databaseName,
collectionBaseName: collectionBaseName,
useDimensionSuffix: shouldUseDimensionSuffix(indexCfg),
shardsNum: defaultIfZero(indexCfg.GetShardsNum(1), 1),
replicasNum: resolveReplicaNumber(indexCfg),
}
}
func resolveReplicaNumber(indexCfg *types.IndexConfig) int {
if indexCfg != nil && indexCfg.ReplicaNumber > 0 {
return indexCfg.ReplicaNumber
}
if raw := strings.TrimSpace(os.Getenv(envTencentVectorDBReplicaNum)); raw != "" {
replicas, err := strconv.Atoi(raw)
if err == nil && replicas >= 0 {
return replicas
}
}
return defaultReplicaNumber
}
func (r *repository) EngineType() types.RetrieverEngineType {
return types.TencentVectorDBRetrieverEngineType
}
func (r *repository) Support() []types.RetrieverType {
return []types.RetrieverType{types.KeywordsRetrieverType, types.VectorRetrieverType}
}
func (r *repository) Save(ctx context.Context, indexInfo *types.IndexInfo, params map[string]any) error {
return r.BatchSave(ctx, []*types.IndexInfo{indexInfo}, params)
}
func (r *repository) BatchSave(ctx context.Context, indexInfoList []*types.IndexInfo, params map[string]any) error {
log := logger.GetLogger(ctx)
if len(indexInfoList) == 0 {
return nil
}
embeddingsByDimension := make(map[int][]*vectorEmbedding)
for _, indexInfo := range indexInfoList {
embedding := toVectorEmbedding(indexInfo, params)
if len(embedding.Embedding) == 0 {
log.Warnf("[TencentVectorDB] skip empty embedding for chunk_id=%s", indexInfo.ChunkID)
continue
}
dim := len(embedding.Embedding)
embeddingsByDimension[dim] = append(embeddingsByDimension[dim], embedding)
}
if len(embeddingsByDimension) == 0 {
return nil
}
bm25, err := r.bm25Encoder()
if err != nil {
return err
}
buildIndex := true
for dim, embeddings := range embeddingsByDimension {
docs, err := r.toDocumentsWithSparseVectors(bm25, embeddings)
if err != nil {
return err
}
if err := r.ensureCollection(ctx, dim); err != nil {
return err
}
collectionName := r.collectionName(dim)
_, err = r.client.Database(r.databaseName).Collection(collectionName).Upsert(
ctx,
docs,
&tcvectordb.UpsertDocumentParams{BuildIndex: &buildIndex},
)
if err != nil {
return fmt.Errorf("tencent vectordb batch save %s: %w", collectionName, err)
}
}
return nil
}
func (r *repository) EstimateStorageSize(ctx context.Context, indexInfoList []*types.IndexInfo, params map[string]any) int64 {
var total int64
for _, indexInfo := range indexInfoList {
embedding := toVectorEmbedding(indexInfo, params)
total += int64(len(embedding.Content))
total += int64(len(embedding.Embedding) * 4)
total += int64(len(embedding.Content) * 2)
total += int64(len(embedding.SourceID) + len(embedding.ChunkID) + len(embedding.KnowledgeID) + len(embedding.KnowledgeBaseID) + 256)
}
logger.GetLogger(ctx).Infof("[TencentVectorDB] estimated storage size for %d indices: %d bytes", len(indexInfoList), total)
return total
}
func (r *repository) DeleteByChunkIDList(ctx context.Context, chunkIDList []string, dimension int, knowledgeType string) error {
return r.deleteByFilter(ctx, dimension, tcvectordb.In(fieldChunkID, chunkIDList))
}
func (r *repository) DeleteBySourceIDList(ctx context.Context, sourceIDList []string, dimension int, knowledgeType string) error {
return r.deleteByFilter(ctx, dimension, tcvectordb.In(fieldSourceID, sourceIDList))
}
func (r *repository) DeleteByKnowledgeIDList(ctx context.Context, knowledgeIDList []string, dimension int, knowledgeType string) error {
return r.deleteByFilter(ctx, dimension, tcvectordb.In(fieldKnowledgeID, knowledgeIDList))
}
func (r *repository) CopyIndices(
ctx context.Context,
sourceKnowledgeBaseID string,
sourceToTargetKBIDMap map[string]string,
sourceToTargetChunkIDMap map[string]string,
targetKnowledgeBaseID string,
dimension int,
knowledgeType string,
) error {
if len(sourceToTargetChunkIDMap) == 0 {
return nil
}
collectionName := r.collectionName(dimension)
ids := slices.Collect(maps.Keys(sourceToTargetChunkIDMap))
var embeddings []*vectorEmbedding
for offset := int64(0); ; offset += copyIndicesQueryPageSize {
query, err := r.client.Database(r.databaseName).Collection(collectionName).Query(
ctx,
nil,
copySourceQueryParams(sourceKnowledgeBaseID, ids, offset),
)
if err != nil {
return fmt.Errorf("tencent vectordb query source indices: %w", err)
}
embeddings = append(embeddings, remapCopiedEmbeddings(
query.Documents,
sourceToTargetKBIDMap,
sourceToTargetChunkIDMap,
targetKnowledgeBaseID,
)...)
if len(query.Documents) < int(copyIndicesQueryPageSize) {
break
}
}
if len(embeddings) == 0 {
return nil
}
bm25, err := r.bm25Encoder()
if err != nil {
return err
}
docs, err := r.toDocumentsWithSparseVectors(bm25, embeddings)
if err != nil {
return err
}
buildIndex := true
_, err = r.client.Database(r.databaseName).Collection(collectionName).Upsert(
ctx,
docs,
&tcvectordb.UpsertDocumentParams{BuildIndex: &buildIndex},
)
if err != nil {
return fmt.Errorf("tencent vectordb copy indices: %w", err)
}
return nil
}
func (r *repository) BatchUpdateChunkEnabledStatus(ctx context.Context, chunkStatusMap map[string]bool) error {
if len(chunkStatusMap) == 0 {
return nil
}
grouped := make(map[bool][]string)
for chunkID, enabled := range chunkStatusMap {
grouped[enabled] = append(grouped[enabled], chunkID)
}
for enabled, chunkIDs := range grouped {
if err := r.updateChunkFields(ctx, chunkIDs, map[string]tcvectordb.Field{fieldIsEnabled: {Val: boolToUint64(enabled)}}); err != nil {
return err
}
}
return nil
}
func (r *repository) BatchUpdateChunkTagID(ctx context.Context, chunkTagMap map[string]string) error {
if len(chunkTagMap) == 0 {
return nil
}
grouped := make(map[string][]string)
for chunkID, tagID := range chunkTagMap {
grouped[tagID] = append(grouped[tagID], chunkID)
}
for tagID, chunkIDs := range grouped {
if err := r.updateChunkFields(ctx, chunkIDs, map[string]tcvectordb.Field{fieldTagID: {Val: tagID}}); err != nil {
return err
}
}
return nil
}
func (r *repository) Retrieve(ctx context.Context, params types.RetrieveParams) ([]*types.RetrieveResult, error) {
switch params.RetrieverType {
case types.VectorRetrieverType:
return r.VectorRetrieve(ctx, params)
case types.KeywordsRetrieverType:
return r.KeywordsRetrieve(ctx, params)
default:
return nil, fmt.Errorf("invalid retriever type: %s", params.RetrieverType)
}
}
func (r *repository) VectorRetrieve(ctx context.Context, params types.RetrieveParams) ([]*types.RetrieveResult, error) {
dimension := len(params.Embedding)
if dimension == 0 {
return r.retrieveResult(nil, types.VectorRetrieverType), nil
}
collectionName := r.collectionName(dimension)
exists, err := r.client.Database(r.databaseName).ExistsCollection(ctx, collectionName)
if err != nil {
return nil, fmt.Errorf("tencent vectordb check collection %s: %w", collectionName, err)
}
if !exists {
return r.retrieveResult(nil, types.VectorRetrieverType), nil
}
limit := int64(params.TopK)
if limit <= 0 {
limit = 10
}
searchParams := &tcvectordb.SearchDocumentParams{
Filter: r.baseFilter(params),
Params: &tcvectordb.SearchDocParams{Ef: 100},
RetrieveVector: false,
OutputFields: outputFields(),
Limit: limit,
}
if params.Threshold > 0 {
radius := float32(params.Threshold)
searchParams.Radius = &radius
}
search, err := r.client.Database(r.databaseName).Collection(collectionName).Search(ctx, [][]float32{params.Embedding}, searchParams)
if err != nil {
return nil, fmt.Errorf("tencent vectordb vector search %s: %w", collectionName, err)
}
if len(search.Documents) == 0 {
return r.retrieveResult(nil, types.VectorRetrieverType), nil
}
results := make([]*types.IndexWithScore, 0, len(search.Documents[0]))
for _, doc := range search.Documents[0] {
embedding := fromDocument(doc)
results = append(results, toIndexWithScore(embedding, types.MatchTypeEmbedding))
}
return r.retrieveResult(results, types.VectorRetrieverType), nil
}
func (r *repository) KeywordsRetrieve(ctx context.Context, params types.RetrieveParams) ([]*types.RetrieveResult, error) {
log := logger.GetLogger(ctx)
query := strings.TrimSpace(params.Query)
if query == "" {
return r.retrieveResult(nil, types.KeywordsRetrieverType), nil
}
bm25, err := r.bm25Encoder()
if err != nil {
return nil, err
}
queryVectors, err := bm25.EncodeQueries([]string{query})
if err != nil {
return nil, fmt.Errorf("tencent vectordb encode keyword query: %w", err)
}
if len(queryVectors) == 0 || len(queryVectors[0]) == 0 {
return r.retrieveResult(nil, types.KeywordsRetrieverType), nil
}
collections, err := r.client.Database(r.databaseName).ListCollection(ctx)
if err != nil {
return nil, fmt.Errorf("tencent vectordb list collections: %w", err)
}
limit := params.TopK
if limit <= 0 {
limit = 10
}
results := make([]*types.IndexWithScore, 0, limit)
matchedCollections := 0
failedCollections := 0
for _, collection := range collections.Collections {
if !r.matchesCollection(collection.CollectionName) {
continue
}
matchedCollections++
search, err := r.client.Database(r.databaseName).Collection(collection.CollectionName).FullTextSearch(
ctx,
tcvectordb.FullTextSearchParams{
Filter: r.baseFilter(params),
RetrieveVector: false,
OutputFields: outputFields(),
Limit: &limit,
Match: &tcvectordb.FullTextSearchMatchOption{
FieldName: fieldSparseVector,
Data: queryVectors,
},
},
)
if err != nil {
failedCollections++
log.Warnf("[TencentVectorDB] keyword search failed in %s: %v", collection.CollectionName, err)
continue
}
if len(search.Documents) == 0 {
continue
}
for _, doc := range search.Documents[0] {
embedding := fromDocument(doc)
results = append(results, toIndexWithScore(embedding, types.MatchTypeKeywords))
}
}
if matchedCollections > 0 && failedCollections == matchedCollections {
return nil, fmt.Errorf("tencent vectordb keyword search failed in all matched collections; ensure collections have the %q sparse vector index and reimport data if they were created before keyword support", fieldSparseVector)
}
sort.SliceStable(results, func(i, j int) bool {
return results[i].Score > results[j].Score
})
if len(results) > limit {
results = results[:limit]
}
return r.retrieveResult(results, types.KeywordsRetrieverType), nil
}
func (r *repository) ensureCollection(ctx context.Context, dimension int) error {
if _, ok := r.initialized.Load(dimension); ok {
return nil
}
if _, err := r.client.CreateDatabaseIfNotExists(ctx, r.databaseName); err != nil {
return fmt.Errorf("tencent vectordb ensure database %s: %w", r.databaseName, err)
}
collectionName := r.collectionName(dimension)
exists, err := r.client.Database(r.databaseName).ExistsCollection(ctx, collectionName)
if err != nil {
return fmt.Errorf("tencent vectordb check collection %s: %w", collectionName, err)
}
if exists {
r.initialized.Store(dimension, true)
return nil
}
indexes := tcvectordb.Indexes{
VectorIndex: []tcvectordb.VectorIndex{
{
FilterIndex: tcvectordb.FilterIndex{
FieldName: fieldVector,
FieldType: tcvectordb.Vector,
IndexType: tcvectordb.HNSW,
},
Dimension: uint32(dimension),
MetricType: tcvectordb.COSINE,
Params: &tcvectordb.HNSWParam{
M: 16,
EfConstruction: 200,
},
},
},
SparseVectorIndex: []tcvectordb.SparseVectorIndex{
{
FieldName: fieldSparseVector,
FieldType: tcvectordb.SparseVector,
IndexType: tcvectordb.SPARSE_INVERTED,
MetricType: tcvectordb.IP,
},
},
FilterIndex: []tcvectordb.FilterIndex{
{FieldName: fieldID, FieldType: tcvectordb.String, IndexType: tcvectordb.PRIMARY},
{FieldName: fieldContent, FieldType: tcvectordb.String, IndexType: tcvectordb.FILTER},
{FieldName: fieldSourceID, FieldType: tcvectordb.String, IndexType: tcvectordb.FILTER},
{FieldName: fieldSourceType, FieldType: tcvectordb.Uint64, IndexType: tcvectordb.FILTER},
{FieldName: fieldChunkID, FieldType: tcvectordb.String, IndexType: tcvectordb.FILTER},
{FieldName: fieldKnowledgeID, FieldType: tcvectordb.String, IndexType: tcvectordb.FILTER},
{FieldName: fieldKnowledgeBaseID, FieldType: tcvectordb.String, IndexType: tcvectordb.FILTER},
{FieldName: fieldTagID, FieldType: tcvectordb.String, IndexType: tcvectordb.FILTER},
{FieldName: fieldIsEnabled, FieldType: tcvectordb.Uint64, IndexType: tcvectordb.FILTER},
},
}
_, err = r.client.Database(r.databaseName).CreateCollection(
ctx,
collectionName,
uint32(r.shardsNum),
uint32(r.replicasNum),
fmt.Sprintf("WeKnora embeddings collection with dimension %d", dimension),
indexes,
)
if err != nil {
if isCollectionAlreadyExistsErr(err) {
logger.GetLogger(ctx).Infof("[TencentVectorDB] collection %s already exists, skip create", collectionName)
r.initialized.Store(dimension, true)
return nil
}
return fmt.Errorf("tencent vectordb create collection %s: %w", collectionName, err)
}
r.initialized.Store(dimension, true)
return nil
}
func isCollectionAlreadyExistsErr(err error) bool {
if err == nil {
return false
}
msg := strings.ToLower(err.Error())
return strings.Contains(msg, "code: 15202") || strings.Contains(msg, "already exist")
}
func (r *repository) deleteByFilter(ctx context.Context, dimension int, cond string) error {
if cond == "" {
return nil
}
collectionName := r.collectionName(dimension)
_, err := r.client.Database(r.databaseName).Collection(collectionName).Delete(ctx, tcvectordb.DeleteDocumentParams{
Filter: tcvectordb.NewFilter(cond),
})
if err != nil {
return fmt.Errorf("tencent vectordb delete from %s: %w", collectionName, err)
}
return nil
}
func (r *repository) updateChunkFields(ctx context.Context, chunkIDs []string, fields map[string]tcvectordb.Field) error {
collections, err := r.client.Database(r.databaseName).ListCollection(ctx)
if err != nil {
return fmt.Errorf("tencent vectordb list collections: %w", err)
}
for _, collection := range collections.Collections {
if !r.matchesCollection(collection.CollectionName) {
continue
}
_, err := r.client.Database(r.databaseName).Collection(collection.CollectionName).Update(ctx, tcvectordb.UpdateDocumentParams{
QueryFilter: tcvectordb.NewFilter(tcvectordb.In(fieldChunkID, chunkIDs)),
UpdateFields: fields,
})
if err != nil {
return fmt.Errorf("tencent vectordb update chunks in %s: %w", collection.CollectionName, err)
}
}
return nil
}
func (r *repository) collectionName(dimension int) string {
if !r.useDimensionSuffix {
return r.collectionBaseName
}
return fmt.Sprintf("%s_%d", r.collectionBaseName, dimension)
}
func (r *repository) matchesCollection(collectionName string) bool {
if !r.useDimensionSuffix {
return collectionName == r.collectionBaseName
}
return strings.HasPrefix(collectionName, r.collectionBaseName+"_")
}
func shouldUseDimensionSuffix(indexCfg *types.IndexConfig) bool {
return indexCfg == nil || indexCfg.CollectionName == ""
}
func (r *repository) baseFilter(params types.RetrieveParams) *tcvectordb.Filter {
conditions := []string{fmt.Sprintf("%s=1", fieldIsEnabled)}
if len(params.KnowledgeBaseIDs) > 0 {
conditions = append(conditions, tcvectordb.In(fieldKnowledgeBaseID, params.KnowledgeBaseIDs))
}
if len(params.KnowledgeIDs) > 0 {
conditions = append(conditions, tcvectordb.In(fieldKnowledgeID, params.KnowledgeIDs))
}
if len(params.TagIDs) > 0 {
conditions = append(conditions, tcvectordb.In(fieldTagID, params.TagIDs))
}
if len(params.ExcludeKnowledgeIDs) > 0 {
conditions = append(conditions, tcvectordb.NotIn(fieldKnowledgeID, params.ExcludeKnowledgeIDs))
}
if len(params.ExcludeChunkIDs) > 0 {
conditions = append(conditions, tcvectordb.NotIn(fieldChunkID, params.ExcludeChunkIDs))
}
return tcvectordb.NewFilter(strings.Join(conditions, " and "))
}
func (r *repository) retrieveResult(results []*types.IndexWithScore, retrieverType types.RetrieverType) []*types.RetrieveResult {
return []*types.RetrieveResult{
{
Results: results,
RetrieverEngineType: types.TencentVectorDBRetrieverEngineType,
RetrieverType: retrieverType,
},
}
}
func (r *repository) bm25Encoder() (encoder.SparseEncoder, error) {
r.bm25Once.Do(func() {
r.bm25, r.bm25Err = encoder.NewBM25Encoder(&encoder.BM25EncoderParams{
Bm25Language: encoder.BM25_ZH_CONTENT,
})
if r.bm25Err != nil {
r.bm25Err = fmt.Errorf("tencent vectordb init BM25 encoder: %w", r.bm25Err)
}
})
return r.bm25, r.bm25Err
}
func (r *repository) toDocumentsWithSparseVectors(
bm25 encoder.SparseEncoder,
embeddings []*vectorEmbedding,
) ([]tcvectordb.Document, error) {
texts := make([]string, 0, len(embeddings))
for _, embedding := range embeddings {
texts = append(texts, embedding.Content)
}
sparseVectors, err := bm25.EncodeTexts(texts)
if err != nil {
return nil, fmt.Errorf("tencent vectordb encode BM25 sparse vectors: %w", err)
}
docs := make([]tcvectordb.Document, 0, len(embeddings))
for i, embedding := range embeddings {
if i >= len(sparseVectors) {
return nil, fmt.Errorf("tencent vectordb encoded sparse vector count mismatch: got %d, want %d", len(sparseVectors), len(embeddings))
}
embedding.SparseVector = sparseVectors[i]
docs = append(docs, toDocument(embedding))
}
return docs, nil
}
func toVectorEmbedding(indexInfo *types.IndexInfo, params map[string]any) *vectorEmbedding {
embedding := &vectorEmbedding{
ID: indexInfo.ID,
Content: cleanInvalidUTF8(indexInfo.Content),
SourceID: indexInfo.SourceID,
SourceType: int(indexInfo.SourceType),
ChunkID: indexInfo.ChunkID,
KnowledgeID: indexInfo.KnowledgeID,
KnowledgeBaseID: indexInfo.KnowledgeBaseID,
TagID: indexInfo.TagID,
IsEnabled: indexInfo.IsEnabled,
}
if embedding.ID == "" {
embedding.ID = indexInfo.SourceID
}
if embedding.ID == "" {
embedding.ID = indexInfo.ChunkID
}
if params != nil && slices.Contains(slices.Collect(maps.Keys(params)), fieldVector) {
if embeddingMap, ok := params[fieldVector].(map[string][]float32); ok {
embedding.Embedding = lookupEmbedding(embeddingMap, indexInfo)
}
}
if params != nil && slices.Contains(slices.Collect(maps.Keys(params)), "embedding") {
if embeddingMap, ok := params["embedding"].(map[string][]float32); ok {
embedding.Embedding = lookupEmbedding(embeddingMap, indexInfo)
}
}
return embedding
}
func lookupEmbedding(embeddingMap map[string][]float32, indexInfo *types.IndexInfo) []float32 {
if embedding, ok := embeddingMap[indexInfo.SourceID]; ok {
return embedding
}
return embeddingMap[indexInfo.ChunkID]
}
func copySourceQueryParams(sourceKnowledgeBaseID string, chunkIDs []string, offset int64) *tcvectordb.QueryDocumentParams {
conditions := []string{tcvectordb.In(fieldChunkID, chunkIDs)}
if sourceKnowledgeBaseID != "" {
conditions = append([]string{tcvectordb.In(fieldKnowledgeBaseID, []string{sourceKnowledgeBaseID})}, conditions...)
}
return &tcvectordb.QueryDocumentParams{
Filter: tcvectordb.NewFilter(strings.Join(conditions, " and ")),
RetrieveVector: true,
OutputFields: outputFields(),
Offset: offset,
Limit: copyIndicesQueryPageSize,
}
}
func remapCopiedEmbeddings(
docs []tcvectordb.Document,
sourceToTargetKBIDMap map[string]string,
sourceToTargetChunkIDMap map[string]string,
targetKnowledgeBaseID string,
) []*vectorEmbedding {
embeddings := make([]*vectorEmbedding, 0, len(docs))
for _, doc := range docs {
embedding := fromDocument(doc)
targetChunkID := sourceToTargetChunkIDMap[embedding.ChunkID]
if targetChunkID == "" {
continue
}
originalSourceID := embedding.SourceID
if originalSourceID == "" {
originalSourceID = embedding.ID
}
targetSourceID := translateSourceID(originalSourceID, embedding.ChunkID, targetChunkID)
embedding.ID = targetSourceID
embedding.SourceID = targetSourceID
embedding.ChunkID = targetChunkID
embedding.KnowledgeBaseID = targetKnowledgeBaseID
if targetKBID := sourceToTargetKBIDMap[embedding.KnowledgeID]; targetKBID != "" {
embedding.KnowledgeID = targetKBID
}
embeddings = append(embeddings, embedding)
}
return embeddings
}
func translateSourceID(originalSourceID, sourceChunkID, targetChunkID string) string {
switch {
case originalSourceID == sourceChunkID:
return targetChunkID
case strings.HasPrefix(originalSourceID, sourceChunkID+"-"):
questionID := strings.TrimPrefix(originalSourceID, sourceChunkID+"-")
return fmt.Sprintf("%s-%s", targetChunkID, questionID)
default:
sum := sha256.Sum256([]byte(targetChunkID + "\x00" + sourceChunkID + "\x00" + originalSourceID))
return fmt.Sprintf("%s-%s", targetChunkID, hex.EncodeToString(sum[:])[:16])
}
}
func cleanInvalidUTF8(s string) string {
var b strings.Builder
b.Grow(len(s))
for i := 0; i < len(s); {
r, size := utf8.DecodeRuneInString(s[i:])
if r == utf8.RuneError && size == 1 {
i++
continue
}
if r == 0 {
i += size
continue
}
b.WriteRune(r)
i += size
}
return b.String()
}
func toDocument(embedding *vectorEmbedding) tcvectordb.Document {
return tcvectordb.Document{
Id: embedding.ID,
Vector: embedding.Embedding,
SparseVector: embedding.SparseVector,
Fields: map[string]tcvectordb.Field{
fieldContent: {Val: embedding.Content},
fieldSourceID: {Val: embedding.SourceID},
fieldSourceType: {Val: uint64(embedding.SourceType)},
fieldChunkID: {Val: embedding.ChunkID},
fieldKnowledgeID: {Val: embedding.KnowledgeID},
fieldKnowledgeBaseID: {Val: embedding.KnowledgeBaseID},
fieldTagID: {Val: embedding.TagID},
fieldIsEnabled: {Val: boolToUint64(embedding.IsEnabled)},
},
}
}
func fromDocument(doc tcvectordb.Document) *vectorEmbedding {
return &vectorEmbedding{
ID: doc.Id,
Content: fieldString(doc, fieldContent),
SourceID: fieldString(doc, fieldSourceID),
SourceType: int(fieldUint64(doc, fieldSourceType)),
ChunkID: fieldString(doc, fieldChunkID),
KnowledgeID: fieldString(doc, fieldKnowledgeID),
KnowledgeBaseID: fieldString(doc, fieldKnowledgeBaseID),
TagID: fieldString(doc, fieldTagID),
Embedding: doc.Vector,
SparseVector: doc.SparseVector,
IsEnabled: fieldUint64(doc, fieldIsEnabled) == 1,
Score: float64(doc.Score),
}
}
func toIndexWithScore(embedding *vectorEmbedding, matchType types.MatchType) *types.IndexWithScore {
return &types.IndexWithScore{
ID: embedding.ID,
Content: embedding.Content,
SourceID: embedding.SourceID,
SourceType: types.SourceType(embedding.SourceType),
ChunkID: embedding.ChunkID,
KnowledgeID: embedding.KnowledgeID,
KnowledgeBaseID: embedding.KnowledgeBaseID,
TagID: embedding.TagID,
Score: embedding.Score,
MatchType: matchType,
IsEnabled: embedding.IsEnabled,
}
}
func outputFields() []string {
return []string{
fieldID,
fieldContent,
fieldSourceID,
fieldSourceType,
fieldChunkID,
fieldKnowledgeID,
fieldKnowledgeBaseID,
fieldTagID,
fieldIsEnabled,
}
}
func fieldString(doc tcvectordb.Document, name string) string {
field, ok := doc.Fields[name]
if !ok {
return ""
}
return field.String()
}
func fieldUint64(doc tcvectordb.Document, name string) uint64 {
field, ok := doc.Fields[name]
if !ok {
return 0
}
return field.Uint64()
}
func boolToUint64(v bool) uint64 {
if v {
return 1
}
return 0
}
func defaultIfZero(v, def int) int {
if v <= 0 {
return def
}
return v
}
@@ -0,0 +1,125 @@
package tencentvectordb
import (
"errors"
"strings"
"testing"
"github.com/Tencent/WeKnora/internal/types"
"github.com/stretchr/testify/assert"
"github.com/tencent/vectordatabase-sdk-go/tcvdbtext/encoder"
)
func TestSupportIncludesKeywordAndVectorRetrieval(t *testing.T) {
repo := NewTencentVectorDBRetrieveEngineRepository(nil, "", nil)
supports := repo.Support()
assert.Contains(t, supports, types.KeywordsRetrieverType)
assert.Contains(t, supports, types.VectorRetrieverType)
}
func TestToDocumentIncludesSparseVector(t *testing.T) {
embedding := &vectorEmbedding{
ID: "chunk-1",
Content: "腾讯云向量数据库支持关键词检索",
SourceID: "source-1",
SourceType: int(types.ChunkSourceType),
ChunkID: "chunk-1",
KnowledgeID: "knowledge-1",
KnowledgeBaseID: "kb-1",
TagID: "tag-1",
Embedding: []float32{0.1, 0.2},
SparseVector: []encoder.SparseVecItem{
{TermId: 10, Score: 0.3},
{TermId: 20, Score: 0.7},
},
IsEnabled: true,
}
doc := toDocument(embedding)
assert.Equal(t, embedding.ID, doc.Id)
assert.Equal(t, embedding.Embedding, doc.Vector)
assert.Equal(t, embedding.SparseVector, doc.SparseVector)
assert.Equal(t, "腾讯云向量数据库支持关键词检索", doc.Fields[fieldContent].String())
assert.Equal(t, uint64(1), doc.Fields[fieldIsEnabled].Uint64())
}
func TestBaseFilterBuildsTencentVectorDBCondition(t *testing.T) {
repo := &repository{}
filter := repo.baseFilter(types.RetrieveParams{
KnowledgeBaseIDs: []string{"kb-1"},
KnowledgeIDs: []string{"knowledge-1", "knowledge-2"},
TagIDs: []string{"tag-1"},
ExcludeKnowledgeIDs: []string{"knowledge-9"},
ExcludeChunkIDs: []string{"chunk-9"},
})
cond := filter.Cond()
for _, want := range []string{
"is_enabled=1",
"knowledge_base_id in (\"kb-1\")",
"knowledge_id in (\"knowledge-1\",\"knowledge-2\")",
"tag_id in (\"tag-1\")",
"knowledge_id not in (\"knowledge-9\")",
"chunk_id not in (\"chunk-9\")",
} {
assert.True(t, strings.Contains(cond, want), "condition %q should contain %q", cond, want)
}
}
func TestTencentVectorDBDefaultsToReplicaNumberOne(t *testing.T) {
repo := NewTencentVectorDBRetrieveEngineRepository(nil, "", nil).(*repository)
assert.Equal(t, 1, repo.replicasNum)
}
func TestTencentVectorDBUsesEnvReplicaNumber(t *testing.T) {
t.Setenv(envTencentVectorDBReplicaNum, "0")
repo := NewTencentVectorDBRetrieveEngineRepository(nil, "", nil).(*repository)
assert.Equal(t, 0, repo.replicasNum)
}
func TestTencentVectorDBUsesPositiveEnvReplicaNumber(t *testing.T) {
t.Setenv(envTencentVectorDBReplicaNum, "3")
repo := NewTencentVectorDBRetrieveEngineRepository(nil, "", nil).(*repository)
assert.Equal(t, 3, repo.replicasNum)
}
func TestTencentVectorDBUsesConfiguredReplicaNumber(t *testing.T) {
t.Setenv(envTencentVectorDBReplicaNum, "0")
repo := NewTencentVectorDBRetrieveEngineRepository(nil, "", &types.IndexConfig{
ReplicaNumber: 2,
}).(*repository)
assert.Equal(t, 2, repo.replicasNum)
}
func TestCollectionNameUsesDimensionSuffixByDefault(t *testing.T) {
repo := NewTencentVectorDBRetrieveEngineRepository(nil, "", nil).(*repository)
assert.Equal(t, "weknora_embeddings_1024", repo.collectionName(1024))
assert.True(t, repo.matchesCollection("weknora_embeddings_1024"))
assert.False(t, repo.matchesCollection("weknora_embeddings"))
}
func TestCollectionNameRespectsExplicitCollectionName(t *testing.T) {
repo := NewTencentVectorDBRetrieveEngineRepository(nil, "", &types.IndexConfig{
CollectionName: "custom_collection",
}).(*repository)
assert.Equal(t, "custom_collection", repo.collectionName(1024))
assert.True(t, repo.matchesCollection("custom_collection"))
assert.False(t, repo.matchesCollection("custom_collection_1024"))
}
func TestCollectionAlreadyExistsErrorDetection(t *testing.T) {
assert.True(t, isCollectionAlreadyExistsErr(errors.New("code: 15202, collection already exists")))
assert.True(t, isCollectionAlreadyExistsErr(errors.New("Collection Already Exist")))
assert.False(t, isCollectionAlreadyExistsErr(errors.New("permission denied")))
assert.False(t, isCollectionAlreadyExistsErr(nil))
}
@@ -0,0 +1,57 @@
package tencentvectordb
import (
"sync"
"github.com/tencent/vectordatabase-sdk-go/tcvdbtext/encoder"
"github.com/tencent/vectordatabase-sdk-go/tcvectordb"
)
const (
envTencentVectorDBDatabase = "TENCENT_VECTORDB_DATABASE"
envTencentVectorDBCollection = "TENCENT_VECTORDB_COLLECTION"
envTencentVectorDBReplicaNum = "TENCENT_VECTORDB_REPLICA_NUMBER"
defaultDatabaseName = "weknora"
defaultCollectionName = "weknora_embeddings"
defaultReplicaNumber = 1
fieldID = "id"
fieldVector = "vector"
fieldSparseVector = "sparse_vector"
fieldContent = "content"
fieldSourceID = "source_id"
fieldSourceType = "source_type"
fieldChunkID = "chunk_id"
fieldKnowledgeID = "knowledge_id"
fieldKnowledgeBaseID = "knowledge_base_id"
fieldTagID = "tag_id"
fieldIsEnabled = "is_enabled"
)
type repository struct {
client *tcvectordb.RpcClient
databaseName string
collectionBaseName string
useDimensionSuffix bool
shardsNum int
replicasNum int
initialized sync.Map
bm25Once sync.Once
bm25 encoder.SparseEncoder
bm25Err error
}
type vectorEmbedding struct {
ID string
Content string
SourceID string
SourceType int
ChunkID string
KnowledgeID string
KnowledgeBaseID string
TagID string
Embedding []float32
SparseVector []encoder.SparseVecItem
IsEnabled bool
Score float64
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,33 @@
package weaviate
import (
"sync"
"github.com/weaviate/weaviate-go-client/v5/weaviate"
)
type weaviateRepository struct {
client *weaviate.Client
collectionBaseName string
replicationFactor int // 0 = use Weaviate server default
desiredShardCount int // 0 = use Weaviate server default
// Cache for initialized collections (dimension -> true)
initializedCollections sync.Map
}
type WeaviateVectorEmbedding struct {
Content string `json:"content"`
SourceID string `json:"source_id"`
SourceType int `json:"source_type"`
ChunkID string `json:"chunk_id"`
KnowledgeID string `json:"knowledge_id"`
KnowledgeBaseID string `json:"knowledge_base_id"`
TagID string `json:"tag_id"`
Embedding []float32 `json:"embedding"`
IsEnabled bool `json:"is_enabled"`
}
type WeaviateVectorEmbeddingWithScore struct {
WeaviateVectorEmbedding
Score float64
}