mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-01 05:23:47 +08:00
Refactor: reformat all code for lefthook using ruff and gofmt (#16585)
This commit is contained in:
Wang Qi
@@ -34,7 +34,6 @@ ATTEMPT_TIME = 2
|
||||
|
||||
@singleton
|
||||
class ESConnection(ESConnectionBase):
|
||||
|
||||
@staticmethod
|
||||
def convert_field_name(field_name: str, use_tokenized_content=False) -> str:
|
||||
match field_name:
|
||||
@@ -111,18 +110,19 @@ class ESConnection(ESConnectionBase):
|
||||
"""
|
||||
|
||||
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],
|
||||
memory_ids: list[str],
|
||||
agg_fields: list[str] | None = None,
|
||||
rank_feature: dict | None = None,
|
||||
hide_forgotten: bool = True
|
||||
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],
|
||||
memory_ids: list[str],
|
||||
agg_fields: list[str] | None = None,
|
||||
rank_feature: dict | None = None,
|
||||
hide_forgotten: bool = True,
|
||||
):
|
||||
"""
|
||||
Refers to https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl.html
|
||||
@@ -154,15 +154,17 @@ class ESConnection(ESConnectionBase):
|
||||
elif isinstance(v, str) or isinstance(v, int):
|
||||
bool_query.filter.append(Q("term", **{field_name: 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:
|
||||
@@ -170,24 +172,31 @@ 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=[self.convert_field_name(f, use_tokenized_content=True) for f in 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=[self.convert_field_name(f, use_tokenized_content=True) for f in 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(self.convert_field_name(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(
|
||||
self.convert_field_name(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():
|
||||
@@ -207,7 +216,7 @@ class ESConnection(ESConnectionBase):
|
||||
if 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})
|
||||
@@ -215,22 +224,24 @@ class ESConnection(ESConnectionBase):
|
||||
|
||||
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)
|
||||
|
||||
if limit > 0:
|
||||
s = s[offset:offset + limit]
|
||||
s = s[offset : offset + limit]
|
||||
q = s.to_dict()
|
||||
self.logger.debug(f"ESConnection.search {str(index_names)} query: " + json.dumps(q))
|
||||
|
||||
for i in range(ATTEMPT_TIME):
|
||||
try:
|
||||
#print(json.dumps(q, ensure_ascii=False))
|
||||
res = self.es.search(index=exist_index_list,
|
||||
body=q,
|
||||
timeout="600s",
|
||||
# search_type="dfs_query_then_fetch",
|
||||
track_total_hits=True,
|
||||
_source=True)
|
||||
# print(json.dumps(q, ensure_ascii=False))
|
||||
res = self.es.search(
|
||||
index=exist_index_list,
|
||||
body=q,
|
||||
timeout="600s",
|
||||
# search_type="dfs_query_then_fetch",
|
||||
track_total_hits=True,
|
||||
_source=True,
|
||||
)
|
||||
if str(res.get("timed_out", "")).lower() == "true":
|
||||
raise Exception("Es Timeout.")
|
||||
self.logger.debug(f"ESConnection.search {str(index_names)} res: " + str(res))
|
||||
@@ -249,7 +260,7 @@ class ESConnection(ESConnectionBase):
|
||||
self.logger.error(f"ESConnection.search timeout for {ATTEMPT_TIME} times!")
|
||||
raise Exception("ESConnection.search timeout.")
|
||||
|
||||
def get_forgotten_messages(self, select_fields: list[str], index_name: str, memory_id: str, limit: int=512):
|
||||
def get_forgotten_messages(self, select_fields: list[str], index_name: str, memory_id: str, limit: int = 512):
|
||||
bool_query = Q("bool", must=[])
|
||||
bool_query.must.append(Q("exists", field="forget_at"))
|
||||
bool_query.filter.append(Q("term", memory_id=memory_id))
|
||||
@@ -292,7 +303,7 @@ class ESConnection(ESConnectionBase):
|
||||
self.logger.error(f"ESConnection.search timeout for {ATTEMPT_TIME} times!")
|
||||
raise Exception("ESConnection.search timeout.")
|
||||
|
||||
def get_missing_field_message(self, select_fields: list[str], index_name: str, memory_id: str, field_name: str, limit: int=512):
|
||||
def get_missing_field_message(self, select_fields: list[str], index_name: str, memory_id: str, field_name: str, limit: int = 512):
|
||||
if not self.index_exist(index_name):
|
||||
return None
|
||||
bool_query = Q("bool", must=[])
|
||||
@@ -340,8 +351,11 @@ class ESConnection(ESConnectionBase):
|
||||
def get(self, doc_id: str, index_name: str, memory_ids: list[str]) -> dict | None:
|
||||
for i in range(ATTEMPT_TIME):
|
||||
try:
|
||||
res = self.es.get(index=index_name,
|
||||
id=doc_id, source=True, )
|
||||
res = self.es.get(
|
||||
index=index_name,
|
||||
id=doc_id,
|
||||
source=True,
|
||||
)
|
||||
if str(res.get("timed_out", "")).lower() == "true":
|
||||
raise Exception("Es Timeout.")
|
||||
message = res["_source"]
|
||||
@@ -365,15 +379,13 @@ class ESConnection(ESConnectionBase):
|
||||
d_copy = self.map_message_to_es_fields(d_copy_raw)
|
||||
d_copy["memory_id"] = memory_id
|
||||
meta_id = d_copy.pop("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=False, timeout="60s")
|
||||
r = self.es.bulk(index=index_name, operations=operations, refresh=False, timeout="60s")
|
||||
if re.search(r"False", str(r["errors"]), re.IGNORECASE):
|
||||
return res
|
||||
|
||||
@@ -409,15 +421,14 @@ class ESConnection(ESConnectionBase):
|
||||
if "feas" != k.split("_")[-1]:
|
||||
continue
|
||||
try:
|
||||
self.es.update(index=index_name, id=message_id, script=f"ctx._source.remove(\"{k}\");")
|
||||
self.es.update(index=index_name, id=message_id, script=f'ctx._source.remove("{k}");')
|
||||
except Exception:
|
||||
self.logger.exception(f"ESConnection.update(index={index_name}, id={message_id}, doc={json.dumps(condition, ensure_ascii=False)}) got exception")
|
||||
try:
|
||||
self.es.update(index=index_name, id=message_id, doc=update_dict)
|
||||
return True
|
||||
except Exception as e:
|
||||
self.logger.exception(
|
||||
f"ESConnection.update(index={index_name}, id={message_id}, doc={json.dumps(condition, ensure_ascii=False)}) got exception: " + str(e))
|
||||
self.logger.exception(f"ESConnection.update(index={index_name}, id={message_id}, doc={json.dumps(condition, ensure_ascii=False)}) got exception: " + str(e))
|
||||
break
|
||||
return False
|
||||
|
||||
@@ -434,8 +445,7 @@ 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.")
|
||||
scripts = []
|
||||
params = {}
|
||||
for k, v in update_dict.items():
|
||||
@@ -465,11 +475,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)
|
||||
@@ -521,10 +528,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")
|
||||
|
||||
@@ -42,7 +42,7 @@ class InfinityConnection(InfinityConnectionBase):
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def convert_message_field_to_infinity(field_name: str, table_fields: list[str]=None):
|
||||
def convert_message_field_to_infinity(field_name: str, table_fields: list[str] = None):
|
||||
match field_name:
|
||||
case "message_type":
|
||||
return "message_type_kwd"
|
||||
@@ -68,7 +68,7 @@ class InfinityConnection(InfinityConnectionBase):
|
||||
return "content_embed"
|
||||
return field_name
|
||||
|
||||
def convert_select_fields(self, output_fields: list[str], table_fields: list[str]=None) -> list[str]:
|
||||
def convert_select_fields(self, output_fields: list[str], table_fields: list[str] = None) -> list[str]:
|
||||
return list({self.convert_message_field_to_infinity(f, table_fields) for f in output_fields})
|
||||
|
||||
@staticmethod
|
||||
@@ -277,7 +277,7 @@ class InfinityConnection(InfinityConnectionBase):
|
||||
self.logger.debug(f"INFINITY search final result: {str(res)}")
|
||||
return res, total_hits_count
|
||||
|
||||
def get_forgotten_messages(self, select_fields: list[str], index_name: str, memory_id: str, limit: int=512):
|
||||
def get_forgotten_messages(self, select_fields: list[str], index_name: str, memory_id: str, limit: int = 512):
|
||||
condition = {"memory_id": memory_id, "exists": "forget_at_flt"}
|
||||
order_by = OrderByExpr()
|
||||
order_by.asc("forget_at_flt")
|
||||
@@ -309,7 +309,7 @@ class InfinityConnection(InfinityConnectionBase):
|
||||
self.connPool.release_conn(inf_conn)
|
||||
return res
|
||||
|
||||
def get_missing_field_message(self, select_fields: list[str], index_name: str, memory_id: str, field_name: str, limit: int=512):
|
||||
def get_missing_field_message(self, select_fields: list[str], index_name: str, memory_id: str, field_name: str, limit: int = 512):
|
||||
condition = {"memory_id": memory_id, "must_not": {"exists": field_name}}
|
||||
order_by = OrderByExpr()
|
||||
order_by.asc("valid_at_flt")
|
||||
|
||||
@@ -27,9 +27,9 @@ def get_json_result_from_llm_response(response_str: str) -> dict:
|
||||
"""
|
||||
try:
|
||||
clean_str = response_str.strip()
|
||||
if clean_str.startswith('```json'):
|
||||
if clean_str.startswith("```json"):
|
||||
clean_str = clean_str[7:] # Remove the starting ```json
|
||||
if clean_str.endswith('```'):
|
||||
if clean_str.endswith("```"):
|
||||
clean_str = clean_str[:-3] # Remove the ending ```
|
||||
|
||||
return json.loads(clean_str.strip())
|
||||
|
||||
@@ -75,7 +75,7 @@ class SearchResult(BaseModel):
|
||||
@singleton
|
||||
class OBConnection(OBConnectionBase):
|
||||
def __init__(self):
|
||||
super().__init__(logger_name='ragflow.memory_ob_conn')
|
||||
super().__init__(logger_name="ragflow.memory_ob_conn")
|
||||
self._fulltext_search_columns = FTS_COLUMNS
|
||||
|
||||
"""
|
||||
@@ -101,10 +101,10 @@ class OBConnection(OBConnectionBase):
|
||||
def _get_vector_column_name_from_table(self, table_name: str) -> Optional[str]:
|
||||
"""Get the vector column name from the table (q_{size}_vec pattern)."""
|
||||
sql = f"""
|
||||
SELECT COLUMN_NAME
|
||||
FROM INFORMATION_SCHEMA.COLUMNS
|
||||
WHERE TABLE_SCHEMA = '{self.db_name}'
|
||||
AND TABLE_NAME = '{table_name}'
|
||||
SELECT COLUMN_NAME
|
||||
FROM INFORMATION_SCHEMA.COLUMNS
|
||||
WHERE TABLE_SCHEMA = '{self.db_name}'
|
||||
AND TABLE_NAME = '{table_name}'
|
||||
AND COLUMN_NAME REGEXP '^q_[0-9]+_vec$'
|
||||
LIMIT 1
|
||||
"""
|
||||
@@ -204,7 +204,7 @@ class OBConnection(OBConnectionBase):
|
||||
memory_ids: list[str],
|
||||
agg_fields: list[str] | None = None,
|
||||
rank_feature: dict | None = None,
|
||||
hide_forgotten: bool = True
|
||||
hide_forgotten: bool = True,
|
||||
):
|
||||
"""Search messages in memory storage."""
|
||||
if isinstance(index_names, str):
|
||||
@@ -273,9 +273,7 @@ class OBConnection(OBConnectionBase):
|
||||
fulltext_query = escape_string(fulltext_query.strip())
|
||||
fulltext_topn = m.topn
|
||||
|
||||
fulltext_search_expr, fulltext_search_weight = self._parse_fulltext_columns(
|
||||
fulltext_query, self._fulltext_search_columns
|
||||
)
|
||||
fulltext_search_expr, fulltext_search_weight = self._parse_fulltext_columns(fulltext_query, self._fulltext_search_columns)
|
||||
elif isinstance(m, MatchDenseExpr):
|
||||
vector_column_name = m.vector_column_name
|
||||
vector_data = m.embedding_data
|
||||
@@ -339,39 +337,27 @@ class OBConnection(OBConnectionBase):
|
||||
)
|
||||
self.logger.debug("OBConnection.search with fusion sql: %s", fusion_sql)
|
||||
rows, elapsed_time = self._execute_search_sql(fusion_sql)
|
||||
self.logger.info(
|
||||
f"OBConnection.search table {table_name}, search type: fusion, elapsed time: {elapsed_time:.3f}s, rows: {len(rows)}"
|
||||
)
|
||||
self.logger.info(f"OBConnection.search table {table_name}, search type: fusion, elapsed time: {elapsed_time:.3f}s, rows: {len(rows)}")
|
||||
|
||||
for row in rows:
|
||||
result.messages.append(self._row_to_entity(row, db_output_fields + ["_score"]))
|
||||
result.total += 1
|
||||
|
||||
elif search_type == "vector":
|
||||
vector_sql = self._build_vector_search_sql(
|
||||
table_name, fields_expr, vector_search_score_expr, filters_expr,
|
||||
vector_search_filter, vector_search_expr, limit, vector_topn, offset
|
||||
)
|
||||
vector_sql = self._build_vector_search_sql(table_name, fields_expr, vector_search_score_expr, filters_expr, vector_search_filter, vector_search_expr, limit, vector_topn, offset)
|
||||
self.logger.debug("OBConnection.search with vector sql: %s", vector_sql)
|
||||
rows, elapsed_time = self._execute_search_sql(vector_sql)
|
||||
self.logger.info(
|
||||
f"OBConnection.search table {table_name}, search type: vector, elapsed time: {elapsed_time:.3f}s, rows: {len(rows)}"
|
||||
)
|
||||
self.logger.info(f"OBConnection.search table {table_name}, search type: vector, elapsed time: {elapsed_time:.3f}s, rows: {len(rows)}")
|
||||
|
||||
for row in rows:
|
||||
result.messages.append(self._row_to_entity(row, db_output_fields + ["_score"]))
|
||||
result.total += 1
|
||||
|
||||
elif search_type == "fulltext":
|
||||
fulltext_sql = self._build_fulltext_search_sql(
|
||||
table_name, fields_expr, fulltext_search_score_expr, filters_expr,
|
||||
fulltext_search_filter, offset, limit, fulltext_topn
|
||||
)
|
||||
fulltext_sql = self._build_fulltext_search_sql(table_name, fields_expr, fulltext_search_score_expr, filters_expr, fulltext_search_filter, offset, limit, fulltext_topn)
|
||||
self.logger.debug("OBConnection.search with fulltext sql: %s", fulltext_sql)
|
||||
rows, elapsed_time = self._execute_search_sql(fulltext_sql)
|
||||
self.logger.info(
|
||||
f"OBConnection.search table {table_name}, search type: fulltext, elapsed time: {elapsed_time:.3f}s, rows: {len(rows)}"
|
||||
)
|
||||
self.logger.info(f"OBConnection.search table {table_name}, search type: fulltext, elapsed time: {elapsed_time:.3f}s, rows: {len(rows)}")
|
||||
|
||||
for row in rows:
|
||||
result.messages.append(self._row_to_entity(row, db_output_fields + ["_score"]))
|
||||
@@ -387,14 +373,10 @@ class OBConnection(OBConnectionBase):
|
||||
|
||||
order_by_expr = ("ORDER BY " + ", ".join(orders)) if orders else ""
|
||||
limit_expr = f"LIMIT {offset}, {limit}" if limit != 0 else ""
|
||||
filter_sql = self._build_filter_search_sql(
|
||||
table_name, fields_expr, filters_expr, order_by_expr, limit_expr
|
||||
)
|
||||
filter_sql = self._build_filter_search_sql(table_name, fields_expr, filters_expr, order_by_expr, limit_expr)
|
||||
self.logger.debug("OBConnection.search with filter sql: %s", filter_sql)
|
||||
rows, elapsed_time = self._execute_search_sql(filter_sql)
|
||||
self.logger.info(
|
||||
f"OBConnection.search table {table_name}, search type: filter, elapsed time: {elapsed_time:.3f}s, rows: {len(rows)}"
|
||||
)
|
||||
self.logger.info(f"OBConnection.search table {table_name}, search type: filter, elapsed time: {elapsed_time:.3f}s, rows: {len(rows)}")
|
||||
|
||||
for row in rows:
|
||||
result.messages.append(self._row_to_entity(row, db_output_fields))
|
||||
@@ -413,13 +395,7 @@ class OBConnection(OBConnectionBase):
|
||||
db_output_fields = [self.convert_field_name(f) for f in select_fields]
|
||||
fields_expr = ", ".join(db_output_fields)
|
||||
|
||||
sql = (
|
||||
f"SELECT {fields_expr}"
|
||||
f" FROM {index_name}"
|
||||
f" WHERE memory_id = {get_value_str(memory_id)} AND forget_at IS NOT NULL"
|
||||
f" ORDER BY forget_at ASC"
|
||||
f" LIMIT {limit}"
|
||||
)
|
||||
sql = f"SELECT {fields_expr} FROM {index_name} WHERE memory_id = {get_value_str(memory_id)} AND forget_at IS NOT NULL ORDER BY forget_at ASC LIMIT {limit}"
|
||||
self.logger.debug("OBConnection.get_forgotten_messages sql: %s", sql)
|
||||
|
||||
res = self.client.perform_raw_text_sql(sql)
|
||||
@@ -431,8 +407,7 @@ class OBConnection(OBConnectionBase):
|
||||
|
||||
return result
|
||||
|
||||
def get_missing_field_message(self, select_fields: list[str], index_name: str, memory_id: str, field_name: str,
|
||||
limit: int = 512):
|
||||
def get_missing_field_message(self, select_fields: list[str], index_name: str, memory_id: str, field_name: str, limit: int = 512):
|
||||
"""Get messages missing a specific field."""
|
||||
if not self._check_table_exists_cached(index_name):
|
||||
return None
|
||||
@@ -441,13 +416,7 @@ class OBConnection(OBConnectionBase):
|
||||
db_output_fields = [self.convert_field_name(f) for f in select_fields]
|
||||
fields_expr = ", ".join(db_output_fields)
|
||||
|
||||
sql = (
|
||||
f"SELECT {fields_expr}"
|
||||
f" FROM {index_name}"
|
||||
f" WHERE memory_id = {get_value_str(memory_id)} AND {db_field_name} IS NULL"
|
||||
f" ORDER BY valid_at ASC"
|
||||
f" LIMIT {limit}"
|
||||
)
|
||||
sql = f"SELECT {fields_expr} FROM {index_name} WHERE memory_id = {get_value_str(memory_id)} AND {db_field_name} IS NULL ORDER BY valid_at ASC LIMIT {limit}"
|
||||
self.logger.debug("OBConnection.get_missing_field_message sql: %s", sql)
|
||||
|
||||
res = self.client.perform_raw_text_sql(sql)
|
||||
@@ -531,11 +500,7 @@ class OBConnection(OBConnectionBase):
|
||||
if not set_values:
|
||||
return True
|
||||
|
||||
update_sql = (
|
||||
f"UPDATE {index_name}"
|
||||
f" SET {', '.join(set_values)}"
|
||||
f" WHERE {' AND '.join(filters)}"
|
||||
)
|
||||
update_sql = f"UPDATE {index_name} SET {', '.join(set_values)} WHERE {' AND '.join(filters)}"
|
||||
self.logger.debug("OBConnection.update sql: %s", update_sql)
|
||||
|
||||
try:
|
||||
@@ -557,14 +522,14 @@ class OBConnection(OBConnectionBase):
|
||||
def get_total(self, res) -> int:
|
||||
if isinstance(res, tuple):
|
||||
return res[1]
|
||||
if hasattr(res, 'total'):
|
||||
if hasattr(res, "total"):
|
||||
return res.total
|
||||
return 0
|
||||
|
||||
def get_doc_ids(self, res) -> list[str]:
|
||||
if isinstance(res, tuple):
|
||||
res = res[0]
|
||||
if hasattr(res, 'messages'):
|
||||
if hasattr(res, "messages"):
|
||||
return [row.get("id") for row in res.messages if row.get("id")]
|
||||
return []
|
||||
|
||||
@@ -577,7 +542,7 @@ class OBConnection(OBConnectionBase):
|
||||
if not fields:
|
||||
return {}
|
||||
|
||||
messages = res.messages if hasattr(res, 'messages') else []
|
||||
messages = res.messages if hasattr(res, "messages") else []
|
||||
|
||||
for doc in messages:
|
||||
message = self.get_message_from_ob_doc(doc)
|
||||
@@ -588,10 +553,7 @@ class OBConnection(OBConnectionBase):
|
||||
if isinstance(v, list):
|
||||
m[n] = v
|
||||
continue
|
||||
if n in ["message_id", "source_id", "valid_at", "invalid_at", "forget_at", "status"] and isinstance(v,
|
||||
(int,
|
||||
float,
|
||||
bool)):
|
||||
if n in ["message_id", "source_id", "valid_at", "invalid_at", "forget_at", "status"] and isinstance(v, (int, float, bool)):
|
||||
m[n] = v
|
||||
continue
|
||||
if not isinstance(v, str):
|
||||
@@ -610,9 +572,7 @@ class OBConnection(OBConnectionBase):
|
||||
if isinstance(res, tuple):
|
||||
res = res[0]
|
||||
messages = getattr(res, "messages", None)
|
||||
return get_highlight_from_messages(
|
||||
messages, keywords, field_name, is_english_fn=lambda s: is_english([s])
|
||||
)
|
||||
return get_highlight_from_messages(messages, keywords, field_name, is_english_fn=lambda s: is_english([s]))
|
||||
|
||||
def get_aggregation(self, res, field_name: str):
|
||||
"""Get aggregation for search results."""
|
||||
|
||||
@@ -18,8 +18,8 @@ from typing import Optional, List
|
||||
from common.constants import MemoryType
|
||||
from common.time_utils import current_timestamp
|
||||
|
||||
class PromptAssembler:
|
||||
|
||||
class PromptAssembler:
|
||||
SYSTEM_BASE_TEMPLATE = """**Memory Extraction Specialist**
|
||||
You are an expert at analyzing conversations to extract structured memory.
|
||||
|
||||
@@ -47,7 +47,6 @@ You are an expert at analyzing conversations to extract structured memory.
|
||||
- invalid_at: When it becomes false (e.g., repeal, disproven) or empty if still true
|
||||
- Default: valid_at = conversation time, invalid_at = "" for timeless facts
|
||||
""",
|
||||
|
||||
MemoryType.EPISODIC.name.lower(): """
|
||||
**EXTRACT EPISODIC KNOWLEDGE:**
|
||||
- Specific experiences, events, personal stories
|
||||
@@ -59,7 +58,6 @@ You are an expert at analyzing conversations to extract structured memory.
|
||||
- invalid_at: Event end time or empty if instantaneous
|
||||
- Extract explicit times: "at 3 PM", "last Monday", "from X to Y"
|
||||
""",
|
||||
|
||||
MemoryType.PROCEDURAL.name.lower(): """
|
||||
**EXTRACT PROCEDURAL KNOWLEDGE:**
|
||||
- Processes, methods, step-by-step instructions
|
||||
@@ -71,7 +69,7 @@ You are an expert at analyzing conversations to extract structured memory.
|
||||
- invalid_at: When it expires/becomes obsolete or empty if current
|
||||
- For version-specific: use release dates
|
||||
- For best practices: invalid_at = ""
|
||||
"""
|
||||
""",
|
||||
}
|
||||
|
||||
OUTPUT_TEMPLATES = {
|
||||
@@ -84,7 +82,6 @@ You are an expert at analyzing conversations to extract structured memory.
|
||||
}
|
||||
]
|
||||
""",
|
||||
|
||||
MemoryType.EPISODIC.name.lower(): """
|
||||
"episodic": [
|
||||
{
|
||||
@@ -94,7 +91,6 @@ You are an expert at analyzing conversations to extract structured memory.
|
||||
}
|
||||
]
|
||||
""",
|
||||
|
||||
MemoryType.PROCEDURAL.name.lower(): """
|
||||
"procedural": [
|
||||
{
|
||||
@@ -103,7 +99,7 @@ You are an expert at analyzing conversations to extract structured memory.
|
||||
"invalid_at": "procedure expiration timestamp or empty"
|
||||
}
|
||||
]
|
||||
"""
|
||||
""",
|
||||
}
|
||||
|
||||
BASE_USER_PROMPT = """
|
||||
@@ -111,7 +107,7 @@ You are an expert at analyzing conversations to extract structured memory.
|
||||
{conversation}
|
||||
|
||||
**CONVERSATION TIME:** {conversation_time}
|
||||
**CURRENT TIME:** {current_time}
|
||||
**CURRENT TIME:** {current_time}
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
@@ -123,9 +119,7 @@ You are an expert at analyzing conversations to extract structured memory.
|
||||
output_format = cls._generate_output_format(types_to_extract)
|
||||
|
||||
full_prompt = cls.SYSTEM_BASE_TEMPLATE.format(
|
||||
type_specific_instructions=type_instructions,
|
||||
timestamp_format=config.get("timestamp_format", "ISO 8601"),
|
||||
max_items=config.get("max_items_per_type", 5)
|
||||
type_specific_instructions=type_instructions, timestamp_format=config.get("timestamp_format", "ISO 8601"), max_items=config.get("max_items_per_type", 5)
|
||||
)
|
||||
|
||||
full_prompt += f"\n**REQUIRED OUTPUT FORMAT (JSON):**\n```json\n{{\n{output_format}\n}}\n```\n"
|
||||
@@ -140,7 +134,7 @@ You are an expert at analyzing conversations to extract structured memory.
|
||||
def _get_types_to_extract(requested_types: List[str]) -> List[str]:
|
||||
types = set()
|
||||
for rt in requested_types:
|
||||
if rt in [e.name.lower() for e in MemoryType] and rt != MemoryType.RAW.name.lower():
|
||||
if rt in [e.name.lower() for e in MemoryType] and rt != MemoryType.RAW.name.lower():
|
||||
types.add(rt)
|
||||
return list(types)
|
||||
|
||||
@@ -184,12 +178,7 @@ You are an expert at analyzing conversations to extract structured memory.
|
||||
return "\n".join(examples)
|
||||
|
||||
@classmethod
|
||||
def assemble_user_prompt(
|
||||
cls,
|
||||
conversation: str,
|
||||
conversation_time: Optional[str] = None,
|
||||
current_time: Optional[str] = None
|
||||
) -> str:
|
||||
def assemble_user_prompt(cls, conversation: str, conversation_time: Optional[str] = None, current_time: Optional[str] = None) -> str:
|
||||
return cls.BASE_USER_PROMPT.format(
|
||||
conversation=conversation,
|
||||
conversation_time=conversation_time or "Not specified",
|
||||
|
||||
Reference in New Issue
Block a user