From f4e0e6bbc6386e19e1c37e40eb5b5acd4ffde9c7 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Wed, 19 Jun 2024 21:44:27 +0800 Subject: [PATCH] [dev] add emb/gene/rank json config --- config/__init__.py | 0 config/config.json | 34 ++-- config/db/__init__.py | 0 config/db/es.json | 5 - config/model/dash_embedding.json | 5 - config/model/dash_llm.json | 4 - config/model/dash_reranker.json | 5 - config/model/dashscope_embedding.json | 5 + config/model/dashscope_generation.json | 5 + config/model/dashscope_rank.json | 5 + config/monitor/monitor.json | 5 - config/worker.json | 150 ----------------- config/workers.json | 9 ++ memory_scope/chat/base_memory_chat.py | 43 +++++ memory_scope/chat/memory_chat.py | 11 +- memory_scope/chat/memory_service.py | 7 - memory_scope/cli.py | 11 +- memory_scope/db/base_db_client.py | 8 +- memory_scope/handler/global_context.py | 28 ++++ memory_scope/handler/init_handler.py | 67 +++----- memory_scope/handler/pipeline.py | 151 ------------------ memory_scope/handler/pipeline_handler.py | 105 ++++++++++++ memory_scope/job.py | 78 +++++++++ ...rate_model.py => base_generation_model.py} | 2 +- memory_scope/utils/context_handler.py | 31 ---- memory_scope/utils/tool_functions.py | 6 + memory_scope/worker/base_worker.py | 93 +++-------- 27 files changed, 373 insertions(+), 500 deletions(-) delete mode 100644 config/__init__.py delete mode 100644 config/db/__init__.py delete mode 100644 config/db/es.json delete mode 100644 config/model/dash_embedding.json delete mode 100644 config/model/dash_llm.json delete mode 100644 config/model/dash_reranker.json create mode 100644 config/model/dashscope_embedding.json create mode 100644 config/model/dashscope_generation.json create mode 100644 config/model/dashscope_rank.json delete mode 100644 config/monitor/monitor.json delete mode 100644 config/worker.json create mode 100644 config/workers.json create mode 100644 memory_scope/chat/base_memory_chat.py delete mode 100644 memory_scope/chat/memory_service.py create mode 100644 memory_scope/handler/global_context.py delete mode 100644 memory_scope/handler/pipeline.py create mode 100644 memory_scope/handler/pipeline_handler.py create mode 100644 memory_scope/job.py rename memory_scope/models/{base_generate_model.py => base_generation_model.py} (94%) delete mode 100644 memory_scope/utils/context_handler.py diff --git a/config/__init__.py b/config/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/config/config.json b/config/config.json index ea868fe3..92b5e716 100644 --- a/config/config.json +++ b/config/config.json @@ -1,14 +1,26 @@ { - "thread_pool_max_count": "", - "worker": "config/worker.json", - "pipeline": { - "summary_short": "parse_params,[load_profile|es_new_obs|es_insight],[update_insight|get_reflection,get_insight|update_profile],summary_collect,memory_store", - "summary_long": "parse_params,info_filter,[es_today_obs|get_observation|get_observation_with_time],contra_repeat,memory_store", - "retrieve": "parse_params,load_profile,[extract_time|es_similar|es_keyword],semantic_rank,fuse_rerank" + "global_configs": { + "thread_pool_max_count": 5, + "dash_scope_apikey": "", + "open_ai_apikey": "", + "chat_list": [ + "memory_chat" + ] }, - "model_embedding": "config/model/dash_embedding.json", - "model_rerank": "config/model/dash_rerank.json", - "model_generate": "config/model/dash_generate.json", - "db": "", - "dash_api_key": "" + "memory_chat": { + "clazz": "chat.memory_chat", + "retrieve": "parse_params,load_profile,[extract_time|es_similar|es_keyword],semantic_rank,fuse_rerank", + "generation_model": "dashscope_generation" + }, + "db": { + "clazz": "db.base_db_client", + "index_name": "memory_test", + "user_name": "" + }, + "monitor": { + "clazz": "minitor.base_monitor", + "index_name": "memory_test", + "user_name": "" + }, + "workers": "workers" } \ No newline at end of file diff --git a/config/db/__init__.py b/config/db/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/config/db/es.json b/config/db/es.json deleted file mode 100644 index 6d33248b..00000000 --- a/config/db/es.json +++ /dev/null @@ -1,5 +0,0 @@ -{ - "index_name": "test_index", - "es_url": "http://localhost:9200" -} - diff --git a/config/model/dash_embedding.json b/config/model/dash_embedding.json deleted file mode 100644 index 70408e30..00000000 --- a/config/model/dash_embedding.json +++ /dev/null @@ -1,5 +0,0 @@ -{ - "model_name": "text-embedding-v2", - "method_type": "DashScopeEmbedding", - "clazz": "models.base_embedding_model" -} \ No newline at end of file diff --git a/config/model/dash_llm.json b/config/model/dash_llm.json deleted file mode 100644 index 07745be9..00000000 --- a/config/model/dash_llm.json +++ /dev/null @@ -1,4 +0,0 @@ -{ - "method": "DashScope", - "model_name": "qwen-max" -} \ No newline at end of file diff --git a/config/model/dash_reranker.json b/config/model/dash_reranker.json deleted file mode 100644 index b0f5593c..00000000 --- a/config/model/dash_reranker.json +++ /dev/null @@ -1,5 +0,0 @@ -{ - "method": "DashScopeRerank", - "model_name": "gte-rerank", - "top_n": 3 -} \ No newline at end of file diff --git a/config/model/dashscope_embedding.json b/config/model/dashscope_embedding.json new file mode 100644 index 00000000..ed5fb740 --- /dev/null +++ b/config/model/dashscope_embedding.json @@ -0,0 +1,5 @@ +{ + "clazz": "models.base_embedding_model", + "model_name": "text-embedding-v2", + "method_type": "DashScopeEmbedding" +} \ No newline at end of file diff --git a/config/model/dashscope_generation.json b/config/model/dashscope_generation.json new file mode 100644 index 00000000..7a08e58a --- /dev/null +++ b/config/model/dashscope_generation.json @@ -0,0 +1,5 @@ +{ + "clazz": "models.base_generation_model", + "model_name": "qwen-max", + "method": "DashScope" +} \ No newline at end of file diff --git a/config/model/dashscope_rank.json b/config/model/dashscope_rank.json new file mode 100644 index 00000000..a1014d84 --- /dev/null +++ b/config/model/dashscope_rank.json @@ -0,0 +1,5 @@ +{ + "clazz": "models.base_rank_model", + "model_name": "gte-rerank", + "method": "DashScopeRerank" +} \ No newline at end of file diff --git a/config/monitor/monitor.json b/config/monitor/monitor.json deleted file mode 100644 index 70408e30..00000000 --- a/config/monitor/monitor.json +++ /dev/null @@ -1,5 +0,0 @@ -{ - "model_name": "text-embedding-v2", - "method_type": "DashScopeEmbedding", - "clazz": "models.base_embedding_model" -} \ No newline at end of file diff --git a/config/worker.json b/config/worker.json deleted file mode 100644 index 32c31749..00000000 --- a/config/worker.json +++ /dev/null @@ -1,150 +0,0 @@ -{ - "es_insight": { - "name": "EsInsightWorker", - "path": "memory_scope/worker", - "es_insight_top_k": 128 - }, - "es_keyword": { - "name": "EsKeywordWorker", - "path": "memory_scope/worker", - "es_keyword_top_k": 10 - }, - "EsNewObsWorker": { - "name": "EsNewObsWorker", - "path": "memory_scope/worker", - "kwargs": { - "es_new_obs_top_k": 256 - } - }, - "EsNotReflectedWorker": { - "name": "EsNotReflectedWorker", - "path": "memory_scope/worker", - "kwargs": { - "es_not_reflected_top_k": 256 - } - }, - "EsSimilarWorker": { - "name": "EsSimilarWorker", - "path": "memory_scope/worker", - "kwargs": { - "es_similar_top_k": 128 - } - }, - "EsTodayObsWorker": { - "name": "EsTodayObsWorker", - "path": "memory_scope/worker", - "kwargs": { - "es_today_obs_top_k": 128 - } - }, - "GetInsightWorker": { - "name": "GetInsightWorker", - "path": "memory_scope/worker", - "kwargs": { - "es_insight_similar_top_k": 128, - "insight_obs_max_cnt": 10, - "get_insight_model": "qwen_max", - "get_insight_max_token": 500, - "get_insight_temperature": 0.6, - "get_insight_top_k": 1 - } - }, - "ExtractTimeWorker": { - "name": "ExtractTimeWorker", - "clazz": "worker.summary_long.get_insight", - "parse_time_model": "qwen_1_8_parse_time_service" - }, - "InfoFilterWorker": { - "name": "InfoFilterWorker", - "path": "memory_scope/worker", - "kwargs": { - "info_filter_msg_max_size": 200, - "info_filter_model": "qwen_max", - "info_filter_max_token": 200, - "info_filter_temperature": 0.6, - "info_filter_top_k": 1 - } - }, - "GetObservationWithTimeWorker": { - "name": "GetObservationWithTimeWorker", - "path": "memory_scope/worker", - "kwargs": { - "summary_messages_model": "qwen_max", - "summary_messages_max_token": 500, - "summary_messages_temperature": 0.6, - "summary_messages_top_k": 1 - } - }, - "GetObservationWorker": { - "name": "GetObservationWorker", - "path": "memory_scope/worker", - "kwargs": { - "summary_messages_model": "qwen_max", - "summary_messages_max_token": 500, - "summary_messages_temperature": 0.6, - "summary_messages_top_k": 1 - } - }, - "ContraRepeatWorker": { - "name": "ContraRepeatWorker", - "path": "memory_scope/worker", - "kwargs": { - "merge_obs_model": "qwen_max", - "merge_obs_max_token": 500, - "merge_obs_temperature": 0.6, - "merge_obs_top_k": 1 - } - }, - "FuseRerankWorker": { - "name": "FuseRerankWorker", - "path": "memory_scope/worker", - "kwargs": { - "fuse_score_threshold": 0.1, - "fuse_ratio_dict": { - "conversation": 0.8, - "observation": 1.0, - "obs_customized": 1.0, - "insight": 1.5, - "profile": 1.5, - "profile_customized": 1.5 - }, - "fuse_time_ratio": 2.0 - } - }, - "UpdateProfileWorker": { - "name": "UpdateProfileWorker", - "path": "memory_scope/worker", - "kwargs": { - "update_profile_threshold": 0.1, - "update_profile_model": "qwen_max", - "update_profile_max_token": 500, - "update_profile_temperature": 0.6, - "update_profile_top_k": 1, - "update_profile_max_thread": 10 - } - }, - "GetReflectionWorker": { - "name": "GetReflectionWorker", - "path": "memory_scope/worker", - "kwargs": { - "reflect_obs_cnt_threshold": 40, - "reflect_num_questions": 3, - "reflect_obs_model": "qwen_max", - "reflect_obs_max_token": 300, - "reflect_obs_temperature": 0.6, - "reflect_obs_top_k": 1 - } - }, - "UpdateInsightWorker": { - "name": "UpdateInsightWorker", - "path": "memory_scope/worker", - "kwargs": { - "update_insight_threshold": 0.1, - "update_insight_model": "qwen_max", - "update_insight_max_token": 500, - "update_insight_temperature": 0.6, - "update_insight_top_k": 1, - "update_insight_max_thread": 10 - } - } -} \ No newline at end of file diff --git a/config/workers.json b/config/workers.json new file mode 100644 index 00000000..0098cbaf --- /dev/null +++ b/config/workers.json @@ -0,0 +1,9 @@ +{ + "update_insight": { + "clazz": "worker.summary_long.update_insight", + "update_insight_model": "qwen_max", + "generation_model": "dashscope_generation", + "embedding_model": "dashscope_embedding", + "rank_model": "dashscope_rank" + } +} \ No newline at end of file diff --git a/memory_scope/chat/base_memory_chat.py b/memory_scope/chat/base_memory_chat.py new file mode 100644 index 00000000..79a20758 --- /dev/null +++ b/memory_scope/chat/base_memory_chat.py @@ -0,0 +1,43 @@ +from abc import ABCMeta, abstractmethod + +from memory_scope.constants.common_constants import RELATED_MEMORIES +from memory_scope.enumeration.memory_method_enum import MemoryMethodEnum +from memory_scope.handler.pipeline_handler import PipelineHandler + + +class BaseMemoryChat(metaclass=ABCMeta): + + def __init__(self, + user_name: str, + retrieve_pipeline: str, + summary_short_pipeline: str, + summary_long_pipeline: str, + **kwargs): + self.user_name: str = user_name + self.retrieve_pipeline_handler = PipelineHandler(user_name=user_name, + memory_method_type=MemoryMethodEnum.RETRIEVE, + pipeline_str=retrieve_pipeline) + + self.summary_short_pipeline_handler = PipelineHandler(user_name=user_name, + memory_method_type=MemoryMethodEnum.SUMMARY_SHORT, + pipeline_str=summary_short_pipeline) + + self.summary_long_pipeline_handler = PipelineHandler(user_name=user_name, + memory_method_type=MemoryMethodEnum.SUMMARY_LONG, + pipeline_str=summary_long_pipeline) + + def retrieve(self): + self.retrieve_pipeline_handler.run() + return self.retrieve_pipeline_handler.get_context(RELATED_MEMORIES, []) + + def summary_short(self): + self.summary_short_pipeline_handler.run() + + def summary_long(self): + self.summary_long_pipeline_handler.run() + + @abstractmethod + def chat(self): + """ + :return: + """ diff --git a/memory_scope/chat/memory_chat.py b/memory_scope/chat/memory_chat.py index dd95008c..e03dc686 100644 --- a/memory_scope/chat/memory_chat.py +++ b/memory_scope/chat/memory_chat.py @@ -1,8 +1,17 @@ +from memory_scope.chat.memory_service import MemoryService from memory_scope.handler.init_handler import InitHandler + + class MemoryChat(object): def __init__(self, init_handler: InitHandler): - pass + self.init_handler: InitHandler = init_handler + + self.memory_service: MemoryService = MemoryService( + retrieve_pipeline=init_handler.retrieve_pipeline, + summary_short_pipeline=init_handler.retrieve_pipeline, + summary_long_pipeline=init_handler.retrieve_pipeline, + ) def memory_retrieve(self): pass diff --git a/memory_scope/chat/memory_service.py b/memory_scope/chat/memory_service.py deleted file mode 100644 index c750621b..00000000 --- a/memory_scope/chat/memory_service.py +++ /dev/null @@ -1,7 +0,0 @@ - -class MemoryService(object): - - def __init__(self): - pass - - def memory \ No newline at end of file diff --git a/memory_scope/cli.py b/memory_scope/cli.py index fbe969e8..2ab398d3 100644 --- a/memory_scope/cli.py +++ b/memory_scope/cli.py @@ -1,15 +1,12 @@ import fire -from memory_scope.chat.memory_chat import MemoryChat -from memory_scope.handler.init_handler import InitHandler +from memory_scope.job import Job def main(config_path: str): - init_handler = InitHandler(config_path) - init_handler.init() - - memory_chat = MemoryChat(init_handler) - memory_chat.chat() + job = Job(config_path=config_path) + job.init_instance_by_config() + job.run() if __name__ == "__main__": diff --git a/memory_scope/db/base_db_client.py b/memory_scope/db/base_db_client.py index bfca8775..d8f53713 100644 --- a/memory_scope/db/base_db_client.py +++ b/memory_scope/db/base_db_client.py @@ -46,4 +46,10 @@ class BaseDBClient(metaclass=ABCMeta): def delete(self): """ :return: - """ \ No newline at end of file + """ + + @abstractmethod + def flush(self): + """ + :return: + """ diff --git a/memory_scope/handler/global_context.py b/memory_scope/handler/global_context.py new file mode 100644 index 00000000..594d386c --- /dev/null +++ b/memory_scope/handler/global_context.py @@ -0,0 +1,28 @@ +from concurrent.futures import ThreadPoolExecutor +from typing import Dict, Any + +from memory_scope.chat.base_memory_chat import BaseMemoryChat +from memory_scope.db.base_db_client import BaseDBClient +from memory_scope.models.base_model import BaseModel +from memory_scope.monitor.base_monitor import BaseMonitor +from memory_scope.worker.base_worker import BaseWorker + + +class GlobalContext(object): + def __init__(self): + self.global_configs: Dict[str, Any] = {} + + self.worker_dict: Dict[str, BaseWorker] = {} + + self.model_dict: Dict[str, BaseModel] = {} + + self.memory_chat_dict: Dict[str, BaseMemoryChat] = {} + + self.db_client: BaseDBClient | None = None + + self.monitor: BaseMonitor | None = None + + self.thread_pool: ThreadPoolExecutor | None = None + + +GLOBAL_CONTEXT = GlobalContext() diff --git a/memory_scope/handler/init_handler.py b/memory_scope/handler/init_handler.py index 71c25bc9..fa1979de 100644 --- a/memory_scope/handler/init_handler.py +++ b/memory_scope/handler/init_handler.py @@ -1,6 +1,6 @@ import json import os -import re +from concurrent.futures import ThreadPoolExecutor from typing import Dict, Any from memory_scope.db.base_db_client import BaseDBClient @@ -23,40 +23,41 @@ class InitHandler(object): self.model_dict: Dict[str, BaseModel] = {} self.db_client: BaseDBClient | None = None self.monitor: BaseMonitor | None = None + self.thread_pool: ThreadPoolExecutor | None = None self.worker_base_dir: str = "" self.model_base_dir: str = "" self.db_base_dir: str = "" self.minitor_base_dir: str = "" - self.retrieve_pipeline: list = [] - self.summary_short_pipeline: list = [] - self.summary_long_pipeline: list = [] + self.retrieve_pipeline: str = "" + self.summary_short_pipeline: str = "" + self.summary_long_pipeline: str = "" def init(self): with open(self.path) as f: self.config = json.load(f) - self.init_global_config(self.config["global"]) + self.retrieve_pipeline = self.config["pipelines"]["retrieve"] + self.summary_short_pipeline = self.config["pipelines"]["summary_short"] + self.summary_long_pipeline = self.config["pipelines"]["summary_long"] + self.global_configs = self.config["global"] + self.set_global_config() self.init_workers(self.config["workers"]) self.init_db(self.config["db"]) self.init_chat_model(self.config["chat_model"]) self.init_monitor(self.config["monitor"]) - self.retrieve_pipeline = self.parse_pipeline(self.config["pipelines"]["retrieve"]) - self.summary_short_pipeline = self.parse_pipeline(self.config["pipelines"]["summary_short"]) - self.summary_long_pipeline = self.parse_pipeline(self.config["pipelines"]["summary_long"]) - - def init_global_config(self, global_configs: Dict[str, str]): + def set_global_config(self): """set global_configs & set apikey into env """ - self.worker_base_dir = global_configs["worker_base_dir"] - self.model_base_dir = global_configs["model_base_dir"] - self.db_base_dir = global_configs["db_base_dir"] - self.minitor_base_dir = global_configs["minitor_base_dir"] - - # TODO sen + self.worker_base_dir = self.global_configs["worker_base_dir"] + self.model_base_dir = self.global_configs["model_base_dir"] + self.db_base_dir = self.global_configs["db_base_dir"] + self.minitor_base_dir = self.global_configs["minitor_base_dir"] + self.thread_pool = ThreadPoolExecutor(max_workers=int(self.global_configs["max_workers"])) + # TODO @ sen def init_workers(self, worker_config_name: str): """ load worker config & init workers @@ -66,7 +67,8 @@ class InitHandler(object): for worker_name, worker_config in worker_config_dict.items(): if worker_name in self.worker_dict: - continue + raise RuntimeError(f"worker_name={worker_name} is repeated!") + self.worker_dict[worker_name] = init_instance_by_config_v2(worker_config, default_clazz_path=self.worker_base_dir, suffix_name="worker", @@ -76,14 +78,6 @@ class InitHandler(object): self.init_model(worker_config.get("generation_model")) self.init_model(worker_config.get("rank_model")) - def init_model(self, model_name: str): - if not model_name or model_name in self.model_dict: - return - - with open(os.path.join(self.config_base_dir, "model", model_name)) as f: - model_config = json.load(f) - self.model_dict[model_name] = init_instance_by_config_v2(model_config, - default_clazz_path=self.model_base_dir) def init_db(self, db_config: dict): self.db_client = init_instance_by_config_v2(db_config, default_clazz_path=self.db_base_dir) @@ -94,26 +88,3 @@ class InitHandler(object): def init_monitor(self, monitor_config: dict): self.monitor = init_instance_by_config_v2(monitor_config, default_clazz_path=self.db_base_dir) - - @staticmethod - def parse_pipeline(pipeline_str: str) -> list: - # re-match e.g., [a|b],c,[d,e,f|g,h],j - pattern = r'(\[[^\]]*\]|[^,]+)' - pipeline_split = re.findall(pattern, pipeline_str) - - pipeline_list = [] - for pipeline_part in pipeline_split: - # e.g., [d,e,f|g,h] - pipeline_part = pipeline_part.strip() - if '[' in pipeline_part or ']' in pipeline_part: - pipeline_part = pipeline_part.replace('[', '').replace(']', '') - - # e.g., ["d,e,f", "g,h"] - line_split = [x.strip() for x in pipeline_part.split("|") if x] - if len(line_split) <= 0: - continue - - # e.g., ["d","e","f"] - pipeline_list.append([x.split(",") for x in line_split]) - - return pipeline_list diff --git a/memory_scope/handler/pipeline.py b/memory_scope/handler/pipeline.py deleted file mode 100644 index 0bd2dc55..00000000 --- a/memory_scope/handler/pipeline.py +++ /dev/null @@ -1,151 +0,0 @@ -import json -import re -import time -from concurrent.futures import ThreadPoolExecutor, as_completed -from importlib import import_module -from itertools import zip_longest -from typing import Dict, Any - -from worker.base_worker import BaseWorker -from utils.context_handler import ContextHandler -from utils.logger import Logger -from utils.timer import timer, Timer -from common.tool_functions import under_line_to_hump -from constants import common_constants -from constants.common_constants import RESPONSE_EXT_INFO, MAX_WORKERS, PIPELINE -from enumeration.memory_method_enum import MemoryMethodEnum -from pipeline.memory import MemoryServiceRequestModel -from cli.cli_config import C -from utils.tool_functions import init_instance_by_config - - -class Pipeline(object): - def __init__(self, method: MemoryMethodEnum): - self.method = method - self.context_handler = ContextHandler() - - # 线程池 - self.thread_pool = ThreadPoolExecutor(max_workers=C.thread_pool_max_count) - - # 全部初始化的worker - self.worker_dict: Dict[str, BaseWorker] = {} - - # 日志 - self.logger: Logger = Logger.get_memory_logger() - - # 初始化pipeline - self.pipeline_list = self.get_pipeline() - self.print_and_init_worker(self.pipeline_list) - - def get_worker(self, worker_name: str, is_multi_thread: bool = False) -> BaseWorker: - return init_instance_by_config( - config = C.worker.get(worker_name), - try_kwargs={ - "is_multi_thread": is_multi_thread, - "thread_pool": self.thread_pool - } - ) - - def worker_run(self, worker_list: list[str]) -> bool: - for worker_name in worker_list: - worker = self.worker_dict[worker_name] - # 执行子类实现的_run函数 - worker.run() - # 保存worker的运行信息 - self.run_infos.append(worker.run_info_dict) - # 结束pipeline - if not worker.continue_run: - return False - return True - - @timer - def print_and_init_worker(self, pipeline_list: list[list]): - self.logger.info("----- Pipeline Begin -----") - i: int = 0 - for pipeline_part in pipeline_list: - if len(pipeline_part) == 1: - for w in pipeline_part[0]: - self.logger.info(f"stage{i}: {w}") - self.worker_dict[w] = self.get_worker(w) - i += 1 - else: - for w_zip in zip_longest(*pipeline_part, fillvalue="-"): - self.logger.info(f"stage{i}: {' | '.join(w_zip)}") - i += 1 - for w in w_zip: - if w == "-": - continue - self.worker_dict[w] = self.get_worker(w, is_multi_thread=True) - self.logger.info("----- Pipeline End -----") - - def get_context(self, key: str, default=None) -> Any: - return self.context_handler.get_context(key, default) - - def flush(self, request: MemoryServiceRequestModel): - # 全局上下文，worker之间交换参数和变量 - self.context_handler.flush() - - # 运行信息 - self.run_infos = [] - self.context_handler.set_context(common_constants.REQUEST, request) - for pipeline_part in self.pipeline_list: - pipeline_part.flush(self.context_handler) - - @timer - def get_pipeline(self) -> list[list]: - pipeline_str = C.pipeline.get(self.method) - self.logger.info(f"pipeline={pipeline_str}") - - # re-match e.g., [a|b],c,[d,e,f|g,h],j - pattern = r'(\[[^\]]*\]|[^,]+)' - pipeline_split = re.findall(pattern, pipeline_str) - - pipeline_list = [] - for pipeline_part in pipeline_split: - # e.g., [d,e,f|g,h] - pipeline_part = pipeline_part.strip() - if '[' in pipeline_part or ']' in pipeline_part: - pipeline_part = pipeline_part.replace('[', '').replace(']', '') - - # e.g., ["d,e,f", "g,h"] - line_split = [x.strip() for x in pipeline_part.split("|") if x] - if len(line_split) <= 0: - continue - - # e.g., ["d","e","f"] - pipeline_list.append([x.split(",") for x in line_split]) - - return pipeline_list - - def run(self): - # run workers in multi threads - with self.thread_pool, Timer("ALL_PIPELINE"): - for pipeline_part in self.pipeline_list: - if len(pipeline_part) == 1: - if not self.worker_run(pipeline_part[0]): - break - elif self.max_workers == 1: - for worker_list in pipeline_part: - self.worker_run(worker_list) - else: - t_list = [] - for worker_list in pipeline_part: - time.sleep(0.001) - t_list.append(self.thread_pool.submit(self.worker_run, worker_list)) - - flag = True - for future in as_completed(t_list): - if not future.result(): - flag = False - break - if not flag: - break - - # 获取ext_info - ext_info = self.get_context(RESPONSE_EXT_INFO) - if ext_info is None: - ext_info = {} - self.context_handler.set_context(RESPONSE_EXT_INFO, ext_info) - - # 保存 run_info_list - ext_info["run_infos"] = json.dumps(self.run_infos, ensure_ascii=False) diff --git a/memory_scope/handler/pipeline_handler.py b/memory_scope/handler/pipeline_handler.py new file mode 100644 index 00000000..9982cb3d --- /dev/null +++ b/memory_scope/handler/pipeline_handler.py @@ -0,0 +1,105 @@ +import re +import threading +from concurrent.futures import as_completed +from itertools import zip_longest +from typing import Dict, Any + +from memory_scope.enumeration.memory_method_enum import MemoryMethodEnum +from memory_scope.handler.global_context import GLOBAL_CONTEXT +from memory_scope.utils.logger import Logger +from memory_scope.utils.timer import Timer + + +class PipelineHandler(object): + def __init__(self, user_name: str, memory_method_type: MemoryMethodEnum, pipeline_str: str): + + self.user_name: str = user_name + self.memory_method_type: MemoryMethodEnum = memory_method_type + self.pipeline_str: str = pipeline_str + + # 日志 + self.logger: Logger = Logger.get_logger() + + # pipeline上下文和锁 + self.context: Dict[str, Any] = {} + self.context_lock = threading.Lock() + + # 解析和打印 pipeline + self.pipeline_list: list[list] = [] + self._parse_pipeline() + self._print_pipeline() + + def _parse_pipeline(self): + # re-match e.g., [a|b],c,[d,e,f|g,h],j + pattern = r'(\[[^\]]*\]|[^,]+)' + pipeline_split = re.findall(pattern, self.pipeline_str) + + self.pipeline_list = [] + for pipeline_part in pipeline_split: + # e.g., [d,e,f|g,h] + pipeline_part = pipeline_part.strip() + if '[' in pipeline_part or ']' in pipeline_part: + pipeline_part = pipeline_part.replace('[', '').replace(']', '') + + # e.g., ["d,e,f", "g,h"] + line_split = [x.strip() for x in pipeline_part.split("|") if x] + if len(line_split) <= 0: + continue + + # e.g., ["d","e","f"] + self.pipeline_list.append([x.split(",") for x in line_split]) + + def _print_pipeline(self): + self.logger.info(f"----- {self.user_name} {self.memory_method_type.value} Pipeline Begin -----") + i: int = 0 + for pipeline_part in self.pipeline_list: + if len(pipeline_part) == 1: + for w in pipeline_part[0]: + self.logger.info(f"stage{i}: {w}") + i += 1 + GLOBAL_CONTEXT.worker_dict[w].context = self.context + else: + for w_zip in zip_longest(*pipeline_part, fillvalue="-"): + self.logger.info(f"stage{i}: {' | '.join(w_zip)}") + i += 1 + for w in w_zip: + if w == "-": + continue + GLOBAL_CONTEXT.worker_dict[w].is_multi_thread = True + GLOBAL_CONTEXT.worker_dict[w].context_lock = self.context_lock + GLOBAL_CONTEXT.worker_dict[w].context = self.context + self.logger.info(f"----- {self.user_name} {self.memory_method_type.value} Pipeline End -----") + + def get_context(self, key: str, default=None): + return self.context.get(key, default) + + def clear_context(self): + self.context.clear() + + @staticmethod + def worker_run(worker_list: list[str]) -> bool: + for worker_name in worker_list: + worker = GLOBAL_CONTEXT.worker_dict[worker_name] + worker.run() + if not worker.continue_run: + return False + return True + + def run(self): + with Timer(f"pipeline_{self.user_name}_{self.memory_method_type.value}"): + for pipeline_part in self.pipeline_list: + if len(pipeline_part) == 1: + if not self.worker_run(pipeline_part[0]): + break + else: + t_list = [] + for worker_list in pipeline_part: + t_list.append(GLOBAL_CONTEXT.thread_pool.submit(self.worker_run, worker_list)) + + flag = True + for future in as_completed(t_list): + if not future.result(): + flag = False + break + if not flag: + break diff --git a/memory_scope/job.py b/memory_scope/job.py new file mode 100644 index 00000000..253c85ca --- /dev/null +++ b/memory_scope/job.py @@ -0,0 +1,78 @@ +import json +import os +from concurrent.futures import ThreadPoolExecutor +from typing import Dict, Any + +from handler.global_context import GLOBAL_CONTEXT +from memory_scope.utils.logger import Logger +from memory_scope.utils.timer import Timer +from memory_scope.utils.tool_functions import complete_config_name, init_instance_by_config_v2 + + +class Job(object): + + def __init__(self, config_path: str): + self.config_path: str = config_path + self.config_base_dir: str = os.path.dirname(config_path) + + self.config: Dict[str, Any] = {} + + self.logger: Logger = Logger.get_logger("memory_chat") + + def init_memory_chat(self): + for chat in self.config["chat_list"]: + memory_chat_config = self.config[chat] + memory_chat = init_instance_by_config_v2(memory_chat_config) + GLOBAL_CONTEXT.memory_chat_dict[chat] = memory_chat + + generation_model = memory_chat_config["generation_model"] + self.init_model(generation_model) + + def init_model(self, model_name: str): + if not model_name or model_name in GLOBAL_CONTEXT.model_dict: + return + + with open(os.path.join(self.config_base_dir, "model", complete_config_name(model_name))) as f: + model_config = json.load(f) + GLOBAL_CONTEXT.model_dict[model_name] = init_instance_by_config_v2(model_config) + + def init_workers(self): + """ load worker config & init workers + """ + worker_config_name: str = self.config["workers"] + with open(os.path.join(self.config_base_dir, complete_config_name(worker_config_name))) as f: + worker_config_dict = json.load(f) + + for worker_name, worker_config in worker_config_dict.items(): + if worker_name in GLOBAL_CONTEXT.worker_dict: + raise RuntimeError(f"worker_name={worker_name} is repeated!") + + GLOBAL_CONTEXT.worker_dict[worker_name] = init_instance_by_config_v2(worker_config, + suffix_name="worker", + **GLOBAL_CONTEXT.global_configs) + + self.init_model(worker_config.get("embedding_model")) + self.init_model(worker_config.get("generation_model")) + self.init_model(worker_config.get("rank_model")) + + def set_global_config(self): + """set global_configs & set apikey into env + """ + + def init_instance_by_config(self): + with open(complete_config_name(self.config_path)) as f: + self.config = json.load(f) + + GLOBAL_CONTEXT.global_configs = self.config["global_configs"] + self.set_global_config() + GLOBAL_CONTEXT.thread_pool = ThreadPoolExecutor(max_workers=int(GLOBAL_CONTEXT.global_configs["max_workers"])) + + self.init_workers() + GLOBAL_CONTEXT.db_client = init_instance_by_config_v2(self.config["db"]) + GLOBAL_CONTEXT.monitor = init_instance_by_config_v2(self.config["monitor"]) + + self.init_memory_chat() + + def run(self): + with GLOBAL_CONTEXT.thread_pool, Timer("job"): + pass diff --git a/memory_scope/models/base_generate_model.py b/memory_scope/models/base_generation_model.py similarity index 94% rename from memory_scope/models/base_generate_model.py rename to memory_scope/models/base_generation_model.py index 2f0c3dbd..de38466f 100644 --- a/memory_scope/models/base_generate_model.py +++ b/memory_scope/models/base_generation_model.py @@ -5,7 +5,7 @@ from memory_scope.models.base_model import BaseModel from memory_scope.models.response import ModelResponse, ModelResponseGen -class BaseGenerateModel(BaseModel): +class BaseGenerationModel(BaseModel): MODEL_REGISTRY.batch_register([ DashScopeLLM ]) diff --git a/memory_scope/utils/context_handler.py b/memory_scope/utils/context_handler.py deleted file mode 100644 index a4ee67d4..00000000 --- a/memory_scope/utils/context_handler.py +++ /dev/null @@ -1,31 +0,0 @@ -import os -import threading -from typing import Dict, Any - -from utils.logger import Logger - -class ContextHandler(object): - def __init__(self): - # 上下文 所有worker共享 - self.context_dict: Dict[str, Any] = {} - - # 日志 - self.logger = Logger.get_logger() - - # 全局锁 - self.context_lock = threading.Lock() - - def flush(self): - self.context_dict: Dict[str, Any] = {} - - def get_context(self, key: str, default=None): - # 多线程环境下，如果是指针下修改，不安全 - return self.context_dict.get(key, default) - - def set_context(self, key: str, value: Any, is_multi_thread: bool = False): - if is_multi_thread: - # add lock to multi thread - with self.context_lock: - self.context_dict[key] = value - else: - self.context_dict[key] = value diff --git a/memory_scope/utils/tool_functions.py b/memory_scope/utils/tool_functions.py index 387d8052..6cd284be 100644 --- a/memory_scope/utils/tool_functions.py +++ b/memory_scope/utils/tool_functions.py @@ -175,3 +175,9 @@ def init_instance_by_config_v2(config: dict, default_clazz_path: str = "", suffi cls_name = under_line_to_hump(clazz_name) return getattr(module, cls_name)(**config, **kwargs) + + +def complete_config_name(config_name: str, suffix: str = ".json"): + if not config_name.endswith(suffix): + config_name += suffix + return config_name diff --git a/memory_scope/worker/base_worker.py b/memory_scope/worker/base_worker.py index 0cac7505..42feb437 100644 --- a/memory_scope/worker/base_worker.py +++ b/memory_scope/worker/base_worker.py @@ -1,33 +1,15 @@ -import time -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Any, List +from typing import Any, Dict -from utils.context_handler import ContextHandler -from utils.logger import Logger -from utils.timer import Timer +from memory_scope.utils.logger import Logger +from memory_scope.utils.timer import Timer class BaseWorker(object): - def __init__(self, - is_multi_thread: bool = False, - raise_exception: bool = True, - logger: Logger = None, - **kwargs): + def __init__(self, raise_exception: bool = True, **kwargs): super(BaseWorker, self).__init__(**kwargs) - - # 原始参数 - self.is_multi_thread: bool = is_multi_thread + # 异常是否继续执行 self.raise_exception: bool = raise_exception - self.logger: Logger = logger - - # 日志 - if not self.logger: - self.logger: Logger = Logger.get_logger() - self.logger.debug(f"init {self.__class__.__name__} is_multi_thread={is_multi_thread}") - - # 提交的线程池 - self.thread_list: list = [] # True 为正常运行，False会结束整个pipeline self.continue_run: bool = True @@ -35,33 +17,21 @@ class BaseWorker(object): # 短name self._name_simple: str = "" - def flush(self, context_handler: ContextHandler, thread_pool: ThreadPoolExecutor): - # 原始参数 - self.context_handler = context_handler - self.thread_pool: ThreadPoolExecutor = thread_pool + # 是否多线程环境 + self.is_multi_thread: bool = False - # 运行信息，保存到ext_info - self.run_infos: List[str] = [] + # pipeline 上下文 + self.context: Dict[str, Any] | None = None + self.context_lock = None - # 运行时间 - self.run_cost: float = 0 + # 日志 + self.logger: Logger = Logger.get_logger() + + # worker 参数 + self.kwargs: dict = kwargs def _run(self): - pass - - def submit_thread(self, fn, /, *args, sleep_time: float = 0, **kwargs): - if self.thread_list: - time.sleep(sleep_time) - t = self.thread_pool.submit(fn, *args, **kwargs) - self.thread_list.append(t) - return t - - def join_threads(self): - result_list = [] - for future in as_completed(self.thread_list): - result_list.append(future.result()) - self.thread_list.clear() - return result_list + raise NotImplementedError def run(self): self.logger.info(f"----- Begin {self.name_simple} -----") @@ -72,16 +42,20 @@ class BaseWorker(object): try: self._run() except Exception as e: - self.add_run_info(f"run {self.name_simple} failed! args={e.args}") + self.logger.exception(f"run {self.name_simple} failed! args={e.args}") - self.run_cost = t.cost - self.logger.info(f"----- End {self.name_simple} {t.get_cost_info()}-----") + self.logger.info(f"----- End {self.name_simple} cost={t.cost_str}-----") def get_context(self, key: str, default=None): - return self.context_handler.get_context(key, default) + return self.context.get(key, default) def set_context(self, key: str, value: Any): - self.context_handler.set_context(key, value, self.is_multi_thread) + if self.is_multi_thread: + # add lock to multi thread + with self.context_lock: + self.context[key] = value + else: + self.context[key] = value @property def name_simple(self) -> str: @@ -89,20 +63,3 @@ class BaseWorker(object): self._name_simple = self.__class__.__name__.replace("Worker", "") return self._name_simple - def add_run_info(self, msg: str, log_warning: bool = True, continue_run: bool = True): - if not continue_run: - self.continue_run = False - msg = f"{msg} pipeline is ended by {self.name_simple}!" - - if log_warning: - self.logger.warning(msg, stacklevel=2) - - self.run_infos.append(msg) - - @property - def run_info_dict(self): - return { - "name": self.name_simple, - "cost": self.run_cost, - "info": self.run_infos, - }