Feat: message manage (#12083)

### What problem does this PR solve?

Message CRUD.

Issue #4213 

### Type of change

- [x] New Feature (non-breaking change which adds functionality)

.github/workflows/tests.yml

@@ -219,7 +219,7 @@ jobs:
             sleep 5
           done
           source .venv/bin/activate && pytest -s --tb=short sdk/python/test/test_frontend_api/get_email.py sdk/python/test/test_frontend_api/test_dataset.py 2>&1 | tee es_api_test.log
-          
+
       - name: Run http api tests against Elasticsearch
         run: |
           export http_proxy=""; export https_proxy=""; export no_proxy=""; export HTTP_PROXY=""; export HTTPS_PROXY=""; export NO_PROXY=""

Dockerfile

@@ -192,6 +192,7 @@ COPY pyproject.toml uv.lock ./
 COPY mcp mcp
 COPY plugin plugin
 COPY common common
+COPY memory memory
 
 COPY docker/service_conf.yaml.template ./conf/service_conf.yaml.template
 COPY docker/entrypoint.sh ./

admin/server/routes.py

@@ -29,6 +29,11 @@ from common.versions import get_ragflow_version
 admin_bp = Blueprint('admin', __name__, url_prefix='/api/v1/admin')
 
 
+@admin_bp.route('/ping', methods=['GET'])
+def ping():
+    return success_response('PONG')
+
+
 @admin_bp.route('/login', methods=['POST'])
 def login():
     if not request.json:

agent/tools/retrieval.py

@@ -25,10 +25,12 @@ from api.db.services.document_service import DocumentService
 from common.metadata_utils import apply_meta_data_filter
 from api.db.services.knowledgebase_service import KnowledgebaseService
 from api.db.services.llm_service import LLMBundle
+from api.db.services.memory_service import MemoryService
+from api.db.joint_services.memory_message_service import query_message
 from common import settings
 from common.connection_utils import timeout
 from rag.app.tag import label_question
-from rag.prompts.generator import cross_languages, kb_prompt
+from rag.prompts.generator import cross_languages, kb_prompt, memory_prompt
 
 
 class RetrievalParam(ToolParamBase):
@@ -57,6 +59,7 @@ class RetrievalParam(ToolParamBase):
         self.top_n = 8
         self.top_k = 1024
         self.kb_ids = []
+        self.memory_ids = []
         self.kb_vars = []
         self.rerank_id = ""
         self.empty_response = ""
@@ -81,15 +84,7 @@ class RetrievalParam(ToolParamBase):
 class Retrieval(ToolBase, ABC):
     component_name = "Retrieval"
 
-    @timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 12)))
-    async def _invoke_async(self, **kwargs):
-        if self.check_if_canceled("Retrieval processing"):
-            return
-
-        if not kwargs.get("query"):
-            self.set_output("formalized_content", self._param.empty_response)
-            return
-
+    async def _retrieve_kb(self, query_text: str):
         kb_ids: list[str] = []
         for id in self._param.kb_ids:
             if id.find("@") < 0:
@@ -124,12 +119,12 @@ class Retrieval(ToolBase, ABC):
         if self._param.rerank_id:
             rerank_mdl = LLMBundle(kbs[0].tenant_id, LLMType.RERANK, self._param.rerank_id)
 
-        vars = self.get_input_elements_from_text(kwargs["query"])
-        vars = {k:o["value"] for k,o in vars.items()}
-        query = self.string_format(kwargs["query"], vars)
+        vars = self.get_input_elements_from_text(query_text)
+        vars = {k: o["value"] for k, o in vars.items()}
+        query = self.string_format(query_text, vars)
 
-        doc_ids=[]
-        if self._param.meta_data_filter!={}:
+        doc_ids = []
+        if self._param.meta_data_filter != {}:
             metas = DocumentService.get_meta_by_kbs(kb_ids)
 
             def _resolve_manual_filter(flt: dict) -> dict:
@@ -198,18 +193,20 @@ class Retrieval(ToolBase, ABC):
 
             if self._param.toc_enhance:
                 chat_mdl = LLMBundle(self._canvas._tenant_id, LLMType.CHAT)
-                cks = settings.retriever.retrieval_by_toc(query, kbinfos["chunks"], [kb.tenant_id for kb in kbs], chat_mdl, self._param.top_n)
+                cks = settings.retriever.retrieval_by_toc(query, kbinfos["chunks"], [kb.tenant_id for kb in kbs],
+                                                          chat_mdl, self._param.top_n)
                 if self.check_if_canceled("Retrieval processing"):
                     return
                 if cks:
                     kbinfos["chunks"] = cks
-            kbinfos["chunks"] = settings.retriever.retrieval_by_children(kbinfos["chunks"], [kb.tenant_id for kb in kbs])
+            kbinfos["chunks"] = settings.retriever.retrieval_by_children(kbinfos["chunks"],
+                                                                         [kb.tenant_id for kb in kbs])
             if self._param.use_kg:
                 ck = settings.kg_retriever.retrieval(query,
-                                                       [kb.tenant_id for kb in kbs],
-                                                       kb_ids,
-                                                       embd_mdl,
-                                                       LLMBundle(self._canvas.get_tenant_id(), LLMType.CHAT))
+                                                     [kb.tenant_id for kb in kbs],
+                                                     kb_ids,
+                                                     embd_mdl,
+                                                     LLMBundle(self._canvas.get_tenant_id(), LLMType.CHAT))
                 if self.check_if_canceled("Retrieval processing"):
                     return
                 if ck["content_with_weight"]:
@@ -218,7 +215,8 @@ class Retrieval(ToolBase, ABC):
             kbinfos = {"chunks": [], "doc_aggs": []}
 
         if self._param.use_kg and kbs:
-            ck = settings.kg_retriever.retrieval(query, [kb.tenant_id for kb in kbs], filtered_kb_ids, embd_mdl, LLMBundle(kbs[0].tenant_id, LLMType.CHAT))
+            ck = settings.kg_retriever.retrieval(query, [kb.tenant_id for kb in kbs], filtered_kb_ids, embd_mdl,
+                                                 LLMBundle(kbs[0].tenant_id, LLMType.CHAT))
             if self.check_if_canceled("Retrieval processing"):
                 return
             if ck["content_with_weight"]:
@@ -248,6 +246,54 @@ class Retrieval(ToolBase, ABC):
 
         return form_cnt
 
+    async def _retrieve_memory(self, query_text: str):
+        memory_ids: list[str] = [memory_id for memory_id in self._param.memory_ids]
+        memory_list = MemoryService.get_by_ids(memory_ids)
+        if not memory_list:
+            raise Exception("No memory is selected.")
+
+        embd_names = list({memory.embd_id for memory in memory_list})
+        assert len(embd_names) == 1, "Memory use different embedding models."
+
+        vars = self.get_input_elements_from_text(query_text)
+        vars = {k: o["value"] for k, o in vars.items()}
+        query = self.string_format(query_text, vars)
+        # query message
+        message_list = query_message({"memory_id": memory_ids}, {
+            "query": query,
+            "similarity_threshold": self._param.similarity_threshold,
+            "keywords_similarity_weight": self._param.keywords_similarity_weight,
+            "top_n": self._param.top_n
+        })
+        print(f"found {len(message_list)} messages.")
+
+        if not message_list:
+            self.set_output("formalized_content", self._param.empty_response)
+            return
+        formated_content = "\n".join(memory_prompt(message_list, 200000))
+
+        # set formalized_content output
+        self.set_output("formalized_content", formated_content)
+        print(f"formated_content {formated_content}")
+        return formated_content
+
+    @timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 12)))
+    async def _invoke_async(self, **kwargs):
+        if self.check_if_canceled("Retrieval processing"):
+            return
+        print(f"debug retrieval, query is {kwargs.get('query')}.", flush=True)
+        if not kwargs.get("query"):
+            self.set_output("formalized_content", self._param.empty_response)
+            return
+
+        if self._param.kb_ids:
+            return await self._retrieve_kb(kwargs["query"])
+        elif self._param.memory_ids:
+            return await self._retrieve_memory(kwargs["query"])
+        else:
+            self.set_output("formalized_content", self._param.empty_response)
+            return
+
     @timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 12)))
     def _invoke(self, **kwargs):
         return asyncio.run(self._invoke_async(**kwargs))

api/apps/canvas_app.py

@@ -192,7 +192,7 @@ async def rerun():
     if 0 < doc["progress"] < 1:
         return get_data_error_result(message=f"`{doc['name']}` is processing...")
 
-    if settings.docStoreConn.indexExist(search.index_name(current_user.id), doc["kb_id"]):
+    if settings.docStoreConn.index_exist(search.index_name(current_user.id), doc["kb_id"]):
         settings.docStoreConn.delete({"doc_id": doc["id"]}, search.index_name(current_user.id), doc["kb_id"])
     doc["progress_msg"] = ""
     doc["chunk_num"] = 0

api/apps/document_app.py

@@ -564,7 +564,7 @@ async def run():
                 DocumentService.update_by_id(id, info)
                 if req.get("delete", False):
                     TaskService.filter_delete([Task.doc_id == id])
-                    if settings.docStoreConn.indexExist(search.index_name(tenant_id), doc.kb_id):
+                    if settings.docStoreConn.index_exist(search.index_name(tenant_id), doc.kb_id):
                         settings.docStoreConn.delete({"doc_id": id}, search.index_name(tenant_id), doc.kb_id)
 
                 if str(req["run"]) == TaskStatus.RUNNING.value:
@@ -615,7 +615,7 @@ async def rename():
                 "title_tks": title_tks,
                 "title_sm_tks": rag_tokenizer.fine_grained_tokenize(title_tks),
             }
-            if settings.docStoreConn.indexExist(search.index_name(tenant_id), doc.kb_id):
+            if settings.docStoreConn.index_exist(search.index_name(tenant_id), doc.kb_id):
                 settings.docStoreConn.update(
                     {"doc_id": req["doc_id"]},
                     es_body,
@@ -696,7 +696,7 @@ async def change_parser():
             tenant_id = DocumentService.get_tenant_id(req["doc_id"])
             if not tenant_id:
                 return get_data_error_result(message="Tenant not found!")
-            if settings.docStoreConn.indexExist(search.index_name(tenant_id), doc.kb_id):
+            if settings.docStoreConn.index_exist(search.index_name(tenant_id), doc.kb_id):
                 settings.docStoreConn.delete({"doc_id": doc.id}, search.index_name(tenant_id), doc.kb_id)
         return None
 

api/apps/kb_app.py

@@ -39,9 +39,9 @@ from api.utils.api_utils import get_json_result
 from rag.nlp import search
 from api.constants import DATASET_NAME_LIMIT
 from rag.utils.redis_conn import REDIS_CONN
-from rag.utils.doc_store_conn import OrderByExpr
 from common.constants import RetCode, PipelineTaskType, StatusEnum, VALID_TASK_STATUS, FileSource, LLMType, PAGERANK_FLD
 from common import settings
+from common.doc_store.doc_store_base import OrderByExpr
 from api.apps import login_required, current_user
 
 
@@ -285,7 +285,7 @@ async def rm():
                     message="Database error (Knowledgebase removal)!")
             for kb in kbs:
                 settings.docStoreConn.delete({"kb_id": kb.id}, search.index_name(kb.tenant_id), kb.id)
-                settings.docStoreConn.deleteIdx(search.index_name(kb.tenant_id), kb.id)
+                settings.docStoreConn.delete_idx(search.index_name(kb.tenant_id), kb.id)
                 if hasattr(settings.STORAGE_IMPL, 'remove_bucket'):
                     settings.STORAGE_IMPL.remove_bucket(kb.id)
             return get_json_result(data=True)
@@ -386,7 +386,7 @@ def knowledge_graph(kb_id):
     }
 
     obj = {"graph": {}, "mind_map": {}}
-    if not settings.docStoreConn.indexExist(search.index_name(kb.tenant_id), kb_id):
+    if not settings.docStoreConn.index_exist(search.index_name(kb.tenant_id), kb_id):
         return get_json_result(data=obj)
     sres = settings.retriever.search(req, search.index_name(kb.tenant_id), [kb_id])
     if not len(sres.ids):
@@ -858,11 +858,11 @@ async def check_embedding():
         index_nm = search.index_name(tenant_id)
 
         res0 = docStoreConn.search(
-            selectFields=[], highlightFields=[],
+            select_fields=[], highlight_fields=[],
             condition={"kb_id": kb_id, "available_int": 1},
-            matchExprs=[], orderBy=OrderByExpr(),
+            match_expressions=[], order_by=OrderByExpr(),
             offset=0, limit=1,
-            indexNames=index_nm, knowledgebaseIds=[kb_id]
+            index_names=index_nm, knowledgebase_ids=[kb_id]
         )
         total = docStoreConn.get_total(res0)
         if total <= 0:
@@ -874,14 +874,14 @@ async def check_embedding():
 
         for off in offsets:
             res1 = docStoreConn.search(
-                selectFields=list(base_fields),
-                highlightFields=[],
+                select_fields=list(base_fields),
+                highlight_fields=[],
                 condition={"kb_id": kb_id, "available_int": 1},
-                matchExprs=[], orderBy=OrderByExpr(),
+                match_expressions=[], order_by=OrderByExpr(),
                 offset=off, limit=1,
-                indexNames=index_nm, knowledgebaseIds=[kb_id]
+                index_names=index_nm, knowledgebase_ids=[kb_id]
             )
-            ids = docStoreConn.get_chunk_ids(res1)
+            ids = docStoreConn.get_doc_ids(res1)
             if not ids:
                 continue
 

api/apps/memories_app.py

@@ -20,10 +20,12 @@ from api.apps import login_required, current_user
 from api.db import TenantPermission
 from api.db.services.memory_service import MemoryService
 from api.db.services.user_service import UserTenantService
+from api.db.services.canvas_service import UserCanvasService
 from api.utils.api_utils import validate_request, get_request_json, get_error_argument_result, get_json_result, \
     not_allowed_parameters
 from api.utils.memory_utils import format_ret_data_from_memory, get_memory_type_human
 from api.constants import MEMORY_NAME_LIMIT, MEMORY_SIZE_LIMIT
+from memory.services.messages import MessageService
 from common.constants import MemoryType, RetCode, ForgettingPolicy
 
 
@@ -57,7 +59,6 @@ async def create_memory():
 
         if res:
             return get_json_result(message=True, data=format_ret_data_from_memory(memory))
-
         else:
             return get_json_result(message=memory, code=RetCode.SERVER_ERROR)
 
@@ -124,7 +125,7 @@ async def update_memory(memory_id):
         return get_json_result(message=True, data=memory_dict)
 
     try:
-        MemoryService.update_memory(memory_id, to_update)
+        MemoryService.update_memory(current_memory.tenant_id, memory_id, to_update)
         updated_memory = MemoryService.get_by_memory_id(memory_id)
         return get_json_result(message=True, data=format_ret_data_from_memory(updated_memory))
 
@@ -133,7 +134,7 @@ async def update_memory(memory_id):
         return get_json_result(message=str(e), code=RetCode.SERVER_ERROR)
 
 
-@manager.route("/<memory_id>", methods=["DELETE"]) # noqa: F821
+@manager.route("/<memory_id>", methods=["DELETE"])  # noqa: F821
 @login_required
 async def delete_memory(memory_id):
     memory = MemoryService.get_by_memory_id(memory_id)
@@ -141,13 +142,14 @@ async def delete_memory(memory_id):
         return get_json_result(message=True, code=RetCode.NOT_FOUND)
     try:
         MemoryService.delete_memory(memory_id)
+        MessageService.delete_message({"memory_id": memory_id}, memory.tenant_id, memory_id)
         return get_json_result(message=True)
     except Exception as e:
         logging.error(e)
         return get_json_result(message=str(e), code=RetCode.SERVER_ERROR)
 
 
-@manager.route("", methods=["GET"]) # noqa: F821
+@manager.route("", methods=["GET"])  # noqa: F821
 @login_required
 async def list_memory():
     args = request.args
@@ -183,3 +185,26 @@ async def get_memory_config(memory_id):
     if not memory:
         return get_json_result(code=RetCode.NOT_FOUND, message=f"Memory '{memory_id}' not found.")
     return get_json_result(message=True, data=format_ret_data_from_memory(memory))
+
+
+@manager.route("/<memory_id>", methods=["GET"])  # noqa: F821
+@login_required
+async def get_memory_detail(memory_id):
+    args = request.args
+    agent_ids = args.getlist("agent_id")
+    keywords = args.get("keywords", "")
+    keywords = keywords.strip()
+    page = int(args.get("page", 1))
+    page_size = int(args.get("page_size", 50))
+    memory = MemoryService.get_by_memory_id(memory_id)
+    if not memory:
+        return get_json_result(code=RetCode.NOT_FOUND, message=f"Memory '{memory_id}' not found.")
+    messages = MessageService.list_message(
+        memory.tenant_id, memory_id, agent_ids, keywords, page, page_size)
+    agent_name_mapping = {}
+    if messages["message_list"]:
+        agent_list = UserCanvasService.get_basic_info_by_canvas_ids([message["agent_id"] for message in messages["message_list"]])
+        agent_name_mapping = {agent["id"]: agent["title"] for agent in agent_list}
+    for message in messages["message_list"]:
+        message["agent_name"] = agent_name_mapping.get(message["agent_id"], "Unknown")
+    return get_json_result(data={"messages": messages, "storage_type": memory.storage_type}, message=True)

api/apps/messages_app.py

@@ -0,0 +1,169 @@
+#
+#  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.
+#
+from quart import request
+from api.apps import login_required
+from api.db.services.memory_service import MemoryService
+from common.time_utils import current_timestamp, timestamp_to_date
+
+from memory.services.messages import MessageService
+from api.db.joint_services import memory_message_service
+from api.db.joint_services.memory_message_service import query_message
+from api.utils.api_utils import validate_request, get_request_json, get_error_argument_result, get_json_result
+from common.constants import RetCode
+
+
+@manager.route("", methods=["POST"]) # noqa: F821
+@login_required
+@validate_request("memory_id", "agent_id", "session_id", "user_input", "agent_response")
+async def add_message():
+
+    req = await get_request_json()
+    memory_ids = req["memory_id"]
+    agent_id = req["agent_id"]
+    session_id = req["session_id"]
+    user_id = req["user_id"] if req.get("user_id") else ""
+    user_input = req["user_input"]
+    agent_response = req["agent_response"]
+
+    res = []
+    for memory_id in memory_ids:
+        success, msg = await memory_message_service.save_to_memory(
+            memory_id,
+            {
+                "user_id": user_id,
+                "agent_id": agent_id,
+                "session_id": session_id,
+                "user_input": user_input,
+                "agent_response": agent_response
+            }
+        )
+        res.append({
+            "memory_id": memory_id,
+            "success": success,
+            "message": msg
+        })
+
+    if all([r["success"] for r in res]):
+        return get_json_result(message="Successfully added to memories.")
+
+    return get_json_result(code=RetCode.SERVER_ERROR, message="Some messages failed to add.", data=res)
+
+
+@manager.route("/<memory_id>:<message_id>", methods=["DELETE"]) # noqa: F821
+@login_required
+async def forget_message(memory_id: str, message_id: int):
+
+    memory = MemoryService.get_by_memory_id(memory_id)
+    if not memory:
+        return get_json_result(code=RetCode.NOT_FOUND, message=f"Memory '{memory_id}' not found.")
+
+    forget_time = timestamp_to_date(current_timestamp())
+    update_succeed = MessageService.update_message(
+        {"memory_id": memory_id, "message_id": int(message_id)},
+        {"forget_at": forget_time},
+        memory.tenant_id, memory_id)
+    if update_succeed:
+        return get_json_result(message=update_succeed)
+    else:
+        return get_json_result(code=RetCode.SERVER_ERROR, message=f"Failed to forget message '{message_id}' in memory '{memory_id}'.")
+
+
+@manager.route("/<memory_id>:<message_id>", methods=["PUT"]) # noqa: F821
+@login_required
+@validate_request("status")
+async def update_message(memory_id: str, message_id: int):
+    req = await get_request_json()
+    status = req["status"]
+    if not isinstance(status, bool):
+        return get_error_argument_result("Status must be a boolean.")
+
+    memory = MemoryService.get_by_memory_id(memory_id)
+    if not memory:
+        return get_json_result(code=RetCode.NOT_FOUND, message=f"Memory '{memory_id}' not found.")
+
+    update_succeed = MessageService.update_message({"memory_id": memory_id, "message_id": int(message_id)}, {"status": status}, memory.tenant_id, memory_id)
+    if update_succeed:
+        return get_json_result(message=update_succeed)
+    else:
+        return get_json_result(code=RetCode.SERVER_ERROR, message=f"Failed to set status for message '{message_id}' in memory '{memory_id}'.")
+
+
+@manager.route("/search", methods=["GET"]) # noqa: F821
+@login_required
+async def search_message():
+    args = request.args
+    print(args, flush=True)
+    empty_fields = [f for f in ["memory_id", "query"] if not args.get(f)]
+    if empty_fields:
+        return get_error_argument_result(f"{', '.join(empty_fields)} can't be empty.")
+
+    memory_ids = args.getlist("memory_id")
+    query = args.get("query")
+    similarity_threshold = float(args.get("similarity_threshold", 0.2))
+    keywords_similarity_weight = float(args.get("keywords_similarity_weight", 0.7))
+    top_n = int(args.get("top_n", 5))
+    agent_id = args.get("agent_id", "")
+    session_id = args.get("session_id", "")
+
+    filter_dict = {
+        "memory_id": memory_ids,
+        "agent_id": agent_id,
+        "session_id": session_id
+    }
+    params = {
+        "query": query,
+        "similarity_threshold": similarity_threshold,
+        "keywords_similarity_weight": keywords_similarity_weight,
+        "top_n": top_n
+    }
+    res = query_message(filter_dict, params)
+    return get_json_result(message=True, data=res)
+
+
+@manager.route("", methods=["GET"]) # noqa: F821
+@login_required
+async def get_messages():
+    args = request.args
+    memory_ids = args.getlist("memory_id")
+    agent_id = args.get("agent_id", "")
+    session_id = args.get("session_id", "")
+    limit = int(args.get("limit", 10))
+    if not memory_ids:
+        return get_error_argument_result("memory_ids is required.")
+    memory_list = MemoryService.get_by_ids(memory_ids)
+    uids = [memory.tenant_id for memory in memory_list]
+    res = MessageService.get_recent_messages(
+        uids,
+        memory_ids,
+        agent_id,
+        session_id,
+        limit
+    )
+    return get_json_result(message=True, data=res)
+
+
+@manager.route("/<memory_id>:<message_id>/content", methods=["GET"]) # noqa: F821
+@login_required
+async def get_message_content(memory_id:str, message_id: int):
+    memory = MemoryService.get_by_memory_id(memory_id)
+    if not memory:
+        return get_json_result(code=RetCode.NOT_FOUND, message=f"Memory '{memory_id}' not found.")
+
+    res = MessageService.get_by_message_id(memory_id, message_id, memory.tenant_id)
+    if res:
+        return get_json_result(message=True, data=res)
+    else:
+        return get_json_result(code=RetCode.NOT_FOUND, message=f"Message '{message_id}' in memory '{memory_id}' not found.")

api/apps/sdk/dataset.py

@@ -495,7 +495,7 @@ def knowledge_graph(tenant_id, dataset_id):
     }
 
     obj = {"graph": {}, "mind_map": {}}
-    if not settings.docStoreConn.indexExist(search.index_name(kb.tenant_id), dataset_id):
+    if not settings.docStoreConn.index_exist(search.index_name(kb.tenant_id), dataset_id):
         return get_result(data=obj)
     sres = settings.retriever.search(req, search.index_name(kb.tenant_id), [dataset_id])
     if not len(sres.ids):

api/apps/sdk/doc.py

@@ -1080,7 +1080,7 @@ def list_chunks(tenant_id, dataset_id, document_id):
         res["chunks"].append(final_chunk)
         _ = Chunk(**final_chunk)
 
-    elif settings.docStoreConn.indexExist(search.index_name(tenant_id), dataset_id):
+    elif settings.docStoreConn.index_exist(search.index_name(tenant_id), dataset_id):
         sres = settings.retriever.search(query, search.index_name(tenant_id), [dataset_id], emb_mdl=None, highlight=True)
         res["total"] = sres.total
         for id in sres.ids:

api/db/init_data.py

@@ -30,6 +30,7 @@ from api.db.services.knowledgebase_service import KnowledgebaseService
 from api.db.services.tenant_llm_service import LLMFactoriesService, TenantLLMService
 from api.db.services.llm_service import LLMService, LLMBundle, get_init_tenant_llm
 from api.db.services.user_service import TenantService, UserTenantService
