mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-19 14:55:40 +08:00
c6ba54bc72
Follow-up to #17526 ("Refactor: merge dataset scope graph"), which introduced two code paths that touch Infinity columns the deployed schema does not declare. This PR makes the runtime robust against the old schema while also adding the new column to the new schema so freshly created tables are correct.
977 lines
45 KiB
Python
977 lines
45 KiB
Python
#
|
|
# 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.
|
|
#
|
|
|
|
import re
|
|
import json
|
|
import copy
|
|
from infinity.common import InfinityException, SortType
|
|
from infinity.errors import ErrorCode
|
|
from common.decorator import singleton
|
|
import pandas as pd
|
|
from common.constants import PAGERANK_FLD, TAG_FLD
|
|
from common.doc_store.doc_store_base import MatchExpr, MatchTextExpr, MatchDenseExpr, FusionExpr, OrderByExpr
|
|
from common.doc_store.infinity_conn_base import InfinityConnectionBase
|
|
from common.float_utils import get_float
|
|
|
|
|
|
DENSE_FILTER_FULLTEXT_WEIGHT_THRESHOLD = 0.8
|
|
DEFAULT_VECTOR_SIMILARITY_WEIGHT = 0.5
|
|
_JSON_LIST_FIELDS = frozenset(
|
|
(
|
|
"source_chunk_ids",
|
|
"source_doc_ids",
|
|
"compilation_template_ids",
|
|
"doc_ids_kwd",
|
|
"entity_names_kwd",
|
|
"outlinks_kwd",
|
|
"related_kb_pages_kwd",
|
|
"rechunked_from_chunk_ids",
|
|
)
|
|
)
|
|
|
|
|
|
def _vector_similarity_weight(match_expressions: list[MatchExpr]) -> float:
|
|
vector_similarity_weight = DEFAULT_VECTOR_SIMILARITY_WEIGHT
|
|
for matchExpr in match_expressions:
|
|
if not isinstance(matchExpr, FusionExpr) or matchExpr.method != "weighted_sum":
|
|
continue
|
|
fusion_params = matchExpr.fusion_params or {}
|
|
weights = fusion_params.get("weights")
|
|
if not weights:
|
|
continue
|
|
weight_parts = str(weights).split(",")
|
|
if len(weight_parts) > 1:
|
|
vector_similarity_weight = get_float(weight_parts[1])
|
|
return vector_similarity_weight
|
|
|
|
|
|
def _build_dense_filter(filter_cond: str | None, filter_fulltext: str | None, vector_similarity_weight: float) -> str:
|
|
if vector_similarity_weight > DENSE_FILTER_FULLTEXT_WEIGHT_THRESHOLD:
|
|
return filter_cond or ""
|
|
return filter_fulltext or filter_cond or ""
|
|
|
|
|
|
@singleton
|
|
class InfinityConnection(InfinityConnectionBase):
|
|
"""
|
|
Dataframe and fields convert
|
|
"""
|
|
|
|
@staticmethod
|
|
def field_keyword(field_name: str):
|
|
# Treat "*_kwd" tag-like columns as keyword lists except for the fields in the exclusion list; source_id is also keyword-like.
|
|
# source_doc_ids / source_chunk_ids are multi-valued provenance lists (artifact/wiki rows) and must be
|
|
# stored/read/updated as keyword lists so the delete-time ref-count (remove one id, drop row when empty) works.
|
|
if field_name in ("source_id", "source_doc_ids", "source_chunk_ids") or (
|
|
field_name.endswith("_kwd") and field_name not in ["knowledge_graph_kwd", "docnm_kwd", "important_kwd", "question_kwd", "parent_kwd"]
|
|
):
|
|
return True
|
|
return False
|
|
|
|
def convert_select_fields(self, output_fields: list[str]) -> list[str]:
|
|
need_empty_count = "important_kwd" in output_fields
|
|
for i, field in enumerate(output_fields):
|
|
if field in ["docnm_kwd", "title_tks", "title_sm_tks"]:
|
|
output_fields[i] = "docnm"
|
|
elif field in ["important_kwd", "important_tks"]:
|
|
output_fields[i] = "important_keywords"
|
|
elif field in ["question_kwd", "question_tks"]:
|
|
output_fields[i] = "questions"
|
|
elif field in ["content_with_weight", "content_ltks", "content_sm_ltks"]:
|
|
output_fields[i] = "content"
|
|
elif field in ["authors_tks", "authors_sm_tks"]:
|
|
output_fields[i] = "authors"
|
|
if need_empty_count and "important_kwd_empty_count" not in output_fields:
|
|
output_fields.append("important_kwd_empty_count")
|
|
return list(set(output_fields))
|
|
|
|
@staticmethod
|
|
def convert_matching_field(field_weight_str: str) -> str:
|
|
tokens = field_weight_str.split("^")
|
|
field = tokens[0]
|
|
if field == "docnm_kwd" or field == "title_tks":
|
|
field = "docnm@ft_docnm_rag_coarse"
|
|
elif field == "title_sm_tks":
|
|
field = "docnm@ft_docnm_rag_fine"
|
|
elif field == "important_kwd":
|
|
field = "important_keywords@ft_important_keywords_rag_coarse"
|
|
elif field == "important_tks":
|
|
field = "important_keywords@ft_important_keywords_rag_fine"
|
|
elif field == "question_kwd":
|
|
field = "questions@ft_questions_rag_coarse"
|
|
elif field == "question_tks":
|
|
field = "questions@ft_questions_rag_fine"
|
|
elif field == "content_with_weight" or field == "content_ltks":
|
|
field = "content@ft_content_rag_coarse"
|
|
elif field == "content_sm_ltks":
|
|
field = "content@ft_content_rag_fine"
|
|
elif field == "authors_tks":
|
|
field = "authors@ft_authors_rag_coarse"
|
|
elif field == "authors_sm_tks":
|
|
field = "authors@ft_authors_rag_fine"
|
|
elif field == "tag_kwd":
|
|
field = "tag_kwd@ft_tag_kwd_whitespace__"
|
|
tokens[0] = field
|
|
return "^".join(tokens)
|
|
|
|
"""
|
|
CRUD operations
|
|
"""
|
|
|
|
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],
|
|
knowledgebase_ids: list[str],
|
|
agg_fields: list[str] | None = None,
|
|
rank_feature: dict | None = None,
|
|
) -> tuple[pd.DataFrame, int]:
|
|
"""
|
|
BUG: Infinity returns empty for a highlight field if the query string doesn't use that field.
|
|
"""
|
|
if isinstance(index_names, str):
|
|
index_names = index_names.split(",")
|
|
assert isinstance(index_names, list) and len(index_names) > 0
|
|
inf_conn = self.connPool.get_conn()
|
|
try:
|
|
db_instance = inf_conn.get_database(self.dbName)
|
|
df_list = list()
|
|
table_list = list()
|
|
output = select_fields.copy()
|
|
output = self.convert_select_fields(output)
|
|
if agg_fields is None:
|
|
agg_fields = []
|
|
for essential_field in ["id"] + agg_fields:
|
|
if essential_field not in output:
|
|
output.append(essential_field)
|
|
score_func = ""
|
|
score_column = ""
|
|
for matchExpr in match_expressions:
|
|
if isinstance(matchExpr, MatchTextExpr):
|
|
score_func = "score()"
|
|
score_column = "SCORE"
|
|
break
|
|
if not score_func:
|
|
for matchExpr in match_expressions:
|
|
if isinstance(matchExpr, MatchDenseExpr):
|
|
score_func = "similarity()"
|
|
score_column = "SIMILARITY"
|
|
break
|
|
if match_expressions:
|
|
if score_func and score_func not in output:
|
|
output.append(score_func)
|
|
if PAGERANK_FLD not in output:
|
|
output.append(PAGERANK_FLD)
|
|
output = [f for f in output if f and f != "_score"]
|
|
if limit <= 0:
|
|
# ElasticSearch default limit is 10000
|
|
limit = 10000
|
|
|
|
# Prepare expressions common to all tables
|
|
filter_cond = None
|
|
filter_fulltext = ""
|
|
if condition:
|
|
# For metadata table (ragflow_doc_meta_), keep kb_id filter
|
|
# For chunk tables, remove kb_id filter as they use table separation per KB
|
|
is_meta_table = any(indexName.startswith("ragflow_doc_meta_") for indexName in index_names)
|
|
if not is_meta_table:
|
|
condition = {k: v for k, v in condition.items() if k != "kb_id"}
|
|
|
|
table_found = False
|
|
for indexName in index_names:
|
|
if indexName.startswith("ragflow_doc_meta_"):
|
|
table_names_to_search = [indexName]
|
|
else:
|
|
table_names_to_search = [f"{indexName}_{kb_id}" for kb_id in knowledgebase_ids]
|
|
for table_name in table_names_to_search:
|
|
try:
|
|
filter_cond = self.equivalent_condition_to_str(condition, db_instance.get_table(table_name))
|
|
table_found = True
|
|
break
|
|
except Exception:
|
|
pass
|
|
if table_found:
|
|
break
|
|
if not table_found:
|
|
self.logger.error(f"No valid tables found for indexNames {index_names} and knowledgebaseIds {knowledgebase_ids}")
|
|
return pd.DataFrame(), 0
|
|
|
|
# vector_similarity_weight = _vector_similarity_weight(match_expressions)
|
|
for matchExpr in match_expressions:
|
|
if isinstance(matchExpr, MatchTextExpr):
|
|
if filter_cond and "filter" not in matchExpr.extra_options:
|
|
matchExpr.extra_options.update({"filter": filter_cond})
|
|
matchExpr.fields = [self.convert_matching_field(field) for field in matchExpr.fields]
|
|
fields = ",".join(matchExpr.fields)
|
|
filter_fulltext = f"filter_fulltext('{fields}', '{matchExpr.matching_text}')"
|
|
if filter_cond:
|
|
filter_fulltext = f"({filter_cond}) AND {filter_fulltext}"
|
|
minimum_should_match = matchExpr.extra_options.get("minimum_should_match", 0.0)
|
|
if isinstance(minimum_should_match, float):
|
|
str_minimum_should_match = str(int(minimum_should_match * 100)) + "%"
|
|
matchExpr.extra_options["minimum_should_match"] = str_minimum_should_match
|
|
|
|
# Add rank_feature support
|
|
if rank_feature and "rank_features" not in matchExpr.extra_options:
|
|
# Convert rank_feature dict to Infinity's rank_features string format
|
|
# Format: "field^feature_name^weight,field^feature_name^weight"
|
|
rank_features_list = []
|
|
for feature_name, weight in rank_feature.items():
|
|
# Use TAG_FLD as the field containing rank features
|
|
rank_features_list.append(f"{TAG_FLD}^{feature_name}^{weight}")
|
|
if rank_features_list:
|
|
matchExpr.extra_options["rank_features"] = ",".join(rank_features_list)
|
|
|
|
for k, v in matchExpr.extra_options.items():
|
|
if not isinstance(v, str):
|
|
matchExpr.extra_options[k] = str(v)
|
|
self.logger.debug(f"INFINITY search MatchTextExpr: {json.dumps(matchExpr.__dict__)}")
|
|
elif isinstance(matchExpr, MatchDenseExpr):
|
|
if filter_fulltext and "filter" not in matchExpr.extra_options:
|
|
matchExpr.extra_options.update({"filter": filter_fulltext})
|
|
# dense_filter = _build_dense_filter(filter_cond, filter_fulltext, vector_similarity_weight)
|
|
# if dense_filter and "filter" not in matchExpr.extra_options:
|
|
# matchExpr.extra_options.update({"filter": dense_filter})
|
|
for k, v in matchExpr.extra_options.items():
|
|
if not isinstance(v, str):
|
|
matchExpr.extra_options[k] = str(v)
|
|
similarity = matchExpr.extra_options.get("similarity")
|
|
if similarity:
|
|
matchExpr.extra_options["threshold"] = similarity
|
|
del matchExpr.extra_options["similarity"]
|
|
self.logger.debug(f"INFINITY search MatchDenseExpr: {json.dumps(matchExpr.__dict__)}")
|
|
elif isinstance(matchExpr, FusionExpr):
|
|
if matchExpr.method == "weighted_sum":
|
|
# The default is "minmax" which gives a zero score for the last doc.
|
|
matchExpr.fusion_params["normalize"] = "atan"
|
|
self.logger.debug(f"INFINITY search FusionExpr: {json.dumps(matchExpr.__dict__)}")
|
|
|
|
order_by_expr_list = list()
|
|
if order_by.fields:
|
|
for order_field in order_by.fields:
|
|
if order_field[1] == 0:
|
|
order_by_expr_list.append((order_field[0], SortType.Asc))
|
|
else:
|
|
order_by_expr_list.append((order_field[0], SortType.Desc))
|
|
|
|
total_hits_count = 0
|
|
# Scatter search tables and gather the results
|
|
for indexName in index_names:
|
|
if indexName.startswith("ragflow_doc_meta_"):
|
|
table_names_to_search = [indexName]
|
|
else:
|
|
table_names_to_search = [f"{indexName}_{kb_id}" for kb_id in knowledgebase_ids]
|
|
for table_name in table_names_to_search:
|
|
try:
|
|
table_instance = db_instance.get_table(table_name)
|
|
except Exception:
|
|
continue
|
|
table_list.append(table_name)
|
|
builder = table_instance.output(output)
|
|
if len(match_expressions) > 0:
|
|
for matchExpr in match_expressions:
|
|
if isinstance(matchExpr, MatchTextExpr):
|
|
fields = ",".join(matchExpr.fields)
|
|
self.logger.info(f"INFINITY search match_text: {matchExpr.matching_text}")
|
|
builder = builder.match_text(
|
|
fields,
|
|
matchExpr.matching_text,
|
|
matchExpr.topn,
|
|
matchExpr.extra_options.copy(),
|
|
)
|
|
elif isinstance(matchExpr, MatchDenseExpr):
|
|
builder = builder.match_dense(
|
|
matchExpr.vector_column_name,
|
|
matchExpr.embedding_data,
|
|
matchExpr.embedding_data_type,
|
|
matchExpr.distance_type,
|
|
matchExpr.topn,
|
|
matchExpr.extra_options.copy(),
|
|
)
|
|
elif isinstance(matchExpr, FusionExpr):
|
|
builder = builder.fusion(matchExpr.method, matchExpr.topn, matchExpr.fusion_params)
|
|
else:
|
|
if filter_cond and len(filter_cond) > 0:
|
|
builder.filter(filter_cond)
|
|
if order_by.fields:
|
|
builder.sort(order_by_expr_list)
|
|
builder.offset(offset).limit(limit)
|
|
kb_res, extra_result = builder.option({"total_hits_count": True}).to_df()
|
|
if extra_result:
|
|
total_hits_count += int(extra_result["total_hits_count"])
|
|
self.logger.debug(f"INFINITY search table: {str(table_name)}, result: {str(kb_res)}")
|
|
df_list.append(kb_res)
|
|
res = self.concat_dataframes(df_list, output)
|
|
if match_expressions and score_column:
|
|
res["_score"] = res[score_column] + res[PAGERANK_FLD]
|
|
res = res.sort_values(by="_score", ascending=False).reset_index(drop=True)
|
|
res = res.head(limit)
|
|
self.logger.debug(f"INFINITY search final result: {str(res)}")
|
|
return res, total_hits_count
|
|
finally:
|
|
self.connPool.release_conn(inf_conn)
|
|
|
|
def get(self, chunk_id: str, index_name: str, knowledgebase_ids: list[str]) -> dict | None:
|
|
# Doc-meta tables are per-tenant, not per-kb: they have no `_kb_id`
|
|
# suffix. Match the special-casing used by index_exist/insert/delete
|
|
# in InfinityConnectionBase so callers can pass either a chunk
|
|
# index (``ragflow_<tenant>``) or a doc-meta index
|
|
# (``ragflow_doc_meta_<tenant>``) without us logging a bogus
|
|
# "blank knowledgebase_ids" warning.
|
|
is_meta_table = index_name.startswith("ragflow_doc_meta_")
|
|
|
|
# Validate the per-kb list BEFORE acquiring a connection — the
|
|
# blank-list case is a caller bug and shouldn't burn a connection
|
|
# from the pool. For meta tables the list is unused, so an empty
|
|
# list is fine.
|
|
if not is_meta_table:
|
|
if not knowledgebase_ids:
|
|
self.logger.warning("INFINITY get called with empty knowledgebase_ids for index %s", index_name)
|
|
return None
|
|
kb_table_names = [f"{index_name}_{kb_id}" for kb_id in knowledgebase_ids if kb_id]
|
|
if not kb_table_names:
|
|
self.logger.warning("INFINITY get has only blank knowledgebase_ids for index %s", index_name)
|
|
return None
|
|
|
|
inf_conn = self.connPool.get_conn()
|
|
try:
|
|
db_instance = inf_conn.get_database(self.dbName)
|
|
df_list = list()
|
|
assert isinstance(knowledgebase_ids, list)
|
|
table_list = list()
|
|
if is_meta_table:
|
|
table_names_to_search = [index_name]
|
|
else:
|
|
table_names_to_search = kb_table_names
|
|
for table_name in table_names_to_search:
|
|
table_list.append(table_name)
|
|
try:
|
|
table_instance = db_instance.get_table(table_name)
|
|
except Exception:
|
|
self.logger.warning(f"Table not found: {table_name}, this dataset isn't created in Infinity. Maybe it is created in other document engine.")
|
|
continue
|
|
kb_res, _ = table_instance.output(["*"]).filter(f"id = '{chunk_id}'").to_df()
|
|
self.logger.debug(f"INFINITY get table: {str(table_list)}, result: {str(kb_res)}")
|
|
df_list.append(kb_res)
|
|
finally:
|
|
self.connPool.release_conn(inf_conn)
|
|
res = self.concat_dataframes(df_list, ["id"])
|
|
fields = set(res.columns.tolist())
|
|
for field in [
|
|
"docnm_kwd",
|
|
"title_tks",
|
|
"title_sm_tks",
|
|
"important_kwd",
|
|
"important_tks",
|
|
"question_kwd",
|
|
"question_tks",
|
|
"content_with_weight",
|
|
"content_ltks",
|
|
"content_sm_ltks",
|
|
"authors_tks",
|
|
"authors_sm_tks",
|
|
]:
|
|
fields.add(field)
|
|
res_fields = self.get_fields(res, list(fields))
|
|
chunk = res_fields.get(chunk_id, None)
|
|
if chunk is not None:
|
|
chunk["id"] = chunk_id
|
|
return chunk
|
|
|
|
def ensure_columns(self, index_name: str, knowledgebase_id: str, column_defs: dict) -> None:
|
|
"""Make sure the per-KB chunk table carries the given columns.
|
|
|
|
Infinity's ``add_columns`` is idempotent for already-present columns,
|
|
so this is safe to call repeatedly. Used by callers that write new
|
|
marker rows referencing columns which were introduced after the
|
|
chunk-table schema was last updated (e.g. ``deleted_doc_id``, added
|
|
in #17685). For per-tenant doc-meta tables, pass ``knowledgebase_id``
|
|
as the empty string.
|
|
|
|
Logs and swallows any failure — this is a best-effort upgrade helper,
|
|
not a hard requirement of the calling write path. New tables are
|
|
created with the current ``conf/infinity_mapping.json`` schema so
|
|
callers will not need to invoke this in the steady state.
|
|
"""
|
|
if index_name.startswith("ragflow_doc_meta_"):
|
|
table_name = index_name
|
|
else:
|
|
table_name = f"{index_name}_{knowledgebase_id}" if knowledgebase_id else None
|
|
if not table_name:
|
|
return
|
|
inf_conn = self.connPool.get_conn()
|
|
try:
|
|
db_instance = inf_conn.get_database(self.dbName)
|
|
try:
|
|
table_instance = db_instance.get_table(table_name)
|
|
except InfinityException as e:
|
|
# src/common/status.cppm, kTableNotExist = 3022
|
|
if e.error_code != ErrorCode.TABLE_NOT_EXIST:
|
|
raise
|
|
# Table doesn't exist yet — the next insert() will create it
|
|
# with the current schema, so we have nothing to upgrade.
|
|
return
|
|
existing = {n for n, *_ in table_instance.show_columns().rows()}
|
|
missing = {c: d for c, d in column_defs.items() if c not in existing}
|
|
if not missing:
|
|
return
|
|
self.logger.info(
|
|
"INFINITY adding %d missing column(s) [%s] to %s",
|
|
len(missing),
|
|
", ".join(sorted(missing)),
|
|
table_name,
|
|
)
|
|
table_instance.add_columns(missing)
|
|
except Exception:
|
|
self.logger.exception(
|
|
"INFINITY failed to upgrade columns on %s; the next insert() may fail",
|
|
table_name,
|
|
)
|
|
finally:
|
|
self.connPool.release_conn(inf_conn)
|
|
|
|
def insert(self, documents: list[dict], index_name: str, knowledgebase_id: str = None, refresh: str | bool = "wait_for") -> list[str]:
|
|
"""
|
|
# Save input to file to test inserting from file in GO
|
|
import datetime
|
|
import os
|
|
debug_file = os.path.join("/var/infinity/tmp", f"insert_{index_name}_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S_%f')}.json")
|
|
with open(debug_file, 'w') as f:
|
|
json.dump({
|
|
"table_name": index_name,
|
|
"knowledgebase_id": knowledgebase_id,
|
|
"chunks": documents
|
|
}, f, indent=2)
|
|
self.logger.debug(f"Saved insert input to {debug_file}")
|
|
"""
|
|
|
|
inf_conn = self.connPool.get_conn()
|
|
try:
|
|
db_instance = inf_conn.get_database(self.dbName)
|
|
if index_name.startswith("ragflow_doc_meta_"):
|
|
table_name = index_name
|
|
else:
|
|
table_name = f"{index_name}_{knowledgebase_id}"
|
|
try:
|
|
table_instance = db_instance.get_table(table_name)
|
|
except InfinityException as e:
|
|
# src/common/status.cppm, kTableNotExist = 3022
|
|
if e.error_code != ErrorCode.TABLE_NOT_EXIST:
|
|
raise
|
|
vector_size = 0
|
|
patt = re.compile(r"q_(?P<vector_size>\d+)_vec")
|
|
for k in documents[0].keys():
|
|
m = patt.match(k)
|
|
if m:
|
|
vector_size = int(m.group("vector_size"))
|
|
break
|
|
if vector_size == 0:
|
|
raise ValueError("Cannot infer vector size from documents")
|
|
|
|
# Determine parser_id from document structure
|
|
# Table parser documents have 'chunk_data' field
|
|
parser_id = None
|
|
if "chunk_data" in documents[0] and isinstance(documents[0].get("chunk_data"), dict):
|
|
from common.constants import ParserType
|
|
|
|
parser_id = ParserType.TABLE.value
|
|
self.logger.debug("Detected TABLE parser from document structure")
|
|
|
|
# Fallback: Create table with base schema (shouldn't normally happen as init_kb() creates it)
|
|
self.logger.debug(f"Fallback: Creating table {table_name} with base schema, parser_id: {parser_id}")
|
|
self.create_idx(index_name, knowledgebase_id, vector_size, parser_id)
|
|
table_instance = db_instance.get_table(table_name)
|
|
|
|
# embedding fields can't have a default value....
|
|
embedding_clmns = []
|
|
clmns = table_instance.show_columns().rows()
|
|
for n, ty, _, _ in clmns:
|
|
r = re.search(r"Embedding\([a-z]+,([0-9]+)\)", ty)
|
|
if not r:
|
|
continue
|
|
embedding_clmns.append((n, int(r.group(1))))
|
|
|
|
docs = copy.deepcopy(documents)
|
|
for d in docs:
|
|
assert "_id" not in d
|
|
assert "id" in d
|
|
for k, v in list(d.items()):
|
|
if k == "docnm_kwd":
|
|
d["docnm"] = v
|
|
elif k == "title_kwd":
|
|
if not d.get("docnm_kwd"):
|
|
d["docnm"] = self.list2str(v)
|
|
elif k == "title_sm_tks":
|
|
if not d.get("docnm_kwd"):
|
|
d["docnm"] = self.list2str(v)
|
|
elif k == "important_kwd":
|
|
if isinstance(v, list):
|
|
empty_count = sum(1 for kw in v if kw == "")
|
|
tokens = [kw for kw in v if kw != ""]
|
|
d["important_keywords"] = self.list2str(tokens, ",")
|
|
d["important_kwd_empty_count"] = empty_count
|
|
else:
|
|
d["important_keywords"] = self.list2str(v, ",")
|
|
elif k == "important_tks":
|
|
if not d.get("important_kwd"):
|
|
d["important_keywords"] = v
|
|
elif k == "content_with_weight":
|
|
d["content"] = v
|
|
elif k == "content_ltks":
|
|
if not d.get("content_with_weight"):
|
|
d["content"] = v
|
|
elif k == "content_sm_ltks":
|
|
if not d.get("content_with_weight"):
|
|
d["content"] = v
|
|
elif k == "authors_tks":
|
|
d["authors"] = v
|
|
elif k == "authors_sm_tks":
|
|
if not d.get("authors_tks"):
|
|
d["authors"] = v
|
|
elif k == "question_kwd":
|
|
d["questions"] = self.list2str(v, "\n")
|
|
elif k == "question_tks":
|
|
if not d.get("question_kwd"):
|
|
d["questions"] = self.list2str(v)
|
|
elif k in _JSON_LIST_FIELDS:
|
|
d[k] = json.dumps(list(v) if isinstance(v, (list, tuple, set)) else [], ensure_ascii=False)
|
|
elif self.field_keyword(k):
|
|
if isinstance(v, list):
|
|
d[k] = "###".join(v)
|
|
else:
|
|
d[k] = v
|
|
elif re.search(r"_feas$", k):
|
|
d[k] = json.dumps(v)
|
|
elif k == "chunk_data":
|
|
# Convert data dict to JSON string for storage
|
|
if isinstance(v, dict):
|
|
d[k] = json.dumps(v)
|
|
else:
|
|
d[k] = v
|
|
elif k == "extra":
|
|
# RAPTOR writes {"raptor_method": ...} as a dict; Infinity's
|
|
# `extra` column is varchar so we serialize on the write path.
|
|
# The read path (raptor_utils._as_extra_dict) already accepts
|
|
# both dict and JSON-string. Other backends (OceanBase JSON
|
|
# column, ES/OpenSearch) keep dict shape — this is Infinity-only.
|
|
if isinstance(v, dict):
|
|
d[k] = json.dumps(v)
|
|
else:
|
|
d[k] = v if v else ""
|
|
elif k == "kb_id":
|
|
if isinstance(d[k], list):
|
|
d[k] = d[k][0] # since d[k] is a list, but we need a str
|
|
elif k == "position_int":
|
|
assert isinstance(v, list)
|
|
arr = [num for row in v for num in row]
|
|
d[k] = "_".join(f"{num:08x}" for num in arr)
|
|
elif k in ["page_num_int", "top_int"]:
|
|
assert isinstance(v, list)
|
|
d[k] = "_".join(f"{num:08x}" for num in v)
|
|
elif k == "meta_fields":
|
|
if isinstance(v, dict):
|
|
d[k] = json.dumps(v, ensure_ascii=False)
|
|
else:
|
|
d[k] = v if v else "{}"
|
|
else:
|
|
d[k] = v
|
|
# Infinity thrift client does not accept None values.
|
|
for k in list(d.keys()):
|
|
if d[k] is None:
|
|
del d[k]
|
|
for k in [
|
|
"docnm_kwd",
|
|
"title_tks",
|
|
"title_sm_tks",
|
|
"important_kwd",
|
|
"important_tks",
|
|
"content_with_weight",
|
|
"content_ltks",
|
|
"content_sm_ltks",
|
|
"authors_tks",
|
|
"authors_sm_tks",
|
|
"question_kwd",
|
|
"question_tks",
|
|
]:
|
|
if k in d:
|
|
del d[k]
|
|
|
|
for n, vs in embedding_clmns:
|
|
if n in d:
|
|
continue
|
|
d[n] = [0] * vs
|
|
ids = ["'{}'".format(d["id"]) for d in docs]
|
|
str_ids = ", ".join(ids)
|
|
str_filter = f"id IN ({str_ids})"
|
|
table_instance.delete(str_filter)
|
|
# for doc in documents:
|
|
# logger.info(f"insert position_int: {doc['position_int']}")
|
|
# logger.info(f"InfinityConnection.insert {json.dumps(documents)}")
|
|
table_instance.insert(docs)
|
|
finally:
|
|
self.connPool.release_conn(inf_conn)
|
|
self.logger.debug(f"INFINITY inserted into {table_name} {str_ids}.")
|
|
return []
|
|
|
|
def update(self, condition: dict, new_value: dict, index_name: str, knowledgebase_id: str) -> bool:
|
|
# if 'position_int' in newValue:
|
|
# logger.info(f"update position_int: {newValue['position_int']}")
|
|
inf_conn = self.connPool.get_conn()
|
|
try:
|
|
db_instance = inf_conn.get_database(self.dbName)
|
|
if index_name.startswith("ragflow_doc_meta_"):
|
|
table_name = index_name
|
|
else:
|
|
table_name = f"{index_name}_{knowledgebase_id}"
|
|
try:
|
|
table_instance = db_instance.get_table(table_name)
|
|
except InfinityException as e:
|
|
# src/common/status.cppm, kTableNotExist = 3022
|
|
if e.error_code == ErrorCode.TABLE_NOT_EXIST:
|
|
self.logger.warning(f"Table {table_name} does not exist, skipping update.")
|
|
return False
|
|
raise
|
|
# if "exists" in condition:
|
|
# del condition["exists"]
|
|
|
|
clmns = {}
|
|
if table_instance:
|
|
for n, ty, de, _ in table_instance.show_columns().rows():
|
|
clmns[n] = (ty, de)
|
|
filter = self.equivalent_condition_to_str(condition, table_instance)
|
|
removeValue = {}
|
|
for k, v in list(new_value.items()):
|
|
if k == "docnm_kwd":
|
|
new_value["docnm"] = self.list2str(v)
|
|
elif k == "title_kwd":
|
|
if not new_value.get("docnm_kwd"):
|
|
new_value["docnm"] = self.list2str(v)
|
|
elif k == "title_sm_tks":
|
|
if not new_value.get("docnm_kwd"):
|
|
new_value["docnm"] = v
|
|
elif k == "important_kwd":
|
|
if isinstance(v, list):
|
|
empty_count = sum(1 for kw in v if kw == "")
|
|
tokens = [kw for kw in v if kw != ""]
|
|
new_value["important_keywords"] = self.list2str(tokens, ",")
|
|
new_value["important_kwd_empty_count"] = empty_count
|
|
else:
|
|
new_value["important_keywords"] = self.list2str(v, ",")
|
|
elif k == "important_tks":
|
|
if not new_value.get("important_kwd"):
|
|
new_value["important_keywords"] = v
|
|
elif k == "content_with_weight":
|
|
new_value["content"] = v
|
|
elif k == "content_ltks":
|
|
if not new_value.get("content_with_weight"):
|
|
new_value["content"] = v
|
|
elif k == "content_sm_ltks":
|
|
if not new_value.get("content_with_weight"):
|
|
new_value["content"] = v
|
|
elif k == "authors_tks":
|
|
new_value["authors"] = v
|
|
elif k == "authors_sm_tks":
|
|
if not new_value.get("authors_tks"):
|
|
new_value["authors"] = v
|
|
elif k == "question_kwd":
|
|
new_value["questions"] = "\n".join(v)
|
|
elif k == "question_tks":
|
|
if not new_value.get("question_kwd"):
|
|
new_value["questions"] = self.list2str(v)
|
|
elif k in _JSON_LIST_FIELDS:
|
|
new_value[k] = json.dumps(list(v) if isinstance(v, (list, tuple, set)) else [], ensure_ascii=False)
|
|
elif self.field_keyword(k):
|
|
if isinstance(v, list):
|
|
new_value[k] = "###".join(v)
|
|
else:
|
|
new_value[k] = v
|
|
elif re.search(r"_feas$", k):
|
|
new_value[k] = json.dumps(v)
|
|
elif k == "kb_id":
|
|
if isinstance(new_value[k], list):
|
|
new_value[k] = new_value[k][0] # since d[k] is a list, but we need a str
|
|
elif k == "position_int":
|
|
assert isinstance(v, list)
|
|
arr = [num for row in v for num in row]
|
|
new_value[k] = "_".join(f"{num:08x}" for num in arr)
|
|
elif k in ["page_num_int", "top_int"]:
|
|
assert isinstance(v, list)
|
|
new_value[k] = "_".join(f"{num:08x}" for num in v)
|
|
elif k == "remove":
|
|
if isinstance(v, str):
|
|
assert v in clmns, f"'{v}' should be in '{clmns}'."
|
|
ty, de = clmns[v]
|
|
if ty.lower().find("cha"):
|
|
if not de:
|
|
de = ""
|
|
new_value[v] = de
|
|
else:
|
|
for kk, vv in v.items():
|
|
removeValue[kk] = vv
|
|
del new_value[k]
|
|
else:
|
|
new_value[k] = v
|
|
for k in [
|
|
"docnm_kwd",
|
|
"title_tks",
|
|
"title_sm_tks",
|
|
"important_kwd",
|
|
"important_tks",
|
|
"content_with_weight",
|
|
"content_ltks",
|
|
"content_sm_ltks",
|
|
"authors_tks",
|
|
"authors_sm_tks",
|
|
"question_kwd",
|
|
"question_tks",
|
|
]:
|
|
if k in new_value:
|
|
del new_value[k]
|
|
|
|
# The Infinity Python client inspects value[0] for list values
|
|
# while building an update expression. An empty list therefore
|
|
# raises IndexError before the request reaches Infinity. Keep
|
|
# JSON and keyword-list columns clearable, but do not send empty
|
|
# values for other columns (for example, an empty vector returned
|
|
# by a partial row read).
|
|
for k, v in list(new_value.items()):
|
|
if not isinstance(v, list) or v:
|
|
continue
|
|
if k in _JSON_LIST_FIELDS:
|
|
new_value[k] = json.dumps([], ensure_ascii=False)
|
|
elif self.field_keyword(k):
|
|
new_value[k] = ""
|
|
else:
|
|
del new_value[k]
|
|
|
|
remove_opt = {} # "[k,new_value]": [id_to_update, ...]
|
|
if removeValue:
|
|
col_to_remove = list(removeValue.keys())
|
|
row_to_opt = table_instance.output(col_to_remove + ["id"]).filter(filter).to_df()
|
|
self.logger.debug(f"INFINITY search table {str(table_name)}, filter {filter}, result: {str(row_to_opt[0])}")
|
|
row_to_opt = self.get_fields(row_to_opt, col_to_remove)
|
|
for id, old_v in row_to_opt.items():
|
|
for k, remove_v in removeValue.items():
|
|
if remove_v in old_v[k]:
|
|
new_v = old_v[k].copy()
|
|
new_v.remove(remove_v)
|
|
kv_key = json.dumps([k, new_v])
|
|
if kv_key not in remove_opt:
|
|
remove_opt[kv_key] = [id]
|
|
else:
|
|
remove_opt[kv_key].append(id)
|
|
|
|
self.logger.debug(f"INFINITY update table {table_name}, filter {filter}, newValue {new_value}.")
|
|
for update_kv, ids in remove_opt.items():
|
|
k, v = json.loads(update_kv)
|
|
table_instance.update(filter + " AND id in ({0})".format(",".join([f"'{id}'" for id in ids])), {k: "###".join(v)})
|
|
|
|
table_instance.update(filter, new_value)
|
|
finally:
|
|
self.connPool.release_conn(inf_conn)
|
|
return True
|
|
|
|
def adjust_chunk_pagerank_fea(
|
|
self,
|
|
chunk_id: str,
|
|
index_name: str,
|
|
knowledgebase_id: str,
|
|
delta: int,
|
|
min_weight: int,
|
|
max_weight: int,
|
|
row_id: int | None = None,
|
|
max_retries: int = 2,
|
|
) -> bool:
|
|
"""Adjust pagerank_fea on one chunk row in Infinity.
|
|
|
|
Uses row_id for a targeted update when available. If the row_id is
|
|
stale (concurrent update changed it), re-reads the current row_id and
|
|
retries up to *max_retries* times.
|
|
"""
|
|
table_name = f"{index_name}_{knowledgebase_id}"
|
|
for attempt in range(max_retries + 1):
|
|
inf_conn = self.connPool.get_conn()
|
|
try:
|
|
db_instance = inf_conn.get_database(self.dbName)
|
|
table_instance = db_instance.get_table(table_name)
|
|
|
|
if row_id is None:
|
|
df, _ = table_instance.output([PAGERANK_FLD, "row_id()"]).filter(f"id = '{chunk_id}'").to_df()
|
|
if df.empty:
|
|
self.logger.warning(
|
|
"adjust_chunk_pagerank_fea: chunk %s not found in %s",
|
|
chunk_id,
|
|
table_name,
|
|
)
|
|
return False
|
|
current_weight = int(float(df[PAGERANK_FLD].iloc[0] or 0))
|
|
row_id = int(df["row_id"].iloc[0])
|
|
else:
|
|
df, _ = table_instance.output([PAGERANK_FLD]).filter(f"id = '{chunk_id}'").to_df()
|
|
if df.empty:
|
|
return False
|
|
current_weight = int(float(df[PAGERANK_FLD].iloc[0] or 0))
|
|
|
|
new_weight = max(min_weight, min(max_weight, current_weight + delta))
|
|
|
|
table_instance.update(
|
|
f"_row_id = {row_id}",
|
|
{PAGERANK_FLD: new_weight},
|
|
)
|
|
self.logger.info(
|
|
"adjust_chunk_pagerank_fea(chunk=%s, table=%s): %s -> %s via row_id=%s",
|
|
chunk_id,
|
|
table_name,
|
|
current_weight,
|
|
new_weight,
|
|
row_id,
|
|
)
|
|
return True
|
|
|
|
except InfinityException as e:
|
|
if attempt < max_retries:
|
|
self.logger.warning(
|
|
"adjust_chunk_pagerank_fea stale row_id=%s for chunk %s (attempt %s/%s): %s",
|
|
row_id,
|
|
chunk_id,
|
|
attempt + 1,
|
|
max_retries,
|
|
e,
|
|
)
|
|
row_id = None
|
|
continue
|
|
self.logger.error(
|
|
"adjust_chunk_pagerank_fea failed for chunk %s after %s attempts: %s",
|
|
chunk_id,
|
|
max_retries + 1,
|
|
e,
|
|
)
|
|
return False
|
|
except Exception as e:
|
|
self.logger.error(
|
|
"adjust_chunk_pagerank_fea error for chunk %s: %s",
|
|
chunk_id,
|
|
e,
|
|
)
|
|
return False
|
|
finally:
|
|
self.connPool.release_conn(inf_conn)
|
|
return False
|
|
|
|
"""
|
|
Helper functions for search result
|
|
"""
|
|
|
|
def get_fields(self, res: tuple[pd.DataFrame, int] | pd.DataFrame, fields: list[str]) -> dict[str, dict]:
|
|
if isinstance(res, tuple):
|
|
res = res[0]
|
|
if not fields:
|
|
return {}
|
|
fields_all = fields.copy()
|
|
fields_all.append("id")
|
|
fields_all = set(fields_all)
|
|
if "docnm" in res.columns:
|
|
for field in ["docnm_kwd", "title_tks", "title_sm_tks"]:
|
|
if field in fields_all:
|
|
res[field] = res["docnm"]
|
|
if "important_keywords" in res.columns:
|
|
if "important_kwd" in fields_all:
|
|
if "important_kwd_empty_count" in res.columns:
|
|
base = res["important_keywords"].apply(lambda raw: raw.split(",") if raw else [])
|
|
counts = res["important_kwd_empty_count"].fillna(0).astype(int)
|
|
res["important_kwd"] = [tokens + [""] * empty_count for tokens, empty_count in zip(base.tolist(), counts.tolist())]
|
|
else:
|
|
res["important_kwd"] = res["important_keywords"].apply(lambda v: v.split(",") if v else [])
|
|
if "important_tks" in fields_all:
|
|
res["important_tks"] = res["important_keywords"]
|
|
if "questions" in res.columns:
|
|
if "question_kwd" in fields_all:
|
|
res["question_kwd"] = res["questions"].apply(lambda v: v.splitlines())
|
|
if "question_tks" in fields_all:
|
|
res["question_tks"] = res["questions"]
|
|
if "content" in res.columns:
|
|
for field in ["content_with_weight", "content_ltks", "content_sm_ltks"]:
|
|
if field in fields_all:
|
|
res[field] = res["content"]
|
|
if "authors" in res.columns:
|
|
for field in ["authors_tks", "authors_sm_tks"]:
|
|
if field in fields_all:
|
|
res[field] = res["authors"]
|
|
|
|
column_map = {col.lower(): col for col in res.columns}
|
|
# row_id() is returned by infinity as "row_id", add mapping for lookup
|
|
if "row_id()" in fields_all and "row_id" in column_map:
|
|
column_map["row_id()"] = column_map["row_id"]
|
|
matched_columns = {column_map[col.lower()]: col for col in fields_all if col.lower() in column_map}
|
|
none_columns = [col for col in fields_all if col.lower() not in column_map]
|
|
|
|
res2 = res[matched_columns.keys()]
|
|
res2 = res2.rename(columns=matched_columns)
|
|
res2.drop_duplicates(subset=["id"], inplace=True)
|
|
|
|
for column in list(res2.columns):
|
|
k = column.lower()
|
|
if k in _JSON_LIST_FIELDS:
|
|
|
|
def parse_json_list(value):
|
|
if isinstance(value, list):
|
|
return value
|
|
if not value:
|
|
return []
|
|
try:
|
|
parsed = json.loads(value)
|
|
return parsed if isinstance(parsed, list) else [parsed]
|
|
except (TypeError, json.JSONDecodeError):
|
|
# Read rows written by the previous varchar encoding.
|
|
return [item for item in str(value).split("###") if item]
|
|
|
|
res2[column] = res2[column].apply(parse_json_list)
|
|
elif self.field_keyword(k):
|
|
res2[column] = res2[column].apply(lambda v: [kwd for kwd in v.split("###") if kwd])
|
|
elif re.search(r"_feas$", k):
|
|
res2[column] = res2[column].apply(lambda v: json.loads(v) if v else {})
|
|
elif k == "chunk_data":
|
|
# Parse JSON data back to dict for table parser fields
|
|
res2[column] = res2[column].apply(lambda v: json.loads(v) if v and isinstance(v, str) else v)
|
|
elif k == "position_int":
|
|
|
|
def to_position_int(v):
|
|
if v:
|
|
arr = [int(hex_val, 16) for hex_val in v.split("_")]
|
|
v = [arr[i : i + 5] for i in range(0, len(arr), 5)]
|
|
else:
|
|
v = []
|
|
return v
|
|
|
|
res2[column] = res2[column].apply(to_position_int)
|
|
elif k in ["page_num_int", "top_int"]:
|
|
res2[column] = res2[column].apply(lambda v: [int(hex_val, 16) for hex_val in v.split("_")] if v else [])
|
|
else:
|
|
pass
|
|
for column in ["docnm", "important_keywords", "questions", "content", "authors"]:
|
|
if column in res2:
|
|
del res2[column]
|
|
for column in none_columns:
|
|
res2[column] = None
|
|
|
|
return res2.set_index("id").to_dict(orient="index")
|