mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-19 14:11:04 +08:00
Feat: Add knowledge compilation workflows (#16515)
## Summary - Add knowledge compilation template APIs, services, and builtin template seed data - Add advanced knowledge compile structure/artifact/RAPTOR workflow support - Update parsing, dataset/document APIs, and supporting services for compilation workflows
This commit is contained in:
Kevin Hu
@@ -65,12 +65,7 @@ class ESConnection(ESConnectionBase):
|
||||
"""
|
||||
|
||||
def _es_search_once(self, index_names: list[str], query: dict, track_total_hits: bool):
|
||||
return self.es.search(
|
||||
index=index_names,
|
||||
body=query,
|
||||
timeout="600s",
|
||||
track_total_hits=track_total_hits
|
||||
)
|
||||
return self.es.search(index=index_names, body=query, timeout="600s", track_total_hits=track_total_hits)
|
||||
|
||||
def _search_with_search_after(self, index_names: list[str], query: dict, offset: int, limit: int):
|
||||
q_base = copy.deepcopy(query)
|
||||
@@ -139,17 +134,18 @@ class ESConnection(ESConnectionBase):
|
||||
return template_res
|
||||
|
||||
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
|
||||
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,
|
||||
):
|
||||
"""
|
||||
Refers to https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl.html
|
||||
@@ -166,19 +162,22 @@ class ESConnection(ESConnectionBase):
|
||||
if v == 0:
|
||||
bool_query.filter.append(Q("range", available_int={"lt": 1}))
|
||||
else:
|
||||
bool_query.filter.append(
|
||||
Q("bool", must_not=Q("range", available_int={"lt": 1})))
|
||||
bool_query.filter.append(Q("bool", must_not=Q("range", available_int={"lt": 1})))
|
||||
continue
|
||||
if k == "id":
|
||||
if not v:
|
||||
continue
|
||||
if isinstance(v, list):
|
||||
bool_query.filter.append(
|
||||
Q("bool", should=[Q("terms", id=v), Q("terms", _id=v)], minimum_should_match=1))
|
||||
bool_query.filter.append(Q("bool", should=[Q("terms", id=v), Q("terms", _id=v)], minimum_should_match=1))
|
||||
elif isinstance(v, str) or isinstance(v, int):
|
||||
bool_query.filter.append(
|
||||
Q("bool", should=[Q("term", id=v), Q("term", _id=v)], minimum_should_match=1))
|
||||
bool_query.filter.append(Q("bool", should=[Q("term", id=v), Q("term", _id=v)], minimum_should_match=1))
|
||||
continue
|
||||
if 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))
|
||||
continue
|
||||
if not v:
|
||||
continue
|
||||
if isinstance(v, list):
|
||||
@@ -186,17 +185,18 @@ class ESConnection(ESConnectionBase):
|
||||
elif isinstance(v, str) or isinstance(v, int):
|
||||
bool_query.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.")
|
||||
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 match_expressions:
|
||||
if isinstance(m, FusionExpr) and m.method == "weighted_sum" and "weights" in m.fusion_params:
|
||||
assert len(match_expressions) == 3 and isinstance(match_expressions[0], MatchTextExpr) and isinstance(
|
||||
match_expressions[1],
|
||||
MatchDenseExpr) and isinstance(
|
||||
match_expressions[2], FusionExpr)
|
||||
assert (
|
||||
len(match_expressions) == 3
|
||||
and isinstance(match_expressions[0], MatchTextExpr)
|
||||
and isinstance(match_expressions[1], MatchDenseExpr)
|
||||
and isinstance(match_expressions[2], FusionExpr)
|
||||
)
|
||||
weights = m.fusion_params["weights"]
|
||||
vector_similarity_weight = get_float(weights.split(",")[1])
|
||||
for m in match_expressions:
|
||||
@@ -204,24 +204,22 @@ class ESConnection(ESConnectionBase):
|
||||
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)) + "%"
|
||||
bool_query.must.append(Q("query_string", fields=m.fields,
|
||||
type="best_fields", query=m.matching_text,
|
||||
minimum_should_match=minimum_should_match,
|
||||
boost=1))
|
||||
bool_query.must.append(Q("query_string", fields=m.fields, type="best_fields", query=m.matching_text, minimum_should_match=minimum_should_match, boost=1))
|
||||
bool_query.boost = 1.0 - vector_similarity_weight
|
||||
|
||||
elif isinstance(m, MatchDenseExpr):
|
||||
assert (bool_query is not None)
|
||||
assert bool_query is not None
|
||||
similarity = 0.0
|
||||
if "similarity" in m.extra_options:
|
||||
similarity = m.extra_options["similarity"]
|
||||
s = s.knn(m.vector_column_name,
|
||||
m.topn,
|
||||
m.topn * 2,
|
||||
query_vector=list(m.embedding_data),
|
||||
filter=bool_query.to_dict(),
|
||||
similarity=similarity,
|
||||
)
|
||||
s = s.knn(
|
||||
m.vector_column_name,
|
||||
m.topn,
|
||||
m.topn * 2,
|
||||
query_vector=list(m.embedding_data),
|
||||
filter=bool_query.to_dict(),
|
||||
similarity=similarity,
|
||||
)
|
||||
|
||||
if bool_query and rank_feature:
|
||||
for fld, sc in rank_feature.items():
|
||||
@@ -239,31 +237,25 @@ class ESConnection(ESConnectionBase):
|
||||
for field, order in order_by.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"}
|
||||
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"}
|
||||
elif field == "id":
|
||||
continue # id as "text", not a "keyword", order by it will cause error
|
||||
continue # id as "text", not a "keyword", order by it will cause error
|
||||
else:
|
||||
order_info = {"order": order, "unmapped_type": "keyword"}
|
||||
orders.append({field: order_info})
|
||||
s = s.sort(*orders)
|
||||
if agg_fields:
|
||||
for fld in agg_fields:
|
||||
s.aggs.bucket(f'aggs_{fld}', 'terms', field=fld, size=1000000)
|
||||
s.aggs.bucket(f"aggs_{fld}", "terms", field=fld, size=1000000)
|
||||
|
||||
has_dense = any(isinstance(m, MatchDenseExpr) for m in match_expressions)
|
||||
has_explicit_sort = bool(order_by and order_by.fields)
|
||||
use_search_after = (
|
||||
limit > 0
|
||||
and (offset + limit > MAX_RESULT_WINDOW)
|
||||
and has_explicit_sort
|
||||
and not has_dense
|
||||
)
|
||||
use_search_after = limit > 0 and (offset + limit > MAX_RESULT_WINDOW) and has_explicit_sort and not has_dense
|
||||
|
||||
if limit > 0 and not use_search_after:
|
||||
s = s[offset:offset + limit]
|
||||
s = s[offset : offset + limit]
|
||||
# Filter _source to only requested fields for efficiency, and add vector
|
||||
# fields to "fields" param so they appear in hit.fields when ES 9.x
|
||||
# exclude_source_vectors is enabled (dense_vector not in _source).
|
||||
@@ -295,7 +287,7 @@ class ESConnection(ESConnectionBase):
|
||||
continue
|
||||
except Exception as e:
|
||||
# Only log debug for NotFoundError(accepted when metadata index doesn't exist)
|
||||
if 'NotFound' in str(e):
|
||||
if "NotFound" in str(e):
|
||||
self.logger.debug(f"ESConnection.search {str(index_names)} query: " + str(q) + " - " + str(e))
|
||||
else:
|
||||
self.logger.exception(f"ESConnection.search {str(index_names)} query: " + str(q) + str(e))
|
||||
@@ -314,16 +306,14 @@ class ESConnection(ESConnectionBase):
|
||||
d_copy["kb_id"] = knowledgebase_id
|
||||
# Use id as _id for uniqueness, also keep "id" as a regular field for sorting
|
||||
meta_id = d_copy.get("id", "")
|
||||
operations.append(
|
||||
{"index": {"_index": index_name, "_id": meta_id}})
|
||||
operations.append({"index": {"_index": index_name, "_id": meta_id}})
|
||||
operations.append(d_copy)
|
||||
|
||||
res = []
|
||||
for _ in range(ATTEMPT_TIME):
|
||||
try:
|
||||
res = []
|
||||
r = self.es.bulk(index=index_name, operations=operations,
|
||||
refresh="wait_for", timeout="60s")
|
||||
r = self.es.bulk(index=index_name, operations=operations, refresh="wait_for", timeout="60s")
|
||||
if re.search(r"False", str(r["errors"]), re.IGNORECASE):
|
||||
return res
|
||||
|
||||
@@ -359,10 +349,9 @@ class ESConnection(ESConnectionBase):
|
||||
if "feas" != k.split("_")[-1]:
|
||||
continue
|
||||
try:
|
||||
self.es.update(index=index_name, id=chunk_id, script=f"ctx._source.remove(\"{k}\");")
|
||||
self.es.update(index=index_name, id=chunk_id, script=f'ctx._source.remove("{k}");')
|
||||
except Exception:
|
||||
self.logger.exception(
|
||||
f"ESConnection.update(index={index_name}, id={chunk_id}, doc={json.dumps(condition, ensure_ascii=False)}) got exception")
|
||||
self.logger.exception(f"ESConnection.update(index={index_name}, id={chunk_id}, doc={json.dumps(condition, ensure_ascii=False)}) got exception")
|
||||
try:
|
||||
if remove_field is not None:
|
||||
self.es.update(
|
||||
@@ -375,9 +364,7 @@ class ESConnection(ESConnectionBase):
|
||||
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); }} }}"
|
||||
f"if (ctx._source.containsKey('{kk}') && ctx._source.{kk} != null) {{ int i = ctx._source.{kk}.indexOf(params.p_{kk}); if (i >= 0) {{ ctx._source.{kk}.remove(i); }} }}"
|
||||
)
|
||||
params[f"p_{kk}"] = vv
|
||||
if scripts:
|
||||
@@ -391,9 +378,7 @@ class ESConnection(ESConnectionBase):
|
||||
if remove_field is not None or remove_dict is not None or doc_part:
|
||||
return True
|
||||
except Exception as e:
|
||||
self.logger.exception(
|
||||
f"ESConnection.update(index={index_name}, id={chunk_id}, doc={json.dumps(condition, ensure_ascii=False)}) got exception: " + str(
|
||||
e))
|
||||
self.logger.exception(f"ESConnection.update(index={index_name}, id={chunk_id}, doc={json.dumps(condition, ensure_ascii=False)}) got exception: " + str(e))
|
||||
break
|
||||
return False
|
||||
|
||||
@@ -405,13 +390,18 @@ class ESConnection(ESConnectionBase):
|
||||
if k == "exists":
|
||||
bool_query.filter.append(Q("exists", field=v))
|
||||
continue
|
||||
if 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))
|
||||
continue
|
||||
if isinstance(v, list):
|
||||
bool_query.filter.append(Q("terms", **{k: v}))
|
||||
elif isinstance(v, str) or isinstance(v, int):
|
||||
bool_query.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.")
|
||||
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 new_value.items():
|
||||
@@ -441,11 +431,8 @@ class ESConnection(ESConnectionBase):
|
||||
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=index_name).using(
|
||||
self.es).query(bool_query)
|
||||
raise Exception(f"newValue `{str(k)}={str(v)}` value type is {str(type(v))}, expected to be int, str.")
|
||||
ubq = UpdateByQuery(index=index_name).using(self.es).query(bool_query)
|
||||
ubq = ubq.script(source="".join(scripts), params=params)
|
||||
ubq = ubq.params(refresh=True)
|
||||
ubq = ubq.params(slices=5)
|
||||
@@ -563,10 +550,7 @@ class ESConnection(ESConnectionBase):
|
||||
self.logger.debug("ESConnection.delete query: " + json.dumps(qry.to_dict()))
|
||||
for _ in range(ATTEMPT_TIME):
|
||||
try:
|
||||
res = self.es.delete_by_query(
|
||||
index=index_name,
|
||||
body=Search().query(qry).to_dict(),
|
||||
refresh=True)
|
||||
res = self.es.delete_by_query(index=index_name, body=Search().query(qry).to_dict(), refresh=True)
|
||||
return res["deleted"]
|
||||
except ConnectionTimeout:
|
||||
self.logger.exception("ES request timeout")
|
||||
|
||||
@@ -26,8 +26,7 @@ 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 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
|
||||
@@ -58,7 +57,7 @@ if (nw <= 0.0) {
|
||||
}
|
||||
"""
|
||||
|
||||
logger = logging.getLogger('ragflow.opensearch_conn')
|
||||
logger = logging.getLogger("ragflow.opensearch_conn")
|
||||
|
||||
|
||||
@singleton
|
||||
@@ -70,10 +69,9 @@ class OSConnection(DocStoreConnection):
|
||||
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,
|
||||
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
|
||||
timeout=600,
|
||||
)
|
||||
if self.os:
|
||||
self.info = self.os.info()
|
||||
@@ -118,8 +116,7 @@ class OSConnection(DocStoreConnection):
|
||||
warning, leave it off, and search() keeps doing vector-only.
|
||||
"""
|
||||
self.hybrid_search_enabled = False
|
||||
self._hybrid_pipeline = os.environ.get("OS_HYBRID_PIPELINE") \
|
||||
or settings.OS.get("hybrid_search_pipeline") or "ragflow_hybrid_pipeline"
|
||||
self._hybrid_pipeline = os.environ.get("OS_HYBRID_PIPELINE") or settings.OS.get("hybrid_search_pipeline") or "ragflow_hybrid_pipeline"
|
||||
|
||||
version_number = self.info.get("version", {}).get("number", "")
|
||||
try:
|
||||
@@ -127,34 +124,32 @@ class OSConnection(DocStoreConnection):
|
||||
except (ValueError, AttributeError):
|
||||
version = (0, 0)
|
||||
if version < self.HYBRID_MIN_VERSION:
|
||||
logger.warning(f"OpenSearch {version_number or 'unknown'} does not support the "
|
||||
f"normalization-processor (requires >= {self.HYBRID_MIN_VERSION[0]}."
|
||||
f"{self.HYBRID_MIN_VERSION[1]}); hybrid search is disabled and "
|
||||
f"queries fall back to vector-only.")
|
||||
logger.warning(
|
||||
f"OpenSearch {version_number or 'unknown'} does not support the "
|
||||
f"normalization-processor (requires >= {self.HYBRID_MIN_VERSION[0]}."
|
||||
f"{self.HYBRID_MIN_VERSION[1]}); hybrid search is disabled and "
|
||||
f"queries fall back to vector-only."
|
||||
)
|
||||
return
|
||||
|
||||
weights = settings.OS.get("hybrid_search_weights", [0.5, 0.5])
|
||||
pipeline_body = {
|
||||
"description": "RAGFlow hybrid search normalization pipeline (BM25 + KNN).",
|
||||
"phase_results_processors": [
|
||||
{"normalization-processor": {
|
||||
"normalization": {"technique": "min_max"},
|
||||
"combination": {"technique": "arithmetic_mean",
|
||||
"parameters": {"weights": weights}}}}
|
||||
],
|
||||
"phase_results_processors": [{"normalization-processor": {"normalization": {"technique": "min_max"}, "combination": {"technique": "arithmetic_mean", "parameters": {"weights": weights}}}}],
|
||||
}
|
||||
try:
|
||||
self.os.transport.perform_request(
|
||||
"PUT", f"/_search/pipeline/{self._hybrid_pipeline}", body=pipeline_body)
|
||||
self.os.transport.perform_request("PUT", f"/_search/pipeline/{self._hybrid_pipeline}", body=pipeline_body)
|
||||
self.hybrid_search_enabled = True
|
||||
logger.info(f"OpenSearch hybrid search enabled via pipeline "
|
||||
f"'{self._hybrid_pipeline}' (weights {weights}).")
|
||||
logger.info(f"OpenSearch hybrid search enabled via pipeline '{self._hybrid_pipeline}' (weights {weights}).")
|
||||
except Exception:
|
||||
logger.warning(f"Could not create OpenSearch search pipeline '{self._hybrid_pipeline}'; "
|
||||
f"hybrid search is disabled and queries fall back to vector-only. "
|
||||
f"Creating a search pipeline needs the "
|
||||
f"'cluster:admin/search/pipeline/put' privilege (relevant on "
|
||||
f"locked-down or managed OpenSearch).", exc_info=True)
|
||||
logger.warning(
|
||||
f"Could not create OpenSearch search pipeline '{self._hybrid_pipeline}'; "
|
||||
f"hybrid search is disabled and queries fall back to vector-only. "
|
||||
f"Creating a search pipeline needs the "
|
||||
f"'cluster:admin/search/pipeline/put' privilege (relevant on "
|
||||
f"locked-down or managed OpenSearch).",
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
"""
|
||||
Database operations
|
||||
@@ -177,8 +172,8 @@ class OSConnection(DocStoreConnection):
|
||||
return True
|
||||
try:
|
||||
from opensearchpy.client import IndicesClient
|
||||
return IndicesClient(self.os).create(index=indexName,
|
||||
body=self.mapping)
|
||||
|
||||
return IndicesClient(self.os).create(index=indexName, body=self.mapping)
|
||||
except Exception:
|
||||
logger.exception("OSConnection.createIndex error %s" % (indexName))
|
||||
|
||||
@@ -215,6 +210,7 @@ class OSConnection(DocStoreConnection):
|
||||
mappings = {**mappings, "dynamic": True}
|
||||
|
||||
from opensearchpy.client import IndicesClient
|
||||
|
||||
body = {
|
||||
"settings": doc_meta_mapping["settings"],
|
||||
"mappings": mappings,
|
||||
@@ -316,17 +312,18 @@ class OSConnection(DocStoreConnection):
|
||||
"""
|
||||
|
||||
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] = [],
|
||||
rank_feature: dict | None = None
|
||||
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] = [],
|
||||
rank_feature: dict | None = None,
|
||||
):
|
||||
"""
|
||||
Refers to https://github.com/opensearch-project/opensearch-py/blob/main/guides/dsl.md
|
||||
@@ -345,9 +342,22 @@ class OSConnection(DocStoreConnection):
|
||||
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})))
|
||||
bqry.filter.append(Q("bool", must_not=Q("range", available_int={"lt": 1})))
|
||||
continue
|
||||
if k == "id":
|
||||
if not v:
|
||||
continue
|
||||
if isinstance(v, list):
|
||||
bqry.filter.append(Q("bool", should=[Q("terms", id=v), Q("ids", values=v)], minimum_should_match=1))
|
||||
elif isinstance(v, str) or isinstance(v, int):
|
||||
bqry.filter.append(Q("bool", should=[Q("term", id=v), Q("ids", values=[v])], minimum_should_match=1))
|
||||
continue
|
||||
if k == "must_not":
|
||||
if isinstance(v, dict):
|
||||
for kk, vv in v.items():
|
||||
if kk == "exists":
|
||||
bqry.must_not.append(Q("exists", field=vv))
|
||||
continue
|
||||
if not v:
|
||||
continue
|
||||
if isinstance(v, list):
|
||||
@@ -355,16 +365,18 @@ class OSConnection(DocStoreConnection):
|
||||
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.")
|
||||
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 match_expressions:
|
||||
if isinstance(m, FusionExpr) and m.method == "weighted_sum" and "weights" in m.fusion_params:
|
||||
assert len(match_expressions) == 3 and isinstance(match_expressions[0], MatchTextExpr) and isinstance(match_expressions[1],
|
||||
MatchDenseExpr) and isinstance(
|
||||
match_expressions[2], FusionExpr)
|
||||
assert (
|
||||
len(match_expressions) == 3
|
||||
and isinstance(match_expressions[0], MatchTextExpr)
|
||||
and isinstance(match_expressions[1], MatchDenseExpr)
|
||||
and isinstance(match_expressions[2], FusionExpr)
|
||||
)
|
||||
weights = m.fusion_params["weights"]
|
||||
vector_similarity_weight = float(weights.split(",")[1])
|
||||
knn_query = {}
|
||||
@@ -374,10 +386,7 @@ class OSConnection(DocStoreConnection):
|
||||
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.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
|
||||
@@ -385,7 +394,7 @@ class OSConnection(DocStoreConnection):
|
||||
# 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)
|
||||
assert bqry is not None
|
||||
similarity = 0.0
|
||||
if "similarity" in m.extra_options:
|
||||
similarity = m.extra_options["similarity"]
|
||||
@@ -419,8 +428,7 @@ class OSConnection(DocStoreConnection):
|
||||
for field, order in order_by.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"}
|
||||
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:
|
||||
@@ -429,10 +437,10 @@ class OSConnection(DocStoreConnection):
|
||||
s = s.sort(*orders)
|
||||
|
||||
for fld in agg_fields:
|
||||
s.aggs.bucket(f'aggs_{fld}', 'terms', field=fld, size=1000000)
|
||||
s.aggs.bucket(f"aggs_{fld}", "terms", field=fld, size=1000000)
|
||||
|
||||
if limit > 0:
|
||||
s = s[offset:offset + limit]
|
||||
s = s[offset : offset + limit]
|
||||
q = s.to_dict()
|
||||
logger.debug(f"OSConnection.search {str(index_names)} query: " + json.dumps(q))
|
||||
|
||||
@@ -455,13 +463,15 @@ class OSConnection(DocStoreConnection):
|
||||
|
||||
for i in range(ATTEMPT_TIME):
|
||||
try:
|
||||
res = self.os.search(index=index_names,
|
||||
body=q,
|
||||
timeout=600,
|
||||
# search_type="dfs_query_then_fetch",
|
||||
track_total_hits=True,
|
||||
_source=True,
|
||||
**search_kwargs)
|
||||
res = self.os.search(
|
||||
index=index_names,
|
||||
body=q,
|
||||
timeout=600,
|
||||
# search_type="dfs_query_then_fetch",
|
||||
track_total_hits=True,
|
||||
_source=True,
|
||||
**search_kwargs,
|
||||
)
|
||||
if str(res.get("timed_out", "")).lower() == "true":
|
||||
raise Exception("OpenSearch Timeout.")
|
||||
logger.debug(f"OSConnection.search {str(index_names)} res: " + str(res))
|
||||
@@ -477,8 +487,11 @@ class OSConnection(DocStoreConnection):
|
||||
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, )
|
||||
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"]
|
||||
@@ -505,16 +518,14 @@ class OSConnection(DocStoreConnection):
|
||||
# doc-meta read path (DocMetadataService filters on / sorts by the
|
||||
# "id" field) can find it, mirroring ESConnection.insert().
|
||||
meta_id = d_copy.get("id", "")
|
||||
operations.append(
|
||||
{"index": {"_index": indexName, "_id": meta_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)
|
||||
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
|
||||
|
||||
@@ -556,9 +567,7 @@ class OSConnection(DocStoreConnection):
|
||||
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); }} }}"
|
||||
f"if (ctx._source.containsKey('{kk}') && ctx._source.{kk} != null) {{ int i = ctx._source.{kk}.indexOf(params.p_{kk}); if (i >= 0) {{ ctx._source.{kk}.remove(i); }} }}"
|
||||
)
|
||||
params[f"p_{kk}"] = vv
|
||||
if scripts:
|
||||
@@ -572,8 +581,7 @@ class OSConnection(DocStoreConnection):
|
||||
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")
|
||||
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
|
||||
@@ -592,8 +600,7 @@ class OSConnection(DocStoreConnection):
|
||||
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.")
|
||||
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():
|
||||
@@ -623,11 +630,8 @@ class OSConnection(DocStoreConnection):
|
||||
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)
|
||||
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)
|
||||
@@ -734,10 +738,7 @@ class OSConnection(DocStoreConnection):
|
||||
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)
|
||||
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))
|
||||
@@ -820,8 +821,7 @@ class OSConnection(DocStoreConnection):
|
||||
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)
|
||||
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)
|
||||
@@ -847,14 +847,8 @@ class OSConnection(DocStoreConnection):
|
||||
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))
|
||||
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)
|
||||
@@ -862,8 +856,7 @@ class OSConnection(DocStoreConnection):
|
||||
|
||||
for i in range(ATTEMPT_TIME):
|
||||
try:
|
||||
res = self.os.sql.query(body={"query": sql, "fetch_size": fetch_size}, format=format,
|
||||
request_timeout="2s")
|
||||
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")
|
||||
|
||||
@@ -216,6 +216,14 @@ class RedisDB:
|
||||
self.__open__()
|
||||
return False
|
||||
|
||||
def set_if_absent(self, k, v, exp=3600):
|
||||
try:
|
||||
return bool(self.REDIS.set(k, v, exp, nx=True))
|
||||
except Exception as e:
|
||||
logging.warning("RedisDB.set_if_absent " + str(k) + " got exception: " + str(e))
|
||||
self.__open__()
|
||||
return False
|
||||
|
||||
def sadd(self, key: str, member: str):
|
||||
try:
|
||||
self.REDIS.sadd(key, member)
|
||||
|
||||
Reference in New Issue
Block a user