+from api.db.joint_services.memory_message_service import init_message_id_sequence, init_memory_size_cache
 from common.constants import LLMType
 from common.file_utils import get_project_base_directory
 from common import settings
@@ -169,6 +170,8 @@ def init_web_data():
     #    init_superuser()
 
     add_graph_templates()
+    init_message_id_sequence()
+    init_memory_size_cache()
     logging.info("init web data success:{}".format(time.time() - start_time))
 
 

api/db/joint_services/memory_message_service.py

@@ -0,0 +1,233 @@
+#
+#  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
+from typing import List
+
+from common.time_utils import current_timestamp, timestamp_to_date, format_iso_8601_to_ymd_hms
+from common.constants import MemoryType, LLMType
+from common.doc_store.doc_store_base import FusionExpr
+from api.db.services.memory_service import MemoryService
+from api.db.services.tenant_llm_service import TenantLLMService
+from api.db.services.llm_service import LLMBundle
+from api.utils.memory_utils import get_memory_type_human
+from memory.services.messages import MessageService
+from memory.services.query import MsgTextQuery, get_vector
+from memory.utils.prompt_util import PromptAssembler
+from memory.utils.msg_util import get_json_result_from_llm_response
+from rag.utils.redis_conn import REDIS_CONN
+
+
+async def save_to_memory(memory_id: str, message_dict: dict):
+    """
+    :param memory_id:
+    :param message_dict: {
+        "user_id": str,
+        "agent_id": str,
+        "session_id": str,
+        "user_input": str,
+        "agent_response": str
+    }
+    """
+    memory = MemoryService.get_by_memory_id(memory_id)
+    if not memory:
+        return False, f"Memory '{memory_id}' not found."
+
+    tenant_id = memory.tenant_id
+    extracted_content = await extract_by_llm(
+        tenant_id,
+        memory.llm_id,
+        {"temperature": memory.temperature},
+        get_memory_type_human(memory.memory_type),
+        message_dict.get("user_input", ""),
+        message_dict.get("agent_response", "")
+    ) if memory.memory_type != MemoryType.RAW.value else []  # if only RAW, no need to extract
+    raw_message_id = REDIS_CONN.generate_auto_increment_id(namespace="memory")
+    message_list = [{
+        "message_id": raw_message_id,
+        "message_type": MemoryType.RAW.name.lower(),
+        "source_id": 0,
+        "memory_id": memory_id,
+        "user_id": "",
+        "agent_id": message_dict["agent_id"],
+        "session_id": message_dict["session_id"],
+        "content": f"User Input: {message_dict.get('user_input')}\nAgent Response: {message_dict.get('agent_response')}",
+        "valid_at": timestamp_to_date(current_timestamp()),
+        "invalid_at": None,
+        "forget_at": None,
+        "status": True
+    }, *[{
+        "message_id": REDIS_CONN.generate_auto_increment_id(namespace="memory"),
+        "message_type": content["message_type"],
+        "source_id": raw_message_id,
+        "memory_id": memory_id,
+        "user_id": "",
+        "agent_id": message_dict["agent_id"],
+        "session_id": message_dict["session_id"],
+        "content": content["content"],
+        "valid_at": content["valid_at"],
+        "invalid_at": content["invalid_at"] if content["invalid_at"] else None,
+        "forget_at": None,
+        "status": True
+    } for content in extracted_content]]
+    embedding_model = LLMBundle(tenant_id, llm_type=LLMType.EMBEDDING, llm_name=memory.embd_id)
+    vector_list, _ = embedding_model.encode([msg["content"] for msg in message_list])
+    for idx, msg in enumerate(message_list):
+        msg["content_embed"] = vector_list[idx]
+    vector_dimension = len(vector_list[0])
+    if not MessageService.has_index(tenant_id, memory_id):
+        created = MessageService.create_index(tenant_id, memory_id, vector_size=vector_dimension)
+        if not created:
+            return False, "Failed to create message index."
+
+    new_msg_size = sum([MessageService.calculate_message_size(m) for m in message_list])
+    current_memory_size = get_memory_size_cache(memory_id, tenant_id)
+    if new_msg_size + current_memory_size > memory.memory_size:
+        size_to_delete = current_memory_size + new_msg_size - memory.memory_size
+        if memory.forgetting_policy == "fifo":
+            message_ids_to_delete, delete_size = MessageService.pick_messages_to_delete_by_fifo(memory_id, tenant_id, size_to_delete)
+            MessageService.delete_message({"message_id": message_ids_to_delete}, tenant_id, memory_id)
+            decrease_memory_size_cache(memory_id, tenant_id, delete_size)
+        else:
+            return False, "Failed to insert message into memory. Memory size reached limit and cannot decide which to delete."
+    fail_cases = MessageService.insert_message(message_list, tenant_id, memory_id)
+    if fail_cases:
+        return False, "Failed to insert message into memory. Details: " + "; ".join(fail_cases)
+
+    increase_memory_size_cache(memory_id, tenant_id, new_msg_size)
+    return True, "Message saved successfully."
+
+
+async def extract_by_llm(tenant_id: str, llm_id: str, extract_conf: dict, memory_type: List[str], user_input: str,
+                         agent_response: str, system_prompt: str = "", user_prompt: str="") -> List[dict]:
+    llm_type = TenantLLMService.llm_id2llm_type(llm_id)
+    if not llm_type:
+        raise RuntimeError(f"Unknown type of LLM '{llm_id}'")
+    if not system_prompt:
+        system_prompt = PromptAssembler.assemble_system_prompt({"memory_type": memory_type})
+    conversation_content = f"User Input: {user_input}\nAgent Response: {agent_response}"
+    conversation_time = timestamp_to_date(current_timestamp())
+    user_prompts = []
+    if user_prompt:
+        user_prompts.append({"role": "user", "content": user_prompt})
+        user_prompts.append({"role": "user", "content": f"Conversation: {conversation_content}\nConversation Time: {conversation_time}\nCurrent Time: {conversation_time}"})
+    else:
+        user_prompts.append({"role": "user", "content": PromptAssembler.assemble_user_prompt(conversation_content, conversation_time, conversation_time)})
+    llm = LLMBundle(tenant_id, llm_type, llm_id)
+    res = await llm.async_chat(system_prompt, user_prompts, extract_conf)
+    res_json = get_json_result_from_llm_response(res)
+    return [{
+        "content": extracted_content["content"],
+        "valid_at": format_iso_8601_to_ymd_hms(extracted_content["valid_at"]),
+        "invalid_at": format_iso_8601_to_ymd_hms(extracted_content["invalid_at"]) if extracted_content.get("invalid_at") else "",
+        "message_type": message_type
+    } for message_type, extracted_content_list in res_json.items() for extracted_content in extracted_content_list]
+
+
+def query_message(filter_dict: dict, params: dict):
+    """
+    :param filter_dict: {
+        "memory_id": List[str],
+        "agent_id": optional
+        "session_id": optional
+    }
+    :param params: {
+        "query": question str,
+        "similarity_threshold": float,
+        "keywords_similarity_weight": float,
+        "top_n": int
+    }
+    """
+    memory_ids = filter_dict["memory_id"]
+    memory_list = MemoryService.get_by_ids(memory_ids)
+    if not memory_list:
+        return []
+
+    condition_dict = {k: v for k, v in filter_dict.items() if v}
+    uids = [memory.tenant_id for memory in memory_list]
+
+    question = params["query"]
+    question = question.strip()
+    memory = memory_list[0]
+    embd_model = LLMBundle(memory.tenant_id, llm_type=LLMType.EMBEDDING, llm_name=memory.embd_id)
+    match_dense = get_vector(question, embd_model, similarity=params["similarity_threshold"])
+    match_text, _ = MsgTextQuery().question(question, min_match=0.3)
+    keywords_similarity_weight = params.get("keywords_similarity_weight", 0.7)
+    fusion_expr = FusionExpr("weighted_sum", params["top_n"], {"weights": ",".join([str(keywords_similarity_weight), str(1 - keywords_similarity_weight)])})
+
+    return MessageService.search_message(memory_ids, condition_dict, uids, [match_text, match_dense, fusion_expr], params["top_n"])
+
+
+def init_message_id_sequence():
+    message_id_redis_key = "id_generator:memory"
+    if REDIS_CONN.exist(message_id_redis_key):
+        current_max_id = REDIS_CONN.get(message_id_redis_key)
+        logging.info(f"No need to init message_id sequence, current max id is {current_max_id}.")
+    else:
+        max_id = 1
+        exist_memory_list = MemoryService.get_all_memory()
+        if not exist_memory_list:
+            REDIS_CONN.set(message_id_redis_key, max_id)
+        else:
+            max_id = MessageService.get_max_message_id(
+                uid_list=[m.tenant_id for m in exist_memory_list],
+                memory_ids=[m.id for m in exist_memory_list]
+            )
+            REDIS_CONN.set(message_id_redis_key, max_id)
+        logging.info(f"Init message_id sequence done, current max id is {max_id}.")
+
+
+def get_memory_size_cache(memory_id: str, uid: str):
+    redis_key = f"memory_{memory_id}"
+    if REDIS_CONN.exists(redis_key):
+        return REDIS_CONN.get(redis_key)
+    else:
+        memory_size_map = MessageService.calculate_memory_size(
+            [memory_id],
+            [uid]
+        )
+        memory_size = memory_size_map.get(memory_id, 0)
+        set_memory_size_cache(memory_id, memory_size)
+        return memory_size
+
+
+def set_memory_size_cache(memory_id: str, size: int):
+    redis_key = f"memory_{memory_id}"
+    return REDIS_CONN.set(redis_key, size)
+
+
+def increase_memory_size_cache(memory_id: str, uid: str, size: int):
+    current_value = get_memory_size_cache(memory_id, uid)
+    return set_memory_size_cache(memory_id, current_value + size)
+
+
+def decrease_memory_size_cache(memory_id: str, uid: str, size: int):
+    current_value = get_memory_size_cache(memory_id, uid)
+    return set_memory_size_cache(memory_id, max(current_value - size, 0))
+
+
+def init_memory_size_cache():
+    memory_list = MemoryService.get_all_memory()
+    if not memory_list:
+        logging.info("No memory found, no need to init memory size.")
+    else:
+        memory_size_map = MessageService.calculate_memory_size(
+            memory_ids=[m.id for m in memory_list],
+            uid_list=[m.tenant_id for m in memory_list],
+        )
+        for memory in memory_list:
+            memory_size = memory_size_map.get(memory.id, 0)
+            set_memory_size_cache(memory.id, memory_size)
+        logging.info("Memory size cache init done.")

api/db/joint_services/user_account_service.py

@@ -34,6 +34,8 @@ from api.db.services.task_service import TaskService
 from api.db.services.tenant_llm_service import TenantLLMService
 from api.db.services.user_canvas_version import UserCanvasVersionService
 from api.db.services.user_service import TenantService, UserService, UserTenantService
+from api.db.services.memory_service import MemoryService
+from memory.services.messages import MessageService
 from rag.nlp import search
 from common.constants import ActiveEnum
 from common import settings
@@ -200,7 +202,16 @@ def delete_user_data(user_id: str) -> dict:
             done_msg += f"- Deleted {llm_delete_res} tenant-LLM records.\n"
             langfuse_delete_res = TenantLangfuseService.delete_ty_tenant_id(tenant_id)
             done_msg += f"- Deleted {langfuse_delete_res} langfuse records.\n"
-            # step1.3 delete own tenant
+            # step1.3 delete memory and messages
+            user_memory = MemoryService.get_by_tenant_id(tenant_id)
+            if user_memory:
+                for memory in user_memory:
+                    if MessageService.has_index(tenant_id, memory.id):
+                        MessageService.delete_index(tenant_id, memory.id)
+                done_msg += " Deleted memory index."
+                memory_delete_res = MemoryService.delete_by_ids([m.id for m in user_memory])
+                done_msg += f"Deleted {memory_delete_res} memory datasets."
+            # step1.4 delete own tenant
             tenant_delete_res = TenantService.delete_by_id(tenant_id)
             done_msg += f"- Deleted {tenant_delete_res} tenant.\n"
         # step2 delete user-tenant relation

api/db/services/canvas_service.py

@@ -123,6 +123,19 @@ class UserCanvasService(CommonService):
             logging.exception(e)
             return False, None
 
+    @classmethod
+    @DB.connection_context()
+    def get_basic_info_by_canvas_ids(cls, canvas_id):
+        fields = [
+            cls.model.id,
+            cls.model.avatar,
+            cls.model.user_id,
+            cls.model.title,
+            cls.model.permission,
+            cls.model.canvas_category
+        ]
+        return cls.model.select(*fields).where(cls.model.id.in_(canvas_id)).dicts()
+
     @classmethod
     @DB.connection_context()
     def get_by_tenant_ids(cls, joined_tenant_ids, user_id,

api/db/services/document_service.py

@@ -38,7 +38,7 @@ from common.time_utils import current_timestamp, get_format_time
 from common.constants import LLMType, ParserType, StatusEnum, TaskStatus, SVR_CONSUMER_GROUP_NAME
 from rag.nlp import rag_tokenizer, search
 from rag.utils.redis_conn import REDIS_CONN
-from rag.utils.doc_store_conn import OrderByExpr
+from common.doc_store.doc_store_base import OrderByExpr
 from common import settings
 
 
@@ -345,7 +345,7 @@ class DocumentService(CommonService):
                 chunks = settings.docStoreConn.search(["img_id"], [], {"doc_id": doc.id}, [], OrderByExpr(),
                                                       page * page_size, page_size, search.index_name(tenant_id),
                                                       [doc.kb_id])
-                chunk_ids = settings.docStoreConn.get_chunk_ids(chunks)
+                chunk_ids = settings.docStoreConn.get_doc_ids(chunks)
                 if not chunk_ids:
                     break
                 all_chunk_ids.extend(chunk_ids)
@@ -1230,8 +1230,8 @@ def doc_upload_and_parse(conversation_id, file_objs, user_id):
             d["q_%d_vec" % len(v)] = v
         for b in range(0, len(cks), es_bulk_size):
             if try_create_idx:
-                if not settings.docStoreConn.indexExist(idxnm, kb_id):
-                    settings.docStoreConn.createIdx(idxnm, kb_id, len(vectors[0]))
+                if not settings.docStoreConn.index_exist(idxnm, kb_id):
+                    settings.docStoreConn.create_idx(idxnm, kb_id, len(vectors[0]))
                 try_create_idx = False
             settings.docStoreConn.insert(cks[b:b + es_bulk_size], idxnm, kb_id)
 

api/db/services/memory_service.py

@@ -15,7 +15,6 @@
 #
 from typing import List
 
-from api.apps import current_user
 from api.db.db_models import DB, Memory, User
 from api.db.services import duplicate_name
 from api.db.services.common_service import CommonService
@@ -23,6 +22,7 @@ from api.utils.memory_utils import calculate_memory_type
 from api.constants import MEMORY_NAME_LIMIT
 from common.misc_utils import get_uuid
 from common.time_utils import get_format_time, current_timestamp
+from memory.utils.prompt_util import PromptAssembler
 
 
 class MemoryService(CommonService):
@@ -34,6 +34,17 @@ class MemoryService(CommonService):
     def get_by_memory_id(cls, memory_id: str):
         return cls.model.select().where(cls.model.id == memory_id).first()
 
+    @classmethod
+    @DB.connection_context()
+    def get_by_tenant_id(cls, tenant_id: str):
+        return cls.model.select().where(cls.model.tenant_id == tenant_id)
+
+    @classmethod
+    @DB.connection_context()
+    def get_all_memory(cls):
+        memory_list = cls.model.select()
+        return list(memory_list)
+
     @classmethod
     @DB.connection_context()
     def get_with_owner_name_by_id(cls, memory_id: str):
@@ -53,7 +64,9 @@ class MemoryService(CommonService):
             cls.model.forgetting_policy,
             cls.model.temperature,
             cls.model.system_prompt,
-            cls.model.user_prompt
+            cls.model.user_prompt,
+            cls.model.create_date,
+            cls.model.create_time
         ]
         memory = cls.model.select(*fields).join(User, on=(cls.model.tenant_id == User.id)).where(
             cls.model.id == memory_id
@@ -72,7 +85,9 @@ class MemoryService(CommonService):
             cls.model.memory_type,
             cls.model.storage_type,
             cls.model.permissions,
-            cls.model.description
+            cls.model.description,
+            cls.model.create_time,
+            cls.model.create_date
         ]
         memories = cls.model.select(*fields).join(User, on=(cls.model.tenant_id == User.id))
         if filter_dict.get("tenant_id"):
@@ -110,6 +125,7 @@ class MemoryService(CommonService):
             "tenant_id": tenant_id,
             "embd_id": embd_id,
             "llm_id": llm_id,
+            "system_prompt": PromptAssembler.assemble_system_prompt({"memory_type": memory_type}),
             "create_time": current_timestamp(),
             "create_date": get_format_time(),
             "update_time": current_timestamp(),
@@ -126,7 +142,7 @@ class MemoryService(CommonService):
 
     @classmethod
     @DB.connection_context()
-    def update_memory(cls, memory_id: str, update_dict: dict):
+    def update_memory(cls, tenant_id: str, memory_id: str, update_dict: dict):
         if not update_dict:
             return 0
         if "temperature" in update_dict and isinstance(update_dict["temperature"], str):
