diff --git a/memory_scope/cli.py b/memory_scope/cli.py index 23519e6f..c11a58bb 100644 --- a/memory_scope/cli.py +++ b/memory_scope/cli.py @@ -69,12 +69,13 @@ class CliJob(object): def run(self, config: str): self.load_config(config) - with G_CONTEXT.thread_pool: - memory_chat = list(G_CONTEXT.memory_chat_dict.values())[0] - memory_chat.run() + # with G_CONTEXT.thread_pool: + memory_chat = list(G_CONTEXT.memory_chat_dict.values())[0] + memory_chat.run() G_CONTEXT.memory_store.close() G_CONTEXT.monitor.close() + G_CONTEXT.thread_pool.shutdown() if __name__ == "__main__": diff --git a/memory_scope/memory/operation/base_operation.py b/memory_scope/memory/operation/base_operation.py index 24f3f536..64cc23c1 100644 --- a/memory_scope/memory/operation/base_operation.py +++ b/memory_scope/memory/operation/base_operation.py @@ -12,9 +12,15 @@ class BaseOperation(metaclass=ABCMeta): self.description: str = description self.kwargs: dict = kwargs - def init_workflow(self): + def init_workflow(self, **kwargs): pass @abstractmethod def run_operation(self, **kwargs): raise NotImplementedError + + def run_operation_backend(self): + pass + + def stop_operation_backend(self): + pass diff --git a/memory_scope/memory/operation/base_workflow.py b/memory_scope/memory/operation/base_workflow.py index 5f3fedb7..7144fda2 100644 --- a/memory_scope/memory/operation/base_workflow.py +++ b/memory_scope/memory/operation/base_workflow.py @@ -82,7 +82,7 @@ class BaseWorkflow(object): continue self.logger.info(f"----- print_workflow_{self.name}_end -----") - def init_workers(self): + def init_workers(self, is_backend: bool = False, **kwargs): for name in list(self.worker_dict.keys()): if name not in G_CONTEXT.worker_config: raise RuntimeError(f"worker={name} is not exists in worker_config!") @@ -91,9 +91,10 @@ class BaseWorkflow(object): config=G_CONTEXT.worker_config[name], suffix_name="worker", name=name, - is_multi_thread=self.worker_dict[name], + is_multi_thread=is_backend or self.worker_dict[name], context=self.context, - context_lock=self.context_lock) + context_lock=self.context_lock, + **kwargs) def _run_sub_workflow(self, worker_list: List[str]) -> bool: for name in worker_list: @@ -113,8 +114,7 @@ class BaseWorkflow(object): else: t_list = [] for sub_workflow in workflow_part: - t_list.append(G_CONTEXT.thread_pool.submit( - self._run_sub_workflow, sub_workflow)) + t_list.append(G_CONTEXT.thread_pool.submit(self._run_sub_workflow, sub_workflow)) flag = True for future in as_completed(t_list): diff --git a/memory_scope/memory/operation/read_memory.py b/memory_scope/memory/operation/read_memory.py index 0b8cd944..bb318a1e 100644 --- a/memory_scope/memory/operation/read_memory.py +++ b/memory_scope/memory/operation/read_memory.py @@ -21,8 +21,8 @@ class ReadMemory(BaseWorkflow, BaseOperation): self.chat_messages: List[Message] = chat_messages self.his_msg_count: int = his_msg_count - def init_workflow(self): - self.init_workers() + def init_workflow(self, **kwargs): + self.init_workers(**kwargs) def run_operation(self, **kwargs): max_count = 1 + self.his_msg_count diff --git a/memory_scope/memory/operation/summary_memory.py b/memory_scope/memory/operation/summary_memory.py index b713ebea..b8ddddbc 100644 --- a/memory_scope/memory/operation/summary_memory.py +++ b/memory_scope/memory/operation/summary_memory.py @@ -11,8 +11,8 @@ class SummaryMemory(BaseWorkflow, BaseBackendOperation): super().__init__(**kwargs) BaseBackendOperation.__init__(self, **kwargs) - def init_workflow(self): - self.init_workers() + def init_workflow(self, **kwargs): + self.init_workers(is_backend=True, **kwargs) def _run_operation(self, **kwargs): self.context[CHAT_KWARGS] = kwargs diff --git a/memory_scope/memory/operation/write_memory.py b/memory_scope/memory/operation/write_memory.py index e32eb7a3..120ba918 100644 --- a/memory_scope/memory/operation/write_memory.py +++ b/memory_scope/memory/operation/write_memory.py @@ -35,8 +35,8 @@ class WriteMemory(BaseWorkflow, BaseBackendOperation): for msg in self.chat_messages: msg.memorized = True - def init_workflow(self): - self.init_workers() + def init_workflow(self, **kwargs): + self.init_workers(is_backend=True, **kwargs) def _run_operation(self, **kwargs): self.context[CHAT_KWARGS] = kwargs diff --git a/memory_scope/memory/service/base_memory_service.py b/memory_scope/memory/service/base_memory_service.py index 17e18530..5076920f 100644 --- a/memory_scope/memory/service/base_memory_service.py +++ b/memory_scope/memory/service/base_memory_service.py @@ -33,7 +33,7 @@ class BaseMemoryService(metaclass=ABCMeta): def add_messages(self, messages: List[Message] | Message): raise NotImplementedError - def start_service(self): + def start_service(self, **kwargs): pass @abstractmethod diff --git a/memory_scope/memory/service/chat_memory_service.py b/memory_scope/memory/service/chat_memory_service.py index 118a1f9a..1da819e3 100644 --- a/memory_scope/memory/service/chat_memory_service.py +++ b/memory_scope/memory/service/chat_memory_service.py @@ -37,9 +37,9 @@ class ChatMemoryService(BaseMemoryService): for _ in range(gap_size): self.chat_messages.pop(0) - def start_service(self): + def start_service(self, **kwargs): for _, operation in self._operation_dict.items(): - operation.init_workflow() + operation.init_workflow(**kwargs) if operation.operation_type == "backend": operation.run_operation_backend() diff --git a/memory_scope/memory/worker/base_worker.py b/memory_scope/memory/worker/base_worker.py index d4812ada..aea78d37 100644 --- a/memory_scope/memory/worker/base_worker.py +++ b/memory_scope/memory/worker/base_worker.py @@ -32,14 +32,11 @@ class BaseWorker(metaclass=ABCMeta): raise RuntimeError(f"async_task is not allowed in multi_thread condition") self.task_list.append((fn, args, kwargs)) + async def _async_gather(self): + return await asyncio.gather(*[fn(*args, **kwargs) for fn, args, kwargs in self.task_list]) + def gather_async_result(self): - if self.is_multi_thread: - raise RuntimeError(f"async_task is not allowed in multi_thread condition") - - async def async_gather(): - return await asyncio.gather(*[fn(*args, **kwargs) for fn, args, kwargs in self.task_list]) - - results = asyncio.run(async_gather()) + results = asyncio.run(self._async_gather()) self.task_list.clear() return results diff --git a/test.py b/test.py deleted file mode 100644 index d24aa3fd..00000000 --- a/test.py +++ /dev/null @@ -1,11 +0,0 @@ -from memory_scope.cli import CliJob - - -def main(config_path: str): - job = CliJob() - job.run(config=config_path) - - -if __name__ == "__main__": - # fire.Fire(main) - main("config/config.yaml") diff --git a/tt.py b/tt.py deleted file mode 100644 index 8b315e87..00000000 --- a/tt.py +++ /dev/null @@ -1,31 +0,0 @@ -import asyncio - - -class TT(object): - def __init__(self): - self.task_list = [] - - async def async_func(self, i: int): - await asyncio.sleep(i) # 模拟异步操作 - print(f"函数{i}的结果") - - def submit_async_task(self, fn, *args, **kwargs): - self.task_list.append((fn, args, kwargs)) - - def gather_async_result(self): - async def async_gather(): - return await asyncio.gather(*[fn(*args, **kwargs) for fn, args, kwargs in self.task_list]) - - results = asyncio.run(async_gather()) - self.task_list.clear() - return results - - def run(self): - self.submit_async_task(self.async_func, i=1) - self.submit_async_task(self.async_func, i=2) - self.submit_async_task(self.async_func, i=3) - - self.gather_async_result() - - -TT().run()