scriptoriumadmin/ReMe
1
0
Fork 0
mirror of https://github.com/agentscope-ai/ReMe.git synced 2026-09-07 08:26:06 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 9

[dev] add emb/gene/rank json config

Browse source
This commit is contained in:
jinli.yl 2024-06-19 21:44:27 +08:00
parent 2a5878b578
commit f4e0e6bbc6
27 changed files with 373 additions and 500 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

0
config/__init__.py
View file

34
config/config.json
View file

@ -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"
}

0
config/db/__init__.py
View file

5
config/db/es.json
View file

@ -1,5 +0,0 @@
{
"index_name": "test_index",
"es_url": "http://localhost:9200"
}

5
config/model/dash_embedding.json
View file

@ -1,5 +0,0 @@
{
"model_name": "text-embedding-v2",
"method_type": "DashScopeEmbedding",
"clazz": "models.base_embedding_model"
}

4
config/model/dash_llm.json
View file

@ -1,4 +0,0 @@
{
"method": "DashScope",
"model_name": "qwen-max"
}

5
config/model/dash_reranker.json
View file

@ -1,5 +0,0 @@
{
"method": "DashScopeRerank",
"model_name": "gte-rerank",
"top_n": 3
}

5
config/model/dashscope_embedding.json Normal file
View file

@ -0,0 +1,5 @@
{
"clazz": "models.base_embedding_model",
"model_name": "text-embedding-v2",
"method_type": "DashScopeEmbedding"
}

5
config/model/dashscope_generation.json Normal file
View file

@ -0,0 +1,5 @@
{
"clazz": "models.base_generation_model",
"model_name": "qwen-max",
"method": "DashScope"
}

5
config/model/dashscope_rank.json Normal file
View file

@ -0,0 +1,5 @@
{
"clazz": "models.base_rank_model",
"model_name": "gte-rerank",
"method": "DashScopeRerank"
}

5
config/monitor/monitor.json
View file

@ -1,5 +0,0 @@
{
"model_name": "text-embedding-v2",
"method_type": "DashScopeEmbedding",
"clazz": "models.base_embedding_model"
}

150
config/worker.json
View file

@ -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
}
}
}

9
config/workers.json Normal file
View file

@ -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"
}
}

43
memory_scope/chat/base_memory_chat.py Normal file
View file

@ -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:
"""

11
memory_scope/chat/memory_chat.py
View file

@ -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

7
memory_scope/chat/memory_service.py
View file

@ -1,7 +0,0 @@
class MemoryService(object):
def __init__(self):
pass
def memory

11
memory_scope/cli.py
View file

@ -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__":

8
memory_scope/db/base_db_client.py
View file

@ -46,4 +46,10 @@ class BaseDBClient(metaclass=ABCMeta):
def delete(self):
"""
:return:
"""
"""
@abstractmethod
def flush(self):
"""
:return:
"""

28
memory_scope/handler/global_context.py Normal file
View file

@ -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()

67
memory_scope/handler/init_handler.py
View file

@ -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

151
memory_scope/handler/pipeline.py
View file

@ -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)

105
memory_scope/handler/pipeline_handler.py Normal file
View file

@ -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

78
memory_scope/job.py Normal file
View file

@ -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

2
memory_scope/models/base_generate_model.py → memory_scope/models/base_generation_model.py
View file

@ -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
])

31
memory_scope/utils/context_handler.py
View file

@ -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

6
memory_scope/utils/tool_functions.py
View file

@ -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

93
memory_scope/worker/base_worker.py
View file

@ -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,
}