mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-06-08 08:07:21 +08:00
663fc1d42c
### What problem does this PR solve? Fixes #14570. On OpenSearch backends (`DOC_ENGINE=opensearch`) every document-metadata write failed with `'OSConnection' object has no attribute 'create_doc_meta_idx'`, so both `PATCH /api/v1/datasets/{ds}/documents/{doc}` with `meta_fields` and `POST /api/v1/datasets/{ds}/metadata/update` were unusable while every other document operation (retrieval, parsing, name update, chunk management) worked correctly on the same OpenSearch cluster. The bug runs deeper than the missing method name in the error message suggests. `DocMetadataService` also reached into `settings.docStoreConn.es.*` directly for the index refresh, the scripted partial update, and the count call, which means that even after adding `create_doc_meta_idx` to `OSConnection` the very next call in the same metadata flow would still raise `AttributeError` because `OSConnection` exposes `self.os` rather than `self.es`. Fixing only the reported symptom would have moved the failure one line down without restoring the feature. This PR adds a uniform document-metadata dispatch surface to both connection classes so they present the same abstract API, and routes the service layer through that surface via `getattr` guards instead of poking at backend-specific attributes. The four new methods on `OSConnection` and `ESConnectionBase` are `create_doc_meta_idx`, `refresh_idx`, `count_idx`, and `replace_meta_fields`. `OSConnection.create_doc_meta_idx` reuses the existing `conf/doc_meta_es_mapping.json` schema in the OpenSearch `body=` form because OpenSearch and Elasticsearch share the same index-creation payload, and `replace_meta_fields` emits a full scripted assignment (`ctx._source.meta_fields = params.meta_fields`) on both backends so removed keys actually disappear instead of being preserved by deep-merge semantics. The `getattr`-guarded dispatch in `DocMetadataService` keeps the existing fall-through paths intact for Infinity and OceanBase, which continue to rely on their search-based count fallback and on the delete-then-insert metadata replacement they used before, so this change is strictly additive for those two backends. Verification: `pytest test/unit_test/rag/utils/test_opensearch_doc_meta.py` runs 16 new unit tests that pass locally and pin the `OSConnection` dispatch surface, the `create_doc_meta_idx` short-circuit when the index already exists, the mapping-file payload routing, the `IndicesClient.create` failure path, the `refresh_idx` and `count_idx` success and error sentinels, and the full-assignment script emitted by `replace_meta_fields`. The test module stubs `common.settings` and `rag.nlp` at import time so the suite runs without the heavy backend SDKs that the rest of the repository pulls in transitively. ### Type of change - [x] Bug Fix (non-breaking change which fixes an issue) --------- Co-authored-by: tmimmanuel <tmimmanuel@users.noreply.github.com>
765 lines
31 KiB
Python
765 lines
31 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 logging
|
|
import re
|
|
import json
|
|
import time
|
|
import os
|
|
|
|
import copy
|
|
from opensearchpy import OpenSearch, NotFoundError
|
|
from opensearchpy import UpdateByQuery, Q, Search, Index
|
|
from opensearchpy import ConnectionTimeout
|
|
from common.decorator import singleton
|
|
from common.file_utils import get_project_base_directory
|
|
from common.doc_store.doc_store_base import DocStoreConnection, MatchExpr, OrderByExpr, MatchTextExpr, MatchDenseExpr, \
|
|
FusionExpr
|
|
from rag.nlp import is_english, rag_tokenizer
|
|
from common.constants import PAGERANK_FLD, TAG_FLD
|
|
from common import settings
|
|
|
|
ATTEMPT_TIME = 2
|
|
|
|
_PAGERANK_FEA_ADJUST_SCRIPT = """
|
|
double cur = 0.0;
|
|
if (ctx._source.containsKey(params.pf)) {
|
|
Object v = ctx._source[params.pf];
|
|
if (v != null) {
|
|
if (v instanceof Number) {
|
|
cur = ((Number)v).doubleValue();
|
|
} else {
|
|
try { cur = Double.parseDouble(v.toString()); } catch (Exception e) { cur = 0.0; }
|
|
}
|
|
}
|
|
}
|
|
double nw = cur + params.delta;
|
|
if (nw < params.min_w) { nw = params.min_w; }
|
|
if (nw > params.max_w) { nw = params.max_w; }
|
|
if (nw <= 0.0) {
|
|
if (ctx._source.containsKey(params.pf)) {
|
|
ctx._source.remove(params.pf);
|
|
}
|
|
} else {
|
|
ctx._source[params.pf] = nw;
|
|
}
|
|
"""
|
|
|
|
logger = logging.getLogger('ragflow.opensearch_conn')
|
|
|
|
|
|
@singleton
|
|
class OSConnection(DocStoreConnection):
|
|
def __init__(self):
|
|
self.info = {}
|
|
logger.info(f"Use OpenSearch {settings.OS['hosts']} as the doc engine.")
|
|
for _ in range(ATTEMPT_TIME):
|
|
try:
|
|
self.os = OpenSearch(
|
|
settings.OS["hosts"].split(","),
|
|
http_auth=(settings.OS["username"], settings.OS[
|
|
"password"]) if "username" in settings.OS and "password" in settings.OS else None,
|
|
verify_certs=False,
|
|
timeout=600
|
|
)
|
|
if self.os:
|
|
self.info = self.os.info()
|
|
break
|
|
except Exception as e:
|
|
logger.warning(f"{str(e)}. Waiting OpenSearch {settings.OS['hosts']} to be healthy.")
|
|
time.sleep(5)
|
|
if not self.os.ping():
|
|
msg = f"OpenSearch {settings.OS['hosts']} is unhealthy in 120s."
|
|
logger.error(msg)
|
|
raise Exception(msg)
|
|
v = self.info.get("version", {"number": "2.18.0"})
|
|
v = v["number"].split(".")[0]
|
|
if int(v) < 2:
|
|
msg = f"OpenSearch version must be greater than or equal to 2, current version: {v}"
|
|
logger.error(msg)
|
|
raise Exception(msg)
|
|
fp_mapping = os.path.join(get_project_base_directory(), "conf", "os_mapping.json")
|
|
if not os.path.exists(fp_mapping):
|
|
msg = f"OpenSearch mapping file not found at {fp_mapping}"
|
|
logger.error(msg)
|
|
raise Exception(msg)
|
|
with open(fp_mapping, "r") as f:
|
|
self.mapping = json.load(f)
|
|
logger.info(f"OpenSearch {settings.OS['hosts']} is healthy.")
|
|
|
|
"""
|
|
Database operations
|
|
"""
|
|
|
|
def db_type(self) -> str:
|
|
return "opensearch"
|
|
|
|
def health(self) -> dict:
|
|
health_dict = dict(self.os.cluster.health())
|
|
health_dict["type"] = "opensearch"
|
|
return health_dict
|
|
|
|
"""
|
|
Table operations
|
|
"""
|
|
|
|
def create_idx(self, indexName: str, knowledgebaseId: str, vectorSize: int, parser_id: str = None):
|
|
if self.index_exist(indexName, knowledgebaseId):
|
|
return True
|
|
try:
|
|
from opensearchpy.client import IndicesClient
|
|
return IndicesClient(self.os).create(index=indexName,
|
|
body=self.mapping)
|
|
except Exception:
|
|
logger.exception("OSConnection.createIndex error %s" % (indexName))
|
|
|
|
def create_doc_meta_idx(self, index_name: str):
|
|
"""
|
|
Create a per-tenant document metadata index on OpenSearch.
|
|
|
|
Mirrors ESConnectionBase.create_doc_meta_idx so that the
|
|
DocMetadataService dispatches uniformly across ES and OS backends.
|
|
Index name pattern: ragflow_doc_meta_{tenant_id}
|
|
"""
|
|
if self.index_exist(index_name, ""):
|
|
return True
|
|
try:
|
|
fp_mapping = os.path.join(get_project_base_directory(), "conf", "doc_meta_es_mapping.json")
|
|
if not os.path.exists(fp_mapping):
|
|
logger.error(f"Document metadata mapping file not found at {fp_mapping}")
|
|
return False
|
|
|
|
with open(fp_mapping, "r") as f:
|
|
doc_meta_mapping = json.load(f)
|
|
|
|
from opensearchpy.client import IndicesClient
|
|
body = {
|
|
"settings": doc_meta_mapping["settings"],
|
|
"mappings": doc_meta_mapping["mappings"],
|
|
}
|
|
return IndicesClient(self.os).create(index=index_name, body=body)
|
|
except Exception as e:
|
|
logger.exception(f"OSConnection.create_doc_meta_idx error creating {index_name}: {e}")
|
|
return False
|
|
|
|
def refresh_idx(self, index_name: str) -> bool:
|
|
"""
|
|
Refresh an index so that recently inserted documents become searchable.
|
|
|
|
DocMetadataService used to call ``settings.docStoreConn.es.indices.refresh``
|
|
directly, which raised AttributeError on the OpenSearch backend because
|
|
OSConnection exposes ``self.os`` rather than ``self.es``. This wrapper
|
|
gives both backends a uniform abstract entry point.
|
|
"""
|
|
try:
|
|
self.os.indices.refresh(index=index_name)
|
|
return True
|
|
except NotFoundError:
|
|
return False
|
|
except Exception as e:
|
|
logger.warning(f"OSConnection.refresh_idx({index_name}) failed: {e}")
|
|
return False
|
|
|
|
def count_idx(self, index_name: str) -> int:
|
|
"""
|
|
Return the document count for an index, or -1 if the call fails.
|
|
|
|
Used by DocMetadataService._drop_empty_metadata_table to decide whether
|
|
a per-tenant metadata index is empty without paying a full search.
|
|
"""
|
|
try:
|
|
response = self.os.count(index=index_name)
|
|
return int(response.get("count", 0))
|
|
except NotFoundError:
|
|
return 0
|
|
except Exception as e:
|
|
logger.warning(f"OSConnection.count_idx({index_name}) failed: {e}")
|
|
return -1
|
|
|
|
def replace_meta_fields(self, index_name: str, doc_id: str, meta_fields: dict) -> bool:
|
|
"""
|
|
Replace the ``meta_fields`` object on a single document.
|
|
|
|
ES.update with a ``doc`` body deep-merges object fields, which retains
|
|
old keys that should be removed. The fix in ESConnection is a script
|
|
that fully assigns the new meta_fields. We provide the same primitive
|
|
on OpenSearch so the service layer never reaches into ``self.es`` or
|
|
``self.os`` directly.
|
|
"""
|
|
body = {
|
|
"script": {
|
|
"source": "ctx._source.meta_fields = params.meta_fields",
|
|
"params": {"meta_fields": meta_fields},
|
|
}
|
|
}
|
|
for _ in range(ATTEMPT_TIME):
|
|
try:
|
|
self.os.update(index=index_name, id=doc_id, body=body, refresh=True)
|
|
return True
|
|
except NotFoundError:
|
|
return False
|
|
except Exception as e:
|
|
logger.warning(f"OSConnection.replace_meta_fields({index_name}, {doc_id}) failed: {e}")
|
|
if re.search(r"(timeout|connection)", str(e).lower()):
|
|
time.sleep(1)
|
|
continue
|
|
return False
|
|
return False
|
|
|
|
def delete_idx(self, indexName: str, knowledgebaseId: str):
|
|
if len(knowledgebaseId) > 0:
|
|
# The index need to be alive after any kb deletion since all kb under this tenant are in one index.
|
|
return
|
|
try:
|
|
self.os.indices.delete(index=indexName, allow_no_indices=True)
|
|
except NotFoundError:
|
|
pass
|
|
except Exception:
|
|
logger.exception("OSConnection.deleteIdx error %s" % (indexName))
|
|
|
|
def index_exist(self, indexName: str, knowledgebaseId: str = None) -> bool:
|
|
s = Index(indexName, self.os)
|
|
for i in range(ATTEMPT_TIME):
|
|
try:
|
|
return s.exists()
|
|
except Exception as e:
|
|
logger.exception("OSConnection.indexExist got exception")
|
|
if str(e).find("Timeout") > 0 or str(e).find("Conflict") > 0:
|
|
continue
|
|
break
|
|
return False
|
|
|
|
"""
|
|
CRUD operations
|
|
"""
|
|
|
|
def search(
|
|
self, selectFields: list[str],
|
|
highlightFields: list[str],
|
|
condition: dict,
|
|
matchExprs: list[MatchExpr],
|
|
orderBy: OrderByExpr,
|
|
offset: int,
|
|
limit: int,
|
|
indexNames: str | list[str],
|
|
knowledgebaseIds: list[str],
|
|
aggFields: list[str] = [],
|
|
rank_feature: dict | None = None
|
|
):
|
|
"""
|
|
Refers to https://github.com/opensearch-project/opensearch-py/blob/main/guides/dsl.md
|
|
"""
|
|
use_knn = False
|
|
if isinstance(indexNames, str):
|
|
indexNames = indexNames.split(",")
|
|
assert isinstance(indexNames, list) and len(indexNames) > 0
|
|
assert "_id" not in condition
|
|
|
|
bqry = Q("bool", must=[])
|
|
condition["kb_id"] = knowledgebaseIds
|
|
for k, v in condition.items():
|
|
if k == "available_int":
|
|
if v == 0:
|
|
bqry.filter.append(Q("range", available_int={"lt": 1}))
|
|
else:
|
|
bqry.filter.append(
|
|
Q("bool", must_not=Q("range", available_int={"lt": 1})))
|
|
continue
|
|
if not v:
|
|
continue
|
|
if isinstance(v, list):
|
|
bqry.filter.append(Q("terms", **{k: v}))
|
|
elif isinstance(v, str) or isinstance(v, int):
|
|
bqry.filter.append(Q("term", **{k: v}))
|
|
else:
|
|
raise Exception(
|
|
f"Condition `{str(k)}={str(v)}` value type is {str(type(v))}, expected to be int, str or list.")
|
|
|
|
s = Search()
|
|
vector_similarity_weight = 0.5
|
|
for m in matchExprs:
|
|
if isinstance(m, FusionExpr) and m.method == "weighted_sum" and "weights" in m.fusion_params:
|
|
assert len(matchExprs) == 3 and isinstance(matchExprs[0], MatchTextExpr) and isinstance(matchExprs[1],
|
|
MatchDenseExpr) and isinstance(
|
|
matchExprs[2], FusionExpr)
|
|
weights = m.fusion_params["weights"]
|
|
vector_similarity_weight = float(weights.split(",")[1])
|
|
knn_query = {}
|
|
for m in matchExprs:
|
|
if isinstance(m, MatchTextExpr):
|
|
minimum_should_match = m.extra_options.get("minimum_should_match", 0.0)
|
|
if isinstance(minimum_should_match, float):
|
|
minimum_should_match = str(int(minimum_should_match * 100)) + "%"
|
|
bqry.must.append(Q("query_string", fields=m.fields,
|
|
type="best_fields", query=m.matching_text,
|
|
minimum_should_match=minimum_should_match,
|
|
boost=1))
|
|
bqry.boost = 1.0 - vector_similarity_weight
|
|
|
|
# Elasticsearch has the encapsulation of KNN_search in python sdk
|
|
# while the Python SDK for OpenSearch does not provide encapsulation for KNN_search,
|
|
# the following codes implement KNN_search in OpenSearch using DSL
|
|
# Besides, Opensearch's DSL for KNN_search query syntax differs from that in Elasticsearch, I also made some adaptions for it
|
|
elif isinstance(m, MatchDenseExpr):
|
|
assert (bqry is not None)
|
|
similarity = 0.0
|
|
if "similarity" in m.extra_options:
|
|
similarity = m.extra_options["similarity"]
|
|
use_knn = True
|
|
vector_column_name = m.vector_column_name
|
|
knn_query[vector_column_name] = {}
|
|
knn_query[vector_column_name]["vector"] = list(m.embedding_data)
|
|
knn_query[vector_column_name]["k"] = m.topn
|
|
knn_query[vector_column_name]["filter"] = bqry.to_dict()
|
|
knn_query[vector_column_name]["boost"] = similarity
|
|
|
|
if bqry and rank_feature:
|
|
for fld, sc in rank_feature.items():
|
|
if fld != PAGERANK_FLD:
|
|
fld = f"{TAG_FLD}.{fld}"
|
|
bqry.should.append(Q("rank_feature", field=fld, linear={}, boost=sc))
|
|
|
|
if bqry:
|
|
s = s.query(bqry)
|
|
for field in highlightFields:
|
|
s = s.highlight(field, force_source=True, no_match_size=30, require_field_match=False)
|
|
|
|
if orderBy:
|
|
orders = list()
|
|
for field, order in orderBy.fields:
|
|
order = "asc" if order == 0 else "desc"
|
|
if field in ["page_num_int", "top_int"]:
|
|
order_info = {"order": order, "unmapped_type": "float",
|
|
"mode": "avg", "numeric_type": "double"}
|
|
elif field.endswith("_int") or field.endswith("_flt"):
|
|
order_info = {"order": order, "unmapped_type": "float"}
|
|
else:
|
|
order_info = {"order": order, "unmapped_type": "text"}
|
|
orders.append({field: order_info})
|
|
s = s.sort(*orders)
|
|
|
|
for fld in aggFields:
|
|
s.aggs.bucket(f'aggs_{fld}', 'terms', field=fld, size=1000000)
|
|
|
|
if limit > 0:
|
|
s = s[offset:offset + limit]
|
|
q = s.to_dict()
|
|
logger.debug(f"OSConnection.search {str(indexNames)} query: " + json.dumps(q))
|
|
|
|
if use_knn:
|
|
del q["query"]
|
|
q["query"] = {"knn": knn_query}
|
|
|
|
for i in range(ATTEMPT_TIME):
|
|
try:
|
|
res = self.os.search(index=indexNames,
|
|
body=q,
|
|
timeout=600,
|
|
# search_type="dfs_query_then_fetch",
|
|
track_total_hits=True,
|
|
_source=True)
|
|
if str(res.get("timed_out", "")).lower() == "true":
|
|
raise Exception("OpenSearch Timeout.")
|
|
logger.debug(f"OSConnection.search {str(indexNames)} res: " + str(res))
|
|
return res
|
|
except Exception as e:
|
|
logger.exception(f"OSConnection.search {str(indexNames)} query: " + str(q))
|
|
if str(e).find("Timeout") > 0:
|
|
continue
|
|
raise e
|
|
logger.error(f"OSConnection.search timeout for {ATTEMPT_TIME} times!")
|
|
raise Exception("OSConnection.search timeout.")
|
|
|
|
def get(self, chunkId: str, indexName: str, knowledgebaseIds: list[str]) -> dict | None:
|
|
for i in range(ATTEMPT_TIME):
|
|
try:
|
|
res = self.os.get(index=(indexName),
|
|
id=chunkId, _source=True, )
|
|
if str(res.get("timed_out", "")).lower() == "true":
|
|
raise Exception("Es Timeout.")
|
|
chunk = res["_source"]
|
|
chunk["id"] = chunkId
|
|
return chunk
|
|
except NotFoundError:
|
|
return None
|
|
except Exception as e:
|
|
logger.exception(f"OSConnection.get({chunkId}) got exception")
|
|
if str(e).find("Timeout") > 0:
|
|
continue
|
|
raise e
|
|
logger.error(f"OSConnection.get timeout for {ATTEMPT_TIME} times!")
|
|
raise Exception("OSConnection.get timeout.")
|
|
|
|
def insert(self, documents: list[dict], indexName: str, knowledgebaseId: str = None) -> list[str]:
|
|
# Refers to https://opensearch.org/docs/latest/api-reference/document-apis/bulk/
|
|
operations = []
|
|
for d in documents:
|
|
assert "_id" not in d
|
|
assert "id" in d
|
|
d_copy = copy.deepcopy(d)
|
|
meta_id = d_copy.pop("id", "")
|
|
operations.append(
|
|
{"index": {"_index": indexName, "_id": meta_id}})
|
|
operations.append(d_copy)
|
|
|
|
res = []
|
|
for _ in range(ATTEMPT_TIME):
|
|
try:
|
|
res = []
|
|
r = self.os.bulk(index=(indexName), body=operations,
|
|
refresh="wait_for", timeout=60)
|
|
if re.search(r"False", str(r["errors"]), re.IGNORECASE):
|
|
return res
|
|
|
|
for item in r["items"]:
|
|
for action in ["create", "delete", "index", "update"]:
|
|
if action in item and "error" in item[action]:
|
|
res.append(str(item[action]["_id"]) + ":" + str(item[action]["error"]))
|
|
return res
|
|
except Exception as e:
|
|
res.append(str(e))
|
|
logger.warning("OSConnection.insert got exception: " + str(e))
|
|
res = []
|
|
if re.search(r"(Timeout|time out)", str(e), re.IGNORECASE):
|
|
res.append(str(e))
|
|
time.sleep(3)
|
|
continue
|
|
return res
|
|
|
|
def update(self, condition: dict, newValue: dict, indexName: str, knowledgebaseId: str) -> bool:
|
|
doc = copy.deepcopy(newValue)
|
|
doc.pop("id", None)
|
|
if "id" in condition and isinstance(condition["id"], str):
|
|
# update specific single document
|
|
chunkId = condition["id"]
|
|
for i in range(ATTEMPT_TIME):
|
|
doc_part = copy.deepcopy(doc)
|
|
remove_value = doc_part.pop("remove", None)
|
|
remove_field = remove_value if isinstance(remove_value, str) else None
|
|
remove_dict = remove_value if isinstance(remove_value, dict) else None
|
|
try:
|
|
if remove_field is not None:
|
|
self.os.update(
|
|
index=indexName,
|
|
id=chunkId,
|
|
body={"script": {"source": f"ctx._source.remove('{remove_field}');"}},
|
|
)
|
|
if remove_dict is not None:
|
|
scripts = []
|
|
params = {}
|
|
for kk, vv in remove_dict.items():
|
|
scripts.append(
|
|
f"if (ctx._source.containsKey('{kk}') && ctx._source.{kk} != null) "
|
|
f"{{ int i = ctx._source.{kk}.indexOf(params.p_{kk}); "
|
|
f"if (i >= 0) {{ ctx._source.{kk}.remove(i); }} }}"
|
|
)
|
|
params[f"p_{kk}"] = vv
|
|
if scripts:
|
|
self.os.update(
|
|
index=indexName,
|
|
id=chunkId,
|
|
body={"script": {"source": "".join(scripts), "params": params}},
|
|
)
|
|
if doc_part:
|
|
self.os.update(index=indexName, id=chunkId, body={"doc": doc_part})
|
|
if remove_field is not None or remove_dict is not None or doc_part:
|
|
return True
|
|
except Exception as e:
|
|
logger.exception(
|
|
f"OSConnection.update(index={indexName}, id={id}, doc={json.dumps(condition, ensure_ascii=False)}) got exception")
|
|
if re.search(r"(timeout|connection)", str(e).lower()):
|
|
continue
|
|
break
|
|
return False
|
|
|
|
# update unspecific maybe-multiple documents
|
|
bqry = Q("bool")
|
|
for k, v in condition.items():
|
|
if not isinstance(k, str) or not v:
|
|
continue
|
|
if k == "exists":
|
|
bqry.filter.append(Q("exists", field=v))
|
|
continue
|
|
if isinstance(v, list):
|
|
bqry.filter.append(Q("terms", **{k: v}))
|
|
elif isinstance(v, str) or isinstance(v, int):
|
|
bqry.filter.append(Q("term", **{k: v}))
|
|
else:
|
|
raise Exception(
|
|
f"Condition `{str(k)}={str(v)}` value type is {str(type(v))}, expected to be int, str or list.")
|
|
scripts = []
|
|
params = {}
|
|
for k, v in newValue.items():
|
|
if k == "remove":
|
|
if isinstance(v, str):
|
|
scripts.append(f"ctx._source.remove('{v}');")
|
|
if isinstance(v, dict):
|
|
for kk, vv in v.items():
|
|
scripts.append(f"int i=ctx._source.{kk}.indexOf(params.p_{kk});ctx._source.{kk}.remove(i);")
|
|
params[f"p_{kk}"] = vv
|
|
continue
|
|
if k == "add":
|
|
if isinstance(v, dict):
|
|
for kk, vv in v.items():
|
|
scripts.append(f"ctx._source.{kk}.add(params.pp_{kk});")
|
|
params[f"pp_{kk}"] = vv.strip()
|
|
continue
|
|
if (not isinstance(k, str) or not v) and k != "available_int":
|
|
continue
|
|
if isinstance(v, str):
|
|
v = re.sub(r"(['\n\r]|\\.)", " ", v)
|
|
params[f"pp_{k}"] = v
|
|
scripts.append(f"ctx._source.{k}=params.pp_{k};")
|
|
elif isinstance(v, int) or isinstance(v, float):
|
|
scripts.append(f"ctx._source.{k}={v};")
|
|
elif isinstance(v, list):
|
|
scripts.append(f"ctx._source.{k}=params.pp_{k};")
|
|
params[f"pp_{k}"] = json.dumps(v, ensure_ascii=False)
|
|
else:
|
|
raise Exception(
|
|
f"newValue `{str(k)}={str(v)}` value type is {str(type(v))}, expected to be int, str.")
|
|
ubq = UpdateByQuery(
|
|
index=indexName).using(
|
|
self.os).query(bqry)
|
|
ubq = ubq.script(source="".join(scripts), params=params)
|
|
ubq = ubq.params(refresh=True)
|
|
ubq = ubq.params(slices=5)
|
|
ubq = ubq.params(conflicts="proceed")
|
|
|
|
for _ in range(ATTEMPT_TIME):
|
|
try:
|
|
_ = ubq.execute()
|
|
return True
|
|
except Exception as e:
|
|
logger.error("OSConnection.update got exception: " + str(e) + "\n".join(scripts))
|
|
if re.search(r"(timeout|connection|conflict)", str(e).lower()):
|
|
continue
|
|
break
|
|
return False
|
|
|
|
def adjust_chunk_pagerank_fea(
|
|
self,
|
|
chunk_id: str,
|
|
indexName: str,
|
|
knowledgebaseId: str,
|
|
delta: float,
|
|
min_w: float = 0.0,
|
|
max_w: float = 100.0,
|
|
row_id: int | None = None,
|
|
) -> bool:
|
|
"""Atomically adjust pagerank_fea on one chunk (painless script)."""
|
|
_ = row_id
|
|
try:
|
|
self.os.update(
|
|
index=indexName,
|
|
id=chunk_id,
|
|
retry_on_conflict=3,
|
|
body={
|
|
"script": {
|
|
"source": _PAGERANK_FEA_ADJUST_SCRIPT.strip(),
|
|
"lang": "painless",
|
|
"params": {
|
|
"pf": PAGERANK_FLD,
|
|
"delta": float(delta),
|
|
"min_w": float(min_w),
|
|
"max_w": float(max_w),
|
|
},
|
|
}
|
|
},
|
|
)
|
|
logger.debug(
|
|
"OSConnection.adjust_chunk_pagerank_fea(index=%s, id=%s, delta=%s) succeeded",
|
|
indexName,
|
|
chunk_id,
|
|
delta,
|
|
)
|
|
return True
|
|
except Exception as e:
|
|
logger.exception(
|
|
"OSConnection.adjust_chunk_pagerank_fea(index=%s, id=%s): %s",
|
|
indexName,
|
|
chunk_id,
|
|
e,
|
|
)
|
|
return False
|
|
|
|
def delete(self, condition: dict, indexName: str, knowledgebaseId: str) -> int:
|
|
assert "_id" not in condition
|
|
condition["kb_id"] = knowledgebaseId
|
|
|
|
# Build a bool query that combines id filter with other conditions
|
|
bool_query = Q("bool")
|
|
|
|
# Handle chunk IDs if present
|
|
if "id" in condition:
|
|
chunk_ids = condition["id"]
|
|
if not isinstance(chunk_ids, list):
|
|
chunk_ids = [chunk_ids]
|
|
if chunk_ids:
|
|
# Filter by specific chunk IDs
|
|
bool_query.filter.append(Q("ids", values=chunk_ids))
|
|
# If chunk_ids is empty, we don't add an ids filter - rely on other conditions
|
|
|
|
# Add all other conditions as filters
|
|
for k, v in condition.items():
|
|
if k == "id":
|
|
continue # Already handled above
|
|
if k == "exists":
|
|
bool_query.filter.append(Q("exists", field=v))
|
|
elif k == "must_not":
|
|
if isinstance(v, dict):
|
|
for kk, vv in v.items():
|
|
if kk == "exists":
|
|
bool_query.must_not.append(Q("exists", field=vv))
|
|
elif isinstance(v, list):
|
|
bool_query.must.append(Q("terms", **{k: v}))
|
|
elif isinstance(v, str) or isinstance(v, int):
|
|
bool_query.must.append(Q("term", **{k: v}))
|
|
elif v is not None:
|
|
raise Exception("Condition value must be int, str or list.")
|
|
|
|
# If no filters were added, use match_all (for tenant-wide operations)
|
|
if not bool_query.filter and not bool_query.must and not bool_query.must_not:
|
|
qry = Q("match_all")
|
|
else:
|
|
qry = bool_query
|
|
logger.debug("OSConnection.delete query: " + json.dumps(qry.to_dict()))
|
|
for _ in range(ATTEMPT_TIME):
|
|
try:
|
|
# print(Search().query(qry).to_dict(), flush=True)
|
|
res = self.os.delete_by_query(
|
|
index=indexName,
|
|
body=Search().query(qry).to_dict(),
|
|
refresh=True)
|
|
return res["deleted"]
|
|
except Exception as e:
|
|
logger.warning("OSConnection.delete got exception: " + str(e))
|
|
if re.search(r"(timeout|connection)", str(e).lower()):
|
|
time.sleep(3)
|
|
continue
|
|
if re.search(r"(not_found)", str(e), re.IGNORECASE):
|
|
return 0
|
|
return 0
|
|
|
|
"""
|
|
Helper functions for search result
|
|
"""
|
|
|
|
def get_total(self, res):
|
|
if isinstance(res["hits"]["total"], type({})):
|
|
return res["hits"]["total"]["value"]
|
|
return res["hits"]["total"]
|
|
|
|
def get_doc_ids(self, res):
|
|
return [d["_id"] for d in res["hits"]["hits"]]
|
|
|
|
def __getSource(self, res):
|
|
rr = []
|
|
for d in res["hits"]["hits"]:
|
|
d["_source"]["id"] = d["_id"]
|
|
d["_source"]["_score"] = d["_score"]
|
|
rr.append(d["_source"])
|
|
return rr
|
|
|
|
def get_fields(self, res, fields: list[str]) -> dict[str, dict]:
|
|
res_fields = {}
|
|
if not fields:
|
|
return {}
|
|
for d in self.__getSource(res):
|
|
m = {n: d.get(n) for n in fields if d.get(n) is not None}
|
|
for n, v in m.items():
|
|
if isinstance(v, list):
|
|
m[n] = v
|
|
continue
|
|
if not isinstance(v, str):
|
|
m[n] = str(m[n])
|
|
# if n.find("tks") > 0:
|
|
# m[n] = remove_redundant_spaces(m[n])
|
|
|
|
if m:
|
|
res_fields[d["id"]] = m
|
|
return res_fields
|
|
|
|
def get_highlight(self, res, keywords: list[str], fieldnm: str):
|
|
ans = {}
|
|
for d in res["hits"]["hits"]:
|
|
hlts = d.get("highlight")
|
|
if not hlts:
|
|
continue
|
|
txt = "...".join([a for a in list(hlts.items())[0][1]])
|
|
if not is_english(txt.split()):
|
|
ans[d["_id"]] = txt
|
|
continue
|
|
|
|
txt = d["_source"][fieldnm]
|
|
txt = re.sub(r"[\r\n]", " ", txt, flags=re.IGNORECASE | re.MULTILINE)
|
|
txts = []
|
|
for t in re.split(r"[.?!;\n]", txt):
|
|
for w in keywords:
|
|
t = re.sub(r"(^|[ .?/'\"\(\)!,:;-])(%s)([ .?/'\"\(\)!,:;-])" % re.escape(w), r"\1<em>\2</em>\3", t,
|
|
flags=re.IGNORECASE | re.MULTILINE)
|
|
if not re.search(r"<em>[^<>]+</em>", t, flags=re.IGNORECASE | re.MULTILINE):
|
|
continue
|
|
txts.append(t)
|
|
ans[d["_id"]] = "...".join(txts) if txts else "...".join([a for a in list(hlts.items())[0][1]])
|
|
|
|
return ans
|
|
|
|
def get_aggregation(self, res, fieldnm: str):
|
|
agg_field = "aggs_" + fieldnm
|
|
if "aggregations" not in res or agg_field not in res["aggregations"]:
|
|
return list()
|
|
bkts = res["aggregations"][agg_field]["buckets"]
|
|
return [(b["key"], b["doc_count"]) for b in bkts]
|
|
|
|
"""
|
|
SQL
|
|
"""
|
|
|
|
def sql(self, sql: str, fetch_size: int, format: str):
|
|
logger.debug(f"OSConnection.sql get sql: {sql}")
|
|
sql = re.sub(r"[ `]+", " ", sql)
|
|
sql = sql.replace("%", "")
|
|
replaces = []
|
|
for r in re.finditer(r" ([a-z_]+_l?tks)( like | ?= ?)'([^']+)'", sql):
|
|
fld, v = r.group(1), r.group(3)
|
|
match = " MATCH({}, '{}', 'operator=OR;minimum_should_match=30%') ".format(
|
|
fld, rag_tokenizer.fine_grained_tokenize(rag_tokenizer.tokenize(v)))
|
|
replaces.append(
|
|
("{}{}'{}'".format(
|
|
r.group(1),
|
|
r.group(2),
|
|
r.group(3)),
|
|
match))
|
|
|
|
for p, r in replaces:
|
|
sql = sql.replace(p, r, 1)
|
|
logger.debug(f"OSConnection.sql to os: {sql}")
|
|
|
|
for i in range(ATTEMPT_TIME):
|
|
try:
|
|
res = self.os.sql.query(body={"query": sql, "fetch_size": fetch_size}, format=format,
|
|
request_timeout="2s")
|
|
return res
|
|
except ConnectionTimeout:
|
|
logger.exception("OSConnection.sql timeout")
|
|
continue
|
|
except Exception:
|
|
logger.exception("OSConnection.sql got exception")
|
|
return None
|
|
logger.error(f"OSConnection.sql timeout for {ATTEMPT_TIME} times!")
|
|
return None
|