@@ -135,7 +151,7 @@ class MemoryService(CommonService):
             update_dict["name"] = duplicate_name(
                 cls.query,
                 name=update_dict["name"],
-                tenant_id=current_user.id
+                tenant_id=tenant_id
             )
         update_dict.update({
             "update_time": current_timestamp(),

common/doc_store/doc_store_base.py

@@ -13,7 +13,6 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
-
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 import numpy as np
@@ -22,7 +21,6 @@ DEFAULT_MATCH_VECTOR_TOPN = 10
 DEFAULT_MATCH_SPARSE_TOPN = 10
 VEC = list | np.ndarray
 
-
 @dataclass
 class SparseVector:
     indices: list[int]
@@ -55,14 +53,13 @@ class SparseVector:
     def __repr__(self):
         return str(self)
 
-
-class MatchTextExpr(ABC):
+class MatchTextExpr:
     def __init__(
         self,
         fields: list[str],
         matching_text: str,
         topn: int,
-        extra_options: dict = dict(),
+        extra_options: dict | None = None,
     ):
         self.fields = fields
         self.matching_text = matching_text
@@ -70,7 +67,7 @@ class MatchTextExpr(ABC):
         self.extra_options = extra_options
 
 
-class MatchDenseExpr(ABC):
+class MatchDenseExpr:
     def __init__(
         self,
         vector_column_name: str,
@@ -78,7 +75,7 @@ class MatchDenseExpr(ABC):
         embedding_data_type: str,
         distance_type: str,
         topn: int = DEFAULT_MATCH_VECTOR_TOPN,
-        extra_options: dict = dict(),
+        extra_options: dict | None = None,
     ):
         self.vector_column_name = vector_column_name
         self.embedding_data = embedding_data
@@ -88,7 +85,7 @@ class MatchDenseExpr(ABC):
         self.extra_options = extra_options
 
 
-class MatchSparseExpr(ABC):
+class MatchSparseExpr:
     def __init__(
         self,
         vector_column_name: str,
@@ -104,7 +101,7 @@ class MatchSparseExpr(ABC):
         self.opt_params = opt_params
 
 
-class MatchTensorExpr(ABC):
+class MatchTensorExpr:
     def __init__(
         self,
         column_name: str,
@@ -120,7 +117,7 @@ class MatchTensorExpr(ABC):
         self.extra_option = extra_option
 
 
-class FusionExpr(ABC):
+class FusionExpr:
     def __init__(self, method: str, topn: int, fusion_params: dict | None = None):
         self.method = method
         self.topn = topn
@@ -129,7 +126,8 @@ class FusionExpr(ABC):
 
 MatchExpr = MatchTextExpr | MatchDenseExpr | MatchSparseExpr | MatchTensorExpr | FusionExpr
 
-class OrderByExpr(ABC):
+
+class OrderByExpr:
     def __init__(self):
         self.fields = list()
     def asc(self, field: str):
@@ -141,13 +139,14 @@ class OrderByExpr(ABC):
     def fields(self):
         return self.fields
 
+
 class DocStoreConnection(ABC):
     """
     Database operations
     """
 
     @abstractmethod
-    def dbType(self) -> str:
+    def db_type(self) -> str:
         """
         Return the type of the database.
         """
@@ -165,21 +164,21 @@ class DocStoreConnection(ABC):
     """
 
     @abstractmethod
-    def createIdx(self, indexName: str, knowledgebaseId: str, vectorSize: int):
+    def create_idx(self, index_name: str, dataset_id: str, vector_size: int):
         """
         Create an index with given name
         """
         raise NotImplementedError("Not implemented")
 
     @abstractmethod
-    def deleteIdx(self, indexName: str, knowledgebaseId: str):
+    def delete_idx(self, index_name: str, dataset_id: str):
         """
         Delete an index with given name
         """
         raise NotImplementedError("Not implemented")
 
     @abstractmethod
-    def indexExist(self, indexName: str, knowledgebaseId: str) -> bool:
+    def index_exist(self, index_name: str, dataset_id: str) -> bool:
         """
         Check if an index with given name exists
         """
@@ -191,16 +190,16 @@ class DocStoreConnection(ABC):
 
     @abstractmethod
     def search(
-        self, selectFields: list[str],
-            highlightFields: list[str],
+        self, select_fields: list[str],
+            highlight_fields: list[str],
             condition: dict,
-            matchExprs: list[MatchExpr],
-            orderBy: OrderByExpr,
+            match_expressions: list[MatchExpr],
+            order_by: OrderByExpr,
             offset: int,
             limit: int,
-            indexNames: str|list[str],
-            knowledgebaseIds: list[str],
-            aggFields: list[str] = [],
+            index_names: str|list[str],
+            dataset_ids: list[str],
+            agg_fields: list[str] | None = None,
             rank_feature: dict | None = None
     ):
         """
@@ -209,28 +208,28 @@ class DocStoreConnection(ABC):
         raise NotImplementedError("Not implemented")
 
     @abstractmethod
-    def get(self, chunkId: str, indexName: str, knowledgebaseIds: list[str]) -> dict | None:
+    def get(self, data_id: str, index_name: str, dataset_ids: list[str]) -> dict | None:
         """
         Get single chunk with given id
         """
         raise NotImplementedError("Not implemented")
 
     @abstractmethod
-    def insert(self, rows: list[dict], indexName: str, knowledgebaseId: str = None) -> list[str]:
+    def insert(self, rows: list[dict], index_name: str, dataset_id: str = None) -> list[str]:
         """
         Update or insert a bulk of rows
         """
         raise NotImplementedError("Not implemented")
 
     @abstractmethod
-    def update(self, condition: dict, newValue: dict, indexName: str, knowledgebaseId: str) -> bool:
+    def update(self, condition: dict, new_value: dict, index_name: str, dataset_id: str) -> bool:
         """
         Update rows with given conjunctive equivalent filtering condition
         """
         raise NotImplementedError("Not implemented")
 
     @abstractmethod
-    def delete(self, condition: dict, indexName: str, knowledgebaseId: str) -> int:
+    def delete(self, condition: dict, index_name: str, dataset_id: str) -> int:
         """
         Delete rows with given conjunctive equivalent filtering condition
         """
@@ -245,7 +244,7 @@ class DocStoreConnection(ABC):
         raise NotImplementedError("Not implemented")
 
     @abstractmethod
-    def get_chunk_ids(self, res):
+    def get_doc_ids(self, res):
         raise NotImplementedError("Not implemented")
 
     @abstractmethod
@@ -253,18 +252,18 @@ class DocStoreConnection(ABC):
         raise NotImplementedError("Not implemented")
 
     @abstractmethod
-    def get_highlight(self, res, keywords: list[str], fieldnm: str):
+    def get_highlight(self, res, keywords: list[str], field_name: str):
         raise NotImplementedError("Not implemented")
 
     @abstractmethod
-    def get_aggregation(self, res, fieldnm: str):
+    def get_aggregation(self, res, field_name: str):
         raise NotImplementedError("Not implemented")
 
     """
     SQL
     """
     @abstractmethod
-    def sql(sql: str, fetch_size: int, format: str):
+    def sql(self, sql: str, fetch_size: int, format: str):
         """
         Run the sql generated by text-to-sql
         """

common/doc_store/es_conn_base.py

@@ -0,0 +1,326 @@
+#
+#  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
+from abc import abstractmethod
+
+from elasticsearch import Elasticsearch, NotFoundError
+from elasticsearch_dsl import Index
+from elastic_transport import ConnectionTimeout
+from common.file_utils import get_project_base_directory
+from common.misc_utils import convert_bytes
+from common.doc_store.doc_store_base import DocStoreConnection, OrderByExpr, MatchExpr
+from rag.nlp import is_english, rag_tokenizer
+from common import settings
+
+ATTEMPT_TIME = 2
+
+
+class ESConnectionBase(DocStoreConnection):
+    def __init__(self, mapping_file_name: str="mapping.json", logger_name: str='ragflow.es_conn'):
+        self.logger = logging.getLogger(logger_name)
+
+        self.info = {}
+        self.logger.info(f"Use Elasticsearch {settings.ES['hosts']} as the doc engine.")
+        for _ in range(ATTEMPT_TIME):
+            try:
+                if self._connect():
+                    break
+            except Exception as e:
+                self.logger.warning(f"{str(e)}. Waiting Elasticsearch {settings.ES['hosts']} to be healthy.")
+                time.sleep(5)
+
+        if not self.es.ping():
+            msg = f"Elasticsearch {settings.ES['hosts']} is unhealthy in 120s."
+            self.logger.error(msg)
+            raise Exception(msg)
+        v = self.info.get("version", {"number": "8.11.3"})
+        v = v["number"].split(".")[0]
+        if int(v) < 8:
+            msg = f"Elasticsearch version must be greater than or equal to 8, current version: {v}"
+            self.logger.error(msg)
+            raise Exception(msg)
+        fp_mapping = os.path.join(get_project_base_directory(), "conf", mapping_file_name)
+        if not os.path.exists(fp_mapping):
+            msg = f"Elasticsearch mapping file not found at {fp_mapping}"
+            self.logger.error(msg)
+            raise Exception(msg)
+        self.mapping = json.load(open(fp_mapping, "r"))
+        self.logger.info(f"Elasticsearch {settings.ES['hosts']} is healthy.")
+
+    def _connect(self):
+        self.es = Elasticsearch(
+            settings.ES["hosts"].split(","),
+            basic_auth=(settings.ES["username"], settings.ES[
+                "password"]) if "username" in settings.ES and "password" in settings.ES else None,
+            verify_certs= settings.ES.get("verify_certs", False),
+            timeout=600 )
+        if self.es:
+            self.info = self.es.info()
+            return True
+        return False
+
+    """
+    Database operations
+    """
+
+    def db_type(self) -> str:
+        return "elasticsearch"
+
+    def health(self) -> dict:
+        health_dict = dict(self.es.cluster.health())
+        health_dict["type"] = "elasticsearch"
+        return health_dict
+
+    def get_cluster_stats(self):
+        """
+        curl -XGET "http://{es_host}/_cluster/stats" -H "kbn-xsrf: reporting" to view raw stats.
+        """
+        raw_stats = self.es.cluster.stats()
+        self.logger.debug(f"ESConnection.get_cluster_stats: {raw_stats}")
+        try:
+            res = {
+                'cluster_name': raw_stats['cluster_name'],
+                'status': raw_stats['status']
+            }
+            indices_status = raw_stats['indices']
+            res.update({
+                'indices': indices_status['count'],
+                'indices_shards': indices_status['shards']['total']
+            })
+            doc_info = indices_status['docs']
+            res.update({
+                'docs': doc_info['count'],
+                'docs_deleted': doc_info['deleted']
+            })
+            store_info = indices_status['store']
+            res.update({
+                'store_size': convert_bytes(store_info['size_in_bytes']),
+                'total_dataset_size': convert_bytes(store_info['total_data_set_size_in_bytes'])
+            })
+            mappings_info = indices_status['mappings']
+            res.update({
+                'mappings_fields': mappings_info['total_field_count'],
+                'mappings_deduplicated_fields': mappings_info['total_deduplicated_field_count'],
+                'mappings_deduplicated_size': convert_bytes(mappings_info['total_deduplicated_mapping_size_in_bytes'])
+            })
+            node_info = raw_stats['nodes']
+            res.update({
+                'nodes': node_info['count']['total'],
+                'nodes_version': node_info['versions'],
+                'os_mem': convert_bytes(node_info['os']['mem']['total_in_bytes']),
+                'os_mem_used': convert_bytes(node_info['os']['mem']['used_in_bytes']),
+                'os_mem_used_percent': node_info['os']['mem']['used_percent'],
+                'jvm_versions': node_info['jvm']['versions'][0]['vm_version'],
+                'jvm_heap_used': convert_bytes(node_info['jvm']['mem']['heap_used_in_bytes']),
+                'jvm_heap_max': convert_bytes(node_info['jvm']['mem']['heap_max_in_bytes'])
+            })
+            return res
+
+        except Exception as e:
+            self.logger.exception(f"ESConnection.get_cluster_stats: {e}")
+            return None
+
+    """
+    Table operations
+    """
+
+    def create_idx(self, index_name: str, dataset_id: str, vector_size: int):
+        if self.index_exist(index_name, dataset_id):
+            return True
+        try:
+            from elasticsearch.client import IndicesClient
+            return IndicesClient(self.es).create(index=index_name,
+                                                 settings=self.mapping["settings"],
+                                                 mappings=self.mapping["mappings"])
+        except Exception:
+            self.logger.exception("ESConnection.createIndex error %s" % index_name)
+
+    def delete_idx(self, index_name: str, dataset_id: str):
+        if len(dataset_id) > 0:
+            # The index need to be alive after any kb deletion since all kb under this tenant are in one index.
+            return
+        try:
+            self.es.indices.delete(index=index_name, allow_no_indices=True)
+        except NotFoundError:
+            pass
+        except Exception:
+            self.logger.exception("ESConnection.deleteIdx error %s" % index_name)
+
+    def index_exist(self, index_name: str, dataset_id: str = None) -> bool:
+        s = Index(index_name, self.es)
+        for i in range(ATTEMPT_TIME):
+            try:
+                return s.exists()
+            except ConnectionTimeout:
+                self.logger.exception("ES request timeout")
+                time.sleep(3)
+                self._connect()
+                continue
+            except Exception as e:
+                self.logger.exception(e)
+                break
+        return False
+
+    """
+    CRUD operations
+    """
+
+    def get(self, doc_id: str, index_name: str, dataset_ids: list[str]) -> dict | None:
+        for i in range(ATTEMPT_TIME):
+            try:
+                res = self.es.get(index=index_name,
+                                  id=doc_id, source=True, )
+                if str(res.get("timed_out", "")).lower() == "true":
+                    raise Exception("Es Timeout.")
+                doc = res["_source"]
+                doc["id"] = doc_id
+                return doc
+            except NotFoundError:
+                return None
+            except Exception as e:
+                self.logger.exception(f"ESConnection.get({doc_id}) got exception")
+                raise e
+        self.logger.error(f"ESConnection.get timeout for {ATTEMPT_TIME} times!")
+        raise Exception("ESConnection.get timeout.")
+
+    @abstractmethod
+    def search(
+            self, select_fields: list[str],
+            highlight_fields: list[str],
+            condition: dict,
+            match_expressions: list[MatchExpr],
+            order_by: OrderByExpr,
+            offset: int,
+            limit: int,
+            index_names: str | list[str],
+            dataset_ids: list[str],
+            agg_fields: list[str] | None = None,
+            rank_feature: dict | None = None
+    ):
+        raise NotImplementedError("Not implemented")
+
+    @abstractmethod
+    def insert(self, documents: list[dict], index_name: str, dataset_id: str = None) -> list[str]:
+        raise NotImplementedError("Not implemented")
+
+    @abstractmethod
+    def update(self, condition: dict, new_value: dict, index_name: str, dataset_id: str) -> bool:
+        raise NotImplementedError("Not implemented")
+
+    @abstractmethod
+    def delete(self, condition: dict, index_name: str, dataset_id: str) -> int:
+        raise NotImplementedError("Not implemented")
+
+    """
+    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 _get_source(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
+
+    @abstractmethod
+    def get_fields(self, res, fields: list[str]) -> dict[str, dict]:
+        raise NotImplementedError("Not implemented")
+
+    def get_highlight(self, res, keywords: list[str], field_name: str):
+        ans = {}
+        for d in res["hits"]["hits"]:
+            highlights = d.get("highlight")
+            if not highlights:
+                continue
+            txt = "...".join([a for a in list(highlights.items())[0][1]])
+            if not is_english(txt.split()):
+                ans[d["_id"]] = txt
+                continue
+
+            txt = d["_source"][field_name]
+            txt = re.sub(r"[\r\n]", " ", txt, flags=re.IGNORECASE | re.MULTILINE)
+            txt_list = []
+            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
+                txt_list.append(t)
+            ans[d["_id"]] = "...".join(txt_list) if txt_list else "...".join([a for a in list(highlights.items())[0][1]])
+
+        return ans
+
+    def get_aggregation(self, res, field_name: str):
+        agg_field = "aggs_" + field_name
+        if "aggregations" not in res or agg_field not in res["aggregations"]:
+            return list()
+        buckets = res["aggregations"][agg_field]["buckets"]
+        return [(b["key"], b["doc_count"]) for b in buckets]
+
+    """
+    SQL
+    """
+
+    def sql(self, sql: str, fetch_size: int, format: str):
+        self.logger.debug(f"ESConnection.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)
+        self.logger.debug(f"ESConnection.sql to es: {sql}")
+
+        for i in range(ATTEMPT_TIME):
+            try:
+                res = self.es.sql.query(body={"query": sql, "fetch_size": fetch_size}, format=format,
+                                        request_timeout="2s")
+                return res
+            except ConnectionTimeout:
+                self.logger.exception("ES request timeout")
+                time.sleep(3)
+                self._connect()
+                continue
+            except Exception as e:
+                self.logger.exception(f"ESConnection.sql got exception. SQL:\n{sql}")
+                raise Exception(f"SQL error: {e}\n\nSQL: {sql}")
+        self.logger.error(f"ESConnection.sql timeout for {ATTEMPT_TIME} times!")
+        return None

common/doc_store/infinity_conn_base.py

@@ -0,0 +1,451 @@
+#
+#  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 os
+import re
+import json
+import time
+from abc import abstractmethod
+
+import infinity
+from infinity.common import ConflictType
+from infinity.index import IndexInfo, IndexType
+from infinity.connection_pool import ConnectionPool
+from infinity.errors import ErrorCode
+import pandas as pd
+from common.file_utils import get_project_base_directory
+from rag.nlp import is_english
+from common import settings
+from common.doc_store.doc_store_base import DocStoreConnection, MatchExpr, OrderByExpr
+
+
+class InfinityConnectionBase(DocStoreConnection):
+    def __init__(self, mapping_file_name: str="infinity_mapping.json", logger_name: str="ragflow.infinity_conn"):
+        self.dbName = settings.INFINITY.get("db_name", "default_db")
+        self.mapping_file_name = mapping_file_name
+        self.logger = logging.getLogger(logger_name)
+        infinity_uri = settings.INFINITY["uri"]
+        if ":" in infinity_uri:
+            host, port = infinity_uri.split(":")
+            infinity_uri = infinity.common.NetworkAddress(host, int(port))
+        self.connPool = None
+        self.logger.info(f"Use Infinity {infinity_uri} as the doc engine.")
+        for _ in range(24):
+            try:
+                conn_pool = ConnectionPool(infinity_uri, max_size=4)
+                inf_conn = conn_pool.get_conn()
+                res = inf_conn.show_current_node()
+                if res.error_code == ErrorCode.OK and res.server_status in ["started", "alive"]:
+                    self._migrate_db(inf_conn)
+                    self.connPool = conn_pool
+                    conn_pool.release_conn(inf_conn)
+                    break
+                conn_pool.release_conn(inf_conn)
+                self.logger.warning(f"Infinity status: {res.server_status}. Waiting Infinity {infinity_uri} to be healthy.")
+                time.sleep(5)
+            except Exception as e:
+                self.logger.warning(f"{str(e)}. Waiting Infinity {infinity_uri} to be healthy.")
+                time.sleep(5)
+        if self.connPool is None:
+            msg = f"Infinity {infinity_uri} is unhealthy in 120s."
+            self.logger.error(msg)
+            raise Exception(msg)
+        self.logger.info(f"Infinity {infinity_uri} is healthy.")
+
+    def _migrate_db(self, inf_conn):
+        inf_db = inf_conn.create_database(self.dbName, ConflictType.Ignore)
+        fp_mapping = os.path.join(get_project_base_directory(), "conf", self.mapping_file_name)
+        if not os.path.exists(fp_mapping):
+            raise Exception(f"Mapping file not found at {fp_mapping}")
+        schema = json.load(open(fp_mapping))
+        table_names = inf_db.list_tables().table_names
+        for table_name in table_names:
+            inf_table = inf_db.get_table(table_name)
+            index_names = inf_table.list_indexes().index_names
+            if "q_vec_idx" not in index_names:
+                # Skip tables not created by me
+                continue
+            column_names = inf_table.show_columns()["name"]
+            column_names = set(column_names)
+            for field_name, field_info in schema.items():
+                if field_name in column_names:
+                    continue
+                res = inf_table.add_columns({field_name: field_info})
+                assert res.error_code == infinity.ErrorCode.OK
+                self.logger.info(f"INFINITY added following column to table {table_name}: {field_name} {field_info}")
+                if field_info["type"] != "varchar" or "analyzer" not in field_info:
+                    continue
+                analyzers = field_info["analyzer"]
+                if isinstance(analyzers, str):
+                    analyzers = [analyzers]
+                for analyzer in analyzers:
+                    inf_table.create_index(
+                        f"ft_{re.sub(r'[^a-zA-Z0-9]', '_', field_name)}_{re.sub(r'[^a-zA-Z0-9]', '_', analyzer)}",
+                        IndexInfo(field_name, IndexType.FullText, {"ANALYZER": analyzer}),
+                        ConflictType.Ignore,
+                    )
+
+    """
+    Dataframe and fields convert
+    """
+
+    @staticmethod
+    @abstractmethod
+    def field_keyword(field_name: str):
+        # judge keyword or not, such as "*_kwd" tag-like columns.
+        raise NotImplementedError("Not implemented")
+
+    @abstractmethod
+    def convert_select_fields(self, output_fields: list[str]) -> list[str]:
+        # rm _kwd, _tks, _sm_tks, _with_weight suffix in field name.
+        raise NotImplementedError("Not implemented")
+
+    @staticmethod
+    @abstractmethod
+    def convert_matching_field(field_weight_str: str) -> str:
+        # convert matching field to
+        raise NotImplementedError("Not implemented")
+
+    @staticmethod
+    def list2str(lst: str | list, sep: str = " ") -> str:
+        if isinstance(lst, str):
+            return lst
+        return sep.join(lst)
+
+    def equivalent_condition_to_str(self, condition: dict, table_instance=None) -> str | None:
+        assert "_id" not in condition
+        columns = {}
+        if table_instance:
+            for n, ty, de, _ in table_instance.show_columns().rows():
+                columns[n] = (ty, de)
+
+        def exists(cln):
+            nonlocal columns
+            assert cln in columns, f"'{cln}' should be in '{columns}'."
+            ty, de = columns[cln]
+            if ty.lower().find("cha"):
+                if not de:
+                    de = ""
+                return f" {cln}!='{de}' "
+            return f"{cln}!={de}"
+
+        cond = list()
+        for k, v in condition.items():
+            if not isinstance(k, str) or not v:
+                continue
+            if self.field_keyword(k):
+                if isinstance(v, list):
+                    inCond = list()
+                    for item in v:
+                        if isinstance(item, str):
+                            item = item.replace("'", "''")
+                        inCond.append(f"filter_fulltext('{self.convert_matching_field(k)}', '{item}')")
+                    if inCond:
+                        strInCond = " or ".join(inCond)
+                        strInCond = f"({strInCond})"
+                        cond.append(strInCond)
+                else:
+                    cond.append(f"filter_fulltext('{self.convert_matching_field(k)}', '{v}')")
+            elif isinstance(v, list):
+                inCond = list()
+                for item in v:
+                    if isinstance(item, str):
+                        item = item.replace("'", "''")
+                        inCond.append(f"'{item}'")
+                    else:
+                        inCond.append(str(item))
+                if inCond:
+                    strInCond = ", ".join(inCond)
+                    strInCond = f"{k} IN ({strInCond})"
+                    cond.append(strInCond)
+            elif k == "must_not":
+                if isinstance(v, dict):
+                    for kk, vv in v.items():
+                        if kk == "exists":
+                            cond.append("NOT (%s)" % exists(vv))
+            elif isinstance(v, str):
+                cond.append(f"{k}='{v}'")
+            elif k == "exists":
+                cond.append(exists(v))
+            else:
+                cond.append(f"{k}={str(v)}")
+        return " AND ".join(cond) if cond else "1=1"
+
+    @staticmethod
+    def concat_dataframes(df_list: list[pd.DataFrame], select_fields: list[str]) -> pd.DataFrame:
+        df_list2 = [df for df in df_list if not df.empty]
+        if df_list2:
+            return pd.concat(df_list2, axis=0).reset_index(drop=True)
+
+        schema = []
+        for field_name in select_fields:
+            if field_name == "score()":  # Workaround: fix schema is changed to score()
+                schema.append("SCORE")
+            elif field_name == "similarity()":  # Workaround: fix schema is changed to similarity()
+                schema.append("SIMILARITY")
+            else:
+                schema.append(field_name)
+        return pd.DataFrame(columns=schema)
+
+    """
+    Database operations
+    """
+
+    def db_type(self) -> str:
+        return "infinity"
+
+    def health(self) -> dict:
+        """
+        Return the health status of the database.
+        """
+        inf_conn = self.connPool.get_conn()
+        res = inf_conn.show_current_node()
+        self.connPool.release_conn(inf_conn)
+        res2 = {
+            "type": "infinity",
+            "status": "green" if res.error_code == 0 and res.server_status in ["started", "alive"] else "red",
+            "error": res.error_msg,
+        }
+        return res2
+
+    """
+    Table operations
+    """
+
+    def create_idx(self, index_name: str, dataset_id: str, vector_size: int):
+        table_name = f"{index_name}_{dataset_id}"
+        inf_conn = self.connPool.get_conn()
+        inf_db = inf_conn.create_database(self.dbName, ConflictType.Ignore)
+
+        fp_mapping = os.path.join(get_project_base_directory(), "conf", self.mapping_file_name)
+        if not os.path.exists(fp_mapping):
+            raise Exception(f"Mapping file not found at {fp_mapping}")
+        schema = json.load(open(fp_mapping))
+        vector_name = f"q_{vector_size}_vec"
+        schema[vector_name] = {"type": f"vector,{vector_size},float"}
+        inf_table = inf_db.create_table(
+            table_name,
+            schema,
+            ConflictType.Ignore,
+        )
+        inf_table.create_index(
+            "q_vec_idx",
+            IndexInfo(
+                vector_name,
+                IndexType.Hnsw,
+                {
+                    "M": "16",
+                    "ef_construction": "50",
+                    "metric": "cosine",
+                    "encode": "lvq",
+                },
+            ),
+            ConflictType.Ignore,
+        )
+        for field_name, field_info in schema.items():
+            if field_info["type"] != "varchar" or "analyzer" not in field_info:
+                continue
+            analyzers = field_info["analyzer"]
+            if isinstance(analyzers, str):
+                analyzers = [analyzers]
+            for analyzer in analyzers:
+                inf_table.create_index(
+                    f"ft_{re.sub(r'[^a-zA-Z0-9]', '_', field_name)}_{re.sub(r'[^a-zA-Z0-9]', '_', analyzer)}",
+                    IndexInfo(field_name, IndexType.FullText, {"ANALYZER": analyzer}),
+                    ConflictType.Ignore,
+                )
+        self.connPool.release_conn(inf_conn)
+        self.logger.info(f"INFINITY created table {table_name}, vector size {vector_size}")
+        return True
+
+    def delete_idx(self, index_name: str, dataset_id: str):
+        table_name = f"{index_name}_{dataset_id}"
+        inf_conn = self.connPool.get_conn()
+        db_instance = inf_conn.get_database(self.dbName)
+        db_instance.drop_table(table_name, ConflictType.Ignore)
+        self.connPool.release_conn(inf_conn)
+        self.logger.info(f"INFINITY dropped table {table_name}")
+
+    def index_exist(self, index_name: str, dataset_id: str) -> bool:
+        table_name = f"{index_name}_{dataset_id}"
+        try:
+            inf_conn = self.connPool.get_conn()
+            db_instance = inf_conn.get_database(self.dbName)
+            _ = db_instance.get_table(table_name)
+            self.connPool.release_conn(inf_conn)
+            return True
+        except Exception as e:
+            self.logger.warning(f"INFINITY indexExist {str(e)}")
+        return False
+
+    """
+    CRUD operations
+    """
+
+    @abstractmethod
+    def search(
+        self,
+        select_fields: list[str],
+        highlight_fields: list[str],
+        condition: dict,
+        match_expressions: list[MatchExpr],
+        order_by: OrderByExpr,
+        offset: int,
+        limit: int,
+        index_names: str | list[str],
+        dataset_ids: list[str],
+        agg_fields: list[str] | None = None,
+        rank_feature: dict | None = None,
+    ) -> tuple[pd.DataFrame, int]:
+        raise NotImplementedError("Not implemented")
+
+    @abstractmethod
+    def get(self, doc_id: str, index_name: str, knowledgebase_ids: list[str]) -> dict | None:
+        raise NotImplementedError("Not implemented")
+
+    @abstractmethod
+    def insert(self, documents: list[dict], index_name: str, dataset_ids: str = None) -> list[str]:
+        raise NotImplementedError("Not implemented")
+
+    @abstractmethod
+    def update(self, condition: dict, new_value: dict, index_name: str, dataset_id: str) -> bool:
+        raise NotImplementedError("Not implemented")
+
+    def delete(self, condition: dict, index_name: str, dataset_id: str) -> int:
+        inf_conn = self.connPool.get_conn()
+        db_instance = inf_conn.get_database(self.dbName)
+        table_name = f"{index_name}_{dataset_id}"
+        try:
+            table_instance = db_instance.get_table(table_name)
+        except Exception:
+            self.logger.warning(f"Skipped deleting from table {table_name} since the table doesn't exist.")
+            return 0
+        filter = self.equivalent_condition_to_str(condition, table_instance)
+        self.logger.debug(f"INFINITY delete table {table_name}, filter {filter}.")
+        res = table_instance.delete(filter)
+        self.connPool.release_conn(inf_conn)
+        return res.deleted_rows
+
+    """
+    Helper functions for search result
+    """
+
+    def get_total(self, res: tuple[pd.DataFrame, int] | pd.DataFrame) -> int:
+        if isinstance(res, tuple):
+            return res[1]
+        return len(res)
+
+    def get_doc_ids(self, res: tuple[pd.DataFrame, int] | pd.DataFrame) -> list[str]:
+        if isinstance(res, tuple):
+            res = res[0]
+        return list(res["id"])
+
+    @abstractmethod
+    def get_fields(self, res: tuple[pd.DataFrame, int] | pd.DataFrame, fields: list[str]) -> dict[str, dict]:
+        raise NotImplementedError("Not implemented")
+
+    def get_highlight(self, res: tuple[pd.DataFrame, int] | pd.DataFrame, keywords: list[str], field_name: str):
+        if isinstance(res, tuple):
+            res = res[0]
+        ans = {}
+        num_rows = len(res)
+        column_id = res["id"]
+        if field_name not in res:
+            return {}
+        for i in range(num_rows):
+            id = column_id[i]
+            txt = res[field_name][i]
+            if re.search(r"<em>[^<>]+</em>", txt, flags=re.IGNORECASE | re.MULTILINE):
+                ans[id] = txt
+                continue
+            txt = re.sub(r"[\r\n]", " ", txt, flags=re.IGNORECASE | re.MULTILINE)
+            txt_list = []
+            for t in re.split(r"[.?!;\n]", txt):
+                if is_english([t]):
+                    for w in keywords:
+                        t = re.sub(
+                            r"(^|[ .?/'\"\(\)!,:;-])(%s)([ .?/'\"\(\)!,:;-])" % re.escape(w),
+                            r"\1<em>\2</em>\3",
+                            t,
+                            flags=re.IGNORECASE | re.MULTILINE,
+                        )
+                else:
+                    for w in sorted(keywords, key=len, reverse=True):
+                        t = re.sub(
+                            re.escape(w),
+                            f"<em>{w}</em>",
+                            t,
+                            flags=re.IGNORECASE | re.MULTILINE,
+                        )
+                if not re.search(r"<em>[^<>]+</em>", t, flags=re.IGNORECASE | re.MULTILINE):
+                    continue
+                txt_list.append(t)
+            if txt_list:
+                ans[id] = "...".join(txt_list)
+            else:
+                ans[id] = txt
+        return ans
+
+    def get_aggregation(self, res: tuple[pd.DataFrame, int] | pd.DataFrame, field_name: str):
+        """
+        Manual aggregation for tag fields since Infinity doesn't provide native aggregation
+        """
+        from collections import Counter
+
+        # Extract DataFrame from result
+        if isinstance(res, tuple):
+            df, _ = res
+        else:
+            df = res
+
+        if df.empty or field_name not in df.columns:
+            return []
+
+        # Aggregate tag counts
+        tag_counter = Counter()
+
+        for value in df[field_name]:
+            if pd.isna(value) or not value:
+                continue
+
+            # Handle different tag formats
+            if isinstance(value, str):
+                # Split by ### for tag_kwd field or comma for other formats
+                if field_name == "tag_kwd" and "###" in value:
+                    tags = [tag.strip() for tag in value.split("###") if tag.strip()]
+                else:
+                    # Try comma separation as fallback
+                    tags = [tag.strip() for tag in value.split(",") if tag.strip()]
+
+                for tag in tags:
+                    if tag:  # Only count non-empty tags
+                        tag_counter[tag] += 1
+            elif isinstance(value, list):
+                # Handle list format
+                for tag in value:
+                    if tag and isinstance(tag, str):
+                        tag_counter[tag.strip()] += 1
+
+        # Return as list of [tag, count] pairs, sorted by count descending
+        return [[tag, count] for tag, count in tag_counter.most_common()]
+
+    """
+    SQL
+    """
+
+    def sql(self, sql: str, fetch_size: int, format: str):
+        raise NotImplementedError("Not implemented")

common/query_base.py

@@ -0,0 +1,72 @@
+#
+#  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
+from abc import ABC, abstractmethod
+
+
+class QueryBase(ABC):
+
+    @staticmethod
+    def is_chinese(line):
+        arr = re.split(r"[ \t]+", line)
+        if len(arr) <= 3:
+            return True
+        e = 0
+        for t in arr:
+            if not re.match(r"[a-zA-Z]+$", t):
+                e += 1
+        return e * 1.0 / len(arr) >= 0.7
+
+    @staticmethod
+    def sub_special_char(line):
+        return re.sub(r"([:\{\}/\[\]\-\*\"\(\)\|\+~\^])", r"\\\1", line).strip()
+
+    @staticmethod
+    def rmWWW(txt):
+        patts = [
+            (
+                r"是*(怎么办|什么样的|哪家|一下|那家|请问|啥样|咋样了|什么时候|何时|何地|何人|是否|是不是|多少|哪里|怎么|哪儿|怎么样|如何|哪些|是啥|啥是|啊|吗|呢|吧|咋|什么|有没有|呀|谁|哪位|哪个)是*",
+                "",
+            ),
+            (r"(^| )(what|who|how|which|where|why)('re|'s)? ", " "),
+            (
+                r"(^| )('s|'re|is|are|were|was|do|does|did|don't|doesn't|didn't|has|have|be|there|you|me|your|my|mine|just|please|may|i|should|would|wouldn't|will|won't|done|go|for|with|so|the|a|an|by|i'm|it's|he's|she's|they|they're|you're|as|by|on|in|at|up|out|down|of|to|or|and|if) ",
+                " ")
+        ]
+        otxt = txt
+        for r, p in patts:
+            txt = re.sub(r, p, txt, flags=re.IGNORECASE)
+        if not txt:
+            txt = otxt
+        return txt
+
+    @staticmethod
+    def add_space_between_eng_zh(txt):
+        # (ENG/ENG+NUM) + ZH
+        txt = re.sub(r'([A-Za-z]+[0-9]+)([\u4e00-\u9fa5]+)', r'\1 \2', txt)
+        # ENG + ZH
+        txt = re.sub(r'([A-Za-z])([\u4e00-\u9fa5]+)', r'\1 \2', txt)
+        # ZH + (ENG/ENG+NUM)
+        txt = re.sub(r'([\u4e00-\u9fa5]+)([A-Za-z]+[0-9]+)', r'\1 \2', txt)
+        txt = re.sub(r'([\u4e00-\u9fa5]+)([A-Za-z])', r'\1 \2', txt)
+        return txt
+
+    @abstractmethod
+    def question(self, text, tbl, min_match):
+        """
+        Returns a query object based on the input text, table, and minimum match criteria.
+        """
+        raise NotImplementedError("Not implemented")

common/settings.py

@@ -39,6 +39,9 @@ from rag.utils.oss_conn import RAGFlowOSS
 
 from rag.nlp import search
 
+import memory.utils.es_conn as memory_es_conn
+import memory.utils.infinity_conn as memory_infinity_conn
+
 LLM = None
 LLM_FACTORY = None
 LLM_BASE_URL = None
@@ -76,9 +79,11 @@ FEISHU_OAUTH = None
 OAUTH_CONFIG = None
 DOC_ENGINE = os.getenv('DOC_ENGINE', 'elasticsearch')
 DOC_ENGINE_INFINITY = (DOC_ENGINE.lower() == "infinity")
+MSG_ENGINE = DOC_ENGINE
 
 
 docStoreConn = None
+msgStoreConn = None
 
 retriever = None
 kg_retriever = None
@@ -256,6 +261,15 @@ def init_settings():
     else:
         raise Exception(f"Not supported doc engine: {DOC_ENGINE}")
 
+    global MSG_ENGINE, msgStoreConn
+    MSG_ENGINE = DOC_ENGINE  # use the same engine for message store
+    if MSG_ENGINE == "elasticsearch":
+        ES = get_base_config("es", {})
+        msgStoreConn = memory_es_conn.ESConnection()
+    elif MSG_ENGINE == "infinity":
+        INFINITY = get_base_config("infinity", {"uri": "infinity:23817"})
+        msgStoreConn = memory_infinity_conn.InfinityConnection()
+
     global AZURE, S3, MINIO, OSS, GCS
     if STORAGE_IMPL_TYPE in ['AZURE_SPN', 'AZURE_SAS']:
         AZURE = get_base_config("azure", {})

common/time_utils.py

@@ -14,6 +14,7 @@
 #  limitations under the License.
 
 import datetime
+import logging
 import time
 
 def current_timestamp():
@@ -123,4 +124,31 @@ def delta_seconds(date_string: str):
         3600.0  # If current time is 2024-01-01 13:00:00
     """
     dt = datetime.datetime.strptime(date_string, "%Y-%m-%d %H:%M:%S")
-    return (datetime.datetime.now() - dt).total_seconds()
+    return (datetime.datetime.now() - dt).total_seconds()
+
+
+def format_iso_8601_to_ymd_hms(time_str: str) -> str:
+    """
+    Convert ISO 8601 formatted string to "YYYY-MM-DD HH:MM:SS" format.
+
+    Args:
+        time_str: ISO 8601 date string (e.g. "2024-01-01T12:00:00Z")
+
+    Returns:
+        str: Date string in "YYYY-MM-DD HH:MM:SS" format
+
+    Example:
+        >>> format_iso_8601_to_ymd_hms("2024-01-01T12:00:00Z")
+        '2024-01-01 12:00:00'
+    """
+    from dateutil import parser
+
+    try:
+        if parser.isoparse(time_str):
+            dt = datetime.datetime.fromisoformat(time_str.replace("Z", "+00:00"))
+            return dt.strftime("%Y-%m-%d %H:%M:%S")
+        else:
+            return time_str
+    except Exception as e:
+        logging.error(str(e))
+        return time_str

conf/message_infinity_mapping.json

@@ -0,0 +1,19 @@
+{
+  "id": {"type": "varchar", "default": ""},
+  "message_id": {"type": "integer", "default": 0},
+  "message_type_kwd": {"type": "varchar", "default": ""},
+  "source_id": {"type":  "integer", "default": 0},
+  "memory_id": {"type": "varchar", "default": ""},
+  "user_id": {"type": "varchar", "default": ""},
+  "agent_id": {"type": "varchar", "default": ""},
+  "session_id": {"type": "varchar", "default": ""},
+  "valid_at": {"type":  "varchar", "default": ""},
+  "valid_at_flt": {"type": "float", "default": 0.0},
+  "invalid_at": {"type":  "varchar", "default": ""},
+  "invalid_at_flt": {"type": "float", "default": 0.0},
+  "forget_at": {"type":  "varchar", "default": ""},
+  "forget_at_flt": {"type": "float", "default": 0.0},
+  "status_int": {"type": "integer", "default": 1},
+  "zone_id": {"type": "integer", "default": 0},
+  "content": {"type": "varchar", "default": "", "analyzer": ["rag-coarse", "rag-fine"], "comment": "content_ltks"}
+}

graphrag/search.py

@@ -24,11 +24,11 @@ from common.misc_utils import get_uuid
 from graphrag.query_analyze_prompt import PROMPTS
 from graphrag.utils import get_entity_type2samples, get_llm_cache, set_llm_cache, get_relation
 from common.token_utils import num_tokens_from_string
-from rag.utils.doc_store_conn import OrderByExpr
 
 from rag.nlp.search import Dealer, index_name
 from common.float_utils import get_float
 from common import settings
+from common.doc_store.doc_store_base import OrderByExpr
 
 
 class KGSearch(Dealer):

graphrag/utils.py

@@ -26,9 +26,9 @@ from networkx.readwrite import json_graph
 from common.misc_utils import get_uuid
 from common.connection_utils import timeout
 from rag.nlp import rag_tokenizer, search
-from rag.utils.doc_store_conn import OrderByExpr
 from rag.utils.redis_conn import REDIS_CONN
 from common import settings
+from common.doc_store.doc_store_base import OrderByExpr
 
 GRAPH_FIELD_SEP = "<SEP>"
 

memory/services/messages.py

@@ -0,0 +1,240 @@
+#
+#  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 sys
+from typing import List
+
+from common import settings
+from common.doc_store.doc_store_base import OrderByExpr, MatchExpr
+
+
+def index_name(uid: str): return f"memory_{uid}"
+
+
+class MessageService:
+
+    @classmethod
+    def has_index(cls, uid: str, memory_id: str):
+        index = index_name(uid)
+        return settings.msgStoreConn.index_exist(index, memory_id)
+
+    @classmethod
+    def create_index(cls, uid: str, memory_id: str, vector_size: int):
+        index = index_name(uid)
+        return settings.msgStoreConn.create_idx(index, memory_id, vector_size)
+
+    @classmethod
+    def delete_index(cls, uid: str, memory_id: str):
+        index = index_name(uid)
+        return settings.msgStoreConn.delete_idx(index, memory_id)
+
+    @classmethod
+    def insert_message(cls, messages: List[dict], uid: str, memory_id: str):
+        index = index_name(uid)
+        [m.update({
+            "id": f'{memory_id}_{m["message_id"]}',
+            "status": 1 if m["status"] else 0
+        }) for m in messages]
+        return settings.msgStoreConn.insert(messages, index, memory_id)
+
+    @classmethod
+    def update_message(cls, condition: dict, update_dict: dict, uid: str, memory_id: str):
+        index = index_name(uid)
+        if "status" in update_dict:
+            update_dict["status"] = 1 if update_dict["status"] else 0
+        return settings.msgStoreConn.update(condition, update_dict, index, memory_id)
+
+    @classmethod
+    def delete_message(cls, condition: dict, uid: str, memory_id: str):
+        index = index_name(uid)
+        return settings.msgStoreConn.delete(condition, index, memory_id)
+
+    @classmethod
+    def list_message(cls, uid: str, memory_id: str, agent_ids: List[str]=None, keywords: str=None, page: int=1, page_size: int=50):
+        index = index_name(uid)
+        filter_dict = {}
+        if agent_ids:
+            filter_dict["agent_id"] = agent_ids
+        if keywords:
+            filter_dict["session_id"] = keywords
+        order_by = OrderByExpr()
+        order_by.desc("valid_at")
+        res = settings.msgStoreConn.search(
+            select_fields=[
+                "message_id", "message_type", "source_id", "memory_id", "user_id", "agent_id", "session_id", "valid_at",
+                "invalid_at", "forget_at", "status"
+            ],
+            highlight_fields=[],
+            condition=filter_dict,
+            match_expressions=[], order_by=order_by,
+            offset=(page-1)*page_size, limit=page_size,
+            index_names=index, memory_ids=[memory_id], agg_fields=[], hide_forgotten=False
+        )
+        total_count = settings.msgStoreConn.get_total(res)
+        doc_mapping = settings.msgStoreConn.get_fields(res, [
+            "message_id", "message_type", "source_id", "memory_id", "user_id", "agent_id", "session_id",
+            "valid_at", "invalid_at", "forget_at", "status"
+        ])
+        return {
+            "message_list": list(doc_mapping.values()),
+            "total_count": total_count
+        }
+
+    @classmethod
+    def get_recent_messages(cls, uid_list: List[str], memory_ids: List[str], agent_id: str, session_id: str, limit: int):
+        index_names = [index_name(uid) for uid in uid_list]
+        condition_dict = {
+            "agent_id": agent_id,
+            "session_id": session_id
+        }
+        order_by = OrderByExpr()
+        order_by.desc("valid_at")
+        res = settings.msgStoreConn.search(
+            select_fields=[
+                "message_id", "message_type", "source_id", "memory_id", "user_id", "agent_id", "session_id", "valid_at",
+                "invalid_at", "forget_at", "status", "content"
+            ],
+            highlight_fields=[],
+            condition=condition_dict,
+            match_expressions=[], order_by=order_by,
+            offset=0, limit=limit,
+            index_names=index_names, memory_ids=memory_ids, agg_fields=[]
+        )
+        doc_mapping = settings.msgStoreConn.get_fields(res, [
+            "message_id", "message_type", "source_id", "memory_id","user_id", "agent_id", "session_id",
+            "valid_at", "invalid_at", "forget_at", "status", "content"
+        ])
+        return list(doc_mapping.values())
+
+    @classmethod
+    def search_message(cls, memory_ids: List[str], condition_dict: dict, uid_list: List[str], match_expressions:list[MatchExpr], top_n: int):
+        index_names = [index_name(uid) for uid in uid_list]
+        # filter only valid messages by default
+        if "status" not in condition_dict:
+            condition_dict["status"] = 1
+
+        order_by = OrderByExpr()
+        order_by.desc("valid_at")
+        res = settings.msgStoreConn.search(
+            select_fields=[
+                "message_id", "message_type", "source_id", "memory_id", "user_id", "agent_id", "session_id",
+                "valid_at",
+                "invalid_at", "forget_at", "status", "content"
+            ],
+            highlight_fields=[],
+            condition=condition_dict,
+            match_expressions=match_expressions,
+            order_by=order_by,
+            offset=0, limit=top_n,
+            index_names=index_names, memory_ids=memory_ids, agg_fields=[]
+        )
+        docs = settings.msgStoreConn.get_fields(res, [
+            "message_id", "message_type", "source_id", "memory_id", "user_id", "agent_id", "session_id", "valid_at",
+            "invalid_at", "forget_at", "status", "content"
+        ])
+        return list(docs.values())
+
+    @staticmethod
+    def calculate_message_size(message: dict):
+        return sys.getsizeof(message["content"]) + sys.getsizeof(message["content_embed"][0]) * len(message["content_embed"])
+
+    @classmethod
+    def calculate_memory_size(cls, memory_ids: List[str], uid_list: List[str]):
+        index_names = [index_name(uid) for uid in uid_list]
+        order_by = OrderByExpr()
+        order_by.desc("valid_at")
+
+        res = settings.msgStoreConn.search(
+            select_fields=["memory_id", "content", "content_embed"],
+            highlight_fields=[],
+            condition={},
+            match_expressions=[],
+            order_by=order_by,
+            offset=0, limit=2000*len(memory_ids),
+            index_names=index_names, memory_ids=memory_ids, agg_fields=[], hide_forgotten=False
+        )
+        docs = settings.msgStoreConn.get_fields(res, ["memory_id", "content", "content_embed"])
+        size_dict = {}
+        for doc in docs.values():
+            if size_dict.get(doc["memory_id"]):
+                size_dict[doc["memory_id"]] += cls.calculate_message_size(doc)
+            else:
+                size_dict[doc["memory_id"]] = cls.calculate_message_size(doc)
+        return size_dict
+
+    @classmethod
+    def pick_messages_to_delete_by_fifo(cls, memory_id: str, uid: str, size_to_delete: int):
+        select_fields = ["message_id", "content", "content_embed"]
+        _index_name = index_name(uid)
+        res = settings.msgStoreConn.get_forgotten_messages(select_fields, _index_name, memory_id)
+        message_list = settings.msgStoreConn.get_fields(res, select_fields)
+        current_size = 0
+        ids_to_remove = []
+        for message in message_list:
+            if current_size < size_to_delete:
+                current_size += cls.calculate_message_size(message)
+                ids_to_remove.append(message["message_id"])
+            else:
+                return ids_to_remove, current_size
+        if current_size >= size_to_delete:
+            return ids_to_remove, current_size
+
+        order_by = OrderByExpr()
+        order_by.asc("valid_at")
+        res = settings.msgStoreConn.search(
+            select_fields=["memory_id", "content", "content_embed"],
+            highlight_fields=[],
+            condition={},
+            match_expressions=[],
+            order_by=order_by,
+            offset=0, limit=2000,
+            index_names=[_index_name], memory_ids=[memory_id], agg_fields=[]
+        )
+        docs = settings.msgStoreConn.get_fields(res, select_fields)
+        for doc in docs.values():
+            if current_size < size_to_delete:
+                current_size += cls.calculate_message_size(doc)
+                ids_to_remove.append(doc["memory_id"])
+            else:
+                return ids_to_remove, current_size
+        return ids_to_remove, current_size
+
+    @classmethod
+    def get_by_message_id(cls, memory_id: str, message_id: int, uid: str):
+        index = index_name(uid)
+        doc_id = f'{memory_id}_{message_id}'
+        return settings.msgStoreConn.get(doc_id, index, [memory_id])
+
+    @classmethod
+    def get_max_message_id(cls, uid_list: List[str], memory_ids: List[str]):
+        order_by = OrderByExpr()
+        order_by.desc("message_id")
+        index_names = [index_name(uid) for uid in uid_list]
+        res = settings.msgStoreConn.search(
+            select_fields=["message_id"],
+            highlight_fields=[],
+            condition={},
+            match_expressions=[],
+            order_by=order_by,
+            offset=0, limit=1,
+            index_names=index_names, memory_ids=memory_ids,
+            agg_fields=[], hide_forgotten=False
+        )
+        docs = settings.msgStoreConn.get_fields(res, ["message_id"])
+        if not docs:
+            return 1
+        else:
+            latest_msg = list(docs.values())[0]
+            return int(latest_msg["message_id"])

memory/services/query.py

@@ -0,0 +1,185 @@
+#
+#  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 logging
+import json
+import numpy as np
+from common.query_base import QueryBase
+from common.doc_store.doc_store_base import MatchDenseExpr, MatchTextExpr
+from common.float_utils import get_float
+from rag.nlp import rag_tokenizer, term_weight, synonym
+
+
+def get_vector(txt, emb_mdl, topk=10, similarity=0.1):
+    if isinstance(similarity, str) and len(similarity) > 0:
+        try:
+            similarity = float(similarity)
+        except Exception as e:
+            logging.warning(f"Convert similarity '{similarity}' to float failed: {e}. Using default 0.1")
+            similarity = 0.1
+    qv, _ = emb_mdl.encode_queries(txt)
+    shape = np.array(qv).shape
+    if len(shape) > 1:
+        raise Exception(
+            f"Dealer.get_vector returned array's shape {shape} doesn't match expectation(exact one dimension).")
+    embedding_data = [get_float(v) for v in qv]
+    vector_column_name = f"q_{len(embedding_data)}_vec"
+    return MatchDenseExpr(vector_column_name, embedding_data, 'float', 'cosine', topk, {"similarity": similarity})
+
+
+class MsgTextQuery(QueryBase):
+
+    def __init__(self):
+        self.tw = term_weight.Dealer()
+        self.syn = synonym.Dealer()
+        self.query_fields = [
+            "content"
+        ]
+
+    def question(self, txt, tbl="messages", min_match: float=0.6):
+        original_query = txt
+        txt = MsgTextQuery.add_space_between_eng_zh(txt)
+        txt = re.sub(
+            r"[ :|\r\n\t,，。？?/`!！&^%%()\[\]{}<>]+",
+            " ",
+            rag_tokenizer.tradi2simp(rag_tokenizer.strQ2B(txt.lower())),
+        ).strip()
+        otxt = txt
+        txt = MsgTextQuery.rmWWW(txt)
+
+        if not self.is_chinese(txt):
+            txt = self.rmWWW(txt)
+            tks = rag_tokenizer.tokenize(txt).split()
+            keywords = [t for t in tks if t]
+            tks_w = self.tw.weights(tks, preprocess=False)
+            tks_w = [(re.sub(r"[ \\\"'^]", "", tk), w) for tk, w in tks_w]
+            tks_w = [(re.sub(r"^[a-z0-9]$", "", tk), w) for tk, w in tks_w if tk]
+            tks_w = [(re.sub(r"^[\+-]", "", tk), w) for tk, w in tks_w if tk]
+            tks_w = [(tk.strip(), w) for tk, w in tks_w if tk.strip()]
+            syns = []
+            for tk, w in tks_w[:256]:
+                syn = self.syn.lookup(tk)
+                syn = rag_tokenizer.tokenize(" ".join(syn)).split()
+                keywords.extend(syn)
+                syn = ["\"{}\"^{:.4f}".format(s, w / 4.) for s in syn if s.strip()]
+                syns.append(" ".join(syn))
+
+            q = ["({}^{:.4f}".format(tk, w) + " {})".format(syn) for (tk, w), syn in zip(tks_w, syns) if
+                 tk and not re.match(r"[.^+\(\)-]", tk)]
+            for i in range(1, len(tks_w)):
+                left, right = tks_w[i - 1][0].strip(), tks_w[i][0].strip()
+                if not left or not right:
+                    continue
+                q.append(
+                    '"%s %s"^%.4f'
+                    % (
+                        tks_w[i - 1][0],
+                        tks_w[i][0],
+                        max(tks_w[i - 1][1], tks_w[i][1]) * 2,
+                    )
+                )
+            if not q:
+                q.append(txt)
+            query = " ".join(q)
+            return MatchTextExpr(
+                self.query_fields, query, 100, {"original_query": original_query}
+            ), keywords
+
+        def need_fine_grained_tokenize(tk):
+            if len(tk) < 3:
+                return False
+            if re.match(r"[0-9a-z\.\+#_\*-]+$", tk):
+                return False
+            return True
+
+        txt = self.rmWWW(txt)
+        qs, keywords = [], []
+        for tt in self.tw.split(txt)[:256]:  # .split():
+            if not tt:
+                continue
+            keywords.append(tt)
+            twts = self.tw.weights([tt])
+            syns = self.syn.lookup(tt)
+            if syns and len(keywords) < 32:
+                keywords.extend(syns)
+            logging.debug(json.dumps(twts, ensure_ascii=False))
+            tms = []
+            for tk, w in sorted(twts, key=lambda x: x[1] * -1):
+                sm = (
+                    rag_tokenizer.fine_grained_tokenize(tk).split()
+                    if need_fine_grained_tokenize(tk)
+                    else []
+                )
+                sm = [
+                    re.sub(
+                        r"[ ,\./;'\[\]\\`~!@#$%\^&\*\(\)=\+_<>\?:\"\{\}\|，。；‘’【】、！￥……（）——《》？：“”-]+",
+                        "",
+                        m,
+                    )
+                    for m in sm
+                ]
+                sm = [self.sub_special_char(m) for m in sm if len(m) > 1]
+                sm = [m for m in sm if len(m) > 1]
+
+                if len(keywords) < 32:
+                    keywords.append(re.sub(r"[ \\\"']+", "", tk))
+                    keywords.extend(sm)
+
+                tk_syns = self.syn.lookup(tk)
+                tk_syns = [self.sub_special_char(s) for s in tk_syns]
+                if len(keywords) < 32:
+                    keywords.extend([s for s in tk_syns if s])
+                tk_syns = [rag_tokenizer.fine_grained_tokenize(s) for s in tk_syns if s]
+                tk_syns = [f"\"{s}\"" if s.find(" ") > 0 else s for s in tk_syns]
+
+                if len(keywords) >= 32:
+                    break
+
+                tk = self.sub_special_char(tk)
+                if tk.find(" ") > 0:
+                    tk = '"%s"' % tk
+                if tk_syns:
+                    tk = f"({tk} OR (%s)^0.2)" % " ".join(tk_syns)
+                if sm:
+                    tk = f'{tk} OR "%s" OR ("%s"~2)^0.5' % (" ".join(sm), " ".join(sm))
+                if tk.strip():
+                    tms.append((tk, w))
+
+            tms = " ".join([f"({t})^{w}" for t, w in tms])
+
+            if len(twts) > 1:
+                tms += ' ("%s"~2)^1.5' % rag_tokenizer.tokenize(tt)
+
+            syns = " OR ".join(
+                [
+                    '"%s"'
+                    % rag_tokenizer.tokenize(self.sub_special_char(s))
+                    for s in syns
+                ]
+            )
+            if syns and tms:
+                tms = f"({tms})^5 OR ({syns})^0.7"
+
+            qs.append(tms)
+
+        if qs:
+            query = " OR ".join([f"({t})" for t in qs if t])
+            if not query:
+                query = otxt
+            return MatchTextExpr(
+                self.query_fields, query, 100, {"minimum_should_match": min_match, "original_query": original_query}
+            ), keywords
+        return None, keywords

memory/utils/es_conn.py

@@ -0,0 +1,494 @@
+#
+#  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 time
+
+import copy
+from elasticsearch import NotFoundError
+from elasticsearch_dsl import UpdateByQuery, Q, Search
+from elastic_transport import ConnectionTimeout
+from common.decorator import singleton
+from common.doc_store.doc_store_base import MatchExpr, OrderByExpr, MatchTextExpr, MatchDenseExpr, FusionExpr
+from common.doc_store.es_conn_base import ESConnectionBase
+from common.float_utils import get_float
+from common.constants import PAGERANK_FLD, TAG_FLD
+
+ATTEMPT_TIME = 2
+
+
+@singleton
+class ESConnection(ESConnectionBase):
+
+    @staticmethod
+    def convert_field_name(field_name: str) -> str:
+        match field_name:
+            case "message_type":
+                return "message_type_kwd"
+            case "status":
+                return "status_int"
+            case "content":
+                return "content_ltks"
+            case _:
+                return field_name
+
+    @staticmethod
+    def map_message_to_es_fields(message: dict) -> dict:
+        """
+        Map message dictionary fields to Elasticsearch document/Infinity fields.
+
+        :param message: A dictionary containing message details.
+        :return: A dictionary formatted for Elasticsearch/Infinity indexing.
+        """
+        storage_doc = {
+            "id": message.get("id"),
+            "message_id": message["message_id"],
+            "message_type_kwd": message["message_type"],
+            "source_id": message["source_id"],
+            "memory_id": message["memory_id"],
+            "user_id": message["user_id"],
+            "agent_id": message["agent_id"],
+            "session_id": message["session_id"],
+            "valid_at": message["valid_at"],
+            "invalid_at": message["invalid_at"],
+            "forget_at": message["forget_at"],
+            "status_int": 1 if message["status"] else 0,
+            "zone_id": message.get("zone_id", 0),
+            "content_ltks": message["content"],
+            f"q_{len(message['content_embed'])}_vec": message["content_embed"],
+        }
+        return storage_doc
+
+    @staticmethod
+    def get_message_from_es_doc(doc: dict) -> dict:
+        """
+        Convert an Elasticsearch/Infinity document back to a message dictionary.
+
+        :param doc: A dictionary representing the Elasticsearch/Infinity document.
+        :return: A dictionary formatted as a message.
+        """
+        embd_field_name = next((key for key in doc.keys() if re.match(r"q_\d+_vec", key)), None)
+        message = {
+            "message_id": doc["message_id"],
+            "message_type": doc["message_type_kwd"],
+            "source_id": doc["source_id"] if doc["source_id"] else None,
+            "memory_id": doc["memory_id"],
+            "user_id": doc.get("user_id", ""),
+            "agent_id": doc["agent_id"],
+            "session_id": doc["session_id"],
+            "zone_id": doc.get("zone_id", 0),
+            "valid_at": doc["valid_at"],
+            "invalid_at": doc.get("invalid_at", "-"),
+            "forget_at": doc.get("forget_at", "-"),
+            "status": bool(int(doc["status_int"])),
+            "content": doc.get("content_ltks", ""),
+            "content_embed": doc.get(embd_field_name, []) if embd_field_name else [],
+        }
+        if doc.get("id"):
+            message["id"] = doc["id"]
+        return message
+
+    """
+    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],
+            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
+        """
+        if isinstance(index_names, str):
+            index_names = index_names.split(",")
+        assert isinstance(index_names, list) and len(index_names) > 0
+        assert "_id" not in condition
+        bool_query = Q("bool", must=[], must_not=[])
+        if hide_forgotten:
+            # filter not forget
+            bool_query.must_not.append(Q("exists", field="forget_at"))
+
+        condition["memory_id"] = memory_ids
+        for k, v in condition.items():
+            if k == "session_id" and v:
+                bool_query.filter.append(Q("query_string", **{"query": f"*{v}*", "fields": ["session_id"], "analyze_wildcard": True}))
+                continue
+            if not v:
+                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.")
+        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)
+                weights = m.fusion_params["weights"]
+                vector_similarity_weight = get_float(weights.split(",")[1])
+        for m in match_expressions:
+            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)) + "%"
+                bool_query.must.append(Q("query_string", fields=[self.convert_field_name(f) 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)
+                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,
+                          )
+
+        if bool_query and rank_feature:
+            for fld, sc in rank_feature.items():
+                if fld != PAGERANK_FLD:
+                    fld = f"{TAG_FLD}.{fld}"
+                bool_query.should.append(Q("rank_feature", field=fld, linear={}, boost=sc))
+
+        if bool_query:
+            s = s.query(bool_query)
+        for field in highlight_fields:
+            s = s.highlight(field)
+
+        if order_by:
+            orders = list()
+            for field, order in order_by.fields:
+                order = "asc" if order == 0 else "desc"
+                if 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)
+
+        if agg_fields:
+            for fld in agg_fields:
+                s.aggs.bucket(f'aggs_{fld}', 'terms', field=fld, size=1000000)
+
+        if limit > 0:
+            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=index_names,
+                                     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))
+                return res
+            except ConnectionTimeout:
+                self.logger.exception("ES request timeout")
+                self._connect()
+                continue
+            except Exception as e:
+                self.logger.exception(f"ESConnection.search {str(index_names)} query: " + str(q) + str(e))
+                raise e
+
+        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=2000):
+        bool_query = Q("bool", must_not=[])
+        bool_query.must_not.append(Q("term", forget_at=None))
+        bool_query.filter.append(Q("term", memory_id=memory_id))
+        # from old to new
+        order_by = OrderByExpr()
+        order_by.asc("forget_at")
+        # build search
+        s = Search()
+        s = s.query(bool_query)
+        s = s.sort(order_by)
+        s = s[:limit]
+        q = s.to_dict()
+        # search
+        for i in range(ATTEMPT_TIME):
+            try:
+                res = self.es.search(index=index_name, body=q, timeout="600s", 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_name)} res: " + str(res))
+                return res
+            except ConnectionTimeout:
+                self.logger.exception("ES request timeout")
+                self._connect()
+                continue
+            except Exception as e:
+                self.logger.exception(f"ESConnection.search {str(index_name)} query: " + str(q) + str(e))
+                raise e
+
+        self.logger.error(f"ESConnection.search timeout for {ATTEMPT_TIME} times!")
+        raise Exception("ESConnection.search timeout.")
+
+    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, )
+                if str(res.get("timed_out", "")).lower() == "true":
+                    raise Exception("Es Timeout.")
+                message = res["_source"]
+                message["id"] = doc_id
+                return self.get_message_from_es_doc(message)
+            except NotFoundError:
+                return None
+            except Exception as e:
+                self.logger.exception(f"ESConnection.get({doc_id}) got exception")
+                raise e
+        self.logger.error(f"ESConnection.get timeout for {ATTEMPT_TIME} times!")
+        raise Exception("ESConnection.get timeout.")
+
+    def insert(self, documents: list[dict], index_name: str, memory_id: str = None) -> list[str]:
+        # Refers to https://www.elastic.co/guide/en/elasticsearch/reference/current/docs-bulk.html
+        operations = []
+        for d in documents:
+            assert "_id" not in d
+            assert "id" in d
+            d_copy_raw = copy.deepcopy(d)
+            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(d_copy)
+        res = []
+        for _ in range(ATTEMPT_TIME):
+            try:
+                res = []
+                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
+
+                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 ConnectionTimeout:
+                self.logger.exception("ES request timeout")
+                time.sleep(3)
+                self._connect()
+                continue
+            except Exception as e:
+                res.append(str(e))
+                self.logger.warning("ESConnection.insert got exception: " + str(e))
+
+        return res
+
+    def update(self, condition: dict, new_value: dict, index_name: str, memory_id: str) -> bool:
+        doc = copy.deepcopy(new_value)
+        update_dict = {self.convert_field_name(k): v for k, v in doc.items()}
+        update_dict.pop("id", None)
+        condition_dict = {self.convert_field_name(k): v for k, v in condition.items()}
+        condition_dict["memory_id"] = memory_id
+        if "id" in condition_dict and isinstance(condition_dict["id"], str):
+            # update specific single document
+            message_id = condition_dict["id"]
+            for i in range(ATTEMPT_TIME):
+                for k in update_dict.keys():
+                    if "feas" != k.split("_")[-1]:
+                        continue
+                    try:
+                        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))
+                    break
+            return False
+
+        # update unspecific maybe-multiple documents
+        bool_query = Q("bool")
+        for k, v in condition_dict.items():
+            if not isinstance(k, str) or not v:
+                continue
+            if k == "exists":
+                bool_query.filter.append(Q("exists", field=v))
+                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.")
+        scripts = []
+        params = {}
+        for k, v in update_dict.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 != "status_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=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)
+        ubq = ubq.params(conflicts="proceed")
+        for _ in range(ATTEMPT_TIME):
+            try:
+                _ = ubq.execute()
+                return True
+            except ConnectionTimeout:
+                self.logger.exception("ES request timeout")
+                time.sleep(3)
+                self._connect()
+                continue
+            except Exception as e:
+                self.logger.error("ESConnection.update got exception: " + str(e) + "\n".join(scripts))
+                break
+        return False
+
+    def delete(self, condition: dict, index_name: str, memory_id: str) -> int:
+        assert "_id" not in condition
+        condition_dict = {self.convert_field_name(k): v for k, v in condition.items()}
+        condition_dict["memory_id"] = memory_id
+        if "id" in condition_dict:
+            message_ids = condition_dict["id"]
+            if not isinstance(message_ids, list):
+                message_ids = [message_ids]
+            if not message_ids:  # when message_ids is empty, delete all
+                qry = Q("match_all")
+            else:
+                qry = Q("ids", values=message_ids)
+        else:
+            qry = Q("bool")
+            for k, v in condition_dict.items():
+                if k == "exists":
+                    qry.filter.append(Q("exists", field=v))
+
+                elif k == "must_not":
+                    if isinstance(v, dict):
+                        for kk, vv in v.items():
+                            if kk == "exists":
+                                qry.must_not.append(Q("exists", field=vv))
+
+                elif isinstance(v, list):
+                    qry.must.append(Q("terms", **{k: v}))
+                elif isinstance(v, str) or isinstance(v, int):
+                    qry.must.append(Q("term", **{k: v}))
+                else:
+                    raise Exception("Condition value must be int, str or list.")
+        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)
+                return res["deleted"]
+            except ConnectionTimeout:
+                self.logger.exception("ES request timeout")
+                time.sleep(3)
+                self._connect()
+                continue
+            except Exception as e:
+                self.logger.warning("ESConnection.delete got exception: " + str(e))
+                if re.search(r"(not_found)", str(e), re.IGNORECASE):
+                    return 0
+        return 0
+
+    """
+    Helper functions for search result
+    """
+
+    def get_fields(self, res, fields: list[str]) -> dict[str, dict]:
+        res_fields = {}
+        if not fields:
+            return {}
+        for doc in self._get_source(res):
+            message = self.get_message_from_es_doc(doc)
+            m = {}
+            for n, v in message.items():
+                if n not in fields:
+                    continue
+                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)):
+                    m[n] = v
+                    continue
+                if not isinstance(v, str):
+                    m[n] = str(v)
+                else:
+                    m[n] = v
+
+            if m:
+                res_fields[doc["id"]] = m
+        return res_fields

