From 197b142cefd0a35464a8b5d613153034cbfdeb5a Mon Sep 17 00:00:00 2001 From: deadtrickster Date: Tue, 4 Aug 2026 08:16:39 +0200 Subject: [PATCH] feat(serenedb): add SereneDB doc-store engine (Go + Python connectors) (#17375) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## What Adds [**SereneDB**](https://serenedb.com) as a selectable doc-store engine on **both** RAGFlow paths: - the **Go** `DocEngine` (`internal/engine/serenedb`), alongside Elasticsearch and Infinity; - the **Python** `DocStoreConnection` (`rag/utils/serenedb_conn.py`) + `DOC_ENGINE=serenedb` registration. SereneDB is a PostgreSQL-wire engine (DuckDB execution) whose single inverted index carries **both** a scored text column (`@@`, BM25) and an IVF vector column (`<#>`, inner product), so hybrid search is one SQL statement. The Go engine connects with `database/sql` + `lib/pq` (already a dependency, no new module); the Python connector uses psycopg2 (already a dependency). ## Storage model One table per tenant with `kb_id` as a filter column - the **Elasticsearch / OceanBase** model, not Infinity's per-dataset tables. This keeps BM25 statistics (IDF, avgdl) computed over the whole tenant corpus (global IDF). Both connectors use this identical layout, so they are storage- and retrieval-compatible: `hybrid` proxy routing and Python↔Go switching are safe. On the Python side the connector is wired as OceanBase's plain-SQL sibling (chunk_data JSON metadata, inline chunk vectors, verbatim ES field names); the ES tokenizer path is unchanged. Metadata stays one table per tenant (`ragflow_doc_meta_`). The query shapes mirror the Python connector, including the five empirically-found landmines: the scored dictionary needs `frequency + norm` (else `BM25()` silently returns 0.0), the `@@` query is the tokenized query, the scored lexical branch matches one column, vectors use an L2-normalized shadow column with `ip`/`sq8`, and the similarity threshold goes directly in the ANN scan's `WHERE`. **Minimum engine version: SereneDB 26.07.4.** --------- Co-authored-by: Claude Opus 4.8 --- .github/workflows/serenedb.yml | 92 +++ api/db/services/dialog_service.py | 2 +- api/db/services/doc_metadata_service.py | 2 +- common/settings.py | 11 +- docker/.env | 8 + docker/docker-compose-base.yml | 23 + docker/service_conf.yaml.template | 6 + internal/engine/engine.go | 1 + internal/engine/global.go | 3 + internal/engine/serenedb/chunk.go | 476 ++++++++++++ internal/engine/serenedb/client.go | 220 ++++++ internal/engine/serenedb/common.go | 308 ++++++++ internal/engine/serenedb/document.go | 113 +++ internal/engine/serenedb/integration_test.go | 269 +++++++ internal/engine/serenedb/metadata.go | 272 +++++++ internal/engine/serenedb/schema.go | 167 ++++ internal/engine/serenedb/search.go | 587 ++++++++++++++ internal/engine/serenedb/serenedb.go | 34 + internal/engine/serenedb/serenedb_test.go | 401 ++++++++++ internal/engine/serenedb/sql.go | 58 ++ internal/server/config/doc_engine_config.go | 59 ++ rag/app/table.py | 8 +- rag/nlp/search.py | 4 +- rag/utils/serenedb_conn.py | 719 ++++++++++++++++++ rag/utils/table_es_metadata.py | 6 +- .../unit_test/rag/utils/test_serenedb_conn.py | 96 +++ 26 files changed, 3933 insertions(+), 12 deletions(-) create mode 100644 .github/workflows/serenedb.yml create mode 100644 internal/engine/serenedb/chunk.go create mode 100644 internal/engine/serenedb/client.go create mode 100644 internal/engine/serenedb/common.go create mode 100644 internal/engine/serenedb/document.go create mode 100644 internal/engine/serenedb/integration_test.go create mode 100644 internal/engine/serenedb/metadata.go create mode 100644 internal/engine/serenedb/schema.go create mode 100644 internal/engine/serenedb/search.go create mode 100644 internal/engine/serenedb/serenedb.go create mode 100644 internal/engine/serenedb/serenedb_test.go create mode 100644 internal/engine/serenedb/sql.go create mode 100644 rag/utils/serenedb_conn.py create mode 100644 test/unit_test/rag/utils/test_serenedb_conn.py diff --git a/.github/workflows/serenedb.yml b/.github/workflows/serenedb.yml new file mode 100644 index 0000000000..bbbe6f342a --- /dev/null +++ b/.github/workflows/serenedb.yml @@ -0,0 +1,92 @@ +name: serenedb + +# SereneDB has no dedicated CI capacity upstream, so this workflow never runs on +# push or pull_request - it only runs when triggered by hand, from the Actions +# tab (workflow_dispatch) or locally with the nektos/act emulator. +# +# It brings up a real SereneDB and runs the Go engine's integration tier against +# it (the unit tier and the Python connector unit tests already run in the normal +# test workflows; those need no live SereneDB). +# +# The steps use docker-in-docker plus the repo's CGO toolchain, so run act in +# host mode - map the self-hosted labels to act's "-self-hosted" so the steps run +# on your host (which has the toolchain + docker), and let act check out into its +# own workspace (do not pass --bind, or the nested build container mounts an empty +# dir). Override host_port if 7890 is taken locally. Verified with: +# +# act workflow_dispatch -W .github/workflows/serenedb.yml \ +# -P self-hosted=-self-hosted -P ragflow-test=-self-hosted \ +# --input host_port=17890 + +on: + workflow_dispatch: + inputs: + serenedb_image: + description: SereneDB image to test against + required: false + default: serenedb/serenedb:26.07.5 + host_port: + description: Host port to publish SereneDB on (override to avoid a local clash) + required: false + default: "7890" + +permissions: + contents: read + +jobs: + serenedb_go_integration: + name: serenedb_go_integration + runs-on: [ "self-hosted", "ragflow-test" ] + env: + SERENEDB_IMAGE: ${{ github.event.inputs.serenedb_image || 'serenedb/serenedb:26.07.5' }} + SERENEDB_CONTAINER: serenedb-ci-${{ github.run_id }} + SERENEDB_HOST_PORT: ${{ github.event.inputs.host_port || '7890' }} + steps: + - name: Check out code + uses: actions/checkout@v6 + with: + fetch-depth: 0 + + - name: Start SereneDB + run: | + set -euo pipefail + docker rm -f "${SERENEDB_CONTAINER}" >/dev/null 2>&1 || true + docker run -d --name "${SERENEDB_CONTAINER}" \ + -p "127.0.0.1:${SERENEDB_HOST_PORT}:7890" \ + -e POSTGRES_PASSWORD=infini_rag_flow \ + "${SERENEDB_IMAGE}" + for i in $(seq 1 60); do + if docker exec "${SERENEDB_CONTAINER}" pg_isready -h 127.0.0.1 -p 7890 -U postgres >/dev/null 2>&1; then + echo "SereneDB is accepting connections" + break + fi + echo "Waiting for SereneDB... ($i/60)" + sleep 5 + done + + - name: Build native tokenizer library + run: | + set -euo pipefail + BUILDER_CONTAINER=serenedb_build_${GITHUB_RUN_ID}_$(od -An -N4 -tx4 /dev/urandom | tr -d ' ') + cleanup_builder() { + docker rm -f -v "${BUILDER_CONTAINER}" >/dev/null 2>&1 || true + } + trap cleanup_builder EXIT + docker run --privileged -d --name "${BUILDER_CONTAINER}" \ + -v "${PWD}:/ragflow" \ + -v "${PWD}/internal/binding/cpp/resource:/usr/share/infinity/resource" \ + infiniflow/infinity_builder:ubuntu22_clang20 + docker exec "${BUILDER_CONTAINER}" bash -c 'git config --global safe.directory "*" && cd /ragflow && ./build.sh --cpp' + # The builder runs as root; hand the emitted artifacts back to the runner + # user so the workspace stays cleanable on the next checkout/cleanup. + docker exec "${BUILDER_CONTAINER}" chown -R "$(id -u):$(id -g)" /ragflow/internal/binding/cpp + + - name: Run Go integration tests against SereneDB + run: | + set -euo pipefail + export SERENEDB_TEST_DSN="host=127.0.0.1 port=${SERENEDB_HOST_PORT} user=postgres password=infini_rag_flow dbname=postgres sslmode=disable" + ./build.sh --test-integration ./internal/engine/serenedb/... + + - name: Stop SereneDB + if: always() + run: docker rm -f "${SERENEDB_CONTAINER}" >/dev/null 2>&1 || true diff --git a/api/db/services/dialog_service.py b/api/db/services/dialog_service.py index f385521d69..47e65b2a8f 100644 --- a/api/db/services/dialog_service.py +++ b/api/db/services/dialog_service.py @@ -71,7 +71,7 @@ async def _hydrate_chunk_vectors(retriever, chunks, tenant_ids, kb_ids): search results) keep whatever placeholder they were given. Other backends still carry vectors in the chunk, so we skip the round-trip. """ - if settings.DOC_ENGINE_INFINITY or settings.DOC_ENGINE_OCEANBASE: + if settings.DOC_ENGINE_INFINITY or settings.DOC_ENGINE_OCEANBASE or settings.DOC_ENGINE_SERENEDB: return if not chunks: return diff --git a/api/db/services/doc_metadata_service.py b/api/db/services/doc_metadata_service.py index 4c8300a691..0b38643a36 100644 --- a/api/db/services/doc_metadata_service.py +++ b/api/db/services/doc_metadata_service.py @@ -470,7 +470,7 @@ class DocMetadataService: logging.debug(f"[update_document_metadata] Updating doc_id: {doc_id}, kb_id: {kb_id}, meta_fields: {processed_meta}") # For Elasticsearch, use efficient partial update - if not settings.DOC_ENGINE_INFINITY and not settings.DOC_ENGINE_OCEANBASE: + if not settings.DOC_ENGINE_INFINITY and not settings.DOC_ENGINE_OCEANBASE and not settings.DOC_ENGINE_SERENEDB: # Check if index exists first index_exists = settings.docStoreConn.index_exist(index_name, "") if not index_exists: diff --git a/common/settings.py b/common/settings.py index 0bccddfceb..2eff4a7be7 100644 --- a/common/settings.py +++ b/common/settings.py @@ -85,6 +85,7 @@ OAUTH_CONFIG = None DOC_ENGINE = os.getenv("DOC_ENGINE", "elasticsearch") DOC_ENGINE_INFINITY = DOC_ENGINE.lower() == "infinity" DOC_ENGINE_OCEANBASE = DOC_ENGINE.lower() == "oceanbase" +DOC_ENGINE_SERENEDB = DOC_ENGINE.lower() == "serenedb" docStoreConn = None @@ -123,6 +124,7 @@ OB = {} OSS = {} OS = {} GCS = {} +SERENEDB = {} DOC_MAXIMUM_SIZE: int = 128 * 1024 * 1024 DOC_BULK_SIZE: int = 32 @@ -301,10 +303,11 @@ def init_settings(): FEISHU_OAUTH = get_base_config("oauth", {}).get("feishu") OAUTH_CONFIG = get_base_config("oauth", {}) - global DOC_ENGINE, DOC_ENGINE_INFINITY, DOC_ENGINE_OCEANBASE, docStoreConn, ES, OB, OS, INFINITY + global DOC_ENGINE, DOC_ENGINE_INFINITY, DOC_ENGINE_OCEANBASE, DOC_ENGINE_SERENEDB, docStoreConn, ES, OB, OS, INFINITY, SERENEDB DOC_ENGINE = os.environ.get("DOC_ENGINE", "elasticsearch").strip() DOC_ENGINE_INFINITY = DOC_ENGINE.lower() == "infinity" DOC_ENGINE_OCEANBASE = DOC_ENGINE.lower() == "oceanbase" + DOC_ENGINE_SERENEDB = DOC_ENGINE.lower() == "serenedb" lower_case_doc_engine = DOC_ENGINE.lower() if lower_case_doc_engine == "elasticsearch": ES = get_base_config("es", {}) @@ -321,6 +324,12 @@ def init_settings(): elif lower_case_doc_engine == "seekdb": OB = get_base_config("seekdb", {}) docStoreConn = rag.utils.ob_conn.OBConnection() + elif lower_case_doc_engine == "serenedb": + SERENEDB = get_base_config("serenedb", {}) + # Imported lazily so psycopg2/SereneDB is only touched when selected. + from rag.utils import serenedb_conn + + docStoreConn = serenedb_conn.SereneDBConnection() else: raise Exception(f"Not supported doc engine: {DOC_ENGINE}") diff --git a/docker/.env b/docker/.env index 008398d4a1..0940323198 100644 --- a/docker/.env +++ b/docker/.env @@ -72,6 +72,14 @@ INFINITY_THRIFT_PORT=23817 INFINITY_HTTP_PORT=23820 INFINITY_PSQL_PORT=5432 +# The hostname where the SereneDB service is exposed. SereneDB speaks the +# PostgreSQL wire protocol, so DOC_ENGINE=serenedb connects over lib/pq/psycopg2. +SERENEDB_HOST=serenedb +# Port to expose SereneDB to the host +SERENEDB_PORT=7890 +# The password for SereneDB (POSTGRES_PASSWORD in the container) +SERENEDB_PASSWORD=infini_rag_flow + # The hostname where the OceanBase service is exposed OCEANBASE_HOST=oceanbase # The port used to expose the OceanBase service diff --git a/docker/docker-compose-base.yml b/docker/docker-compose-base.yml index 286294f823..835a521518 100644 --- a/docker/docker-compose-base.yml +++ b/docker/docker-compose-base.yml @@ -100,6 +100,27 @@ services: retries: 120 restart: unless-stopped + serenedb: + profiles: + - serenedb + image: serenedb/serenedb:26.07.5 + env_file: .env + environment: + - POSTGRES_PASSWORD=${SERENEDB_PASSWORD} + ports: + - ${SERENEDB_PORT}:7890 + volumes: + - serenedb_data:/var/lib/serenedb + mem_limit: ${MEM_LIMIT} + networks: + - ragflow + healthcheck: + test: ["CMD", "pg_isready", "-h", "127.0.0.1", "-p", "7890", "-U", "postgres"] + interval: 10s + timeout: 10s + retries: 120 + restart: unless-stopped + oceanbase: profiles: - oceanbase @@ -380,6 +401,8 @@ volumes: driver: local infinity_data: driver: local + serenedb_data: + driver: local ob_data: driver: local seekdb_data: diff --git a/docker/service_conf.yaml.template b/docker/service_conf.yaml.template index c6ba55fd22..5b77399425 100644 --- a/docker/service_conf.yaml.template +++ b/docker/service_conf.yaml.template @@ -41,6 +41,12 @@ infinity: uri: '${INFINITY_HOST:-infinity}:23817' postgres_port: 5432 db_name: 'default_db' +serenedb: + host: '${SERENEDB_HOST:-serenedb}' + port: 7890 + user: '${SERENEDB_USER:-postgres}' + password: '${SERENEDB_PASSWORD:-infini_rag_flow}' + db_name: '${SERENEDB_DBNAME:-postgres}' oceanbase: scheme: 'oceanbase' # set 'mysql' to create connection using mysql config config: diff --git a/internal/engine/engine.go b/internal/engine/engine.go index 24c5f80e78..c9a903d535 100644 --- a/internal/engine/engine.go +++ b/internal/engine/engine.go @@ -32,6 +32,7 @@ type EngineType string const ( EngineElasticsearch EngineType = "elasticsearch" EngineInfinity EngineType = "infinity" + EngineSereneDB EngineType = "serenedb" ) // DocEngine document storage engine interface diff --git a/internal/engine/global.go b/internal/engine/global.go index 9f4564a118..af6ab046d3 100644 --- a/internal/engine/global.go +++ b/internal/engine/global.go @@ -25,6 +25,7 @@ import ( "ragflow/internal/engine/elasticsearch" "ragflow/internal/engine/infinity" + "ragflow/internal/engine/serenedb" "ragflow/internal/tokenizer" @@ -52,6 +53,8 @@ func InitDocEngine() error { globalEngine, err = elasticsearch.NewEngine(globalConfig.GetElasticsearchConfig()) case "infinity": globalEngine, err = infinity.NewEngine(globalConfig.GetInfinityConfig()) + case "serenedb": + globalEngine, err = serenedb.NewEngine(globalConfig.GetSereneDBConfig()) default: err = fmt.Errorf("unsupported doc engine type: %s", engineType) } diff --git a/internal/engine/serenedb/chunk.go b/internal/engine/serenedb/chunk.go new file mode 100644 index 0000000000..9f82255c28 --- /dev/null +++ b/internal/engine/serenedb/chunk.go @@ -0,0 +1,476 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package serenedb + +import ( + "context" + "encoding/json" + "fmt" + "sort" + "strconv" + "strings" + + "ragflow/internal/common" + + "github.com/lib/pq" + "go.uber.org/zap" +) + +const defaultVectorSize = 1024 + +// chunkTableDDL returns the statements that create a chunk table, its scored +// dictionary, and the hybrid inverted index (text columns + normalized vector +// shadow). Pure so it can be asserted in tests. +func chunkTableDDL(tableName string, vectorSize int) []string { + if vectorSize <= 0 { + vectorSize = defaultVectorSize + } + cols := make([]string, 0, len(columnOrder)+2) + for _, c := range columnOrder { + cols = append(cols, fmt.Sprintf("%s %s", c, columnDDL[c])) + } + vec, vecN := rawVectorColumn(vectorSize), normColumn(vectorSize) + cols = append(cols, + fmt.Sprintf("%s FLOAT[%d]", vec, vectorSize), + fmt.Sprintf("%s FLOAT[%d]", vecN, vectorSize)) + + fts := make([]string, 0, len(ftsColumns)) + for _, c := range ftsColumns { + fts = append(fts, fmt.Sprintf("%s %s", c, dictionaryName)) + } + return []string{ + fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s (%s)", tableName, strings.Join(cols, ", ")), + dictionaryDDL, + fmt.Sprintf("CREATE INDEX IF NOT EXISTS %s ON %s "+ + "USING inverted (id, %s, %s ivf (metric = 'ip', quant = 'sq8')) "+ + "WITH (optimize_top_k = 'bm25(1.2, 0.75)')", + indexRelation(tableName), tableName, strings.Join(fts, ", "), vecN), + } +} + +// CreateChunkStore ensures the tenant chunk table and its indexes exist. All +// datasets in the tenant share this table, so it is created once and is a +// no-op for later datasets. +func (e *serenedbEngine) CreateChunkStore(ctx context.Context, baseName, datasetID string, vectorSize int, parserID string) error { + tableName := chunkTableName(baseName) + for _, stmt := range chunkTableDDL(tableName, vectorSize) { + if err := e.exec(ctx, stmt); err != nil { + return fmt.Errorf("serenedb: create chunk store %s: %w", tableName, err) + } + } + common.Info("SereneDB created chunk store", zap.String("table", tableName)) + return nil +} + +// DropChunkStore removes a dataset from the tenant table. Because the table is +// shared, a dataset drop deletes that dataset's rows; only a whole-tenant drop +// (empty datasetID) drops the table itself. +func (e *serenedbEngine) DropChunkStore(ctx context.Context, baseName, datasetID string) error { + tableName := chunkTableName(baseName) + exists, err := e.tableExists(ctx, tableName) + if err != nil { + return err + } + if !exists { + return nil + } + if datasetID != "" { + return e.exec(ctx, fmt.Sprintf("DELETE FROM %s WHERE kb_id = $1", tableName), datasetID) + } + if err := e.exec(ctx, fmt.Sprintf("DROP INDEX IF EXISTS %s", indexRelation(tableName))); err != nil { + return err + } + return e.exec(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %s", tableName)) +} + +// ChunkStoreExists reports whether the tenant table exists. +func (e *serenedbEngine) ChunkStoreExists(ctx context.Context, baseName, datasetID string) (bool, error) { + return e.tableExists(ctx, chunkTableName(baseName)) +} + +// toFloatSlice coerces a chunk vector value to []float64. +func toFloatSlice(v interface{}) ([]float64, bool) { + switch val := v.(type) { + case []float64: + return val, true + case []float32: + out := make([]float64, len(val)) + for i, f := range val { + out[i] = float64(f) + } + return out, true + case []interface{}: + out := make([]float64, 0, len(val)) + for _, x := range val { + switch n := x.(type) { + case float64: + out = append(out, n) + case float32: + out = append(out, float64(n)) + case int: + out = append(out, float64(n)) + case json.Number: + f, _ := n.Float64() + out = append(out, f) + default: + return nil, false + } + } + return out, true + default: + return nil, false + } +} + +// prepareChunkRow flattens a chunk into ordered (columns, values) for insert. +// Vector columns get their L2-normalized shadow; unknown fields collapse into +// the `extra` JSON column; JSON columns are serialized. ES field names are +// preserved verbatim. +func prepareChunkRow(chunk map[string]interface{}, defaultKbID string) ([]string, []interface{}) { + d := map[string]interface{}{} + extra := map[string]interface{}{} + vecCols := map[string][]float64{} + + for k, v := range chunk { + if m := vectorColumnPattern.FindStringSubmatch(k); m != nil { + if vec, ok := toFloatSlice(v); ok { + vecCols[k] = vec + } + continue + } + if _, known := columnDDL[k]; !known { + extra[k] = v + continue + } + if k == "kb_id" { + if list, ok := v.([]interface{}); ok && len(list) > 0 { + v = list[0] + } else if list, ok := v.([]string); ok && len(list) > 0 { + v = list[0] + } + } + if isJSONColumn(k) { + if _, ok := v.(string); !ok { + b, _ := json.Marshal(v) + v = string(b) + } + } + d[k] = v + } + + if len(extra) > 0 { + merged := map[string]interface{}{} + if s, ok := d["extra"].(string); ok && s != "" { + _ = json.Unmarshal([]byte(s), &merged) + } + for k, v := range extra { + merged[k] = v + } + b, _ := json.Marshal(merged) + d["extra"] = string(b) + } + // The table is shared across datasets, so kb_id identifies the row's + // dataset. Honour an explicit kb_id, otherwise stamp the target dataset. + if defaultKbID != "" { + if _, ok := d["kb_id"]; !ok { + d["kb_id"] = defaultKbID + } + } + for k, dv := range columnDefaults { + if _, ok := d[k]; !ok { + d[k] = dv + } + } + + cols := make([]string, 0, len(d)+2*len(vecCols)) + for c := range d { + cols = append(cols, c) + } + sort.Strings(cols) + vals := make([]interface{}, 0, len(cols)+2*len(vecCols)) + for _, c := range cols { + if isArrayColumn(c) { + // VARCHAR[] columns must be bound through pq.Array; database/sql + // cannot convert a bare Go slice. + vals = append(vals, pq.Array(toStringArray(d[c]))) + } else { + vals = append(vals, d[c]) + } + } + vnames := make([]string, 0, len(vecCols)) + for vc := range vecCols { + vnames = append(vnames, vc) + } + sort.Strings(vnames) + for _, vc := range vnames { + size, _ := strconv.Atoi(vectorColumnPattern.FindStringSubmatch(vc)[1]) + cols = append(cols, vc, normColumn(size)) + vals = append(vals, pq.Array(vecCols[vc]), pq.Array(l2Normalize(vecCols[vc]))) + } + return cols, vals +} + +// InsertChunks upserts chunks by id. Rows are grouped by identical column set +// and inserted in one multi-row statement per group. Returns an empty slice on +// success (ids are the caller's own). +func (e *serenedbEngine) InsertChunks(ctx context.Context, chunks []map[string]interface{}, baseName, datasetID string) ([]string, error) { + if len(chunks) == 0 { + return []string{}, nil + } + tableName := chunkTableName(baseName) + exists, err := e.tableExists(ctx, tableName) + if err != nil { + return nil, err + } + if !exists { + size := 0 + for k := range chunks[0] { + if m := vectorColumnPattern.FindStringSubmatch(k); m != nil { + size, _ = strconv.Atoi(m[1]) + } + } + if err := e.CreateChunkStore(ctx, baseName, datasetID, size, ""); err != nil { + return nil, err + } + } + + type group struct { + cols []string + rows [][]interface{} + } + groups := map[string]*group{} + for _, chunk := range chunks { + cols, vals := prepareChunkRow(chunk, datasetID) + key := strings.Join(cols, ",") + g := groups[key] + if g == nil { + g = &group{cols: cols} + groups[key] = g + } + g.rows = append(g.rows, vals) + } + + for _, g := range groups { + query, args := buildUpsert(tableName, g.cols, g.rows) + if err := e.exec(ctx, query, args...); err != nil { + return nil, fmt.Errorf("serenedb: insert into %s: %w", tableName, err) + } + } + return []string{}, nil +} + +// buildUpsert renders a multi-row INSERT ... ON CONFLICT (id) DO UPDATE for one +// column set. Pure so it can be asserted in tests. +func buildUpsert(tableName string, cols []string, rows [][]interface{}) (string, []interface{}) { + var args []interface{} + valueGroups := make([]string, 0, len(rows)) + ph := 1 + for _, row := range rows { + placeholders := make([]string, len(cols)) + for i := range cols { + placeholders[i] = "$" + strconv.Itoa(ph) + ph++ + args = append(args, row[i]) + } + valueGroups = append(valueGroups, "("+strings.Join(placeholders, ", ")+")") + } + updates := make([]string, 0, len(cols)) + for _, c := range cols { + if c == "id" { + continue + } + updates = append(updates, fmt.Sprintf("%s = EXCLUDED.%s", c, c)) + } + query := fmt.Sprintf("INSERT INTO %s (%s) VALUES %s ON CONFLICT (id) DO UPDATE SET %s", + tableName, strings.Join(cols, ", "), strings.Join(valueGroups, ", "), strings.Join(updates, ", ")) + return query, args +} + +// UpdateChunks applies field updates to rows matching condition. add/remove on +// array columns map to list_append/array_remove; other fields are set directly. +func (e *serenedbEngine) UpdateChunks(ctx context.Context, condition, newValue map[string]interface{}, baseName, datasetID string) error { + tableName := chunkTableName(baseName) + exists, err := e.tableExists(ctx, tableName) + if err != nil { + return err + } + if !exists { + return fmt.Errorf("serenedb: table %s does not exist", tableName) + } + cond := copyCondition(condition) + cond["kb_id"] = datasetID + if unknown := unrecognizedFilterKeys(cond); len(unknown) > 0 { + return fmt.Errorf("serenedb: refusing to update %s with unrecognized filter keys %v", tableName, unknown) + } + filters := buildFilters(cond) + if len(filters) == 0 { + return fmt.Errorf("serenedb: refusing to update %s without a filter", tableName) + } + sets := buildUpdateSets(newValue) + if len(sets) == 0 { + return nil + } + query := fmt.Sprintf("UPDATE %s SET %s WHERE %s", tableName, strings.Join(sets, ", "), strings.Join(filters, " AND ")) + return e.exec(ctx, query) +} + +// buildUpdateSets renders the SET clause fragments. Pure for tests. +func buildUpdateSets(newValue map[string]interface{}) []string { + var sets []string + for k, v := range newValue { + switch k { + case "remove": + items := map[string]interface{}{} + if s, ok := v.(string); ok { + items[s] = nil + } else if m, ok := v.(map[string]interface{}); ok { + items = m + } + for kk, vv := range items { + if _, known := columnDDL[kk]; !known { + continue + } + if vv == nil { + sets = append(sets, fmt.Sprintf("%s = NULL", kk)) + } else if isArrayColumn(kk) { + sets = append(sets, fmt.Sprintf("%s = array_remove(%s, %s)", kk, kk, escapeLiteral(vv))) + } + } + case "add": + if m, ok := v.(map[string]interface{}); ok { + for kk, vv := range m { + if isArrayColumn(kk) { + sets = append(sets, fmt.Sprintf("%s = list_append(%s, %s)", kk, kk, escapeLiteral(vv))) + } + } + } + default: + if isJSONColumn(k) { + if s, ok := v.(string); ok { + sets = append(sets, fmt.Sprintf("%s = %s", k, escapeLiteral(s))) + } else { + b, _ := json.Marshal(v) + sets = append(sets, fmt.Sprintf("%s = %s", k, escapeLiteral(string(b)))) + } + } else if _, known := columnDDL[k]; known { + sets = append(sets, fmt.Sprintf("%s = %s", k, escapeLiteral(v))) + } + } + } + sort.Strings(sets) + return sets +} + +// DeleteChunks removes rows matching condition and returns the count. A missing +// table is not an error. +func (e *serenedbEngine) DeleteChunks(ctx context.Context, condition map[string]interface{}, baseName, datasetID string) (int64, error) { + tableName := chunkTableName(baseName) + exists, err := e.tableExists(ctx, tableName) + if err != nil { + return 0, err + } + if !exists { + return 0, nil + } + cond := copyCondition(condition) + cond["kb_id"] = datasetID + if unknown := unrecognizedFilterKeys(cond); len(unknown) > 0 { + return 0, fmt.Errorf("serenedb: refusing to delete from %s with unrecognized filter keys %v", tableName, unknown) + } + filters := buildFilters(cond) + if len(filters) == 0 { + return 0, nil + } + where := strings.Join(filters, " AND ") + res, err := e.db.ExecContext(ctx, fmt.Sprintf("DELETE FROM %s WHERE %s", tableName, where)) + if err != nil { + return 0, err + } + n, _ := res.RowsAffected() + return n, nil +} + +// GetChunk looks a chunk up by id in the tenant table, optionally scoped to +// the caller's datasets. +func (e *serenedbEngine) GetChunk(ctx context.Context, baseName, chunkID string, datasetIDs []string) (interface{}, error) { + tableName := chunkTableName(baseName) + exists, err := e.tableExists(ctx, tableName) + if err != nil { + return nil, err + } + if !exists { + return nil, nil + } + query := fmt.Sprintf("SELECT * FROM %s WHERE id = $1", tableName) + if kbs := stringSlice(datasetIDs); len(kbs) > 0 { + if kb := buildFilters(map[string]interface{}{"kb_id": kbs}); len(kb) > 0 { + query += " AND " + kb[0] + } + } + rows, err := e.queryMaps(ctx, query, chunkID) + if err != nil { + return nil, err + } + if len(rows) == 0 { + return nil, nil + } + return rows[0], nil +} + +// stringSlice drops empty ids so an all-blank dataset list yields no kb filter. +func stringSlice(ids []string) []string { + out := make([]string, 0, len(ids)) + for _, id := range ids { + if id != "" { + out = append(out, id) + } + } + return out +} + +// toStringArray coerces an array-column value into the []string that pq.Array +// binds as a VARCHAR[]. +func toStringArray(v interface{}) []string { + switch val := v.(type) { + case []string: + return val + case []interface{}: + out := make([]string, 0, len(val)) + for _, x := range val { + if s, ok := x.(string); ok { + out = append(out, s) + } else { + out = append(out, fmt.Sprintf("%v", x)) + } + } + return out + case nil: + return nil + case string: + return []string{val} + default: + return []string{fmt.Sprintf("%v", val)} + } +} + +func copyCondition(condition map[string]interface{}) map[string]interface{} { + out := make(map[string]interface{}, len(condition)+1) + for k, v := range condition { + out[k] = v + } + return out +} diff --git a/internal/engine/serenedb/client.go b/internal/engine/serenedb/client.go new file mode 100644 index 0000000000..4a695142c1 --- /dev/null +++ b/internal/engine/serenedb/client.go @@ -0,0 +1,220 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package serenedb + +import ( + "context" + "database/sql" + "fmt" + "os" + "regexp" + "strings" + "time" + + "ragflow/internal/common" + "ragflow/internal/server/config" + + _ "github.com/lib/pq" + "go.uber.org/zap" +) + +const ( + defaultHost = "serenedb" + defaultPort = 7890 + defaultUser = "postgres" + defaultDBName = "postgres" + defaultSSLMode = "disable" + poolMaxOpen = 8 + poolMaxIdle = 2 + connMaxLifetime = 30 * time.Minute +) + +var ( + dsnKeywordPwRe = regexp.MustCompile(`(password=)('(?:[^'\\]|\\.)*'|\S+)`) + dsnURLPwRe = regexp.MustCompile(`(://[^:/@\s]+:)[^@\s]+(@)`) +) + +// redactDSN hides the password in both the lib/pq keyword form (password=...) +// and the URL form (scheme://user:password@host) used by SERENEDB_DSN. +func redactDSN(dsn string) string { + dsn = dsnKeywordPwRe.ReplaceAllString(dsn, "${1}***") + dsn = dsnURLPwRe.ReplaceAllString(dsn, "${1}***${2}") + return dsn +} + +// quoteDSNValue wraps a lib/pq keyword-DSN value in single quotes, escaping +// backslashes and quotes, so hosts/users/passwords with spaces or special +// characters do not break the DSN. +func quoteDSNValue(v string) string { + r := strings.ReplaceAll(v, `\`, `\\`) + r = strings.ReplaceAll(r, `'`, `\'`) + return "'" + r + "'" +} + +// serenedbEngine implements engine.DocEngine backed by SereneDB. +type serenedbEngine struct { + db *sql.DB + dsnSafe string +} + +// NewEngine constructs the engine from the SereneDB config, mirroring the +// elasticsearch/infinity factories in engine/global.go. +func NewEngine(cfg config.SereneDBConfig) (*serenedbEngine, error) { + dsn := buildDSN(cfg) + + db, err := sql.Open("postgres", dsn) + if err != nil { + return nil, fmt.Errorf("serenedb: open: %w", err) + } + db.SetMaxOpenConns(poolMaxOpen) + db.SetMaxIdleConns(poolMaxIdle) + db.SetConnMaxLifetime(connMaxLifetime) + + e := &serenedbEngine{db: db, dsnSafe: redactDSN(dsn)} + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + if err := db.PingContext(ctx); err != nil { + _ = db.Close() + return nil, fmt.Errorf("serenedb: ping %s: %w", e.dsnSafe, err) + } + common.Info("SereneDB engine initialized", zap.String("dsn", e.dsnSafe)) + return e, nil +} + +// buildDSN assembles a lib/pq keyword DSN. SERENEDB_DSN overrides everything +// (matching the Python connector); otherwise config values fill in, then the +// documented defaults. +func buildDSN(cfg config.SereneDBConfig) string { + if env := os.Getenv("SERENEDB_DSN"); env != "" { + return env + } + host, port, user, password, dbName := defaultHost, defaultPort, defaultUser, "", defaultDBName + sslMode := defaultSSLMode + if cfg.Host != "" { + host = cfg.Host + } + if cfg.Port != 0 { + port = cfg.Port + } + if cfg.User != "" { + user = cfg.User + } + password = cfg.Password + if cfg.DBName != "" { + dbName = cfg.DBName + } + if cfg.SSLMode != "" { + sslMode = cfg.SSLMode + } + parts := []string{ + fmt.Sprintf("host=%s", quoteDSNValue(host)), + fmt.Sprintf("port=%d", port), + fmt.Sprintf("user=%s", quoteDSNValue(user)), + fmt.Sprintf("dbname=%s", quoteDSNValue(dbName)), + fmt.Sprintf("sslmode=%s", quoteDSNValue(sslMode)), + } + if password != "" { + parts = append(parts, fmt.Sprintf("password=%s", quoteDSNValue(password))) + } + return strings.Join(parts, " ") +} + +// GetType returns the engine type string. +func (e *serenedbEngine) GetType() string { + return "serenedb" +} + +// SupportsPageRank reports dataset-level pagerank support. Like Elasticsearch, +// SereneDB folds pagerank_fea into every scored query and stores it in a real +// column that UpdateChunks can set, so the dataset-level toggle is supported. +func (e *serenedbEngine) SupportsPageRank() bool { + return true +} + +// Ping verifies connectivity. +func (e *serenedbEngine) Ping(ctx context.Context) error { + return e.db.PingContext(ctx) +} + +// Close releases the connection pool. +func (e *serenedbEngine) Close() error { + if e.db == nil { + return nil + } + return e.db.Close() +} + +// exec runs a statement that returns no rows. +func (e *serenedbEngine) exec(ctx context.Context, query string, args ...interface{}) error { + _, err := e.db.ExecContext(ctx, query, args...) + return err +} + +// queryMaps runs a query and shapes each row into a map keyed by column name, +// decoding JSON and array columns to structured values. +func (e *serenedbEngine) queryMaps(ctx context.Context, query string, args ...interface{}) ([]map[string]interface{}, error) { + rows, err := e.db.QueryContext(ctx, query, args...) + if err != nil { + return nil, err + } + defer rows.Close() + + cols, err := rows.Columns() + if err != nil { + return nil, err + } + var out []map[string]interface{} + for rows.Next() { + vals := make([]interface{}, len(cols)) + ptrs := make([]interface{}, len(cols)) + for i := range vals { + ptrs[i] = &vals[i] + } + if err := rows.Scan(ptrs...); err != nil { + return nil, err + } + entity := make(map[string]interface{}, len(cols)) + for i, col := range cols { + v := decodeValue(col, vals[i]) + if v == nil { + continue + } + entity[col] = v + } + out = append(out, entity) + } + return out, rows.Err() +} + +// tableExists reports whether a relation exists. It validates the identifier +// (table names cannot be parameterized) and queries the catalog, so a missing +// table returns (false, nil) while a connectivity or permission failure is +// surfaced as an error instead of being masked as "absent". +func (e *serenedbEngine) tableExists(ctx context.Context, tableName string) (bool, error) { + if !validIdentifier(tableName) { + return false, fmt.Errorf("serenedb: invalid table name %q", tableName) + } + rows, err := e.db.QueryContext(ctx, + "SELECT 1 FROM information_schema.tables WHERE table_name = $1 AND table_schema = current_schema() LIMIT 1", + tableName) + if err != nil { + return false, err + } + defer rows.Close() + return rows.Next(), rows.Err() +} diff --git a/internal/engine/serenedb/common.go b/internal/engine/serenedb/common.go new file mode 100644 index 0000000000..d504c1dcdb --- /dev/null +++ b/internal/engine/serenedb/common.go @@ -0,0 +1,308 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package serenedb + +import ( + "encoding/json" + "fmt" + "math" + "regexp" + "strconv" + "strings" +) + +var esBoostRe = regexp.MustCompile(`\^[0-9.]+`) +var esSyntaxRe = regexp.MustCompile(`["()~*?:+\-]|\bAND\b|\bOR\b|\bNOT\b`) + +// escapeLiteral renders a Go value as a SQL literal for the templated search +// statements (filters, aggregation). Parameterized placeholders are used on +// the write path; the read path templates values that are engine-internal. +func escapeLiteral(v interface{}) string { + switch val := v.(type) { + case nil: + return "NULL" + case bool: + if val { + return "true" + } + return "false" + case int: + return strconv.Itoa(val) + case int64: + return strconv.FormatInt(val, 10) + case float64: + return strconv.FormatFloat(val, 'f', -1, 64) + case string: + return "'" + strings.ReplaceAll(val, "'", "''") + "'" + case []string, []interface{}, map[string]interface{}: + b, _ := json.Marshal(val) + return "'" + strings.ReplaceAll(string(b), "'", "''") + "'" + default: + return "'" + strings.ReplaceAll(fmt.Sprintf("%v", val), "'", "''") + "'" + } +} + +// asString coerces a filter value to its scalar string form when possible. +func asString(v interface{}) (string, bool) { + switch val := v.(type) { + case string: + return val, true + default: + return "", false + } +} + +// toStringSlice normalizes a filter value into a slice of scalars. +func toStringSlice(v interface{}) []interface{} { + switch val := v.(type) { + case []interface{}: + return val + case []string: + out := make([]interface{}, len(val)) + for i, s := range val { + out[i] = s + } + return out + default: + return []interface{}{v} + } +} + +// buildFilters translates a RAGFlow condition map into SQL predicates. kb_id is +// scoped by table name (like the other engines), so callers strip it before +// calling. Array columns filter with list_contains; exists/must_not map to +// IS NULL / IS NOT NULL. +// neverMatch is a predicate that matches no rows. An empty IN list (or empty +// array-contains) reduces to it, and it keeps a DELETE/UPDATE from widening +// scope when the caller asked to match "nothing". +const neverMatch = "1 = 0" + +// recognizedFilterKey reports whether buildFilters emits a predicate for a key. +// Callers that must not silently widen scope (DeleteChunks/UpdateChunks/ +// DeleteMetadata) reject conditions carrying unrecognized keys. +func recognizedFilterKey(k string) bool { + return k == "exists" || k == "must_not" || isArrayColumn(k) || isKnownColumn(k) +} + +// unrecognizedFilterKeys returns any condition keys buildFilters would drop. +func unrecognizedFilterKeys(condition map[string]interface{}) []string { + var unknown []string + for k := range condition { + if !recognizedFilterKey(k) { + unknown = append(unknown, k) + } + } + return unknown +} + +func inList(column string, vals []interface{}) string { + if len(vals) == 0 { + return neverMatch + } + parts := make([]string, 0, len(vals)) + for _, x := range vals { + parts = append(parts, escapeLiteral(x)) + } + return fmt.Sprintf("%s IN (%s)", column, strings.Join(parts, ", ")) +} + +func buildFilters(condition map[string]interface{}) []string { + var filters []string + for k, v := range condition { + if v == nil { + continue + } + if s, ok := v.(string); ok && s == "" { + continue + } + switch { + case k == "exists": + if col, ok := asString(v); ok { + if _, known := columnDDL[col]; known { + filters = append(filters, col+" IS NOT NULL") + } + } + case k == "must_not": + if mn, ok := v.(map[string]interface{}); ok { + if ex, ok := mn["exists"]; ok { + if col, ok := asString(ex); ok { + if _, known := columnDDL[col]; known { + filters = append(filters, col+" IS NULL") + } + } + } + } + case isArrayColumn(k): + vals := toStringSlice(v) + if len(vals) == 0 { + filters = append(filters, neverMatch) + continue + } + ors := make([]string, 0, len(vals)) + for _, x := range vals { + ors = append(ors, fmt.Sprintf("list_contains(%s, %s)", k, escapeLiteral(x))) + } + filters = append(filters, "("+strings.Join(ors, " OR ")+")") + case isKnownColumn(k): + switch list := v.(type) { + case []interface{}: + filters = append(filters, inList(k, list)) + case []string: + filters = append(filters, inList(k, toStringSlice(list))) + default: + filters = append(filters, fmt.Sprintf("%s = %s", k, escapeLiteral(v))) + } + } + } + return filters +} + +func isArrayColumn(name string) bool { + _, ok := arraySet[name] + return ok +} + +func isJSONColumn(name string) bool { + _, ok := jsonSet[name] + return ok +} + +// filtersExpr joins predicates, defaulting to TRUE when there are none. +func filtersExpr(filters []string) string { + if len(filters) == 0 { + return "TRUE" + } + return strings.Join(filters, " AND ") +} + +// stripESQuery reduces RAGFlow's tokenized, ^-weighted query_string to a plain +// space-separated token bag for @@. The stored *_ltks columns are tokenized +// the same way, so @@ must see those tokens, not the raw human question. +func stripESQuery(matchingText string) string { + txt := esBoostRe.ReplaceAllString(matchingText, " ") + txt = esSyntaxRe.ReplaceAllString(txt, " ") + seen := map[string]struct{}{} + var out []string + for _, t := range strings.Fields(txt) { + if _, ok := seen[t]; ok { + continue + } + seen[t] = struct{}{} + out = append(out, t) + } + return strings.Join(out, " ") +} + +// l2Normalize returns the unit vector; ip on the unit column is exact cosine. +func l2Normalize(vec []float64) []float64 { + var s float64 + for _, v := range vec { + s += v * v + } + s = math.Sqrt(s) + if s == 0 { + return vec + } + out := make([]float64, len(vec)) + for i, v := range vec { + out[i] = v / s + } + return out +} + +// vectorLiteral renders a normalized query vector as a typed SQL array literal. +func vectorLiteral(vec []float64) string { + norm := l2Normalize(vec) + parts := make([]string, len(norm)) + for i, v := range norm { + parts[i] = strconv.FormatFloat(v, 'f', -1, 64) + } + return fmt.Sprintf("ARRAY[%s]::FLOAT[%d]", strings.Join(parts, ","), len(norm)) +} + +// parsePgArray decodes a PostgreSQL text-array literal (e.g. {a,"b,c"}) into a +// slice. lib/pq returns array columns as this literal when scanned dynamically. +func parsePgArray(literal string) []string { + if len(literal) < 2 || literal[0] != '{' || literal[len(literal)-1] != '}' { + return nil + } + body := literal[1 : len(literal)-1] + if body == "" { + return []string{} + } + var out []string + var buf strings.Builder + inQuote := false + for i := 0; i < len(body); i++ { + c := body[i] + switch { + case c == '"': + if inQuote && i+1 < len(body) && body[i+1] == '"' { + buf.WriteByte('"') + i++ + continue + } + inQuote = !inQuote + case c == '\\' && i+1 < len(body): + buf.WriteByte(body[i+1]) + i++ + case c == ',' && !inQuote: + out = append(out, buf.String()) + buf.Reset() + default: + buf.WriteByte(c) + } + } + out = append(out, buf.String()) + return out +} + +// decodeValue turns a raw database/sql scan value into the entity value for a +// column: JSON columns are parsed, array columns are split, scalars pass +// through. ES field names are stored verbatim so no renaming is needed. +func decodeValue(column string, raw interface{}) interface{} { + if raw == nil { + return nil + } + asBytes := func() (string, bool) { + switch b := raw.(type) { + case []byte: + return string(b), true + case string: + return b, true + } + return "", false + } + if isJSONColumn(column) { + if s, ok := asBytes(); ok { + var v interface{} + if err := json.Unmarshal([]byte(s), &v); err == nil { + return v + } + return s + } + } + if isArrayColumn(column) { + if s, ok := asBytes(); ok { + return parsePgArray(s) + } + } + if s, ok := raw.([]byte); ok { + return string(s) + } + return raw +} diff --git a/internal/engine/serenedb/document.go b/internal/engine/serenedb/document.go new file mode 100644 index 0000000000..3c07114d12 --- /dev/null +++ b/internal/engine/serenedb/document.go @@ -0,0 +1,113 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package serenedb + +import ( + "context" + "fmt" + "strconv" +) + +// asChunkMap coerces a skill document to a chunk map. +func asChunkMap(doc interface{}) (map[string]interface{}, error) { + m, ok := doc.(map[string]interface{}) + if !ok { + return nil, fmt.Errorf("serenedb: document must be a map, got %T", doc) + } + return m, nil +} + +// ensureSkillTable creates the skill table (named by indexName) if absent, +// inferring the vector size from the documents. +func (e *serenedbEngine) ensureSkillTable(ctx context.Context, indexName string, docs []map[string]interface{}) error { + exists, err := e.tableExists(ctx, indexName) + if err != nil { + return err + } + if exists { + return nil + } + size := 0 + for _, doc := range docs { + for k := range doc { + if m := vectorColumnPattern.FindStringSubmatch(k); m != nil { + size, _ = strconv.Atoi(m[1]) + } + } + } + for _, stmt := range chunkTableDDL(indexName, size) { + if err := e.exec(ctx, stmt); err != nil { + return fmt.Errorf("serenedb: create skill table %s: %w", indexName, err) + } + } + return nil +} + +// IndexDocument upserts a single skill document. +func (e *serenedbEngine) IndexDocument(ctx context.Context, indexName, docID string, doc interface{}) error { + m, err := asChunkMap(doc) + if err != nil { + return err + } + if _, ok := m["id"]; !ok { + m["id"] = docID + } + if err := e.ensureSkillTable(ctx, indexName, []map[string]interface{}{m}); err != nil { + return err + } + cols, vals := prepareChunkRow(m, "") + query, args := buildUpsert(indexName, cols, [][]interface{}{vals}) + return e.exec(ctx, query, args...) +} + +// BulkIndex upserts a batch of skill documents. +func (e *serenedbEngine) BulkIndex(ctx context.Context, indexName string, docs []interface{}) (interface{}, error) { + maps := make([]map[string]interface{}, 0, len(docs)) + for _, d := range docs { + m, err := asChunkMap(d) + if err != nil { + return nil, err + } + maps = append(maps, m) + } + if len(maps) == 0 { + return nil, nil + } + if err := e.ensureSkillTable(ctx, indexName, maps); err != nil { + return nil, err + } + for _, m := range maps { + cols, vals := prepareChunkRow(m, "") + query, args := buildUpsert(indexName, cols, [][]interface{}{vals}) + if err := e.exec(ctx, query, args...); err != nil { + return nil, fmt.Errorf("serenedb: bulk index %s: %w", indexName, err) + } + } + return len(maps), nil +} + +// DeleteDocument removes a skill document by id. A missing table is not an error. +func (e *serenedbEngine) DeleteDocument(ctx context.Context, indexName, docID string) error { + exists, err := e.tableExists(ctx, indexName) + if err != nil { + return err + } + if !exists { + return nil + } + return e.exec(ctx, fmt.Sprintf("DELETE FROM %s WHERE id = $1", indexName), docID) +} diff --git a/internal/engine/serenedb/integration_test.go b/internal/engine/serenedb/integration_test.go new file mode 100644 index 0000000000..9f8fc7a4c8 --- /dev/null +++ b/internal/engine/serenedb/integration_test.go @@ -0,0 +1,269 @@ +//go:build integration + +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +// This end-to-end test drives the engine against a real SereneDB. It is +// skipped unless SERENEDB_TEST_DSN points at a live instance (>= 26.07.4), so +// the default test run stays pure and CI-safe. To run it: +// +// docker run -d --name serenedb-gotest -p 127.0.0.1:7899:7890 \ +// -e POSTGRES_PASSWORD=gotest serenedb/serenedb:26.07.4 +// SERENEDB_TEST_DSN='host=127.0.0.1 port=7899 user=postgres password=gotest dbname=postgres sslmode=disable' \ +// go test -run Integration -v ./internal/engine/serenedb/ +package serenedb + +import ( + "context" + "os" + "path/filepath" + "sync" + "testing" + "time" + + "ragflow/internal/common" + "ragflow/internal/engine/types" + "ragflow/internal/server/config" +) + +// logOnce initializes the shared logger the engine's Search path expects. In +// production the server does this at startup; a bare `go test` does not. +var logOnce sync.Once + +func liveEngine(t *testing.T) *serenedbEngine { + t.Helper() + dsn := os.Getenv("SERENEDB_TEST_DSN") + if dsn == "" { + t.Skip("SERENEDB_TEST_DSN not set; skipping live SereneDB integration test") + } + logOnce.Do(func() { + _ = common.InitLogger("info", common.FileOutput{Path: filepath.Join(t.TempDir(), "serenedb-it.log")}, "serenedb-it") + }) + t.Setenv("SERENEDB_DSN", dsn) + e, err := NewEngine(config.SereneDBConfig{}) + if err != nil { + t.Fatalf("NewEngine: %v", err) + } + return e +} + +// waitForFulltext polls until the async inverted index has caught up with the +// last write (SereneDB refreshes the index ~1s after insert, like ES). +func waitForFulltext(t *testing.T, e *serenedbEngine, req *types.SearchRequest) *types.SearchResult { + t.Helper() + ctx := context.Background() + deadline := time.Now().Add(15 * time.Second) + for { + res, err := e.Search(ctx, req) + if err != nil { + t.Fatalf("Search: %v", err) + } + if len(res.Chunks) > 0 || time.Now().After(deadline) { + return res + } + time.Sleep(500 * time.Millisecond) + } +} + +func chunkIDs(res *types.SearchResult) []string { + ids := make([]string, 0, len(res.Chunks)) + for _, c := range res.Chunks { + if id, ok := c["id"].(string); ok { + ids = append(ids, id) + } + } + return ids +} + +func contains(ids []string, want string) bool { + for _, id := range ids { + if id == want { + return true + } + } + return false +} + +func TestIntegrationChunkLifecycle(t *testing.T) { + e := liveEngine(t) + defer e.Close() + ctx := context.Background() + + const base = "ragflow_gotest" + const kb = "kb1" + // Start clean and always tear down the throwaway table. + _ = e.DropChunkStore(ctx, base, kb) + defer func() { _ = e.DropChunkStore(ctx, base, kb) }() + + if err := e.CreateChunkStore(ctx, base, kb, 4, ""); err != nil { + t.Fatalf("CreateChunkStore: %v", err) + } + if ok, _ := e.ChunkStoreExists(ctx, base, kb); !ok { + t.Fatal("ChunkStoreExists = false after create") + } + + chunks := []map[string]interface{}{ + {"id": "a", "doc_id": "d1", "kb_id": kb, "content_ltks": "alpha beta", "content_with_weight": "alpha beta", "q_4_vec": []float64{1, 0, 0, 0}, "important_kwd": []interface{}{"alpha"}}, + {"id": "b", "doc_id": "d1", "kb_id": kb, "content_ltks": "gamma delta", "content_with_weight": "gamma delta", "q_4_vec": []float64{0, 1, 0, 0}, "important_kwd": []interface{}{"gamma"}}, + {"id": "c", "doc_id": "d2", "kb_id": kb, "content_ltks": "alpha gamma", "content_with_weight": "alpha gamma", "q_4_vec": []float64{0.9, 0.1, 0, 0}, "important_kwd": []interface{}{"alpha", "gamma"}}, + } + if _, err := e.InsertChunks(ctx, chunks, base, kb); err != nil { + t.Fatalf("InsertChunks: %v", err) + } + + req := func(exprs []interface{}, filter map[string]interface{}) *types.SearchRequest { + return &types.SearchRequest{ + IndexNames: []string{base}, KbIDs: []string{kb}, + Limit: 10, MatchExprs: exprs, Filter: filter, + } + } + + t.Run("fulltext", func(t *testing.T) { + res := waitForFulltext(t, e, req([]interface{}{ + &types.MatchTextExpr{MatchingText: "alpha", TopN: 10}, + }, nil)) + ids := chunkIDs(res) + if !contains(ids, "a") || !contains(ids, "c") { + t.Fatalf("fulltext 'alpha' should match a and c, got %v", ids) + } + if contains(ids, "b") { + t.Fatalf("fulltext 'alpha' should not match b, got %v", ids) + } + for _, ch := range res.Chunks { + if _, ok := ch["_score"].(float64); !ok { + t.Errorf("chunk %v missing float _score", ch["id"]) + } + } + }) + + t.Run("vector", func(t *testing.T) { + res := waitForFulltext(t, e, req([]interface{}{ + &types.MatchDenseExpr{VectorColumnName: "q_4_vec", EmbeddingData: []float64{1, 0, 0, 0}, TopN: 10}, + }, nil)) + ids := chunkIDs(res) + if len(ids) == 0 || ids[0] != "a" { + t.Fatalf("vector query [1,0,0,0] should rank 'a' first, got %v", ids) + } + }) + + t.Run("fusion", func(t *testing.T) { + res := waitForFulltext(t, e, req([]interface{}{ + &types.MatchTextExpr{MatchingText: "alpha", TopN: 10}, + &types.MatchDenseExpr{VectorColumnName: "q_4_vec", EmbeddingData: []float64{1, 0, 0, 0}, TopN: 10}, + &types.FusionExpr{Method: "weighted_sum", FusionParams: map[string]interface{}{"weights": "0.3,0.7"}}, + }, nil)) + ids := chunkIDs(res) + if len(ids) == 0 || ids[0] != "a" { + t.Fatalf("fusion(alpha, [1,0,0,0]) should rank 'a' first, got %v", ids) + } + }) + + t.Run("filter_only", func(t *testing.T) { + res, err := e.Search(ctx, req(nil, map[string]interface{}{"doc_id": "d2"})) + if err != nil { + t.Fatalf("filter search: %v", err) + } + ids := chunkIDs(res) + if len(ids) != 1 || ids[0] != "c" { + t.Fatalf("filter doc_id=d2 should return only c, got %v", ids) + } + }) + + t.Run("get_and_scores", func(t *testing.T) { + got, err := e.GetChunk(ctx, base, "a", []string{kb}) + if err != nil || got == nil { + t.Fatalf("GetChunk(a) = %v, %v", got, err) + } + m := got.(map[string]interface{}) + if m["content_ltks"] != "alpha beta" { + t.Errorf("GetChunk content = %v", m["content_ltks"]) + } + // important_kwd is a native array column. + if arr, ok := m["important_kwd"].([]string); !ok || len(arr) == 0 || arr[0] != "alpha" { + t.Errorf("important_kwd not decoded as array: %v", m["important_kwd"]) + } + res := waitForFulltext(t, e, req([]interface{}{ + &types.MatchTextExpr{MatchingText: "alpha", TopN: 10}, + }, nil)) + knn, _ := e.KNNScores(ctx, res.Chunks, nil, 10) + scores := e.GetScores(knn) + if _, ok := scores["a"]; !ok { + t.Errorf("GetScores missing 'a': %v", scores) + } + }) + + t.Run("update_and_delete", func(t *testing.T) { + if err := e.UpdateChunks(ctx, + map[string]interface{}{"id": "b"}, + map[string]interface{}{"add": map[string]interface{}{"tag_kwd": "x"}}, + base, kb); err != nil { + t.Fatalf("UpdateChunks add: %v", err) + } + n, err := e.DeleteChunks(ctx, map[string]interface{}{"id": "b"}, base, kb) + if err != nil || n != 1 { + t.Fatalf("DeleteChunks(b) = %d, %v (want 1)", n, err) + } + }) +} + +func TestIntegrationMetadata(t *testing.T) { + e := liveEngine(t) + defer e.Close() + ctx := context.Background() + + const tenant = "gotest_tenant" + _ = e.DropMetadataStore(ctx, tenant) + defer func() { _ = e.DropMetadataStore(ctx, tenant) }() + + if err := e.CreateMetadataStore(ctx, tenant); err != nil { + t.Fatalf("CreateMetadataStore: %v", err) + } + if _, err := e.InsertMetadata(ctx, []map[string]interface{}{ + {"id": "doc1", "kb_id": "kb1", "meta_fields": map[string]interface{}{"author": "ann", "year": float64(2026)}}, + }, tenant); err != nil { + t.Fatalf("InsertMetadata: %v", err) + } + + // Merge update preserves untouched keys. + if err := e.UpdateMetadata(ctx, "doc1", "kb1", map[string]interface{}{"author": "bob"}, tenant); err != nil { + t.Fatalf("UpdateMetadata: %v", err) + } + res, err := e.SearchMetadata(ctx, &types.SearchMetadataRequest{TenantID: tenant, Limit: 10}) + if err != nil { + t.Fatalf("SearchMetadata: %v", err) + } + if res.Total != 1 || len(res.MetadataRecords) != 1 { + t.Fatalf("SearchMetadata total=%d records=%d", res.Total, len(res.MetadataRecords)) + } + mf, ok := res.MetadataRecords[0]["meta_fields"].(map[string]interface{}) + if !ok { + t.Fatalf("meta_fields not decoded to map: %v", res.MetadataRecords[0]["meta_fields"]) + } + if mf["author"] != "bob" { + t.Errorf("merge should set author=bob, got %v", mf["author"]) + } + if _, ok := mf["year"]; !ok { + t.Errorf("merge should preserve year, got %v", mf) + } + + if err := e.DeleteMetadataKeys(ctx, "doc1", "kb1", []string{"year"}, tenant); err != nil { + t.Fatalf("DeleteMetadataKeys: %v", err) + } + after, _ := e.loadMetaFields(ctx, buildMetadataTableName(tenant), "doc1", "kb1") + if _, ok := after["year"]; ok { + t.Errorf("year should be removed, got %v", after) + } +} diff --git a/internal/engine/serenedb/metadata.go b/internal/engine/serenedb/metadata.go new file mode 100644 index 0000000000..4457af1105 --- /dev/null +++ b/internal/engine/serenedb/metadata.go @@ -0,0 +1,272 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package serenedb + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + "ragflow/internal/engine/types" + + "gorm.io/gorm" +) + +// metadataTableDDL creates the per-tenant metadata table and its lookup +// indexes. Pure for tests. +func metadataTableDDL(tableName string) []string { + cols := make([]string, 0, len(docMetaColumnOrder)) + for _, c := range docMetaColumnOrder { + cols = append(cols, fmt.Sprintf("%s %s", c, docMetaDDL[c])) + } + return []string{ + fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s (%s)", tableName, strings.Join(cols, ", ")), + fmt.Sprintf("CREATE INDEX IF NOT EXISTS idx_%s_kb_id ON %s (kb_id)", tableName, tableName), + } +} + +// CreateMetadataStore creates the tenant's document metadata table. +func (e *serenedbEngine) CreateMetadataStore(ctx context.Context, tenantID string) error { + tableName := buildMetadataTableName(tenantID) + for _, stmt := range metadataTableDDL(tableName) { + if err := e.exec(ctx, stmt); err != nil { + return fmt.Errorf("serenedb: create metadata store %s: %w", tableName, err) + } + } + return nil +} + +// DropMetadataStore drops the tenant metadata table. +func (e *serenedbEngine) DropMetadataStore(ctx context.Context, tenantID string) error { + return e.exec(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %s", buildMetadataTableName(tenantID))) +} + +// MetadataStoreExists reports whether the tenant metadata table exists. +func (e *serenedbEngine) MetadataStoreExists(ctx context.Context, tenantID string) (bool, error) { + return e.tableExists(ctx, buildMetadataTableName(tenantID)) +} + +func metaFieldsJSON(v interface{}) string { + if s, ok := v.(string); ok { + return s + } + if v == nil { + return "{}" + } + b, err := json.Marshal(v) + if err != nil { + return "{}" + } + return string(b) +} + +// InsertMetadata upserts metadata records by id. meta_fields is stored as a +// JSON string. +func (e *serenedbEngine) InsertMetadata(ctx context.Context, metadata []map[string]interface{}, tenantID string) ([]string, error) { + if len(metadata) == 0 { + return []string{}, nil + } + tableName := buildMetadataTableName(tenantID) + exists, err := e.tableExists(ctx, tableName) + if err != nil { + return nil, err + } + if !exists { + if err := e.CreateMetadataStore(ctx, tenantID); err != nil { + return nil, err + } + } + query := fmt.Sprintf("INSERT INTO %s (id, kb_id, meta_fields) VALUES ($1, $2, $3) "+ + "ON CONFLICT (id) DO UPDATE SET kb_id = EXCLUDED.kb_id, meta_fields = EXCLUDED.meta_fields", tableName) + for _, rec := range metadata { + if err := e.exec(ctx, query, rec["id"], rec["kb_id"], metaFieldsJSON(rec["meta_fields"])); err != nil { + return nil, fmt.Errorf("serenedb: insert metadata into %s: %w", tableName, err) + } + } + return []string{}, nil +} + +// UpdateMetadata merges metaFields into the stored record, preserving keys the +// caller did not send. Missing rows are inserted. +func (e *serenedbEngine) UpdateMetadata(ctx context.Context, docID, datasetID string, metaFields map[string]interface{}, tenantID string) error { + tableName := buildMetadataTableName(tenantID) + exists, err := e.tableExists(ctx, tableName) + if err != nil { + return err + } + if !exists { + if err := e.CreateMetadataStore(ctx, tenantID); err != nil { + return err + } + } + existing, err := e.loadMetaFields(ctx, tableName, docID, datasetID) + if err != nil { + return err + } + if existing == nil { + return e.exec(ctx, + fmt.Sprintf("INSERT INTO %s (id, kb_id, meta_fields) VALUES ($1, $2, $3)", tableName), + docID, datasetID, metaFieldsJSON(metaFields)) + } + for k, v := range metaFields { + existing[k] = v + } + return e.exec(ctx, + fmt.Sprintf("UPDATE %s SET meta_fields = $1 WHERE id = $2 AND kb_id = $3", tableName), + metaFieldsJSON(existing), docID, datasetID) +} + +// DeleteMetadataKeys removes specific keys from a record's meta_fields, dropping +// the whole row if none remain. +func (e *serenedbEngine) DeleteMetadataKeys(ctx context.Context, docID, datasetID string, keys []string, tenantID string) error { + tableName := buildMetadataTableName(tenantID) + existing, err := e.loadMetaFields(ctx, tableName, docID, datasetID) + if err != nil { + return err + } + if existing == nil { + return fmt.Errorf("serenedb: metadata document not found: %s", docID) + } + changed := false + for _, k := range keys { + if _, ok := existing[k]; ok { + delete(existing, k) + changed = true + } + } + if !changed { + return nil + } + if len(existing) == 0 { + return e.exec(ctx, + fmt.Sprintf("DELETE FROM %s WHERE id = $1 AND kb_id = $2", tableName), docID, datasetID) + } + return e.exec(ctx, + fmt.Sprintf("UPDATE %s SET meta_fields = $1 WHERE id = $2 AND kb_id = $3", tableName), + metaFieldsJSON(existing), docID, datasetID) +} + +// DeleteMetadata removes records matching condition. A missing table is not an +// error. +func (e *serenedbEngine) DeleteMetadata(ctx context.Context, condition map[string]interface{}, tenantID string) (int64, error) { + tableName := buildMetadataTableName(tenantID) + exists, err := e.tableExists(ctx, tableName) + if err != nil { + return 0, err + } + if !exists { + return 0, nil + } + if unknown := unrecognizedFilterKeys(condition); len(unknown) > 0 { + return 0, fmt.Errorf("serenedb: refusing to delete metadata from %s with unrecognized filter keys %v", tableName, unknown) + } + filters := buildFilters(condition) + if len(filters) == 0 { + return 0, nil + } + res, err := e.db.ExecContext(ctx, + fmt.Sprintf("DELETE FROM %s WHERE %s", tableName, strings.Join(filters, " AND "))) + if err != nil { + return 0, err + } + n, _ := res.RowsAffected() + return n, nil +} + +// SearchMetadata returns metadata records matching the request. A missing table +// yields a non-nil empty result so callers do not fall back to in-memory scans. +func (e *serenedbEngine) SearchMetadata(ctx context.Context, req *types.SearchMetadataRequest) (*types.SearchMetadataResult, error) { + if req.TenantID == "" { + return nil, fmt.Errorf("serenedb: SearchMetadata requires a tenant id") + } + tableName := buildMetadataTableName(req.TenantID) + empty := &types.SearchMetadataResult{MetadataRecords: []map[string]interface{}{}, Total: 0} + exists, err := e.tableExists(ctx, tableName) + if err != nil { + return nil, err + } + if !exists { + return empty, nil + } + // SelectFields is interpolated into the projection, so keep only real + // metadata columns. + fields := "*" + if len(req.SelectFields) > 0 { + valid := make([]string, 0, len(req.SelectFields)) + for _, f := range req.SelectFields { + if _, ok := docMetaDDL[f]; ok { + valid = append(valid, f) + } + } + if len(valid) > 0 { + fields = strings.Join(valid, ", ") + } + } + filters := buildFilters(req.Filter) + where := filtersExpr(filters) + limit := req.Limit + if limit <= 0 { + limit = defaultPageSize + } + offset := req.Offset + if offset < 0 { + offset = 0 + } + query := buildFilterSQL(tableName, fields, where, req.OrderBy, limit, offset) + records, err := e.queryMaps(ctx, query) + if err != nil { + return nil, fmt.Errorf("serenedb: search metadata %s: %w", tableName, err) + } + total, err := e.countRows(ctx, tableName, where) + if err != nil { + return nil, err + } + if records == nil { + records = []map[string]interface{}{} + } + return &types.SearchMetadataResult{MetadataRecords: records, Total: total}, nil +} + +// FilterDocIdsByMetaPushdown returns nil, which tells the caller to filter +// document metadata in memory. Pushing metadata predicates into SQL is a +// deferred optimization for this engine; nil is the interface's defined +// fall-back path and is always correct. +func (e *serenedbEngine) FilterDocIdsByMetaPushdown(ctx context.Context, sqlDB *gorm.DB, kbIDs []string, conditions []map[string]interface{}, logic string) []string { + return nil +} + +// loadMetaFields returns the parsed meta_fields map for a record, or nil when +// the record does not exist. +func (e *serenedbEngine) loadMetaFields(ctx context.Context, tableName, docID, datasetID string) (map[string]interface{}, error) { + rows, err := e.queryMaps(ctx, + fmt.Sprintf("SELECT meta_fields FROM %s WHERE id = $1 AND kb_id = $2", tableName), docID, datasetID) + if err != nil { + return nil, err + } + if len(rows) == 0 { + return nil, nil + } + switch mf := rows[0]["meta_fields"].(type) { + case map[string]interface{}: + return mf, nil + case nil: + return map[string]interface{}{}, nil + default: + return map[string]interface{}{}, nil + } +} diff --git a/internal/engine/serenedb/schema.go b/internal/engine/serenedb/schema.go new file mode 100644 index 0000000000..34bdf674ce --- /dev/null +++ b/internal/engine/serenedb/schema.go @@ -0,0 +1,167 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package serenedb + +import ( + "fmt" + "regexp" +) + +// pagerankField is folded into every scored search and is always selected. +const pagerankField = "pagerank_fea" + +// docMetaPrefix marks the per-tenant metadata tables. A datasetID-scoped +// delete must not touch these, and inserts route to the metadata path. +const docMetaPrefix = "ragflow_doc_meta_" + +// dictionaryName is the text-search dictionary the inverted index uses. +// frequency and norm are what make BM25() score at all: without frequency the +// scorer silently returns 0.0 for every row. +const dictionaryName = "rf_scored_delim" + +const dictionaryDDL = "CREATE TEXT SEARCH DICTIONARY IF NOT EXISTS " + dictionaryName + + " (template = 'delimiter', delimiter = ' ', frequency = true, position = true, norm = true)" + +// vectorColumnPattern matches the ES vector field name, e.g. q_1024_vec. +var vectorColumnPattern = regexp.MustCompile(`^q_(\d+)_vec$`) + +// identifierPattern is the shape a table name (derived from a tenant index +// name) must have before it is interpolated into DDL/DML. Table names are not +// parameterizable, so they are validated instead. +var identifierPattern = regexp.MustCompile(`^[A-Za-z_][A-Za-z0-9_]*$`) + +func validIdentifier(name string) bool { + return identifierPattern.MatchString(name) +} + +// Column groups keep the ES mapping names verbatim so the read path needs no +// renames. The write path adapts Go values to SQL; the read path decodes JSON +// columns back to structured values. +var ( + textColumns = []string{ + "docnm_kwd", "doc_type_kwd", "title_tks", "title_sm_tks", "content_with_weight", + "content_ltks", "content_sm_ltks", "important_tks", "question_tks", "create_time", + "img_id", "knowledge_graph_kwd", "entity_kwd", "entity_type_kwd", "from_entity_kwd", + "to_entity_kwd", "removed_kwd", "raptor_kwd", "group_id", "mom_id", "n_hop_with_weight", + } + arrayColumns = []string{"important_kwd", "question_kwd", "tag_kwd", "source_id", "entities_kwd"} + intColumns = []string{"pagerank_fea", "available_int", "weight_int", "raptor_layer_int", "_order_id"} + floatColumns = []string{"create_timestamp_flt", "weight_flt", "rank_flt"} + jsonColumns = []string{ + "tag_feas", "position_int", "page_num_int", "top_int", "chunk_data", + "metadata", "extra", "meta_fields", + } +) + +// columnDDL maps every stored column to its SQL type. columnOrder preserves a +// stable CREATE TABLE order (Go maps do not). +var ( + columnDDL = map[string]string{} + columnOrder []string + arraySet = toSet(arrayColumns) + jsonSet = toSet(jsonColumns) +) + +// ftsColumns are the text columns the inverted index carries. +var ftsColumns = []string{"title_tks", "important_tks", "question_tks", "content_ltks"} + +// lexScoredCol is the single column the scored lexical branch matches. +// ORDER BY BM25() over a multi-column @@ OR returns an empty set, so per-field +// boosts must be summed in application code, never as a SQL-level OR. +// content_ltks is the dominant field and is what the parity eval scored on. +const lexScoredCol = "content_ltks" + +// docMetaColumnOrder / docMetaDDL define the per-tenant metadata table. +var docMetaColumnOrder = []string{"id", "kb_id", "meta_fields"} + +var docMetaDDL = map[string]string{ + "id": "VARCHAR PRIMARY KEY", + "kb_id": "VARCHAR", + "meta_fields": "JSON", +} + +// columnDefaults are applied on insert when the caller omits them. +var columnDefaults = map[string]interface{}{ + "available_int": 1, + "removed_kwd": "N", + "_order_id": 0, +} + +func init() { + columnOrder = append(columnOrder, "id", "kb_id", "doc_id") + columnDDL["id"] = "VARCHAR PRIMARY KEY" + columnDDL["kb_id"] = "VARCHAR" + columnDDL["doc_id"] = "VARCHAR" + add := func(cols []string, typ string) { + for _, c := range cols { + if _, seen := columnDDL[c]; seen { + continue + } + columnDDL[c] = typ + columnOrder = append(columnOrder, c) + } + } + add(textColumns, "TEXT") + add(arrayColumns, "VARCHAR[]") + add(intColumns, "INTEGER") + add(floatColumns, "DOUBLE PRECISION") + add(jsonColumns, "JSON") +} + +func toSet(cols []string) map[string]struct{} { + s := make(map[string]struct{}, len(cols)) + for _, c := range cols { + s[c] = struct{}{} + } + return s +} + +// isKnownColumn reports whether a field is a stored column or a vector column. +func isKnownColumn(name string) bool { + if _, ok := columnDDL[name]; ok { + return true + } + return vectorColumnPattern.MatchString(name) +} + +// normColumn is the L2-normalized shadow of a vector column. ip on the unit +// column is exact cosine and is what the IVF index quantizes. +func normColumn(vectorSize int) string { + return fmt.Sprintf("q_%d_vec_n", vectorSize) +} + +func rawVectorColumn(vectorSize int) string { + return fmt.Sprintf("q_%d_vec", vectorSize) +} + +// indexRelation is the inverted index name for a table. +func indexRelation(tableName string) string { + return "idx_" + tableName +} + +// chunkTableName returns the tenant's chunk table. All of a tenant's datasets +// share one table (the Elasticsearch/OceanBase model), with kb_id as a filter +// column, so BM25 statistics are computed over the whole tenant corpus rather +// than per dataset. baseName is already the tenant index name. +func chunkTableName(baseName string) string { + return baseName +} + +// buildMetadataTableName returns the per-tenant metadata table name. +func buildMetadataTableName(tenantID string) string { + return docMetaPrefix + tenantID +} diff --git a/internal/engine/serenedb/search.go b/internal/engine/serenedb/search.go new file mode 100644 index 0000000000..e032a2b941 --- /dev/null +++ b/internal/engine/serenedb/search.go @@ -0,0 +1,587 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package serenedb + +import ( + "context" + "fmt" + "math" + "regexp" + "sort" + "strconv" + "strings" + + "ragflow/internal/engine/types" +) + +const ( + defaultPageSize = 30 + defaultBranchTopN = 200 +) + +// pagerankExpr folds the dataset pagerank into a scored search, matching the ES +// fusion math (score + pagerank_fea/100). +const pagerankExpr = "COALESCE(" + pagerankField + ", 0) / 100.0" + +type parsedMatch struct { + textQuery string + textTopN int + vectorData []float64 + vectorTopN int + vecThreshold float64 + vectorWeight float64 + hasText bool + hasVector bool +} + +// Search runs the fulltext, vector, hybrid-fusion, or filter-only query implied +// by the request's match expressions against the tenant table(s), scoping to +// req.KbIDs with a kb_id filter. Because all datasets share one table, BM25 is +// scored over the whole tenant corpus. +func (e *serenedbEngine) Search(ctx context.Context, req *types.SearchRequest) (*types.SearchResult, error) { + types.LogSearchRequest("serenedb", req) + + pm := parseMatchExprs(req.MatchExprs) + outputFields := resolveOutputFields(req.SelectFields) + fieldsExpr := strings.Join(outputFields, ", ") + + filters := searchFilters(req.Filter, req.KbIDs, pm.hasText || pm.hasVector) + where := filtersExpr(filters) + + offset := req.Offset + if offset < 0 { + offset = 0 + } + limit := req.Limit + if limit <= 0 { + limit = defaultPageSize + } + scored := pm.hasText || pm.hasVector + + result := &types.SearchResult{Chunks: []map[string]interface{}{}} + for _, tableName := range req.IndexNames { + exists, err := e.tableExists(ctx, tableName) + if err != nil { + return nil, err + } + if !exists { + continue + } + var query string + switch { + case pm.hasText && pm.hasVector: + query = buildFusionSQL(tableName, fieldsExpr, outputFields, where, pm, offset, limit) + case pm.hasText: + query = buildFulltextSQL(tableName, fieldsExpr, where, pm.textQuery, branchLimit(pm.textTopN, limit), offset) + case pm.hasVector: + query = buildVectorSQL(tableName, fieldsExpr, where, pm, branchLimit(pm.vectorTopN, limit), offset) + default: + total, err := e.countRows(ctx, tableName, where) + if err != nil { + return nil, err + } + result.Total += total + query = buildFilterSQL(tableName, fieldsExpr, where, req.OrderBy, limit, offset) + } + rows, err := e.queryMaps(ctx, query) + if err != nil { + return nil, fmt.Errorf("serenedb: search %s: %w", tableName, err) + } + result.Chunks = append(result.Chunks, rows...) + } + + if scored && len(result.Chunks) > 1 { + sortByScore(result.Chunks) + } + if limit > 0 && len(result.Chunks) > limit { + result.Chunks = result.Chunks[:limit] + } + if result.Total == 0 { + result.Total = int64(len(result.Chunks)) + } + return result, nil +} + +// resolveOutputFields keeps id and pagerank in the projection and drops the +// synthetic _score and any unknown fields. +func resolveOutputFields(selectFields []string) []string { + var fields []string + seen := map[string]struct{}{} + add := func(f string) { + if _, ok := seen[f]; ok { + return + } + seen[f] = struct{}{} + fields = append(fields, f) + } + add("id") + src := selectFields + useAll := len(src) == 0 + for _, f := range src { + if f == "*" { + useAll = true + } + } + if useAll { + for _, c := range columnOrder { + add(c) + } + } else { + for _, f := range src { + if f == "_score" || f == "*" { + continue + } + if isKnownColumn(f) { + add(f) + } + } + } + add(pagerankField) + return fields +} + +// searchFilters builds the SQL predicates for a search. kb_id scopes the query +// to the requested datasets within the shared tenant table, and scored queries +// default to available_int=1 when the caller did not set it. +func searchFilters(filter map[string]interface{}, kbIDs []string, scored bool) []string { + cond := map[string]interface{}{} + for k, v := range filter { + cond[k] = v + } + if kbs := stringSlice(kbIDs); len(kbs) > 0 { + cond["kb_id"] = kbs + } + if scored { + _, hasAvail := cond["available_int"] + _, hasStatus := cond["status"] + if !hasAvail && !hasStatus { + cond["available_int"] = 1 + } + } + return buildFilters(cond) +} + +// branchLimit is the row cap for a single-mode query. +func branchLimit(topN, limit int) int { + if limit > 0 { + return limit + } + if topN > 0 { + return topN + } + return defaultPageSize +} + +// buildFulltextSQL scores the single lexical column with BM25 plus pagerank. +func buildFulltextSQL(tableName, fieldsExpr, where, textQuery string, limit, offset int) string { + idx := indexRelation(tableName) + match := fmt.Sprintf("%s @@ %s", lexScoredCol, escapeLiteral(textQuery)) + return fmt.Sprintf( + "SELECT %s, BM25(%s.tableoid) + %s AS _score FROM %s WHERE %s AND (%s) "+ + "ORDER BY _score DESC LIMIT %d OFFSET %d", + fieldsExpr, idx, pagerankExpr, idx, where, match, limit, offset) +} + +// buildVectorSQL runs the ANN scan on the normalized shadow column. The +// similarity threshold goes straight in the WHERE (relies on SereneDB 26.07.4). +func buildVectorSQL(tableName, fieldsExpr, where string, pm parsedMatch, limit, offset int) string { + idx := indexRelation(tableName) + vecN := normColumn(len(pm.vectorData)) + qv := vectorLiteral(pm.vectorData) + sim := fmt.Sprintf("-(%s <#> %s)", vecN, qv) + return fmt.Sprintf( + "SELECT %s, %s + %s AS _score FROM %s WHERE %s AND %s >= %s "+ + "ORDER BY %s <#> %s LIMIT %d OFFSET %d", + fieldsExpr, sim, pagerankExpr, idx, where, sim, formatFloat(pm.vecThreshold), + vecN, qv, limit, offset) +} + +// buildFusionSQL is the one-statement hybrid over a single tenant table: the +// BM25 branch normalized against the whole-table max with a window function, +// FULL OUTER JOINed with the ANN branch, weighted-summed with pagerank. Because +// the table holds the whole tenant corpus, the BM25 normalization is global. +func buildFusionSQL(tableName, fieldsExpr string, outputFields []string, where string, pm parsedMatch, offset, limit int) string { + idx := indexRelation(tableName) + vecN := normColumn(len(pm.vectorData)) + qv := vectorLiteral(pm.vectorData) + match := fmt.Sprintf("%s @@ %s", lexScoredCol, escapeLiteral(pm.textQuery)) + lexN := pm.textTopN + if lexN <= 0 { + lexN = defaultBranchTopN + } + vN := pm.vectorTopN + if vN <= 0 { + vN = defaultBranchTopN + } + n := limit + if n <= 0 { + n = lexN + vN + } + prefixed := make([]string, len(outputFields)) + for i, f := range outputFields { + prefixed[i] = "t." + f + } + vw := pm.vectorWeight + return fmt.Sprintf(`WITH lex AS ( + SELECT id, BM25(%s.tableoid) AS s + FROM %s WHERE %s AND (%s) + ORDER BY s DESC LIMIT %d), +lexn AS (SELECT id, s / NULLIF(MAX(s) OVER (), 0) AS sn FROM lex), +vec AS ( + SELECT id, -(%s <#> %s) AS sim + FROM %s WHERE %s AND -(%s <#> %s) >= %s + ORDER BY %s <#> %s LIMIT %d), +fused AS ( + SELECT COALESCE(l.id, v.id) AS id, + COALESCE(l.sn, 0) * %s + COALESCE(v.sim, 0) * %s AS fs + FROM lexn l FULL OUTER JOIN vec v ON l.id = v.id) +SELECT %s, f.fs + COALESCE(t.%s, 0) / 100.0 AS _score +FROM fused f JOIN %s t ON t.id = f.id +ORDER BY _score DESC LIMIT %d OFFSET %d`, + idx, idx, where, match, lexN, + vecN, qv, idx, where, vecN, qv, formatFloat(pm.vecThreshold), vecN, qv, vN, + formatWeight(1.0-vw), formatWeight(vw), + strings.Join(prefixed, ", "), pagerankField, tableName, n, offset) +} + +// buildFilterSQL is the metadata/browse path: no scoring, optional ordering. +func buildFilterSQL(tableName, fieldsExpr, where string, orderBy *types.OrderByExpr, limit, offset int) string { + var order string + if orderBy != nil && len(orderBy.Fields) > 0 { + var parts []string + for _, f := range orderBy.Fields { + if _, known := columnDDL[f.Field]; !known { + continue + } + dir := "ASC" + if f.Type == types.SortDesc { + dir = "DESC" + } + parts = append(parts, fmt.Sprintf("%s %s", f.Field, dir)) + } + if len(parts) > 0 { + order = " ORDER BY " + strings.Join(parts, ", ") + } + } + return fmt.Sprintf("SELECT %s FROM %s WHERE %s%s LIMIT %d OFFSET %d", + fieldsExpr, tableName, where, order, limit, offset) +} + +func (e *serenedbEngine) countRows(ctx context.Context, tableName, where string) (int64, error) { + rows, err := e.queryMaps(ctx, fmt.Sprintf("SELECT count(*) AS c FROM %s WHERE %s", tableName, where)) + if err != nil { + return 0, err + } + if len(rows) == 0 { + return 0, nil + } + return toInt64(rows[0]["c"]), nil +} + +// parseMatchExprs extracts the text query, dense vector, and fusion weight from +// the ordered match expressions. +func parseMatchExprs(exprs []interface{}) parsedMatch { + pm := parsedMatch{vectorWeight: 0.5} + for _, m := range exprs { + switch expr := m.(type) { + case string: + if expr != "" { + pm.textQuery = stripESQuery(expr) + pm.hasText = true + } + case *types.MatchTextExpr: + raw := expr.MatchingText + if raw == "" && expr.ExtraOptions != nil { + if oq, ok := expr.ExtraOptions["original_query"].(string); ok { + raw = oq + } + } + if raw != "" { + pm.textQuery = stripESQuery(raw) + pm.textTopN = expr.TopN + pm.hasText = true + } + case *types.MatchDenseExpr: + if len(expr.EmbeddingData) > 0 { + pm.vectorData = expr.EmbeddingData + pm.vectorTopN = expr.TopN + pm.vecThreshold = denseThreshold(expr.ExtraOptions) + pm.hasVector = true + } + case *types.FusionExpr: + if w, ok := fusionVectorWeight(expr.FusionParams); ok { + pm.vectorWeight = w + } + } + } + return pm +} + +func denseThreshold(opts map[string]interface{}) float64 { + if opts == nil { + return 0.0 + } + switch v := opts["similarity"].(type) { + case float64: + return v + case string: + f, _ := strconv.ParseFloat(v, 64) + return f + } + if s, ok := opts["threshold"].(string); ok { + f, _ := strconv.ParseFloat(s, 64) + return f + } + return 0.0 +} + +// fusionVectorWeight reads the vector weight (second element of the weights +// pair) from the fusion params. +func fusionVectorWeight(params map[string]interface{}) (float64, bool) { + if params == nil { + return 0, false + } + w, ok := params["weights"].(string) + if !ok { + return 0, false + } + parts := strings.Split(w, ",") + if len(parts) < 2 { + return 0, false + } + f, err := strconv.ParseFloat(strings.TrimSpace(parts[1]), 64) + if err != nil { + return 0, false + } + return f, true +} + +func formatFloat(f float64) string { + return strconv.FormatFloat(f, 'f', -1, 64) +} + +// formatWeight renders a fusion weight without float-subtraction noise +// (e.g. 1.0 - 0.95 renders as 0.05, not 0.050000000000000044). Weights are +// low-precision by nature, so rounding to 1e-6 is exact enough. +func formatWeight(f float64) string { + return strconv.FormatFloat(math.Round(f*1e6)/1e6, 'f', -1, 64) +} + +func toInt64(v interface{}) int64 { + switch n := v.(type) { + case int64: + return n + case int: + return int64(n) + case float64: + return int64(n) + case string: + i, _ := strconv.ParseInt(n, 10, 64) + return i + } + return 0 +} + +func toFloat64(v interface{}) float64 { + switch n := v.(type) { + case float64: + return n + case int64: + return float64(n) + case int: + return float64(n) + case string: + f, _ := strconv.ParseFloat(n, 64) + return f + } + return 0 +} + +func sortByScore(chunks []map[string]interface{}) { + sort.SliceStable(chunks, func(i, j int) bool { + return toFloat64(chunks[i]["_score"]) > toFloat64(chunks[j]["_score"]) + }) +} + +// GetChunkIDs returns the ids of the given chunks in order. +func (e *serenedbEngine) GetChunkIDs(chunks []map[string]interface{}) []string { + ids := make([]string, 0, len(chunks)) + for _, c := range chunks { + if id, ok := c["id"].(string); ok { + ids = append(ids, id) + } + } + return ids +} + +// GetScores maps chunk id to its recovered score, reading the structure +// KNNScores produces. +func (e *serenedbEngine) GetScores(searchResult map[string]interface{}) map[string]float64 { + scores := map[string]float64{} + hits, ok := searchResult["hits"].(map[string]interface{}) + if !ok { + return scores + } + hitList, ok := hits["hits"].([]interface{}) + if !ok { + return scores + } + for _, h := range hitList { + hit, ok := h.(map[string]interface{}) + if !ok { + continue + } + id, ok := hit["_id"].(string) + if !ok { + continue + } + scores[id] = toFloat64(hit["_score"]) + } + return scores +} + +// KNNScores repackages the per-chunk _score into the hits structure GetScores +// consumes. +func (e *serenedbEngine) KNNScores(ctx context.Context, chunks []map[string]interface{}, queryVector []float64, topK int) (map[string]interface{}, error) { + if len(chunks) == 0 { + return nil, nil + } + hits := make([]interface{}, 0, len(chunks)) + for _, c := range chunks { + id, _ := c["id"].(string) + hits = append(hits, map[string]interface{}{"_id": id, "_score": toFloat64(c["_score"])}) + } + return map[string]interface{}{"hits": map[string]interface{}{"hits": hits}}, nil +} + +// GetFields returns the requested fields per chunk id, omitting nil values. +func (e *serenedbEngine) GetFields(chunks []map[string]interface{}, fields []string) map[string]map[string]interface{} { + out := map[string]map[string]interface{}{} + if len(chunks) == 0 || len(fields) == 0 { + return out + } + for _, c := range chunks { + id, ok := c["id"].(string) + if !ok { + continue + } + row := map[string]interface{}{} + for _, f := range fields { + if v, ok := c[f]; ok && v != nil { + row[f] = v + } + } + out[id] = row + } + return out +} + +// GetAggregation counts distinct values of a field across chunks, ordered by +// count descending. +func (e *serenedbEngine) GetAggregation(chunks []map[string]interface{}, fieldName string) []map[string]interface{} { + counts := map[string]int{} + for _, c := range chunks { + // Aggregation-style chunks carry an explicit value/count. + if val, ok := c["value"]; ok { + if s, ok := val.(string); ok { + counts[s] += toInt(c["count"]) + continue + } + } + v, ok := c[fieldName] + if !ok { + continue + } + for _, item := range asList(v) { + if s, ok := item.(string); ok { + if strings.TrimSpace(s) != "" { + counts[s]++ + } + } + } + } + out := make([]map[string]interface{}, 0, len(counts)) + for k, n := range counts { + out = append(out, map[string]interface{}{"key": k, "count": n}) + } + sort.SliceStable(out, func(i, j int) bool { + return out[i]["count"].(int) > out[j]["count"].(int) + }) + return out +} + +var nonWordBoundary = regexp.MustCompile(`\s*`) + +// GetHighlight emphasizes keyword hits in the stored text, client-side. +func (e *serenedbEngine) GetHighlight(chunks []map[string]interface{}, keywords []string, fieldName string) map[string]string { + ans := map[string]string{} + if len(chunks) == 0 || len(keywords) == 0 { + return ans + } + var pats []*regexp.Regexp + for _, k := range keywords { + if k == "" { + continue + } + pats = append(pats, regexp.MustCompile(`(?i)(^|\W)(`+regexp.QuoteMeta(k)+`)(\W|$)`)) + } + for _, c := range chunks { + id, ok := c["id"].(string) + if !ok { + continue + } + txt, ok := c[fieldName].(string) + if !ok || txt == "" { + continue + } + marked := txt + for _, p := range pats { + marked = p.ReplaceAllString(marked, "$1$2$3") + } + if strings.Contains(marked, "") { + ans[id] = nonWordBoundary.ReplaceAllString(marked, " ") + } + } + return ans +} + +func asList(v interface{}) []interface{} { + switch val := v.(type) { + case []interface{}: + return val + case []string: + out := make([]interface{}, len(val)) + for i, s := range val { + out[i] = s + } + return out + default: + return []interface{}{v} + } +} + +func toInt(v interface{}) int { + switch n := v.(type) { + case int: + return n + case int64: + return int(n) + case float64: + return int(n) + } + return 0 +} diff --git a/internal/engine/serenedb/serenedb.go b/internal/engine/serenedb/serenedb.go new file mode 100644 index 0000000000..8192dc5b93 --- /dev/null +++ b/internal/engine/serenedb/serenedb.go @@ -0,0 +1,34 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +// Package serenedb implements the doc-store DocEngine backed by SereneDB, a +// PostgreSQL-wire engine whose single inverted index carries both a scored +// text column (@@, BM25) and an IVF vector column (<#>, inner product), so +// hybrid search is one SQL statement. It connects with database/sql + lib/pq +// and emits SQL directly; the query shapes mirror the Python +// SereneDBConnection (rag/utils/serenedb_conn.py) that reached Elasticsearch +// retrieval parity. +// +// Table layout follows the Elasticsearch/OceanBase model: one chunk table per +// tenant (the index name) with kb_id as a filter column, so BM25 statistics are +// computed over the whole tenant corpus. Metadata is one table per tenant +// (ragflow_doc_meta_{tenantID}). +// +// Minimum engine version: SereneDB 26.07.4. The vector-branch and fusion +// queries use the natural forms that rely on the 26.07.4 fixes for the +// vector-op predicate in an ANN scan's WHERE and the multi-reference index +// CTE; on earlier builds both silently returned empty. +package serenedb diff --git a/internal/engine/serenedb/serenedb_test.go b/internal/engine/serenedb/serenedb_test.go new file mode 100644 index 0000000000..2f81356639 --- /dev/null +++ b/internal/engine/serenedb/serenedb_test.go @@ -0,0 +1,401 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +// These tests exercise the deterministic SQL builders and value codecs. They +// need no live SereneDB, matching the pure-unit-test style of the infinity +// engine, so they run under build.sh --test. +package serenedb + +import ( + "reflect" + "strings" + "testing" + + "ragflow/internal/engine/types" + "ragflow/internal/server/config" +) + +func mustContain(t *testing.T, got, want string) { + t.Helper() + if !strings.Contains(got, want) { + t.Errorf("expected substring %q in:\n%s", want, got) + } +} + +func mustNotContain(t *testing.T, got, want string) { + t.Helper() + if strings.Contains(got, want) { + t.Errorf("did not expect substring %q in:\n%s", want, got) + } +} + +func TestEscapeLiteral(t *testing.T) { + cases := map[string]struct { + in interface{} + want string + }{ + "nil": {nil, "NULL"}, + "bool": {true, "true"}, + "int": {7, "7"}, + "float": {1.5, "1.5"}, + "string": {"a'b", "'a''b'"}, + "list": {[]string{"x", "y"}, `'["x","y"]'`}, + } + for name, c := range cases { + if got := escapeLiteral(c.in); got != c.want { + t.Errorf("%s: escapeLiteral(%v) = %q, want %q", name, c.in, got, c.want) + } + } +} + +func TestBuildFiltersArrayUsesListContains(t *testing.T) { + got := buildFilters(map[string]interface{}{"tag_kwd": []interface{}{"a", "b"}}) + if len(got) != 1 { + t.Fatalf("want 1 filter, got %v", got) + } + mustContain(t, got[0], "list_contains(tag_kwd, 'a')") + mustContain(t, got[0], " OR ") + mustContain(t, got[0], "list_contains(tag_kwd, 'b')") +} + +func TestBuildFiltersScalarAndIn(t *testing.T) { + scalar := buildFilters(map[string]interface{}{"doc_id": "d1"}) + if len(scalar) != 1 || scalar[0] != "doc_id = 'd1'" { + t.Errorf("scalar filter = %v", scalar) + } + in := buildFilters(map[string]interface{}{"doc_id": []interface{}{"a", "b"}}) + if len(in) != 1 || in[0] != "doc_id IN ('a', 'b')" { + t.Errorf("IN filter = %v", in) + } +} + +func TestBuildFiltersExists(t *testing.T) { + ex := buildFilters(map[string]interface{}{"exists": "img_id"}) + if len(ex) != 1 || ex[0] != "img_id IS NOT NULL" { + t.Errorf("exists = %v", ex) + } + mn := buildFilters(map[string]interface{}{"must_not": map[string]interface{}{"exists": "img_id"}}) + if len(mn) != 1 || mn[0] != "img_id IS NULL" { + t.Errorf("must_not exists = %v", mn) + } +} + +func TestBuildFiltersSkipsUnknownAndEmpty(t *testing.T) { + got := buildFilters(map[string]interface{}{"not_a_column": "x", "doc_id": ""}) + if len(got) != 0 { + t.Errorf("expected no filters, got %v", got) + } +} + +func TestStripESQuery(t *testing.T) { + got := stripESQuery(`(auto^0.5) (ptr^0.4) "auto _"^0.9 auto`) + // boosts and punctuation removed; tokens de-duplicated preserving order. + if got != "auto ptr _" { + t.Errorf("stripESQuery = %q", got) + } +} + +func TestL2NormalizeUnitLength(t *testing.T) { + out := l2Normalize([]float64{3, 4}) + if !reflect.DeepEqual(out, []float64{0.6, 0.8}) { + t.Errorf("l2Normalize = %v", out) + } + if got := l2Normalize([]float64{0, 0}); !reflect.DeepEqual(got, []float64{0, 0}) { + t.Errorf("zero vector should pass through, got %v", got) + } +} + +func TestVectorLiteralTyped(t *testing.T) { + got := vectorLiteral([]float64{3, 4}) + mustContain(t, got, "ARRAY[0.6,0.8]::FLOAT[2]") +} + +func TestParsePgArray(t *testing.T) { + cases := map[string]struct { + in string + want []string + }{ + "simple": {"{a,b,c}", []string{"a", "b", "c"}}, + "empty": {"{}", []string{}}, + "quoted": {`{"a,b","c"}`, []string{"a,b", "c"}}, + "nonarr": {"plain", nil}, + } + for name, c := range cases { + if got := parsePgArray(c.in); !reflect.DeepEqual(got, c.want) { + t.Errorf("%s: parsePgArray(%q) = %v, want %v", name, c.in, got, c.want) + } + } +} + +func TestDecodeValue(t *testing.T) { + if got := decodeValue("tag_kwd", []byte("{a,b}")); !reflect.DeepEqual(got, []string{"a", "b"}) { + t.Errorf("array decode = %v", got) + } + got := decodeValue("position_int", []byte(`{"p":1}`)) + m, ok := got.(map[string]interface{}) + if !ok || m["p"].(float64) != 1 { + t.Errorf("json decode = %v", got) + } + if got := decodeValue("doc_id", []byte("d1")); got != "d1" { + t.Errorf("scalar decode = %v", got) + } + if got := decodeValue("doc_id", nil); got != nil { + t.Errorf("nil decode = %v", got) + } +} + +func TestChunkTableDDLLandmines(t *testing.T) { + stmts := chunkTableDDL("ragflow_t1_kb1", 1024) + joined := strings.Join(stmts, "\n") + // dictionary must declare frequency+norm or BM25() silently scores 0. + mustContain(t, joined, "frequency = true") + mustContain(t, joined, "norm = true") + // normalized shadow column indexed with ip/sq8. + mustContain(t, joined, "q_1024_vec_n FLOAT[1024]") + mustContain(t, joined, "q_1024_vec_n ivf (metric = 'ip', quant = 'sq8')") + mustContain(t, joined, "optimize_top_k = 'bm25(1.2, 0.75)'") + mustContain(t, joined, "CREATE TABLE IF NOT EXISTS ragflow_t1_kb1") +} + +func TestBuildFulltextSQLSingleColumn(t *testing.T) { + sql := buildFulltextSQL("t", "id, content_ltks", "TRUE", "hello world", 10, 0) + mustContain(t, sql, "content_ltks @@ 'hello world'") + mustContain(t, sql, "BM25(idx_t.tableoid)") + mustContain(t, sql, "ORDER BY _score DESC LIMIT 10 OFFSET 0") + // The scored branch must match one column only, never an OR across fields. + mustNotContain(t, sql, "title_tks @@") +} + +func TestBuildVectorSQLThresholdInWhere(t *testing.T) { + pm := parsedMatch{vectorData: []float64{3, 4}, vecThreshold: 0.2} + sql := buildVectorSQL("t", "id", "TRUE", pm, 10, 5) + mustContain(t, sql, "-(q_2_vec_n <#> ARRAY[0.6,0.8]::FLOAT[2]) >= 0.2") + mustContain(t, sql, "ORDER BY q_2_vec_n <#> ARRAY[0.6,0.8]::FLOAT[2] LIMIT 10 OFFSET 5") +} + +func TestBuildFusionSQLShape(t *testing.T) { + pm := parsedMatch{ + textQuery: "q", vectorData: []float64{3, 4}, + textTopN: 200, vectorTopN: 200, vectorWeight: 0.95, + } + sql := buildFusionSQL("t", "id, content_ltks", []string{"id", "content_ltks"}, "TRUE", pm, 0, 30) + mustContain(t, sql, "s / NULLIF(MAX(s) OVER (), 0) AS sn") // window-fn normalization over the whole tenant table + mustContain(t, sql, "FULL OUTER JOIN vec v ON l.id = v.id") + mustContain(t, sql, "COALESCE(l.sn, 0) * 0.05 + COALESCE(v.sim, 0) * 0.95 AS fs") + mustContain(t, sql, "f.fs + COALESCE(t.pagerank_fea, 0) / 100.0 AS _score") + mustContain(t, sql, "ORDER BY _score DESC") + mustContain(t, sql, "content_ltks @@ 'q'") + mustContain(t, sql, "t.id, t.content_ltks") // output fields prefixed with t. +} + +func TestBuildUpsert(t *testing.T) { + q, args := buildUpsert("t", []string{"id", "doc_id"}, [][]interface{}{{"1", "d1"}, {"2", "d2"}}) + mustContain(t, q, "INSERT INTO t (id, doc_id) VALUES ($1, $2), ($3, $4)") + mustContain(t, q, "ON CONFLICT (id) DO UPDATE SET doc_id = EXCLUDED.doc_id") + mustNotContain(t, q, "id = EXCLUDED.id") // never update the conflict key + if !reflect.DeepEqual(args, []interface{}{"1", "d1", "2", "d2"}) { + t.Errorf("args = %v", args) + } +} + +func TestBuildUpdateSetsAddRemove(t *testing.T) { + sets := buildUpdateSets(map[string]interface{}{ + "add": map[string]interface{}{"tag_kwd": "x"}, + "remove": map[string]interface{}{"tag_kwd": "y"}, + }) + joined := strings.Join(sets, " | ") + mustContain(t, joined, "tag_kwd = list_append(tag_kwd, 'x')") + mustContain(t, joined, "tag_kwd = array_remove(tag_kwd, 'y')") +} + +func TestBuildUpdateSetsRejectsUnknownColumns(t *testing.T) { + // Every branch must whitelist the identifier; an unknown key (including a + // remove-to-NULL, which is the only branch that emits a bare identifier) + // must never reach the SQL. + sets := buildUpdateSets(map[string]interface{}{ + "remove": map[string]interface{}{"evil = 1; DROP TABLE t; --": nil, "tag_kwd": nil}, + "add": map[string]interface{}{"not_a_column": "x"}, + "bogus_field": "v", + "1=1; DROP": "v", + }) + joined := strings.Join(sets, " | ") + mustContain(t, joined, "tag_kwd = NULL") + mustNotContain(t, joined, "DROP") + mustNotContain(t, joined, "not_a_column") + mustNotContain(t, joined, "bogus_field") +} + +func TestChunkTableNameIsTenantScoped(t *testing.T) { + // All datasets share the tenant table; datasetID never enters the name. + if got := chunkTableName("ragflow_t1"); got != "ragflow_t1" { + t.Errorf("chunkTableName = %q, want ragflow_t1", got) + } +} + +func TestResolveOutputFields(t *testing.T) { + got := resolveOutputFields([]string{"content_ltks", "_score", "bogus_field"}) + // id first, _score and unknown dropped, pagerank appended. + if got[0] != "id" { + t.Errorf("id must be first: %v", got) + } + joined := strings.Join(got, ",") + mustContain(t, joined, "content_ltks") + mustContain(t, joined, pagerankField) + mustNotContain(t, joined, "_score") + mustNotContain(t, joined, "bogus_field") +} + +func TestSearchFiltersScoping(t *testing.T) { + // KbIDs become an IN predicate over the shared tenant table. + scored := searchFilters(map[string]interface{}{}, []string{"kb1", "kb2"}, true) + joined := strings.Join(scored, " AND ") + mustContain(t, joined, "kb_id IN ('kb1', 'kb2')") + mustContain(t, joined, "available_int = 1") + // No KbIDs and no match expr -> no predicates. + if got := searchFilters(map[string]interface{}{}, nil, false); len(got) != 0 { + t.Errorf("expected no filters, got %v", got) + } + // A scored query with no available_int/status defaults available_int=1. + one := searchFilters(map[string]interface{}{}, nil, true) + if len(one) != 1 || one[0] != "available_int = 1" { + t.Errorf("scored default = %v", one) + } + // Blank dataset ids are dropped, so no empty IN () is emitted. + if got := searchFilters(map[string]interface{}{}, []string{""}, false); len(got) != 0 { + t.Errorf("blank kb ids should yield no filter, got %v", got) + } +} + +func TestFusionVectorWeight(t *testing.T) { + w, ok := fusionVectorWeight(map[string]interface{}{"weights": "0.05,0.95"}) + if !ok || w != 0.95 { + t.Errorf("weight = %v, ok = %v", w, ok) + } + if _, ok := fusionVectorWeight(map[string]interface{}{}); ok { + t.Error("missing weights should not parse") + } +} + +func TestParseMatchExprs(t *testing.T) { + pm := parseMatchExprs([]interface{}{ + &types.MatchTextExpr{MatchingText: "auto^1", TopN: 50}, + &types.MatchDenseExpr{EmbeddingData: []float64{1, 2}, TopN: 60, ExtraOptions: map[string]interface{}{"similarity": 0.3}}, + &types.FusionExpr{FusionParams: map[string]interface{}{"weights": "0.1,0.9"}}, + }) + if !pm.hasText || !pm.hasVector { + t.Fatalf("expected text+vector, got %+v", pm) + } + if pm.textQuery != "auto" || pm.textTopN != 50 { + t.Errorf("text parse = %q/%d", pm.textQuery, pm.textTopN) + } + if pm.vecThreshold != 0.3 || pm.vectorTopN != 60 { + t.Errorf("vector parse = %v/%d", pm.vecThreshold, pm.vectorTopN) + } + if pm.vectorWeight != 0.9 { + t.Errorf("fusion weight = %v", pm.vectorWeight) + } +} + +func TestBoundRunSQL(t *testing.T) { + if got := boundRunSQL("SELECT * FROM t;", 1024); got != "SELECT * FROM t LIMIT 1024" { + t.Errorf("bound = %q", got) + } + if got := boundRunSQL("SELECT * FROM t LIMIT 5", 1024); got != "SELECT * FROM t LIMIT 5" { + t.Errorf("existing limit must be kept: %q", got) + } + if got := boundRunSQL("UPDATE t SET x=1", 1024); got != "UPDATE t SET x=1" { + t.Errorf("non-select must be untouched: %q", got) + } +} + +func TestMetadataTableDDL(t *testing.T) { + stmts := metadataTableDDL("ragflow_doc_meta_t1") + joined := strings.Join(stmts, "\n") + mustContain(t, joined, "CREATE TABLE IF NOT EXISTS ragflow_doc_meta_t1") + mustContain(t, joined, "meta_fields JSON") + mustContain(t, joined, "id VARCHAR PRIMARY KEY") +} + +func TestGetScoresFromKNNStructure(t *testing.T) { + e := &serenedbEngine{} + knn := map[string]interface{}{"hits": map[string]interface{}{"hits": []interface{}{ + map[string]interface{}{"_id": "a", "_score": 1.5}, + map[string]interface{}{"_id": "b", "_score": 0.0}, + }}} + got := e.GetScores(knn) + if got["a"] != 1.5 || got["b"] != 0.0 { + t.Errorf("GetScores = %v", got) + } +} + +func TestBuildFiltersEmptyListNeverMatches(t *testing.T) { + // An empty IN / array list must not emit invalid `col IN ()`; it reduces to + // a never-match predicate so a DELETE/UPDATE cannot widen scope. + if got := buildFilters(map[string]interface{}{"doc_id": []interface{}{}}); len(got) != 1 || got[0] != neverMatch { + t.Errorf("empty IN filter = %v", got) + } + if got := buildFilters(map[string]interface{}{"tag_kwd": []interface{}{}}); len(got) != 1 || got[0] != neverMatch { + t.Errorf("empty array filter = %v", got) + } +} + +func TestUnrecognizedFilterKeys(t *testing.T) { + unknown := unrecognizedFilterKeys(map[string]interface{}{ + "doc_id": "x", "exists": "img_id", "must_not": nil, "tag_kwd": "t", "bogus": "y", + }) + if len(unknown) != 1 || unknown[0] != "bogus" { + t.Errorf("unrecognized keys = %v, want [bogus]", unknown) + } +} + +func TestValidIdentifier(t *testing.T) { + for _, ok := range []string{"ragflow_t1", "ragflow_doc_meta_abc", "_x"} { + if !validIdentifier(ok) { + t.Errorf("%q should be valid", ok) + } + } + for _, bad := range []string{"a; DROP TABLE t", "1abc", "a b", "a-b", "", "t';--"} { + if validIdentifier(bad) { + t.Errorf("%q should be invalid", bad) + } + } +} + +func TestRedactDSN(t *testing.T) { + mustContain(t, redactDSN("host=h user=u password=secret sslmode=disable"), "password=***") + mustNotContain(t, redactDSN("host=h password=secret"), "secret") + // URL form (SERENEDB_DSN) + got := redactDSN("postgresql://u:secret@host:7890/db") + mustContain(t, got, "://u:***@host") + mustNotContain(t, got, "secret") +} + +func TestQuoteDSNValue(t *testing.T) { + if got := quoteDSNValue("p'a ss"); got != `'p\'a ss'` { + t.Errorf("quoteDSNValue = %q", got) + } + if got := quoteDSNValue(`a\b`); got != `'a\\b'` { + t.Errorf("quoteDSNValue backslash = %q", got) + } +} + +func TestBuildDSNSSLMode(t *testing.T) { + t.Setenv("SERENEDB_DSN", "") // force the config path, not the env override + def := buildDSN(config.SereneDBConfig{Host: "h", Port: 7890, User: "u"}) + mustContain(t, def, "sslmode='disable'") // safe default preserved + enc := buildDSN(config.SereneDBConfig{Host: "h", User: "u", SSLMode: "require"}) + mustContain(t, enc, "sslmode='require'") +} diff --git a/internal/engine/serenedb/sql.go b/internal/engine/serenedb/sql.go new file mode 100644 index 0000000000..254ab280d6 --- /dev/null +++ b/internal/engine/serenedb/sql.go @@ -0,0 +1,58 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package serenedb + +import ( + "context" + "fmt" + "regexp" + "strings" +) + +const runSQLDefaultLimit = 1024 + +var ( + selectPrefixRe = regexp.MustCompile(`(?i)^(select|with)\b`) + hasLimitRe = regexp.MustCompile(`(?i)\blimit\b`) +) + +// boundRunSQL trims a statement and appends a default LIMIT to unbounded +// SELECT/WITH queries so the text-to-SQL path cannot stream an entire table. +// Pure for tests. +func boundRunSQL(sqlText string, limit int) string { + txt := strings.TrimSpace(sqlText) + txt = strings.TrimRight(txt, ";") + if limit > 0 && selectPrefixRe.MatchString(txt) && !hasLimitRe.MatchString(txt) { + txt = fmt.Sprintf("%s LIMIT %d", txt, limit) + } + return txt +} + +// RunSQL executes a text-to-SQL query directly against SereneDB and returns the +// rows as maps. SereneDB speaks SQL natively, so unlike the Infinity engine no +// psql subprocess or field-alias rewrite is needed. +func (e *serenedbEngine) RunSQL(ctx context.Context, tableName string, sqlText string, kbIDs []string, format string) ([]map[string]interface{}, error) { + query := boundRunSQL(sqlText, runSQLDefaultLimit) + rows, err := e.queryMaps(ctx, query) + if err != nil { + return nil, fmt.Errorf("serenedb: run sql: %w", err) + } + if rows == nil { + rows = []map[string]interface{}{} + } + return rows, nil +} diff --git a/internal/server/config/doc_engine_config.go b/internal/server/config/doc_engine_config.go index 6a929a4a9b..30c1232e48 100644 --- a/internal/server/config/doc_engine_config.go +++ b/internal/server/config/doc_engine_config.go @@ -23,6 +23,7 @@ import ( type DocEngineConfig struct { ES ElasticsearchConfig `mapstructure:"es"` Infinity InfinityConfig `mapstructure:"infinity"` + SereneDB SereneDBConfig `mapstructure:"serenedb"` } // ElasticsearchConfig Elasticsearch configuration @@ -41,12 +42,66 @@ type InfinityConfig struct { DocMetaMappingFileName string `mapstructure:"doc_meta_mapping_file_name"` } +// SereneDBConfig SereneDB configuration. SereneDB speaks the PostgreSQL wire +// protocol, so the engine connects with database/sql + lib/pq. +type SereneDBConfig struct { + Host string `mapstructure:"host"` + Port int `mapstructure:"port"` + User string `mapstructure:"user"` + Password string `mapstructure:"password"` + DBName string `mapstructure:"db_name"` + // SSLMode is the lib/pq sslmode; empty defaults to "disable" for a trusted + // local deployment. Set it (e.g. "require") to encrypt the connection. + SSLMode string `mapstructure:"ssl_mode"` +} + func (c *Config) ParseDocEngineConfig(v *viper.Viper) error { c.parseInfinityConfig(v) c.parseElasticsearchConfig(v) + c.parseSereneDBConfig(v) return nil } +func (c *Config) parseSereneDBConfig(v *viper.Viper) { + // Default SereneDB config + c.docEngine.SereneDB.Host = "localhost" + c.docEngine.SereneDB.Port = 5432 + c.docEngine.SereneDB.User = "postgres" + c.docEngine.SereneDB.DBName = "default_db" + + if !v.IsSet("serenedb") { + return + } + sub := v.Sub("serenedb") + if sub == nil { + return + } + + if sub.IsSet("host") { + c.docEngine.SereneDB.Host = sub.GetString("host") + } + + if sub.IsSet("port") { + c.docEngine.SereneDB.Port = sub.GetInt("port") + } + + if sub.IsSet("user") { + c.docEngine.SereneDB.User = sub.GetString("user") + } + + if sub.IsSet("password") { + c.docEngine.SereneDB.Password = sub.GetString("password") + } + + if sub.IsSet("db_name") { + c.docEngine.SereneDB.DBName = sub.GetString("db_name") + } + + if sub.IsSet("ssl_mode") { + c.docEngine.SereneDB.SSLMode = sub.GetString("ssl_mode") + } +} + func (c *Config) parseInfinityConfig(v *viper.Viper) { // Default Infinity config c.docEngine.Infinity.URI = "localhost:23817" @@ -142,3 +197,7 @@ func (i InfinityConfig) ExportConfigs() map[string]interface{} { infinityConfigs["doc_meta_mapping_file_name"] = i.DocMetaMappingFileName return infinityConfigs } + +func (c *Config) GetSereneDBConfig() SereneDBConfig { + return c.docEngine.SereneDB +} diff --git a/rag/app/table.py b/rag/app/table.py index bff9bcd613..feac9524b9 100644 --- a/rag/app/table.py +++ b/rag/app/table.py @@ -490,7 +490,7 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_TASK_PAGE_NUMBER, clmns_map = [(py_clmns[i].lower() + fields_map[clmn_tys[i]], str(clmns[i]).replace("_", " ")) for i in range(len(clmns))] # field_map: only columns stored in chunk_data (metadata or both) — used for retrieval/SQL stored_indices = [i for i in range(len(clmns)) if column_roles.get(clmns[i], "both") in ("metadata", "both")] - if settings.DOC_ENGINE_INFINITY or settings.DOC_ENGINE_OCEANBASE: + if settings.DOC_ENGINE_INFINITY or settings.DOC_ENGINE_OCEANBASE or settings.DOC_ENGINE_SERENEDB: field_map = {py_clmns[i].lower(): str(clmns[i]).replace("_", " ") for i in stored_indices} else: field_map = {clmns_map[i][0]: clmns_map[i][1] for i in stored_indices} @@ -551,7 +551,7 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_TASK_PAGE_NUMBER, if role in ("indexing", "vectorize", "both"): text_fields.append((col_name, row[col_name])) if role in ("metadata", "both"): - if settings.DOC_ENGINE_INFINITY or settings.DOC_ENGINE_OCEANBASE: + if settings.DOC_ENGINE_INFINITY or settings.DOC_ENGINE_OCEANBASE or settings.DOC_ENGINE_SERENEDB: stored[str(col_name)] = row[col_name] else: fld = clmns_map[j][0] @@ -565,7 +565,7 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_TASK_PAGE_NUMBER, stored[f"{py_clmns[j].lower()}_raw"] = raw_s if not text_fields and not stored: continue - if settings.DOC_ENGINE_INFINITY or settings.DOC_ENGINE_OCEANBASE: + if settings.DOC_ENGINE_INFINITY or settings.DOC_ENGINE_OCEANBASE or settings.DOC_ENGINE_SERENEDB: if stored: d["chunk_data"] = stored else: @@ -576,7 +576,7 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_TASK_PAGE_NUMBER, logger.debug(f"[TABLE_PARSER_DEBUG] Chunk content_with_weight length: {len(d.get('content_with_weight', '') or '')}") _cd = d.get("chunk_data") logger.debug(f"[TABLE_PARSER_DEBUG] Chunk chunk_data keys: {list(_cd.keys()) if isinstance(_cd, dict) else 'N/A'}") - if not (settings.DOC_ENGINE_INFINITY or settings.DOC_ENGINE_OCEANBASE): + if not (settings.DOC_ENGINE_INFINITY or settings.DOC_ENGINE_OCEANBASE or settings.DOC_ENGINE_SERENEDB): _extra = [k for k in d if k not in ("docnm_kwd", "title_tks", "content_with_weight", "content_ltks", "content_sm_ltks")] logger.debug(f"[TABLE_PARSER_DEBUG] Chunk ES extra field keys (sample): {_extra[:20]}") res.append(d) diff --git a/rag/nlp/search.py b/rag/nlp/search.py index 4c46b5f5fe..32940fdff5 100644 --- a/rag/nlp/search.py +++ b/rag/nlp/search.py @@ -204,7 +204,7 @@ class Dealer: # citations (see Dealer.fetch_chunk_vectors). OceanBase # still relies on local rerank against chunk vectors, so # keep pulling them for that backend. - if settings.DOC_ENGINE_OCEANBASE: + if settings.DOC_ENGINE_OCEANBASE or settings.DOC_ENGINE_SERENEDB: src.append(f"q_{len(q_vec)}_vec") fusionExpr = FusionExpr("weighted_sum", topk, {"weights": "0.001,1"}) @@ -629,7 +629,7 @@ class Dealer: sim = [s if s is not None else 0.0 for s in sim] tsim = sim vsim = sim - elif settings.DOC_ENGINE_OCEANBASE: + elif settings.DOC_ENGINE_OCEANBASE or settings.DOC_ENGINE_SERENEDB: # OceanBase still returns chunk vectors in the result; use # the historical local rerank that depends on them. sim, tsim, vsim = self.rerank( diff --git a/rag/utils/serenedb_conn.py b/rag/utils/serenedb_conn.py new file mode 100644 index 0000000000..d983eba2b0 --- /dev/null +++ b/rag/utils/serenedb_conn.py @@ -0,0 +1,719 @@ +# +# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""SereneDB DocStoreConnection for RAGFlow. + +Shape follows ob_conn.py (the plain-SQL sibling): ONE table per tenant index +(`ragflow_{tenant_id}`), kb_id is a filtered column, ES field names are kept verbatim so the +read path needs no renames (unlike infinity_conn). The engine-specific notes below were +verified against SereneDB 26.07.4: + + P1 ONE inverted index carries several text columns + the vector; `@@` works per column and + OR across columns SUMS the per-column BM25 scores (natural field boosting — short + title/keyword columns score higher via length norm). + P2 VARCHAR[] columns + list_contains() filters work on base table and index relation; + unnest() powers tag aggregation. + P3 Weighted hybrid fusion is ONE SQL statement: BM25 branch normalized with + `s / MAX(s) OVER ()` (window fn — on <26.07.4 a second CTE reference trips the + iresearch_scan plan-copy bug; FIXED in 26.07.4 #962, but the window form is simpler AND + version-portable so we keep it), FULL OUTER JOIN with the vector branch, weighted sum + pagerank. + P4 Writes are immediately durable in the base table; the inverted index refreshes + asynchronously (~1s converged in probe) — same contract as Elasticsearch's + refresh_interval, which RAGFlow already tolerates. + P6 psycopg2 returns native Python lists for arrays/vectors and dicts for JSON columns. + P7 update()'s add/remove semantics map to list_append() / array_remove(). + +MINIMUM ENGINE VERSION: SereneDB 26.07.4. The vector-branch and fusion queries use the natural +forms that rely on the 26.07.4 fixes #964 (vector-op predicate in an ANN scan's WHERE) and #962 +(multi-reference index CTE); on <26.07.4 both silently returned empty. Verified against the +released image 2026-07-23. + +BM25 REQUIRES the dictionary to declare `frequency = true, norm = true` — without frequency the +scorer silently returns 0.0 for every row. Vectors are stored raw (ES mirror) AND as a +normalized shadow column indexed with metric='ip', quant='sq8': bge-m3 norms are 0.935-0.976, +not 1.0, so ip on unit vectors is the only way to get exact cosine AND quantization. + +Deliberate first-cut simplifications (documented, revisit on in-app eval): + - ES per-field boosts (^10 etc.) are approximated by the P1 summed-OR match over + title/important/question/content token columns; content-only scoring reached ES-parity MRR + on the gold eval, so this is headroom, not debt. + - minimum_should_match is not enforced (OR semantics); IDF-noise is tamed by BM25 itself. + - rank_feature tag boosting is skipped (parity: ob_conn TODOs it as well); pagerank IS applied. +""" + +import json +import logging +import os +import re +import threading +import time +from dataclasses import dataclass, field +from urllib.parse import quote + +import psycopg2 +import psycopg2.extras +import psycopg2.pool + +from common.doc_store.doc_store_base import ( + DocStoreConnection, + FusionExpr, + MatchDenseExpr, + MatchExpr, + MatchTextExpr, + OrderByExpr, +) + +logger = logging.getLogger("ragflow.serenedb_conn") + +PAGERANK_FLD = "pagerank_fea" +ATTEMPT_TIME = 2 + +vector_column_pattern = re.compile(r"q_(?P\d+)_vec$") + +# Column type map — ES mapping names verbatim. psycopg2 adapts python lists to arrays and +# json.dumps handles JSON columns on the write path; the read path gets native types (P6). +TEXT_COLUMNS = [ + "docnm_kwd", + "doc_type_kwd", + "title_tks", + "title_sm_tks", + "content_with_weight", + "content_ltks", + "content_sm_ltks", + "important_tks", + "question_tks", + "create_time", + "img_id", + "knowledge_graph_kwd", + "entity_kwd", + "entity_type_kwd", + "from_entity_kwd", + "to_entity_kwd", + "removed_kwd", + "raptor_kwd", + "group_id", + "mom_id", + "n_hop_with_weight", +] +ARRAY_COLUMNS = ["important_kwd", "question_kwd", "tag_kwd", "source_id", "entities_kwd"] +INT_COLUMNS = ["pagerank_fea", "available_int", "weight_int", "raptor_layer_int", "_order_id"] +FLOAT_COLUMNS = ["create_timestamp_flt", "weight_flt", "rank_flt"] +JSON_COLUMNS = ["tag_feas", "position_int", "page_num_int", "top_int", "chunk_data", "metadata", "extra", "meta_fields"] + +COLUMN_DDL: dict[str, str] = { + "id": "VARCHAR PRIMARY KEY", + "kb_id": "VARCHAR", + "doc_id": "VARCHAR", + **{c: "TEXT" for c in TEXT_COLUMNS}, + **{c: "VARCHAR[]" for c in ARRAY_COLUMNS}, + **{c: "INTEGER" for c in INT_COLUMNS}, + **{c: "DOUBLE PRECISION" for c in FLOAT_COLUMNS}, + **{c: "JSON" for c in JSON_COLUMNS}, +} +COLUMN_NAMES = list(COLUMN_DDL.keys()) + +# Text columns the inverted index carries (all indexable, used for @@ existence filters). +FTS_COLUMNS = ["title_tks", "important_tks", "question_tks", "content_ltks"] +# The SCORED lexical branch matches this ONE column. `ORDER BY BM25(idx.tableoid)` over a +# multi-column `@@` OR returns EMPTY (the WAND top-k iterator can't score a cross-column +# disjunction — same family as the other silent-empty scorer bugs). content_ltks is the +# dominant field and is what the parity eval scored on; ES-style field boosts (docnm^10 etc.) +# are deferred — reintroducing them needs per-column BM25 summed in Python, not an OR, +# precisely because of this bug. +LEX_SCORED_COL = "content_ltks" + +DOC_META_DDL = {"id": "VARCHAR PRIMARY KEY", "kb_id": "VARCHAR", "meta_fields": "JSON"} + +DEFAULTS = {"available_int": 1, "removed_kwd": "N", "_order_id": 0} + +DICTIONARY_NAME = "rf_scored_delim" +# frequency/norm are what make BM25() score at all — see module docstring. +DICTIONARY_DDL = f"CREATE TEXT SEARCH DICTIONARY IF NOT EXISTS {DICTIONARY_NAME} (template = 'delimiter', delimiter = ' ', frequency = true, position = true, norm = true)" + + +def _index_relation(table_name: str) -> str: + return f"idx_{table_name}" + + +def _norm_column(vector_size: int) -> str: + return f"q_{vector_size}_vec_n" + + +def _escape(value) -> str: + """SQL-literal encoding for the templated search statements (filters, aggregation).""" + if value is None: + return "NULL" + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, (int, float)): + return str(value) + if isinstance(value, (list, dict)): + return "'" + json.dumps(value, ensure_ascii=False).replace("'", "''") + "'" + return "'" + str(value).replace("'", "''") + "'" + + +def _strip_es_query(matching_text: str) -> str: + """Fallback tokenization when extra_options lacks original_query: strip ES query_string + syntax (boosts, quotes, boolean sugar) down to plain space-separated tokens for `@@`.""" + txt = re.sub(r"\^[0-9.]+", " ", matching_text) + txt = re.sub(r'["()~*?:+\-]|\bAND\b|\bOR\b|\bNOT\b', " ", txt) + toks = [t for t in txt.split() if t] + seen, out = set(), [] + for t in toks: + if t not in seen: + seen.add(t) + out.append(t) + return " ".join(out) + + +def _l2_normalize(vec: list[float]) -> list[float]: + s = sum(v * v for v in vec) ** 0.5 + if s == 0: + return list(vec) + return [v / s for v in vec] + + +@dataclass +class SearchResult: + total: int = 0 + chunks: list[dict] = field(default_factory=list) + + +class SereneDBConnection(DocStoreConnection): + _instance = None + _instance_lock = threading.Lock() + + def __new__(cls, *args, **kwargs): + with cls._instance_lock: + if cls._instance is None: + cls._instance = super().__new__(cls) + return cls._instance + + def __init__(self): + if getattr(self, "_initialized", False): + return + self._initialized = True + dsn = None + try: + from common.settings import get_base_config # inside RAGFlow + + cfg = get_base_config("serenedb", {}) or {} + host = cfg.get("host", "serenedb") + port = int(cfg.get("port", 7890)) + user = cfg.get("user", "postgres") + password = cfg.get("password", "") + dbname = cfg.get("db_name", "postgres") + # URL-encode credentials so a password with @ / : / space does not + # corrupt the DSN. + dsn = f"postgresql://{quote(user, safe='')}:{quote(password, safe='')}@{host}:{port}/{quote(dbname, safe='')}" + ssl_mode = cfg.get("ssl_mode") + if ssl_mode: + dsn += f"?sslmode={quote(str(ssl_mode), safe='')}" + except Exception: + pass + dsn = os.environ.get("SERENEDB_DSN", dsn or "postgresql://postgres@serenedb:7890/") + self._pool = psycopg2.pool.ThreadedConnectionPool(minconn=1, maxconn=8, dsn=dsn) + self._known_tables: set[str] = set() + self._known_lock = threading.Lock() + self._dsn_display = re.sub(r":[^:@/]+@", ":***@", dsn) + logger.info(f"SereneDB connection initialized: {self._dsn_display}") + + def _run(self, sql: str, params=None, fetch: bool = True): + conn = self._pool.getconn() + try: + conn.autocommit = True + with conn.cursor() as cur: + cur.execute(sql, params) + if fetch and cur.description is not None: + return cur.fetchall(), [d[0] for d in cur.description] + return [], [] + finally: + self._pool.putconn(conn) + + """ + Database operations + """ + + def db_type(self) -> str: + return "serenedb" + + def health(self) -> dict: + rows, _ = self._run("SELECT version()") + return {"type": "serenedb", "status": "green", "dsn": self._dsn_display, "version": rows[0][0] if rows else "unknown"} + + """ + Table operations + """ + + def create_idx(self, index_name: str, dataset_id: str, vector_size: int, parser_id: str = None): + if index_name.startswith("ragflow_doc_meta_"): + cols = ", ".join(f"{k} {t}" for k, t in DOC_META_DDL.items()) + self._run(f"CREATE TABLE IF NOT EXISTS {index_name} ({cols})", fetch=False) + return + cols = ", ".join(f"{k} {t}" for k, t in COLUMN_DDL.items()) + vec, vec_n = f"q_{vector_size}_vec", _norm_column(vector_size) + self._run(f"CREATE TABLE IF NOT EXISTS {index_name} ({cols}, {vec} FLOAT[{vector_size}], {vec_n} FLOAT[{vector_size}])", fetch=False) + self._run(DICTIONARY_DDL, fetch=False) + fts = ", ".join(f"{c} {DICTIONARY_NAME}" for c in FTS_COLUMNS) + self._run( + f"CREATE INDEX IF NOT EXISTS {_index_relation(index_name)} ON {index_name} " + f"USING inverted (id, {fts}, {vec_n} ivf (metric = 'ip', quant = 'sq8')) " + f"WITH (optimize_top_k = 'bm25(1.2, 0.75)')", + fetch=False, + ) + with self._known_lock: + self._known_tables.add(index_name) + + def create_doc_meta_idx(self, index_name: str): + # RAGFlow calls this directly for the per-tenant metadata table (not in the ABC, but the + # retriever/document service expects it — same as ob_conn_base.create_doc_meta_idx). + self.create_idx(index_name, None, 0) + return True + + def delete_idx(self, index_name: str, dataset_id: str): + if dataset_id and not index_name.startswith("ragflow_doc_meta_"): + # all KBs of a tenant share one table; a KB deletion must not drop the index + return + self._run(f"DROP INDEX IF EXISTS {_index_relation(index_name)}", fetch=False) + self._run(f"DROP TABLE IF EXISTS {index_name}", fetch=False) + with self._known_lock: + self._known_tables.discard(index_name) + + def index_exist(self, index_name: str, dataset_id: str = None) -> bool: + if index_name in self._known_tables: + return True + try: + self._run(f"SELECT 1 FROM {index_name} LIMIT 0") + with self._known_lock: + self._known_tables.add(index_name) + return True + except Exception: + return False + + """ + Filters + """ + + def _get_filters(self, condition: dict) -> list[str]: + filters = [] + for k, v in condition.items(): + if not v: + continue + if k == "exists": + if v in COLUMN_DDL: + filters.append(f"{v} IS NOT NULL") + elif k == "must_not" and isinstance(v, dict) and "exists" in v: + if v["exists"] in COLUMN_DDL: + filters.append(f"{v['exists']} IS NULL") + elif k in ARRAY_COLUMNS: + vals = v if isinstance(v, list) else [v] + ors = " OR ".join(f"list_contains({k}, {_escape(x)})" for x in vals) + filters.append(f"({ors})") + elif k in COLUMN_DDL or vector_column_pattern.match(k): + if isinstance(v, list): + filters.append(f"{k} IN ({', '.join(_escape(x) for x in v)})") + else: + filters.append(f"{k} = {_escape(v)}") + return filters + + """ + CRUD + """ + + def search( + self, + select_fields: list[str], + highlight_fields: list[str], + condition: dict, + match_expressions: list[MatchExpr], + order_by: OrderByExpr, + offset: int, + limit: int, + index_names: str | list[str], + dataset_ids: list[str], + agg_fields: list[str] | None = None, + rank_feature: dict | None = None, + **kwargs, + ) -> SearchResult: + if isinstance(index_names, str): + index_names = index_names.split(",") + agg_fields = agg_fields or [] + + output_fields = [f for f in (select_fields or []) if f != "_score"] + if not output_fields or "*" in output_fields: + output_fields = COLUMN_NAMES.copy() + if "id" not in output_fields: + output_fields = ["id"] + output_fields + for f in highlight_fields or []: + if f not in output_fields: + output_fields.append(f) + output_fields = [f for f in output_fields if f in COLUMN_DDL or vector_column_pattern.match(f)] + if PAGERANK_FLD not in output_fields: + output_fields.append(PAGERANK_FLD) + fields_expr = ", ".join(output_fields) + + condition = dict(condition or {}) + condition["kb_id"] = dataset_ids + filters = self._get_filters(condition) + filters_expr = " AND ".join(filters) if filters else "TRUE" + + text_query = text_topn = None + vec_col = vec_data = vec_topn = None + vec_threshold = 0.0 + vector_weight = 0.5 + for m in match_expressions: + if isinstance(m, MatchTextExpr): + # matching_text is RAGFlow's TOKENIZED, ^-weighted query_string (e.g. + # "(auto^0.5) (_^0.3) (ptr^0.4) "auto _"^0.9 ..."). The stored *_ltks columns are + # tokenized the same way, so `@@` must see those tokens — the raw human question + # in original_query would miss anything the tokenizer splits (auto_ptr -> auto _ ptr). + text_query = _strip_es_query(m.matching_text or (m.extra_options or {}).get("original_query", "")) + text_topn = m.topn + elif isinstance(m, MatchDenseExpr): + vec_col = m.vector_column_name + vec_data = list(m.embedding_data) + vec_topn = m.topn + vec_threshold = float((m.extra_options or {}).get("similarity", 0.0)) + elif isinstance(m, FusionExpr): + if m.method == "weighted_sum" and "weights" in (m.fusion_params or {}): + vector_weight = float(m.fusion_params["weights"].split(",")[1]) + + result = SearchResult() + pagerank_expr = f"COALESCE({PAGERANK_FLD}, 0) / 100.0" + + for index_name in index_names: + if not self.index_exist(index_name): + continue + idx = _index_relation(index_name) + t0 = time.time() + + if text_query and vec_data: + rows = self._fusion_search(index_name, idx, fields_expr, output_fields, filters_expr, text_query, text_topn, vec_col, vec_data, vec_topn, vec_threshold, vector_weight, offset, limit) + search_type = "fusion" + elif text_query: + match = f"{LEX_SCORED_COL} @@ {_escape(text_query)}" + n = limit if limit > 0 else (text_topn or 10000) + rows, _ = self._run( + f"SELECT {fields_expr}, BM25({idx}.tableoid) + {pagerank_expr} AS _score FROM {idx} WHERE {filters_expr} AND ({match}) ORDER BY _score DESC LIMIT {n} OFFSET {offset}" + ) + search_type = "fulltext" + elif vec_data: + # Similarity threshold goes straight in the ANN scan's WHERE (relies on the + # 26.07.4 fix #964 — on <26.07.4 a vector-op predicate here silently emptied the + # result and had to be applied outside the scan). + vec_n = _norm_column(len(vec_data)) + qv = "ARRAY[" + ",".join(str(float(x)) for x in _l2_normalize(vec_data)) + f"]::FLOAT[{len(vec_data)}]" + n = limit if limit > 0 else (vec_topn or 10) + rows, _ = self._run( + f"SELECT {fields_expr}, -({vec_n} <#> {qv}) + {pagerank_expr} AS _score " + f"FROM {idx} WHERE {filters_expr} AND -({vec_n} <#> {qv}) >= {vec_threshold} " + f"ORDER BY {vec_n} <#> {qv} LIMIT {n} OFFSET {offset}" + ) + search_type = "vector" + elif agg_fields: + self._aggregation(result, index_name, agg_fields, filters_expr) + logger.info(f"SereneDB search {index_name} type=aggregation took={time.time() - t0:.3f}s groups={result.total}") + continue + else: + orders = [] + for f, o in order_by.fields if order_by else []: + if f in COLUMN_DDL: + orders.append(f"{f} {'ASC' if o == 0 else 'DESC'}") + order_expr = ("ORDER BY " + ", ".join(orders)) if orders else "" + limit_expr = f"LIMIT {limit} OFFSET {offset}" if limit else "" + crows, _ = self._run(f"SELECT count(*) FROM {index_name} WHERE {filters_expr}") + result.total += crows[0][0] + rows, _ = self._run(f"SELECT {fields_expr} FROM {index_name} WHERE {filters_expr} {order_expr} {limit_expr}") + for row in rows: + result.chunks.append(self._row_to_entity(row, output_fields)) + logger.info(f"SereneDB search {index_name} type=filter took={time.time() - t0:.3f}s rows={len(rows)}") + continue + + for row in rows: + result.chunks.append(self._row_to_entity(row, output_fields + ["_score"])) + logger.info(f"SereneDB search {index_name} type={search_type} took={time.time() - t0:.3f}s rows={len(rows)} q={text_query!r}") + + if result.total == 0: + result.total = len(result.chunks) + return result + + def _fusion_search(self, table, idx, fields_expr, output_fields, filters_expr, text_query, text_topn, vec_col, vec_data, vec_topn, vec_threshold, vector_weight, offset, limit): + """The P3 shape: one statement, window-normalized BM25 branch (a scalar-subquery + normalizer would trip the iresearch_scan plan-copy bug), FULL OUTER JOIN, weighted sum. + RAGFlow parity math: (1-vw) * bm25_norm + vw * cosine + pagerank/100.""" + vec_n = _norm_column(len(vec_data)) + qv = "ARRAY[" + ",".join(str(float(x)) for x in _l2_normalize(vec_data)) + f"]::FLOAT[{len(vec_data)}]" + match = f"{LEX_SCORED_COL} @@ {_escape(text_query)}" + lex_n = text_topn or 200 + v_n = vec_topn or 200 + n = limit if limit > 0 else (lex_n + v_n) + prefixed = ", ".join(f"t.{f}" for f in output_fields) + sql = f""" +WITH lex AS ( + SELECT id, BM25({idx}.tableoid) AS s + FROM {idx} WHERE {filters_expr} AND ({match}) + ORDER BY s DESC LIMIT {lex_n}), +lexn AS (SELECT id, s / NULLIF(MAX(s) OVER (), 0) AS sn FROM lex), +vec AS ( + SELECT id, -({vec_n} <#> {qv}) AS sim + FROM {idx} WHERE {filters_expr} AND -({vec_n} <#> {qv}) >= {vec_threshold} + ORDER BY {vec_n} <#> {qv} LIMIT {v_n}), +fused AS ( + SELECT COALESCE(l.id, v.id) AS id, + COALESCE(l.sn, 0) * {1.0 - vector_weight} + COALESCE(v.sim, 0) * {vector_weight} AS fs + FROM lexn l FULL OUTER JOIN vec v ON l.id = v.id) +SELECT {prefixed}, f.fs + COALESCE(t.{PAGERANK_FLD}, 0) / 100.0 AS _score +FROM fused f JOIN {table} t ON t.id = f.id +ORDER BY _score DESC LIMIT {n} OFFSET {offset}""" + rows, _ = self._run(sql) + return rows + + def _aggregation(self, result: SearchResult, index_name: str, agg_fields, filters_expr): + for agg_field in agg_fields: + if agg_field not in COLUMN_DDL: + # agg_field is interpolated into SQL; only aggregate real columns. + continue + if agg_field in ARRAY_COLUMNS: + rows, _ = self._run(f"SELECT u.v, count(*) FROM (SELECT unnest({agg_field}) AS v FROM {index_name} WHERE {filters_expr} AND {agg_field} IS NOT NULL) u GROUP BY u.v") + else: + rows, _ = self._run(f"SELECT {agg_field}, count(*) FROM {index_name} WHERE {filters_expr} AND {agg_field} IS NOT NULL GROUP BY {agg_field}") + for value, count in rows: + result.chunks.append({"value": value, "count": int(count)}) + result.total += 1 + + def get(self, chunk_id: str, index_name: str, dataset_ids: list[str]) -> dict | None: + if not self.index_exist(index_name): + return None + rows, cols = self._run(f"SELECT * FROM {index_name} WHERE id = %s", (chunk_id,)) + if not rows: + return None + return self._row_to_entity(rows[0], cols) + + def insert(self, rows: list[dict], index_name: str, dataset_id: str = None) -> list[str]: + if not rows: + return [] + if index_name.startswith("ragflow_doc_meta_"): + return self._insert_doc_meta(rows, index_name) + if not self.index_exist(index_name): + size = 0 + for k in rows[0]: + m = vector_column_pattern.match(k) + if m: + size = int(m.group("vector_size")) + self.create_idx(index_name, dataset_id, size or 1024) + + # Batch by identical column tuple and multi-row upsert; per-row INSERTs would make + # bulk ingest interminable. + errors = [] + groups: dict[tuple, list[list]] = {} + for doc in rows: + d, extra = {}, {} + vec_cols = {} + for k, v in doc.items(): + m = vector_column_pattern.match(k) + if m: + vec_cols[k] = v + continue + if k not in COLUMN_DDL: + extra[k] = v + continue + if k == "kb_id" and isinstance(v, list): + v = v[0] + if k in JSON_COLUMNS and not isinstance(v, str): + v = json.dumps(v, ensure_ascii=False) + if k == "content_with_weight" and isinstance(v, dict): + v = json.dumps(v, ensure_ascii=False) + d[k] = v + if extra: + merged = d.get("extra") + base = json.loads(merged) if isinstance(merged, str) and merged else {} + base.update(extra) + d["extra"] = json.dumps(base, ensure_ascii=False) + for k, dv in DEFAULTS.items(): + d.setdefault(k, dv) + + cols, vals = list(d.keys()), list(d.values()) + for vc, vv in vec_cols.items(): + size = int(vector_column_pattern.match(vc).group("vector_size")) + cols += [vc, _norm_column(size)] + vals += [vv, _l2_normalize(vv)] + groups.setdefault(tuple(cols), []).append(vals) + + conn = self._pool.getconn() + try: + conn.autocommit = True + with conn.cursor() as cur: + for cols, val_rows in groups.items(): + updates = ", ".join(f"{c} = EXCLUDED.{c}" for c in cols if c != "id") + try: + psycopg2.extras.execute_values(cur, f"INSERT INTO {index_name} ({', '.join(cols)}) VALUES %s ON CONFLICT (id) DO UPDATE SET {updates}", val_rows, page_size=500) + except Exception as e: + logger.error(f"SereneDB insert error on {index_name}: {e}") + errors.append(str(e)) + finally: + self._pool.putconn(conn) + return errors + + def _insert_doc_meta(self, rows: list[dict], index_name: str) -> list[str]: + if not self.index_exist(index_name): + self.create_idx(index_name, None, 0) + errors = [] + for doc in rows: + meta = doc.get("meta_fields") or {} + if not isinstance(meta, str): + meta = json.dumps(meta, ensure_ascii=False) + try: + self._run( + f"INSERT INTO {index_name} (id, kb_id, meta_fields) VALUES (%s, %s, %s) ON CONFLICT (id) DO UPDATE SET kb_id = EXCLUDED.kb_id, meta_fields = EXCLUDED.meta_fields", + (doc.get("id"), doc.get("kb_id"), meta), + fetch=False, + ) + except Exception as e: + errors.append(str(e)) + return errors + + def update(self, condition: dict, new_value: dict, index_name: str, dataset_id: str) -> bool: + if not self.index_exist(index_name): + return True + condition = dict(condition or {}) + if not index_name.startswith("ragflow_doc_meta_"): + condition["kb_id"] = dataset_id + filters = self._get_filters(condition) + if not filters: + return False + sets = [] + for k, v in new_value.items(): + if k == "remove": + items = {v: None} if isinstance(v, str) else v + for kk, vv in items.items(): + if kk not in COLUMN_DDL: + continue + if vv is None: + sets.append(f"{kk} = NULL") + elif kk in ARRAY_COLUMNS: + sets.append(f"{kk} = array_remove({kk}, {_escape(vv)})") + elif k == "add": + for kk, vv in v.items(): + if kk in ARRAY_COLUMNS: + sets.append(f"{kk} = list_append({kk}, {_escape(vv)})") + elif k in JSON_COLUMNS: + sets.append(f"{k} = {_escape(json.dumps(v, ensure_ascii=False) if not isinstance(v, str) else v)}") + elif k in COLUMN_DDL: + sets.append(f"{k} = {_escape(v)}") + if not sets: + return True + try: + self._run(f"UPDATE {index_name} SET {', '.join(sets)} WHERE {' AND '.join(filters)}", fetch=False) + return True + except Exception as e: + logger.error(f"SereneDB update error on {index_name}: {e}") + return False + + def delete(self, condition: dict, index_name: str, dataset_id: str) -> int: + if not self.index_exist(index_name): + return 0 + condition = dict(condition or {}) + if not index_name.startswith("ragflow_doc_meta_"): + condition["kb_id"] = dataset_id + filters = self._get_filters(condition) + if not filters: + return 0 + where = " AND ".join(filters) + rows, _ = self._run(f"SELECT count(*) FROM {index_name} WHERE {where}") + n = rows[0][0] + if n: + self._run(f"DELETE FROM {index_name} WHERE {where}", fetch=False) + return n + + """ + Result helpers + """ + + def _row_to_entity(self, row, cols) -> dict: + entity = {} + for c, v in zip(cols, row): + if v is None: + continue + if c in JSON_COLUMNS and isinstance(v, str): + try: + v = json.loads(v) + except json.JSONDecodeError: + pass + entity[c] = v + return entity + + def get_scores(self, res: SearchResult) -> dict[str, float]: + # chunk id -> fused/vector/bm25 score. NOT in the ABC, but RAGFlow's retriever calls it to + # recover the first-stage score without re-reading vectors (see es_conn_base.get_scores). + # search() stamps each chunk with "_score"; default 0.0 for filter-only results. + return {c["id"]: float(c.get("_score", 0.0)) for c in res.chunks if "id" in c} + + def get_total(self, res: SearchResult) -> int: + return res.total + + def get_doc_ids(self, res: SearchResult) -> list[str]: + return [c["id"] for c in res.chunks] + + def get_fields(self, res: SearchResult, fields: list[str]) -> dict[str, dict]: + out = {} + for c in res.chunks: + out[c["id"]] = {f: c[f] for f in fields if c.get(f) is not None} + return out + + def get_highlight(self, res: SearchResult, keywords: list[str], field_name: str): + # Same client-side strategy as ob_conn: emphasize keyword hits in the stored text. + ans = {} + if not res.chunks or not keywords: + return ans + pats = [re.compile(r"(^|\W)(%s)(\W|$)" % re.escape(k), re.IGNORECASE | re.MULTILINE) for k in keywords if k] + for c in res.chunks: + txt = c.get(field_name) + if not txt: + continue + marked = txt + for p in pats: + marked = p.sub(r"\1\2\3", marked) + if "" in marked: + ans[c["id"]] = re.sub(r"\s*", " ", marked) + return ans + + def get_aggregation(self, res: SearchResult, field_name: str): + out = [] + counts = {} + for c in res.chunks: + if "value" in c and "count" in c: + out.append((c["value"], c["count"])) + elif field_name in c: + v = c[field_name] + for vv in v if isinstance(v, list) else [v]: + if isinstance(vv, str) and vv.strip(): + counts[vv] = counts.get(vv, 0) + 1 + out.extend(counts.items()) + return out + + """ + SQL passthrough (text-to-SQL feature) + """ + + def sql(self, sql: str, fetch_size: int = 1024, format: str = "json"): + txt = sql.strip().rstrip(";") + if fetch_size and re.match(r"^(select|with)\b", txt, re.IGNORECASE) and not re.search(r"\blimit\b", txt, re.IGNORECASE): + txt = f"{txt} LIMIT {int(fetch_size)}" + try: + rows, cols = self._run(txt) + except Exception: + logger.exception("SereneDB sql passthrough failed") + raise + return {"columns": [{"name": c, "type": "text"} for c in cols], "rows": [list(r) for r in rows]} diff --git a/rag/utils/table_es_metadata.py b/rag/utils/table_es_metadata.py index 7f24e7775c..c06155c225 100644 --- a/rag/utils/table_es_metadata.py +++ b/rag/utils/table_es_metadata.py @@ -210,7 +210,7 @@ def aggregate_table_doc_metadata(chunks: list, task: dict) -> dict: e, exc_info=True, ) - if not fm and not (settings.DOC_ENGINE_INFINITY or settings.DOC_ENGINE_OCEANBASE): + if not fm and not (settings.DOC_ENGINE_INFINITY or settings.DOC_ENGINE_OCEANBASE or settings.DOC_ENGINE_SERENEDB): logging.debug(f"[TABLE_META_DEBUG] field_map empty on task snapshot — will use ES key probe on chunk dicts; kb_parser_config keys={list((task.get('kb_parser_config') or {}).keys())}") logging.debug(f"[TABLE_META_DEBUG] meta_cols={meta_cols}, field_map entries={len(fm)}, infinity={settings.DOC_ENGINE_INFINITY}, oceanbase={settings.DOC_ENGINE_OCEANBASE}") sample_ck = next((c for c in chunks if isinstance(c, dict)), None) @@ -219,7 +219,7 @@ def aggregate_table_doc_metadata(chunks: list, task: dict) -> dict: logging.debug(f"[TABLE_META_DEBUG] first chunk non-vector keys (sample): {sk}") es_col_keys: dict[str, tuple[str | None, str]] = {} - if not (settings.DOC_ENGINE_INFINITY or settings.DOC_ENGINE_OCEANBASE): + if not (settings.DOC_ENGINE_INFINITY or settings.DOC_ENGINE_OCEANBASE or settings.DOC_ENGINE_SERENEDB): for col in meta_cols: tk, src = _resolve_es_chunk_field_key(col, fm, sample_ck) es_col_keys[col] = (tk, src) @@ -230,7 +230,7 @@ def aggregate_table_doc_metadata(chunks: list, task: dict) -> dict: for i, ck in enumerate(chunks): if not isinstance(ck, dict): continue - if settings.DOC_ENGINE_INFINITY or settings.DOC_ENGINE_OCEANBASE: + if settings.DOC_ENGINE_INFINITY or settings.DOC_ENGINE_OCEANBASE or settings.DOC_ENGINE_SERENEDB: cd = ck.get("chunk_data") if not isinstance(cd, dict): continue diff --git a/test/unit_test/rag/utils/test_serenedb_conn.py b/test/unit_test/rag/utils/test_serenedb_conn.py new file mode 100644 index 0000000000..b925512dc7 --- /dev/null +++ b/test/unit_test/rag/utils/test_serenedb_conn.py @@ -0,0 +1,96 @@ +# +# Copyright 2025 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +""" +Unit tests for the pure SQL-building helpers of the SereneDB connector. These +need no live database: the module-level helpers are imported directly, and +_get_filters is a method that does not touch self, so it is called unbound. +""" + +from __future__ import annotations + +from rag.utils.serenedb_conn import ( + SereneDBConnection, + _escape, + _index_relation, + _l2_normalize, + _norm_column, + _strip_es_query, +) + + +class TestEscape: + def test_none(self): + assert _escape(None) == "NULL" + + def test_bool(self): + assert _escape(True) == "true" + assert _escape(False) == "false" + + def test_numbers(self): + assert _escape(7) == "7" + assert _escape(1.5) == "1.5" + + def test_string_is_quoted_and_escaped(self): + assert _escape("a'b") == "'a''b'" + + def test_list_and_dict_serialize_to_json(self): + assert _escape(["x", "y"]) == '\'["x", "y"]\'' + assert _escape({"a": 1}) == "'{\"a\": 1}'" + + +class TestStripESQuery: + def test_boosts_and_syntax_removed_and_deduped(self): + # weights, quotes and boolean sugar are stripped; tokens de-duplicated + # in first-seen order. + assert _strip_es_query('(auto^0.5) (ptr^0.4) "auto _"^0.9 auto') == "auto ptr _" + + def test_plain_query_untouched(self): + assert _strip_es_query("alpha beta") == "alpha beta" + + +class TestVectorHelpers: + def test_l2_normalize_unit_length(self): + assert _l2_normalize([3.0, 4.0]) == [0.6, 0.8] + + def test_l2_normalize_zero_vector_passthrough(self): + assert _l2_normalize([0.0, 0.0]) == [0.0, 0.0] + + def test_norm_and_index_names(self): + assert _norm_column(1024) == "q_1024_vec_n" + assert _index_relation("ragflow_t1") == "idx_ragflow_t1" + + +class TestGetFilters: + # _get_filters uses only module-level state, so it can be called unbound. + def filters(self, condition): + return SereneDBConnection._get_filters(None, condition) + + def test_scalar_equals(self): + assert self.filters({"doc_id": "d1"}) == ["doc_id = 'd1'"] + + def test_list_is_in(self): + assert self.filters({"doc_id": ["a", "b"]}) == ["doc_id IN ('a', 'b')"] + + def test_array_column_uses_list_contains(self): + got = self.filters({"tag_kwd": ["a", "b"]}) + assert got == ["(list_contains(tag_kwd, 'a') OR list_contains(tag_kwd, 'b'))"] + + def test_exists_and_must_not(self): + assert self.filters({"exists": "img_id"}) == ["img_id IS NOT NULL"] + assert self.filters({"must_not": {"exists": "img_id"}}) == ["img_id IS NULL"] + + def test_unknown_and_empty_are_skipped(self): + assert self.filters({"not_a_column": "x", "doc_id": ""}) == []