From d201c3e621154d7fdf504b5ff462e60fdb438c7f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9D=92=E8=BD=A9?= Date: Wed, 21 Aug 2024 11:23:06 +0800 Subject: [PATCH] merge fixes of elastic search --- Dockerfile | 11 +- examples/api/simple_usages_cn.ipynb | 284 +++++++++--------- examples/cli/README.md | 0 examples/docker/entry_point.sh | 2 +- memoryscope/core/config/arguments.py | 6 +- memoryscope/core/config/config_manager.py | 4 - memoryscope/core/memoryscope.py | 7 +- memoryscope/core/memoryscope_context.py | 5 +- memoryscope/core/models/base_model.py | 2 +- .../models/llama_index_embedding_model.py | 3 +- .../models/llama_index_generation_model.py | 4 +- .../core/models/llama_index_rank_model.py | 3 +- memoryscope/core/operation/base_workflow.py | 2 +- .../storage/llama_index_es_memory_store.py | 54 +--- .../storage/llama_index_sync_elasticsearch.py | 73 ++++- memoryscope/core/utils/datetime_handler.py | 8 +- memoryscope/core/utils/logger.py | 38 ++- memoryscope/core/worker/memory_manager.py | 2 +- quick-start-demo.py | 7 + quick_start_demo.py | 7 - 20 files changed, 270 insertions(+), 252 deletions(-) create mode 100644 examples/cli/README.md create mode 100644 quick-start-demo.py delete mode 100644 quick_start_demo.py diff --git a/Dockerfile b/Dockerfile index 59db3639..614842c0 100644 --- a/Dockerfile +++ b/Dockerfile @@ -22,9 +22,10 @@ RUN echo '[global]' > /etc/pip.conf && \ echo 'index-url = https://mirrors.aliyun.com/pypi/simple/' >> /etc/pip.conf && \ echo 'trusted-host = mirrors.aliyun.com' >> /etc/pip.conf -# (Not necessary) Install the majority of deps, using docker build cache to accelerate future building -COPY requirements.txt ./ -RUN pip3 install -r requirements.txt +# # (Not necessary) Install the majority of deps, using docker build cache to accelerate future building +# COPY requirements.txt ./ +# RUN pip3 install poetry +# RUN poetry install # Install Elastic Search RUN useradd -m elastic_search_user @@ -46,7 +47,9 @@ WORKDIR /memory_scope_project # Enter working dir WORKDIR /memory_scope_project COPY . . -RUN pip install -r requirements.txt +RUN pip3 install poetry +# RUN pip install -r requirements.txt +RUN poetry install # Launch! # CMD ["bash"] diff --git a/examples/api/simple_usages_cn.ipynb b/examples/api/simple_usages_cn.ipynb index ca00098a..059ad530 100644 --- a/examples/api/simple_usages_cn.ipynb +++ b/examples/api/simple_usages_cn.ipynb @@ -2,29 +2,38 @@ "cells": [ { "cell_type": "markdown", + "metadata": { + "collapsed": false + }, "source": [ "# chat 和 service 接口的示例用法\n", "这个笔记本展示了 MemoryScope 的 **chat** 和 **service** 接口的简单用法，以及它的主要功能。\n", "\n", "在运行这个笔记本之前，请先按照 Readme 中的 **Installation** 指南进行安装，并启动 Docker 镜像。\n" - ], - "metadata": { - "collapsed": false - } + ] }, { "cell_type": "markdown", - "source": [ - "## 初始化一个 MemoryScope 实例\n", - "首先，我们需要指定一个配置并初始化一个 MemoryScope 实例。" - ], "metadata": { "collapsed": false - } + }, + "source": [ + "## 初始化一个 MemoryScope 实例\n", + "首先，我们需要指定一个配置并初始化一个 MemoryScope 实例。\n", + "\n", + "(需要首先启动elastic search服务，详见https://www.elastic.co/guide/en/elasticsearch/reference/current/getting-started.html)" + ] }, { "cell_type": "code", "execution_count": 1, + "metadata": { + "ExecuteTime": { + "end_time": "2024-08-02T14:42:19.303078Z", + "start_time": "2024-08-02T14:42:17.332785Z" + }, + "collapsed": false + }, "outputs": [], "source": [ "import sys\n", @@ -47,28 +56,28 @@ ")\n", "\n", "ms = MemoryScope(arguments=arguments)\n" - ], - "metadata": { - "collapsed": false, - "ExecuteTime": { - "end_time": "2024-08-02T14:42:19.303078Z", - "start_time": "2024-08-02T14:42:17.332785Z" - } - } + ] }, { "cell_type": "markdown", + "metadata": { + "collapsed": false + }, "source": [ "## 聊天（不含记忆）\n", "MemoryScope 配有默认的 chat 接口，因此开始聊天非常容易，就像使用任何大型语言模型聊天机器人一样。" - ], - "metadata": { - "collapsed": false - } + ] }, { "cell_type": "code", "execution_count": 2, + "metadata": { + "ExecuteTime": { + "end_time": "2024-08-02T14:42:24.843842Z", + "start_time": "2024-08-02T14:42:19.304038Z" + }, + "collapsed": false + }, "outputs": [ { "name": "stdout", @@ -84,28 +93,28 @@ "memory_chat.run_service_operation(\"delete_all\")\n", "response = memory_chat.chat_with_memory(query=\"我的爱好是弹琴。\")\n", "print(\"回答1：\\n\" + response.message.content)" - ], - "metadata": { - "collapsed": false, - "ExecuteTime": { - "end_time": "2024-08-02T14:42:24.843842Z", - "start_time": "2024-08-02T14:42:19.304038Z" - } - } + ] }, { "cell_type": "markdown", + "metadata": { + "collapsed": false + }, "source": [ "----\n", "你可以选择进行含有或不含有多轮对话上下文的聊天。然而，由于尚未调用**记忆巩固**功能，系统中还没有任何记忆片段。" - ], - "metadata": { - "collapsed": false - } + ] }, { "cell_type": "code", "execution_count": 3, + "metadata": { + "ExecuteTime": { + "end_time": "2024-08-02T14:42:33.777924Z", + "start_time": "2024-08-02T14:42:24.845133Z" + }, + "collapsed": false + }, "outputs": [ { "name": "stdout", @@ -124,28 +133,28 @@ "response = memory_chat.chat_with_memory(query=\"你知道我有什么乐器爱好吗？\",\n", " history_message_strategy=None)\n", "print(\"回答3：\\n\" + response.message.content)" - ], - "metadata": { - "collapsed": false, - "ExecuteTime": { - "end_time": "2024-08-02T14:42:33.777924Z", - "start_time": "2024-08-02T14:42:24.845133Z" - } - } + ] }, { "cell_type": "markdown", + "metadata": { + "collapsed": false + }, "source": [ "## **记忆巩固**\n", "现在，我们再聊多几句，然后尝试**记忆巩固**功能。" - ], - "metadata": { - "collapsed": false - } + ] }, { "cell_type": "code", "execution_count": 4, + "metadata": { + "ExecuteTime": { + "end_time": "2024-08-02T14:42:52.517936Z", + "start_time": "2024-08-02T14:42:33.780200Z" + }, + "collapsed": false + }, "outputs": [ { "name": "stdout", @@ -171,18 +180,18 @@ "print(\"回答6：\\n\" + response.message.content)\n", "response = memory_chat.chat_with_memory(query=\"帮我写一句给朋友的生日祝福语，简短一点。\")\n", "print(\"回答7：\\n\" + response.message.content)" - ], - "metadata": { - "collapsed": false, - "ExecuteTime": { - "end_time": "2024-08-02T14:42:52.517936Z", - "start_time": "2024-08-02T14:42:33.780200Z" - } - } + ] }, { "cell_type": "code", "execution_count": 5, + "metadata": { + "ExecuteTime": { + "end_time": "2024-08-02T14:43:14.475175Z", + "start_time": "2024-08-02T14:42:52.518868Z" + }, + "collapsed": false + }, "outputs": [ { "name": "stdout", @@ -200,30 +209,30 @@ "memory_service.init_service()\n", "result = memory_service.consolidate_memory()\n", "print(f\"consolidate_memory result={result}\")" - ], - "metadata": { - "collapsed": false, - "ExecuteTime": { - "end_time": "2024-08-02T14:43:14.475175Z", - "start_time": "2024-08-02T14:42:52.518868Z" - } - } + ] }, { "cell_type": "markdown", + "metadata": { + "collapsed": false + }, "source": [ "----\n", "**记忆巩固**从用户的7条聊天消息中提取了3条 *observations* ，其余无效的信息被过滤掉了。\n", "\n", "我们尝试更多的情况，以测试其时间感知能力和过滤用户虚构的内容的能力。" - ], - "metadata": { - "collapsed": false - } + ] }, { "cell_type": "code", "execution_count": 6, + "metadata": { + "ExecuteTime": { + "end_time": "2024-08-02T14:43:57.226685Z", + "start_time": "2024-08-02T14:43:14.475977Z" + }, + "collapsed": false + }, "outputs": [ { "name": "stdout", @@ -271,18 +280,18 @@ "print(\"回答11：\\n\" + response.message.content)\n", "response = memory_chat.chat_with_memory(query=\"SMCI是什么公司，做什么的？\")\n", "print(\"回答12：\\n\" + response.message.content)" - ], - "metadata": { - "collapsed": false, - "ExecuteTime": { - "end_time": "2024-08-02T14:43:57.226685Z", - "start_time": "2024-08-02T14:43:14.475977Z" - } - } + ] }, { "cell_type": "code", "execution_count": 7, + "metadata": { + "ExecuteTime": { + "end_time": "2024-08-02T14:44:26.290389Z", + "start_time": "2024-08-02T14:43:57.227597Z" + }, + "collapsed": false + }, "outputs": [ { "name": "stdout", @@ -297,30 +306,30 @@ "source": [ "result = memory_service.consolidate_memory()\n", "print(f\"consolidate_memory result={result}\")" - ], - "metadata": { - "collapsed": false, - "ExecuteTime": { - "end_time": "2024-08-02T14:44:26.290389Z", - "start_time": "2024-08-02T14:43:57.227597Z" - } - } + ] }, { "cell_type": "markdown", + "metadata": { + "collapsed": false + }, "source": [ "----\n", "我们可以看到，**记忆巩固**成功过滤掉了虚假内容，并展示了良好的时间敏感性。\n", "\n", "我们尝试更多的情况，以测试其解决冲突内容的能力。" - ], - "metadata": { - "collapsed": false - } + ] }, { "cell_type": "code", "execution_count": 8, + "metadata": { + "ExecuteTime": { + "end_time": "2024-08-02T14:44:57.748723Z", + "start_time": "2024-08-02T14:44:26.292543Z" + }, + "collapsed": false + }, "outputs": [ { "name": "stdout", @@ -358,18 +367,18 @@ "print(\"回答18：\\n\" + response.message.content)\n", "response = memory_chat.chat_with_memory(query=\"明天是我生日。\")\n", "print(\"回答19：\\n\" + response.message.content)" - ], - "metadata": { - "collapsed": false, - "ExecuteTime": { - "end_time": "2024-08-02T14:44:57.748723Z", - "start_time": "2024-08-02T14:44:26.292543Z" - } - } + ] }, { "cell_type": "code", "execution_count": 9, + "metadata": { + "ExecuteTime": { + "end_time": "2024-08-02T14:45:38.454853Z", + "start_time": "2024-08-02T14:44:57.750558Z" + }, + "collapsed": false + }, "outputs": [ { "name": "stdout", @@ -389,28 +398,28 @@ "source": [ "result = memory_service.consolidate_memory()\n", "print(f\"consolidate_memory result={result}\")" - ], - "metadata": { - "collapsed": false, - "ExecuteTime": { - "end_time": "2024-08-02T14:45:38.454853Z", - "start_time": "2024-08-02T14:44:57.750558Z" - } - } + ] }, { "cell_type": "markdown", + "metadata": { + "collapsed": false + }, "source": [ "## **反思与再巩固**\n", "现在，我们在系统中已经积累了足够多的新的 *observations* ，因此我们可以调用**反思与再巩固**功能，让我们看看会得到什么。" - ], - "metadata": { - "collapsed": false - } + ] }, { "cell_type": "code", "execution_count": 10, + "metadata": { + "ExecuteTime": { + "end_time": "2024-08-02T14:45:52.599528Z", + "start_time": "2024-08-02T14:45:38.455556Z" + }, + "collapsed": false + }, "outputs": [ { "name": "stdout", @@ -435,29 +444,25 @@ "source": [ "result = memory_service.reflect_and_reconsolidate()\n", "print(f\"consolidate_memory result={result}\")" - ], - "metadata": { - "collapsed": false, - "ExecuteTime": { - "end_time": "2024-08-02T14:45:52.599528Z", - "start_time": "2024-08-02T14:45:38.455556Z" - } - } + ] }, { "cell_type": "markdown", + "metadata": { + "collapsed": false + }, "source": [ "## 低用户时延（RT）\n", "\n", "最后，我们测试 MemoryScope 系统对用户的响应时间 (RT)。具体来说，我们测试在有和没有从系统中检索记忆片段时聊天的响应时间的差异。" - ], - "metadata": { - "collapsed": false - } + ] }, { "cell_type": "code", - "execution_count": 0, + "execution_count": null, + "metadata": { + "collapsed": false + }, "outputs": [], "source": [ "import time\n", @@ -475,14 +480,18 @@ "end_time = time.time()\n", "total_time = end_time - start_time\n", "print(\"使用记忆检索\\n回答21：\\n\" + response.message.content + f\"\\n 耗时：{total_time}秒\\n\")" - ], - "metadata": { - "collapsed": false - } + ] }, { "cell_type": "code", "execution_count": 11, + "metadata": { + "ExecuteTime": { + "end_time": "2024-08-02T14:46:08.814486Z", + "start_time": "2024-08-02T14:45:52.601688Z" + }, + "collapsed": false + }, "outputs": [ { "name": "stdout", @@ -528,45 +537,38 @@ "end_time = time.time()\n", "total_time = end_time - start_time\n", "print(\"不使用记忆检索\\n回答21：\\n\" + response.message.content + f\"\\n 耗时：{total_time}秒\")" - ], - "metadata": { - "collapsed": false, - "ExecuteTime": { - "end_time": "2024-08-02T14:46:08.814486Z", - "start_time": "2024-08-02T14:45:52.601688Z" - } - } + ] }, { "cell_type": "markdown", + "metadata": { + "collapsed": false + }, "source": [ "----\n", "我们可以看到，从 MemoryScope 检索记忆片段不会增加聊天的响应时间。" - ], - "metadata": { - "collapsed": false - } + ] }, { "cell_type": "markdown", + "metadata": { + "collapsed": false + }, "source": [ "## 更多用法\n", "我们建议读者参考[进阶自定义用法](./examples/api/advanced_customization_cn.ipynb)来对MemoryScope系统进行各种自定义设置。您还可以通过自定义**workflow**和对应的**worker**来创建或定制满足您特定需求的**operation**。\n", "\n", "此外，您还可以尝试使用[在命令行与MemoryScope聊天机器人交互](./examples/cli/dash_cli_cn1.sh)。我们在这里实现了始终在后台异步运行**记忆巩固**和**反思与再巩固**这两个操作，从而使得它们不会增加聊天的响应时间。" - ], - "metadata": { - "collapsed": false - } + ] }, { "cell_type": "code", "execution_count": null, - "outputs": [], - "source": [], "metadata": { "collapsed": false - } + }, + "outputs": [], + "source": [] } ], "metadata": { diff --git a/examples/cli/README.md b/examples/cli/README.md new file mode 100644 index 00000000..e69de29b diff --git a/examples/docker/entry_point.sh b/examples/docker/entry_point.sh index e296a67d..e75e0619 100644 --- a/examples/docker/entry_point.sh +++ b/examples/docker/entry_point.sh @@ -1,2 +1,2 @@ sh examples/docker/run_elastic_search.sh -python quick_start_demo.py --config_path=memoryscope/core/config/demo_config.yaml \ No newline at end of file +python quick-start-demo.py --config_path=memoryscope/core/config/demo_config.yaml \ No newline at end of file diff --git a/memoryscope/core/config/arguments.py b/memoryscope/core/config/arguments.py index 20651ef8..e13c6ec5 100644 --- a/memoryscope/core/config/arguments.py +++ b/memoryscope/core/config/arguments.py @@ -3,7 +3,7 @@ from typing import Literal, Dict @dataclass class Arguments(object): - language: Literal["cn", "en"] = field(default="en", metadata={"help": "support en & cn now"}) + language: Literal["cn", "en"] = field(default="cn", metadata={"help": "support en & cn now"}) thread_pool_max_workers: int = field(default=5, metadata={"help": "thread pool max workers"}) @@ -35,7 +35,7 @@ class Arguments(object): worker_params: Dict[str, dict] = field(default_factory=lambda: {}, metadata={ "help": "dict format: worker_name -> param_key -> param_value"}) - generation_backend: str = field(default="openai_generation", metadata={ + generation_backend: str = field(default="dashscope_generation", metadata={ "help": "global generation backend: openai_generation, dashscope_generation, etc."}) generation_model: str = field(default="gpt-4o", metadata={ @@ -44,7 +44,7 @@ class Arguments(object): generation_params: dict = field(default_factory=lambda: {}, metadata={ "help": "global generation params: max_tokens, top_p, temperature, etc."}) - embedding_backend: str = field(default="openai_embedding", metadata={ + embedding_backend: str = field(default="dashscope_generation", metadata={ "help": "global embedding backend: openai_embedding, dashscope_embedding, etc."}) embedding_model: str = field(default="text-embedding-3-small", metadata={ diff --git a/memoryscope/core/config/config_manager.py b/memoryscope/core/config/config_manager.py index 8f4c75c8..74312708 100644 --- a/memoryscope/core/config/config_manager.py +++ b/memoryscope/core/config/config_manager.py @@ -52,10 +52,6 @@ class ConfigManager(object): def _init_logger(self) -> Logger: global_config = self.config["global"] logger_name = global_config["logger_name"] - logger_name_time_suffix = global_config["logger_name_time_suffix"] - if logger_name_time_suffix: - suffix = datetime.now().strftime(logger_name_time_suffix) - logger_name = f"{logger_name}_{suffix}" return Logger.get_logger(logger_name, to_stream=global_config["logger_to_screen"]) def read_config(self, config_path: str): diff --git a/memoryscope/core/memoryscope.py b/memoryscope/core/memoryscope.py index 818bfe12..94f1b7b2 100644 --- a/memoryscope/core/memoryscope.py +++ b/memoryscope/core/memoryscope.py @@ -13,8 +13,9 @@ from memoryscope.enumeration.model_enum import ModelEnum class MemoryScope(ConfigManager): def __init__(self, **kwargs): - super().__init__(**kwargs) self._context: MemoryscopeContext = MemoryscopeContext() + self._context.memory_scope_uuid = datetime.now().strftime(r"%Y%m%d_%H%M%S") + super().__init__(**kwargs) self._init_context_by_config() def _init_context_by_config(self): @@ -33,10 +34,6 @@ class MemoryScope(ConfigManager): self.logger.warning("If a semantic ranking model is not available, MemoryScope will use cosine similarity " "scoring as a substitute. However, the ranking effectiveness will be somewhat " "compromised.") - self._context.memory_scope_uuid = datetime.now().strftime(global_conf["logger_name_time_suffix"]) - - # set context_initialized - self._context.context_initialized = True # init memory_chat memory_chat_conf_dict = self.config["memory_chat"] diff --git a/memoryscope/core/memoryscope_context.py b/memoryscope/core/memoryscope_context.py index 99a163ae..f8be3cd7 100644 --- a/memoryscope/core/memoryscope_context.py +++ b/memoryscope/core/memoryscope_context.py @@ -33,11 +33,10 @@ class MemoryscopeContext(object): print_workflow_dynamic: bool = False - context_initialized: bool = False def get_ms_context(): ms_context = MemoryscopeContext() - if ms_context.context_initialized: - return ms_context + if ms_context.memory_scope_uuid: + return ms_context.memory_scope_uuid else: raise RuntimeError("MemoryscopeContext is not initialized yet. Please initialize it first.") diff --git a/memoryscope/core/models/base_model.py b/memoryscope/core/models/base_model.py index 4fc6befe..8f46871c 100644 --- a/memoryscope/core/models/base_model.py +++ b/memoryscope/core/models/base_model.py @@ -38,7 +38,7 @@ class BaseModel(metaclass=ABCMeta): self.kwargs: dict = kwargs self._model: Any = None - self.logger = Logger.get_logger(Logger.append_timestamp("base_model")) + self.logger = Logger.get_logger("base_model") @property def model(self): diff --git a/memoryscope/core/models/llama_index_embedding_model.py b/memoryscope/core/models/llama_index_embedding_model.py index 742b9558..88a733b5 100644 --- a/memoryscope/core/models/llama_index_embedding_model.py +++ b/memoryscope/core/models/llama_index_embedding_model.py @@ -6,6 +6,7 @@ from llama_index.embeddings.openai import OpenAIEmbedding from memoryscope.core.models.base_model import BaseModel, MODEL_REGISTRY from memoryscope.enumeration.model_enum import ModelEnum from memoryscope.scheme.model_response import ModelResponse +from memoryscope.core.utils.logger import Logger class LlamaIndexEmbeddingModel(BaseModel): @@ -17,7 +18,7 @@ class LlamaIndexEmbeddingModel(BaseModel): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self.logger = self.logger.get_logger(self.logger.append_timestamp("llama_index_embedding_model")) + self.logger = Logger.get_logger("llama_index_embedding_model") @classmethod def register_model(cls, model_name: str, model_class: type): diff --git a/memoryscope/core/models/llama_index_generation_model.py b/memoryscope/core/models/llama_index_generation_model.py index 739b6de3..f38f08d0 100644 --- a/memoryscope/core/models/llama_index_generation_model.py +++ b/memoryscope/core/models/llama_index_generation_model.py @@ -9,7 +9,7 @@ from memoryscope.enumeration.message_role_enum import MessageRoleEnum from memoryscope.enumeration.model_enum import ModelEnum from memoryscope.scheme.message import Message from memoryscope.scheme.model_response import ModelResponse, ModelResponseGen - +from memoryscope.core.utils.logger import Logger class LlamaIndexGenerationModel(BaseModel): """ @@ -27,7 +27,7 @@ class LlamaIndexGenerationModel(BaseModel): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self.logger = self.logger.get_logger(self.logger.append_timestamp("llama_index_generation_model")) + self.logger = Logger.get_logger("llama_index_generation_model") def before_call(self, model_response: ModelResponse, **kwargs): """ diff --git a/memoryscope/core/models/llama_index_rank_model.py b/memoryscope/core/models/llama_index_rank_model.py index 6e4c4245..e290e0cc 100644 --- a/memoryscope/core/models/llama_index_rank_model.py +++ b/memoryscope/core/models/llama_index_rank_model.py @@ -7,6 +7,7 @@ from llama_index.postprocessor.dashscope_rerank import DashScopeRerank from memoryscope.core.models.base_model import BaseModel, MODEL_REGISTRY from memoryscope.enumeration.model_enum import ModelEnum from memoryscope.scheme.model_response import ModelResponse +from memoryscope.core.utils.logger import Logger class LlamaIndexRankModel(BaseModel): @@ -22,7 +23,7 @@ class LlamaIndexRankModel(BaseModel): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self.logger = self.logger.get_logger(self.logger.append_timestamp("llama_index_rank_model")) + self.logger = Logger.get_logger("llama_index_rank_model") def before_call(self, model_response: ModelResponse, **kwargs): """ diff --git a/memoryscope/core/operation/base_workflow.py b/memoryscope/core/operation/base_workflow.py index e1fd1495..4f8b6156 100644 --- a/memoryscope/core/operation/base_workflow.py +++ b/memoryscope/core/operation/base_workflow.py @@ -31,7 +31,7 @@ class BaseWorkflow(object): self.workflow_context: Dict[str, Any] = {} self.context_lock = threading.Lock() - self.logger: Logger = Logger.get_logger(Logger.append_timestamp("workflow")) + self.logger: Logger = Logger.get_logger("workflow") if self.workflow: self.workflow_worker_list = self._parse_workflow() diff --git a/memoryscope/core/storage/llama_index_es_memory_store.py b/memoryscope/core/storage/llama_index_es_memory_store.py index 5593b558..b77523f4 100644 --- a/memoryscope/core/storage/llama_index_es_memory_store.py +++ b/memoryscope/core/storage/llama_index_es_memory_store.py @@ -9,8 +9,7 @@ from memoryscope.core.models.base_model import BaseModel from memoryscope.core.storage.base_memory_store import BaseMemoryStore from memoryscope.core.storage.llama_index_sync_elasticsearch import (SyncElasticsearchStore, ESCombinedRetrieveStrategy, - _to_elasticsearch_filter, - SPECIAL_QUERY) + _to_elasticsearch_filter) from memoryscope.core.utils.logger import Logger from memoryscope.scheme.memory_node import MemoryNode @@ -38,7 +37,7 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): self.index = VectorStoreIndex.from_vector_store(vector_store=self.es_store, embed_model=self.embedding_model.model) - self.logger = Logger.get_logger(Logger.append_timestamp("es_memory_store")) + self.logger = Logger.get_logger("es_memory_store") def retrieve_memories(self, query: str = "", @@ -57,15 +56,12 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): similarity_top_k=top_k, sparse_top_k=top_k) - if not query: - query = SPECIAL_QUERY - - if not query and self.emb_dims: - query = QueryBundle(query_str=SPECIAL_QUERY, embedding=self.dummy_query_vector()) - - text_nodes = retriever.retrieve(query) - if text_nodes and text_nodes[0].embedding: - self.emb_dims = len(text_nodes[0].embedding) + if query: + text_nodes = retriever.retrieve(query) + if text_nodes and text_nodes[0].embedding: + self.emb_dims = len(text_nodes[0].embedding) + else: + text_nodes = self.es_store.sync_search_all_with_filter(es_filter, ['embedding']) self.logger.log_dictionary_info({ "action": "retrieve_memories", "query": query, @@ -77,31 +73,7 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): query: str = "", top_k: int = 3, filter_dict: Dict[str, List[str]] | Dict[str, str] = None) -> List[MemoryNode]: - # if index is not created, return [] - exists = self.es_store.client.indices.exists(index=self.index_name) - if not exists: - return [] - - if filter_dict is None: - filter_dict = {} - - es_filter = _to_elasticsearch_filter(filter_dict) - retriever = self.index.as_retriever(vector_store_kwargs={"es_filter": es_filter, "fields": ['embedding']}, - similarity_top_k=top_k, - sparse_top_k=top_k) - - if not query: - query = SPECIAL_QUERY - - if not query: - query = QueryBundle(query_str=SPECIAL_QUERY, embedding=self.dummy_query_vector()) - - text_nodes: List[NodeWithScore] = await retriever.aretrieve(query) - - if text_nodes and text_nodes[0].embedding: - self.emb_dims = len(text_nodes[0].embedding) - - return [self._text_node_2_memory_node(n) for n in text_nodes] + raise NotImplementedError def batch_insert(self, nodes: List[MemoryNode]): self.index.insert_nodes([self._memory_node_2_text_node(node) for node in nodes]) @@ -191,13 +163,15 @@ class LlamaIndexEsMemoryStore(BaseMemoryStore): MemoryNode: The converted MemoryNode with text and metadata from the NodeWithScore. """ - if text_node.metadata["key_vector"]: + if text_node.metadata.get("key_vector", None): key_vector = pickle.loads(text_node.metadata["key_vector"].encode('latin1')) else: key_vector = [] text_node.metadata["key_vector"] = key_vector text_node.metadata["vector"] = text_node.embedding if text_node.embedding else [] - text_node.metadata["score_recall"] = text_node.score - + + if hasattr(text_node, "score"): + text_node.metadata["score_recall"] = text_node.score + return MemoryNode(content=text_node.text, **text_node.metadata) diff --git a/memoryscope/core/storage/llama_index_sync_elasticsearch.py b/memoryscope/core/storage/llama_index_sync_elasticsearch.py index 41881a69..ab156d06 100644 --- a/memoryscope/core/storage/llama_index_sync_elasticsearch.py +++ b/memoryscope/core/storage/llama_index_sync_elasticsearch.py @@ -36,8 +36,6 @@ DISTANCE_STRATEGIES = Literal[ "EUCLIDEAN_DISTANCE", ] -SPECIAL_QUERY: str = "**--**" - def get_elasticsearch_client( url: Optional[str] = None, @@ -159,7 +157,7 @@ class ESCombinedRetrieveStrategy(AsyncDenseVectorStrategy): # RRF is used to even the score from the knn query and text query # RRF has two optional parameters: {'rank_constant':int, 'window_size':int} # https://www.elastic.co/guide/en/elasticsearch/reference/current/rrf.html - if query == SPECIAL_QUERY: + if not query: query_body = { "query": { "bool": { @@ -430,7 +428,7 @@ class SyncElasticsearchStore(BasePydanticVectorStore): retrieval_strategy=retrieval_strategy, ) - self.logger = Logger.get_logger(Logger.append_timestamp("elastic_search")) + self.logger = Logger.get_logger("elastic_search") @property def client(self) -> Any: @@ -472,11 +470,10 @@ class SyncElasticsearchStore(BasePydanticVectorStore): Note: This method delegates the actual operation to the `sync_add` method. """ - self.logger.log_dictionary_info({ - "action": "add", - "node_count": len(nodes), - }) - return self.sync_add(nodes, create_index_if_not_exists=create_index_if_not_exists) + add_res = self.sync_add(nodes, create_index_if_not_exists=create_index_if_not_exists) + self.log_vector_store_brief(title='after add') + return add_res + def sync_add( self, @@ -555,11 +552,9 @@ class SyncElasticsearchStore(BasePydanticVectorStore): This method internally calls a synchronous delete method (`sync_delete`) to execute the deletion operation against Elasticsearch. """ - self.logger.log_dictionary_info({ - "action": "delete", - "id": ref_doc_id, - }) - return self.sync_delete(ref_doc_id, **delete_kwargs) + del_res = self.sync_delete(ref_doc_id, **delete_kwargs) + self.log_vector_store_brief(title='after delete') + return del_res def sync_delete(self, ref_doc_id: str, **delete_kwargs: Any) -> None: """ @@ -613,19 +608,59 @@ class SyncElasticsearchStore(BasePydanticVectorStore): Exception: If an error occurs during the Elasticsearch query execution. """ + + q_res = self.sync_query(query, custom_query, es_filter, **kwargs) self.logger.log_dictionary_info({ "action": "query", "query": query.query_str, + "result": [tn.text for tn in q_res.nodes] }) - return self.sync_query(query, custom_query, es_filter, **kwargs) + return q_res def sync_delete_all(self): self._store.client.delete_by_query(index=[self.index_name], body={"query": {"match_all": {}}}) def sync_search_all(self): search_res = self._store.client.search(index=[self.index_name], body={"query": {"match_all": {}}}) - raise search_res + return search_res + def log_vector_store_brief(self, title="current vector store content"): + search_res = self.sync_search_all() + + brief = { + f"{hit['_source']['metadata']['memory_id']}({hit['_source']['metadata']['user_name']}/{hit['_source']['metadata']['target_name']}/{hit['_source']['metadata']['memory_type']})": + hit['_source']['content'] + for hit in search_res["hits"]["hits"] + } + self.logger.log_dictionary_info(brief, title=title) + + return brief + + def sync_search_all_with_filter(self, es_filter, fields): + query_body = {'query': {'bool': {'filter': es_filter}}} + k = 1000 + fields = ['embedding', 'metadata', 'content'] + response = self.client.search( + index=self.index_name, + **query_body, + size=k, + source=True, + source_includes=fields, + ) + res = [] + for hit in response["hits"]["hits"]: + tn = TextNode( + id_=hit['_id'], + text=hit['_source']['content'], + embedding=hit['_source']['embedding'], + text_template="{content}", + metadata=hit['_source']['metadata'] + ) + res.append( + tn + ) + return res + def sync_query( self, query: VectorStoreQuery, @@ -676,9 +711,15 @@ class SyncElasticsearchStore(BasePydanticVectorStore): custom_query=custom_query, fields=fields, ) + + return self.post_process_hits(hits) + + + def post_process_hits(self, hits: List[Dict[str, Any]]) -> VectorStoreQueryResult: top_k_nodes = [] top_k_ids = [] top_k_scores = [] + for hit in hits: source = hit["_source"] metadata = source.get("metadata", None) diff --git a/memoryscope/core/utils/datetime_handler.py b/memoryscope/core/utils/datetime_handler.py index 9a702144..29189277 100644 --- a/memoryscope/core/utils/datetime_handler.py +++ b/memoryscope/core/utils/datetime_handler.py @@ -3,7 +3,6 @@ import re from typing import List from memoryscope.constants.language_constants import WEEKDAYS, DATATIME_WORD_LIST, MONTH_DICT -from memoryscope.core.utils.logger import Logger from memoryscope.enumeration.language_enum import LanguageEnum @@ -14,7 +13,6 @@ class DatetimeHandler(object): specialized text parsing for date components. """ - logger = Logger.get_logger() def __init__(self, dt: datetime.datetime | str | int | float = None): """ @@ -224,7 +222,7 @@ class DatetimeHandler(object): """ func_name = f"extract_date_parts_{language.value}" if not hasattr(cls, func_name): - cls.logger.warning(f"language={language.value} needs to complete extract_date_parts func!") + # cls.logger.warning(f"language={language.value} needs to complete extract_date_parts func!") return {} return getattr(cls, func_name)(input_string=input_string) @@ -274,11 +272,11 @@ class DatetimeHandler(object): def has_time_word(cls, query: str, language: LanguageEnum) -> bool: func_name = f"has_time_word_{language.value}" if not hasattr(cls, func_name): - cls.logger.warning(f"language={language.value} needs to complete has_time_word function!") + # cls.logger.warning(f"language={language.value} needs to complete has_time_word function!") return False if language not in DATATIME_WORD_LIST: - cls.logger.warning(f"language={language.value} is missing in DATATIME_WORD_LIST!") + # cls.logger.warning(f"language={language.value} is missing in DATATIME_WORD_LIST!") return False datetime_word_list = DATATIME_WORD_LIST[language] diff --git a/memoryscope/core/utils/logger.py b/memoryscope/core/utils/logger.py index 8c3bf251..4a79e122 100644 --- a/memoryscope/core/utils/logger.py +++ b/memoryscope/core/utils/logger.py @@ -1,3 +1,4 @@ +import os import logging import pprint from logging.handlers import RotatingFileHandler @@ -6,7 +7,7 @@ from rich.console import Console from rich.panel import Panel from rich.text import Text -LOG_FORMAT = "%(asctime)s %(levelname)s %(threadName)s %(module)s:%(lineno)d] %(message)s" +LOG_FORMAT = "%(asctime)s %(levelname)s [%(module)s:%(lineno)d] %(message)s" DATE_FORMAT = "%Y-%m-%d %H:%M:%S" LOGGER_DICT = {} @@ -17,6 +18,11 @@ def rich2text(rich_table): console.print(rich_table) return '\n' + str(Text.from_ansi(capture.get())) +def append_memoryscope_uuid(dir_path): + from memoryscope.core.memoryscope_context import get_ms_context # pylint: disable=import-outside-toplevel + dir_path = os.path.join(dir_path, get_ms_context()) + return dir_path + class Logger(logging.Logger): """ The `Logger` class handle the stream of information or errors in activities. @@ -72,13 +78,15 @@ class Logger(logging.Logger): self.info(f"logger={name} is inited.") # Logs an initialization message - def log_dictionary_info(self, dictionary): - self.info(self.format_current_context(dictionary)) + def log_dictionary_info(self, dictionary, title=""): + self.info(self.format_current_context(dictionary, title)) - def format_current_context(self, context): + def format_current_context(self, context, title=""): pp = pprint.PrettyPrinter() pretty_string = pp.pformat(context) - return rich2text(Panel(pretty_string, width=128)) + if title: + pretty_string = f"{title}\n{pretty_string}" + return self.wrap_in_box(pretty_string) def wrap_in_box(self, context): return rich2text(Panel(context, width=128)) @@ -90,7 +98,7 @@ class Logger(logging.Logger): for chat_message in message.meta_data['data']['messages']: buf.append(chat_message.content) buf.append('\n') - buf.append(f"--------------------------------------------------------------\n") + buf.append(f"------------------------------------------\n") buf.append(f"LM Output:\n") buf.append(message.message.content) buf.append('\n') @@ -103,7 +111,7 @@ class Logger(logging.Logger): buf.append(f"Query Input:\n") buf.append(model_response.meta_data['data']['query_str']) buf.append('\n') - buf.append(f"--------------------------------------------------------------\n") + buf.append(f"------------------------------------------\n") buf.append(f"Rank:\n") rank = 0 for index, score in model_response.rank_scores.items(): @@ -126,10 +134,11 @@ class Logger(logging.Logger): for consistent log message formatting. """ file_path = Path().joinpath(self.dir_path, f"{self.name}.{self.file_type}") - file_path.parent.mkdir(exist_ok=True) # Ensure the directory exists + os.makedirs(file_path.parent, exist_ok=True) # Ensure the directory exists file_name = file_path.as_posix() # Get the absolute path as a string - Console().print(f"[{self.name}] Registering logger to file at: " + file_name, style="bold blue") - + if not hasattr(Logger, 'notice_print'): + Console().print(f"\nRegistering loggers at: {os.path.abspath(os.path.dirname(file_name))}. System logs can be found in this directory.\n", style="bold red") + Logger.notice_print = True # Instantiate a rotating file handler with specified parameters file_handler = RotatingFileHandler( filename=file_name, @@ -232,11 +241,8 @@ class Logger(logging.Logger): name = "default" if name not in LOGGER_DICT: - LOGGER_DICT[name] = Logger(name=name, **kwargs) + logger_dir = kwargs.get('dir_path', 'log') + logger_dir = append_memoryscope_uuid(logger_dir) + LOGGER_DICT[name] = Logger(name=name, dir_path=logger_dir, **kwargs) return LOGGER_DICT[name] - - @staticmethod - def append_timestamp(name: str) -> str: - from memoryscope.core.memoryscope_context import get_ms_context - return f"{name}_{get_ms_context().memory_scope_uuid}" \ No newline at end of file diff --git a/memoryscope/core/worker/memory_manager.py b/memoryscope/core/worker/memory_manager.py index e43714b9..014fa280 100644 --- a/memoryscope/core/worker/memory_manager.py +++ b/memoryscope/core/worker/memory_manager.py @@ -24,7 +24,7 @@ class MemoryManager(object): # dict: key -> memory_id self._key_id_dict: Dict[str, List[str]] = {} - self.logger = Logger.get_logger(Logger.append_timestamp("memory_manager")) + self.logger = Logger.get_logger("memory_manager") self.workerflow_name = workerflow_name diff --git a/quick-start-demo.py b/quick-start-demo.py new file mode 100644 index 00000000..4987e6ba --- /dev/null +++ b/quick-start-demo.py @@ -0,0 +1,7 @@ +import os + +if os.environ.get('DASHSCOPE_API_KEY', None) is None and os.environ.get('API_KEY', None) is None: + os.environ['DASHSCOPE_API_KEY'] = input('Missing api key from dashscope ( `https://help.aliyun.com/zh/model-studio/developer-reference/get-api-key` ), please input key and press Enter:') + +from memoryscope import cli +cli() \ No newline at end of file diff --git a/quick_start_demo.py b/quick_start_demo.py deleted file mode 100644 index b6093536..00000000 --- a/quick_start_demo.py +++ /dev/null @@ -1,7 +0,0 @@ -import os - -if os.environ.get('DASHSCOPE_API_KEY', None) is None: - os.environ['DASHSCOPE_API_KEY'] = input('Missing api key from dashscope ( https:// ???? ), please input key and press Enter:') - -from memoryscope import cli -cli() \ No newline at end of file