memory/utils/infinity_conn.py

@@ -0,0 +1,467 @@
+#
+#  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.time_utils import date_string_to_timestamp
+
+
+@singleton
+class InfinityConnection(InfinityConnectionBase):
+    def __init__(self):
+        super().__init__()
+        self.mapping_file_name = "message_infinity_mapping.json"
+
+    """
+    Dataframe and fields convert
+    """
+
+    @staticmethod
+    def field_keyword(field_name: str):
+        # no keywords right now
+        return False
+
+    @staticmethod
+    def convert_message_field_to_infinity(field_name: str):
+        match field_name:
+            case "message_type":
+                return "message_type_kwd"
+            case "status":
+                return "status_int"
+            case _:
+                return field_name
+
+    @staticmethod
+    def convert_infinity_field_to_message(field_name: str):
+        if field_name.startswith("message_type"):
+            return "message_type"
+        if field_name.startswith("status"):
+            return "status"
+        if re.match(r"q_\d+_vec", field_name):
+            return "content_embed"
+        return field_name
+
+    def convert_select_fields(self, output_fields: list[str]) -> list[str]:
+        return list({self.convert_message_field_to_infinity(f) for f in output_fields})
+
+    @staticmethod
+    def convert_matching_field(field_weight_str: str) -> str:
+        tokens = field_weight_str.split("^")
+        field = tokens[0]
+        if field == "content":
+            field = "content@ft_contentm_rag_fine"
+        tokens[0] = field
+        return "^".join(tokens)
+
+    @staticmethod
+    def convert_condition_and_order_field(field_name: str):
+        match field_name:
+            case "message_type":
+                return "message_type_kwd"
+            case "status":
+                return "status_int"
+            case "valid_at":
+                return "valid_at_flt"
+            case "invalid_at":
+                return "invalid_at_flt"
+            case "forget_at":
+                return "forget_at_flt"
+            case _:
+                return field_name
+
+    """
+    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],
+        memory_ids: list[str],
+        agg_fields: list[str] | None = None,
+        rank_feature: dict | None = None,
+        hide_forgotten: bool = True,
+    ) -> 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()
+        db_instance = inf_conn.get_database(self.dbName)
+        df_list = list()
+        table_list = list()
+        if hide_forgotten:
+            condition.update({"must_not": {"exists": "forget_at_flt"}})
+        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 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 != "_score"]
+        if limit <= 0:
+            # ElasticSearch default limit is 10000
+            limit = 10000
+
+        # Prepare expressions common to all tables
+        filter_cond = None
+        filter_fulltext = ""
+        if condition:
+            condition_dict = {self.convert_condition_and_order_field(k): v for k, v in condition.items()}
+            table_found = False
+            for indexName in index_names:
+                for mem_id in memory_ids:
+                    table_name = f"{indexName}_{mem_id}"
+                    try:
+                        filter_cond = self.equivalent_condition_to_str(condition_dict, 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 memoryIds {memory_ids}")
+                return pd.DataFrame(), 0
+
+        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})
+                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):
+                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:
+                order_field_name = self.convert_condition_and_order_field(order_field[0])
+                if order_field[1] == 0:
+                    order_by_expr_list.append((order_field_name, SortType.Asc))
+                else:
+                    order_by_expr_list.append((order_field_name, SortType.Desc))
+
+        total_hits_count = 0
+        # Scatter search tables and gather the results
+        for indexName in index_names:
+            for memory_id in memory_ids:
+                table_name = f"{indexName}_{memory_id}"
+                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)
+                            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)
+                mem_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(mem_res)}")
+                df_list.append(mem_res)
+        self.connPool.release_conn(inf_conn)
+        res = self.concat_dataframes(df_list, output)
+        if match_expressions:
+            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
+
+    def get_forgotten_messages(self, select_fields: list[str], index_name: str, memory_id: str, limit: int=2000):
+        condition = {"memory_id": memory_id, "exists": "forget_at_flt"}
+        order_by = OrderByExpr()
+        order_by.asc("forget_at_flt")
+        # query
+        inf_conn = self.connPool.get_conn()
+        db_instance = inf_conn.get_database(self.dbName)
+        table_name = f"{index_name}_{memory_id}"
+        table_instance = db_instance.get_table(table_name)
+        output_fields = [self.convert_message_field_to_infinity(f) for f in select_fields]
+        builder = table_instance.output(output_fields)
+        filter_cond = self.equivalent_condition_to_str(condition, db_instance.get_table(table_name))
+        builder.filter(filter_cond)
+        order_by_expr_list = list()
+        if order_by.fields:
+            for order_field in order_by.fields:
+                order_field_name = self.convert_condition_and_order_field(order_field[0])
+                if order_field[1] == 0:
+                    order_by_expr_list.append((order_field_name, SortType.Asc))
+                else:
+                    order_by_expr_list.append((order_field_name, SortType.Desc))
+        builder.sort(order_by_expr_list)
+        builder.offset(0).limit(limit)
+        mem_res, _ = builder.option({"total_hits_count": True}).to_df()
+        res = self.concat_dataframes(mem_res, output_fields)
+        res.head(limit)
+        self.connPool.release_conn(inf_conn)
+        return res
+
+    def get(self, message_id: str, index_name: str, memory_ids: list[str]) -> dict | None:
+        inf_conn = self.connPool.get_conn()
+        db_instance = inf_conn.get_database(self.dbName)
+        df_list = list()
+        assert isinstance(memory_ids, list)
+        table_list = list()
+        for memoryId in memory_ids:
+            table_name = f"{index_name}_{memoryId}"
+            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 memory isn't created in Infinity. Maybe it is created in other document engine.")
+                continue
+            mem_res, _ = table_instance.output(["*"]).filter(f"id = '{message_id}'").to_df()
+            self.logger.debug(f"INFINITY get table: {str(table_list)}, result: {str(mem_res)}")
+            df_list.append(mem_res)
+        self.connPool.release_conn(inf_conn)
+        res = self.concat_dataframes(df_list, ["id"])
+        fields = set(res.columns.tolist())
+        res_fields = self.get_fields(res, list(fields))
+        return res_fields.get(message_id, None)
+
+    def insert(self, documents: list[dict], index_name: str, memory_id: str = None) -> list[str]:
+        if not documents:
+            return []
+        inf_conn = self.connPool.get_conn()
+        db_instance = inf_conn.get_database(self.dbName)
+        table_name = f"{index_name}_{memory_id}"
+        vector_size = int(len(documents[0]["content_embed"]))
+        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
+            if vector_size == 0:
+                raise ValueError("Cannot infer vector size from documents")
+            self.create_idx(index_name, memory_id, vector_size)
+            table_instance = db_instance.get_table(table_name)
+
+        # embedding fields can't have a default value....
+        embedding_columns = []
+        table_columns = table_instance.show_columns().rows()
+        for n, ty, _, _ in table_columns:
+            r = re.search(r"Embedding\([a-z]+,([0-9]+)\)", ty)
+            if not r:
+                continue
+            embedding_columns.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()):
+                field_name = self.convert_message_field_to_infinity(k)
+                if field_name in ["valid_at", "invalid_at", "forget_at"]:
+                    d[f"{field_name}_flt"] = date_string_to_timestamp(v) if v else 0
+                    if v is None:
+                        d[field_name] = ""
+                elif self.field_keyword(k):
+                    if isinstance(v, list):
+                        d[k] = "###".join(v)
+                    else:
+                        d[k] = v
+                elif k == "memory_id":
+                    if isinstance(d[k], list):
+                        d[k] = d[k][0]  # since d[k] is a list, but we need a str
+                elif field_name == "content_embed":
+                    d[f"q_{vector_size}_vec"] = d["content_embed"]
+                    d.pop("content_embed")
+                else:
+                    d[field_name] = v
+                if k != field_name:
+                    d.pop(k)
+
+            for n, vs in embedding_columns:
+                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)
+        table_instance.insert(docs)
+        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, memory_id: str) -> bool:
+        inf_conn = self.connPool.get_conn()
+        db_instance = inf_conn.get_database(self.dbName)
+        table_name = f"{index_name}_{memory_id}"
+        table_instance = db_instance.get_table(table_name)
+
+        columns = {}
+        if table_instance:
+            for n, ty, de, _ in table_instance.show_columns().rows():
+                columns[n] = (ty, de)
+        condition_dict = {self.convert_condition_and_order_field(k): v for k, v in condition.items()}
+        filter = self.equivalent_condition_to_str(condition_dict, table_instance)
+        update_dict = {self.convert_message_field_to_infinity(k): v for k, v in new_value.items()}
+        date_floats = {}
+        for k, v in update_dict.items():
+            if k in ["valid_at", "invalid_at", "forget_at"]:
+                date_floats[f"{k}_flt"] = date_string_to_timestamp(v) if v else 0
+            elif self.field_keyword(k):
+                if isinstance(v, list):
+                    update_dict[k] = "###".join(v)
+                else:
+                    update_dict[k] = v
+            elif k == "memory_id":
+                if isinstance(update_dict[k], list):
+                    update_dict[k] = update_dict[k][0]  # since d[k] is a list, but we need a str
+            else:
+                update_dict[k] = v
+        if date_floats:
+            update_dict.update(date_floats)
+
+        self.logger.debug(f"INFINITY update table {table_name}, filter {filter}, newValue {new_value}.")
+        table_instance.update(filter, update_dict)
+        self.connPool.release_conn(inf_conn)
+        return True
+
+    """
+    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 = {self.convert_message_field_to_infinity(f) for f in fields_all}
+
+        column_map = {col.lower(): col for col in res.columns}
+        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 self.field_keyword(k):
+                res2[column] = res2[column].apply(lambda v: [kwd for kwd in v.split("###") if kwd])
+            else:
+                pass
+        for column in ["content"]:
+            if column in res2:
+                del res2[column]
+        for column in none_columns:
+            res2[column] = None
+
+        res_dict = res2.set_index("id").to_dict(orient="index")
+        return {_id: {self.convert_infinity_field_to_message(k): v for k, v in doc.items()} for _id, doc in res_dict.items()}

memory/utils/msg_util.py

@@ -0,0 +1,37 @@
+#
+#  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 json
+
+
+def get_json_result_from_llm_response(response_str: str) -> dict:
+    """
+    Parse the LLM response string to extract JSON content.
+    The function looks for the first and last curly braces to identify the JSON part.
+    If parsing fails, it returns an empty dictionary.
+
+    :param response_str: The response string from the LLM.
+    :return: A dictionary parsed from the JSON content in the response.
+    """
+    try:
+        clean_str = response_str.strip()
+        if clean_str.startswith('```json'):
+            clean_str = clean_str[7:]  # Remove the starting ```json
+        if clean_str.endswith('```'):
+            clean_str = clean_str[:-3]  # Remove the ending ```
+
+        return json.loads(clean_str.strip())
+    except (ValueError, json.JSONDecodeError):
+        return {}

memory/utils/prompt_util.py

@@ -0,0 +1,201 @@
+#
+#  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.
+#
+from typing import Optional, List
+
+from common.constants import MemoryType
+from common.time_utils import current_timestamp
+
+class PromptAssembler:
+
+    SYSTEM_BASE_TEMPLATE = """**Memory Extraction Specialist**
+You are an expert at analyzing conversations to extract structured memory.
+
+{type_specific_instructions}
+
+
+**OUTPUT REQUIREMENTS:**
+1. Output MUST be valid JSON
+2. Follow the specified output format exactly
+3. Each extracted item MUST have: content, valid_at, invalid_at
+4. Timestamps in {timestamp_format} format
+5. Only extract memory types specified above
+6. Maximum {max_items} items per type
+"""
+
+    TYPE_INSTRUCTIONS = {
+        MemoryType.SEMANTIC.name.lower(): """
+        **EXTRACT SEMANTIC KNOWLEDGE:**
+        - Universal facts, definitions, concepts, relationships
+        - Time-invariant, generally true information
+        - Examples: "The capital of France is Paris", "Water boils at 100°C"
+
+        **Timestamp Rules for Semantic Knowledge:**
+        - valid_at: When the fact became true (e.g., law enactment, discovery)
+        - 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
+        - Time-bound, person-specific, contextual
+        - Examples: "Yesterday I fixed the bug", "User reported issue last week"
+
+        **Timestamp Rules for Episodic Knowledge:**
+        - valid_at: Event start/occurrence time
+        - 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
+        - Goal-oriented, actionable, often includes conditions
+        - Examples: "To reset password, click...", "Debugging steps: 1)..."
+
+        **Timestamp Rules for Procedural Knowledge:**
+        - valid_at: When procedure becomes valid/effective
+        - invalid_at: When it expires/becomes obsolete or empty if current
+        - For version-specific: use release dates
+        - For best practices: invalid_at = ""
+        """
+    }
+
+    OUTPUT_TEMPLATES = {
+        MemoryType.SEMANTIC.name.lower(): """
+        "semantic": [
+            {
+                "content": "Clear factual statement",
+                "valid_at": "timestamp or empty",
+                "invalid_at": "timestamp or empty"
+            }
+        ]
+        """,
+
+        MemoryType.EPISODIC.name.lower(): """
+        "episodic": [
+            {
+                "content": "Narrative event description",
+                "valid_at": "event start timestamp",
+                "invalid_at": "event end timestamp or empty"
+            }
+        ]
+        """,
+
+        MemoryType.PROCEDURAL.name.lower(): """
+        "procedural": [
+            {
+                "content": "Actionable instructions",
+                "valid_at": "procedure effective timestamp",
+                "invalid_at": "procedure expiration timestamp or empty"
+            }
+        ]
+        """
+    }
+
+    BASE_USER_PROMPT = """
+**CONVERSATION:**
+{conversation}
+
+**CONVERSATION TIME:** {conversation_time}
+**CURRENT TIME:** {current_time}    
+"""
+
+    @classmethod
+    def assemble_system_prompt(cls, config: dict) -> str:
+        types_to_extract = cls._get_types_to_extract(config["memory_type"])
+
+        type_instructions = cls._generate_type_instructions(types_to_extract)
+
+        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)
+        )
+
+        full_prompt += f"\n**REQUIRED OUTPUT FORMAT (JSON):**\n```json\n{{\n{output_format}\n}}\n```\n"
+
+        examples = cls._generate_examples(types_to_extract)
+        if examples:
+            full_prompt += f"\n**EXAMPLES:**\n{examples}\n"
+
+        return full_prompt
+
+    @staticmethod
+    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():
+                types.add(rt)
+        return list(types)
+
+    @classmethod
+    def _generate_type_instructions(cls, types_to_extract: List[str]) -> str:
+        target_types = set(types_to_extract)
+        instructions = [cls.TYPE_INSTRUCTIONS[mt] for mt in target_types]
+        return "\n".join(instructions)
+
+    @classmethod
+    def _generate_output_format(cls, types_to_extract: List[str]) -> str:
+        target_types = set(types_to_extract)
+        output_parts = [cls.OUTPUT_TEMPLATES[mt] for mt in target_types]
+        return ",\n".join(output_parts)
+
+    @staticmethod
+    def _generate_examples(types_to_extract: list[str]) -> str:
+        examples = []
+
+        if MemoryType.SEMANTIC.name.lower() in types_to_extract:
+            examples.append("""
+            **Semantic Example:**
+            Input: "Python lists are mutable and support various operations."
+            Output: {"semantic": [{"content": "Python lists are mutable data structures", "valid_at": "2024-01-15T10:00:00", "invalid_at": ""}]}
+            """)
+
+        if MemoryType.EPISODIC.name.lower() in types_to_extract:
+            examples.append("""
+            **Episodic Example:**
+            Input: "I deployed the new feature yesterday afternoon."
+            Output: {"episodic": [{"content": "User deployed new feature", "valid_at": "2024-01-14T14:00:00", "invalid_at": "2024-01-14T18:00:00"}]}
+            """)
+
+        if MemoryType.PROCEDURAL.name.lower() in types_to_extract:
+            examples.append("""
+            **Procedural Example:**
+            Input: "To debug API errors: 1) Check logs 2) Verify endpoints 3) Test connectivity."
+            Output: {"procedural": [{"content": "API error debugging: 1. Check logs 2. Verify endpoints 3. Test connectivity", "valid_at": "2024-01-15T10:00:00", "invalid_at": ""}]}
+            """)
+
+        return "\n".join(examples)
+
+    @classmethod
+    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",
+            current_time=current_time or current_timestamp(),
+        )
+
+    @classmethod
+    def get_raw_user_prompt(cls):
+        return cls.BASE_USER_PROMPT

rag/benchmark.py

@@ -77,9 +77,9 @@ class Benchmark:
     def init_index(self, vector_size: int):
         if self.initialized_index:
             return
-        if settings.docStoreConn.indexExist(self.index_name, self.kb_id):
-            settings.docStoreConn.deleteIdx(self.index_name, self.kb_id)
-        settings.docStoreConn.createIdx(self.index_name, self.kb_id, vector_size)
+        if settings.docStoreConn.index_exist(self.index_name, self.kb_id):
+            settings.docStoreConn.delete_idx(self.index_name, self.kb_id)
+        settings.docStoreConn.create_idx(self.index_name, self.kb_id, vector_size)
         self.initialized_index = True
 
     def ms_marco_index(self, file_path, index_name):

rag/nlp/query.py

@@ -19,11 +19,12 @@ import json
 import re
 from collections import defaultdict
 
-from rag.utils.doc_store_conn import MatchTextExpr
+from common.query_base import QueryBase
+from common.doc_store.doc_store_base import MatchTextExpr
 from rag.nlp import rag_tokenizer, term_weight, synonym
 
 
-class FulltextQueryer:
+class FulltextQueryer(QueryBase):
     def __init__(self):
         self.tw = term_weight.Dealer()
         self.syn = synonym.Dealer()
@@ -37,64 +38,19 @@ class FulltextQueryer:
             "content_sm_ltks",
         ]
 
-    @staticmethod
-    def sub_special_char(line):
-        return re.sub(r"([:\{\}/\[\]\-\*\"\(\)\|\+~\^])", r"\\\1", line).strip()
-
-    @staticmethod
-    def is_chinese(line):
-        arr = re.split(r"[ \t]+", line)
-        if len(arr) <= 3:
-            return True
-        e = 0
-        for t in arr:
-            if not re.match(r"[a-zA-Z]+$", t):
-                e += 1
-        return e * 1.0 / len(arr) >= 0.7
-
-    @staticmethod
-    def rmWWW(txt):
-        patts = [
-            (
-                r"是*(怎么办|什么样的|哪家|一下|那家|请问|啥样|咋样了|什么时候|何时|何地|何人|是否|是不是|多少|哪里|怎么|哪儿|怎么样|如何|哪些|是啥|啥是|啊|吗|呢|吧|咋|什么|有没有|呀|谁|哪位|哪个)是*",
-                "",
-            ),
-            (r"(^| )(what|who|how|which|where|why)('re|'s)? ", " "),
-            (
-                r"(^| )('s|'re|is|are|were|was|do|does|did|don't|doesn't|didn't|has|have|be|there|you|me|your|my|mine|just|please|may|i|should|would|wouldn't|will|won't|done|go|for|with|so|the|a|an|by|i'm|it's|he's|she's|they|they're|you're|as|by|on|in|at|up|out|down|of|to|or|and|if) ",
-                " ")
-        ]
-        otxt = txt
-        for r, p in patts:
-            txt = re.sub(r, p, txt, flags=re.IGNORECASE)
-        if not txt:
-            txt = otxt
-        return txt
-
-    @staticmethod
-    def add_space_between_eng_zh(txt):
-        # (ENG/ENG+NUM) + ZH
-        txt = re.sub(r'([A-Za-z]+[0-9]+)([\u4e00-\u9fa5]+)', r'\1 \2', txt)
-        # ENG + ZH
-        txt = re.sub(r'([A-Za-z])([\u4e00-\u9fa5]+)', r'\1 \2', txt)
-        # ZH + (ENG/ENG+NUM)
-        txt = re.sub(r'([\u4e00-\u9fa5]+)([A-Za-z]+[0-9]+)', r'\1 \2', txt)
-        txt = re.sub(r'([\u4e00-\u9fa5]+)([A-Za-z])', r'\1 \2', txt)
-        return txt
-
     def question(self, txt, tbl="qa", min_match: float = 0.6):
         original_query = txt
-        txt = FulltextQueryer.add_space_between_eng_zh(txt)
+        txt = self.add_space_between_eng_zh(txt)
         txt = re.sub(
             r"[ :|\r\n\t,，。？?/`!！&^%%()\[\]{}<>]+",
             " ",
             rag_tokenizer.tradi2simp(rag_tokenizer.strQ2B(txt.lower())),
         ).strip()
         otxt = txt
-        txt = FulltextQueryer.rmWWW(txt)
+        txt = self.rmWWW(txt)
 
         if not self.is_chinese(txt):
-            txt = FulltextQueryer.rmWWW(txt)
+            txt = self.rmWWW(txt)
             tks = rag_tokenizer.tokenize(txt).split()
             keywords = [t for t in tks if t]
             tks_w = self.tw.weights(tks, preprocess=False)
@@ -138,7 +94,7 @@ class FulltextQueryer:
                 return False
             return True
 
-        txt = FulltextQueryer.rmWWW(txt)
+        txt = self.rmWWW(txt)
         qs, keywords = [], []
         for tt in self.tw.split(txt)[:256]:  # .split():
             if not tt:
@@ -164,7 +120,7 @@ class FulltextQueryer:
                     )
                     for m in sm
                 ]
-                sm = [FulltextQueryer.sub_special_char(m) for m in sm if len(m) > 1]
+                sm = [self.sub_special_char(m) for m in sm if len(m) > 1]
                 sm = [m for m in sm if len(m) > 1]
 
                 if len(keywords) < 32:
@@ -172,7 +128,7 @@ class FulltextQueryer:
                     keywords.extend(sm)
 
                 tk_syns = self.syn.lookup(tk)
-                tk_syns = [FulltextQueryer.sub_special_char(s) for s in tk_syns]
+                tk_syns = [self.sub_special_char(s) for s in tk_syns]
                 if len(keywords) < 32:
                     keywords.extend([s for s in tk_syns if s])
                 tk_syns = [rag_tokenizer.fine_grained_tokenize(s) for s in tk_syns if s]
@@ -181,7 +137,7 @@ class FulltextQueryer:
                 if len(keywords) >= 32:
                     break
 
-                tk = FulltextQueryer.sub_special_char(tk)
+                tk = self.sub_special_char(tk)
                 if tk.find(" ") > 0:
                     tk = '"%s"' % tk
                 if tk_syns:
@@ -199,7 +155,7 @@ class FulltextQueryer:
             syns = " OR ".join(
                 [
                     '"%s"'
-                    % rag_tokenizer.tokenize(FulltextQueryer.sub_special_char(s))
+                    % rag_tokenizer.tokenize(self.sub_special_char(s))
                     for s in syns
                 ]
             )
@@ -264,10 +220,10 @@ class FulltextQueryer:
         keywords = [f'"{k.strip()}"' for k in keywords]
         for tk, w in sorted(tks_w, key=lambda x: x[1] * -1)[:keywords_topn]:
             tk_syns = self.syn.lookup(tk)
-            tk_syns = [FulltextQueryer.sub_special_char(s) for s in tk_syns]
+            tk_syns = [self.sub_special_char(s) for s in tk_syns]
             tk_syns = [rag_tokenizer.fine_grained_tokenize(s) for s in tk_syns if s]
             tk_syns = [f"\"{s}\"" if s.find(" ") > 0 else s for s in tk_syns]
-            tk = FulltextQueryer.sub_special_char(tk)
+            tk = self.sub_special_char(tk)
             if tk.find(" ") > 0:
                 tk = '"%s"' % tk
             if tk_syns:

rag/nlp/search.py

@@ -24,7 +24,7 @@ from dataclasses import dataclass
 from rag.prompts.generator import relevant_chunks_with_toc
 from rag.nlp import rag_tokenizer, query
 import numpy as np
-from rag.utils.doc_store_conn import DocStoreConnection, MatchDenseExpr, FusionExpr, OrderByExpr
+from common.doc_store.doc_store_base import MatchDenseExpr, FusionExpr, OrderByExpr, DocStoreConnection
 from common.string_utils import remove_redundant_spaces
 from common.float_utils import get_float
 from common.constants import PAGERANK_FLD, TAG_FLD
@@ -155,7 +155,7 @@ class Dealer:
                     kwds.add(kk)
 
         logging.debug(f"TOTAL: {total}")
-        ids = self.dataStore.get_chunk_ids(res)
+        ids = self.dataStore.get_doc_ids(res)
         keywords = list(kwds)
         highlight = self.dataStore.get_highlight(res, keywords, "content_with_weight")
         aggs = self.dataStore.get_aggregation(res, "docnm_kwd")
@@ -545,7 +545,7 @@ class Dealer:
         return res
 
     def all_tags(self, tenant_id: str, kb_ids: list[str], S=1000):
-        if not self.dataStore.indexExist(index_name(tenant_id), kb_ids[0]):
+        if not self.dataStore.index_exist(index_name(tenant_id), kb_ids[0]):
             return []
         res = self.dataStore.search([], [], {}, [], OrderByExpr(), 0, 0, index_name(tenant_id), kb_ids, ["tag_kwd"])
         return self.dataStore.get_aggregation(res, "tag_kwd")

rag/prompts/generator.py

@@ -136,6 +136,19 @@ def kb_prompt(kbinfos, max_tokens, hash_id=False):
     return knowledges
 
 
+def memory_prompt(message_list, max_tokens):
+    used_token_count = 0
+    content_list = []
+    for message in message_list:
+        current_content_tokens = num_tokens_from_string(message["content"])
+        if used_token_count + current_content_tokens > max_tokens * 0.97:
+            logging.warning(f"Not all the retrieval into prompt: {len(content_list)}/{len(message_list)}")
+            break
+        content_list.append(message["content"])
+        used_token_count += current_content_tokens
+    return content_list
+
+
 CITATION_PROMPT_TEMPLATE = load_prompt("citation_prompt")
 CITATION_PLUS_TEMPLATE = load_prompt("citation_plus")
 CONTENT_TAGGING_PROMPT_TEMPLATE = load_prompt("content_tagging_prompt")

rag/svr/task_executor.py

@@ -506,7 +506,7 @@ def build_TOC(task, docs, progress_callback):
 
 def init_kb(row, vector_size: int):
     idxnm = search.index_name(row["tenant_id"])
-    return settings.docStoreConn.createIdx(idxnm, row.get("kb_id", ""), vector_size)
+    return settings.docStoreConn.create_idx(idxnm, row.get("kb_id", ""), vector_size)
 
 
 async def embedding(docs, mdl, parser_config=None, callback=None):

rag/utils/es_conn.py

@@ -14,194 +14,92 @@
 #  limitations under the License.
 #
 
-import logging
 import re
 import json
 import time
-import os
 
 import copy
-from elasticsearch import Elasticsearch, NotFoundError
-from elasticsearch_dsl import UpdateByQuery, Q, Search, Index
+from elasticsearch_dsl import UpdateByQuery, Q, Search
 from elastic_transport import ConnectionTimeout
 from common.decorator import singleton
-from common.file_utils import get_project_base_directory
-from common.misc_utils import convert_bytes
-from rag.utils.doc_store_conn import DocStoreConnection, MatchExpr, OrderByExpr, MatchTextExpr, MatchDenseExpr, \
-    FusionExpr
-from rag.nlp import is_english, rag_tokenizer
+from common.doc_store.doc_store_base import MatchTextExpr, OrderByExpr, MatchExpr, MatchDenseExpr, FusionExpr
+from common.doc_store.es_conn_base import ESConnectionBase
 from common.float_utils import get_float
-from common import settings
 from common.constants import PAGERANK_FLD, TAG_FLD
 
 ATTEMPT_TIME = 2
 
-logger = logging.getLogger('ragflow.es_conn')
-
 
 @singleton
-class ESConnection(DocStoreConnection):
-    def __init__(self):
-        self.info = {}
-        logger.info(f"Use Elasticsearch {settings.ES['hosts']} as the doc engine.")
-        for _ in range(ATTEMPT_TIME):
-            try:
-                if self._connect():
-                    break
-            except Exception as e:
-                logger.warning(f"{str(e)}. Waiting Elasticsearch {settings.ES['hosts']} to be healthy.")
-                time.sleep(5)
-
-        if not self.es.ping():
-            msg = f"Elasticsearch {settings.ES['hosts']} is unhealthy in 120s."
-            logger.error(msg)
-            raise Exception(msg)
-        v = self.info.get("version", {"number": "8.11.3"})
-        v = v["number"].split(".")[0]
-        if int(v) < 8:
-            msg = f"Elasticsearch version must be greater than or equal to 8, current version: {v}"
-            logger.error(msg)
-            raise Exception(msg)
-        fp_mapping = os.path.join(get_project_base_directory(), "conf", "mapping.json")
-        if not os.path.exists(fp_mapping):
-            msg = f"Elasticsearch mapping file not found at {fp_mapping}"
-            logger.error(msg)
-            raise Exception(msg)
-        self.mapping = json.load(open(fp_mapping, "r"))
-        logger.info(f"Elasticsearch {settings.ES['hosts']} is healthy.")
-
-    def _connect(self):
-        self.es = Elasticsearch(
-            settings.ES["hosts"].split(","),
-            basic_auth=(settings.ES["username"], settings.ES[
-                "password"]) if "username" in settings.ES and "password" in settings.ES else None,
-            verify_certs= settings.ES.get("verify_certs", False),
-            timeout=600 )
-        if self.es:
-            self.info = self.es.info()
-            return True
-        return False
-
-    """
-    Database operations
-    """
-
-    def dbType(self) -> str:
-        return "elasticsearch"
-
-    def health(self) -> dict:
-        health_dict = dict(self.es.cluster.health())
-        health_dict["type"] = "elasticsearch"
-        return health_dict
-
-    """
-    Table operations
-    """
-
-    def createIdx(self, indexName: str, knowledgebaseId: str, vectorSize: int):
-        if self.indexExist(indexName, knowledgebaseId):
-            return True
-        try:
-            from elasticsearch.client import IndicesClient
-            return IndicesClient(self.es).create(index=indexName,
-                                                 settings=self.mapping["settings"],
-                                                 mappings=self.mapping["mappings"])
-        except Exception:
-            logger.exception("ESConnection.createIndex error %s" % (indexName))
-
-    def deleteIdx(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.es.indices.delete(index=indexName, allow_no_indices=True)
-        except NotFoundError:
-            pass
-        except Exception:
-            logger.exception("ESConnection.deleteIdx error %s" % (indexName))
-
-    def indexExist(self, indexName: str, knowledgebaseId: str = None) -> bool:
-        s = Index(indexName, self.es)
-        for i in range(ATTEMPT_TIME):
-            try:
-                return s.exists()
-            except ConnectionTimeout:
-                logger.exception("ES request timeout")
-                time.sleep(3)
-                self._connect()
-                continue
-            except Exception as e:
-                logger.exception(e)
-                break
-        return False
+class ESConnection(ESConnectionBase):
 
     """
     CRUD operations
     """
 
     def search(
-            self, selectFields: list[str],
-            highlightFields: list[str],
+            self, select_fields: list[str],
+            highlight_fields: list[str],
             condition: dict,
-            matchExprs: list[MatchExpr],
-            orderBy: OrderByExpr,
+            match_expressions: list[MatchExpr],
+            order_by: OrderByExpr,
             offset: int,
             limit: int,
-            indexNames: str | list[str],
-            knowledgebaseIds: list[str],
-            aggFields: list[str] = [],
+            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
         """
-        if isinstance(indexNames, str):
-            indexNames = indexNames.split(",")
-        assert isinstance(indexNames, list) and len(indexNames) > 0
+        if isinstance(index_names, str):
+            index_names = index_names.split(",")
+        assert isinstance(index_names, list) and len(index_names) > 0
         assert "_id" not in condition
 
-        bqry = Q("bool", must=[])
-        condition["kb_id"] = knowledgebaseIds
+        bool_query = Q("bool", must=[])
+        condition["kb_id"] = knowledgebase_ids
         for k, v in condition.items():
             if k == "available_int":
                 if v == 0:
-                    bqry.filter.append(Q("range", available_int={"lt": 1}))
+                    bool_query.filter.append(Q("range", available_int={"lt": 1}))
                 else:
-                    bqry.filter.append(
+                    bool_query.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}))
+                bool_query.filter.append(Q("terms", **{k: v}))
             elif isinstance(v, str) or isinstance(v, int):
-                bqry.filter.append(Q("term", **{k: v}))
+                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.")
 
         s = Search()
         vector_similarity_weight = 0.5
-        for m in matchExprs:
+        for m in match_expressions:
             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)
+                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 matchExprs:
+        for m in match_expressions:
             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,
+                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))
-                bqry.boost = 1.0 - vector_similarity_weight
+                bool_query.boost = 1.0 - vector_similarity_weight
 
             elif isinstance(m, MatchDenseExpr):
-                assert (bqry is not None)
+                assert (bool_query is not None)
                 similarity = 0.0
                 if "similarity" in m.extra_options:
                     similarity = m.extra_options["similarity"]
@@ -209,24 +107,24 @@ class ESConnection(DocStoreConnection):
                           m.topn,
                           m.topn * 2,
                           query_vector=list(m.embedding_data),
-                          filter=bqry.to_dict(),
+                          filter=bool_query.to_dict(),
                           similarity=similarity,
                           )
 
-        if bqry and rank_feature:
+        if bool_query 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))
+                bool_query.should.append(Q("rank_feature", field=fld, linear={}, boost=sc))
 
-        if bqry:
-            s = s.query(bqry)
-        for field in highlightFields:
+        if bool_query:
+            s = s.query(bool_query)
+        for field in highlight_fields:
             s = s.highlight(field)
 
-        if orderBy:
+        if order_by:
             orders = list()
-            for field, order in orderBy.fields:
+            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",
@@ -237,19 +135,19 @@ class ESConnection(DocStoreConnection):
                     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 agg_fields:
+            for fld in agg_fields:
+                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"ESConnection.search {str(indexNames)} query: " + json.dumps(q))
+        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=indexNames,
+                res = self.es.search(index=index_names,
                                      body=q,
                                      timeout="600s",
                                      # search_type="dfs_query_then_fetch",
@@ -257,55 +155,37 @@ class ESConnection(DocStoreConnection):
                                      _source=True)
                 if str(res.get("timed_out", "")).lower() == "true":
                     raise Exception("Es Timeout.")
-                logger.debug(f"ESConnection.search {str(indexNames)} res: " + str(res))
+                self.logger.debug(f"ESConnection.search {str(index_names)} res: " + str(res))
                 return res
             except ConnectionTimeout:
-                logger.exception("ES request timeout")
+                self.logger.exception("ES request timeout")
                 self._connect()
                 continue
             except Exception as e:
-                logger.exception(f"ESConnection.search {str(indexNames)} query: " + str(q) + str(e))
+                self.logger.exception(f"ESConnection.search {str(index_names)} query: " + str(q) + str(e))
                 raise e
 
-        logger.error(f"ESConnection.search timeout for {ATTEMPT_TIME} times!")
+        self.logger.error(f"ESConnection.search timeout for {ATTEMPT_TIME} times!")
         raise Exception("ESConnection.search timeout.")
 
-    def get(self, chunkId: str, indexName: str, knowledgebaseIds: list[str]) -> dict | None:
-        for i in range(ATTEMPT_TIME):
-            try:
-                res = self.es.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"ESConnection.get({chunkId}) got exception")
-                raise e
-        logger.error(f"ESConnection.get timeout for {ATTEMPT_TIME} times!")
-        raise Exception("ESConnection.get timeout.")
-
-    def insert(self, documents: list[dict], indexName: str, knowledgebaseId: str = None) -> list[str]:
+    def insert(self, documents: list[dict], index_name: str, knowledgebase_id: str = None) -> list[str]:
         # Refers to https://www.elastic.co/guide/en/elasticsearch/reference/current/docs-bulk.html
         operations = []
         for d in documents:
             assert "_id" not in d
             assert "id" in d
             d_copy = copy.deepcopy(d)
-            d_copy["kb_id"] = knowledgebaseId
+            d_copy["kb_id"] = knowledgebase_id
             meta_id = d_copy.pop("id", "")
             operations.append(
-                {"index": {"_index": indexName, "_id": meta_id}})
+                {"index": {"_index": index_name, "_id": meta_id}})
             operations.append(d_copy)
 
         res = []
         for _ in range(ATTEMPT_TIME):
             try:
                 res = []
-                r = self.es.bulk(index=(indexName), operations=operations,
+                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
@@ -316,58 +196,58 @@ class ESConnection(DocStoreConnection):
                             res.append(str(item[action]["_id"]) + ":" + str(item[action]["error"]))
                 return res
             except ConnectionTimeout:
-                logger.exception("ES request timeout")
+                self.logger.exception("ES request timeout")
                 time.sleep(3)
                 self._connect()
                 continue
             except Exception as e:
                 res.append(str(e))
-                logger.warning("ESConnection.insert got exception: " + str(e))
+                self.logger.warning("ESConnection.insert got exception: " + str(e))
 
         return res
 
-    def update(self, condition: dict, newValue: dict, indexName: str, knowledgebaseId: str) -> bool:
-        doc = copy.deepcopy(newValue)
+    def update(self, condition: dict, new_value: dict, index_name: str, knowledgebase_id: str) -> bool:
+        doc = copy.deepcopy(new_value)
         doc.pop("id", None)
-        condition["kb_id"] = knowledgebaseId
+        condition["kb_id"] = knowledgebase_id
         if "id" in condition and isinstance(condition["id"], str):
             # update specific single document
-            chunkId = condition["id"]
+            chunk_id = condition["id"]
             for i in range(ATTEMPT_TIME):
                 for k in doc.keys():
                     if "feas" != k.split("_")[-1]:
                         continue
                     try:
-                        self.es.update(index=indexName, id=chunkId, script=f"ctx._source.remove(\"{k}\");")
+                        self.es.update(index=index_name, id=chunk_id, script=f"ctx._source.remove(\"{k}\");")
                     except Exception:
-                        logger.exception(f"ESConnection.update(index={indexName}, id={chunkId}, 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:
-                    self.es.update(index=indexName, id=chunkId, doc=doc)
+                    self.es.update(index=index_name, id=chunk_id, doc=doc)
                     return True
                 except Exception as e:
-                    logger.exception(
-                        f"ESConnection.update(index={indexName}, id={chunkId}, 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
 
         # update unspecific maybe-multiple documents
-        bqry = Q("bool")
+        bool_query = 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))
+                bool_query.filter.append(Q("exists", field=v))
                 continue
             if isinstance(v, list):
-                bqry.filter.append(Q("terms", **{k: v}))
+                bool_query.filter.append(Q("terms", **{k: v}))
             elif isinstance(v, str) or isinstance(v, int):
-                bqry.filter.append(Q("term", **{k: v}))
+                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.")
         scripts = []
         params = {}
-        for k, v in newValue.items():
+        for k, v in new_value.items():
             if k == "remove":
                 if isinstance(v, str):
                     scripts.append(f"ctx._source.remove('{v}');")
@@ -397,8 +277,8 @@ class ESConnection(DocStoreConnection):
                 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.es).query(bqry)
+            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)
@@ -409,19 +289,18 @@ class ESConnection(DocStoreConnection):
                 _ = ubq.execute()
                 return True
             except ConnectionTimeout:
-                logger.exception("ES request timeout")
+                self.logger.exception("ES request timeout")
                 time.sleep(3)
                 self._connect()
                 continue
             except Exception as e:
-                logger.error("ESConnection.update got exception: " + str(e) + "\n".join(scripts))
+                self.logger.error("ESConnection.update got exception: " + str(e) + "\n".join(scripts))
                 break
         return False
 
-    def delete(self, condition: dict, indexName: str, knowledgebaseId: str) -> int:
-        qry = None
+    def delete(self, condition: dict, index_name: str, knowledgebase_id: str) -> int:
         assert "_id" not in condition
-        condition["kb_id"] = knowledgebaseId
+        condition["kb_id"] = knowledgebase_id
         if "id" in condition:
             chunk_ids = condition["id"]
             if not isinstance(chunk_ids, list):
@@ -448,21 +327,21 @@ class ESConnection(DocStoreConnection):
                     qry.must.append(Q("term", **{k: v}))
                 else:
                     raise Exception("Condition value must be int, str or list.")
-        logger.debug("ESConnection.delete query: " + json.dumps(qry.to_dict()))
+        self.logger.debug("ESConnection.delete query: " + json.dumps(qry.to_dict()))
         for _ in range(ATTEMPT_TIME):
             try:
                 res = self.es.delete_by_query(
-                    index=indexName,
+                    index=index_name,
                     body=Search().query(qry).to_dict(),
                     refresh=True)
                 return res["deleted"]
             except ConnectionTimeout:
-                logger.exception("ES request timeout")
+                self.logger.exception("ES request timeout")
                 time.sleep(3)
                 self._connect()
                 continue
             except Exception as e:
-                logger.warning("ESConnection.delete got exception: " + str(e))
+                self.logger.warning("ESConnection.delete got exception: " + str(e))
                 if re.search(r"(not_found)", str(e), re.IGNORECASE):
                     return 0
         return 0
@@ -471,27 +350,11 @@ class ESConnection(DocStoreConnection):
     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_chunk_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):
+        for d in self._get_source(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):
@@ -508,124 +371,3 @@ class ESConnection(DocStoreConnection):
             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"ESConnection.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"ESConnection.sql to es: {sql}")
-
-        for i in range(ATTEMPT_TIME):
-            try:
-                res = self.es.sql.query(body={"query": sql, "fetch_size": fetch_size}, format=format,
-                                        request_timeout="2s")
-                return res
-            except ConnectionTimeout:
-                logger.exception("ES request timeout")
-                time.sleep(3)
-                self._connect()
-                continue
-            except Exception as e:
-                logger.exception(f"ESConnection.sql got exception. SQL:\n{sql}")
-                raise Exception(f"SQL error: {e}\n\nSQL: {sql}")
-        logger.error(f"ESConnection.sql timeout for {ATTEMPT_TIME} times!")
-        return None
-
-    def get_cluster_stats(self):
-        """
-        curl -XGET "http://{es_host}/_cluster/stats" -H "kbn-xsrf: reporting" to view raw stats.
-        """
-        raw_stats = self.es.cluster.stats()
-        logger.debug(f"ESConnection.get_cluster_stats: {raw_stats}")
-        try:
-            res = {
-                'cluster_name': raw_stats['cluster_name'],
-                'status': raw_stats['status']
-            }
-            indices_status = raw_stats['indices']
-            res.update({
-                'indices': indices_status['count'],
-                'indices_shards': indices_status['shards']['total']
-            })
-            doc_info = indices_status['docs']
-            res.update({
-                'docs': doc_info['count'],
-                'docs_deleted': doc_info['deleted']
-            })
-            store_info = indices_status['store']
-            res.update({
-                'store_size': convert_bytes(store_info['size_in_bytes']),
-                'total_dataset_size': convert_bytes(store_info['total_data_set_size_in_bytes'])
-            })
-            mappings_info = indices_status['mappings']
-            res.update({
-                'mappings_fields': mappings_info['total_field_count'],
-                'mappings_deduplicated_fields': mappings_info['total_deduplicated_field_count'],
-                'mappings_deduplicated_size': convert_bytes(mappings_info['total_deduplicated_mapping_size_in_bytes'])
-            })
-            node_info = raw_stats['nodes']
-            res.update({
-                'nodes': node_info['count']['total'],
-                'nodes_version': node_info['versions'],
-                'os_mem': convert_bytes(node_info['os']['mem']['total_in_bytes']),
-                'os_mem_used': convert_bytes(node_info['os']['mem']['used_in_bytes']),
-                'os_mem_used_percent': node_info['os']['mem']['used_percent'],
-                'jvm_versions': node_info['jvm']['versions'][0]['vm_version'],
-                'jvm_heap_used': convert_bytes(node_info['jvm']['mem']['heap_used_in_bytes']),
-                'jvm_heap_max': convert_bytes(node_info['jvm']['mem']['heap_max_in_bytes'])
-            })
-            return res
-
-        except Exception as e:
-            logger.exception(f"ESConnection.get_cluster_stats: {e}")
-            return None

rag/utils/infinity_conn.py

@@ -14,365 +14,125 @@
 #  limitations under the License.
 #
 
-import logging
-import os
 import re
 import json
-import time
 import copy
-import infinity
-from infinity.common import ConflictType, InfinityException, SortType
-from infinity.index import IndexInfo, IndexType
-from infinity.connection_pool import ConnectionPool
+from infinity.common import InfinityException, SortType
 from infinity.errors import ErrorCode
 from common.decorator import singleton
 import pandas as pd
-from common.file_utils import get_project_base_directory
-from rag.nlp import is_english
 from common.constants import PAGERANK_FLD, TAG_FLD
-from common import settings
-from rag.utils.doc_store_conn import (
-    DocStoreConnection,
-    MatchExpr,
-    MatchTextExpr,
-    MatchDenseExpr,
-    FusionExpr,
-    OrderByExpr,
-)
-
-logger = logging.getLogger("ragflow.infinity_conn")
-
-
-def field_keyword(field_name: str):
-    # Treat "*_kwd" tag-like columns as keyword lists except knowledge_graph_kwd; source_id is also keyword-like.
-    if field_name == "source_id" or (field_name.endswith("_kwd") and field_name not in ["knowledge_graph_kwd", "docnm_kwd", "important_kwd", "question_kwd"]):
-        return True
-    return False
-
-def convert_select_fields(output_fields: list[str]) -> list[str]:
-    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"
-    return list(set(output_fields))
-
-def convert_matching_field(field_weightstr: str) -> str:
-    tokens = field_weightstr.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"
-    tokens[0] = field
-    return "^".join(tokens)
-
-def list2str(lst: str|list, sep: str = " ") -> str:
-    if isinstance(lst, str):
-        return lst
-    return sep.join(lst)
-
-
-def equivalent_condition_to_str(condition: dict, table_instance=None) -> str | None:
-    assert "_id" not in condition
-    clmns = {}
-    if table_instance:
-        for n, ty, de, _ in table_instance.show_columns().rows():
-            clmns[n] = (ty, de)
-
-    def exists(cln):
-        nonlocal clmns
-        assert cln in clmns, f"'{cln}' should be in '{clmns}'."
-        ty, de = clmns[cln]
-        if ty.lower().find("cha"):
-            if not de:
-                de = ""
-            return f" {cln}!='{de}' "
-        return f"{cln}!={de}"
-
-    cond = list()
-    for k, v in condition.items():
-        if not isinstance(k, str) or not v:
-            continue
-        if field_keyword(k):
-            if isinstance(v, list):
-                inCond = list()
-                for item in v:
-                    if isinstance(item, str):
-                        item = item.replace("'", "''")
-                    inCond.append(f"filter_fulltext('{convert_matching_field(k)}', '{item}')")
-                if inCond:
-                    strInCond = " or ".join(inCond)
-                    strInCond = f"({strInCond})"
-                    cond.append(strInCond)
-            else:
-                cond.append(f"filter_fulltext('{convert_matching_field(k)}', '{v}')")
-        elif isinstance(v, list):
-            inCond = list()
-            for item in v:
-                if isinstance(item, str):
-                    item = item.replace("'", "''")
-                    inCond.append(f"'{item}'")
-                else:
-                    inCond.append(str(item))
-            if inCond:
-                strInCond = ", ".join(inCond)
-                strInCond = f"{k} IN ({strInCond})"
-                cond.append(strInCond)
-        elif k == "must_not":
-            if isinstance(v, dict):
-                for kk, vv in v.items():
-                    if kk == "exists":
-                        cond.append("NOT (%s)" % exists(vv))
-        elif isinstance(v, str):
-            cond.append(f"{k}='{v}'")
-        elif k == "exists":
-            cond.append(exists(v))
-        else:
-            cond.append(f"{k}={str(v)}")
-    return " AND ".join(cond) if cond else "1=1"
-
-
-def concat_dataframes(df_list: list[pd.DataFrame], selectFields: list[str]) -> pd.DataFrame:
-    df_list2 = [df for df in df_list if not df.empty]
-    if df_list2:
-        return pd.concat(df_list2, axis=0).reset_index(drop=True)
-
-    schema = []
-    for field_name in selectFields:
-        if field_name == "score()":  # Workaround: fix schema is changed to score()
-            schema.append("SCORE")
-        elif field_name == "similarity()":  # Workaround: fix schema is changed to similarity()
-            schema.append("SIMILARITY")
-        else:
-            schema.append(field_name)
-    return pd.DataFrame(columns=schema)
+from common.doc_store.doc_store_base import MatchExpr, MatchTextExpr, MatchDenseExpr, FusionExpr, OrderByExpr
+from common.doc_store.infinity_conn_base import InfinityConnectionBase
 
 
 @singleton
-class InfinityConnection(DocStoreConnection):
-    def __init__(self):
-        self.dbName = settings.INFINITY.get("db_name", "default_db")
-        infinity_uri = settings.INFINITY["uri"]
-        if ":" in infinity_uri:
-            host, port = infinity_uri.split(":")
-            infinity_uri = infinity.common.NetworkAddress(host, int(port))
-        self.connPool = None
-        logger.info(f"Use Infinity {infinity_uri} as the doc engine.")
-        for _ in range(24):
-            try:
-                connPool = ConnectionPool(infinity_uri, max_size=4)
-                inf_conn = connPool.get_conn()
-                res = inf_conn.show_current_node()
-                if res.error_code == ErrorCode.OK and res.server_status in ["started", "alive"]:
-                    self._migrate_db(inf_conn)
-                    self.connPool = connPool
-                    connPool.release_conn(inf_conn)
-                    break
-                connPool.release_conn(inf_conn)
-                logger.warn(f"Infinity status: {res.server_status}. Waiting Infinity {infinity_uri} to be healthy.")
-                time.sleep(5)
-            except Exception as e:
-                logger.warning(f"{str(e)}. Waiting Infinity {infinity_uri} to be healthy.")
-                time.sleep(5)
-        if self.connPool is None:
-            msg = f"Infinity {infinity_uri} is unhealthy in 120s."
-            logger.error(msg)
-            raise Exception(msg)
-        logger.info(f"Infinity {infinity_uri} is healthy.")
-
-    def _migrate_db(self, inf_conn):
-        inf_db = inf_conn.create_database(self.dbName, ConflictType.Ignore)
-        fp_mapping = os.path.join(get_project_base_directory(), "conf", "infinity_mapping.json")
-        if not os.path.exists(fp_mapping):
-            raise Exception(f"Mapping file not found at {fp_mapping}")
-        schema = json.load(open(fp_mapping))
-        table_names = inf_db.list_tables().table_names
-        for table_name in table_names:
-            inf_table = inf_db.get_table(table_name)
-            index_names = inf_table.list_indexes().index_names
-            if "q_vec_idx" not in index_names:
-                # Skip tables not created by me
-                continue
-            column_names = inf_table.show_columns()["name"]
-            column_names = set(column_names)
-            for field_name, field_info in schema.items():
-                if field_name in column_names:
-                    continue
-                res = inf_table.add_columns({field_name: field_info})
-                assert res.error_code == infinity.ErrorCode.OK
-                logger.info(f"INFINITY added following column to table {table_name}: {field_name} {field_info}")
-                if field_info["type"] != "varchar" or "analyzer" not in field_info:
-                    continue
-                analyzers = field_info["analyzer"]
-                if isinstance(analyzers, str):
-                    analyzers = [analyzers]
-                for analyzer in analyzers:
-                    inf_table.create_index(
-                        f"ft_{re.sub(r'[^a-zA-Z0-9]', '_', field_name)}_{re.sub(r'[^a-zA-Z0-9]', '_', analyzer)}",
-                        IndexInfo(field_name, IndexType.FullText, {"ANALYZER": analyzer}),
-                        ConflictType.Ignore,
-                    )
+class InfinityConnection(InfinityConnectionBase):
 
     """
-    Database operations
+    Dataframe and fields convert
     """
 
-    def dbType(self) -> str:
-        return "infinity"
-
-    def health(self) -> dict:
-        """
-        Return the health status of the database.
-        """
-        inf_conn = self.connPool.get_conn()
-        res = inf_conn.show_current_node()
-        self.connPool.release_conn(inf_conn)
-        res2 = {
-            "type": "infinity",
-            "status": "green" if res.error_code == 0 and res.server_status in ["started", "alive"] else "red",
-            "error": res.error_msg,
-        }
-        return res2
-
-    """
-    Table operations
-    """
-
-    def createIdx(self, indexName: str, knowledgebaseId: str, vectorSize: int):
-        table_name = f"{indexName}_{knowledgebaseId}"
-        inf_conn = self.connPool.get_conn()
-        inf_db = inf_conn.create_database(self.dbName, ConflictType.Ignore)
-
-        fp_mapping = os.path.join(get_project_base_directory(), "conf", "infinity_mapping.json")
-        if not os.path.exists(fp_mapping):
-            raise Exception(f"Mapping file not found at {fp_mapping}")
-        schema = json.load(open(fp_mapping))
-        vector_name = f"q_{vectorSize}_vec"
-        schema[vector_name] = {"type": f"vector,{vectorSize},float"}
-        inf_table = inf_db.create_table(
-            table_name,
-            schema,
-            ConflictType.Ignore,
-        )
-        inf_table.create_index(
-            "q_vec_idx",
-            IndexInfo(
-                vector_name,
-                IndexType.Hnsw,
-                {
-                    "M": "16",
-                    "ef_construction": "50",
-                    "metric": "cosine",
-                    "encode": "lvq",
-                },
-            ),
-            ConflictType.Ignore,
-        )
-        for field_name, field_info in schema.items():
-            if field_info["type"] != "varchar" or "analyzer" not in field_info:
-                continue
-            analyzers = field_info["analyzer"]
-            if isinstance(analyzers, str):
-                analyzers = [analyzers]
-            for analyzer in analyzers:
-                inf_table.create_index(
-                    f"ft_{re.sub(r'[^a-zA-Z0-9]', '_', field_name)}_{re.sub(r'[^a-zA-Z0-9]', '_', analyzer)}",
-                    IndexInfo(field_name, IndexType.FullText, {"ANALYZER": analyzer}),
-                    ConflictType.Ignore,
-                )
-        self.connPool.release_conn(inf_conn)
-        logger.info(f"INFINITY created table {table_name}, vector size {vectorSize}")
-
-    def deleteIdx(self, indexName: str, knowledgebaseId: str):
-        table_name = f"{indexName}_{knowledgebaseId}"
-        inf_conn = self.connPool.get_conn()
-        db_instance = inf_conn.get_database(self.dbName)
-        db_instance.drop_table(table_name, ConflictType.Ignore)
-        self.connPool.release_conn(inf_conn)
-        logger.info(f"INFINITY dropped table {table_name}")
-
-    def indexExist(self, indexName: str, knowledgebaseId: str) -> bool:
-        table_name = f"{indexName}_{knowledgebaseId}"
-        try:
-            inf_conn = self.connPool.get_conn()
-            db_instance = inf_conn.get_database(self.dbName)
-            _ = db_instance.get_table(table_name)
-            self.connPool.release_conn(inf_conn)
+    @staticmethod
+    def field_keyword(field_name: str):
+        # Treat "*_kwd" tag-like columns as keyword lists except knowledge_graph_kwd; source_id is also keyword-like.
+        if field_name == "source_id" or (
+                field_name.endswith("_kwd") and field_name not in ["knowledge_graph_kwd", "docnm_kwd", "important_kwd",
+                                                                   "question_kwd"]):
             return True
-        except Exception as e:
-            logger.warning(f"INFINITY indexExist {str(e)}")
         return False
 
+    def convert_select_fields(self, output_fields: list[str]) -> list[str]:
+        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"
+        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"
+        tokens[0] = field
+        return "^".join(tokens)
+
+
     """
     CRUD operations
     """
 
     def search(
         self,
-        selectFields: list[str],
-        highlightFields: list[str],
+        select_fields: list[str],
+        highlight_fields: list[str],
         condition: dict,
-        matchExprs: list[MatchExpr],
-        orderBy: OrderByExpr,
+        match_expressions: list[MatchExpr],
+        order_by: OrderByExpr,
         offset: int,
         limit: int,
-        indexNames: str | list[str],
-        knowledgebaseIds: list[str],
-        aggFields: list[str] = [],
+        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(indexNames, str):
-            indexNames = indexNames.split(",")
-        assert isinstance(indexNames, list) and len(indexNames) > 0
+        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()
         db_instance = inf_conn.get_database(self.dbName)
         df_list = list()
         table_list = list()
-        output = selectFields.copy()
-        output = convert_select_fields(output)
-        for essential_field in ["id"] + aggFields:
+        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 matchExprs:
+        for matchExpr in match_expressions:
             if isinstance(matchExpr, MatchTextExpr):
                 score_func = "score()"
                 score_column = "SCORE"
                 break
         if not score_func:
-            for matchExpr in matchExprs:
+            for matchExpr in match_expressions:
                 if isinstance(matchExpr, MatchDenseExpr):
                     score_func = "similarity()"
                     score_column = "SIMILARITY"
                     break
-        if matchExprs:
+        if match_expressions:
             if score_func not in output:
                 output.append(score_func)
             if PAGERANK_FLD not in output:
@@ -387,11 +147,11 @@ class InfinityConnection(DocStoreConnection):
         filter_fulltext = ""
         if condition:
             table_found = False
-            for indexName in indexNames:
-                for kb_id in knowledgebaseIds:
+            for indexName in index_names:
+                for kb_id in knowledgebase_ids:
                     table_name = f"{indexName}_{kb_id}"
                     try:
-                        filter_cond = equivalent_condition_to_str(condition, db_instance.get_table(table_name))
+                        filter_cond = self.equivalent_condition_to_str(condition, db_instance.get_table(table_name))
                         table_found = True
                         break
                     except Exception:
@@ -399,14 +159,14 @@ class InfinityConnection(DocStoreConnection):
                 if table_found:
                     break
             if not table_found:
-                logger.error(f"No valid tables found for indexNames {indexNames} and knowledgebaseIds {knowledgebaseIds}")
+                self.logger.error(f"No valid tables found for indexNames {index_names} and knowledgebaseIds {knowledgebase_ids}")
                 return pd.DataFrame(), 0
 
-        for matchExpr in matchExprs:
+        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 = [convert_matching_field(field) for field in matchExpr.fields]
+                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:
@@ -430,7 +190,7 @@ class InfinityConnection(DocStoreConnection):
                 for k, v in matchExpr.extra_options.items():
                     if not isinstance(v, str):
                         matchExpr.extra_options[k] = str(v)
-                logger.debug(f"INFINITY search MatchTextExpr: {json.dumps(matchExpr.__dict__)}")
+                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})
@@ -441,16 +201,16 @@ class InfinityConnection(DocStoreConnection):
                 if similarity:
                     matchExpr.extra_options["threshold"] = similarity
                     del matchExpr.extra_options["similarity"]
-                logger.debug(f"INFINITY search MatchDenseExpr: {json.dumps(matchExpr.__dict__)}")
+                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"
-                logger.debug(f"INFINITY search FusionExpr: {json.dumps(matchExpr.__dict__)}")
+                self.logger.debug(f"INFINITY search FusionExpr: {json.dumps(matchExpr.__dict__)}")
 
         order_by_expr_list = list()
-        if orderBy.fields:
-            for order_field in orderBy.fields:
+        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:
@@ -458,8 +218,8 @@ class InfinityConnection(DocStoreConnection):
 
         total_hits_count = 0
         # Scatter search tables and gather the results
-        for indexName in indexNames:
-            for knowledgebaseId in knowledgebaseIds:
+        for indexName in index_names:
+            for knowledgebaseId in knowledgebase_ids:
                 table_name = f"{indexName}_{knowledgebaseId}"
                 try:
                     table_instance = db_instance.get_table(table_name)
@@ -467,8 +227,8 @@ class InfinityConnection(DocStoreConnection):
                     continue
                 table_list.append(table_name)
                 builder = table_instance.output(output)
-                if len(matchExprs) > 0:
-                    for matchExpr in matchExprs:
+                if len(match_expressions) > 0:
+                    for matchExpr in match_expressions:
                         if isinstance(matchExpr, MatchTextExpr):
                             fields = ",".join(matchExpr.fields)
                             builder = builder.match_text(
@@ -491,53 +251,52 @@ class InfinityConnection(DocStoreConnection):
                 else:
                     if filter_cond and len(filter_cond) > 0:
                         builder.filter(filter_cond)
-                if orderBy.fields:
+                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"])
-                logger.debug(f"INFINITY search table: {str(table_name)}, result: {str(kb_res)}")
+                self.logger.debug(f"INFINITY search table: {str(table_name)}, result: {str(kb_res)}")
                 df_list.append(kb_res)
         self.connPool.release_conn(inf_conn)
-        res = concat_dataframes(df_list, output)
-        if matchExprs:
+        res = self.concat_dataframes(df_list, output)
+        if match_expressions:
             res["_score"] = res[score_column] + res[PAGERANK_FLD]
             res = res.sort_values(by="_score", ascending=False).reset_index(drop=True)
             res = res.head(limit)
-        logger.debug(f"INFINITY search final result: {str(res)}")
+        self.logger.debug(f"INFINITY search final result: {str(res)}")
         return res, total_hits_count
 
-    def get(self, chunkId: str, indexName: str, knowledgebaseIds: list[str]) -> dict | None:
+    def get(self, chunk_id: str, index_name: str, knowledgebase_ids: list[str]) -> dict | None:
         inf_conn = self.connPool.get_conn()
         db_instance = inf_conn.get_database(self.dbName)
         df_list = list()
-        assert isinstance(knowledgebaseIds, list)
+        assert isinstance(knowledgebase_ids, list)
         table_list = list()
-        for knowledgebaseId in knowledgebaseIds:
-            table_name = f"{indexName}_{knowledgebaseId}"
+        for knowledgebaseId in knowledgebase_ids:
+            table_name = f"{index_name}_{knowledgebaseId}"
             table_list.append(table_name)
-            table_instance = None
             try:
                 table_instance = db_instance.get_table(table_name)
             except Exception:
-                logger.warning(f"Table not found: {table_name}, this dataset isn't created in Infinity. Maybe it is created in other document engine.")
+                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 = '{chunkId}'").to_df()
-            logger.debug(f"INFINITY get table: {str(table_list)}, result: {str(kb_res)}")
+            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)
         self.connPool.release_conn(inf_conn)
-        res = concat_dataframes(df_list, ["id"])
+        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))
-        return res_fields.get(chunkId, None)
+        return res_fields.get(chunk_id, None)
 
-    def insert(self, documents: list[dict], indexName: str, knowledgebaseId: str = None) -> list[str]:
+    def insert(self, documents: list[dict], index_name: str, knowledgebase_id: str = None) -> list[str]:
         inf_conn = self.connPool.get_conn()
         db_instance = inf_conn.get_database(self.dbName)
-        table_name = f"{indexName}_{knowledgebaseId}"
+        table_name = f"{index_name}_{knowledgebase_id}"
         try:
             table_instance = db_instance.get_table(table_name)
         except InfinityException as e:
@@ -553,7 +312,7 @@ class InfinityConnection(DocStoreConnection):
                     break
             if vector_size == 0:
                 raise ValueError("Cannot infer vector size from documents")
-            self.createIdx(indexName, knowledgebaseId, vector_size)
+            self.create_idx(index_name, knowledgebase_id, vector_size)
             table_instance = db_instance.get_table(table_name)
 
         # embedding fields can't have a default value....
@@ -574,12 +333,12 @@ class InfinityConnection(DocStoreConnection):
                     d["docnm"] = v
                 elif k == "title_kwd":
                     if not d.get("docnm_kwd"):
-                        d["docnm"] = list2str(v)
+                        d["docnm"] = self.list2str(v)
                 elif k == "title_sm_tks":
                     if not d.get("docnm_kwd"):
-                        d["docnm"] = list2str(v)
+                        d["docnm"] = self.list2str(v)
                 elif k == "important_kwd":
-                    d["important_keywords"] = list2str(v)
+                    d["important_keywords"] = self.list2str(v)
                 elif k == "important_tks":
                     if not d.get("important_kwd"):
                         d["important_keywords"] = v
@@ -597,11 +356,11 @@ class InfinityConnection(DocStoreConnection):
                     if not d.get("authors_tks"):
                         d["authors"] = v
                 elif k == "question_kwd":
-                    d["questions"] = list2str(v, "\n")
+                    d["questions"] = self.list2str(v, "\n")
                 elif k == "question_tks":
                     if not d.get("question_kwd"):
-                        d["questions"] = list2str(v)
-                elif field_keyword(k):
+                        d["questions"] = self.list2str(v)
+                elif self.field_keyword(k):
                     if isinstance(v, list):
                         d[k] = "###".join(v)
                     else:
@@ -637,15 +396,15 @@ class InfinityConnection(DocStoreConnection):
         # logger.info(f"InfinityConnection.insert {json.dumps(documents)}")
         table_instance.insert(docs)
         self.connPool.release_conn(inf_conn)
-        logger.debug(f"INFINITY inserted into {table_name} {str_ids}.")
+        self.logger.debug(f"INFINITY inserted into {table_name} {str_ids}.")
         return []
 
-    def update(self, condition: dict, newValue: dict, indexName: str, knowledgebaseId: str) -> bool:
+    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()
         db_instance = inf_conn.get_database(self.dbName)
-        table_name = f"{indexName}_{knowledgebaseId}"
+        table_name = f"{index_name}_{knowledgebase_id}"
         table_instance = db_instance.get_table(table_name)
         # if "exists" in condition:
         #    del condition["exists"]
@@ -654,57 +413,57 @@ class InfinityConnection(DocStoreConnection):
         if table_instance:
             for n, ty, de, _ in table_instance.show_columns().rows():
                 clmns[n] = (ty, de)
-        filter = equivalent_condition_to_str(condition, table_instance)
+        filter = self.equivalent_condition_to_str(condition, table_instance)
         removeValue = {}
-        for k, v in list(newValue.items()):
+        for k, v in list(new_value.items()):
             if k == "docnm_kwd":
-                newValue["docnm"] = list2str(v)
+                new_value["docnm"] = self.list2str(v)
             elif k == "title_kwd":
-                if not newValue.get("docnm_kwd"):
-                    newValue["docnm"] = list2str(v)
+                if not new_value.get("docnm_kwd"):
+                    new_value["docnm"] = self.list2str(v)
             elif k == "title_sm_tks":
-                if not newValue.get("docnm_kwd"):
-                    newValue["docnm"] = v
+                if not new_value.get("docnm_kwd"):
+                    new_value["docnm"] = v
             elif k == "important_kwd":
-                newValue["important_keywords"] = list2str(v)
+                new_value["important_keywords"] = self.list2str(v)
             elif k == "important_tks":
-                if not newValue.get("important_kwd"):
-                    newValue["important_keywords"] = v
+                if not new_value.get("important_kwd"):
+                    new_value["important_keywords"] = v
             elif k == "content_with_weight":
-                newValue["content"] = v
+                new_value["content"] = v
             elif k == "content_ltks":
-                if not newValue.get("content_with_weight"):
-                    newValue["content"] = v
+                if not new_value.get("content_with_weight"):
+                    new_value["content"] = v
             elif k == "content_sm_ltks":
-                if not newValue.get("content_with_weight"):
-                    newValue["content"] = v
+                if not new_value.get("content_with_weight"):
+                    new_value["content"] = v
             elif k == "authors_tks":
-                newValue["authors"] = v
+                new_value["authors"] = v
             elif k == "authors_sm_tks":
-                if not newValue.get("authors_tks"):
-                    newValue["authors"] = v
+                if not new_value.get("authors_tks"):
+                    new_value["authors"] = v
             elif k == "question_kwd":
-                newValue["questions"] = "\n".join(v)
+                new_value["questions"] = "\n".join(v)
             elif k == "question_tks":
-                if not newValue.get("question_kwd"):
-                    newValue["questions"] = list2str(v)
-            elif field_keyword(k):
+                if not new_value.get("question_kwd"):
+                    new_value["questions"] = self.list2str(v)
+            elif self.field_keyword(k):
                 if isinstance(v, list):
-                    newValue[k] = "###".join(v)
+                    new_value[k] = "###".join(v)
                 else:
-                    newValue[k] = v
+                    new_value[k] = v
             elif re.search(r"_feas$", k):
-                newValue[k] = json.dumps(v)
+                new_value[k] = json.dumps(v)
             elif k == "kb_id":
-                if isinstance(newValue[k], list):
-                    newValue[k] = newValue[k][0]  # since d[k] is a list, but we need a str
+                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]
-                newValue[k] = "_".join(f"{num:08x}" for num in arr)
+                new_value[k] = "_".join(f"{num:08x}" for num in arr)
             elif k in ["page_num_int", "top_int"]:
                 assert isinstance(v, list)
-                newValue[k] = "_".join(f"{num:08x}" for num in v)
+                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}'."
@@ -712,22 +471,22 @@ class InfinityConnection(DocStoreConnection):
                     if ty.lower().find("cha"):
                         if not de:
                             de = ""
-                    newValue[v] = de
+                    new_value[v] = de
                 else:
                     for kk, vv in v.items():
                         removeValue[kk] = vv
-                    del newValue[k]
+                    del new_value[k]
             else:
-                newValue[k] = v
+                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 newValue:
-                del newValue[k]
+            if k in new_value:
+                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()
-            logger.debug(f"INFINITY search table {str(table_name)}, filter {filter}, result: {str(row_to_opt[0])}")
+            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():
@@ -740,78 +499,53 @@ class InfinityConnection(DocStoreConnection):
                         else:
                             remove_opt[kv_key].append(id)
 
-        logger.debug(f"INFINITY update table {table_name}, filter {filter}, newValue {newValue}.")
+        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, newValue)
+        table_instance.update(filter, new_value)
         self.connPool.release_conn(inf_conn)
         return True
 
-    def delete(self, condition: dict, indexName: str, knowledgebaseId: str) -> int:
-        inf_conn = self.connPool.get_conn()
-        db_instance = inf_conn.get_database(self.dbName)
-        table_name = f"{indexName}_{knowledgebaseId}"
-        try:
-            table_instance = db_instance.get_table(table_name)
-        except Exception:
-            logger.warning(f"Skipped deleting from table {table_name} since the table doesn't exist.")
-            return 0
-        filter = equivalent_condition_to_str(condition, table_instance)
-        logger.debug(f"INFINITY delete table {table_name}, filter {filter}.")
-        res = table_instance.delete(filter)
-        self.connPool.release_conn(inf_conn)
-        return res.deleted_rows
-
     """
     Helper functions for search result
     """
 
-    def get_total(self, res: tuple[pd.DataFrame, int] | pd.DataFrame) -> int:
-        if isinstance(res, tuple):
-            return res[1]
-        return len(res)
-
-    def get_chunk_ids(self, res: tuple[pd.DataFrame, int] | pd.DataFrame) -> list[str]:
-        if isinstance(res, tuple):
-            res = res[0]
-        return list(res["id"])
-
     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 {}
-        fieldsAll = fields.copy()
-        fieldsAll.append("id")
-        fieldsAll = set(fieldsAll)
+        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 fieldsAll:
+                if field in fields_all:
                     res[field] = res["docnm"]
         if "important_keywords" in res.columns:
-            if "important_kwd" in fieldsAll:
+            if "important_kwd" in fields_all:
                 res["important_kwd"] = res["important_keywords"].apply(lambda v: v.split())
-            if "important_tks" in fieldsAll:
+            if "important_tks" in fields_all:
                 res["important_tks"] = res["important_keywords"]
         if "questions" in res.columns:
-            if "question_kwd" in fieldsAll:
+            if "question_kwd" in fields_all:
                 res["question_kwd"] = res["questions"].apply(lambda v: v.splitlines())
-            if "question_tks" in fieldsAll:
+            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 fieldsAll:
+                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 fieldsAll:
+                if field in fields_all:
                     res[field] = res["authors"]
 
         column_map = {col.lower(): col for col in res.columns}
-        matched_columns = {column_map[col.lower()]: col for col in fieldsAll if col.lower() in column_map}
-        none_columns = [col for col in fieldsAll if col.lower() not in column_map]
+        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)
@@ -819,7 +553,7 @@ class InfinityConnection(DocStoreConnection):
 
         for column in list(res2.columns):
             k = column.lower()
-            if field_keyword(k):
+            if 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 {})
@@ -844,95 +578,3 @@ class InfinityConnection(DocStoreConnection):
             res2[column] = None
 
         return res2.set_index("id").to_dict(orient="index")
-
-    def get_highlight(self, res: tuple[pd.DataFrame, int] | pd.DataFrame, keywords: list[str], fieldnm: str):
-        if isinstance(res, tuple):
-            res = res[0]
-        ans = {}
-        num_rows = len(res)
-        column_id = res["id"]
-        if fieldnm not in res:
-            return {}
-        for i in range(num_rows):
-            id = column_id[i]
-            txt = res[fieldnm][i]
-            if re.search(r"<em>[^<>]+</em>", txt, flags=re.IGNORECASE | re.MULTILINE):
-                ans[id] = txt
-                continue
-            txt = re.sub(r"[\r\n]", " ", txt, flags=re.IGNORECASE | re.MULTILINE)
-            txts = []
-            for t in re.split(r"[.?!;\n]", txt):
-                if is_english([t]):
-                    for w in keywords:
-                        t = re.sub(
-                            r"(^|[ .?/'\"\(\)!,:;-])(%s)([ .?/'\"\(\)!,:;-])" % re.escape(w),
-                            r"\1<em>\2</em>\3",
-                            t,
-                            flags=re.IGNORECASE | re.MULTILINE,
-                        )
-                else:
-                    for w in sorted(keywords, key=len, reverse=True):
-                        t = re.sub(
-                            re.escape(w),
-                            f"<em>{w}</em>",
-                            t,
-                            flags=re.IGNORECASE | re.MULTILINE,
-                        )
-                if not re.search(r"<em>[^<>]+</em>", t, flags=re.IGNORECASE | re.MULTILINE):
-                    continue
-                txts.append(t)
-            if txts:
-                ans[id] = "...".join(txts)
-            else:
-                ans[id] = txt
-        return ans
-
-    def get_aggregation(self, res: tuple[pd.DataFrame, int] | pd.DataFrame, fieldnm: str):
-        """
-        Manual aggregation for tag fields since Infinity doesn't provide native aggregation
-        """
-        from collections import Counter
-
-        # Extract DataFrame from result
-        if isinstance(res, tuple):
-            df, _ = res
-        else:
-            df = res
-
-        if df.empty or fieldnm not in df.columns:
-            return []
-
-        # Aggregate tag counts
-        tag_counter = Counter()
-
-        for value in df[fieldnm]:
-            if pd.isna(value) or not value:
-                continue
-
-            # Handle different tag formats
-            if isinstance(value, str):
-                # Split by ### for tag_kwd field or comma for other formats
-                if fieldnm == "tag_kwd" and "###" in value:
-                    tags = [tag.strip() for tag in value.split("###") if tag.strip()]
-                else:
-                    # Try comma separation as fallback
-                    tags = [tag.strip() for tag in value.split(",") if tag.strip()]
-
-                for tag in tags:
-                    if tag:  # Only count non-empty tags
-                        tag_counter[tag] += 1
-            elif isinstance(value, list):
-                # Handle list format
-                for tag in value:
-                    if tag and isinstance(tag, str):
-                        tag_counter[tag.strip()] += 1
-
-        # Return as list of [tag, count] pairs, sorted by count descending
-        return [[tag, count] for tag, count in tag_counter.most_common()]
-
-    """
-    SQL
-    """
-
-    def sql(sql: str, fetch_size: int, format: str):
-        raise NotImplementedError("Not implemented")

rag/utils/ob_conn.py

@@ -37,9 +37,8 @@ from common import settings
 from common.constants import PAGERANK_FLD, TAG_FLD
 from common.decorator import singleton
 from common.float_utils import get_float
+from common.doc_store.doc_store_base import DocStoreConnection, MatchExpr, OrderByExpr, FusionExpr, MatchTextExpr, MatchDenseExpr
 from rag.nlp import rag_tokenizer
-from rag.utils.doc_store_conn import DocStoreConnection, MatchExpr, OrderByExpr, FusionExpr, MatchTextExpr, \
-    MatchDenseExpr
 
 ATTEMPT_TIME = 2
 OB_QUERY_TIMEOUT = int(os.environ.get("OB_QUERY_TIMEOUT", "100_000_000"))
@@ -497,7 +496,7 @@ class OBConnection(DocStoreConnection):
     Database operations
     """
 
-    def dbType(self) -> str:
+    def db_type(self) -> str:
         return "oceanbase"
 
     def health(self) -> dict:
@@ -553,7 +552,7 @@ class OBConnection(DocStoreConnection):
     Table operations
     """
 
-    def createIdx(self, indexName: str, knowledgebaseId: str, vectorSize: int):
+    def create_idx(self, indexName: str, knowledgebaseId: str, vectorSize: int):
         vector_field_name = f"q_{vectorSize}_vec"
         vector_index_name = f"{vector_field_name}_idx"
 
@@ -604,7 +603,7 @@ class OBConnection(DocStoreConnection):
             # always refresh metadata to make sure it contains the latest table structure
             self.client.refresh_metadata([indexName])
 
-    def deleteIdx(self, indexName: str, knowledgebaseId: str):
+    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
@@ -615,7 +614,7 @@ class OBConnection(DocStoreConnection):
         except Exception as e:
             raise Exception(f"OBConnection.deleteIndex error: {str(e)}")
 
-    def indexExist(self, indexName: str, knowledgebaseId: str = None) -> bool:
+    def index_exist(self, indexName: str, knowledgebaseId: str = None) -> bool:
         return self._check_table_exists_cached(indexName)
 
     def _get_count(self, table_name: str, filter_list: list[str] = None) -> int:
@@ -1500,7 +1499,7 @@ class OBConnection(DocStoreConnection):
     def get_total(self, res) -> int:
         return res.total
 
-    def get_chunk_ids(self, res) -> list[str]:
+    def get_doc_ids(self, res) -> list[str]:
         return [row["id"] for row in res.chunks]
 
     def get_fields(self, res, fields: list[str]) -> dict[str, dict]:

rag/utils/opensearch_conn.py

@@ -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 rag.utils.doc_store_conn 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
@@ -79,7 +78,7 @@ class OSConnection(DocStoreConnection):
     Database operations
     """
 
-    def dbType(self) -> str:
+    def db_type(self) -> str:
         return "opensearch"
 
     def health(self) -> dict:
@@ -91,8 +90,8 @@ class OSConnection(DocStoreConnection):
     Table operations
     """
 
-    def createIdx(self, indexName: str, knowledgebaseId: str, vectorSize: int):
-        if self.indexExist(indexName, knowledgebaseId):
+    def create_idx(self, indexName: str, knowledgebaseId: str, vectorSize: int):
+        if self.index_exist(indexName, knowledgebaseId):
             return True
         try:
             from opensearchpy.client import IndicesClient
@@ -101,7 +100,7 @@ class OSConnection(DocStoreConnection):
         except Exception:
             logger.exception("OSConnection.createIndex error %s" % (indexName))
 
-    def deleteIdx(self, indexName: str, knowledgebaseId: str):
+    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
@@ -112,7 +111,7 @@ class OSConnection(DocStoreConnection):
         except Exception:
             logger.exception("OSConnection.deleteIdx error %s" % (indexName))
 
-    def indexExist(self, indexName: str, knowledgebaseId: str = None) -> bool:
+    def index_exist(self, indexName: str, knowledgebaseId: str = None) -> bool:
         s = Index(indexName, self.os)
         for i in range(ATTEMPT_TIME):
             try:
@@ -460,7 +459,7 @@ class OSConnection(DocStoreConnection):
             return res["hits"]["total"]["value"]
         return res["hits"]["total"]
 
-    def get_chunk_ids(self, res):
+    def get_doc_ids(self, res):
         return [d["_id"] for d in res["hits"]["hits"]]
 
     def __getSource(self, res):

rag/utils/redis_conn.py

@@ -272,6 +272,49 @@ class RedisDB:
             self.__open__()
         return None
 
+    def generate_auto_increment_id(self, key_prefix: str = "id_generator", namespace: str = "default", increment: int = 1, ensure_minimum: int | None = None) -> int:
+        redis_key = f"{key_prefix}:{namespace}"
+
+        try:
+            # Use pipeline for atomicity
+            pipe = self.REDIS.pipeline()
+
+            # Check if key exists
+            pipe.exists(redis_key)
+
+            # Get/Increment
+            if ensure_minimum is not None:
+                # Ensure minimum value
+                pipe.get(redis_key)
+                results = pipe.execute()
+
+                if results[0] == 0:  # Key doesn't exist
+                    start_id = max(1, ensure_minimum)
+                    pipe.set(redis_key, start_id)
+                    pipe.execute()
+                    return start_id
+                else:
+                    current = int(results[1])
+                    if current < ensure_minimum:
+                        pipe.set(redis_key, ensure_minimum)
+                        pipe.execute()
+                        return ensure_minimum
+
+            # Increment operation
+            next_id = self.REDIS.incrby(redis_key, increment)
+
+            # If it's the first time, set a reasonable initial value
+            if next_id == increment:
+                self.REDIS.set(redis_key, 1 + increment)
+                return 1 + increment
+
+            return next_id
+
+        except Exception as e:
+            logging.warning("RedisDB.generate_auto_increment_id got exception: " + str(e))
+            self.__open__()
+        return -1
+
     def transaction(self, key, value, exp=3600):
         try:
             pipeline = self.REDIS.pipeline(transaction=True)

test/testcases/test_web_api/test_memory_app/conftest.py

@@ -32,8 +32,8 @@ def add_memory_func(request, WebApiAuth):
         payload = {
             "name": f"test_memory_{i}",
             "memory_type": ["raw"] + random.choices(["semantic", "episodic", "procedural"], k=random.randint(0, 3)),
-            "embd_id": "SILICONFLOW@BAAI/bge-large-zh-v1.5",
-            "llm_id": "ZHIPU-AI@glm-4-flash"
+            "embd_id": "BAAI/bge-large-zh-v1.5@SILICONFLOW",
+            "llm_id": "glm-4-flash@ZHIPU-AI"
         }
         res = create_memory(WebApiAuth, payload)
         memory_ids.append(res["data"]["id"])

test/testcases/test_web_api/test_memory_app/test_create_memory.py

@@ -49,8 +49,8 @@ class TestMemoryCreate:
         payload = {
             "name": name,
             "memory_type": ["raw"] + random.choices(["semantic", "episodic", "procedural"], k=random.randint(0, 3)),
-            "embd_id": "SILICONFLOW@BAAI/bge-large-zh-v1.5",
-            "llm_id": "ZHIPU-AI@glm-4-flash"
+            "embd_id": "BAAI/bge-large-zh-v1.5@SILICONFLOW",
+            "llm_id": "glm-4-flash@ZHIPU-AI"
         }
         res = create_memory(WebApiAuth, payload)
         assert res["code"] == 0, res
@@ -72,8 +72,8 @@ class TestMemoryCreate:
         payload = {
             "name": name,
             "memory_type": ["raw"] + random.choices(["semantic", "episodic", "procedural"], k=random.randint(0, 3)),
-            "embd_id": "SILICONFLOW@BAAI/bge-large-zh-v1.5",
-            "llm_id": "ZHIPU-AI@glm-4-flash"
+            "embd_id": "BAAI/bge-large-zh-v1.5@SILICONFLOW",
+            "llm_id": "glm-4-flash@ZHIPU-AI"
         }
         res = create_memory(WebApiAuth, payload)
         assert res["message"] == expected_message, res
@@ -84,8 +84,8 @@ class TestMemoryCreate:
         payload = {
             "name": name,
             "memory_type": ["something"],
-            "embd_id": "SILICONFLOW@BAAI/bge-large-zh-v1.5",
-            "llm_id": "ZHIPU-AI@glm-4-flash"
+            "embd_id": "BAAI/bge-large-zh-v1.5@SILICONFLOW",
+            "llm_id": "glm-4-flash@ZHIPU-AI"
         }
         res = create_memory(WebApiAuth, payload)
         assert res["message"] == f"Memory type '{ {'something'} }' is not supported.", res
@@ -96,8 +96,8 @@ class TestMemoryCreate:
         payload = {
             "name": name,
             "memory_type": ["raw"] + random.choices(["semantic", "episodic", "procedural"], k=random.randint(0, 3)),
-            "embd_id": "SILICONFLOW@BAAI/bge-large-zh-v1.5",
-            "llm_id": "ZHIPU-AI@glm-4-flash"
+            "embd_id": "BAAI/bge-large-zh-v1.5@SILICONFLOW",
+            "llm_id": "glm-4-flash@ZHIPU-AI"
         }
         res1 = create_memory(WebApiAuth, payload)
         assert res1["code"] == 0, res1

test/testcases/test_web_api/test_memory_app/test_update_memory.py

@@ -101,7 +101,7 @@ class TestMemoryUpdate:
     @pytest.mark.p1
     def test_llm(self, WebApiAuth, add_memory_func):
         memory_ids = add_memory_func
-        llm_id = "ZHIPU-AI@glm-4"
+        llm_id = "glm-4@ZHIPU-AI"
         payload = {"llm_id": llm_id}
         res = update_memory(WebApiAuth, memory_ids[0], payload)
         assert res["code"] == 0, res