From 3cf388b1a4bd5ca5eab440b98ddc0c6bb274cf88 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Mon, 17 Jun 2024 20:12:58 +0800 Subject: [PATCH] [dev] init dir & config --- README.md | 4 +- README_ZH.md | 5 + docs/README.md | 20 + docs/__init__.py | 0 examples/__init__.py | 0 memory_scope/__init__.py | 0 memory_scope/_version.py | 4 + memory_scope/agent/__init__.py | 0 memory_scope/agent/chat_agent.py | 19 + memory_scope/cli/__init__.py | 0 memory_scope/cli/cli.py | 719 +++++++++ memory_scope/cli/cli_config.py | 1294 +++++++++++++++++ memory_scope/config/__init__.py | 0 memory_scope/config/bailian_memory_config.py | 137 ++ memory_scope/constants/__init__.py | 0 memory_scope/constants/common_constants.py | 116 ++ memory_scope/db/__init__.py | 0 memory_scope/db/elastic_search_client.py | 341 +++++ memory_scope/enumeration/__init__.py | 0 memory_scope/enumeration/dash_api_enum.py | 9 + .../enumeration/memory_method_enum.py | 11 + .../enumeration/memory_node_status.py | 7 + .../enumeration/memory_recall_type.py | 9 + memory_scope/enumeration/memory_scene_enum.py | 11 + memory_scope/enumeration/memory_type_enum.py | 15 + memory_scope/enumeration/message_role_enum.py | 9 + memory_scope/models/__init__.py | 0 memory_scope/models/dash_client.py | 91 ++ memory_scope/models/dash_embedding_client.py | 43 + memory_scope/models/dash_generate_client.py | 45 + memory_scope/models/dash_rerank_client.py | 43 + memory_scope/node/__init__.py | 0 memory_scope/node/memory_node.py | 71 + memory_scope/node/memory_wrap_node.py | 33 + memory_scope/node/message.py | 11 + memory_scope/node/user_attribute.py | 36 + memory_scope/parsers/__init__.py | 0 memory_scope/parsers/response_text_parser.py | 33 + memory_scope/pipeline/__init__.py | 0 memory_scope/pipeline/memory.py | 30 + .../pipeline/memory_service_bailian.py | 162 +++ memory_scope/prompts/__init__.py | 0 memory_scope/prompts/bailian_prompt_config.py | 626 ++++++++ memory_scope/utils/__init__.py | 0 memory_scope/utils/context_handler.py | 129 ++ memory_scope/utils/logger.py | 102 ++ memory_scope/utils/timer.py | 76 + memory_scope/utils/tool_functions.py | 164 +++ memory_scope/utils/user_profile_handler.py | 102 ++ memory_scope/worker/__init__.py | 0 memory_scope/worker/base_worker.py | 126 ++ memory_scope/worker/es/__init__.py | 0 memory_scope/worker/es/es_insight_worker.py | 23 + memory_scope/worker/es/es_keyword_worker.py | 53 + memory_scope/worker/es/es_new_obs_worker.py | 24 + .../worker/es/es_not_reflected_worker.py | 24 + .../worker/es/es_retrieve_all_worker.py | 30 + memory_scope/worker/es/es_similar_worker.py | 80 + memory_scope/worker/es/es_today_obs_worker.py | 29 + memory_scope/worker/es/load_profile_worker.py | 34 + memory_scope/worker/memory_base_worker.py | 132 ++ memory_scope/worker/memory_store_worker.py | 34 + memory_scope/worker/parse_params_worker.py | 36 + memory_scope/worker/retrieve/__init__.py | 0 .../worker/retrieve/extract_time_worker.py | 65 + .../worker/retrieve/fuse_rerank_worker.py | 115 ++ .../worker/retrieve/semantic_rank_worker.py | 68 + memory_scope/worker/summary_long/__init__.py | 0 .../worker/summary_long/get_insight_worker.py | 132 ++ .../summary_long/get_reflection_worker.py | 72 + .../summary_long/long_contra_repeat_worker.py | 111 ++ .../summary_long/summary_collect_worker.py | 33 + .../summary_long/update_insight_worker.py | 160 ++ .../summary_long/update_profile_worker.py | 197 +++ memory_scope/worker/summary_short/__init__.py | 0 .../summary_short/contra_repeat_worker.py | 92 ++ .../get_observation_with_time_worker.py | 128 ++ .../summary_short/get_observation_worker.py | 116 ++ .../summary_short/info_filter_worker.py | 57 + tests/__init__.py | 0 tests/es_test2.py | 89 ++ tests/test_dash_api.py | 321 ++++ tests/test_memory.py | 256 ++++ 83 files changed, 7163 insertions(+), 1 deletion(-) create mode 100644 docs/README.md create mode 100644 docs/__init__.py create mode 100644 examples/__init__.py create mode 100644 memory_scope/__init__.py create mode 100644 memory_scope/_version.py create mode 100644 memory_scope/agent/__init__.py create mode 100644 memory_scope/agent/chat_agent.py create mode 100644 memory_scope/cli/__init__.py create mode 100644 memory_scope/cli/cli.py create mode 100644 memory_scope/cli/cli_config.py create mode 100644 memory_scope/config/__init__.py create mode 100644 memory_scope/config/bailian_memory_config.py create mode 100644 memory_scope/constants/__init__.py create mode 100644 memory_scope/constants/common_constants.py create mode 100644 memory_scope/db/__init__.py create mode 100644 memory_scope/db/elastic_search_client.py create mode 100644 memory_scope/enumeration/__init__.py create mode 100644 memory_scope/enumeration/dash_api_enum.py create mode 100644 memory_scope/enumeration/memory_method_enum.py create mode 100644 memory_scope/enumeration/memory_node_status.py create mode 100644 memory_scope/enumeration/memory_recall_type.py create mode 100644 memory_scope/enumeration/memory_scene_enum.py create mode 100644 memory_scope/enumeration/memory_type_enum.py create mode 100644 memory_scope/enumeration/message_role_enum.py create mode 100644 memory_scope/models/__init__.py create mode 100644 memory_scope/models/dash_client.py create mode 100644 memory_scope/models/dash_embedding_client.py create mode 100644 memory_scope/models/dash_generate_client.py create mode 100644 memory_scope/models/dash_rerank_client.py create mode 100644 memory_scope/node/__init__.py create mode 100644 memory_scope/node/memory_node.py create mode 100644 memory_scope/node/memory_wrap_node.py create mode 100644 memory_scope/node/message.py create mode 100644 memory_scope/node/user_attribute.py create mode 100644 memory_scope/parsers/__init__.py create mode 100644 memory_scope/parsers/response_text_parser.py create mode 100644 memory_scope/pipeline/__init__.py create mode 100644 memory_scope/pipeline/memory.py create mode 100644 memory_scope/pipeline/memory_service_bailian.py create mode 100644 memory_scope/prompts/__init__.py create mode 100644 memory_scope/prompts/bailian_prompt_config.py create mode 100644 memory_scope/utils/__init__.py create mode 100644 memory_scope/utils/context_handler.py create mode 100644 memory_scope/utils/logger.py create mode 100644 memory_scope/utils/timer.py create mode 100644 memory_scope/utils/tool_functions.py create mode 100644 memory_scope/utils/user_profile_handler.py create mode 100644 memory_scope/worker/__init__.py create mode 100644 memory_scope/worker/base_worker.py create mode 100644 memory_scope/worker/es/__init__.py create mode 100644 memory_scope/worker/es/es_insight_worker.py create mode 100644 memory_scope/worker/es/es_keyword_worker.py create mode 100644 memory_scope/worker/es/es_new_obs_worker.py create mode 100644 memory_scope/worker/es/es_not_reflected_worker.py create mode 100644 memory_scope/worker/es/es_retrieve_all_worker.py create mode 100644 memory_scope/worker/es/es_similar_worker.py create mode 100644 memory_scope/worker/es/es_today_obs_worker.py create mode 100644 memory_scope/worker/es/load_profile_worker.py create mode 100644 memory_scope/worker/memory_base_worker.py create mode 100644 memory_scope/worker/memory_store_worker.py create mode 100644 memory_scope/worker/parse_params_worker.py create mode 100644 memory_scope/worker/retrieve/__init__.py create mode 100644 memory_scope/worker/retrieve/extract_time_worker.py create mode 100644 memory_scope/worker/retrieve/fuse_rerank_worker.py create mode 100644 memory_scope/worker/retrieve/semantic_rank_worker.py create mode 100644 memory_scope/worker/summary_long/__init__.py create mode 100644 memory_scope/worker/summary_long/get_insight_worker.py create mode 100644 memory_scope/worker/summary_long/get_reflection_worker.py create mode 100644 memory_scope/worker/summary_long/long_contra_repeat_worker.py create mode 100644 memory_scope/worker/summary_long/summary_collect_worker.py create mode 100644 memory_scope/worker/summary_long/update_insight_worker.py create mode 100644 memory_scope/worker/summary_long/update_profile_worker.py create mode 100644 memory_scope/worker/summary_short/__init__.py create mode 100644 memory_scope/worker/summary_short/contra_repeat_worker.py create mode 100644 memory_scope/worker/summary_short/get_observation_with_time_worker.py create mode 100644 memory_scope/worker/summary_short/get_observation_worker.py create mode 100644 memory_scope/worker/summary_short/info_filter_worker.py create mode 100644 tests/__init__.py create mode 100644 tests/es_test2.py create mode 100644 tests/test_dash_api.py create mode 100644 tests/test_memory.py diff --git a/README.md b/README.md index 75177bc6..d68b146e 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,5 @@ English | [**中文**](./README_ZH.md) -# ModelScope \ No newline at end of file +# ModelScope + +https://github.com/modelscope/agentscope/blob/main/README.md \ No newline at end of file diff --git a/README_ZH.md b/README_ZH.md index e69de29b..ac906560 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -0,0 +1,5 @@ +[**English**](./README.md) | 中文 + +# ModelScope + +https://github.com/modelscope/agentscope/blob/main/README_ZH.md \ No newline at end of file diff --git a/docs/README.md b/docs/README.md new file mode 100644 index 00000000..77f713f4 --- /dev/null +++ b/docs/README.md @@ -0,0 +1,20 @@ +# MemoryScope Documentation + +## Build Documentation + +Please use the following commands to build sphinx doc of MemoryScope. + +参考：https://github.com/modelscope/agentscope/blob/main/docs/README.md + +```shell +# step 1: Install dependencies +pip install sphinx sphinx-autobuild sphinx_rtd_theme myst-parser sphinxcontrib-mermaid + +# step 2: go into the sphinx_doc dir +cd sphinx_doc + +# step 3: build the sphinx doc +./build_sphinx_doc.sh + +# step 4: view sphinx_doc/build/html/index.html using your browser +``` \ No newline at end of file diff --git a/docs/__init__.py b/docs/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/__init__.py b/examples/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/memory_scope/__init__.py b/memory_scope/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/memory_scope/_version.py b/memory_scope/_version.py new file mode 100644 index 00000000..eedde6fb --- /dev/null +++ b/memory_scope/_version.py @@ -0,0 +1,4 @@ +# -*- coding: utf-8 -*- +""" Version of MemoryScope.""" + +__version__ = "0.1alpha1" diff --git a/memory_scope/agent/__init__.py b/memory_scope/agent/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/memory_scope/agent/chat_agent.py b/memory_scope/agent/chat_agent.py new file mode 100644 index 00000000..7efa81fa --- /dev/null +++ b/memory_scope/agent/chat_agent.py @@ -0,0 +1,19 @@ +class ChatAgent(object): + + def __init__(self): + pass + + def memory_retrieve(self): + pass + + def memory_summary_short(self): + pass + + def memory_summary_long(self): + pass + + def chat(self): + pass + + def chat_with_memory(self): + pass diff --git a/memory_scope/cli/__init__.py b/memory_scope/cli/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/memory_scope/cli/cli.py b/memory_scope/cli/cli.py new file mode 100644 index 00000000..12ac3a70 --- /dev/null +++ b/memory_scope/cli/cli.py @@ -0,0 +1,719 @@ +import json +import logging +import os +import subprocess +import sys +import uuid +from enum import Enum +from pathlib import Path +from typing import Annotated, Optional + +import questionary +import requests +import typer + +import memgpt.utils as utils +from memgpt.agent import Agent, save_agent +from memgpt.cli.cli_config import configure +from memgpt.config import MemGPTConfig +from memgpt.constants import CLI_WARNING_PREFIX, MEMGPT_DIR +from memgpt.credentials import MemGPTCredentials +from memgpt.data_types import EmbeddingConfig, LLMConfig, User +from memgpt.log import logger +from memgpt.metadata import MetadataStore +from memgpt.migrate import migrate_all_agents, migrate_all_sources +from memgpt.server.constants import WS_DEFAULT_PORT + +# from memgpt.interface import CLIInterface as interface # for printing to terminal +from memgpt.streaming_interface import ( + StreamingRefreshCLIInterface as interface, # for printing to terminal +) +from memgpt.utils import open_folder_in_explorer, printd + + +def migrate( + debug: Annotated[bool, typer.Option(help="Print extra tracebacks for failed migrations")] = False, +): + """Migrate old agents (pre 0.2.12) to the new database system""" + migrate_all_agents(debug=debug) + migrate_all_sources(debug=debug) + + +class QuickstartChoice(Enum): + openai = "openai" + # azure = "azure" + memgpt_hosted = "memgpt" + + +def str_to_quickstart_choice(choice_str: str) -> QuickstartChoice: + try: + return QuickstartChoice[choice_str] + except KeyError: + valid_options = [choice.name for choice in QuickstartChoice] + raise ValueError(f"{choice_str} is not a valid QuickstartChoice. Valid options are: {valid_options}") + + +def set_config_with_dict(new_config: dict) -> (MemGPTConfig, bool): + """_summary_ + + Args: + new_config (dict): Dict of new config values + + Returns: + new_config MemGPTConfig, modified (bool): Returns the new config and a boolean indicating if the config was modified + """ + from memgpt.utils import printd + + old_config = MemGPTConfig.load() + modified = False + for k, v in vars(old_config).items(): + if k in new_config: + if v != new_config[k]: + printd(f"Replacing config {k}: {v} -> {new_config[k]}") + modified = True + # old_config[k] = new_config[k] + setattr(old_config, k, new_config[k]) # Set the new value using dot notation + else: + printd(f"Skipping new config {k}: {v} == {new_config[k]}") + + # update embedding config + if old_config.default_embedding_config: + for k, v in vars(old_config.default_embedding_config).items(): + if k in new_config: + if v != new_config[k]: + printd(f"Replacing config {k}: {v} -> {new_config[k]}") + modified = True + # old_config[k] = new_config[k] + setattr(old_config.default_embedding_config, k, new_config[k]) + else: + printd(f"Skipping new config {k}: {v} == {new_config[k]}") + else: + modified = True + fields = ["embedding_model", "embedding_dim", "embedding_chunk_size", "embedding_endpoint", "embedding_endpoint_type"] + args = {} + for field in fields: + if field in new_config: + args[field] = new_config[field] + printd(f"Setting new config {field}: {new_config[field]}") + old_config.default_embedding_config = EmbeddingConfig(**args) + + # update llm config + if old_config.default_llm_config: + for k, v in vars(old_config.default_llm_config).items(): + if k in new_config: + if v != new_config[k]: + printd(f"Replacing config {k}: {v} -> {new_config[k]}") + modified = True + # old_config[k] = new_config[k] + setattr(old_config.default_llm_config, k, new_config[k]) + else: + printd(f"Skipping new config {k}: {v} == {new_config[k]}") + else: + modified = True + fields = ["model", "model_endpoint", "model_endpoint_type", "model_wrapper", "context_window"] + args = {} + for field in fields: + if field in new_config: + args[field] = new_config[field] + printd(f"Setting new config {field}: {new_config[field]}") + old_config.default_llm_config = LLMConfig(**args) + return (old_config, modified) + + +def quickstart( + backend: Annotated[QuickstartChoice, typer.Option(help="Quickstart setup backend")] = "memgpt", + latest: Annotated[bool, typer.Option(help="Use --latest to pull the latest config from online")] = False, + debug: Annotated[bool, typer.Option(help="Use --debug to enable debugging output")] = False, + terminal: bool = True, +): + """Set the base config file with a single command + + This function and `configure` should be the ONLY places where MemGPTConfig.save() is called. + """ + + # setup logger + utils.DEBUG = debug + logging.getLogger().setLevel(logging.CRITICAL) + if debug: + logging.getLogger().setLevel(logging.DEBUG) + + # make sure everything is set up properly + MemGPTConfig.create_config_dir() + credentials = MemGPTCredentials.load() + + config_was_modified = False + if backend == QuickstartChoice.memgpt_hosted: + # if latest, try to pull the config from the repo + # fallback to using local + if latest: + # Download the latest memgpt hosted config + url = "https://raw.githubusercontent.com/cpacker/MemGPT/main/configs/memgpt_hosted.json" + response = requests.get(url) + + # Check if the request was successful + if response.status_code == 200: + # Parse the response content as JSON + config = response.json() + # Output a success message and the first few items in the dictionary as a sample + printd("JSON config file downloaded successfully.") + new_config, config_was_modified = set_config_with_dict(config) + else: + typer.secho(f"Failed to download config from {url}. Status code: {response.status_code}", fg=typer.colors.RED) + + # Load the file from the relative path + script_dir = os.path.dirname(__file__) # Get the directory where the script is located + backup_config_path = os.path.join(script_dir, "..", "configs", "memgpt_hosted.json") + try: + with open(backup_config_path, "r", encoding="utf-8") as file: + backup_config = json.load(file) + printd("Loaded backup config file successfully.") + new_config, config_was_modified = set_config_with_dict(backup_config) + except FileNotFoundError: + typer.secho(f"Backup config file not found at {backup_config_path}", fg=typer.colors.RED) + return + else: + # Load the file from the relative path + script_dir = os.path.dirname(__file__) # Get the directory where the script is located + # print("SCRIPT", script_dir) + backup_config_path = os.path.join(script_dir, "..", "configs", "memgpt_hosted.json") + # print("FILE PATH", backup_config_path) + try: + with open(backup_config_path, "r", encoding="utf-8") as file: + backup_config = json.load(file) + # print(backup_config) + printd("Loaded config file successfully.") + new_config, config_was_modified = set_config_with_dict(backup_config) + except FileNotFoundError: + typer.secho(f"Config file not found at {backup_config_path}", fg=typer.colors.RED) + return + + elif backend == QuickstartChoice.openai: + # Make sure we have an API key + api_key = os.getenv("OPENAI_API_KEY") + while api_key is None or len(api_key) == 0: + # Ask for API key as input + api_key = questionary.password("Enter your OpenAI API key (starts with 'sk-', see https://platform.openai.com/api-keys):").ask() + credentials.openai_key = api_key + credentials.save() + + # if latest, try to pull the config from the repo + # fallback to using local + if latest: + url = "https://raw.githubusercontent.com/cpacker/MemGPT/main/configs/openai.json" + response = requests.get(url) + + # Check if the request was successful + if response.status_code == 200: + # Parse the response content as JSON + config = response.json() + # Output a success message and the first few items in the dictionary as a sample + print("JSON config file downloaded successfully.") + new_config, config_was_modified = set_config_with_dict(config) + else: + typer.secho(f"Failed to download config from {url}. Status code: {response.status_code}", fg=typer.colors.RED) + + # Load the file from the relative path + script_dir = os.path.dirname(__file__) # Get the directory where the script is located + backup_config_path = os.path.join(script_dir, "..", "configs", "openai.json") + try: + with open(backup_config_path, "r", encoding="utf-8") as file: + backup_config = json.load(file) + printd("Loaded backup config file successfully.") + new_config, config_was_modified = set_config_with_dict(backup_config) + except FileNotFoundError: + typer.secho(f"Backup config file not found at {backup_config_path}", fg=typer.colors.RED) + return + else: + # Load the file from the relative path + script_dir = os.path.dirname(__file__) # Get the directory where the script is located + backup_config_path = os.path.join(script_dir, "..", "configs", "openai.json") + try: + with open(backup_config_path, "r", encoding="utf-8") as file: + backup_config = json.load(file) + printd("Loaded config file successfully.") + new_config, config_was_modified = set_config_with_dict(backup_config) + except FileNotFoundError: + typer.secho(f"Config file not found at {backup_config_path}", fg=typer.colors.RED) + return + + else: + raise NotImplementedError(backend) + + if config_was_modified: + printd(f"Saving new config file.") + new_config.save() + typer.secho(f"📖 MemGPT configuration file updated!", fg=typer.colors.GREEN) + typer.secho( + "\n".join( + [ + f"🧠 model\t-> {new_config.default_llm_config.model}", + f"🖥️ endpoint\t-> {new_config.default_llm_config.model_endpoint}", + ] + ), + fg=typer.colors.GREEN, + ) + else: + typer.secho(f"📖 MemGPT configuration file unchanged.", fg=typer.colors.WHITE) + typer.secho( + "\n".join( + [ + f"🧠 model\t-> {new_config.default_llm_config.model}", + f"🖥️ endpoint\t-> {new_config.default_llm_config.model_endpoint}", + ] + ), + fg=typer.colors.WHITE, + ) + + # 'terminal' = quickstart was run alone, in which case we should guide the user on the next command + if terminal: + if config_was_modified: + typer.secho('⚡ Run "memgpt run" to create an agent with the new config.', fg=typer.colors.YELLOW) + else: + typer.secho('⚡ Run "memgpt run" to create an agent.', fg=typer.colors.YELLOW) + + +def open_folder(): + """Open a folder viewer of the MemGPT home directory""" + try: + print(f"Opening home folder: {MEMGPT_DIR}") + open_folder_in_explorer(MEMGPT_DIR) + except Exception as e: + print(f"Failed to open folder with system viewer, error:\n{e}") + + +class ServerChoice(Enum): + rest_api = "rest" + ws_api = "websocket" + + +def create_default_user_or_exit(config: MemGPTConfig, ms: MetadataStore): + user_id = uuid.UUID(config.anon_clientid) + user = ms.get_user(user_id=user_id) + if user is None: + ms.create_user(User(id=user_id)) + user = ms.get_user(user_id=user_id) + if user is None: + typer.secho(f"Failed to create default user in database.", fg=typer.colors.RED) + sys.exit(1) + else: + return user + else: + return user + + +def server( + type: Annotated[ServerChoice, typer.Option(help="Server to run")] = "rest", + port: Annotated[Optional[int], typer.Option(help="Port to run the server on")] = None, + host: Annotated[Optional[str], typer.Option(help="Host to run the server on (default to localhost)")] = None, + use_ssl: Annotated[bool, typer.Option(help="Run the server using HTTPS?")] = False, + ssl_cert: Annotated[Optional[str], typer.Option(help="Path to SSL certificate (if use_ssl is True)")] = None, + ssl_key: Annotated[Optional[str], typer.Option(help="Path to SSL key file (if use_ssl is True)")] = None, + debug: Annotated[bool, typer.Option(help="Turn debugging output on")] = False, +): + """Launch a MemGPT server process""" + + if type == ServerChoice.rest_api: + pass + + if MemGPTConfig.exists(): + config = MemGPTConfig.load() + ms = MetadataStore(config) + create_default_user_or_exit(config, ms) + else: + typer.secho(f"No configuration exists. Run memgpt configure before starting the server.", fg=typer.colors.RED) + sys.exit(1) + + try: + from memgpt.server.rest_api.server import start_server + + start_server( + port=port, + host=host, + use_ssl=use_ssl, + ssl_cert=ssl_cert, + ssl_key=ssl_key, + debug=debug, + ) + + except KeyboardInterrupt: + # Handle CTRL-C + typer.secho("Terminating the server...") + sys.exit(0) + + elif type == ServerChoice.ws_api: + if debug: + from memgpt.server.server import logger as server_logger + + # Set the logging level + server_logger.setLevel(logging.DEBUG) + # Create a StreamHandler + stream_handler = logging.StreamHandler() + # Set the formatter (optional) + formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s") + stream_handler.setFormatter(formatter) + # Add the handler to the logger + server_logger.addHandler(stream_handler) + + if port is None: + port = WS_DEFAULT_PORT + + # Change to the desired directory + script_path = Path(__file__).resolve() + script_dir = script_path.parent + + server_directory = os.path.join(script_dir.parent, "server", "ws_api") + command = f"python server.py {port}" + + # Run the command + typer.secho(f"Running WS (websockets) server: {command} (inside {server_directory})") + + process = None + try: + # Start the subprocess in a new session + process = subprocess.Popen(command, shell=True, start_new_session=True, cwd=server_directory) + process.wait() + except KeyboardInterrupt: + # Handle CTRL-C + if process is not None: + typer.secho("Terminating the server...") + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + typer.secho("Server terminated with kill()") + sys.exit(0) + + +def run( + persona: Annotated[Optional[str], typer.Option(help="Specify persona")] = None, + agent: Annotated[Optional[str], typer.Option(help="Specify agent name")] = None, + human: Annotated[Optional[str], typer.Option(help="Specify human")] = None, + preset: Annotated[Optional[str], typer.Option(help="Specify preset")] = None, + # model flags + model: Annotated[Optional[str], typer.Option(help="Specify the LLM model")] = None, + model_wrapper: Annotated[Optional[str], typer.Option(help="Specify the LLM model wrapper")] = None, + model_endpoint: Annotated[Optional[str], typer.Option(help="Specify the LLM model endpoint")] = None, + model_endpoint_type: Annotated[Optional[str], typer.Option(help="Specify the LLM model endpoint type")] = None, + context_window: Annotated[ + Optional[int], typer.Option(help="The context window of the LLM you are using (e.g. 8k for most Mistral 7B variants)") + ] = None, + # other + first: Annotated[bool, typer.Option(help="Use --first to send the first message in the sequence")] = False, + strip_ui: Annotated[bool, typer.Option(help="Remove all the bells and whistles in CLI output (helpful for testing)")] = False, + debug: Annotated[bool, typer.Option(help="Use --debug to enable debugging output")] = False, + no_verify: Annotated[bool, typer.Option(help="Bypass message verification")] = False, + yes: Annotated[bool, typer.Option("-y", help="Skip confirmation prompt and use defaults")] = False, + # streaming + stream: Annotated[bool, typer.Option(help="Enables message streaming in the CLI (if the backend supports it)")] = False, +): + """Start chatting with an MemGPT agent + + Example usage: `memgpt run --agent myagent --data-source mydata --persona mypersona --human myhuman --model gpt-3.5-turbo` + + :param persona: Specify persona + :param agent: Specify agent name (will load existing state if the agent exists, or create a new one with that name) + :param human: Specify human + :param model: Specify the LLM model + + """ + + # setup logger + # TODO: remove Utils Debug after global logging is complete. + utils.DEBUG = debug + # TODO: add logging command line options for runtime log level + + if debug: + logger.setLevel(logging.DEBUG) + else: + logger.setLevel(logging.CRITICAL) + + from memgpt.migrate import ( + VERSION_CUTOFF, + config_is_compatible, + wipe_config_and_reconfigure, + ) + + if not config_is_compatible(allow_empty=True): + typer.secho(f"\nYour current config file is incompatible with MemGPT versions later than {VERSION_CUTOFF}\n", fg=typer.colors.RED) + choices = [ + "Run the full config setup (recommended)", + "Create a new config using defaults", + "Cancel", + ] + selection = questionary.select( + f"To use MemGPT, you must either downgrade your MemGPT version (<= {VERSION_CUTOFF}), or regenerate your config. Would you like to proceed?", + choices=choices, + default=choices[0], + ).ask() + if selection == choices[0]: + try: + wipe_config_and_reconfigure() + except Exception as e: + typer.secho(f"Fresh config generation failed - error:\n{e}", fg=typer.colors.RED) + raise + elif selection == choices[1]: + try: + # Don't create a config, so that the next block of code asking about quickstart is run + wipe_config_and_reconfigure(run_configure=False, create_config=False) + except Exception as e: + typer.secho(f"Fresh config generation failed - error:\n{e}", fg=typer.colors.RED) + raise + else: + typer.secho("MemGPT config regeneration cancelled", fg=typer.colors.RED) + raise KeyboardInterrupt() + + typer.secho("Note: if you would like to migrate old agents to the new release, please run `memgpt migrate`!", fg=typer.colors.GREEN) + + if not MemGPTConfig.exists(): + # if no config, ask about quickstart + # do you want to do: + # - openai (run quickstart) + # - memgpt hosted (run quickstart) + # - other (run configure) + if yes: + # if user is passing '-y' to bypass all inputs, use memgpt hosted + # since it can't fail out if you don't have an API key + quickstart(backend=QuickstartChoice.memgpt_hosted) + config = MemGPTConfig() + + else: + config_choices = { + "memgpt": "Use the free MemGPT endpoints", + "openai": "Use OpenAI (requires an OpenAI API key)", + "other": "Other (OpenAI Azure, custom LLM endpoint, etc)", + } + print() + config_selection = questionary.select( + "How would you like to set up MemGPT?", + choices=list(config_choices.values()), + default=config_choices["memgpt"], + ).ask() + + if config_selection == config_choices["memgpt"]: + print() + quickstart(backend=QuickstartChoice.memgpt_hosted, debug=debug, terminal=False, latest=False) + elif config_selection == config_choices["openai"]: + print() + quickstart(backend=QuickstartChoice.openai, debug=debug, terminal=False, latest=False) + elif config_selection == config_choices["other"]: + configure() + else: + raise ValueError(config_selection) + + config = MemGPTConfig.load() + + else: # load config + config = MemGPTConfig.load() + + # read user id from config + ms = MetadataStore(config) + user = create_default_user_or_exit(config, ms) + human = human if human else config.human + persona = persona if persona else config.persona + + # determine agent to use, if not provided + if not yes and not agent: + agents = ms.list_agents(user_id=user.id) + agents = [a.name for a in agents] + + if len(agents) > 0: + print() + select_agent = questionary.confirm("Would you like to select an existing agent?").ask() + if select_agent is None: + raise KeyboardInterrupt + if select_agent: + agent = questionary.select("Select agent:", choices=agents).ask() + + # create agent config + agent_state = ms.get_agent(agent_name=agent, user_id=user.id) if agent else None + if agent and agent_state: # use existing agent + typer.secho(f"\n🔁 Using existing agent {agent}", fg=typer.colors.GREEN) + # agent_config = AgentConfig.load(agent) + # agent_state = ms.get_agent(agent_name=agent, user_id=user_id) + printd("Loading agent state:", agent_state.id) + printd("Agent state:", agent_state.state) + # printd("State path:", agent_config.save_state_dir()) + # printd("Persistent manager path:", agent_config.save_persistence_manager_dir()) + # printd("Index path:", agent_config.save_agent_index_dir()) + # persistence_manager = LocalStateManager(agent_config).load() # TODO: implement load + # TODO: load prior agent state + if persona and persona != agent_state.persona: + typer.secho(f"{CLI_WARNING_PREFIX}Overriding existing persona {agent_state.persona} with {persona}", fg=typer.colors.YELLOW) + agent_state.persona = persona + # raise ValueError(f"Cannot override {agent_state.name} existing persona {agent_state.persona} with {persona}") + if human and human != agent_state.human: + typer.secho(f"{CLI_WARNING_PREFIX}Overriding existing human {agent_state.human} with {human}", fg=typer.colors.YELLOW) + agent_state.human = human + # raise ValueError(f"Cannot override {agent_config.name} existing human {agent_config.human} with {human}") + + # Allow overriding model specifics (model, model wrapper, model endpoint IP + type, context_window) + if model and model != agent_state.llm_config.model: + typer.secho( + f"{CLI_WARNING_PREFIX}Overriding existing model {agent_state.llm_config.model} with {model}", fg=typer.colors.YELLOW + ) + agent_state.llm_config.model = model + if context_window is not None and int(context_window) != agent_state.llm_config.context_window: + typer.secho( + f"{CLI_WARNING_PREFIX}Overriding existing context window {agent_state.llm_config.context_window} with {context_window}", + fg=typer.colors.YELLOW, + ) + agent_state.llm_config.context_window = context_window + if model_wrapper and model_wrapper != agent_state.llm_config.model_wrapper: + typer.secho( + f"{CLI_WARNING_PREFIX}Overriding existing model wrapper {agent_state.llm_config.model_wrapper} with {model_wrapper}", + fg=typer.colors.YELLOW, + ) + agent_state.llm_config.model_wrapper = model_wrapper + if model_endpoint and model_endpoint != agent_state.llm_config.model_endpoint: + typer.secho( + f"{CLI_WARNING_PREFIX}Overriding existing model endpoint {agent_state.llm_config.model_endpoint} with {model_endpoint}", + fg=typer.colors.YELLOW, + ) + agent_state.llm_config.model_endpoint = model_endpoint + if model_endpoint_type and model_endpoint_type != agent_state.llm_config.model_endpoint_type: + typer.secho( + f"{CLI_WARNING_PREFIX}Overriding existing model endpoint type {agent_state.llm_config.model_endpoint_type} with {model_endpoint_type}", + fg=typer.colors.YELLOW, + ) + agent_state.llm_config.model_endpoint_type = model_endpoint_type + + # Update the agent with any overrides + ms.update_agent(agent_state) + + # create agent + memgpt_agent = Agent(agent_state=agent_state, interface=interface()) + + else: # create new agent + # create new agent config: override defaults with args if provided + typer.secho("\n🧬 Creating new agent...", fg=typer.colors.WHITE) + + agent_name = agent if agent else utils.create_random_username() + llm_config = config.default_llm_config + embedding_config = config.default_embedding_config # TODO allow overriding embedding params via CLI run + + # Allow overriding model specifics (model, model wrapper, model endpoint IP + type, context_window) + if model and model != llm_config.model: + typer.secho(f"{CLI_WARNING_PREFIX}Overriding default model {llm_config.model} with {model}", fg=typer.colors.YELLOW) + llm_config.model = model + if context_window is not None and int(context_window) != llm_config.context_window: + typer.secho( + f"{CLI_WARNING_PREFIX}Overriding default context window {llm_config.context_window} with {context_window}", + fg=typer.colors.YELLOW, + ) + llm_config.context_window = context_window + if model_wrapper and model_wrapper != llm_config.model_wrapper: + typer.secho( + f"{CLI_WARNING_PREFIX}Overriding existing model wrapper {llm_config.model_wrapper} with {model_wrapper}", + fg=typer.colors.YELLOW, + ) + llm_config.model_wrapper = model_wrapper + if model_endpoint and model_endpoint != llm_config.model_endpoint: + typer.secho( + f"{CLI_WARNING_PREFIX}Overriding existing model endpoint {llm_config.model_endpoint} with {model_endpoint}", + fg=typer.colors.YELLOW, + ) + llm_config.model_endpoint = model_endpoint + if model_endpoint_type and model_endpoint_type != llm_config.model_endpoint_type: + typer.secho( + f"{CLI_WARNING_PREFIX}Overriding existing model endpoint type {llm_config.model_endpoint_type} with {model_endpoint_type}", + fg=typer.colors.YELLOW, + ) + llm_config.model_endpoint_type = model_endpoint_type + + # create agent + try: + preset_obj = ms.get_preset(name=preset if preset else config.preset, user_id=user.id) + human_obj = ms.get_human(human, user.id) + persona_obj = ms.get_persona(persona, user.id) + if preset_obj is None: + # create preset records in metadata store + from memgpt.presets.presets import add_default_presets + + add_default_presets(user.id, ms) + # try again + preset_obj = ms.get_preset(name=preset if preset else config.preset, user_id=user.id) + if preset_obj is None: + typer.secho("Couldn't find presets in database, please run `memgpt configure`", fg=typer.colors.RED) + sys.exit(1) + if human_obj is None: + typer.secho("Couldn't find human {human} in database, please run `memgpt add human`", fg=typer.colors.RED) + if persona_obj is None: + typer.secho("Couldn't find persona {persona} in database, please run `memgpt add persona`", fg=typer.colors.RED) + + # Overwrite fields in the preset if they were specified + preset_obj.human = ms.get_human(human, user.id).text + preset_obj.persona = ms.get_persona(persona, user.id).text + + typer.secho(f"-> 🤖 Using persona profile: '{preset_obj.persona_name}'", fg=typer.colors.WHITE) + typer.secho(f"-> 🧑 Using human profile: '{preset_obj.human_name}'", fg=typer.colors.WHITE) + + memgpt_agent = Agent( + interface=interface(), + name=agent_name, + created_by=user.id, + preset=preset_obj, + llm_config=llm_config, + embedding_config=embedding_config, + # gpt-3.5-turbo tends to omit inner monologue, relax this requirement for now + first_message_verify_mono=True if (model is not None and "gpt-4" in model) else False, + ) + save_agent(agent=memgpt_agent, ms=ms) + + except ValueError as e: + typer.secho(f"Failed to create agent from provided information:\n{e}", fg=typer.colors.RED) + sys.exit(1) + typer.secho(f"🎉 Created new agent '{memgpt_agent.agent_state.name}' (id={memgpt_agent.agent_state.id})", fg=typer.colors.GREEN) + + # start event loop + from memgpt.main import run_agent_loop + + print() # extra space + run_agent_loop( + memgpt_agent=memgpt_agent, config=config, first=first, ms=ms, no_verify=no_verify, stream=stream + ) # TODO: add back no_verify + + +def delete_agent( + agent_name: Annotated[str, typer.Option(help="Specify agent to delete")], + user_id: Annotated[Optional[str], typer.Option(help="User ID to associate with the agent.")] = None, +): + """Delete an agent from the database""" + # use client ID is no user_id provided + config = MemGPTConfig.load() + ms = MetadataStore(config) + if user_id is None: + user = create_default_user_or_exit(config, ms) + else: + user = ms.get_user(user_id=uuid.UUID(user_id)) + + try: + agent = ms.get_agent(agent_name=agent_name, user_id=user.id) + except Exception as e: + typer.secho(f"Failed to get agent {agent_name}\n{e}", fg=typer.colors.RED) + sys.exit(1) + + if agent is None: + typer.secho(f"Couldn't find agent named '{agent_name}' to delete", fg=typer.colors.RED) + sys.exit(1) + + confirm = questionary.confirm(f"Are you sure you want to delete agent '{agent_name}' (id={agent.id})?", default=False).ask() + if confirm is None: + raise KeyboardInterrupt + if not confirm: + typer.secho(f"Cancelled agent deletion '{agent_name}' (id={agent.id})", fg=typer.colors.GREEN) + return + + try: + ms.delete_agent(agent_id=agent.id) + typer.secho(f"🕊️ Successfully deleted agent '{agent_name}' (id={agent.id})", fg=typer.colors.GREEN) + except Exception: + typer.secho(f"Failed to delete agent '{agent_name}' (id={agent.id})", fg=typer.colors.RED) + sys.exit(1) + + +def version(): + import memgpt + + print(memgpt.__version__) + return memgpt.__version__ diff --git a/memory_scope/cli/cli_config.py b/memory_scope/cli/cli_config.py new file mode 100644 index 00000000..55d74ad8 --- /dev/null +++ b/memory_scope/cli/cli_config.py @@ -0,0 +1,1294 @@ +import builtins +import os +import uuid +from enum import Enum +from typing import Annotated, Optional + +import questionary +import typer +from prettytable.colortable import ColorTable, Themes +from tqdm import tqdm + +from memgpt import utils +from memgpt.agent_store.storage import StorageConnector, TableType +from memgpt.config import MemGPTConfig +from memgpt.constants import LLM_MAX_TOKENS, MEMGPT_DIR +from memgpt.credentials import SUPPORTED_AUTH_TYPES, MemGPTCredentials +from memgpt.data_types import EmbeddingConfig, LLMConfig, Source, User +from memgpt.llm_api.anthropic import ( + anthropic_get_model_list, + antropic_get_model_context_window, +) +from memgpt.llm_api.azure_openai import azure_openai_get_model_list +from memgpt.llm_api.cohere import ( + COHERE_VALID_MODEL_LIST, + cohere_get_model_context_window, + cohere_get_model_list, +) +from memgpt.llm_api.google_ai import ( + google_ai_get_model_context_window, + google_ai_get_model_list, +) +from memgpt.llm_api.llm_api_tools import LLM_API_PROVIDER_OPTIONS +from memgpt.llm_api.openai import openai_get_model_list +from memgpt.local_llm.constants import ( + DEFAULT_ENDPOINTS, + DEFAULT_OLLAMA_MODEL, + DEFAULT_WRAPPER_NAME, +) +from memgpt.local_llm.utils import get_available_wrappers +from memgpt.metadata import MetadataStore +from memgpt.models.pydantic_models import HumanModel, PersonaModel +from memgpt.presets.presets import create_preset_from_file +from memgpt.server.utils import shorten_key_middle + +app = typer.Typer() + + +def get_azure_credentials(): + creds = dict( + azure_key=os.getenv("AZURE_OPENAI_KEY"), + azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"), + azure_version=os.getenv("AZURE_OPENAI_VERSION"), + azure_deployment=os.getenv("AZURE_OPENAI_DEPLOYMENT"), + azure_embedding_deployment=os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT"), + ) + # embedding endpoint and version default to non-embedding + creds["azure_embedding_endpoint"] = os.getenv("AZURE_OPENAI_EMBEDDING_ENDPOINT", creds["azure_endpoint"]) + creds["azure_embedding_version"] = os.getenv("AZURE_OPENAI_EMBEDDING_VERSION", creds["azure_version"]) + return creds + + +def get_openai_credentials() -> Optional[str]: + openai_key = os.getenv("OPENAI_API_KEY", None) + return openai_key + + +def get_google_ai_credentials() -> Optional[str]: + google_ai_key = os.getenv("GOOGLE_AI_API_KEY", None) + return google_ai_key + + +def configure_llm_endpoint(config: MemGPTConfig, credentials: MemGPTCredentials): + # configure model endpoint + model_endpoint_type, model_endpoint = None, None + + # get default + default_model_endpoint_type = config.default_llm_config.model_endpoint_type if config.default_embedding_config else None + if ( + config.default_llm_config + and config.default_llm_config.model_endpoint_type is not None + and config.default_llm_config.model_endpoint_type not in [provider for provider in LLM_API_PROVIDER_OPTIONS if provider != "local"] + ): # local model + default_model_endpoint_type = "local" + + provider = questionary.select( + "Select LLM inference provider:", + choices=LLM_API_PROVIDER_OPTIONS, + default=default_model_endpoint_type, + ).ask() + if provider is None: + raise KeyboardInterrupt + + # set: model_endpoint_type, model_endpoint + if provider == "openai": + # check for key + if credentials.openai_key is None: + # allow key to get pulled from env vars + openai_api_key = os.getenv("OPENAI_API_KEY", None) + # if we still can't find it, ask for it as input + if openai_api_key is None: + while openai_api_key is None or len(openai_api_key) == 0: + # Ask for API key as input + openai_api_key = questionary.password( + "Enter your OpenAI API key (starts with 'sk-', see https://platform.openai.com/api-keys):" + ).ask() + if openai_api_key is None: + raise KeyboardInterrupt + credentials.openai_key = openai_api_key + credentials.save() + else: + # Give the user an opportunity to overwrite the key + openai_api_key = None + default_input = ( + shorten_key_middle(credentials.openai_key) if credentials.openai_key.startswith("sk-") else credentials.openai_key + ) + openai_api_key = questionary.password( + "Enter your OpenAI API key (starts with 'sk-', see https://platform.openai.com/api-keys):", + default=default_input, + ).ask() + if openai_api_key is None: + raise KeyboardInterrupt + # If the user modified it, use the new one + if openai_api_key != default_input: + credentials.openai_key = openai_api_key + credentials.save() + + model_endpoint_type = "openai" + model_endpoint = "https://api.openai.com/v1" + model_endpoint = questionary.text("Override default endpoint:", default=model_endpoint).ask() + if model_endpoint is None: + raise KeyboardInterrupt + provider = "openai" + + elif provider == "azure": + # check for necessary vars + azure_creds = get_azure_credentials() + if not all([azure_creds["azure_key"], azure_creds["azure_endpoint"], azure_creds["azure_version"]]): + raise ValueError( + "Missing environment variables for Azure (see https://memgpt.readme.io/docs/endpoints#azure-openai). Please set then run `memgpt configure` again." + ) + else: + credentials.azure_key = azure_creds["azure_key"] + credentials.azure_version = azure_creds["azure_version"] + credentials.azure_endpoint = azure_creds["azure_endpoint"] + if "azure_deployment" in azure_creds: + credentials.azure_deployment = azure_creds["azure_deployment"] + credentials.azure_embedding_version = azure_creds["azure_embedding_version"] + credentials.azure_embedding_endpoint = azure_creds["azure_embedding_endpoint"] + if "azure_embedding_deployment" in azure_creds: + credentials.azure_embedding_deployment = azure_creds["azure_embedding_deployment"] + credentials.save() + + model_endpoint_type = "azure" + model_endpoint = azure_creds["azure_endpoint"] + + elif provider == "google_ai": + + # check for key + if credentials.google_ai_key is None: + # allow key to get pulled from env vars + google_ai_key = get_google_ai_credentials() + # if we still can't find it, ask for it as input + if google_ai_key is None: + while google_ai_key is None or len(google_ai_key) == 0: + # Ask for API key as input + google_ai_key = questionary.password( + "Enter your Google AI (Gemini) API key (see https://aistudio.google.com/app/apikey):" + ).ask() + if google_ai_key is None: + raise KeyboardInterrupt + credentials.google_ai_key = google_ai_key + else: + # Give the user an opportunity to overwrite the key + google_ai_key = None + default_input = shorten_key_middle(credentials.google_ai_key) + + google_ai_key = questionary.password( + "Enter your Google AI (Gemini) API key (see https://aistudio.google.com/app/apikey):", + default=default_input, + ).ask() + if google_ai_key is None: + raise KeyboardInterrupt + # If the user modified it, use the new one + if google_ai_key != default_input: + credentials.google_ai_key = google_ai_key + + default_input = os.getenv("GOOGLE_AI_SERVICE_ENDPOINT", None) + if default_input is None: + default_input = "generativelanguage" + google_ai_service_endpoint = questionary.text( + "Enter your Google AI (Gemini) service endpoint (see https://ai.google.dev/api/rest):", + default=default_input, + ).ask() + credentials.google_ai_service_endpoint = google_ai_service_endpoint + + # write out the credentials + credentials.save() + + model_endpoint_type = "google_ai" + + elif provider == "anthropic": + # check for key + if credentials.anthropic_key is None: + # allow key to get pulled from env vars + anthropic_api_key = os.getenv("ANTHROPIC_API_KEY", None) + # if we still can't find it, ask for it as input + if anthropic_api_key is None: + while anthropic_api_key is None or len(anthropic_api_key) == 0: + # Ask for API key as input + anthropic_api_key = questionary.password( + "Enter your Anthropic API key (starts with 'sk-', see https://console.anthropic.com/settings/keys):" + ).ask() + if anthropic_api_key is None: + raise KeyboardInterrupt + credentials.anthropic_key = anthropic_api_key + credentials.save() + else: + # Give the user an opportunity to overwrite the key + anthropic_api_key = None + default_input = ( + shorten_key_middle(credentials.anthropic_key) if credentials.anthropic_key.startswith("sk-") else credentials.anthropic_key + ) + anthropic_api_key = questionary.password( + "Enter your Anthropic API key (starts with 'sk-', see https://console.anthropic.com/settings/keys):", + default=default_input, + ).ask() + if anthropic_api_key is None: + raise KeyboardInterrupt + # If the user modified it, use the new one + if anthropic_api_key != default_input: + credentials.anthropic_key = anthropic_api_key + credentials.save() + + model_endpoint_type = "anthropic" + model_endpoint = "https://api.anthropic.com/v1" + model_endpoint = questionary.text("Override default endpoint:", default=model_endpoint).ask() + if model_endpoint is None: + raise KeyboardInterrupt + provider = "anthropic" + + elif provider == "cohere": + # check for key + if credentials.cohere_key is None: + # allow key to get pulled from env vars + cohere_api_key = os.getenv("COHERE_API_KEY", None) + # if we still can't find it, ask for it as input + if cohere_api_key is None: + while cohere_api_key is None or len(cohere_api_key) == 0: + # Ask for API key as input + cohere_api_key = questionary.password("Enter your Cohere API key (see https://dashboard.cohere.com/api-keys):").ask() + if cohere_api_key is None: + raise KeyboardInterrupt + credentials.cohere_key = cohere_api_key + credentials.save() + else: + # Give the user an opportunity to overwrite the key + cohere_api_key = None + default_input = ( + shorten_key_middle(credentials.cohere_key) if credentials.cohere_key.startswith("sk-") else credentials.cohere_key + ) + cohere_api_key = questionary.password( + "Enter your Cohere API key (see https://dashboard.cohere.com/api-keys):", + default=default_input, + ).ask() + if cohere_api_key is None: + raise KeyboardInterrupt + # If the user modified it, use the new one + if cohere_api_key != default_input: + credentials.cohere_key = cohere_api_key + credentials.save() + + model_endpoint_type = "cohere" + model_endpoint = "https://api.cohere.ai/v1" + model_endpoint = questionary.text("Override default endpoint:", default=model_endpoint).ask() + if model_endpoint is None: + raise KeyboardInterrupt + provider = "cohere" + + else: # local models + # backend_options_old = ["webui", "webui-legacy", "llamacpp", "koboldcpp", "ollama", "lmstudio", "lmstudio-legacy", "vllm", "openai"] + backend_options = builtins.list(DEFAULT_ENDPOINTS.keys()) + # assert backend_options_old == backend_options, (backend_options_old, backend_options) + default_model_endpoint_type = None + if config.default_llm_config and config.default_llm_config.model_endpoint_type in backend_options: + # set from previous config + default_model_endpoint_type = config.default_llm_config.model_endpoint_type + model_endpoint_type = questionary.select( + "Select LLM backend (select 'openai' if you have an OpenAI compatible proxy):", + backend_options, + default=default_model_endpoint_type, + ).ask() + if model_endpoint_type is None: + raise KeyboardInterrupt + + # set default endpoint + # if OPENAI_API_BASE is set, assume that this is the IP+port the user wanted to use + default_model_endpoint = os.getenv("OPENAI_API_BASE") + # if OPENAI_API_BASE is not set, try to pull a default IP+port format from a hardcoded set + if default_model_endpoint is None: + if model_endpoint_type in DEFAULT_ENDPOINTS: + default_model_endpoint = DEFAULT_ENDPOINTS[model_endpoint_type] + model_endpoint = questionary.text("Enter default endpoint:", default=default_model_endpoint).ask() + if model_endpoint is None: + raise KeyboardInterrupt + while not utils.is_valid_url(model_endpoint): + typer.secho(f"Endpoint must be a valid address", fg=typer.colors.YELLOW) + model_endpoint = questionary.text("Enter default endpoint:", default=default_model_endpoint).ask() + if model_endpoint is None: + raise KeyboardInterrupt + elif config.default_llm_config and config.default_llm_config.model_endpoint: + model_endpoint = questionary.text("Enter default endpoint:", default=config.default_llm_config.model_endpoint).ask() + if model_endpoint is None: + raise KeyboardInterrupt + while not utils.is_valid_url(model_endpoint): + typer.secho(f"Endpoint must be a valid address", fg=typer.colors.YELLOW) + model_endpoint = questionary.text("Enter default endpoint:", default=config.default_llm_config.model_endpoint).ask() + if model_endpoint is None: + raise KeyboardInterrupt + else: + # default_model_endpoint = None + model_endpoint = None + model_endpoint = questionary.text("Enter default endpoint:").ask() + if model_endpoint is None: + raise KeyboardInterrupt + while not utils.is_valid_url(model_endpoint): + typer.secho(f"Endpoint must be a valid address", fg=typer.colors.YELLOW) + model_endpoint = questionary.text("Enter default endpoint:").ask() + if model_endpoint is None: + raise KeyboardInterrupt + else: + model_endpoint = default_model_endpoint + assert model_endpoint, f"Environment variable OPENAI_API_BASE must be set." + + return model_endpoint_type, model_endpoint + + +def get_model_options( + credentials: MemGPTCredentials, + model_endpoint_type: str, + model_endpoint: str, + filter_list: bool = True, + filter_prefix: str = "gpt-", +) -> list: + try: + if model_endpoint_type == "openai": + if credentials.openai_key is None: + raise ValueError("Missing OpenAI API key") + fetched_model_options_response = openai_get_model_list(url=model_endpoint, api_key=credentials.openai_key) + + # Filter the list for "gpt" models only + if filter_list: + model_options = [obj["id"] for obj in fetched_model_options_response["data"] if obj["id"].startswith(filter_prefix)] + else: + model_options = [obj["id"] for obj in fetched_model_options_response["data"]] + + elif model_endpoint_type == "azure": + if credentials.azure_key is None: + raise ValueError("Missing Azure key") + if credentials.azure_version is None: + raise ValueError("Missing Azure version") + fetched_model_options_response = azure_openai_get_model_list( + url=model_endpoint, api_key=credentials.azure_key, api_version=credentials.azure_version + ) + + # Filter the list for "gpt" models only + if filter_list: + model_options = [obj["id"] for obj in fetched_model_options_response["data"] if obj["id"].startswith(filter_prefix)] + else: + model_options = [obj["id"] for obj in fetched_model_options_response["data"]] + + elif model_endpoint_type == "google_ai": + if credentials.google_ai_key is None: + raise ValueError("Missing Google AI API key") + if credentials.google_ai_service_endpoint is None: + raise ValueError("Missing Google AI service endpoint") + model_options = google_ai_get_model_list( + service_endpoint=credentials.google_ai_service_endpoint, api_key=credentials.google_ai_key + ) + model_options = [str(m["name"]) for m in model_options] + model_options = [mo[len("models/") :] if mo.startswith("models/") else mo for mo in model_options] + + # TODO remove manual filtering for gemini-pro + model_options = [mo for mo in model_options if str(mo).startswith("gemini") and "-pro" in str(mo)] + # model_options = ["gemini-pro"] + + elif model_endpoint_type == "anthropic": + if credentials.anthropic_key is None: + raise ValueError("Missing Anthropic API key") + fetched_model_options = anthropic_get_model_list(url=model_endpoint, api_key=credentials.anthropic_key) + model_options = [obj["name"] for obj in fetched_model_options] + + elif model_endpoint_type == "cohere": + if credentials.cohere_key is None: + raise ValueError("Missing Cohere API key") + fetched_model_options = cohere_get_model_list(url=model_endpoint, api_key=credentials.cohere_key) + model_options = [obj for obj in fetched_model_options] + + else: + # Attempt to do OpenAI endpoint style model fetching + # TODO support local auth with api-key header + if credentials.openllm_auth_type == "bearer_token": + api_key = credentials.openllm_key + else: + api_key = None + fetched_model_options_response = openai_get_model_list(url=model_endpoint, api_key=api_key, fix_url=True) + model_options = [obj["id"] for obj in fetched_model_options_response["data"]] + # NOTE no filtering of local model options + + # list + return model_options + + except: + raise Exception(f"Failed to get model list from {model_endpoint}") + + +def configure_model(config: MemGPTConfig, credentials: MemGPTCredentials, model_endpoint_type: str, model_endpoint: str): + # set: model, model_wrapper + model, model_wrapper = None, None + if model_endpoint_type == "openai" or model_endpoint_type == "azure": + # Get the model list from the openai / azure endpoint + hardcoded_model_options = ["gpt-4", "gpt-4-32k", "gpt-4-1106-preview", "gpt-3.5-turbo", "gpt-3.5-turbo-16k"] + fetched_model_options = [] + try: + fetched_model_options = get_model_options( + credentials=credentials, model_endpoint_type=model_endpoint_type, model_endpoint=model_endpoint + ) + except Exception as e: + # NOTE: if this fails, it means the user's key is probably bad + typer.secho( + f"Failed to get model list from {model_endpoint} - make sure your API key and endpoints are correct!", fg=typer.colors.RED + ) + raise e + + # First ask if the user wants to see the full model list (some may be incompatible) + see_all_option_str = "[see all options]" + other_option_str = "[enter model name manually]" + + # Check if the model we have set already is even in the list (informs our default) + valid_model = config.default_llm_config and config.default_llm_config.model in hardcoded_model_options + model = questionary.select( + "Select default model (recommended: gpt-4):", + choices=hardcoded_model_options + [see_all_option_str, other_option_str], + default=config.default_llm_config.model if valid_model else hardcoded_model_options[0], + ).ask() + if model is None: + raise KeyboardInterrupt + + # If the user asked for the full list, show it + if model == see_all_option_str: + typer.secho(f"Warning: not all models shown are guaranteed to work with MemGPT", fg=typer.colors.RED) + model = questionary.select( + "Select default model (recommended: gpt-4):", + choices=fetched_model_options + [other_option_str], + default=config.default_llm_config.model if (valid_model and config.default_llm_config) else fetched_model_options[0], + ).ask() + if model is None: + raise KeyboardInterrupt + + # Finally if the user asked to manually input, allow it + if model == other_option_str: + model = "" + while len(model) == 0: + model = questionary.text( + "Enter custom model name:", + ).ask() + if model is None: + raise KeyboardInterrupt + + elif model_endpoint_type == "google_ai": + try: + fetched_model_options = get_model_options( + credentials=credentials, model_endpoint_type=model_endpoint_type, model_endpoint=model_endpoint + ) + except Exception as e: + # NOTE: if this fails, it means the user's key is probably bad + typer.secho( + f"Failed to get model list from {model_endpoint} - make sure your API key and endpoints are correct!", fg=typer.colors.RED + ) + raise e + + model = questionary.select( + "Select default model:", + choices=fetched_model_options, + default=fetched_model_options[0], + ).ask() + if model is None: + raise KeyboardInterrupt + + elif model_endpoint_type == "anthropic": + try: + fetched_model_options = get_model_options( + credentials=credentials, model_endpoint_type=model_endpoint_type, model_endpoint=model_endpoint + ) + except Exception as e: + # NOTE: if this fails, it means the user's key is probably bad + typer.secho( + f"Failed to get model list from {model_endpoint} - make sure your API key and endpoints are correct!", fg=typer.colors.RED + ) + raise e + + model = questionary.select( + "Select default model:", + choices=fetched_model_options, + default=fetched_model_options[0], + ).ask() + if model is None: + raise KeyboardInterrupt + + elif model_endpoint_type == "cohere": + + fetched_model_options = [] + try: + fetched_model_options = get_model_options( + credentials=credentials, model_endpoint_type=model_endpoint_type, model_endpoint=model_endpoint + ) + except Exception as e: + # NOTE: if this fails, it means the user's key is probably bad + typer.secho( + f"Failed to get model list from {model_endpoint} - make sure your API key and endpoints are correct!", fg=typer.colors.RED + ) + raise e + + fetched_model_options = [m["name"] for m in fetched_model_options] + hardcoded_model_options = [m for m in fetched_model_options if m in COHERE_VALID_MODEL_LIST] + + # First ask if the user wants to see the full model list (some may be incompatible) + see_all_option_str = "[see all options]" + other_option_str = "[enter model name manually]" + + # Check if the model we have set already is even in the list (informs our default) + valid_model = config.default_llm_config.model in hardcoded_model_options + model = questionary.select( + "Select default model (recommended: command-r-plus):", + choices=hardcoded_model_options + [see_all_option_str, other_option_str], + default=config.default_llm_config.model if valid_model else hardcoded_model_options[0], + ).ask() + if model is None: + raise KeyboardInterrupt + + # If the user asked for the full list, show it + if model == see_all_option_str: + typer.secho(f"Warning: not all models shown are guaranteed to work with MemGPT", fg=typer.colors.RED) + model = questionary.select( + "Select default model (recommended: command-r-plus):", + choices=fetched_model_options + [other_option_str], + default=config.default_llm_config.model if valid_model else fetched_model_options[0], + ).ask() + if model is None: + raise KeyboardInterrupt + + # Finally if the user asked to manually input, allow it + if model == other_option_str: + model = "" + while len(model) == 0: + model = questionary.text( + "Enter custom model name:", + ).ask() + if model is None: + raise KeyboardInterrupt + + else: # local models + + # ask about local auth + if model_endpoint_type in ["groq"]: # TODO all llm engines under 'local' that will require api keys + use_local_auth = True + local_auth_type = "bearer_token" + local_auth_key = questionary.password( + "Enter your Groq API key:", + ).ask() + if local_auth_key is None: + raise KeyboardInterrupt + credentials.openllm_auth_type = local_auth_type + credentials.openllm_key = local_auth_key + credentials.save() + else: + use_local_auth = questionary.confirm( + "Is your LLM endpoint authenticated? (default no)", + default=False, + ).ask() + if use_local_auth is None: + raise KeyboardInterrupt + if use_local_auth: + local_auth_type = questionary.select( + "What HTTP authentication method does your endpoint require?", + choices=SUPPORTED_AUTH_TYPES, + default=SUPPORTED_AUTH_TYPES[0], + ).ask() + if local_auth_type is None: + raise KeyboardInterrupt + local_auth_key = questionary.password( + "Enter your authentication key:", + ).ask() + if local_auth_key is None: + raise KeyboardInterrupt + # credentials = MemGPTCredentials.load() + credentials.openllm_auth_type = local_auth_type + credentials.openllm_key = local_auth_key + credentials.save() + + # ollama also needs model type + if model_endpoint_type == "ollama": + default_model = ( + config.default_llm_config.model + if config.default_llm_config and config.default_llm_config.model_endpoint_type == "ollama" + else DEFAULT_OLLAMA_MODEL + ) + model = questionary.text( + "Enter default model name (required for Ollama, see: https://memgpt.readme.io/docs/ollama):", + default=default_model, + ).ask() + if model is None: + raise KeyboardInterrupt + model = None if len(model) == 0 else model + + default_model = ( + config.default_llm_config.model if config.default_llm_config and config.default_llm_config.model_endpoint_type == "vllm" else "" + ) + + # vllm needs huggingface model tag + if model_endpoint_type in ["vllm", "groq"]: + try: + # Don't filter model list for vLLM since model list is likely much smaller than OpenAI/Azure endpoint + # + probably has custom model names + # TODO support local auth + model_options = get_model_options( + credentials=credentials, model_endpoint_type=model_endpoint_type, model_endpoint=model_endpoint + ) + except: + print(f"Failed to get model list from {model_endpoint}, using defaults") + model_options = None + + # If we got model options from vLLM endpoint, allow selection + custom input + if model_options is not None: + other_option_str = "other (enter name)" + valid_model = config.default_llm_config.model in model_options + model_options.append(other_option_str) + model = questionary.select( + "Select default model:", + choices=model_options, + default=config.default_llm_config.model if valid_model else model_options[0], + ).ask() + if model is None: + raise KeyboardInterrupt + + # If we got custom input, ask for raw input + if model == other_option_str: + model = questionary.text( + "Enter HuggingFace model tag (e.g. ehartford/dolphin-2.2.1-mistral-7b):", + default=default_model, + ).ask() + if model is None: + raise KeyboardInterrupt + # TODO allow empty string for input? + model = None if len(model) == 0 else model + + else: + model = questionary.text( + "Enter HuggingFace model tag (e.g. ehartford/dolphin-2.2.1-mistral-7b):", + default=default_model, + ).ask() + if model is None: + raise KeyboardInterrupt + model = None if len(model) == 0 else model + + # model wrapper + available_model_wrappers = builtins.list(get_available_wrappers().keys()) + model_wrapper = questionary.select( + f"Select default model wrapper (recommended: {DEFAULT_WRAPPER_NAME}):", + choices=available_model_wrappers, + default=DEFAULT_WRAPPER_NAME, + ).ask() + if model_wrapper is None: + raise KeyboardInterrupt + + # set: context_window + if str(model) not in LLM_MAX_TOKENS: + + context_length_options = [ + str(2**12), # 4096 + str(2**13), # 8192 + str(2**14), # 16384 + str(2**15), # 32768 + str(2**18), # 262144 + "custom", # enter yourself + ] + + if model_endpoint_type == "google_ai": + try: + fetched_context_window = str( + google_ai_get_model_context_window( + service_endpoint=credentials.google_ai_service_endpoint, api_key=credentials.google_ai_key, model=model + ) + ) + print(f"Got context window {fetched_context_window} for model {model} (from Google API)") + context_length_options = [ + fetched_context_window, + "custom", + ] + except Exception as e: + print(f"Failed to get model details for model '{model}' on Google AI API ({str(e)})") + + context_window_input = questionary.select( + "Select your model's context window (see https://cloud.google.com/vertex-ai/generative-ai/docs/learn/model-versioning#gemini-model-versions):", + choices=context_length_options, + default=context_length_options[0], + ).ask() + if context_window_input is None: + raise KeyboardInterrupt + + elif model_endpoint_type == "anthropic": + try: + fetched_context_window = str( + antropic_get_model_context_window(url=model_endpoint, api_key=credentials.anthropic_key, model=model) + ) + print(f"Got context window {fetched_context_window} for model {model}") + context_length_options = [ + fetched_context_window, + "custom", + ] + except Exception as e: + print(f"Failed to get model details for model '{model}' ({str(e)})") + + context_window_input = questionary.select( + "Select your model's context window (see https://docs.anthropic.com/claude/docs/models-overview):", + choices=context_length_options, + default=context_length_options[0], + ).ask() + if context_window_input is None: + raise KeyboardInterrupt + + elif model_endpoint_type == "cohere": + try: + fetched_context_window = str( + cohere_get_model_context_window(url=model_endpoint, api_key=credentials.cohere_key, model=model) + ) + print(f"Got context window {fetched_context_window} for model {model}") + context_length_options = [ + fetched_context_window, + "custom", + ] + except Exception as e: + print(f"Failed to get model details for model '{model}' ({str(e)})") + + context_window_input = questionary.select( + "Select your model's context window (see https://docs.cohere.com/docs/command-r):", + choices=context_length_options, + default=context_length_options[0], + ).ask() + if context_window_input is None: + raise KeyboardInterrupt + + else: + + # Ask the user to specify the context length + context_window_input = questionary.select( + "Select your model's context window (for Mistral 7B models, this is probably 8k / 8192):", + choices=context_length_options, + default=str(LLM_MAX_TOKENS["DEFAULT"]), + ).ask() + if context_window_input is None: + raise KeyboardInterrupt + + # If custom, ask for input + if context_window_input == "custom": + while True: + context_window_input = questionary.text("Enter context window (e.g. 8192)").ask() + if context_window_input is None: + raise KeyboardInterrupt + try: + context_window = int(context_window_input) + break + except ValueError: + print(f"Context window must be a valid integer") + else: + context_window = int(context_window_input) + else: + # Pull the context length from the models + context_window = int(LLM_MAX_TOKENS[str(model)]) + return model, model_wrapper, context_window + + +def configure_embedding_endpoint(config: MemGPTConfig, credentials: MemGPTCredentials): + # configure embedding endpoint + + default_embedding_endpoint_type = config.default_embedding_config.embedding_endpoint_type if config.default_embedding_config else None + + embedding_endpoint_type, embedding_endpoint, embedding_dim, embedding_model = None, None, None, None + embedding_provider = questionary.select( + "Select embedding provider:", choices=["openai", "azure", "hugging-face", "local"], default=default_embedding_endpoint_type + ).ask() + if embedding_provider is None: + raise KeyboardInterrupt + + if embedding_provider == "openai": + # check for key + if credentials.openai_key is None: + # allow key to get pulled from env vars + openai_api_key = os.getenv("OPENAI_API_KEY", None) + if openai_api_key is None: + # if we still can't find it, ask for it as input + while openai_api_key is None or len(openai_api_key) == 0: + # Ask for API key as input + openai_api_key = questionary.password( + "Enter your OpenAI API key (starts with 'sk-', see https://platform.openai.com/api-keys):" + ).ask() + if openai_api_key is None: + raise KeyboardInterrupt + credentials.openai_key = openai_api_key + credentials.save() + + embedding_endpoint_type = "openai" + embedding_endpoint = "https://api.openai.com/v1" + embedding_dim = 1536 + embedding_model = "text-embedding-ada-002" + + elif embedding_provider == "azure": + # check for necessary vars + azure_creds = get_azure_credentials() + if not all([azure_creds["azure_key"], azure_creds["azure_embedding_endpoint"], azure_creds["azure_embedding_version"]]): + raise ValueError( + "Missing environment variables for Azure (see https://memgpt.readme.io/docs/endpoints#azure-openai). Please set then run `memgpt configure` again." + ) + credentials.azure_key = azure_creds["azure_key"] + credentials.azure_version = azure_creds["azure_version"] + credentials.azure_embedding_endpoint = azure_creds["azure_embedding_endpoint"] + credentials.save() + + embedding_endpoint_type = "azure" + embedding_endpoint = azure_creds["azure_embedding_endpoint"] + embedding_dim = 1536 + embedding_model = "text-embedding-ada-002" + + elif embedding_provider == "hugging-face": + # configure hugging face embedding endpoint (https://github.com/huggingface/text-embeddings-inference) + # supports custom model/endpoints + embedding_endpoint_type = "hugging-face" + embedding_endpoint = None + + # get endpoint + embedding_endpoint = questionary.text("Enter default endpoint:").ask() + if embedding_endpoint is None: + raise KeyboardInterrupt + while not utils.is_valid_url(embedding_endpoint): + typer.secho(f"Endpoint must be a valid address", fg=typer.colors.YELLOW) + embedding_endpoint = questionary.text("Enter default endpoint:").ask() + if embedding_endpoint is None: + raise KeyboardInterrupt + + # get model type + default_embedding_model = ( + config.default_embedding_config.embedding_model if config.default_embedding_config else "BAAI/bge-large-en-v1.5" + ) + embedding_model = questionary.text( + "Enter HuggingFace model tag (e.g. BAAI/bge-large-en-v1.5):", + default=default_embedding_model, + ).ask() + if embedding_model is None: + raise KeyboardInterrupt + + # get model dimentions + default_embedding_dim = config.default_embedding_config.embedding_dim if config.default_embedding_config else "1024" + embedding_dim = questionary.text("Enter embedding model dimentions (e.g. 1024):", default=str(default_embedding_dim)).ask() + if embedding_dim is None: + raise KeyboardInterrupt + try: + embedding_dim = int(embedding_dim) + except Exception: + raise ValueError(f"Failed to cast {embedding_dim} to integer.") + elif embedding_provider == "ollama": + # configure ollama embedding endpoint + embedding_endpoint_type = "ollama" + embedding_endpoint = "http://localhost:11434/api/embeddings" + # Source: https://github.com/ollama/ollama/blob/main/docs/api.md#generate-embeddings:~:text=http%3A//localhost%3A11434/api/embeddings + + # get endpoint (is this necessary?) + embedding_endpoint = questionary.text("Enter Ollama API endpoint:").ask() + if embedding_endpoint is None: + raise KeyboardInterrupt + while not utils.is_valid_url(embedding_endpoint): + typer.secho(f"Endpoint must be a valid address", fg=typer.colors.YELLOW) + embedding_endpoint = questionary.text("Enter Ollama API endpoint:").ask() + if embedding_endpoint is None: + raise KeyboardInterrupt + + # get model type + default_embedding_model = ( + config.default_embedding_config.embedding_model if config.default_embedding_config else "mxbai-embed-large" + ) + embedding_model = questionary.text( + "Enter Ollama model tag (e.g. mxbai-embed-large):", + default=default_embedding_model, + ).ask() + if embedding_model is None: + raise KeyboardInterrupt + + # get model dimensions + default_embedding_dim = config.default_embedding_config.embedding_dim if config.default_embedding_config else "512" + embedding_dim = questionary.text("Enter embedding model dimensions (e.g. 512):", default=str(default_embedding_dim)).ask() + if embedding_dim is None: + raise KeyboardInterrupt + try: + embedding_dim = int(embedding_dim) + except Exception: + raise ValueError(f"Failed to cast {embedding_dim} to integer.") + else: # local models + embedding_endpoint_type = "local" + embedding_endpoint = None + embedding_model = "BAAI/bge-small-en-v1.5" + embedding_dim = 384 + + return embedding_endpoint_type, embedding_endpoint, embedding_dim, embedding_model + + +def configure_archival_storage(config: MemGPTConfig, credentials: MemGPTCredentials): + # Configure archival storage backend + archival_storage_options = ["postgres", "chroma", "milvus", "qdrant"] + archival_storage_type = questionary.select( + "Select storage backend for archival data:", archival_storage_options, default=config.archival_storage_type + ).ask() + if archival_storage_type is None: + raise KeyboardInterrupt + archival_storage_uri, archival_storage_path = config.archival_storage_uri, config.archival_storage_path + + # configure postgres + if archival_storage_type == "postgres": + archival_storage_uri = questionary.text( + "Enter postgres connection string (e.g. postgresql+pg8000://{user}:{password}@{ip}:5432/{database}):", + default=config.archival_storage_uri if config.archival_storage_uri else "", + ).ask() + if archival_storage_uri is None: + raise KeyboardInterrupt + + # TODO: add back + ## configure lancedb + # if archival_storage_type == "lancedb": + # archival_storage_uri = questionary.text( + # "Enter lanncedb connection string (e.g. ./.lancedb", + # default=config.archival_storage_uri if config.archival_storage_uri else "./.lancedb", + # ).ask() + + # configure chroma + if archival_storage_type == "chroma": + chroma_type = questionary.select("Select chroma backend:", ["http", "persistent"], default="persistent").ask() + if chroma_type is None: + raise KeyboardInterrupt + if chroma_type == "http": + archival_storage_uri = questionary.text("Enter chroma ip (e.g. localhost:8000):", default="localhost:8000").ask() + if archival_storage_uri is None: + raise KeyboardInterrupt + if chroma_type == "persistent": + archival_storage_path = os.path.join(MEMGPT_DIR, "chroma") + + if archival_storage_type == "qdrant": + qdrant_type = questionary.select("Select Qdrant backend:", ["local", "server"], default="local").ask() + if qdrant_type is None: + raise KeyboardInterrupt + if qdrant_type == "server": + archival_storage_uri = questionary.text( + "Enter the Qdrant instance URI (Default: localhost:6333):", default="localhost:6333" + ).ask() + if archival_storage_uri is None: + raise KeyboardInterrupt + if qdrant_type == "local": + archival_storage_path = os.path.join(MEMGPT_DIR, "qdrant") + + if archival_storage_type == "milvus": + default_milvus_uri = archival_storage_path = os.path.join(MEMGPT_DIR, "milvus.db") + archival_storage_uri = questionary.text( + f"Enter the Milvus connection URI (Default: {default_milvus_uri}):", default=default_milvus_uri + ).ask() + if archival_storage_uri is None: + raise KeyboardInterrupt + return archival_storage_type, archival_storage_uri, archival_storage_path + + # TODO: allow configuring embedding model + + +def configure_recall_storage(config: MemGPTConfig, credentials: MemGPTCredentials): + # Configure recall storage backend + recall_storage_options = ["sqlite", "postgres"] + recall_storage_type = questionary.select( + "Select storage backend for recall data:", recall_storage_options, default=config.recall_storage_type + ).ask() + if recall_storage_type is None: + raise KeyboardInterrupt + recall_storage_uri, recall_storage_path = config.recall_storage_uri, config.recall_storage_path + # configure postgres + if recall_storage_type == "postgres": + recall_storage_uri = questionary.text( + "Enter postgres connection string (e.g. postgresql+pg8000://{user}:{password}@{ip}:5432/{database}):", + default=config.recall_storage_uri if config.recall_storage_uri else "", + ).ask() + if recall_storage_uri is None: + raise KeyboardInterrupt + + return recall_storage_type, recall_storage_uri, recall_storage_path + + +@app.command() +def configure(): + """Updates default MemGPT configurations + + This function and quickstart should be the ONLY place where MemGPTConfig.save() is called + """ + + # check credentials + credentials = MemGPTCredentials.load() + openai_key = get_openai_credentials() + + MemGPTConfig.create_config_dir() + + # Will pre-populate with defaults, or what the user previously set + config = MemGPTConfig.load() + try: + model_endpoint_type, model_endpoint = configure_llm_endpoint( + config=config, + credentials=credentials, + ) + model, model_wrapper, context_window = configure_model( + config=config, + credentials=credentials, + model_endpoint_type=str(model_endpoint_type), + model_endpoint=str(model_endpoint), + ) + embedding_endpoint_type, embedding_endpoint, embedding_dim, embedding_model = configure_embedding_endpoint( + config=config, + credentials=credentials, + ) + archival_storage_type, archival_storage_uri, archival_storage_path = configure_archival_storage( + config=config, + credentials=credentials, + ) + recall_storage_type, recall_storage_uri, recall_storage_path = configure_recall_storage( + config=config, + credentials=credentials, + ) + except ValueError as e: + typer.secho(str(e), fg=typer.colors.RED) + return + + # openai key might have gotten added along the way + openai_key = credentials.openai_key if credentials.openai_key is not None else openai_key + + # TODO: remove most of this (deplicated with User table) + config = MemGPTConfig( + default_llm_config=LLMConfig( + model=model, + model_endpoint=model_endpoint, + model_endpoint_type=model_endpoint_type, + model_wrapper=model_wrapper, + context_window=context_window, + ), + default_embedding_config=EmbeddingConfig( + embedding_endpoint_type=embedding_endpoint_type, + embedding_endpoint=embedding_endpoint, + embedding_dim=embedding_dim, + embedding_model=embedding_model, + ), + # storage + archival_storage_type=archival_storage_type, + archival_storage_uri=archival_storage_uri, + archival_storage_path=archival_storage_path, + # recall storage + recall_storage_type=recall_storage_type, + recall_storage_uri=recall_storage_uri, + recall_storage_path=recall_storage_path, + # metadata storage (currently forced to match recall storage) + metadata_storage_type=recall_storage_type, + metadata_storage_uri=recall_storage_uri, + metadata_storage_path=recall_storage_path, + ) + + typer.secho(f"📖 Saving config to {config.config_path}", fg=typer.colors.GREEN) + config.save() + + # create user records + ms = MetadataStore(config) + user_id = uuid.UUID(config.anon_clientid) + user = User( + id=uuid.UUID(config.anon_clientid), + ) + if ms.get_user(user_id): + # update user + ms.update_user(user) + else: + ms.create_user(user) + + # create preset records in metadata store + from memgpt.presets.presets import add_default_presets + + add_default_presets(user_id, ms) + + +class ListChoice(str, Enum): + agents = "agents" + humans = "humans" + personas = "personas" + sources = "sources" + presets = "presets" + + +@app.command() +def list(arg: Annotated[ListChoice, typer.Argument]): + config = MemGPTConfig.load() + ms = MetadataStore(config) + user_id = uuid.UUID(config.anon_clientid) + table = ColorTable(theme=Themes.OCEAN) + if arg == ListChoice.agents: + """List all agents""" + table.field_names = ["Name", "LLM Model", "Embedding Model", "Embedding Dim", "Persona", "Human", "Data Source", "Create Time"] + for agent in tqdm(ms.list_agents(user_id=user_id)): + source_ids = ms.list_attached_sources(agent_id=agent.id) + assert all([source_id is not None and isinstance(source_id, uuid.UUID) for source_id in source_ids]) + sources = [ms.get_source(source_id=source_id) for source_id in source_ids] + assert all([source is not None and isinstance(source, Source)] for source in sources) + source_names = [source.name for source in sources if source is not None] + table.add_row( + [ + agent.name, + agent.llm_config.model, + agent.embedding_config.embedding_model, + agent.embedding_config.embedding_dim, + agent.persona, + agent.human, + ",".join(source_names), + utils.format_datetime(agent.created_at), + ] + ) + print(table) + elif arg == ListChoice.humans: + """List all humans""" + table.field_names = ["Name", "Text"] + for human in ms.list_humans(user_id=user_id): + table.add_row([human.name, human.text.replace("\n", "")[:100]]) + print(table) + elif arg == ListChoice.personas: + """List all personas""" + table.field_names = ["Name", "Text"] + for persona in ms.list_personas(user_id=user_id): + table.add_row([persona.name, persona.text.replace("\n", "")[:100]]) + print(table) + elif arg == ListChoice.sources: + """List all data sources""" + + # create table + table.field_names = ["Name", "Description", "Embedding Model", "Embedding Dim", "Created At", "Agents"] + # TODO: eventually look accross all storage connections + # TODO: add data source stats + # TODO: connect to agents + + # get all sources + for source in ms.list_sources(user_id=user_id): + # get attached agents + agent_ids = ms.list_attached_agents(source_id=source.id) + agent_states = [ms.get_agent(agent_id=agent_id) for agent_id in agent_ids] + agent_names = [agent_state.name for agent_state in agent_states if agent_state is not None] + + table.add_row( + [ + source.name, + source.description, + source.embedding_model, + source.embedding_dim, + utils.format_datetime(source.created_at), + ",".join(agent_names), + ] + ) + + print(table) + elif arg == ListChoice.presets: + """List all available presets""" + table.field_names = ["Name", "Description", "Sources", "Functions"] + for preset in ms.list_presets(user_id=user_id): + sources = ms.get_preset_sources(preset_id=preset.id) + table.add_row( + [ + preset.name, + preset.description, + ",".join([source.name for source in sources]), + # json.dumps(preset.functions_schema, indent=4) + ",\n".join([f["name"] for f in preset.functions_schema]), + ] + ) + print(table) + else: + raise ValueError(f"Unknown argument {arg}") + + +@app.command() +def add( + option: str, # [human, persona] + name: Annotated[str, typer.Option(help="Name of human/persona")], + text: Annotated[Optional[str], typer.Option(help="Text of human/persona")] = None, + filename: Annotated[Optional[str], typer.Option("-f", help="Specify filename")] = None, +): + """Add a person/human""" + config = MemGPTConfig.load() + user_id = uuid.UUID(config.anon_clientid) + ms = MetadataStore(config) + if filename: # read from file + assert text is None, "Cannot specify both text and filename" + with open(filename, "r", encoding="utf-8") as f: + text = f.read() + if option == "persona": + persona = ms.get_persona(name=name, user_id=user_id) + if persona: + # config if user wants to overwrite + if not questionary.confirm(f"Persona {name} already exists. Overwrite?").ask(): + return + persona.text = text + ms.update_persona(persona) + else: + persona = PersonaModel(name=name, text=text, user_id=user_id) + ms.add_persona(persona) + + elif option == "human": + human = ms.get_human(name=name, user_id=user_id) + if human: + # config if user wants to overwrite + if not questionary.confirm(f"Human {name} already exists. Overwrite?").ask(): + return + human.text = text + ms.update_human(human) + else: + human = HumanModel(name=name, text=text, user_id=user_id) + ms.add_human(HumanModel(name=name, text=text, user_id=user_id)) + elif option == "preset": + assert filename, "Must specify filename for preset" + create_preset_from_file(filename, name, user_id, ms) + else: + raise ValueError(f"Unknown kind {option}") + + +@app.command() +def delete(option: str, name: str): + """Delete a source from the archival memory.""" + + config = MemGPTConfig.load() + user_id = uuid.UUID(config.anon_clientid) + ms = MetadataStore(config) + assert ms.get_user(user_id=user_id), f"User {user_id} does not exist" + + try: + # delete from metadata + if option == "source": + # delete metadata + source = ms.get_source(source_name=name, user_id=user_id) + assert source is not None, f"Source {name} does not exist" + ms.delete_source(source_id=source.id) + + # delete from passages + conn = StorageConnector.get_storage_connector(TableType.PASSAGES, config, user_id=user_id) + conn.delete({"data_source": name}) + + assert ( + conn.get_all({"data_source": name}) == [] + ), f"Expected no passages with source {name}, but got {conn.get_all({'data_source': name})}" + + # TODO: should we also delete from agents? + elif option == "agent": + agent = ms.get_agent(agent_name=name, user_id=user_id) + assert agent is not None, f"Agent {name} for user_id {user_id} does not exist" + + # recall memory + recall_conn = StorageConnector.get_storage_connector(TableType.RECALL_MEMORY, config, user_id=user_id, agent_id=agent.id) + recall_conn.delete({"agent_id": agent.id}) + + # archival memory + archival_conn = StorageConnector.get_storage_connector(TableType.ARCHIVAL_MEMORY, config, user_id=user_id, agent_id=agent.id) + archival_conn.delete({"agent_id": agent.id}) + + # metadata + ms.delete_agent(agent_id=agent.id) + + elif option == "human": + human = ms.get_human(name=name, user_id=user_id) + assert human is not None, f"Human {name} does not exist" + ms.delete_human(name=name, user_id=user_id) + elif option == "persona": + persona = ms.get_persona(name=name, user_id=user_id) + assert persona is not None, f"Persona {name} does not exist" + ms.delete_persona(name=name, user_id=user_id) + assert ms.get_persona(name=name, user_id=user_id) is None, f"Persona {name} still exists" + elif option == "preset": + preset = ms.get_preset(name=name, user_id=user_id) + assert preset is not None, f"Preset {name} does not exist" + ms.delete_preset(name=name, user_id=user_id) + else: + raise ValueError(f"Option {option} not implemented") + + typer.secho(f"Deleted {option} '{name}'", fg=typer.colors.GREEN) + + except Exception as e: + typer.secho(f"Failed to delete {option}'{name}'\n{e}", fg=typer.colors.RED) diff --git a/memory_scope/config/__init__.py b/memory_scope/config/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/memory_scope/config/bailian_memory_config.py b/memory_scope/config/bailian_memory_config.py new file mode 100644 index 00000000..86857b77 --- /dev/null +++ b/memory_scope/config/bailian_memory_config.py @@ -0,0 +1,137 @@ +from typing import Dict, List + +import dashscope +from pydantic import Field, BaseModel + +from enumeration.memory_type_enum import MemoryTypeEnum + + +class BailianMemoryConfig(BaseModel): + # constant + qwen_max: str = dashscope.Generation.Models.qwen_max + qwen_plus: str = dashscope.Generation.Models.qwen_plus + qwen_turbo: str = dashscope.Generation.Models.qwen_turbo + qwen_long: str = "qwen-long" + + # 关键词映射 + key_word_relate_dict: Dict[str, List[str]] = { + "天气": ["地点", "工作"], + } + + # 用户额外画像 + extra_user_attrs: List[str] = [ + + ] + + # 记忆的系统prompt + default_system_prompt: str = "请判断下面的内容是否可以帮助更好的理解用户问题，如果有用处，请记住下面的内容，如果没有用处，请遗忘下面的内容：" + + # es config + es_index_name: str = Field("memory_index", description="es_index_name") + es_user_name: str = Field("elastic", description="es_user_name") + es_password: str = Field("Beilianmemory_", description="es_password") + + # retry count + dash_generate_retry_cnt: int = Field(2, description="dash_generate retry_cnt") + dash_embedding_retry_cnt: int = Field(5, description="dash_embedding retry_cnt") + dash_rerank_retry_cnt: int = Field(5, description="dash_rerank retry_cnt") + es_retry_cnt: int = Field(5, description="es retry_cnt") + + # other config + seed: int = Field(0, description="global seed") + + # memory request model + messages_pick_n: int = Field(1, description="summary:需要总结的msg个数(偶数);retrieve:传1+History（奇数）") + memory_id: str = Field("", description="base memory id") + workspace_id: str = Field("", description="workspace_id") + api_key: str = Field("", description="api_key") + output_max_count: int = Field(5, description="output_max_count") + + # base request model + trace_id: str = Field("", description="trace_id") + tenant_id: str = Field("", description="tenant_id") + request_id: str = Field("", description="request_id") + uid: str = Field("", description="uid") + account_id: str = Field("", description="account_id") + app_id: str = Field("", description="app_id") + + # top k + es_insight_top_k: int = Field(128, description="es_insight_top_k") + es_keyword_top_k: int = Field(10, description="es_keyword_top_k") + es_new_obs_top_k: int = Field(256, description="es_new_obs_top_k") + es_not_reflected_top_k: int = Field(256, description="es_not_reflected_top_k") + es_similar_top_k: int = Field(128, description="es_similar_top_k") + es_today_obs_top_k: int = Field(128, description="es_today_obs_top_k") + es_insight_similar_top_k: int = Field(128, description="es_insight_similar_top_k") + es_contra_repeat_similar_top_k: int = Field(1, description="es_contra_repeat_similar_top_k") + + # parse_time_model + parse_time_model: str = Field("qwen_1_8_parse_time_service", description="parse_time_model") + parse_time_max_token: int = Field(100, description="parse_time_max_token") + parse_time_temperature: float = Field(0.6, description="parse_time_temperature") + parse_time_top_k: int = Field(1, description="parse_time_top_k") + + # info_score_model + info_filter_msg_max_size: int = Field(200, description="info_filter_msg_max_size") + info_filter_model: str = Field(qwen_max, description="info_filter_model") + info_filter_max_token: int = Field(200, description="info_filter_max_token") + info_filter_temperature: float = Field(0.6, description="info_filter_temperature") + info_filter_top_k: int = Field(1, description="info_filter_top_k") + + # summary_messages_model + summary_messages_model: str = Field(qwen_max, description="summary_messages_model") + summary_messages_max_token: int = Field(500, description="summary_messages_max_token") + summary_messages_temperature: float = Field(0.6, description="summary_messages_temperature") + summary_messages_top_k: int = Field(1, description="summary_messages_top_k") + + # summarize_messages_model + merge_obs_model: str = Field(qwen_max, description="merge_obs_model") + merge_obs_max_token: int = Field(500, description="merge_obs_max_token") + merge_obs_temperature: float = Field(0.6, description="merge_obs_temperature") + merge_obs_top_k: int = Field(1, description="merge_obs_top_k") + + # fuse params + fuse_score_threshold: float = Field(0.1, description="fuse_score_threshold") + fuse_ratio_dict: Dict[str, float] = Field({ + MemoryTypeEnum.CONVERSATION.value: 0.8, + MemoryTypeEnum.OBSERVATION.value: 1.0, + MemoryTypeEnum.OBS_CUSTOMIZED.value: 1.0, + MemoryTypeEnum.INSIGHT.value: 1.5, + MemoryTypeEnum.PROFILE.value: 1.5, + MemoryTypeEnum.PROFILE_CUSTOMIZED.value: 1.5, + }, description="fuse_multiplier_dict") + fuse_time_ratio: float = Field(2.0, description="fuse_ratio_dict") + + # update profile + update_profile_threshold: float = Field(0.1, description="update_profile_threshold") + update_profile_model: str = Field(qwen_max, description="update_profile_model") + update_profile_max_token: int = Field(500, description="update_profile_max_token") + update_profile_temperature: float = Field(0.6, description="update_profile_temperature") + update_profile_top_k: int = Field(1, description="update_profile_top_k") + update_profile_max_thread: int = Field(10, description="update_profile_max_thread") + + # reflection + reflect_obs_cnt_threshold: int = Field(40, description="reflect_obs_cnt_threshold") + reflect_num_questions: int = Field(3, description="reflect_num_questions") + reflect_obs_model: str = Field(qwen_max, description="reflect_obs_model") + reflect_obs_max_token: int = Field(300, description="reflect_obs_max_token") + reflect_obs_temperature: float = Field(0.6, description="reflect_obs_temperature") + reflect_obs_top_k: int = Field(1, description="reflect_obs_top_k") + + # get insight + insight_obs_max_cnt: int = Field(10, description="insight_obs_max_cnt") + get_insight_model: str = Field(qwen_max, description="get_insight_model") + get_insight_max_token: int = Field(500, description="get_insight_max_token") + get_insight_temperature: float = Field(0.6, description="get_insight_temperature") + get_insight_top_k: int = Field(1, description="get_insight_top_k") + + # update insight + update_insight_threshold: float = Field(0.1, description="update_insight_threshold") + update_insight_model: str = Field(qwen_max, description="update_insight_model") + update_insight_max_token: int = Field(500, description="update_insight_max_token") + update_insight_temperature: float = Field(0.6, description="update_insight_temperature") + update_insight_top_k: int = Field(1, description="update_insight_top_k") + update_insight_max_thread: int = Field(10, description="update_insight_max_thread") + + # long contra repeat + long_contra_repeat_threshold: float = Field(0.6, description="long_contra_repeat_threshold") \ No newline at end of file diff --git a/memory_scope/constants/__init__.py b/memory_scope/constants/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/memory_scope/constants/common_constants.py b/memory_scope/constants/common_constants.py new file mode 100644 index 00000000..c276cf02 --- /dev/null +++ b/memory_scope/constants/common_constants.py @@ -0,0 +1,116 @@ +from enumeration.dash_api_enum import DashApiEnum +from enumeration.env_type import EnvType + +APP_ENV = "APP_ENV" + +PIPELINE = "pipeline" + +WORKER = "worker" + +MEMORY = "memory" + +DEFAULT_SYSTEM_PROMPT = "default_system_prompt" + +RELATED_MEMORIES = "related_memories" + +MODIFIED_MEMORIES = "modified_memories" + +RESPONSE_EXT_INFO = "response_ext_info" + +REQUEST = "request" + +CONFIG = "config" + +PROMPT_CONFIG = "prompt_config" + +MESSAGES = "messages" + +EXTRACT_TIME_DICT = "extract_time_dict" + +NEW_OBS_NODES = "new_obs_nodes" + +NEW_OBS_WITH_TIME_NODES = "new_obs_with_time_nodes" + +INSIGHT_NODES = "insight_nodes" + +MERGE_OBS_NODES = "merge_obs_nodes" + +NEW_INSIGHT_NODES = "new_insight_nodes" + +TODAY_OBS_NODES = "today_obs_nodes" + +ALL_NODES = "all_nodes" + +ALL_MEMORIES = "all_memories" + +SIMILAR_OBS_NODES = "similar_obs_nodes" + +KEYWORD_OBS_NODES = "keyword_obs_nodes" + +NOT_REFLECTED_OBS_NODES = "not_reflected_obs_nodes" + +NOT_REFLECTED_MERGE_NODES = "not_reflected_merge_nodes" + +NEW_INSIGHT_KEYS = "new_insight_keys" + +INSIGHT_KEY = "insight_key" + +INSIGHT_VALUE = "insight_value" + +DT = "dt" + +MSG_TIME = "msg_time" + +NEW = "new" + +TIME_INFER = "time_infer" + +KEY_WORD = "key_word" + +REFLECTED = "reflected" + +NEW_USER_PROFILE = "new_user_profile" + +RECALL_TYPE = "recall_type" + +ALL_ONLINE_NODES = "all_online_nodes" + +MAX_WORKERS = "max_workers" + +TIME_MATCHED = "time_matched" + +QUERY_KEYWORDS = "query_keywords" + +DASH_ENV_URL_DICT = { + EnvType.DAILY: "https://dashscope.aliyuncs.com", + # EnvType.PRE: "https://dashscope.aliyuncs.com", + EnvType.PRE: "http://nlb-a3gi6od2xpdx16ezde.cn-beijing.nlb.aliyuncs.com", + EnvType.PROD: "http://ep-2zei3b9a7e2e447bd259.epsrv-2zexnj17q1p8mtjwe3dx.cn-beijing.privatelink.aliyuncs.com", +} + +DASH_API_URL_DICT = { + DashApiEnum.GENERATION: "/api/v1/services/aigc/text-generation/generation", + DashApiEnum.EMBEDDING: "/api/v1/services/embeddings/text-embedding/text-embedding", + DashApiEnum.RERANK: "/api/v1/services/rerank/text-rerank/text-rerank", +} + +ES_ENV_URL_DICT = { + EnvType.DAILY: "http://es-cn-lr53pmrna0002pffb.public.elasticsearch.aliyuncs.com:9200", + EnvType.PRE: "http://ep-bp1i04ae830e377a26a4.epsrv-bp15vzbd1o3umr1girls.cn-hangzhou.privatelink.aliyuncs.com:9200", + EnvType.PROD: "http://ep-2zeibdbbe2904414e741.epsrv-2zet33kwqg8bphgmm36f.cn-beijing.privatelink.aliyuncs.com:9200", +} + +WEEKDAYS = ["周一", "周二", "周三", "周四", "周五", "周六", "周日"] + +DATATIME_WORD_LIST = ["天", "周", "月", "年", "星期", "点", "分钟", "小时", "秒", "上午", "下午", "早上", "早晨", + "晚上", "中午", "日", "夜", "清晨", "傍晚", "凌晨", "岁"] + +TIME_FORMAT_V1 = "{year}年{month}月{day}日{weekday}{hour}点" + +DATATIME_KEY_MAP = { + "年": "year", + "月": "month", + "日": "day", + "周": "week", + "星期几": "weekday", +} diff --git a/memory_scope/db/__init__.py b/memory_scope/db/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/memory_scope/db/elastic_search_client.py b/memory_scope/db/elastic_search_client.py new file mode 100644 index 00000000..74138cf3 --- /dev/null +++ b/memory_scope/db/elastic_search_client.py @@ -0,0 +1,341 @@ +from elasticsearch import Elasticsearch +from elasticsearch.helpers import bulk + +from common.dash_embedding_client import DashEmbeddingClient +from common.logger import Logger +from constants.common_constants import ES_ENV_URL_DICT +from enumeration.env_type import EnvType + + +class ElasticSearchClient(object): + def __init__(self, + es_user_name: str, + es_password: str, + es_index_name: str, + embedding_client: DashEmbeddingClient | None = None, + env_type: EnvType | str = EnvType.DAILY, + content_key: str = "content", + vector_key: str = "vector", + **kwargs): + + self.es_index_name: str = es_index_name + self.embedding_client: DashEmbeddingClient = embedding_client + self.content_key: str = content_key + self.vector_key: str = vector_key + + self.es_client = Elasticsearch( + hosts=[ES_ENV_URL_DICT.get(EnvType(env_type))], + basic_auth=(es_user_name, es_password), + **kwargs) + + self.logger = Logger.get_logger() + self.logger.debug(f"connect es_client info={self.es_client.info()}") + + def log_index_info(self): + index_info = self.es_client.indices.get(index=self.es_index_name) + self.logger.info(f"index={self.es_index_name} exists. index_info={index_info}") + + def insert(self, _id: str, body: dict): + assert body and self.content_key in body, f"body={body} is illegal!" + + # text_type: document + content = body[self.content_key] + vector = self.embedding_client.call(text=content, text_type="document") + if not vector: + self.logger.warning(f"embedding_client call failed, stop es insert!") + return + + body[self.vector_key] = vector + response = self.es_client.index(id=_id, index=self.es_index_name, body=body) + self.logger.info(f"insert response={response}") + + def insert_batch(self, doc_list: list): + """ + doc_list = [ + { + "_id": 2, + "_source": { + "author": "john", + "text": "Elasticsearch: cool.", + "timestamp": "2023-03-23T10:00:00" + } + }, + { + "_id": 3, + "_source": { + "author": "jane", + "text": "Elasticsearch: very cool.", + "timestamp": "2023-03-23T11:00:00" + } + } + ] + """ + text_list = [] + for doc in doc_list: + assert "_id" in doc and "_source" in doc + content = doc["_source"][self.content_key] + text_list.append(content) + + vector_dict = self.embedding_client.call(text=text_list, text_type="document") + if not vector_dict: + self.logger.warning(f"embedding_client call failed, stop es insert!") + return + + # add _index + for i, doc in enumerate(doc_list): + doc["_index"] = self.es_index_name + vector = vector_dict[i] + doc["_source"][self.vector_key] = vector + + # 执行批量插入 + responses = bulk(self.es_client, doc_list) + + # 输出批量插入的响应 + for response in responses[1]: + self.logger.info(f"insert_batch response={response}") + + def print_hits(self, hits: list): + for hit in hits: + print_kwargs = { + "_id": hit['_id'], + "_score": hit['_score'], + } + for k, v in hit['_source'].items(): + # 不打印vector + if k == self.vector_key: + v = len(v) + print_kwargs[k] = v + self.logger.info(" ".join([f"{k}={v}" for k, v in print_kwargs.items()])) + + def exact_search(self, + size: int, + exact_filters: dict = None, + wildcard_filters: dict = None, + print_hits: bool = False, + exclude_vector: bool = True): + """ + { + "match": { + "category": "electronics" # 一级字段过滤 + } + }, + { + "match": { + "product.name": "laptop" # 二级字段过滤 + } + } + { + "terms": { + "product.keyA": ["a", "b", "c"] # 二级字段keyA的精确值必须为a、b、c中的一 + } + } + """ + must_list = [] + for key, value in exact_filters.items(): + if not key: + continue + if isinstance(value, str): + must_list.append({"match": {key: value}}) + elif isinstance(value, list): + must_list.append({"terms": {key: value}}) + + query = { + "size": size, + "query": { + "bool": { + "must": must_list + } + }, + # 添加_source配置以排除vector字段 + "_source": { + "excludes": [self.vector_key] if exclude_vector else [] + } + } + + if wildcard_filters: + should_list = [] + for key, value in wildcard_filters.items(): + if not key: + continue + if isinstance(value, str): + should_list.append({"wildcard": {key: f"*{value}*"}}) + elif isinstance(value, list): + for v in value: + should_list.append({"wildcard": {key: f"*{v}*"}}) + + query["query"]["bool"].update({ + "should": should_list, + "minimum_should_match": 1, + }) + self.logger.info(f"query={query}") + + response = self.es_client.search(index=self.es_index_name, body=query) + hits = response['hits']['hits'] + + # 耗时log + self.logger.info(f"exact_search cost={response['took']}ms " + f"size={len(hits)} " + f"timed_out={response['timed_out']} " + f"shards={response['_shards']} " + f"exact_filters={exact_filters}", stacklevel=2) + + # 每一条结果log一次 + if print_hits: + self.print_hits(hits) + + return hits + + def exact_search_v2(self, + size: int, + term_filters: dict = None, + match_filters: dict = None, + print_hits: bool = False, + exclude_vector: bool = True): + + """ +"bool": { + "must": [ + {"term": {"field1": "固定值"}}, # 一级目录关键字过滤（等于某个值） + {"terms": {"field2": ["a", "b", "c"]}} # 二级目录关键字过滤（等于三个中的任意一个） + ], + "should": [ # 至少匹配其中之一 + {"match": {"key": "ccc"}}, # key包含"ccc" + {"match": {"key": "bbb"}} # 或者key包含"bbb" + ], + "minimum_should_match": 1 # 至少有一个`should`条件匹配 +} + """ + + query = { + "size": size, + "query": { + "bool": { + + } + }, + # 添加_source配置以排除vector字段 + "_source": { + "excludes": [self.vector_key] if exclude_vector else [] + } + } + + if term_filters: + must_list = [] + for k, v in term_filters.items(): + if isinstance(v, list): + must_list.append({"terms": {k: v}}) + elif isinstance(v, str): + must_list.append({"term": {k: v}}) + else: + raise NotImplemented + query["query"]["bool"]["must"] = must_list + + if match_filters: + match_list = [] + for k, v in match_filters.items(): + if isinstance(v, list): + for v_sub in v: + match_list.append({"match": {k: v_sub}}) + elif isinstance(v, str): + match_list.append({"match": {k: v}}) + else: + raise NotImplemented + query["query"]["bool"]["should"] = match_list + query["query"]["bool"]["minimum_should_match"] = 1 + + self.logger.info(query) + response = self.es_client.search(index=self.es_index_name, body=query) + hits = response['hits']['hits'] + + # 耗时log + self.logger.info(f"exact_search cost={response['took']}ms " + f"size={len(hits)} " + f"timed_out={response['timed_out']} " + f"shards={response['_shards']}", stacklevel=2) + + # 每一条结果log一次 + if print_hits: + self.print_hits(hits) + + return hits + + def similar_search(self, + text: str, + size: int, + exact_filters: dict = None, + print_hits: bool = False, + exclude_vector: bool = True): + + if exact_filters is None: + exact_filters = {} + + # 过滤or + or_filters = {} + for k in list(exact_filters.keys()): + v = exact_filters[k] + if isinstance(v, list): + exact_filters.pop(k) + or_filters[k] = v + + vector = self.embedding_client.call(text=text) + if not vector: + self.logger.warning(f"embedding_client call failed, stop select from es!") + return + + query = { + # 返回最相似的top_k个文档 + "size": size, + "query": { + "bool": { + "must": { + "script_score": { + # 对所有文档执行 + "query": { + "match_all": {} + }, + "script": { + # 使用余弦相似度+1,es不能返回负数 + "source": f"cosineSimilarity(params.query_vector, '{self.vector_key}') + 1.0", + "params": {"query_vector": vector} + } + } + }, + "filter": [ + {"term": {k: v}} for k, v in exact_filters.items() + ], + } + }, + # 添加_source配置以排除vector字段 + "_source": { + "excludes": [self.vector_key] if exclude_vector else [] + } + } + + if or_filters: + k_v_pair = [] + for k, v_list in or_filters.items(): + for v in v_list: + k_v_pair.append((k, v)) + query["query"]["bool"]["should"] = [{"term": {k: v}} for k, v in k_v_pair] + query["query"]["bool"]["minimum_should_match"] = 1 + + response = self.es_client.search(index=self.es_index_name, body=query) + hits = response['hits']['hits'] + + # 耗时log + self.logger.info(f"similar_search cost={response['took']}ms " + f"size={len(hits)} " + f"timed_out={response['timed_out']} " + f"shards={response['_shards']} " + f"text={text} " + f"exact_filters={exact_filters}", stacklevel=2) + + # 还原打分 + for hit in hits: + hit['_score'] -= 1 + + # 每一条结果log一次 + if print_hits: + self.print_hits(hits) + + return hits diff --git a/memory_scope/enumeration/__init__.py b/memory_scope/enumeration/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/memory_scope/enumeration/dash_api_enum.py b/memory_scope/enumeration/dash_api_enum.py new file mode 100644 index 00000000..8a58a2af --- /dev/null +++ b/memory_scope/enumeration/dash_api_enum.py @@ -0,0 +1,9 @@ +from enum import Enum + + +class DashApiEnum(str, Enum): + GENERATION = "generation" + + EMBEDDING = "embedding" + + RERANK = "rerank" diff --git a/memory_scope/enumeration/memory_method_enum.py b/memory_scope/enumeration/memory_method_enum.py new file mode 100644 index 00000000..fa8a7092 --- /dev/null +++ b/memory_scope/enumeration/memory_method_enum.py @@ -0,0 +1,11 @@ +from enum import Enum + + +class MemoryMethodEnum(str, Enum): + SUMMARY = "summary" + + RETRIEVE = "retrieve" + + SUMMARY_SHORT = "summary_short" + + SUMMARY_LONG = "summary_long" diff --git a/memory_scope/enumeration/memory_node_status.py b/memory_scope/enumeration/memory_node_status.py new file mode 100644 index 00000000..f8c706b0 --- /dev/null +++ b/memory_scope/enumeration/memory_node_status.py @@ -0,0 +1,7 @@ +from enum import Enum + + +class MemoryNodeStatus(str, Enum): + ACTIVE = "active" + + EXPIRED = "expired" diff --git a/memory_scope/enumeration/memory_recall_type.py b/memory_scope/enumeration/memory_recall_type.py new file mode 100644 index 00000000..956ee4f8 --- /dev/null +++ b/memory_scope/enumeration/memory_recall_type.py @@ -0,0 +1,9 @@ +from enum import Enum + + +class MemoryRecallType(str, Enum): + SIMILAR = "similar" + + KEYWORD = "keyword" + + PROFILE = "profile" diff --git a/memory_scope/enumeration/memory_scene_enum.py b/memory_scope/enumeration/memory_scene_enum.py new file mode 100644 index 00000000..ed9a5fc8 --- /dev/null +++ b/memory_scope/enumeration/memory_scene_enum.py @@ -0,0 +1,11 @@ +from enum import Enum + + +class MemorySceneEnum(str, Enum): + BAILIAN = "BAILIAN" + + TONGYI_MAIN_CHAT = "TONGYI_MAIN_CHAT" + + TONGYI_CHAR_CHAT = "TONGYI_CHAR_CHAT" + + ASSISTANT_API = "ASSISTANT_API" diff --git a/memory_scope/enumeration/memory_type_enum.py b/memory_scope/enumeration/memory_type_enum.py new file mode 100644 index 00000000..e4e72b7e --- /dev/null +++ b/memory_scope/enumeration/memory_type_enum.py @@ -0,0 +1,15 @@ +from enum import Enum + + +class MemoryTypeEnum(str, Enum): + CONVERSATION = "conversation" + + OBSERVATION = "observation" + + INSIGHT = "insight" + + PROFILE = "profile" + + OBS_CUSTOMIZED = "obs_customized" + + PROFILE_CUSTOMIZED = "profile_customized" diff --git a/memory_scope/enumeration/message_role_enum.py b/memory_scope/enumeration/message_role_enum.py new file mode 100644 index 00000000..669ebe01 --- /dev/null +++ b/memory_scope/enumeration/message_role_enum.py @@ -0,0 +1,9 @@ +from enum import Enum + + +class MessageRoleEnum(str, Enum): + USER = "user" + + ASSISTANT = "assistant" + + SYSTEM = "system" diff --git a/memory_scope/models/__init__.py b/memory_scope/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/memory_scope/models/dash_client.py b/memory_scope/models/dash_client.py new file mode 100644 index 00000000..6d07fed4 --- /dev/null +++ b/memory_scope/models/dash_client.py @@ -0,0 +1,91 @@ +import json +import time +from http import HTTPStatus + +import requests + +from common.logger import Logger +from common.timer import Timer +from enumeration.env_type import EnvType + + +class DashClient(object): + + def __init__(self, + request_id: str, + dash_scope_uid: str, + authorization: str, + workspace: str, + model_name: str, + env_type: EnvType | str = EnvType.DAILY, + timeout: int = None, + max_retry_count: int = 2, + retry_sleep_time: float = 1.0, + **kwargs): + + self.model_name: str = model_name + self.env_type: EnvType = EnvType(env_type) + self.timeout: int = timeout + self.max_retry_count: int = max_retry_count + self.retry_sleep_time: float = retry_sleep_time + self.kwargs: dict = kwargs + + # 20240506 update by 泉雨 + # if authorization: + # workspace = "" + # dash_scope_uid = "" + + self.headers = { + 'Content-Type': 'application/json', + 'Authorization': authorization, + 'X-Request-Id': request_id, + 'X-DashScope-Uid': dash_scope_uid, + 'X-DashScope-WorkSpace': workspace, + } + + self.url: str = "" + self.data = {} + + self.logger = Logger.get_logger() + + def before_call(self, model_name: str = None, **kwargs): + pass + + def after_call(self, response_obj, **kwargs): + pass + + def call_once(self, model_name: str = None, retry_cnt: int = 0, **kwargs): + if model_name is None: + model_name = self.model_name + + self.before_call(model_name=model_name, **kwargs) + + with Timer(self.__class__.__name__, log_time=False) as t: + self.logger.debug(f"url={self.url} header={self.headers} data={self.data} timeout={self.timeout}") + response = requests.post(url=self.url, + headers=self.headers, + data=json.dumps(self.data), + timeout=self.timeout) + + if response.status_code == HTTPStatus.OK: + response_obj = json.loads(response.text) + self.logger.info(f"{self.__class__.__name__} env={self.env_type.value} {t.get_cost_info()}, " + f"call model={model_name} success! retry_cnt={retry_cnt}", + stacklevel=3) + return self.after_call(response_obj, **kwargs), True + + else: + self.logger.warning(f"{self.__class__.__name__} env={self.env_type.value} {t.get_cost_info()}, " + f"call model={model_name} failed! retry_cnt={retry_cnt} details={response.text}", + stacklevel=3) + return None, False + + def call(self, model_name: str = None, **kwargs): + for i in range(self.max_retry_count): + result, flag = self.call_once(model_name=model_name, retry_cnt=i, **kwargs) + if flag: + return result + else: + time.sleep(self.retry_sleep_time) + + return None diff --git a/memory_scope/models/dash_embedding_client.py b/memory_scope/models/dash_embedding_client.py new file mode 100644 index 00000000..3529bfae --- /dev/null +++ b/memory_scope/models/dash_embedding_client.py @@ -0,0 +1,43 @@ +from typing import List, Dict + +import dashscope + +from common.dash_client import DashClient +from constants.common_constants import DASH_ENV_URL_DICT, DASH_API_URL_DICT +from enumeration.dash_api_enum import DashApiEnum + + +class DashEmbeddingClient(DashClient): + """ + url: https://help.aliyun.com/document_detail/2782232.html?spm=a2c4g.2782227.0.0.76195b1d9UeBAk#a6a39590fegqx + """ + + def __init__(self, model_name: str = dashscope.TextEmbedding.Models.text_embedding_v2, **kwargs): + super(DashEmbeddingClient, self).__init__(model_name=model_name, **kwargs) + self.url = DASH_ENV_URL_DICT.get(self.env_type) + DASH_API_URL_DICT.get(DashApiEnum.EMBEDDING) + + def before_call(self, model_name: str = None, **kwargs): + text: str | List[str] = kwargs.pop("text", "") + # text_type: query or document + text_type: str = kwargs.pop("text_type", "query") + + if isinstance(text, str): + text = [text] + + self.kwargs["text_type"] = text_type + self.data = { + "model": model_name, + "input": { + "texts": text, + }, + "parameters": {**kwargs, **self.kwargs}, + } + + def after_call(self, response_obj, **kwargs) -> Dict[int, List[float]] | List[float]: + embedding_results = {} + for emb in response_obj["output"]["embeddings"]: + embedding_results[emb["text_index"]] = emb["embedding"] + + if len(embedding_results) == 1: + embedding_results = list(embedding_results.values())[0] + return embedding_results diff --git a/memory_scope/models/dash_generate_client.py b/memory_scope/models/dash_generate_client.py new file mode 100644 index 00000000..de6e11ec --- /dev/null +++ b/memory_scope/models/dash_generate_client.py @@ -0,0 +1,45 @@ +from typing import List, Dict + +import dashscope + +from common.dash_client import DashClient +from constants.common_constants import DASH_ENV_URL_DICT, DASH_API_URL_DICT +from enumeration.dash_api_enum import DashApiEnum + + +class DashGenerateClient(DashClient): + """ + url: https://help.aliyun.com/document_detail/2712576.html + """ + + def __init__(self, model_name: str = dashscope.Generation.Models.qwen_max, **kwargs): + super(DashGenerateClient, self).__init__(model_name=model_name, **kwargs) + self.url = DASH_ENV_URL_DICT.get(self.env_type) + DASH_API_URL_DICT.get(DashApiEnum.GENERATION) + + def before_call(self, model_name: str = None, **kwargs): + prompt: str = kwargs.pop("prompt", "") + messages: List[Dict[str, str]] = kwargs.pop("messages", []) + + input_text = {} + if prompt: + input_text["prompt"] = prompt + elif messages: + input_text["messages"] = messages + else: + raise RuntimeError("prompt and messages is both empty!") + + self.data = { + "model": model_name, + "input": input_text, + "parameters": {**kwargs, **self.kwargs}, + } + + def after_call(self, response_obj, **kwargs): + self.logger.debug(f"response_obj={response_obj}") + output = response_obj["output"] + if "text" in output: + return output["text"] + elif "choices" in output: + return output["choices"][0]["message"]["content"] + else: + raise NotImplementedError diff --git a/memory_scope/models/dash_rerank_client.py b/memory_scope/models/dash_rerank_client.py new file mode 100644 index 00000000..be18766c --- /dev/null +++ b/memory_scope/models/dash_rerank_client.py @@ -0,0 +1,43 @@ +from typing import List + +import dashscope + +from common.dash_client import DashClient +from constants.common_constants import DASH_ENV_URL_DICT, DASH_API_URL_DICT +from enumeration.dash_api_enum import DashApiEnum + + +class DashReRankClient(DashClient): + """ + url: https://help.aliyun.com/document_detail/2780059.html + """ + + def __init__(self, model_name: str = dashscope.TextReRank.Models.gte_rerank, **kwargs): + super(DashReRankClient, self).__init__(model_name=model_name, **kwargs) + self.url = DASH_ENV_URL_DICT.get(self.env_type) + DASH_API_URL_DICT.get(DashApiEnum.RERANK) + + def before_call(self, model_name: str = None, **kwargs): + query: str = kwargs.pop("query", "") + documents: List[str] = kwargs.pop("documents", []) + top_n: int | None = kwargs.pop("top_n", None) + return_documents: bool = kwargs.pop("return_documents", False) + + assert query and documents, f"query or documents is empty! query={query}, documents={len(documents)}" + if top_n is None: + top_n = len(documents) + + self.kwargs.update({ + "top_n": top_n, + "return_documents": return_documents, + }) + self.data = { + "model": model_name, + "input": { + "query": query, + "documents": documents, + }, + "parameters": {**kwargs, **self.kwargs}, + } + + def after_call(self, response_obj, **kwargs): + return response_obj["output"]["results"] diff --git a/memory_scope/node/__init__.py b/memory_scope/node/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/memory_scope/node/memory_node.py b/memory_scope/node/memory_node.py new file mode 100644 index 00000000..4a682fd9 --- /dev/null +++ b/memory_scope/node/memory_node.py @@ -0,0 +1,71 @@ +import re +from typing import Dict, List + +from pydantic import Field, BaseModel + + +class MemoryNode(BaseModel): + """ + 除了 content_modified，其他均和数据库字段保持统一 + 根据code判断，如果code是空，则为新增的memoryNode，如果有值，则为更新 + if content_modified is true，则需要调用embedding服务 + """ + id: str = Field("", description="唯一主键 uuid64") + + code: str = Field("", description="和id保持一致，为空则是新增") + + # 0520新增 + timeCreated: str = Field("", description="Memory创建时间（算法不关注）") + + # 0520新增 + timeModified: str = Field("", description="Memory更新时间（算法不关注）") + + content: str = Field("", description="记忆内容") + + memoryId: str = Field("", description="记忆 id，检索区分字段") + + # 0520新增 + scene: str = Field("", description="source: TONGYI_MAIN_CHAT/TONGYI_CHAR_CHAT/BAILIAN/ASSISTANT") + + # 0520新增 + # NOTE 百炼服务端只召回observation, insight, profile, obs_customized, profile_customized + memoryType: str = Field("", description="conversation, observation, insight, " + "profile, obs_customized, profile_customized") + + # 0520新增，但不是数据库字段 + content_modified: bool = Field(False, description="content是否被更新，if true，则需要调用embedding服务") + + # reflected: 1 is reflected before, 0 has not reflected, 如果是用户自定义，写入空值"". + metaData: Dict[str, str] = Field({}, description="元信息: infoScore, algoVersion, datetime, reflected") + + status: str = Field("active", description="active or expired") + + tenantId: str = Field("", description="request id") + + vector: List[float] = Field([], description="content embedding result, return empty") + + def get_time_info(self, time_format: str): + pattern = re.compile(r'\{([^}]*)}') + keys = pattern.findall(time_format) + + match_flag = True + kv_dict = {} + for k in keys: + if k not in self.metaData: + match_flag = False + break + v = self.metaData[k] + if not v: + match_flag = False + break + + kv_dict[k] = v + + if match_flag: + return time_format.format(**kv_dict) + return "" + + def to_dict(self): + res = {"content": self.content, "memoryId": self.memoryId, "memoryType": self.memoryType, + "status": self.status, "metaData": self.metaData} + return res diff --git a/memory_scope/node/memory_wrap_node.py b/memory_scope/node/memory_wrap_node.py new file mode 100644 index 00000000..84029fe6 --- /dev/null +++ b/memory_scope/node/memory_wrap_node.py @@ -0,0 +1,33 @@ +from pydantic import Field, BaseModel + +from model.memory_node import MemoryNode + + +class MemoryWrapNode(BaseModel): + id: str = Field("", description="uuid64") + + score_similar: float = Field(0, description="相似度打分") + + score_rank: float = Field(0, description="排序打分") + + score_rerank: float = Field(0, description="重排打分") + + memory_node: MemoryNode = Field(None, description="memory node 核心，返回给上游的结构") + + @classmethod + def init_from_es(cls, hit: dict): + memory_node = MemoryNode(**hit['_source']) + return cls(id=hit['_id'], score_similar=hit['_score'], memory_node=memory_node) + + @classmethod + def init_from_attrs(cls, **kwargs): + _id: str = kwargs.get("_id", "") + score_similar: float = kwargs.pop("score_similar", 0) + score_rank: float = kwargs.pop("score_rank", 0) + score_rerank: float = kwargs.pop("score_rerank", 0) + memory_node = MemoryNode(**kwargs) + return cls(id=_id, + score_similar=score_similar, + score_rank=score_rank, + score_rerank=score_rerank, + memory_node=memory_node) diff --git a/memory_scope/node/message.py b/memory_scope/node/message.py new file mode 100644 index 00000000..e458ddab --- /dev/null +++ b/memory_scope/node/message.py @@ -0,0 +1,11 @@ +from pydantic import Field, BaseModel + + +class Message(BaseModel): + role: str = Field(..., description="The role of the message sender (user, assistant, system)") + + content: str = Field(..., description="The body of the message") + + time_created: str = Field("", description="Timestamp when the message was created") + + info_score: str = Field("", description="2 > 1 > 0") diff --git a/memory_scope/node/user_attribute.py b/memory_scope/node/user_attribute.py new file mode 100644 index 00000000..2b9daaa0 --- /dev/null +++ b/memory_scope/node/user_attribute.py @@ -0,0 +1,36 @@ +from typing import Dict, List + +from pydantic import Field, BaseModel + + +class UserAttribute(BaseModel): + """ + 用户画像的一条属性，和数据库保持一致，只会选择status为1的属性透传过来。 + status会透传过来。 + 如果code为空，则为新增，否则是更新。 + 确保请求是10条，返回是原始10条+加上新增的条数（如果可以新增）。只会对正确的请求操作数据库。 + """ + code: str = Field("", description="唯一主键 code") + + memory_id: str = Field("", description="memory id") + + # 上游可能没有传这个参数，可能隐藏在memory_id做区分 + scene: str = Field("", description="source: TONGYI_MAIN_CHAT/TONGYI_CHAR_CHAT/BAILIAN/ASSISTANT") + + # 从key改成memory_key + memory_key: str = Field("", description="memory key") + + value: List[str] = Field([], description="value") + + is_unique: int = Field(1, description="属性是否唯一，if 1 value只有一个，if 0, value 可以很多个") + + is_mutable: int = Field(1, description="是否可变，if 1，value可变，if 1，不可变（用户定义）") + + memory_type: str = Field("", description="profile, profile_customized") + + description: str = Field("", description="memory id") + + status: int = Field(1, + description="0为删除，1为active，状态，算法不感知，只为了保存用户删除的画像，给算法传status为valid的用户画像") + + ext_info: Dict[str, str] = Field({}, description="占位符字典") diff --git a/memory_scope/parsers/__init__.py b/memory_scope/parsers/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/memory_scope/parsers/response_text_parser.py b/memory_scope/parsers/response_text_parser.py new file mode 100644 index 00000000..cfeee808 --- /dev/null +++ b/memory_scope/parsers/response_text_parser.py @@ -0,0 +1,33 @@ +import re + +from common.logger import Logger + + +class ResponseTextParser(object): + pattern_v1 = re.compile(r'<(.*?)>') + + def __init__(self, response_text: str): + self.response_text: str = response_text.strip() + self.logger: Logger = Logger.get_logger() + + def parse_v1(self, prefix: str = ""): + result = [] + for line in self.response_text.split('\n'): + line = line.strip() + if not line: + continue + matches = [match.group(1) for match in self.pattern_v1.finditer(line)] + if matches: + result.append(matches) + self.logger.info(f"{prefix} response_text={self.response_text} result={result}", stacklevel=2) + return result + + def parse_v2(self, prefix: str = ""): + result = [] + for line in self.response_text.split('\n'): + line = line.strip() + if not line or line == "无": + continue + result.append(line) + self.logger.info(f"{prefix} response_text={self.response_text} result={result}", stacklevel=2) + return result diff --git a/memory_scope/pipeline/__init__.py b/memory_scope/pipeline/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/memory_scope/pipeline/memory.py b/memory_scope/pipeline/memory.py new file mode 100644 index 00000000..61258c45 --- /dev/null +++ b/memory_scope/pipeline/memory.py @@ -0,0 +1,30 @@ +from typing import List, Dict + +from pydantic import Field + +from model.message import Message +from model.user_attribute import UserAttribute +from request.base_model import RequestBaseModel + + +class MemoryServiceRequestModel(RequestBaseModel): + messages: List[Message] = Field(..., + description="summary: 多轮对话的list,默认按照时间正序，最后一条是最新的; retrieve: 最后一条是query") + + messages_pick_n: int = Field(1, description="summary：传需要总结的msg的个数;retrieve：不传") + + memory_id: str = Field(..., description="memory id") + + workspace_id: str = Field("", description="workspace id") + + api_key: str = Field("", description="api id") + + scene: str = Field("", description="需要枚举来源: TONGYI_MAIN_CHAT, TONGYI_CHAR_CHAT, BAILIAN, ASSISTANT_API") + + algo_version: str = Field("", description="算法版本，只在做AB实验时透传") + + output_max_count: int = Field(3, description="retrieve时最多的条数，约定3-10条") + + user_profile: List[UserAttribute] = Field([], description="user_profile") + + ext_info: Dict[str, str] = Field({}, description="extra information") diff --git a/memory_scope/pipeline/memory_service_bailian.py b/memory_scope/pipeline/memory_service_bailian.py new file mode 100644 index 00000000..0b122940 --- /dev/null +++ b/memory_scope/pipeline/memory_service_bailian.py @@ -0,0 +1,162 @@ +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 common.context_handler import ContextHandler +from common.logger import Logger +from common.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 request.memory import MemoryServiceRequestModel +from worker.bailian.base_worker import BaseWorker + + +class MemoryServiceBailian(object): + THREAD_POOL_MAX_COUNT: int = 5 + + MEMORY_WORKER_PATH: str = "worker.bailian" + + def __init__(self, request: MemoryServiceRequestModel, method: MemoryMethodEnum): + # 全局上下文，worker之间交换参数和变量 + self.context_handler = ContextHandler(scene=request.scene, + method=method.value, + algo_version=request.algo_version) + self.context_handler.set_context(common_constants.REQUEST, request) + + # 线程池 + max_workers = self.context_handler.get_env_config(MAX_WORKERS) + if max_workers: + self.max_workers = int(max_workers) + else: + self.max_workers = self.THREAD_POOL_MAX_COUNT + self.thread_pool = ThreadPoolExecutor(max_workers=self.max_workers) + + # 运行信息 + self.run_infos = [] + + # 全部初始化的worker + self.worker_dict: Dict[str, BaseWorker] = {} + + # 日志 + self.logger: Logger = Logger.get_memory_logger() + + def get_worker(self, worker_name: str, is_multi_thread: bool = False) -> BaseWorker: + # 更新worker name + worker_name_split = worker_name.split(".") + worker_name = worker_name_split[-1] + if common_constants.WORKER not in worker_name: + worker_name = f"{worker_name}_{common_constants.WORKER}" + + # 构造path + worker_paths = [self.MEMORY_WORKER_PATH] + worker_paths.extend(worker_name_split[:-1]) + worker_paths.append(worker_name) + module = import_module(".".join(worker_paths)) + + worker_clazz_name = under_line_to_hump(worker_name) + return getattr(module, worker_clazz_name)(context_handler=self.context_handler, + 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) + + @timer + def get_pipeline(self) -> list[list]: + pipeline_str = self.context_handler.get_env_config(PIPELINE) + 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): + pipeline_list = self.get_pipeline() + self.print_and_init_worker(pipeline_list) + + # run workers in multi threads + with self.thread_pool, Timer("ALL_PIPELINE"): + for pipeline_part in 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/prompts/__init__.py b/memory_scope/prompts/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/memory_scope/prompts/bailian_prompt_config.py b/memory_scope/prompts/bailian_prompt_config.py new file mode 100644 index 00000000..7d3a3f38 --- /dev/null +++ b/memory_scope/prompts/bailian_prompt_config.py @@ -0,0 +1,626 @@ +from pydantic import BaseModel + + +class BailianPromptConfig(BaseModel): + """ + NOTE 不需要strip() 拼装会统一strip + """ + + info_filter_system: str = """ +任务指令：对所给{batch_size}个句子中所含有的关于用户的信息打分，分数为0,1,2或3。 +注意：其中0表示不包含用户信息，1表示句子中包含用户假设的信息或者用户虚构的内容，2表示可以对用户信息做一些不准确的猜测，3表示明确含有关于用户的有效信息或者可以推断出用户准确的信息或者用户要求记录。 +按如下格式输出, 每一行输出一个打分，一定加<>，一共输出{batch_size}个分数: +结果： +<分数:0或1或2或3> + """ + + info_filter_few_shot: str = """ +示例1 +句子： +1 用户：帮我写一段给同事张三女儿三岁生日的祝福语。 +2 用户：公元1400年至1550年中国历史大事表。 +3 用户：你吃午饭了吗？ +4 用户：我今天心情不好，可以安慰我一下吗？ +5 用户：能给我整理一张如何使用大模型的技巧列表吗，要求内容尽量精简。 +6 用户：明天下午3点提醒我去拿一下文件。 +结果： +<3> +<0> +<0> +<2> +<2> +<3> + +示例2 +句子： +1 用户：我刚刚入职了阿里巴巴。 +2 用户：露天睡觉蚊子多，咋搞。 +3 用户：创造力和外倾性有关？ +4 用户：一个区县的所有的事业人员的档案审核、修改和规范，应该是县委组织部下属的干部档案中心负责还是县人社局负责？ +5 用户：假如我要和一个女人准备要孩子，我作为男人，怎么保护女人和孩子以及怎么备孕确保精子质量高对后代好 +6 用户：我和你一起出去玩，你会感觉开心吗？ +结果： +<3> +<2> +<0> +<0> +<1> +<1> + +示例3 +句子： +1 用户：你的妈妈患有焦虑症，怎么安慰和开导她？ +2 用户：肾脏严重亏空 +3 用户：我很喜欢打篮球，所以我身体很好 +4 用户：篮球明星有哪些？ +结果： +<1> +<1> +<3> +<0> + """ + + info_filter_user_query: str = """ +句子： +{user_query} +结果： + """ + + """ + 思考：思考的依据和过程，不超过20字。 + """ + get_observation_system: str = """ +任务：从下面的{num_obs}句用户句子中依次提取出关于用户的重要信息，与相应的关键词。最多提取{num_obs}条信息。对每一句句子，只提取非常明确的信息，不要进行任何推测。 +不要提取重复的信息，如果句子中的所有信息与已经提取出的信息重复了则回答“重复“，如果没有重要信息则回答“无”。 +请一定要按如下格式依次输出，最后的结果一定要加<>: +信息：<句子序号> <> <明确的重要信息或“重复”或”无“> <关键词> + """ + + """ +思考：从第1句可以得知用户现在没有工作，负债几万，这是关于用户工作与经济状况的重要信息。 +思考：第2句是用户对他人观点的讨论和疑问，没有明确提及用户个人信息。 +思考：第3句含有的信息与第1句重复了。 +思考：从第4句可以得知用户是一个刚毕业的学生，这是关于用户身份背景状况的重要信息。其余信息重要性不足。 +思考：从第1句可以得知张三是用户的同事，这是关于用户的人际关系的重要信息。其余信息重要性不足。 +思考：第2句是用户提出的要求，没有明确提及用户个人信息。 +思考：从第3句可以得知用户前天打羽毛球时把球拍打断了线，但这不是重要的信息。还可以得知用户对猫毛过敏，这是关于用户的健康的重要信息。 +思考：从第4句是用户提出的要求，没有明确提及用户个人信息。 +思考：从第5句可以得知用户在阿里巴巴徐汇滨江园区工作，这是关于用户的工作地点的重要信息。 +思考：从第1句可以得知用户寻求购买新能源汽车的建议或推荐，这是这是关于用户的大宗消费的重要的信息。 +思考：从第2句可以得知用户当前所在城市为上海，这是关于用户的生活地区的重要信息。其余信息与第1句重复了。 +思考：第3句是用户对某个观点的讨论和疑问，没有明确提及用户个人信息。 +思考：第4句是用户提出的要求，没有明确提及用户个人信息。 +思考：从第5句可以得知用户购买了海天味业股票，购买数量为100股，购买金额为5000元，这是关于用户的投资决策的重要信息。 + """ + get_observation_few_shot: str = """ +示例1： +用户句子： +1 用户：我现在处境很糟，没有工作，负债几万，怎么办 +2 用户：有人说兴趣是最好的老师，也建议兴趣和职业联系起来，但我发现喜欢打篮球的人很多，但靠打篮球成职业的稀少，赚钱的更少，此外，怎么分辨兴趣和喜欢 +3 用户：我现在处境很糟，没有工作，负债几万，怎么办 +4 用户：我是一个刚毕业的学生，对社会，行业不了解，给我介绍一下社会系统和行业格局 +思考：从第1句可以得知用户现在没有工作，负债几万，这是关于用户工作与经济状况的重要信息。 +信息：<1> <> <用户当前无工作且负债几万> <无工作, 负债几万> +思考：第2句是用户对他人观点的讨论和疑问，没有明确提及用户个人信息。 +信息：<2> <> <无> <> +思考：第3句含有的信息与第1句重复了。 +信息：<3> <> <重复> <> +思考：从第4句可以得知用户是一个刚毕业的学生，这是关于用户身份背景状况的重要信息。其余信息重要性不足。 +信息：<4> <> <用户是一名刚毕业的学生。> <刚毕业, 学生> + +示例2： +用户句子： +1 用户：帮我写一段给同事张三女儿三岁生日的祝福语。 +2 用户：能给我整理一张如何使用大模型的技巧列表吗，要求内容尽量精简。 +3 用户：两个坏消息，我打羽毛球把拍子打断线了。。。然后我去我朋友家撸猫，结果我猫毛过敏，今天疯狂打喷嚏。。。 +4 用户：公元1400年至1550年中国历史大事表。 +5 用户：谢啦。我中午在公司附近吃，帮我推荐一家阿里巴巴徐汇滨江园区附近的餐厅吧。 +思考：从第1句可以得知张三是用户的同事，这是关于用户的人际关系的重要信息。其余信息重要性不足。 +信息：<1> <> <张三是用户的同事。> <张三, 同事> +思考：第2句是用户提出的要求，没有明确提及用户个人信息。 +信息：<2> <> <无> <> +思考：从第3句可以得知用户前天打羽毛球时把球拍打断了线，但这不是重要的信息。还可以得知用户对猫毛过敏，这是关于用户的健康的重要信息。 +信息：<3> <> <用户对猫毛过敏。> <猫毛, 过敏> +思考：从第4句是用户提出的要求，没有明确提及用户个人信息。 +信息：<4> <> <无> <> +思考：从第5句可以得知用户在阿里巴巴徐汇滨江园区工作，这是关于用户的工作地点的重要信息。 +信息：<5> <> <用户在阿里巴巴徐汇滨江园区工作。> <阿里巴巴, 徐汇滨江园区, 工作> + +示例3： +用户句子： +1 用户：我想买辆新能源汽车，有什么推荐吗？ +2 用户：我在上海，想买辆新能源汽车，有什么推荐吗？ +3 用户：案外人异议审查期间，人民法院不得对执行标的进行处分，不就是中止执行的意思吗？ +4 用户：请写两句藏头诗分别以“胜”和“利”开头。 +5 用户：我花5000元买了100股海天味业。 +思考：从第1句可以得知用户寻求购买新能源汽车的建议或推荐，这是这是关于用户的大宗消费的重要的信息。 +信息：<1> <> <用户寻求购买新能源汽车的建议或推荐。> <购买, 新能源汽车> +思考：从第2句可以得知用户当前所在城市为上海，这是关于用户的生活地区的重要信息。其余信息与第1句重复了。 +信息：<2> <> <用户所在的城市是上海。> <上海> +思考：第3句是用户对某个观点的讨论和疑问，没有明确提及用户个人信息。 +信息：<3> <> <无> <> +思考：第4句是用户提出的要求，没有明确提及用户个人信息。 +信息：<4> <> <无> <> +思考：从第5句可以得知用户购买了海天味业股票，购买数量为100股，购买金额为5000元，这是关于用户的投资决策的重要信息。 +信息：<5> <> <用户购买了海天味业股票，购买数量为100股，购买金额为5000元。> <海天味业, 股票> + """ + + get_observation_user_query: str = """ +用户句子： +{user_query} + """ + + get_observation_with_time_system: str = """ +任务：从下面的{num_obs}句用户句子中依次提取出关于用户的重要信息，相应的关键词与时间信息。 +每一句用户句子的格式是：<序号> <对话时间> 用户：<句子> +对每一句句子，只提取非常明确的重要信息，不要进行任何推测。不要提取重复的信息，如果句子中的所有信息与已经提取出的信息重复了则回答“重复“，如果没有重要信息则回答“无”。 +如果用户信息涉及时间，则结合对话时间推断用户信息的时间信息，没有则不输出。 +请一步步思考，并一定要按如下格式依次输出，最后的结果一定要加<>: +信息：<句子序号> <时间信息或“无”> <明确的重要信息或“重复”或”无“> <关键词> + """ + + get_observation_with_time_few_shot: str = """ +示例1： +句子： +1 2022年5月1日周二3点 用户：帮我写一段给同事张三女儿三岁生日的祝福语。 +2 2022年5月2日周二17点 用户：公元1400年至1550年中国历史大事表。 +3 2022年5月3日周二18点 用户：能给我整理一张如何使用大模型的技巧列表吗，要求内容尽量精简。 +4 2022年7月3日周四12点 用户：上上个月我办了游泳卡。 + +思考：从第1句可以得知张三是用户的同事，这是关于用户的人际关系的重要信息。其余信息重要性不足。用户信息不涉及时间。 +信息：<1> <> <张三是用户的同事。> <张三, 同事> +思考：第2句是用户提出的要求，没有明确提及用户个人信息。 +信息：<2> <> <无> <> +思考：第3句是用户提出的要求，没有明确提及用户个人信息。 +信息：<3> <> <无> <> +思考：从第4句可以得出用户上上个月办了游泳卡。用户信息涉及时间，结合对话时间为2022年7月，推断用户在2022年5月用户办了游泳卡。 +信息：<4> <2022年5月> <用户在2022年5月办了游泳卡。> <游泳卡> + + +示例2： +句子： +1 2020年1月4日周日10点 用户：我花5000元买了100股海天味业。 +2 2023年4月27日周五8点 用户：明天是我和妻子的结婚纪念日，帮我推荐一家餐厅。 +3 2020年1月4日周日10点 用户：我花5000元买了100股海天味业。 +4 2021年6月2日周四23点 用户：谢啦。我中午在公司附近吃，帮我推荐一家阿里巴巴徐汇滨江园区附近的餐厅吧。 +5 2021年7月9日周六11点 用户：两个坏消息，我打羽毛球把拍子打断线了。。。然后我去我朋友家撸猫，结果我猫毛过敏，今天疯狂打喷嚏。。。 + +思考：从第1句可以得知用户购买了海天味业股票，购买数量为100股，购买金额为5000元，这是关于用户的投资决策的重要信息。用户信息不涉及时间。 +信息：<1> <> <用户购买了海天味业股票，购买数量为100股，购买金额为5000元。> <海天味业, 股票> +思考：从第2句可以得知用户与妻子的结婚纪念日是明天，这是关于用户重要纪念日的信息。其余信息重要性不足。用户信息涉及时间，结合对话时间为2023年4月27日， +以及结婚纪念日为周期性日期，推断用户与妻子的结婚纪念日是每年4月28日。 +信息：<2> <每年4月28日> <用户与妻子的结婚纪念日是每年4月28日。> <妻子, 结婚纪念日> +思考：第3句含有的信息与第1句重复了。 +信息：<3> <> <重复> <> +思考：从第4句以得知用户在阿里巴巴徐汇滨江园区工作，这是关于用户的工作的重要信息。其余信息重要性不足。用户信息不涉及时间。 +信息：<4> <> <用户在阿里巴巴徐汇滨江园区工作。> <阿里巴巴, 徐汇滨江园区, 工作> +思考：从第5句可以得知用户前天打羽毛球时把球拍打断了线，但这不是重要的信息。还可以得知用户对猫毛过敏，这是关于用户的健康的重要信息。用户信息不涉及时间。 +信息：<5> <> <用户对猫毛过敏。> <猫毛, 过敏> + + +示例3： +句子： +1 2023年6月30日周五15点 用户：上个月我和家人一起去杭州旅游，景色很不错。 +2 2023年7月2日周二10点 用户：昨天是我生日，一个人过的。 +3 2020年7月3日周四11点 用户：提醒我下周一去体检。 +4 2023年5月21日周六14点 用户：有人说兴趣是最好的老师，也建议兴趣和职业联系起来，但我发现喜欢打篮球的人很多，但靠打篮球成职业的稀少，赚钱的更少，此外，怎么分辨兴趣和喜欢 + +思考：从第1句可以得知用户和家人上个月去杭州旅游了，这是关于用户的经历的重要信息。其余信息重要性不足。用户信息涉及时间，结合对话时间为2023年6月推断用户和家人2023年5月去杭州旅游了。 +信息：<1> <2023年5月> <用户和家人2023年5月去杭州旅游了。> <家人, 杭州, 旅游> +思考：从第2句可以得知用户的生日是昨天，这是关于用户重要纪念日的信息。其余信息重要性不足。用户信息涉及时间，结合对话时间为2023年7月2日， +以及生日为周期性日期，推断用户的生日是每年7月2日。 +信息：<2> <每年7月2日> <用户的生日是每年7月2日。> <生日> +思考：从第3句可以得出用户下周一去体检，这是用户要求记忆的重要信息。用户信息涉及时间，结合对话时间为2020年7月3日周四，推断用户2020年7月6日周一去体检。 +信息：<3> <2020年7月6日周一> <用户2020年7月6日周一去体检。> <体检> +思考：第4句是用户对他人观点的讨论和疑问，没有明确提及用户个人信息。 +信息：<4> <> <无> <> + """ + + get_observation_with_time_user_query: str = """ +用户句子： +{user_query} + """ + + contra_repeat_system: str = """ +对下面的{num_obs}句句子，逐一判断是否与“前面序号”的任意句子存在信息的矛盾，或者句子的主要信息被“前面序号”的任意句子中的信息包含。只判断与“前面序号”的句子的关系。 +请一步步思考，并按如下格式输出: +思考：思考的依据和过程，30字以内。 +判断：<句子序号> <矛盾，被包含，无>，一定加<> +""" + + contra_repeat_few_shot: str = """ +示例1 +句子： +1 用户经常失眠，对安眠药的效果感兴趣，暗示可能考虑使用。 +2 用户经常失眠，寻求缓解方法。 +3 陈伟业是用户的领导 +4 陈伟业是用户的领导 +5 陈伟业是用户的领导，是银行分行行长 + +思考：第1句不会存在与前面序号句子的矛盾或者完全重复。 +判断：<1> <无> +思考：第2句中所有信息都被前面序号中第1句的信息完全包含。 +判断：<2> <被包含> +思考：第3句信息没有在前面序号句子中出现 +判断：<3> <无> +思考：第4句与前面序号中第3句的信息完全重复，即被完全包含。 +判断：<4> <被包含> +思考：第5句中陈伟业是用户的领导的信息被前面序号中第3句的信息包含，但新增了陈伟业是银行分行行长的信息，故不是被完全包含。 +判断：<5> <无> + +示例2 +句子： +1 用户的孩子成绩不太好。 +2 用户的孩子在学校经常逃课。 +3 用户的父亲生日在2024年6月2日，用户打算准备礼物。 +4 用户的父亲生日在2024年5月1日。 +5 用户很喜欢和同班同学打篮球。 +6 用户喜欢打篮球。 + +思考：第1句不会存在与前面序号句子的矛盾或者完全重复。 +判断：<1> <无> +思考：第2句与前面序号句子既不矛盾也不重复。 +判断：<2> <无> +思考：第3句与前面序号句子既不矛盾也不重复。 +判断：<3> <无> +思考：第4句关于用户父亲生日的日期信息与前面序号句子第3句矛盾了。 +判断：<4> <矛盾> +思考：第5句与前面序号句子既不矛盾也不重复。 +判断：<5> <无> +思考：第6句中所有信息都被前面序号中第5句的信息完全包含。 +判断：<2> <被包含> + """ + + contra_repeat_user_query: str = """ +句子： +{user_query} + """ + + long_contra_repeat_system: str = """ + 对下面的{num_obs}句句子，逐一判断是否与“前面序号”的任意句子存在信息的矛盾，或者句子的主要信息被“前面序号”的任意句子中的信息包含。只判断与“前面序号”的句子的关系。 + 请一步步思考，并按如下格式输出: + 思考：思考的依据和过程，30字以内。 + 判断：<句子序号> <矛盾，被包含，无>，一定加<> + """ + + long_contra_repeat_few_shot: str = """ + 示例1 + 句子： + 1 用户经常失眠，对安眠药的效果感兴趣，暗示可能考虑使用。 + 2 用户经常失眠，寻求缓解方法。 + 3 陈伟业是用户的领导 + 4 陈伟业是用户的领导 + 5 陈伟业是用户的领导，是银行分行行长 + + 思考：第1句不会存在与前面序号句子的矛盾或者完全重复。 + 判断：<1> <无> + 思考：第2句中所有信息都被前面序号中第1句的信息完全包含。 + 判断：<2> <被包含> + 思考：第3句信息没有在前面序号句子中出现 + 判断：<3> <无> + 思考：第4句与前面序号中第3句的信息完全重复，即被完全包含。 + 判断：<4> <被包含> + 思考：第5句中陈伟业是用户的领导的信息被前面序号中第3句的信息包含，但新增了陈伟业是银行分行行长的信息，故不是被完全包含。 + 判断：<5> <无> + + 示例2 + 句子： + 1 用户的孩子成绩不太好。 + 2 用户的孩子在学校经常逃课。 + 3 用户的父亲生日在2024年6月2日，用户打算准备礼物。 + 4 用户的父亲生日在2024年5月1日。 + 5 用户很喜欢和同班同学打篮球。 + 6 用户喜欢打篮球。 + + 思考：第1句不会存在与前面序号句子的矛盾或者完全重复。 + 判断：<1> <无> + 思考：第2句与前面序号句子既不矛盾也不重复。 + 判断：<2> <无> + 思考：第3句与前面序号句子既不矛盾也不重复。 + 判断：<3> <无> + 思考：第4句关于用户父亲生日的日期信息与前面序号句子第3句矛盾了。 + 判断：<4> <矛盾> + 思考：第5句与前面序号句子既不矛盾也不重复。 + 判断：<5> <无> + 思考：第6句中所有信息都被前面序号中第5句的信息完全包含。 + 判断：<2> <被包含> + """ + + long_contra_repeat_user_query: str = """ + 句子： + {user_query} + """ + + get_reflect_system: str = """ +任务：从下面的信息中提取出最重要的{num_questions}条用户属性，要求不与已有的用户属性语义重复。 +要求1：用户属性可以是一般的用户偏好，也可以是运动偏好，旅游偏好，饮食偏好等等，也可以是重要事件性质，比如最近重要的事情，也可以是一些高度概括的人生理想，价值观，人生观，性格， 也可以是和朋友的人际关系等等。 +要求2：根据用户属性，我们可以生成“用户的<用户属性>是什么？”的问题，以此可以从下面的信息中提取用户属性对应的值。 +输出格式：每一行输出一个用户属性，每个用户属性推荐4个字，如果没有信息请回答无，最多输出{num_questions}条。 + """ + + get_reflect_few_shot: str = """ +示例1 +信息： +用户想知道明天上海的天气情况。 +用户可能在上海工作，并关心是否需要带伞上班。 +用户在阿里巴巴徐汇滨江园区附近工作。 +用户计划中午在公司附近用餐。 +用户对咖啡因过敏。 +用户喝了咖啡后晚上会出现失眠的情况。 +用户偏好口味较为清淡、不辣的中餐馆。 +用户刚开始了他们的第一份工作。 +用户的工作岗位是阿里巴巴的算法工程师。 +用户希望得到与该岗位相关的职场建议。 +用户面临的问题是在项目进展初期如何有效与上司沟通。 +用户的目标是及时同步项目状态给上司。 +用户希望了解image generation（图像生成）技术的发展概览和最新进展。 +用户对variational auto-encoder、GAN、Diffusion Model等技术及其相互关系感兴趣。 +已有用户属性：性别，工作地点，工作单位，睡眠状况，美食偏好 +新增用户属性： +过敏源 +技术方向 +工作岗位 + + +示例2 +信息： +用户想要了解如何使用torchvision库来可视化深度学习任务的进度信息。 +用户希望了解如何将基于numpy和pytorch的并行计算方案迁移到CUDA支持的GPU上运行。 +用户询问是否需要依赖特定的包来完成这一任务。 +用户希望了解如何在Python中自定义进程和线程以实现并行计算。 +用户在编程中遇到了与并行计算相关的问题。 +用户希望学习如何使用Python（numpy，pytorch）在GPU上实现简单的并行计算。 +用户希望了解并行计算的基本概念，包括threads。 +用户询问有关世界各地著名菜系的信息。 +用户对全球各地的美食非常感兴趣。 +用户关心其体重与运动消耗的额外热量及心率之间的关系。 +用户在询问为了实现这一目标，每天需要额外消耗多少大卡热量。 +用户希望每月减重1kg。 +用户希望得到类似战略类手机游戏的推荐。 +用户喜欢玩三国志系列、文明系列、全面战争、骑马与砍杀等战略类游戏。 +用户希望根据他们的喜好获得新的游戏推荐。 +用户列举了他们喜欢的具体游戏类型，包括：三国志系列、文明系列、全面战争、骑马与砍杀等。 +用户喜欢玩战略类游戏。 +已有用户属性：工作地点，性别，美食偏好 +新增用户属性： +游戏偏好 +运动计划 +技术方向 + +示例3 +信息： +用户寻求推荐一个相关课程或网址以进行学习。 +用户计划去青岛旅游。 +用户正为张三的女儿选购生日礼物。 +用户请求为一位名叫张三的人的女儿撰写一段温馨的祝福语。 +用户的同事名叫张三。 +用户与张三约定讨论阿里云百炼项目。 +用户与同事张三讨论了该项目的PRD（产品需求文档）。 +同事张三计划下周对PRD进行最终确定。 +张三还安排了在再下一周进行POC（Proof of Concept，概念验证）的讨论。 +用户希望获知该项目工程开发工作的负责团队信息，以了解项目执行的组织架构与分工情况。 +已有用户属性： +新增用户属性： +朋友关系 + +示例4 +信息： +用户在寻求有关推拿按摩手法的教程或相关网站推荐。 +用户希望系统地学习正规的推拿按摩手法。 +用户对按摩感兴趣，并且经常去推拿按摩店。 +用户想了解自己在静息状态下一小时大概会消耗多少大卡热量。 +用户年龄为28岁。 +用户体重为70kg。 +用户是男性。 +已有用户属性：性别，年龄，体重，当前学习进展 +新增用户属性： +无 + """ + + get_reflect_user_query: str = """ +信息： +{user_query} +已有用户属性：{exist_keys} +新增用户属性： + """ + + get_insight_system: str = """ +任务：从下面的信息中提取出关于用户属性信息。语言简洁，每条不超过50字。请在一句话内表达，语言简洁，每条不超过50字。 + """ + + get_insight_few_shot: str = """ +示例1 +信息： +用户考虑是否应该给猫咪Sally购买一些猫玩具。 +用户家中有一只名为Sally的宠物，需要在室内自由活动以保障其身心健康。 +用户想要为名为Sally的宠物购买猫粮。 +用户养有一只名叫Sally的猫，并很喜欢它。 +孙二为这只猫取名为“Sally”。 +用户希望得到关于购买何种猫粮的建议。 +用户希望了解不同品种猫咪的基本信息。 +用户对猫毛过敏。 +用户关心领养猫与在宠物店购买猫之间的区别。 +用户正在寻求治疗猫毛过敏的方法。 +问题：用户的过敏源是什么？ +用户对猫毛过敏，正寻求有效的解决方案，以便更好地照顾其宠物猫Sally并减轻过敏症状。 + +示例2 +信息： +用户的工作岗位是阿里巴巴的算法工程师。 +用户希望详细了解和学习图像生成技术。 +用户面临的问题是在项目进展初期如何有效与上司沟通。 +用户希望了解image generation（图像生成）技术的发展概览和最新进展。 +用户希望得到与该岗位相关的职场建议。 +用户刚开始了他们的第一份工作。 +用户的身份可能是初学者或专业人士。 +用户希望学习如何使用Python（numpy，pytorch）在GPU上实现简单的并行计算。 +用户想要了解如何使用torchvision库来可视化深度学习任务的进度信息。 +用户对variational auto-encoder、GAN、Diffusion Model等技术及其相互关系感兴趣。 +问题：用户的技术方向是什么？ +用户是初入职场的阿里巴巴算法工程师，正积极探索图像生成技术和并行计算的知识，并寻求技术学习方面的指导，以提升自己的专业技能。 + """ + + get_insight_user_query: str = """ +信息： +{user_query} +问题：用户的{insight_key}是什么？ + """ + + update_plural_profile_system: str = """ +从下面的句子中提取出给定类别的用户资料信息，并判断和已有信息是否重复。只输出无重复的新信息。若无法提取该类别的用户资料的新信息则回答无。 +请一步步思考，并按如下格式输出: +思考: 思考的依据和过程，150字以内。 +用户资料: <信息>或<无>, 一定加<> + """ + + update_plural_profile_few_shot: str = """ +示例1: +句子：用户上周去了西溪游泳馆游泳，那个游泳馆人非常多。 +句子：用户计划每周六和朋友张三去朝阳体育馆打羽毛球。 +类别：运动（用户喜欢的运动） +已有信息：运动（用户喜欢的运动）：游泳 +思考：从第一句句子可以得出游泳是用户喜欢的运动之一，但与已有信息重复。从第二句句子可以得出羽毛球是用户喜欢的运动之一，是新的信息。 +用户资料: <羽毛球> + +示例2: +句子：用户对咖啡因过敏。 +句子：用户不喜欢吃香菇。 +类别：过敏（用户的已知过敏反应） +已有信息：过敏（用户的已知过敏反应）: 咖啡因 +思考：从第一句句子可以得出咖啡因是用户的已知过敏反应之一，但与已有信息重复。从第二句句子只能得出用户不喜欢香菇而非对香菇过敏，无法得出新的用户已知过敏信息。 +用户资料: <无> + +示例3: +句子：用户热衷于动作类类游戏如只狼、艾尔登法环。 +句子：用户在休闲时间经常长时间玩策略类游戏如文明6。 +句子：用户是音乐发烧友，关注各个品牌的耳机的音质和性价比。 +类别：爱好（用户的业余爱好） +已有信息：爱好（用户的业余爱好）: +思考：从第一句句子可以得出动作类游戏是用户的爱好之一，是新的信息。从第二句句子可以得出策略类游戏是用户的爱好之一，是新的信息。从第三句句子可以得出音乐是用户的爱好之一，是新的信息。 +用户资料: <动作类游戏, 策略类游戏, 音乐> + +示例4: +句子：关于职场沟通你有什么具体的建议吗？最好结合一个实例。我一直听人说要加强沟通，经常和上司沟通，同步项目的进展，但是我总是感觉还有许多事情要做。 +句子：项目并没有达到一个充分的可以汇报的状态，然后准备汇报材料又很费时间，导致有时候我没有及时和上司同步项目状态。针对这个情况你有什么建议？ +类别：职业（用户的职业） +已有信息：职业（用户的职业）：工程师 +思考：句子中虽然提及了职场沟通等工作相关内容，但是并不能推断出用户的职位是什么，只能推知与宽泛的项目实施与管理相关。 +用户资料: <无> + """ + + update_plural_profile_user_query: str = """ +{user_query} +类别：{update_profile} +已有信息：{update_profile_value} + """ + + update_unique_profile_system: str = """ +从下面的句子中提取出给定类别的用户资料信息，并判断与已有信息是否矛盾。若矛盾则输出更新的信息，若不矛盾则保留已有信息，整合已有信息和新信息并输出。 +请一步步思考，并按如下格式输出: +思考: 思考的依据和过程，150字以内。 +用户资料: <信息>, 一定加<> + """ + + update_unique_profile_few_shot: str = """ +示例1: +句子：因为昨天成都下大雨，用户全身都被淋湿了。 +句子：用户关心明天成都的天气预报。 +类别：地区（用户所在地区） +已有信息：地区（用户所在地区）: 杭州 +思考：从第一句句子可以得出用户在成都。第二句句子没有直接透露用户所在地信息，但与第一句句子用户在成都的信息吻合。这与已有信息（用户在杭州）矛盾，输出更新的信息。 +用户资料：<成都> + +示例2: +句子：用户女朋友下个月过生日。 +句子：用户生日在7月15日。 +类别：生日（用户的生日）。 +已有信息：生日（用户的生日）：1987年7月15日。 +思考：第一句句子中提及生日，但并不是用户的生日，无法得出用户生日信息。从第二句句子可以得出用户生日在7月15日，与已有信息不矛盾，整合可以得出用户生日是1987年7月15日。 +用户资料: <1987年7月15日> + +示例3: +句子：用户在招商银行工作。 +句子：用户刚刚毕业，第一份工作是银行前台。 +句子：用户的理想工作是职业游戏选手。 +类别：职业（用户的职业） +已有信息：职业（用户的职业）： +思考：整合第一和第二句句子的信息可以得出用户的现在的职业是招商银行前台。第三句句子说明了用户的理想工作但并不是现在的职业。 +用户资料：<招商银行前台> + +示例4: +句子：用户大学期间接触过优化算法的研究。 +类别：学习专业 (用户大学学习的专业) +已有信息：学习专业 (用户大学学习的专业)：与人工智能相关 +思考：从句子可以得出用户大学学习的专业与优化算法相关，这与已有信息（用户大学学习的专业与人工智能相关）不矛盾，整合可以得出用户大学学习的专业与人工智能和优化算法相关。 +用户资料：<与人工智能和优化算法相关> + +示例5: +句子：今天和同学去打球了。 +句子：明天和女朋友一起去杭州旅游。 +类别：学习专业 (用户大学学习的专业) +已有信息：学习专业 (用户大学学习的专业)： +思考：两个句子和学习专业都没有关联，没有新提取的信息。 +用户资料：<无> + """ + + update_unique_profile_user_query: str = """ +{user_query} +类别：{update_profile} +已有信息：{update_profile_value} + """ + + update_insight_system: str = """ +从下面的句子中提取出给定类别的用户资料信息，并判断与已有信息是否矛盾，若矛盾以新信息为准。整合已有信息和新信息并输出。 +请一步步思考，并按如下格式输出: +思考: 思考的依据和过程，150字以内。 +用户资料: <信息>, 一定加<> + """ + + update_insight_few_shot: str = """ +示例1: +句子：因为昨天成都下大雨，用户全身都被淋湿了。 +句子：用户关心明天成都的天气预报。 +类别：用户所在地区 +已有信息：用户所在地区: 杭州 +思考：从第一句句子可以得出用户在成都。第二句句子没有直接透露用户所在地信息，但与第一句句子用户在成都的信息吻合。这与已有信息（用户在杭州）矛盾，输出更新的信息。 +用户资料：<成都> + +示例2: +句子：用户最近养好了肠胃。 +句子：用户关注中医养生。 +类别：用户健康状况 +已有信息：用户健康状况: 肠胃不好，高血压 +思考：从第一句句子可以得出用户最近养好了肠胃，与已有信息矛盾，以新信息为准。第二句句子与用户健康状况无关。整合已有信息和新信息得到用户健康状况是肠胃健康，高血压。 +用户资料：<肠胃健康，高血压> + +示例3: +句子：用户刚刚毕业，第一份工作是银行前台。 +句子：用户的理想工作是职业游戏选手。 +类别：用户职业 +已有信息：用户职业：在招商银行工作 +思考：整合已有信息和第一句句子的信息可以得出用户的现在的职业是招商银行前台。第二句句子说明了用户的理想工作但并不是现在的职业。 +用户资料：<招商银行前台> + +示例4: +句子：用户大学期间接触过优化算法的研究。 +类别：用户学习专业 +已有信息：用户学习专业：与人工智能相关 +思考：从句子可以得出用户大学学习的专业与优化算法相关，这与已有信息（用户学习专业与人工智能相关）不矛盾，整合可以得出用户大学学习的专业与人工智能和优化算法相关。 +用户资料：<与人工智能和优化算法相关> + +示例5: +句子：用户单身。 +句子：用户受到一名18岁男生的追求，但不想接受又不想伤害他。 +句子：用户喜欢成熟且情绪稳定的男生。 +类别：用户情感状况 +已有信息：用户情感状况：有男朋友 +思考：从第一句句子可以得出用户现在单身，与已有信息矛盾，以新信息为准。从第二句句子得出用户受到一名18岁男生的追求但并不喜欢他。第三句话表达了用户理想的伴侣类型但与用户 +情感状况无关。整合得出用户情感状况为单身，受到一名18岁男生的追求但并不喜欢他。 +用户资料：<单身，受到一名18岁男生的追求但并不喜欢他。> + """ + + update_insight_user_query: str = """ +{user_query} +类别：{insight_key} +已有信息：{insight_key_value} + """ diff --git a/memory_scope/utils/__init__.py b/memory_scope/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/memory_scope/utils/context_handler.py b/memory_scope/utils/context_handler.py new file mode 100644 index 00000000..a55afa00 --- /dev/null +++ b/memory_scope/utils/context_handler.py @@ -0,0 +1,129 @@ +import os +import threading +from typing import Dict, Any + +from common.logger import Logger +from constants.common_constants import APP_ENV +from enumeration.env_type import EnvType + + +class ContextHandler(object): + # 环境类型：日常 预发 开发 + _env_type: EnvType | None = None + + # 所有环境变量 + _env_params = os.environ + + # 环境变量解析 -> config + _env_configs: Dict[str, str] | None = None + + def __init__(self, scene: str, method: str, prefix: str = "memory", algo_version: str = ""): + self.scene: str = scene + self.method: str = method + self.prefix: str = prefix + self.algo_version: str = algo_version + + # 上下文 所有worker共享 + self.context_dict: Dict[str, Any] = {} + + # 日志 + self.logger = Logger.get_logger() + + # 全局锁 + self.context_lock = threading.Lock() + + # algo_version: 上游参数 > 环境变量参数 + self._update_algo_version() + + def _update_algo_version(self): + # 上游传参优先 + if self.algo_version: + return + + p_key = f"{self.prefix}_{self.method}_algo_version" + p_key_with_scene = f"{p_key}_{self.scene}" + + if p_key_with_scene in self._env_params: + self.algo_version = self._env_params.get(p_key_with_scene) + + if p_key in self._env_params: + self.algo_version = self._env_params.get(p_key_with_scene) + + @property + def env_type(self): + if self._env_type is None: + env_type = self._env_params.get(APP_ENV) + assert env_type, f"env_type={env_type} is empty!" + self._env_type = EnvType(env_type.lower()) + + return self._env_type + + @property + def env_configs(self): + if self._env_configs is not None: + return self._env_configs + + prefix: str = f"{self.prefix}_{self.method}_" + all_ket_set = set() + for k, v in self._env_params.items(): + if not k.startswith(prefix): + continue + + if not v: + continue + + # {key}_{scene}_{algo_version} + raw_k = k.removeprefix(prefix) + if self.scene in raw_k: + raw_k_split = [x.strip("_") for x in raw_k.split(self.scene) if x.strip("_")] + if len(raw_k_split) == 1: + raw_k = raw_k_split[0] + elif len(raw_k_split) == 2: + algo_version = raw_k_split[1] + if algo_version != self.algo_version: + continue + + raw_k = raw_k_split[0] + else: + self.logger.info(f"_update_env_configs encounter error! k={k} v={v}") + continue + + all_ket_set.add(raw_k) + + self._env_configs = {} + for k in all_ket_set: + v = self.get_param(k) + if v: + self._env_configs[k] = v + self.logger.info(f"update env_configs={self._env_configs}") + return self._env_configs + + def get_env_config(self, key: str, default=None) -> str: + return self.env_configs.get(key, default) + + def get_param(self, key: str, default=None): + p_key = f"{self.prefix}_{self.method}_{key}" + p_key_with_scene = f"{p_key}_{self.scene}" + # memory_summary_{key}_tongyi_algo_v1 > memory_summary_{key}_tongyi > memory_summary_{key} + + if p_key_with_scene in self._env_params: + p_key = p_key_with_scene + + if self.algo_version: + p_key_with_version = f"{p_key_with_scene}_{self.algo_version}" + if p_key_with_version in self._env_params: + p_key = p_key_with_version + + return self._env_params.get(p_key, default) + + 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/logger.py b/memory_scope/utils/logger.py new file mode 100644 index 00000000..684ba6b7 --- /dev/null +++ b/memory_scope/utils/logger.py @@ -0,0 +1,102 @@ +import logging +from logging.handlers import RotatingFileHandler +from pathlib import Path + +from constants.common_constants import MEMORY + +# remove %(thread)s .%(funcName)s +LOG_FORMAT = "%(asctime)s %(levelname)s %(trace_id)s %(module)s:%(lineno)d] %(message)s" +DATE_FORMAT = "%Y-%m-%d %H:%M:%S" + +LOGGER_DICT = {} + + +class Logger(logging.Logger): + def __init__(self, + name: str, + level: int = logging.INFO, + format_style: str = LOG_FORMAT, + date_format_style: str = DATE_FORMAT, + to_stream: bool = True, + to_file: bool = True, + file_mode: str = "w", + file_type: str = "log", + dir_path: str = "log", + max_bytes: int = 1024 * 1024 * 1024, + backup_count: int = 10): + super(Logger, self).__init__(name, level) + + self.formatter = logging.Formatter(format_style, date_format_style) + self.date_format_style = date_format_style + self.to_stream: bool = to_stream + self.to_file: bool = to_file + self.file_mode: str = file_mode + self.file_type: str = file_type + self.dir_path: str = dir_path + + self.max_bytes: int = max_bytes + self.backup_count: int = backup_count + + self.trace_id: str = "" + + if self.to_stream: + self._add_stream_handler() + if self.to_file: + self._add_file_handler() + + self.info(f"logger={name} is inited.") + + def _add_file_handler(self): + file_path = Path().joinpath(self.dir_path, f"{self.name}.{self.file_type}") + file_path.parent.mkdir(exist_ok=True) + file_name = file_path.as_posix() + + file_handler = RotatingFileHandler( + filename=file_name, + maxBytes=self.max_bytes, + backupCount=self.backup_count, + encoding="utf-8") + file_handler.setFormatter(self.formatter) + self.addHandler(file_handler) + + def _add_stream_handler(self): + stream_handler = logging.StreamHandler() + stream_handler.setFormatter(self.formatter) + stream_handler.encoding = 'utf-8' + self.addHandler(stream_handler) + + def close(self): + for handler in self.handlers: + handler.close() + + def clear(self): + self.handlers.clear() + + def set_trace_id(self, trace_id: str): + self.trace_id: str = trace_id + if len(self.trace_id) >= 8: + self.trace_id = self.trace_id[:8] + + def makeRecord(self, name, level, fn, lno, msg, args, exc_info, + func=None, extra=None, sinfo=None): + if extra is None: + extra = {} + extra["trace_id"] = self.trace_id + return super().makeRecord(name, level, fn, lno, msg, args, exc_info, func, extra, sinfo) + + @classmethod + def get_logger(cls, name: str = None, **kwargs): + if name is None: + if LOGGER_DICT: + name = list(LOGGER_DICT.keys())[0] + else: + name = MEMORY + + if name not in LOGGER_DICT: + LOGGER_DICT[name] = Logger(name=name, **kwargs) + + return LOGGER_DICT[name] + + @classmethod + def get_memory_logger(cls, **kwargs): + return cls.get_logger(MEMORY, **kwargs) diff --git a/memory_scope/utils/timer.py b/memory_scope/utils/timer.py new file mode 100644 index 00000000..4098d120 --- /dev/null +++ b/memory_scope/utils/timer.py @@ -0,0 +1,76 @@ +""" +file: timer.py +author: yuli +date: 20221106 +""" + +import time + +from common.logger import Logger + + +class Timer(object): + + def __init__(self, name: str, log_time: bool = True, use_ms: bool = True, **kwargs): + self.name: str = name + self.log_time: bool = log_time + self.use_ms: bool = use_ms + self.kwargs: dict = kwargs + + self.logger = Logger.get_logger() + + # time record + self.t_start = 0 + self.t_end = 0 + self.cost = 0 + + @classmethod + def kwargs_to_str(cls, float_precision: int = 4, **kwargs): + line_list = [] + for k, v in kwargs.items(): + if isinstance(v, float): + float_style = f".{float_precision}f" + line = f"{k}={v:{float_style}}" + else: + line = f"{k}={v}" + line_list.append(line) + + return " ".join(line_list) + + def __enter__(self): + self.t_start = time.time() + # with Timer("XXX") as t, need return self + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.t_end = time.time() + self.cost = self.t_end - self.t_start + if self.use_ms: + self.cost *= 1000 + + if self.log_time: + line = f"{self.name}.timer" + + if self.use_ms: + line = f"{line} cost={self.cost:.1f}ms" + else: + line = f"{line} cost={self.cost:.4f}s" + + if self.kwargs: + line = f"{line} {self.kwargs_to_str(**self.kwargs)}" + + self.logger.info(line, stacklevel=3) + + def get_cost_info(self): + if self.use_ms: + return f"cost={self.cost:.1f}ms" + else: + return f"cost={self.cost:.4f}s" + + +def timer(func): + def wrapper(*args, **kwargs): + with Timer(name=func.__name__, **kwargs): + return func(*args, **kwargs) + + return wrapper diff --git a/memory_scope/utils/tool_functions.py b/memory_scope/utils/tool_functions.py new file mode 100644 index 00000000..a72f804b --- /dev/null +++ b/memory_scope/utils/tool_functions.py @@ -0,0 +1,164 @@ +import os +import re +from datetime import datetime +from typing import Dict, List + +from constants.common_constants import WEEKDAYS +from enumeration.env_type import EnvType + +global_env_type = None + + +def get_global_env_type(): + global global_env_type + + if global_env_type is None: + env = os.environ.get("APP_ENV", "") + if env is None or not env: + raise EnvironmentError("Environment variable APP_ENV must be set") + env = env.split("-")[-1] + + if env not in EnvType.__members__.values(): + global_env_type = EnvType.DAILY + else: + global_env_type = EnvType(env) + + return global_env_type + + +def under_line_to_hump(underline_str): + sub = re.sub(r'(_\w)', lambda x: x.group(1)[1].upper(), underline_str) + return sub[0:1].upper() + sub[1:] + + +def parse_response_text_v1(response_text: str) -> dict: + """ + parse text like: + <1> + <2> ddd + <4> dddd<555> + + result = {1: "AAA", 2: "BBB", 4: "CCC"} + """ + result_dict: Dict[int, str] = {} + + # 确保第一个数字，后面是string + matches = re.findall(r'<(\d+)>\s*<([^>]+)>', response_text.strip()) + + # matches 为空返回 + for key, value in matches: + result_dict[int(key)] = value + + return result_dict + + +def parse_response_text_v2(response_text: str) -> Dict[int, List[str]]: + """ + parse text like: + XXX + <1> <222> + <2> + <4,5> + + result = {1: ["AAA", "222"], 2: "BBB", 4: "CCC"} + """ + result_dict: Dict[int, List[str]] = {} + for line in response_text.strip().split("\n"): + if "> <" not in line: + continue + ll = [x.removeprefix("<").removesuffix(">") for x in line.strip().split("> <")] + idx: str = ll[0] + values: List[str] = ll[1:] + if idx.isdigit(): + idx_int = int(idx) + else: + idx_split = idx.split(",") + if len(idx_split) == 0: + continue + idx = idx_split[0] + if idx.isdigit(): + idx_int = int(idx) + else: + continue + if values: + result_dict[idx_int] = values + + return result_dict + + +def parse_response_text_v3(response_text: str) -> List[List[str]]: + """ + parse text like: + XXX + <1> + <2c> + <41> + + result = [["1", "AAA"], ["2c", "BBB"], ["41", "CCC", "BBB"]] + """ + result_list: List[List[str]] = [] + for line in response_text.strip().split("\n"): + if "> <" not in line: + continue + ll = [x.removeprefix("<").removesuffix(">") for x in line.strip().split("> <")] + result_list.append(ll) + return result_list + + +def get_datetime_info_dict(parse_dt: datetime): + return { + "year": parse_dt.year, + "month": parse_dt.month, + "day": parse_dt.day, + "hour": parse_dt.hour, + "minute": parse_dt.minute, + "second": parse_dt.second, + "week": parse_dt.isocalendar().week, + "weekday": WEEKDAYS[parse_dt.isocalendar().weekday - 1], + } + + +def extract_date_parts(input_string: str): + # Extending our pattern to handle "每" (every) as a possible value. + patterns = { + 'year': r'(\d+|每)年', + 'month': r'(\d+|每)月', + 'day': r'(\d+|每)日', + 'weekday': r'周([一二三四五六日])?', + 'hour': r'(\d+)点' + } + weekday_dict = {"一": 1, "二": 2, "三": 3, "四": 4, "五": 5, "六": 6, "日": 7} + extracted_data = {} + + # Search for patterns in the input string and populate the dictionary + for key, pattern in patterns.items(): + match = re.search(pattern, input_string) + if match: # If there is a match, include it in the output dictionary + if match.group(1) == "每": + extracted_data[key] = -1 + elif match.group(1) in weekday_dict.keys(): + extracted_data[key] = weekday_dict[match.group(1)] + else: + extracted_data[key] = int(match.group(1)) + return extracted_data + + +def time_to_formatted_str(time: datetime | str | int | float = None, + date_format: str = "%Y%m%d", # e.g. %Y%m%d -> "20240528", add %H:%M:%S + string_format: str = "") -> str: + if isinstance(time, str | int | float): + if isinstance(time, str): + time = float(time) + current_dt = datetime.fromtimestamp(time) + elif isinstance(time, datetime): + current_dt = time + else: + current_dt = datetime.now() + + return_str = "" + if date_format: + return_str = current_dt.strftime(date_format) + elif string_format: + return_str = string_format.format(**get_datetime_info_dict(current_dt)) + + return return_str diff --git a/memory_scope/utils/user_profile_handler.py b/memory_scope/utils/user_profile_handler.py new file mode 100644 index 00000000..9a390f95 --- /dev/null +++ b/memory_scope/utils/user_profile_handler.py @@ -0,0 +1,102 @@ +import json +from typing import List, Dict + +from enumeration.memory_node_status import MemoryNodeStatus +from model.memory_wrap_node import MemoryWrapNode +from model.user_attribute import UserAttribute + + +class UserProfileHandler(object): + @classmethod + def format_content(cls, key: str, description: str, value: str | List[str] = None): + if not key.startswith("用户"): + key = f"用户的{key}" + + if not description.startswith("用户"): + description = f"用户{description}" + + content = f"{key}（{description}）" + + if value: + if isinstance(value, list): + value = "，".join(value) + content = f"{content}：{value}" + + return content + + """ + 提供UserAttribute 和 MemoryWrapNode 的相互转化 + """ + + @classmethod + def to_nodes(cls, + user_profile: List[UserAttribute] | Dict[str, UserAttribute] | None = None, + split_value: bool = False) -> List[MemoryWrapNode]: + + user_profile_dict: Dict[str, UserAttribute] = {} + if user_profile: + if isinstance(user_profile, list): + for user_attr in user_profile: + user_profile_dict[user_attr.memory_key] = user_attr + elif isinstance(user_profile, dict): + user_profile_dict.update(user_profile) + + user_profile_nodes: List[MemoryWrapNode] = [] + for _, user_attr in user_profile_dict.items(): + # 获取id + _id = user_attr.code + if not _id: + _id = f"{user_attr.memory_id}_{user_attr.scene}_profile_{user_attr.memory_key}" + + attr_node = MemoryWrapNode.init_from_attrs(id=_id, + code=_id, + content="", + memoryId=user_attr.memory_id, + scene=user_attr.scene, + memoryType=user_attr.memory_type, + content_modified=True, + metaData={ + "memory_key": user_attr.memory_key, + "value": json.dumps(user_attr.value, ensure_ascii=False), + "is_unique": str(user_attr.is_unique), + "is_mutable": str(user_attr.is_mutable), + "description": user_attr.description, + "status": MemoryNodeStatus.ACTIVE.value, + "ext_info": json.dumps(user_attr.ext_info, + ensure_ascii=False), + }, + status=MemoryNodeStatus.ACTIVE.value) + + if split_value: + for value in user_attr.value: + content = cls.format_content(user_attr.memory_key, user_attr.description, value) + attr_node_copy = attr_node.copy(deep=True) + attr_node_copy.memory_node.content = content + user_profile_nodes.append(attr_node_copy) + else: + content = cls.format_content(user_attr.memory_key, user_attr.description, user_attr.value) + attr_node.memory_node.content = content + user_profile_nodes.append(attr_node) + + return user_profile_nodes + + @classmethod + def to_user_attr(cls, user_profile_nodes: List[MemoryWrapNode]) -> Dict[str, UserAttribute]: + user_profile_dict: Dict[str, UserAttribute] = {} + + for node in user_profile_nodes: + user_attr = UserAttribute( + code=node.id, + memory_id=node.memory_node.memoryId, + scene=node.memory_node.scene, + memory_key=node.memory_node.metaData["memory_key"], + value=json.loads(node.memory_node.metaData["value"]), + is_unique=int(node.memory_node.metaData["is_unique"]), + is_mutable=int(node.memory_node.metaData["is_mutable"]), + memory_type=node.memory_node.memoryType, + description=node.memory_node.metaData["description"], + status=1 if node.memory_node.metaData["status"] == MemoryNodeStatus.ACTIVE.value else 0, + ext_info=json.loads(node.memory_node.metaData["ext_info"]), + ) + user_profile_dict[user_attr.memory_key] = user_attr + return user_profile_dict diff --git a/memory_scope/worker/__init__.py b/memory_scope/worker/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/memory_scope/worker/base_worker.py b/memory_scope/worker/base_worker.py new file mode 100644 index 00000000..81fdbc51 --- /dev/null +++ b/memory_scope/worker/base_worker.py @@ -0,0 +1,126 @@ +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import Any, List + +from common.context_handler import ContextHandler +from common.logger import Logger +from common.timer import Timer + + +class BaseWorker(object): + + def __init__(self, + context_handler: ContextHandler, + is_multi_thread: bool = False, + thread_pool: ThreadPoolExecutor = None, + raise_exception: bool = True, + logger: Logger = None, + **kwargs): + super(BaseWorker, self).__init__(**kwargs) + + # 原始参数 + self.context_handler: ContextHandler = context_handler + self.is_multi_thread: bool = is_multi_thread + self.thread_pool: ThreadPoolExecutor = thread_pool + 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 + + # 运行信息，保存到ext_info + self.run_infos: List[str] = [] + + # 运行时间 + self.run_cost: float = 0 + + # 短name + self._name_simple: str = "" + + 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 + + def run(self): + self.logger.info(f"----- Begin {self.name_simple} -----") + with Timer(self.name_simple, log_time=False) as t: + if self.raise_exception: + self._run() + else: + try: + self._run() + except Exception as e: + self.add_run_info(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()}-----") + + def get_context(self, key: str, default=None): + return self.context_handler.get_context(key, default) + + def set_context(self, key: str, value: Any): + self.context_handler.set_context(key, value, self.is_multi_thread) + + def get_param(self, key: str, default=None): + return self.context_handler.env_configs.get(key, default) + + @property + def env_type(self) -> str: + return self.context_handler.env_type.value + + @property + def scene(self) -> str: + return self.context_handler.scene + + @property + def method(self) -> str: + return self.context_handler.method + + @property + def algo_version(self) -> str: + return self.context_handler.algo_version + + @property + def name_simple(self) -> str: + if not self._name_simple: + 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, + } diff --git a/memory_scope/worker/es/__init__.py b/memory_scope/worker/es/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/memory_scope/worker/es/es_insight_worker.py b/memory_scope/worker/es/es_insight_worker.py new file mode 100644 index 00000000..8239f6a7 --- /dev/null +++ b/memory_scope/worker/es/es_insight_worker.py @@ -0,0 +1,23 @@ +from typing import List + +from constants.common_constants import INSIGHT_NODES +from enumeration.memory_node_status import MemoryNodeStatus +from enumeration.memory_type_enum import MemoryTypeEnum +from model.memory_wrap_node import MemoryWrapNode +from worker.bailian.memory_base_worker import MemoryBaseWorker + + +class EsInsightWorker(MemoryBaseWorker): + + def _run(self): + hits = self.es_client.exact_search_v2(size=self.config.es_insight_top_k, + term_filters={ + "memoryId": self.config.memory_id, + "status": MemoryNodeStatus.ACTIVE.value, + "scene": self.scene.lower(), + "memoryType": MemoryTypeEnum.INSIGHT.value, + }) + + insight_nodes: List[MemoryWrapNode] = [MemoryWrapNode.init_from_es(hit) for hit in hits] + self.logger.info(f"insight_nodes.size={len(insight_nodes)}") + self.set_context(INSIGHT_NODES, insight_nodes) diff --git a/memory_scope/worker/es/es_keyword_worker.py b/memory_scope/worker/es/es_keyword_worker.py new file mode 100644 index 00000000..f251b630 --- /dev/null +++ b/memory_scope/worker/es/es_keyword_worker.py @@ -0,0 +1,53 @@ +from typing import List + +from constants.common_constants import KEY_WORD, KEYWORD_OBS_NODES, RECALL_TYPE, QUERY_KEYWORDS +from enumeration.memory_node_status import MemoryNodeStatus +from enumeration.memory_recall_type import MemoryRecallType +from enumeration.memory_type_enum import MemoryTypeEnum +from model.memory_wrap_node import MemoryWrapNode +from worker.bailian.memory_base_worker import MemoryBaseWorker + + +class EsKeywordWorker(MemoryBaseWorker): + + def _run(self): + query = self.messages[-1].content + # keywords = jieba.analyse.extract_tags(query, topK=3, withWeight=False, allowPOS=()) # 分解关键词 + # keywords = jieba.cut(query, cut_all=False) # 使用精确模式分词 + + # 查询相关关键词 + keywords = set() + query_keywords = set() + for key, values in self.config.key_word_relate_dict.items(): + if key in query: + keywords.add(key) + keywords.update(values) + query_keywords.add(values[0]) + keywords = list(keywords) + + self.set_context(QUERY_KEYWORDS, query_keywords) + + # 任意一个匹配都算 + hits = self.es_client.exact_search_v2(size=self.config.es_keyword_top_k, + term_filters={ + "memoryId": self.config.memory_id, + "status": MemoryNodeStatus.ACTIVE.value, + "scene": self.scene.lower(), + "memoryType": [MemoryTypeEnum.OBSERVATION.value, + MemoryTypeEnum.INSIGHT.value, + MemoryTypeEnum.OBS_CUSTOMIZED.value], + }, + match_filters={ + f"metaData.{KEY_WORD}": keywords, + }) + + # 初始化成MemoryWrapNode，并加入召回源的参数 + keyword_obs_nodes: List[MemoryWrapNode] = [] + for hit in hits: + node = MemoryWrapNode.init_from_es(hit) + node.memory_node.metaData[RECALL_TYPE] = MemoryRecallType.KEYWORD.value + keyword_obs_nodes.append(node) + self.logger.info(f"keyword_obs_nodes size={len(keyword_obs_nodes)}") + for node in keyword_obs_nodes: + self.logger.info(f"node={node.memory_node.content} score_similar={node.score_similar}") + self.set_context(KEYWORD_OBS_NODES, keyword_obs_nodes) diff --git a/memory_scope/worker/es/es_new_obs_worker.py b/memory_scope/worker/es/es_new_obs_worker.py new file mode 100644 index 00000000..fbebaddf --- /dev/null +++ b/memory_scope/worker/es/es_new_obs_worker.py @@ -0,0 +1,24 @@ +from typing import List + +from constants.common_constants import NEW, NEW_OBS_NODES +from enumeration.memory_node_status import MemoryNodeStatus +from enumeration.memory_type_enum import MemoryTypeEnum +from model.memory_wrap_node import MemoryWrapNode +from worker.bailian.memory_base_worker import MemoryBaseWorker + + +class EsNewObsWorker(MemoryBaseWorker): + + def _run(self): + hits = self.es_client.exact_search_v2(size=self.config.es_new_obs_top_k, + term_filters={ + "memoryId": self.config.memory_id, + "status": MemoryNodeStatus.ACTIVE.value, + "scene": self.scene.lower(), + "memoryType": MemoryTypeEnum.OBSERVATION.value, + f"metaData.{NEW}": "1", + }) + + new_obs_nodes: List[MemoryWrapNode] = [MemoryWrapNode.init_from_es(hit) for hit in hits] + self.logger.info(f"es new obs, size={len(new_obs_nodes)}") + self.set_context(NEW_OBS_NODES, new_obs_nodes) diff --git a/memory_scope/worker/es/es_not_reflected_worker.py b/memory_scope/worker/es/es_not_reflected_worker.py new file mode 100644 index 00000000..c031c428 --- /dev/null +++ b/memory_scope/worker/es/es_not_reflected_worker.py @@ -0,0 +1,24 @@ +from typing import List + +from constants.common_constants import REFLECTED, NOT_REFLECTED_OBS_NODES +from enumeration.memory_node_status import MemoryNodeStatus +from enumeration.memory_type_enum import MemoryTypeEnum +from model.memory_wrap_node import MemoryWrapNode +from worker.bailian.memory_base_worker import MemoryBaseWorker + + +class EsNotReflectedWorker(MemoryBaseWorker): + def _run(self): + hits = self.es_client.exact_search_v2(size=self.config.es_not_reflected_top_k, + term_filters={ + "memoryId": self.config.memory_id, + "status": MemoryNodeStatus.ACTIVE.value, + "scene": self.scene.lower(), + "memoryType": [MemoryTypeEnum.OBSERVATION.value, + MemoryTypeEnum.OBS_CUSTOMIZED.value], + f"metaData.{REFLECTED}": "0", + }) + + not_reflected_obs_nodes: List[MemoryWrapNode] = [MemoryWrapNode.init_from_es(hit) for hit in hits] + self.logger.info(f"retrieve_not_reflected_obs.size={len(not_reflected_obs_nodes)}") + self.set_context(NOT_REFLECTED_OBS_NODES, not_reflected_obs_nodes) diff --git a/memory_scope/worker/es/es_retrieve_all_worker.py b/memory_scope/worker/es/es_retrieve_all_worker.py new file mode 100644 index 00000000..e0ba99a8 --- /dev/null +++ b/memory_scope/worker/es/es_retrieve_all_worker.py @@ -0,0 +1,30 @@ +from typing import List + +from constants.common_constants import ALL_NODES, ALL_MEMORIES +from enumeration.memory_node_status import MemoryNodeStatus +from model.memory_wrap_node import MemoryWrapNode +from worker.bailian.memory_base_worker import MemoryBaseWorker + + +class EsRetrieveAllWorker(MemoryBaseWorker): + + def _run(self): + # msg_time_created = self.messages[-1].time_created + hits = self.es_client.exact_search_v2(size=1000, + term_filters={ + "memoryId": self.config.memory_id, + "status": MemoryNodeStatus.ACTIVE.value, + "scene": self.scene.lower(), + # "memoryType": MemoryTypeEnum.OBSERVATION.value, + # f"metaData.{DT}": time_to_formatted_str(msg_time_created), + }) + + all_nodes: List[MemoryWrapNode] = [MemoryWrapNode.init_from_es(hit) for hit in hits] + self.logger.info(f"retrieve_all.size={len(all_nodes)}") + self.set_context(ALL_NODES, all_nodes) + + all_memories = [] + if all_nodes: + for node in all_nodes: + all_memories.append(node.memory_node.to_dict()) + self.set_context(ALL_MEMORIES, all_memories) \ No newline at end of file diff --git a/memory_scope/worker/es/es_similar_worker.py b/memory_scope/worker/es/es_similar_worker.py new file mode 100644 index 00000000..3b0f1f10 --- /dev/null +++ b/memory_scope/worker/es/es_similar_worker.py @@ -0,0 +1,80 @@ +from typing import List + +from constants.common_constants import SIMILAR_OBS_NODES, RECALL_TYPE +from enumeration.memory_node_status import MemoryNodeStatus +from enumeration.memory_recall_type import MemoryRecallType +from enumeration.memory_type_enum import MemoryTypeEnum +from model.memory_wrap_node import MemoryWrapNode +from worker.bailian.memory_base_worker import MemoryBaseWorker + + +class EsSimilarWorker(MemoryBaseWorker): + + def es_similar_obs(self) -> List[MemoryWrapNode]: + query = self.messages[-1].content + hits = self.es_client.similar_search(text=query, + size=self.config.es_similar_top_k, + exact_filters={ + "memoryId": self.config.memory_id, + "status": MemoryNodeStatus.ACTIVE.value, + "scene": self.scene.lower(), + "memoryType": MemoryTypeEnum.OBSERVATION.value}) + + # 初始化成MemoryWrapNode，并加入召回源的参数 + similar_obs_nodes: List[MemoryWrapNode] = [] + for hit in hits: + node = MemoryWrapNode.init_from_es(hit) + node.memory_node.metaData[RECALL_TYPE] = MemoryRecallType.SIMILAR.value + similar_obs_nodes.append(node) + return similar_obs_nodes + + def es_similar_insight(self) -> List[MemoryWrapNode]: + query = self.messages[-1].content + hits = self.es_client.similar_search(text=query, + size=self.config.es_similar_top_k, + exact_filters={ + "memoryId": self.config.memory_id, + "status": MemoryNodeStatus.ACTIVE.value, + "scene": self.scene.lower(), + "memoryType": MemoryTypeEnum.INSIGHT.value}) + + # 初始化成MemoryWrapNode，并加入召回源的参数 + similar_obs_nodes: List[MemoryWrapNode] = [] + for hit in hits: + node = MemoryWrapNode.init_from_es(hit) + node.memory_node.metaData[RECALL_TYPE] = MemoryRecallType.SIMILAR.value + similar_obs_nodes.append(node) + return similar_obs_nodes + + def es_similar_obs_custom(self) -> List[MemoryWrapNode]: + query = self.messages[-1].content + hits = self.es_client.similar_search(text=query, + size=self.config.es_similar_top_k, + exact_filters={ + "memoryId": self.config.memory_id, + "status": MemoryNodeStatus.ACTIVE.value, + "scene": self.scene.lower(), + "memoryType": MemoryTypeEnum.OBS_CUSTOMIZED.value}) + + # 初始化成MemoryWrapNode，并加入召回源的参数 + similar_obs_nodes: List[MemoryWrapNode] = [] + for hit in hits: + node = MemoryWrapNode.init_from_es(hit) + node.memory_node.metaData[RECALL_TYPE] = MemoryRecallType.SIMILAR.value + similar_obs_nodes.append(node) + return similar_obs_nodes + + def _run(self): + for func in [self.es_similar_obs, self.es_similar_insight, self.es_similar_obs_custom]: + self.submit_thread(func, sleep_time=0.01) + + similar_obs_nodes: List[MemoryWrapNode] = [] + for result in self.join_threads(): + similar_obs_nodes.extend(result) + + self.logger.info(f"similar_obs_nodes.size={len(similar_obs_nodes)}") + for node in similar_obs_nodes: + self.logger.info(f"node={node.memory_node.content} " + f"score_similar={node.score_similar} " + f"type={node.memory_node.memoryType}") + self.set_context(SIMILAR_OBS_NODES, similar_obs_nodes) diff --git a/memory_scope/worker/es/es_today_obs_worker.py b/memory_scope/worker/es/es_today_obs_worker.py new file mode 100644 index 00000000..904ed171 --- /dev/null +++ b/memory_scope/worker/es/es_today_obs_worker.py @@ -0,0 +1,29 @@ +from typing import List + +from common.tool_functions import time_to_formatted_str +from constants.common_constants import TODAY_OBS_NODES, DT +from enumeration.memory_node_status import MemoryNodeStatus +from enumeration.memory_type_enum import MemoryTypeEnum +from model.memory_wrap_node import MemoryWrapNode +from worker.bailian.memory_base_worker import MemoryBaseWorker + + +class EsTodayObsWorker(MemoryBaseWorker): + + def _run(self): + if not self.messages: + self.logger.warning("messages is empty!") + return + msg_time_created = self.messages[-1].time_created + hits = self.es_client.exact_search_v2(size=self.config.es_today_obs_top_k, + term_filters={ + "memoryId": self.config.memory_id, + "status": MemoryNodeStatus.ACTIVE.value, + "scene": self.scene.lower(), + "memoryType": MemoryTypeEnum.OBSERVATION.value, + f"metaData.{DT}": time_to_formatted_str(msg_time_created), + }) + + today_obs_nodes: List[MemoryWrapNode] = [MemoryWrapNode.init_from_es(hit) for hit in hits] + self.logger.info(f"retrieve_today_obs.size={len(today_obs_nodes)}") + self.set_context(TODAY_OBS_NODES, today_obs_nodes) diff --git a/memory_scope/worker/es/load_profile_worker.py b/memory_scope/worker/es/load_profile_worker.py new file mode 100644 index 00000000..6065e804 --- /dev/null +++ b/memory_scope/worker/es/load_profile_worker.py @@ -0,0 +1,34 @@ +from typing import List, Dict + +from common.user_profile_handler import UserProfileHandler +from constants import common_constants +from enumeration.memory_node_status import MemoryNodeStatus +from enumeration.memory_type_enum import MemoryTypeEnum +from model.memory_wrap_node import MemoryWrapNode +from model.user_attribute import UserAttribute +from request.memory import MemoryServiceRequestModel +from worker.bailian.memory_base_worker import MemoryBaseWorker + + +class LoadProfileWorker(MemoryBaseWorker): + + def _run(self): + hits = self.es_client.exact_search_v2(size=10000, + term_filters={ + "memoryId": self.config.memory_id, + "status": MemoryNodeStatus.ACTIVE.value, + "scene": self.scene.lower(), + "memoryType": [MemoryTypeEnum.PROFILE.value, + MemoryTypeEnum.PROFILE_CUSTOMIZED.value], + }) + + user_profile_node: List[MemoryWrapNode] = [MemoryWrapNode.init_from_es(hit) for hit in hits] + user_profile_dict: Dict[str, UserAttribute] = UserProfileHandler.to_user_attr(user_profile_node) + + request: MemoryServiceRequestModel = self.get_context(common_constants.REQUEST) + for user_attr in request.user_profile: + user_profile_dict[user_attr.memory_key] = user_attr + request.user_profile = list(user_profile_dict.values()) + self.logger.info(f"retrieve_user_profile.size={len(user_profile_dict)}") + for key, user_attr in user_profile_dict.items(): + self.logger.info(f"{key}: {user_attr.description}: {user_attr.value}") diff --git a/memory_scope/worker/memory_base_worker.py b/memory_scope/worker/memory_base_worker.py new file mode 100644 index 00000000..b9bc8729 --- /dev/null +++ b/memory_scope/worker/memory_base_worker.py @@ -0,0 +1,132 @@ +from typing import List, Dict, Optional + +from common.dash_embedding_client import DashEmbeddingClient +from common.dash_generate_client import DashGenerateClient +from common.dash_rerank_client import DashReRankClient +from common.elastic_search_client import ElasticSearchClient +from config.bailian_memory_config import BailianMemoryConfig +from config.bailian_prompt_config import BailianPromptConfig +from constants import common_constants +from constants.common_constants import CONFIG, MESSAGES, PROMPT_CONFIG +from enumeration.message_role_enum import MessageRoleEnum +from model.message import Message +from model.user_attribute import UserAttribute +from request.memory import MemoryServiceRequestModel +from worker.bailian.base_worker import BaseWorker + + +class MemoryBaseWorker(BaseWorker): + + def __init__(self, **kwargs): + super(MemoryBaseWorker, self).__init__(**kwargs) + + self._user_profile_dict: Dict[str, UserAttribute] = {} + self._request_ext_info: Dict[str, str] = {} + + self._config: Optional[BailianMemoryConfig] = None + self._prompt_config: Optional[BailianPromptConfig] = None + + self._dash_embedding_client: Optional[DashEmbeddingClient] = None + self._dash_generate_client: Optional[DashGenerateClient] = None + self._dash_rerank_client: Optional[DashReRankClient] = None + + self._es_client: Optional[ElasticSearchClient] = None + + @property + def request(self) -> MemoryServiceRequestModel: + return self.get_context(common_constants.REQUEST) + + @property + def messages(self) -> List[Message]: + messages: List[Message] = self.context_handler.get_context(MESSAGES) + if messages is None: + messages = self.request.messages[-self.config.messages_pick_n:] + self.context_handler.set_context(MESSAGES, messages) + return messages + + @messages.setter + def messages(self, value): + self.context_handler.set_context(MESSAGES, value) + + @property + def user_profile_dict(self) -> Dict[str, UserAttribute]: + if not self._user_profile_dict: + self._user_profile_dict = {user_attr.memory_key: user_attr for user_attr in self.request.user_profile} + return self._user_profile_dict + + @property + def request_ext_info(self): + if not self._request_ext_info: + self._request_ext_info = self.request.ext_info + return self._request_ext_info + + @property + def config(self) -> BailianMemoryConfig: + if self._config is None: + self._config = self.get_context(CONFIG) + return self._config + + @property + def prompt_config(self) -> BailianPromptConfig: + if self._prompt_config is None: + self._prompt_config = self.get_context(PROMPT_CONFIG) + if not self._prompt_config: + self._prompt_config = BailianPromptConfig() + self.set_context(PROMPT_CONFIG, self._prompt_config) + return self._prompt_config + + @property + def emb_client(self): + if self._dash_embedding_client is None: + self._dash_embedding_client = DashEmbeddingClient(request_id=self.config.request_id, + dash_scope_uid=self.config.uid, + authorization=self.config.api_key, + workspace=self.config.workspace_id, + env_type=self.env_type, + max_retry_count=self.config.dash_embedding_retry_cnt) + return self._dash_embedding_client + + @property + def gene_client(self): + if self._dash_generate_client is None: + self._dash_generate_client = DashGenerateClient(request_id=self.config.request_id, + dash_scope_uid=self.config.uid, + authorization=self.config.api_key, + workspace=self.config.workspace_id, + env_type=self.env_type, + max_retry_count=self.config.dash_generate_retry_cnt) + return self._dash_generate_client + + @property + def rerank_client(self): + if self._dash_rerank_client is None: + self._dash_rerank_client = DashReRankClient(request_id=self.config.request_id, + dash_scope_uid=self.config.uid, + authorization=self.config.api_key, + workspace=self.config.workspace_id, + env_type=self.env_type, + max_retry_count=self.config.dash_rerank_retry_cnt) + return self._dash_rerank_client + + @property + def es_client(self): + if self._es_client is None: + self._es_client = ElasticSearchClient(es_user_name=self.config.es_user_name, + es_password=self.config.es_password, + es_index_name=self.config.es_index_name, + embedding_client=self.emb_client, + max_retries=self.config.es_retry_cnt) + return self._es_client + + @staticmethod + def prompt_to_msg(system_prompt: str, few_shot: str, user_query: str): + return [ + { + "role": MessageRoleEnum.SYSTEM.value, + "content": system_prompt.strip(), + }, + { + "role": MessageRoleEnum.USER.value, + "content": "\n".join([x.strip() for x in [few_shot, system_prompt, user_query]]) + }, + ] diff --git a/memory_scope/worker/memory_store_worker.py b/memory_scope/worker/memory_store_worker.py new file mode 100644 index 00000000..437258a0 --- /dev/null +++ b/memory_scope/worker/memory_store_worker.py @@ -0,0 +1,34 @@ +from typing import List + +from common.user_profile_handler import UserProfileHandler +from constants.common_constants import MODIFIED_MEMORIES, NEW_USER_PROFILE +from model.memory_node import MemoryNode +from model.memory_wrap_node import MemoryWrapNode +from model.user_attribute import UserAttribute +from worker.bailian.memory_base_worker import MemoryBaseWorker + + +class MemoryStoreWorker(MemoryBaseWorker): + + def _run(self): + modified_memories: List[MemoryWrapNode] | List[MemoryNode] = self.get_context(MODIFIED_MEMORIES) + if modified_memories: + if isinstance(modified_memories[0], MemoryWrapNode): + modified_memories = [n.memory_node for n in modified_memories] + + for n in modified_memories: + if not n.id: + n.id = f"{n.memoryId}_{n.scene}_content_{n.content}" + n.code = n.id + # TODO add batch insert + self.es_client.insert(n.id, body=n.model_dump(exclude=set("content_modified", ))) + else: + self.logger.warning("modified_memories is empty!") + + new_user_profile: List[UserAttribute] = self.get_context(NEW_USER_PROFILE) + if new_user_profile: + new_user_nodes: List[MemoryNode] = [n.memory_node for n in UserProfileHandler.to_nodes(new_user_profile)] + for n in new_user_nodes: + self.es_client.insert(n.id, body=n.model_dump(exclude=set("content_modified", ))) + else: + self.logger.warning("new_user_profile is empty!") diff --git a/memory_scope/worker/parse_params_worker.py b/memory_scope/worker/parse_params_worker.py new file mode 100644 index 00000000..ed92d8fd --- /dev/null +++ b/memory_scope/worker/parse_params_worker.py @@ -0,0 +1,36 @@ +import json + +from config.bailian_memory_config import BailianMemoryConfig +from constants.common_constants import REQUEST, CONFIG +from request.memory import MemoryServiceRequestModel +from worker.bailian.base_worker import BaseWorker + + +class ParseParamsWorker(BaseWorker): + + def _run(self): + # 参数合并 + memory_config = {} + + # 更新环境变量 + memory_config.update(self.context_handler.env_configs) + + # 更新请求参数 + request: MemoryServiceRequestModel = self.context_handler.get_context(REQUEST) + memory_config.update(request.model_dump(exclude=set("ext_info", ))) + + # 更新ext_info + if request.ext_info: + memory_config.update(request.ext_info) + + # 存入上下文 + memory_config_model: BailianMemoryConfig = BailianMemoryConfig(**memory_config) + self.context_handler.set_context(CONFIG, memory_config_model) + + # 打印 + self.logger.info(f"memory_config_model={json.dumps(memory_config_model.model_dump(), ensure_ascii=False)}") + + # 上游可能没有传这个参数，可能隐藏在memory_id做区分 + if request.user_profile: + for user_attr in request.user_profile: + user_attr.scene = request.scene diff --git a/memory_scope/worker/retrieve/__init__.py b/memory_scope/worker/retrieve/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/memory_scope/worker/retrieve/extract_time_worker.py b/memory_scope/worker/retrieve/extract_time_worker.py new file mode 100644 index 00000000..54094089 --- /dev/null +++ b/memory_scope/worker/retrieve/extract_time_worker.py @@ -0,0 +1,65 @@ +import re + +from common.tool_functions import time_to_formatted_str +from constants.common_constants import DATATIME_WORD_LIST, DATATIME_KEY_MAP +from constants.common_constants import EXTRACT_TIME_DICT +from worker.bailian.memory_base_worker import MemoryBaseWorker + + +class ExtractTimeWorker(MemoryBaseWorker): + + @staticmethod + def get_parse_time_prompt(query: str, query_time_str: str): + return f""" +任务指令：从语句与语句发生的时间，推断并提取语句内容中指向的时间段。回答尽可能完整的时间段。 +语句：{query} +时间：{query_time_str} +回答： + """.strip() + + def _run(self): + # save to context + extract_time_dict = {} + self.set_context(EXTRACT_TIME_DICT, extract_time_dict) + + # get query & time_created_dt + query = self.messages[-1].content + time_created = self.messages[-1].time_created + + # find datetime keyword + contain_datetime = False + for datetime_word in DATATIME_WORD_LIST: + if datetime_word in query: + contain_datetime = True + break + if not contain_datetime: + self.logger.info(f"contain_datetime={contain_datetime}") + return + + # prepare prompt + time_format = "{year}年{month}月{day}日，{year}年第{week}周，{weekday}，{hour}时{minute}分{second}秒。" + query_time_str = time_to_formatted_str(time=time_created, + date_format="", + string_format=time_format) + extract_time_prompt = self.get_parse_time_prompt(query=query, query_time_str=query_time_str) + self.logger.info(f"extract_time_prompt={extract_time_prompt}") + + # call sft model + response_text = self.gene_client.call(prompt=extract_time_prompt, + model_name=self.config.parse_time_model, + max_token=self.config.parse_time_max_token, + temperature=self.config.parse_time_temperature, + top_k=self.config.parse_time_top_k) + + # if empty, return + if not response_text: + return + + # re-match time info to dict + pattern = r'-\s*(\S+)：(\d+)' + matches = re.findall(pattern, response_text) + for key, value in matches: + if key in DATATIME_KEY_MAP.keys(): + extract_time_dict[DATATIME_KEY_MAP[key]] = value + + self.logger.info(f"response_text={response_text} filters={extract_time_dict}") diff --git a/memory_scope/worker/retrieve/fuse_rerank_worker.py b/memory_scope/worker/retrieve/fuse_rerank_worker.py new file mode 100644 index 00000000..225ef0b0 --- /dev/null +++ b/memory_scope/worker/retrieve/fuse_rerank_worker.py @@ -0,0 +1,115 @@ +from typing import Dict, List + +from constants.common_constants import RELATED_MEMORIES, EXTRACT_TIME_DICT, ALL_ONLINE_NODES, DEFAULT_SYSTEM_PROMPT, \ + TIME_MATCHED +from model.memory_wrap_node import MemoryWrapNode +from worker.bailian.memory_base_worker import MemoryBaseWorker + + +class FuseRerankWorker(MemoryBaseWorker): + + @staticmethod + def format_time_infer(time_infer: str, extract_time_dict: Dict[str, str], meta_data: Dict[str, str]): + if time_infer: + return time_infer + + time_infer = "" + if "year" in extract_time_dict: + value = meta_data.get("msg_year") + if value: + time_infer += f"{value}年" + elif value == "-1": + time_infer += f"每年" + + if "month" in extract_time_dict: + value = meta_data.get("msg_month") + if value: + time_infer += f"{value}月" + elif value == "-1": + time_infer += f"每月" + + if "day" in extract_time_dict: + value = meta_data.get("msg_day") + if value: + time_infer += f"{value}日" + elif value == "-1": + time_infer += f"每日" + + if "weekday" in extract_time_dict: + value = meta_data.get("msg_weekday") + if value: + time_infer += value + + return time_infer + + def _run(self): + # 解析时间meta信息 + extract_time_dict: Dict[str, str] = self.get_context(EXTRACT_TIME_DICT) + all_online_nodes: List[MemoryWrapNode] = self.get_context(ALL_ONLINE_NODES) + + if not all_online_nodes: + self.add_run_info("all_online_nodes is empty, stop") + return + + filtered_nodes = [] + for node in all_online_nodes: + if node.score_rank < self.config.fuse_score_threshold: + continue + + # 根据类型给ratio + type_ratio: float = self.config.fuse_ratio_dict.get(node.memory_node.memoryType, 0.1) + + # 时间系数，完全匹配才行 + fuse_time_ratio: float = 1.0 + match_event_flag = False + match_msg_flag = False + if extract_time_dict: + match_event_flag = True + for k, v in extract_time_dict.items(): + event_value = node.memory_node.metaData.get(f"event_{k}", "") + if event_value in ["-1", v]: + continue + else: + match_event_flag = False + break + + match_msg_flag = True + for k, v in extract_time_dict.items(): + msg_value = node.memory_node.metaData.get(f"msg_{k}", "") + if msg_value == v: + continue + else: + match_msg_flag = False + break + + if match_event_flag or match_msg_flag: + fuse_time_ratio = self.config.fuse_time_ratio + node.memory_node.metaData[TIME_MATCHED] = "1" + + node.score_rerank = node.score_rank * type_ratio * fuse_time_ratio + self.logger.info(f"content={node.memory_node.content} f_event={int(match_event_flag)} " + f"f_msg={int(match_msg_flag)} score_rerank={node.score_rerank}") + filtered_nodes.append(node) + + # get output & save context + filtered_nodes = sorted(filtered_nodes, key=lambda x: x.score_rerank, reverse=True) + filtered_nodes = filtered_nodes[: self.config.output_max_count] + related_memories: List[str] = [] + for node in filtered_nodes: + content = node.memory_node.content + + # 如果命中时间逻辑 + if node.memory_node.metaData.get(TIME_MATCHED, "") == "1": + # time_infer = node.memory_node.metaData.get(TIME_INFER) + # if not time_infer: + # time_infer = self.format_time_infer(time_infer=time_infer, + # extract_time_dict=extract_time_dict, + # meta_data=node.memory_node.metaData) + time_infer = self.format_time_infer(time_infer="", + extract_time_dict=extract_time_dict, + meta_data=node.memory_node.metaData) + content = f"{time_infer}: {content}" + related_memories.append(content) + + self.set_context(RELATED_MEMORIES, related_memories) + self.set_context(DEFAULT_SYSTEM_PROMPT, self.config.default_system_prompt) diff --git a/memory_scope/worker/retrieve/semantic_rank_worker.py b/memory_scope/worker/retrieve/semantic_rank_worker.py new file mode 100644 index 00000000..08baf118 --- /dev/null +++ b/memory_scope/worker/retrieve/semantic_rank_worker.py @@ -0,0 +1,68 @@ +from typing import List, Dict + +from common.user_profile_handler import UserProfileHandler +from constants.common_constants import SIMILAR_OBS_NODES, RECALL_TYPE, KEYWORD_OBS_NODES, ALL_ONLINE_NODES, \ + QUERY_KEYWORDS +from enumeration.memory_recall_type import MemoryRecallType +from model.memory_wrap_node import MemoryWrapNode +from worker.bailian.memory_base_worker import MemoryBaseWorker + + +class SemanticRankWorker(MemoryBaseWorker): + + def user_profile_to_nodes(self) -> List[MemoryWrapNode]: + user_profile_nodes: List[MemoryWrapNode] = UserProfileHandler.to_nodes(self.user_profile_dict, split_value=True) + for node in user_profile_nodes: + # 从画像侧召回 + node.memory_node.metaData[RECALL_TYPE] = MemoryRecallType.PROFILE + self.logger.info(f"user profile node={node.memory_node.content}") + return user_profile_nodes + + def _run(self): + all_node_dict: Dict[str, MemoryWrapNode] = {} + + # 优先级: similar_obs_nodes < keyword_obs_nodes < profile_nodes + similar_obs_nodes: List[MemoryWrapNode] = self.get_context(SIMILAR_OBS_NODES) + if similar_obs_nodes: + for node in similar_obs_nodes: + all_node_dict[node.memory_node.content] = node + + keyword_obs_nodes: List[MemoryWrapNode] = self.get_context(KEYWORD_OBS_NODES) + if keyword_obs_nodes: + for node in keyword_obs_nodes: + all_node_dict[node.memory_node.content] = node + + profile_nodes: List[MemoryWrapNode] = self.user_profile_to_nodes() + if profile_nodes: + for node in profile_nodes: + all_node_dict[node.memory_node.content] = node + + if not all_node_dict: + self.add_run_info(f"all_node_dict is empty!", continue_run=False) + return + + # call recall model + query_keywords = self.get_context(QUERY_KEYWORDS) + # TODO 根据效果更改 + # query: str = "用户：" + self.messages[-1].content + query: str = self.messages[-1].content + if query_keywords: + query_keyword_join = "，".join(query_keywords) + query = f"{query} 用户的{query_keyword_join}。" + documents = list(all_node_dict.keys()) + result = self.rerank_client.call(query=query, documents=documents) + + if not result: + self.add_run_info(f"semantic call recall model failed!") + return + + # set score + for rank_node in result: + content = documents[rank_node["index"]] + node = all_node_dict[content] + node.score_rank = rank_node["relevance_score"] + self.logger.info(f"query={query} content={node.memory_node.content} score_rank={node.score_rank}") + + # save context + all_online_nodes: List[MemoryWrapNode] = list(all_node_dict.values()) + self.set_context(ALL_ONLINE_NODES, all_online_nodes) diff --git a/memory_scope/worker/summary_long/__init__.py b/memory_scope/worker/summary_long/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/memory_scope/worker/summary_long/get_insight_worker.py b/memory_scope/worker/summary_long/get_insight_worker.py new file mode 100644 index 00000000..a2494db0 --- /dev/null +++ b/memory_scope/worker/summary_long/get_insight_worker.py @@ -0,0 +1,132 @@ +from datetime import datetime +from typing import List + +from common.tool_functions import time_to_formatted_str, get_datetime_info_dict +from constants.common_constants import NEW_INSIGHT_NODES, DT, NOT_REFLECTED_MERGE_NODES, NEW_INSIGHT_KEYS, INSIGHT_KEY, \ + INSIGHT_VALUE, REFLECTED +from enumeration.memory_node_status import MemoryNodeStatus +from enumeration.memory_type_enum import MemoryTypeEnum +from model.memory_wrap_node import MemoryWrapNode +from worker.bailian.memory_base_worker import MemoryBaseWorker + + +class GetInsightWorker(MemoryBaseWorker): + + def new_insight_node(self, insight_key: str, insight_value: str) -> MemoryWrapNode: + created_dt = datetime.now() + dt = time_to_formatted_str(time=created_dt) + + # 组合meta_data + meta_data = { + DT: dt, + INSIGHT_KEY: insight_key, + INSIGHT_VALUE: insight_value, + } + meta_data.update({f"msg_{k}": str(v) for k, v in get_datetime_info_dict(created_dt).items()}) + + content = f"用户的{insight_key}：{insight_value}" + return MemoryWrapNode.init_from_attrs(content=content, + memoryId=self.config.memory_id, + scene=self.scene, + memoryType=MemoryTypeEnum.INSIGHT.value, + content_modified=True, # 新增的insight需要置为true + metaData=meta_data, + status=MemoryNodeStatus.ACTIVE.value, + tenantId=self.config.tenant_id) + + def reflect_new_insight_key(self, + insight_key: str, + not_reflected_merge_nodes: List[MemoryWrapNode]) -> MemoryWrapNode | None: + + # 检索历史memory + hits = self.es_client.similar_search(text=insight_key, + size=self.config.es_insight_similar_top_k, + exact_filters={ + "memoryId": self.config.memory_id, + "status": MemoryNodeStatus.ACTIVE.value, + "scene": self.scene.lower(), + "memoryType": [MemoryTypeEnum.OBSERVATION.value, + MemoryTypeEnum.OBS_CUSTOMIZED.value], + }) + + # 转化成 MemoryNodeWrap 合并新增nodes + related_nodes: List[MemoryWrapNode] = [MemoryWrapNode.init_from_es(x) for x in hits] + related_nodes.extend(not_reflected_merge_nodes) + + # content去重 + related_node_dict = {n.memory_node.content: n for n in related_nodes} + related_nodes = sorted(list(related_node_dict.values()), key=lambda x: x.memory_node.id) + documents = [n.memory_node.content for n in related_nodes] + + # 重排所有记忆 + result = self.rerank_client.call(query=insight_key, documents=documents) + if not result: + self.add_run_info(f"reflect insight_key={insight_key} call rerank client failed!") + return + + # 根据打分过滤 + for rank_node in result: + index = rank_node["index"] + score = rank_node["relevance_score"] + related_nodes[index].score_rank = score + related_nodes_sorted = sorted(related_nodes, key=lambda x: x.score_rank, reverse=True)[ + :self.config.insight_obs_max_cnt] + + # 生成prompt + user_query_list = [x.memory_node.content for x in related_nodes_sorted] + get_insight_message = self.prompt_to_msg( + system_prompt=self.prompt_config.get_insight_system, + few_shot=self.prompt_config.get_insight_few_shot, + user_query=self.prompt_config.get_insight_user_query.format( + insight_key=insight_key, user_query="\n".join(user_query_list))) + self.logger.info(f"get_insight_message={get_insight_message}") + + # call LLM, 提取insight + response_text = self.gene_client.call(messages=get_insight_message, + model_name=self.config.get_insight_model, + max_token=self.config.get_insight_max_token, + temperature=self.config.get_insight_temperature, + top_k=self.config.get_insight_top_k) + + # return if empty + if not response_text: + self.add_run_info("reflect_upon_user_attr call llm failed!") + return + response_text = response_text.strip() + if response_text in ["无"]: + return + return self.new_insight_node(insight_key=insight_key, insight_value=response_text) + + def _run(self): + new_insight_keys: List[MemoryWrapNode] = self.get_context(NEW_INSIGHT_KEYS) + if not new_insight_keys: + self.add_run_info("new_insight_keys is empty! stop insight.") + return + + not_reflected_merge_nodes: List[MemoryWrapNode] = self.get_context(NOT_REFLECTED_MERGE_NODES) + if not not_reflected_merge_nodes: + self.add_run_info("not_reflected_merge_nodes is empty! stop get insight.") + return + + # submit insight task + for insight_key in new_insight_keys: + self.submit_thread(self.reflect_new_insight_key, + sleep_time=1, + insight_key=insight_key, + not_reflected_merge_nodes=not_reflected_merge_nodes) + + # save output + new_insight_nodes: List[MemoryWrapNode] = [] + for result in self.join_threads(): + if result: + new_insight_nodes.append(result) + assert isinstance(result, MemoryWrapNode) + insight_key = result.memory_node.metaData.get(INSIGHT_KEY, "") + insight_value = result.memory_node.metaData.get(INSIGHT_VALUE, "") + self.logger.info(f"after_get_insight insight_key={insight_key} insight_value={insight_value}") + + self.set_context(NEW_INSIGHT_NODES, new_insight_nodes) + + # set REFLECTED + for node in not_reflected_merge_nodes: + node.memory_node.metaData[REFLECTED] = "1" diff --git a/memory_scope/worker/summary_long/get_reflection_worker.py b/memory_scope/worker/summary_long/get_reflection_worker.py new file mode 100644 index 00000000..1b8d4419 --- /dev/null +++ b/memory_scope/worker/summary_long/get_reflection_worker.py @@ -0,0 +1,72 @@ +from typing import List + +from common.response_text_parser import ResponseTextParser +from constants.common_constants import NEW_OBS_NODES, NOT_REFLECTED_OBS_NODES, REFLECTED, INSIGHT_NODES, INSIGHT_KEY, \ + NEW_INSIGHT_KEYS, NOT_REFLECTED_MERGE_NODES +from model.memory_wrap_node import MemoryWrapNode +from worker.bailian.memory_base_worker import MemoryBaseWorker + + +class GetReflectionWorker(MemoryBaseWorker): + + def _run(self): + # 过滤得到 not_reflected_merge_nodes + new_obs_nodes: List[MemoryWrapNode] = self.get_context(NEW_OBS_NODES) + not_reflected_nodes: List[MemoryWrapNode] = self.get_context(NOT_REFLECTED_OBS_NODES) + not_reflected_merge_nodes: List[MemoryWrapNode] = [] + if new_obs_nodes: + not_reflected_merge_nodes.extend(new_obs_nodes) + if not_reflected_nodes: + not_reflected_merge_nodes.extend(not_reflected_nodes) + not_reflected_merge_nodes = [node for node in not_reflected_merge_nodes + if node.memory_node.metaData.get(REFLECTED, "") == "0"] + + # count + not_reflected_count = len(not_reflected_merge_nodes) + if not_reflected_count <= self.config.reflect_obs_cnt_threshold: + self.logger.info(f"not_reflected_count={not_reflected_count} is not enough, stop reflect.") + return + + # save context + self.set_context(NOT_REFLECTED_MERGE_NODES, not_reflected_merge_nodes) + + # get profile_keys + exist_keys: List[str] = [] + profile_keys: List[str] = list(self.user_profile_dict.keys()) + exist_keys.extend(profile_keys) + self.logger.info(f"profile_keys={profile_keys}") + + # get insight_keys + insight_nodes: List[MemoryWrapNode] = self.get_context(INSIGHT_NODES) + if insight_nodes: + insight_keys = [n.memory_node.metaData.get(INSIGHT_KEY) for n in insight_nodes] + insight_keys = [x.strip() for x in insight_keys if x] + exist_keys.extend(insight_keys) + self.logger.info(f"insight_keys={insight_keys}") + + # gen reflect prompt + user_query_list = [n.memory_node.content for n in not_reflected_merge_nodes] + reflect_message = self.prompt_to_msg( + system_prompt=self.prompt_config.get_reflect_system.format( + num_questions=self.config.reflect_num_questions), + few_shot=self.prompt_config.get_reflect_few_shot, + user_query=self.prompt_config.get_reflect_user_query.format(exist_keys="，".join(exist_keys), + user_query="\n".join(user_query_list))) + self.logger.info(f"reflect_message={reflect_message}") + + # call LLM + response_text = self.gene_client.call(messages=reflect_message, + model_name=self.config.reflect_obs_model, + max_token=self.config.reflect_obs_max_token, + temperature=self.config.reflect_obs_temperature, + top_k=self.config.reflect_obs_top_k) + + # return if empty + if not response_text: + self.add_run_info("reflect_obs_questions call llm failed!") + return + + # parse text & save + new_insight_keys = ResponseTextParser(response_text).parse_v2("get_insight_keys") + if new_insight_keys: + self.set_context(NEW_INSIGHT_KEYS, new_insight_keys) diff --git a/memory_scope/worker/summary_long/long_contra_repeat_worker.py b/memory_scope/worker/summary_long/long_contra_repeat_worker.py new file mode 100644 index 00000000..c16af9d4 --- /dev/null +++ b/memory_scope/worker/summary_long/long_contra_repeat_worker.py @@ -0,0 +1,111 @@ +from typing import List + +from common.response_text_parser import ResponseTextParser +from constants.common_constants import NEW_OBS_NODES, TODAY_OBS_NODES, MSG_TIME, NEW_OBS_WITH_TIME_NODES, \ + MODIFIED_MEMORIES +from enumeration.memory_node_status import MemoryNodeStatus +from enumeration.memory_type_enum import MemoryTypeEnum +from model.memory_wrap_node import MemoryWrapNode +from worker.bailian.memory_base_worker import MemoryBaseWorker + + +class LongContraRepeatWorker(MemoryBaseWorker): + + def _run(self): + # 合并当前的obs和今日的obs + new_obs_nodes: List[MemoryWrapNode] = self.get_context(NEW_OBS_NODES) + # new_obs_with_time_nodes: List[MemoryWrapNode] = self.get_context(NEW_OBS_WITH_TIME_NODES) + # oday_obs_nodes: List[MemoryWrapNode] = self.get_context(TODAY_OBS_NODES) + all_obs_nodes: List[MemoryWrapNode] = [] + for new_obs_node in new_obs_nodes: + text = new_obs_node.memory_node.content + hits = self.es_client.similar_search(text=text, + size=self.config.es_contra_repeat_similar_top_k, + exact_filters={ + "memoryId": self.config.memory_id, + "status": MemoryNodeStatus.ACTIVE.value, + "scene": self.scene.lower(), + "memoryType": [MemoryTypeEnum.OBSERVATION.value, + MemoryTypeEnum.OBS_CUSTOMIZED.value], + }) + + related_nodes: List[MemoryWrapNode] = [MemoryWrapNode.init_from_es(x) for x in hits] + + has_match = False + for related_node in related_nodes: + if related_node.score_similar < self.config.long_contra_repeat_threshold: + continue + else: + has_match = True + all_obs_nodes.append(related_node) + if has_match: + all_obs_nodes.append(new_obs_node) + + if not all_obs_nodes: + self.add_run_info("all_obs_nodes is empty!") + return + + # gene prompt + user_query_list = [] + all_obs_nodes = sorted(all_obs_nodes, key=lambda x: x.memory_node.metaData.get(MSG_TIME, ""), reverse=True) + for i, n in enumerate(all_obs_nodes): + user_query_list.append(f"{i + 1} {n.memory_node.content}") + merge_obs_message = self.prompt_to_msg( + system_prompt=self.prompt_config.long_contra_repeat_system.format(num_obs=len(user_query_list)), + few_shot=self.prompt_config.long_contra_repeat_few_shot, + user_query=self.prompt_config.long_contra_repeat_user_query.format(user_query="\n".join(user_query_list))) + self.logger.info(f"merge_obs_message={merge_obs_message}") + + # call LLM + response_text = self.gene_client.call(messages=merge_obs_message, + model_name=self.config.merge_obs_model, + max_token=self.config.merge_obs_max_token, + temperature=self.config.merge_obs_temperature, + top_k=self.config.merge_obs_top_k) + + # return if empty + if not response_text: + self.add_run_info("contra repeat call llm failed!") + return + + # parse text + idx_merge_obs_list = ResponseTextParser(response_text).parse_v1("contra_repeat") + if len(idx_merge_obs_list) <= 0: + self.add_run_info("idx_merge_obs_list is empty!") + return + + # add merged obs + merge_obs_nodes: List[MemoryWrapNode] = [] + for obs_content_list in idx_merge_obs_list: + if not obs_content_list: + continue + + # [6, 逃课] + if len(obs_content_list) != 2: + self.logger.warning(f"obs_content_list={obs_content_list} is invalid!") + continue + + idx, keep_flag = obs_content_list + + if not idx.isdigit(): + self.logger.warning(f"idx={idx} is invalid!") + continue + + # 序号需要修正-1 + idx = int(idx) - 1 + if idx >= len(all_obs_nodes): + self.logger.warning(f"idx={idx} is invalid!") + continue + + if keep_flag not in ["矛盾", "被包含", "无"]: + self.logger.warning(f"keep_flag={keep_flag} is invalid!") + continue + + node: MemoryWrapNode = all_obs_nodes[idx] + if keep_flag != "无": + node.memory_node.status = MemoryNodeStatus.EXPIRED.value + merge_obs_nodes.append(node) + self.logger.info(f"after contra repeat: {node.memory_node.content} {node.memory_node.status}") + + # save context + self.set_context(MODIFIED_MEMORIES, merge_obs_nodes) diff --git a/memory_scope/worker/summary_long/summary_collect_worker.py b/memory_scope/worker/summary_long/summary_collect_worker.py new file mode 100644 index 00000000..e26ab6fa --- /dev/null +++ b/memory_scope/worker/summary_long/summary_collect_worker.py @@ -0,0 +1,33 @@ +from typing import List, Dict + +from constants.common_constants import NEW_INSIGHT_NODES, MODIFIED_MEMORIES, INSIGHT_NODES, NEW_OBS_NODES, \ + NOT_REFLECTED_OBS_NODES, NEW, NOT_REFLECTED_MERGE_NODES +from model.memory_wrap_node import MemoryWrapNode +from worker.bailian.memory_base_worker import MemoryBaseWorker + + +class SummaryCollectWorker(MemoryBaseWorker): + + def _run(self): + insight_nodes: List[MemoryWrapNode] = self.get_context(INSIGHT_NODES) + new_insight_nodes: List[MemoryWrapNode] = self.get_context(NEW_INSIGHT_NODES) + new_obs_nodes: List[MemoryWrapNode] = self.get_context(NEW_OBS_NODES) + not_reflected_nodes: List[MemoryWrapNode] = self.get_context(NOT_REFLECTED_OBS_NODES) + not_reflected_merge_nodes: List[MemoryWrapNode] = self.get_context(NOT_REFLECTED_MERGE_NODES) + + # 合并逻辑，复杂，务必check + all_node_dict: Dict[str, MemoryWrapNode] = {} + if insight_nodes: + all_node_dict.update({n.id: n for n in insight_nodes if n.memory_node.content_modified}) + if new_insight_nodes: + all_node_dict.update({n.memory_node.content: n for n in new_insight_nodes}) + if new_obs_nodes: + # 设置为非新 + for n in new_obs_nodes: + n.memory_node.metaData[NEW] = "0" + all_node_dict.update({n.memory_node.content: n for n in new_obs_nodes}) + if not_reflected_merge_nodes and not_reflected_nodes: + # 进入reflect阶段 + all_node_dict.update({n.id: n for n in not_reflected_nodes}) + + self.set_context(MODIFIED_MEMORIES, list(all_node_dict.values())) diff --git a/memory_scope/worker/summary_long/update_insight_worker.py b/memory_scope/worker/summary_long/update_insight_worker.py new file mode 100644 index 00000000..ee2d889d --- /dev/null +++ b/memory_scope/worker/summary_long/update_insight_worker.py @@ -0,0 +1,160 @@ +from datetime import datetime +from typing import List + +from common.response_text_parser import ResponseTextParser +from common.tool_functions import time_to_formatted_str, get_datetime_info_dict +from constants.common_constants import INSIGHT_NODES, NEW_OBS_NODES, INSIGHT_KEY, INSIGHT_VALUE, DT +from model.memory_wrap_node import MemoryWrapNode +from worker.bailian.memory_base_worker import MemoryBaseWorker + + +class UpdateInsightWorker(MemoryBaseWorker): + + def filter_obs_nodes(self, + insight_node: MemoryWrapNode, + new_obs_nodes: List[MemoryWrapNode]) -> (MemoryWrapNode, List[MemoryWrapNode], float): + max_score: float = 0 + filtered_nodes: List[MemoryWrapNode] = [] + + insight_key = insight_node.memory_node.metaData.get(INSIGHT_KEY, "") + insight_value = insight_node.memory_node.metaData.get(INSIGHT_VALUE, "") + if not insight_key or not insight_value: + self.logger.warning(f"insight_key={insight_key} insight_value={insight_value} is empty!") + return insight_node, filtered_nodes, max_score + + result = self.rerank_client.call(query=insight_key, + documents=[x.memory_node.content for x in new_obs_nodes]) + + if not result: + self.add_run_info(f"update_insight={insight_key} call rerank failed!") + return insight_node, filtered_nodes, max_score + + # 找到大于阈值的obs node + + for rank_node in result: + index = rank_node["index"] + score = rank_node["relevance_score"] + node = new_obs_nodes[index] + keep_flag = "filtered" + if score >= self.config.update_insight_threshold: + filtered_nodes.append(node) + keep_flag = "keep" + max_score = max(max_score, score) + self.logger.info(f"insight_key={insight_key} insight_value={insight_value} " + f"score={score} keep_flag={keep_flag}") + + if not filtered_nodes: + self.logger.info(f"update_insight={insight_key} filtered_nodes is empty!") + + return insight_node, filtered_nodes, max_score + + def update_insight_node(self, insight_node: MemoryWrapNode, insight_key: str, insight_value: str): + created_dt = datetime.now() + dt = time_to_formatted_str(time=created_dt) + meta_data = { + DT: dt, + INSIGHT_KEY: insight_key, + INSIGHT_VALUE: insight_value, + } + meta_data.update({f"msg_{k}": str(v) for k, v in get_datetime_info_dict(created_dt).items()}) + + content = f"用户的{insight_key}：{insight_value}" + insight_node.memory_node.content = content + insight_node.memory_node.content_modified = True + insight_node.memory_node.metaData = meta_data + insight_node.memory_node.tenantId = self.config.tenant_id + return insight_node + + def update_insight(self, + insight_node: MemoryWrapNode, + filtered_nodes: List[MemoryWrapNode]) -> MemoryWrapNode: + + insight_key = insight_node.memory_node.metaData.get(INSIGHT_KEY, "") + insight_value = insight_node.memory_node.metaData.get(INSIGHT_VALUE, "") + self.logger.info(f"update_insight insight_key={insight_key} insight_value={insight_value} " + f"doc.size={len(filtered_nodes)}") + + # gen prompt + user_query_list = [] + for node in filtered_nodes: + user_query_list.append(f"句子：{node.memory_node.content}") + update_insight_message = self.prompt_to_msg( + system_prompt=self.prompt_config.update_insight_system, + few_shot=self.prompt_config.update_insight_few_shot, + user_query=self.prompt_config.update_insight_user_query.format( + user_query="\n".join(user_query_list), + insight_key=insight_key, + insight_key_value=insight_key + "：" + insight_value)) + self.logger.info(f"update_insight_message={update_insight_message}") + + # call LLM + response_text: str = self.gene_client.call(messages=update_insight_message, + model_name=self.config.update_insight_model, + max_token=self.config.update_insight_max_token, + temperature=self.config.update_insight_temperature, + top_k=self.config.update_insight_top_k) + + # return if empty + if not response_text: + self.add_run_info(f"update_insight insight_key={insight_key} call llm failed!") + return insight_node + + profile_list = ResponseTextParser(response_text).parse_v1(f"update_profile {insight_key}") + if not profile_list: + self.add_run_info(f"update_insight insight_key={insight_key} profile_list empty 1!") + return insight_node + profile_list = profile_list[0] + if not profile_list: + self.add_run_info(f"update_insight insight_key={insight_key} profile_list empty 2") + return insight_node + insight_value = profile_list[0] + + if not insight_value or insight_value in ["无", "重复"]: + self.logger.info(f"insight_value={insight_value}, skip.") + return insight_node + + return self.update_insight_node(insight_node, insight_key, insight_value) + + def _run(self): + # 获取新的obs和insight + new_obs_nodes: List[MemoryWrapNode] = self.get_context(NEW_OBS_NODES) + insight_nodes: List[MemoryWrapNode] = self.get_context(INSIGHT_NODES) + if not new_obs_nodes: + self.logger.info("new_obs_nodes is empty, stop update sights!") + return + if not insight_nodes: + self.logger.info("insight_nodes is empty, stop update sights!") + return + + # 提交打分任务 + for node in insight_nodes: + self.submit_thread(self.filter_obs_nodes, + sleep_time=0.1, + insight_node=node, + new_obs_nodes=new_obs_nodes) + + # 选择topN + result_list = [] + for result in self.join_threads(): + insight_node, filtered_nodes, max_score = result + if not filtered_nodes: + continue + result_list.append(result) + result_sorted = sorted(result_list, key=lambda x: x[2], reverse=True) + if len(result_sorted) > self.config.update_insight_max_thread: + result_sorted = result_sorted[:self.config.update_insight_max_thread] + + # 提交LLM update任务 + for insight_node, filtered_nodes, _ in result_sorted: + self.submit_thread(self.update_insight, + sleep_time=1, + insight_node=insight_node, + filtered_nodes=filtered_nodes) + + # 等待结果 + for result in self.join_threads(): + if result: + insight_node: MemoryWrapNode = result + insight_key = insight_node.memory_node.metaData.get(INSIGHT_KEY, "") + insight_value = insight_node.memory_node.metaData.get(INSIGHT_VALUE, "") + self.logger.info(f"after_update_insight insight_key={insight_key} insight_value={insight_value}") diff --git a/memory_scope/worker/summary_long/update_profile_worker.py b/memory_scope/worker/summary_long/update_profile_worker.py new file mode 100644 index 00000000..692f2e50 --- /dev/null +++ b/memory_scope/worker/summary_long/update_profile_worker.py @@ -0,0 +1,197 @@ +from typing import List + +from common.response_text_parser import ResponseTextParser +from constants.common_constants import NEW_OBS_NODES, NEW_USER_PROFILE +from enumeration.memory_type_enum import MemoryTypeEnum +from model.memory_wrap_node import MemoryWrapNode +from model.user_attribute import UserAttribute +from worker.bailian.memory_base_worker import MemoryBaseWorker + + +class UpdateProfileWorker(MemoryBaseWorker): + + def filter_obs_nodes(self, + user_attr: UserAttribute, + new_obs_nodes: List[MemoryWrapNode]) -> (UserAttribute, List[MemoryWrapNode], float): + max_score: float = 0 + filtered_nodes: List[MemoryWrapNode] = [] + result = self.rerank_client.call(query=user_attr.description, + documents=[x.memory_node.content for x in new_obs_nodes]) + + if not result: + self.add_run_info(f"update_user_attr={user_attr.memory_key} call rerank failed!") + return user_attr, filtered_nodes, max_score + + # 找到大于阈值的obs node + filtered_nodes: List[MemoryWrapNode] = [] + for rank_node in result: + index = rank_node["index"] + score = rank_node["relevance_score"] + node = new_obs_nodes[index] + keep_flag = "filtered" + if score >= self.config.update_profile_threshold: + filtered_nodes.append(node) + keep_flag = "keep" + max_score = max(max_score, score) + self.logger.info(f"key={user_attr.memory_key} desc={user_attr.description} " + f"content={node.memory_node.content} score={score} keep_flag={keep_flag}") + + if not filtered_nodes: + self.logger.info(f"update_user_attr={user_attr} filtered_nodes is empty!") + return user_attr, filtered_nodes, max_score + + def update_user_attr(self, user_attr: UserAttribute, filtered_nodes: List[MemoryWrapNode]) -> UserAttribute: + self.logger.info(f"update_user_attr memory_key={user_attr.memory_key} desc={user_attr.description} " + f"value={user_attr.value} doc.size={len(filtered_nodes)}") + + # 根据不同的参数类型是否多值，分别给出prompt + user_query_list = [] + for node in filtered_nodes: + user_query_list.append(f"句子：{node.memory_node.content}") + update_profile = f"{user_attr.memory_key}（{user_attr.description}）" + update_profile_value = update_profile + "：" + "，".join(user_attr.value) + + if user_attr.is_unique == 1: + update_profile_message = self.prompt_to_msg( + system_prompt=self.prompt_config.update_unique_profile_system, + few_shot=self.prompt_config.update_unique_profile_few_shot, + user_query=self.prompt_config.update_unique_profile_user_query.format( + user_query="\n".join(user_query_list), + update_profile=update_profile, + update_profile_value=update_profile_value)) + else: + update_profile_message = self.prompt_to_msg( + system_prompt=self.prompt_config.update_plural_profile_system, + few_shot=self.prompt_config.update_plural_profile_few_shot, + user_query=self.prompt_config.update_plural_profile_user_query.format( + user_query="\n".join(user_query_list), + update_profile=update_profile, + update_profile_value=update_profile_value)) + self.logger.info(f"update_profile_message={update_profile_message}") + + # call LLM + response_text: str = self.gene_client.call(messages=update_profile_message, + model_name=self.config.update_profile_model, + max_token=self.config.update_profile_max_token, + temperature=self.config.update_profile_temperature, + top_k=self.config.update_profile_top_k) + + # return if empty + if not response_text: + self.add_run_info(f"update_one_user_attr key={user_attr.memory_key} call llm failed!") + return user_attr + + profile_list = ResponseTextParser(response_text).parse_v1(f"update_attr {user_attr.memory_key}") + if not profile_list: + self.add_run_info(f"update_one_user_attr key={user_attr.memory_key} profile_list empty 1!") + return user_attr + profile_list = profile_list[0] + if not profile_list: + self.add_run_info(f"update_one_user_attr key={user_attr.memory_key} profile_list empty 2") + return user_attr + profile = profile_list[0] + + if not profile or profile in ["无", "重复"]: + self.logger.info(f"profile={profile}, skip.") + return user_attr + + # check 英文中午逗号 + if user_attr.is_unique == 1: + user_attr.value = [profile.strip()] + else: + attr_value_list = profile.replace("，", ",").split(",") + user_attr.value = [x.strip() for x in sorted(list(set(user_attr.value + attr_value_list)))] + return user_attr + + def add_extra_user_attrs(self): + # 解析为空返回 + extra_user_attr_list = [x.strip() for x in self.config.extra_user_attrs if x.strip()] + if not extra_user_attr_list: + return + + for user_attr_info in extra_user_attr_list: + user_attr_split = user_attr_info.split(":") + + # 格式不对返回 + if len(user_attr_split) < 1: + continue + user_attr_key = user_attr_split[0] + + user_attr_desc = "" + if len(user_attr_split) >= 2: + user_attr_desc = user_attr_split[1] + + user_attr_unique = 0 + if len(user_attr_split) >= 3: + user_attr_unique = int(user_attr_split[2]) + + # 已经包含返回 + if user_attr_key in self.user_profile_dict: + user_attr = self.user_profile_dict[user_attr_key] + # description为空，补充description + if not user_attr.description: + user_attr.description = user_attr_desc + continue + + # 增加新属性 + new_attr = UserAttribute(memory_id=self.config.memory_id, + scene=self.scene, + memory_key=user_attr_key, + is_unique=int(user_attr_unique), + is_mutable=1, + memory_type=MemoryTypeEnum.PROFILE, + description=user_attr_desc, + status=1) + self.user_profile_dict[user_attr_key] = new_attr + + def _run(self): + new_obs_nodes: List[MemoryWrapNode] = self.get_context(NEW_OBS_NODES) + if not new_obs_nodes: + self.logger.info("new_obs_nodes is empty, stop user profile!") + self.set_context(NEW_USER_PROFILE, list(self.user_profile_dict.values())) + return + + # 增加环境变量配置的属性 + if self.config.extra_user_attrs: + self.add_extra_user_attrs() + + new_user_profile: List[UserAttribute] = [] + self.set_context(NEW_USER_PROFILE, new_user_profile) + + for user_attr_key, user_attr in self.user_profile_dict.items(): + # 不可修改直接跳过 + if user_attr.is_mutable != 1: + new_user_profile.append(user_attr) + self.logger.info(f"{user_attr_key} is not mutable! continue") + continue + + self.submit_thread(self.filter_obs_nodes, + sleep_time=0.1, + user_attr=user_attr, + new_obs_nodes=new_obs_nodes) + + # 选择topN + result_list = [] + for result in self.join_threads(): + user_attr, filtered_nodes, max_score = result + if not filtered_nodes: + continue + result_list.append(result) + result_sorted = sorted(result_list, key=lambda x: x[2], reverse=True) + if len(result_sorted) > self.config.update_profile_max_thread: + result_sorted = result_sorted[:self.config.update_profile_max_thread] + + # 提交LLM update任务 + for user_attr, filtered_nodes, _ in result_sorted: + self.submit_thread(self.update_user_attr, + sleep_time=1, + user_attr=user_attr, + filtered_nodes=filtered_nodes) + + # collect result & save + for result in self.join_threads(): + if result: + user_attribute: UserAttribute = result + self.logger.info(f"after_update_profile memory_key={user_attribute.memory_key} " + f"desc={user_attribute.description} value={user_attribute.value}") + new_user_profile.append(user_attribute) diff --git a/memory_scope/worker/summary_short/__init__.py b/memory_scope/worker/summary_short/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/memory_scope/worker/summary_short/contra_repeat_worker.py b/memory_scope/worker/summary_short/contra_repeat_worker.py new file mode 100644 index 00000000..36ccf9af --- /dev/null +++ b/memory_scope/worker/summary_short/contra_repeat_worker.py @@ -0,0 +1,92 @@ +from typing import List + +from common.response_text_parser import ResponseTextParser +from constants.common_constants import NEW_OBS_NODES, TODAY_OBS_NODES, MSG_TIME, NEW_OBS_WITH_TIME_NODES, \ + MODIFIED_MEMORIES +from enumeration.memory_node_status import MemoryNodeStatus +from model.memory_wrap_node import MemoryWrapNode +from worker.bailian.memory_base_worker import MemoryBaseWorker + + +class ContraRepeatWorker(MemoryBaseWorker): + + def _run(self): + # 合并当前的obs和今日的obs + new_obs_nodes: List[MemoryWrapNode] = self.get_context(NEW_OBS_NODES) + new_obs_with_time_nodes: List[MemoryWrapNode] = self.get_context(NEW_OBS_WITH_TIME_NODES) + today_obs_nodes: List[MemoryWrapNode] = self.get_context(TODAY_OBS_NODES) + all_obs_nodes: List[MemoryWrapNode] = [] + if new_obs_nodes: + all_obs_nodes.extend(new_obs_nodes) + if new_obs_with_time_nodes: + all_obs_nodes.extend(new_obs_with_time_nodes) + if today_obs_nodes: + all_obs_nodes.extend(today_obs_nodes) + if not all_obs_nodes: + self.add_run_info("all_obs_nodes is empty!") + return + + # gene prompt + user_query_list = [] + all_obs_nodes = sorted(all_obs_nodes, key=lambda x: x.memory_node.metaData.get(MSG_TIME, ""), reverse=True) + for i, n in enumerate(all_obs_nodes): + user_query_list.append(f"{i + 1} {n.memory_node.content}") + merge_obs_message = self.prompt_to_msg( + system_prompt=self.prompt_config.contra_repeat_system.format(num_obs=len(user_query_list)), + few_shot=self.prompt_config.contra_repeat_few_shot, + user_query=self.prompt_config.contra_repeat_user_query.format(user_query="\n".join(user_query_list))) + self.logger.info(f"merge_obs_message={merge_obs_message}") + + # call LLM + response_text = self.gene_client.call(messages=merge_obs_message, + model_name=self.config.merge_obs_model, + max_token=self.config.merge_obs_max_token, + temperature=self.config.merge_obs_temperature, + top_k=self.config.merge_obs_top_k) + + # return if empty + if not response_text: + self.add_run_info("contra repeat call llm failed!") + return + + # parse text + idx_merge_obs_list = ResponseTextParser(response_text).parse_v1("contra_repeat") + if len(idx_merge_obs_list) <= 0: + self.add_run_info("idx_merge_obs_list is empty!") + return + + # add merged obs + merge_obs_nodes: List[MemoryWrapNode] = [] + for obs_content_list in idx_merge_obs_list: + if not obs_content_list: + continue + + # [6, 逃课] + if len(obs_content_list) != 2: + self.logger.warning(f"obs_content_list={obs_content_list} is invalid!") + continue + + idx, keep_flag = obs_content_list + + if not idx.isdigit(): + self.logger.warning(f"idx={idx} is invalid!") + continue + + # 序号需要修正-1 + idx = int(idx) - 1 + if idx >= len(all_obs_nodes): + self.logger.warning(f"idx={idx} is invalid!") + continue + + if keep_flag not in ["矛盾", "被包含", "无"]: + self.logger.warning(f"keep_flag={keep_flag} is invalid!") + continue + + node: MemoryWrapNode = all_obs_nodes[idx] + if keep_flag != "无": + node.memory_node.status = MemoryNodeStatus.EXPIRED.value + merge_obs_nodes.append(node) + self.logger.info(f"after contra repeat: {node.memory_node.content} {node.memory_node.status}") + + # save context + self.set_context(MODIFIED_MEMORIES, merge_obs_nodes) diff --git a/memory_scope/worker/summary_short/get_observation_with_time_worker.py b/memory_scope/worker/summary_short/get_observation_with_time_worker.py new file mode 100644 index 00000000..63c13efe --- /dev/null +++ b/memory_scope/worker/summary_short/get_observation_with_time_worker.py @@ -0,0 +1,128 @@ +from datetime import datetime +from typing import List + +from common.response_text_parser import ResponseTextParser +from common.tool_functions import time_to_formatted_str, get_datetime_info_dict, extract_date_parts +from constants.common_constants import REFLECTED, DT, TIME_INFER, NEW, MSG_TIME, KEY_WORD, DATATIME_WORD_LIST, \ + NEW_OBS_WITH_TIME_NODES +from enumeration.memory_node_status import MemoryNodeStatus +from enumeration.memory_type_enum import MemoryTypeEnum +from model.memory_wrap_node import MemoryWrapNode +from model.message import Message +from worker.bailian.memory_base_worker import MemoryBaseWorker + + +class GetObservationWithTimeWorker(MemoryBaseWorker): + + def add_observation(self, message: Message, obs_content: str, time_infer: str, keywords: str): + created_dt: datetime = datetime.fromtimestamp(float(message.time_created)) + dt = time_to_formatted_str(time=created_dt) + + # 组合meta_data + meta_data = { + MemoryTypeEnum.CONVERSATION.value: message.content, # 原始对话 + REFLECTED: "0", # reflect标记 + DT: dt, # 当天标记 + NEW: "1", # summary-long标记 + MSG_TIME: message.time_created, # 对话时间 + TIME_INFER: time_infer, # 推断的时间 + KEY_WORD: keywords, # 关键词 + } + + # 事件时间 + meta_data.update({f"event_{k}": str(v) for k, v in extract_date_parts(time_infer).items()}) + # 对话时间 + meta_data.update({f"msg_{k}": str(v) for k, v in get_datetime_info_dict(created_dt).items()}) + + return MemoryWrapNode.init_from_attrs(content=obs_content, + memoryId=self.config.memory_id, + timeCreated=message.time_created, + scene=self.scene, + memoryType=MemoryTypeEnum.OBSERVATION.value, + content_modified=True, # 新增的obs需要置为true + metaData=meta_data, + status=MemoryNodeStatus.ACTIVE.value, + tenantId=self.config.tenant_id) + + def _run(self): + # gene prompt + user_query_list = [] + i = 1 + for msg in self.messages: + match = False + for time_keyword in DATATIME_WORD_LIST: + if time_keyword in msg.content: + match = True + break + if match: + dt = time_to_formatted_str(time=msg.time_created, + date_format="", + string_format="{year}年{month}月{day}日{weekday}{hour}点") + user_query_list.append(f"{i} {dt} 用户：{msg.content}") + i += 1 + + if not user_query_list: + self.add_run_info(f"get obs with time user_query_list={user_query_list} is empty") + return + + obtain_obs_message = self.prompt_to_msg( + system_prompt=self.prompt_config.get_observation_with_time_system.format(num_obs=len(user_query_list)), + few_shot=self.prompt_config.get_observation_with_time_few_shot, + user_query=self.prompt_config.get_observation_with_time_user_query.format( + user_query="\n".join(user_query_list))) + self.logger.info(f"obtain_obs_message={obtain_obs_message}") + + # call LLM + response_text: str = self.gene_client.call(messages=obtain_obs_message, + model_name=self.config.summary_messages_model, + max_token=self.config.summary_messages_max_token, + temperature=self.config.summary_messages_temperature, + top_k=self.config.summary_messages_top_k) + + # return if empty + if not response_text: + self.add_run_info("summary call llm failed!", continue_run=False) + return + + # parse text + idx_obs_list = ResponseTextParser(response_text).parse_v1("get_obs_time") + if len(idx_obs_list) <= 0: + self.add_run_info("idx_obs_list is empty!", continue_run=False) + return + + # gene new obs nodes + new_obs_nodes: List[MemoryWrapNode] = [] + for obs_content_list in idx_obs_list: + if not obs_content_list: + continue + + # [1, 2022年6月, 用户在2022年6月去杭州旅游, 旅游] + if len(obs_content_list) != 4: + self.logger.warning(f"obs_content_list={obs_content_list} is invalid!") + continue + + idx, time_infer, obs_content, keywords = obs_content_list + + if obs_content in ["无", "重复"]: + continue + + if not idx.isdigit(): + self.logger.warning(f"idx={idx} is invalid!") + continue + + if time_infer == "无": + time_infer = "" + + # 序号需要修正-1 + idx = int(idx) - 1 + if idx >= len(self.messages): + self.logger.warning(f"idx={idx} is invalid! messages.size={len(self.messages)}") + continue + + new_obs_nodes.append(self.add_observation(message=self.messages[idx], + obs_content=obs_content, + time_infer=time_infer, + keywords=keywords)) + + # save context + self.set_context(NEW_OBS_WITH_TIME_NODES, new_obs_nodes) diff --git a/memory_scope/worker/summary_short/get_observation_worker.py b/memory_scope/worker/summary_short/get_observation_worker.py new file mode 100644 index 00000000..9ed4a75c --- /dev/null +++ b/memory_scope/worker/summary_short/get_observation_worker.py @@ -0,0 +1,116 @@ +from datetime import datetime +from typing import List + +from common.response_text_parser import ResponseTextParser +from common.tool_functions import time_to_formatted_str, get_datetime_info_dict +from constants.common_constants import REFLECTED, DT, NEW_OBS_NODES, TIME_INFER, NEW, MSG_TIME, KEY_WORD, \ + DATATIME_WORD_LIST +from enumeration.memory_node_status import MemoryNodeStatus +from enumeration.memory_type_enum import MemoryTypeEnum +from model.memory_wrap_node import MemoryWrapNode +from model.message import Message +from worker.bailian.memory_base_worker import MemoryBaseWorker + + +class GetObservationWorker(MemoryBaseWorker): + + def add_observation(self, message: Message, obs_content: str, keywords: str): + created_dt: datetime = datetime.fromtimestamp(float(message.time_created)) + dt = time_to_formatted_str(time=created_dt) + + # 组合meta_data + meta_data = { + MemoryTypeEnum.CONVERSATION.value: message.content, # 原始对话 + REFLECTED: "0", # reflect标记 + DT: dt, # 当天标记 + NEW: "1", # summary-long标记 + MSG_TIME: message.time_created, # 对话时间 + TIME_INFER: "", # 推断的时间 + KEY_WORD: keywords, # 关键词 + } + meta_data.update({f"msg_{k}": str(v) for k, v in get_datetime_info_dict(created_dt).items()}) + + return MemoryWrapNode.init_from_attrs(content=obs_content, + memoryId=self.config.memory_id, + timeCreated=message.time_created, + scene=self.scene, + memoryType=MemoryTypeEnum.OBSERVATION.value, + content_modified=True, # 新增的obs需要置为true + metaData=meta_data, + status=MemoryNodeStatus.ACTIVE.value, + tenantId=self.config.tenant_id) + + def _run(self): + # gene prompt + user_query_list = [] + i = 1 + for msg in self.messages: + match = False + for time_keyword in DATATIME_WORD_LIST: + if time_keyword in msg.content: + match = True + break + if not match: + user_query_list.append(f"{i} 用户：{msg.content}") + i += 1 + + if not user_query_list: + self.add_run_info(f"get obs user_query_list={user_query_list} is empty") + return + + obtain_obs_message = self.prompt_to_msg( + system_prompt=self.prompt_config.get_observation_system.format(num_obs=len(user_query_list)), + few_shot=self.prompt_config.get_observation_few_shot, + user_query=self.prompt_config.get_observation_user_query.format(user_query="\n".join(user_query_list))) + self.logger.info(f"obtain_obs_message={obtain_obs_message}") + + # call LLM + response_text: str = self.gene_client.call(messages=obtain_obs_message, + model_name=self.config.summary_messages_model, + max_token=self.config.summary_messages_max_token, + temperature=self.config.summary_messages_temperature, + top_k=self.config.summary_messages_top_k) + + # return if empty + if not response_text: + self.add_run_info("summary call llm failed!", continue_run=False) + return + + # parse text + idx_obs_list = ResponseTextParser(response_text).parse_v1("get obs") + if len(idx_obs_list) <= 0: + self.add_run_info("idx_obs_list is empty!", continue_run=False) + return + + # gene new obs nodes + new_obs_nodes: List[MemoryWrapNode] = [] + for obs_content_list in idx_obs_list: + if not obs_content_list: + continue + + # [1, 2022年6月, 用户在2022年6月去杭州旅游, 旅游] + if len(obs_content_list) != 4: + self.logger.warning(f"obs_content_list={obs_content_list} is invalid!") + continue + + idx, time_infer, obs_content, keywords = obs_content_list + + if obs_content in ["无", "重复"]: + continue + + if not idx.isdigit(): + self.logger.warning(f"idx={idx} is invalid!") + continue + + # 序号需要修正-1 + idx = int(idx) - 1 + if idx >= len(self.messages): + self.logger.warning(f"idx={idx} is invalid! messages.size={len(self.messages)}") + continue + + new_obs_nodes.append(self.add_observation(message=self.messages[idx], + obs_content=obs_content, + keywords=keywords)) + + # save context + self.set_context(NEW_OBS_NODES, new_obs_nodes) diff --git a/memory_scope/worker/summary_short/info_filter_worker.py b/memory_scope/worker/summary_short/info_filter_worker.py new file mode 100644 index 00000000..f369ca4e --- /dev/null +++ b/memory_scope/worker/summary_short/info_filter_worker.py @@ -0,0 +1,57 @@ +from common.response_text_parser import ResponseTextParser +from enumeration.message_role_enum import MessageRoleEnum +from worker.bailian.memory_base_worker import MemoryBaseWorker + + +class InfoFilterWorker(MemoryBaseWorker): + + def _run(self): + # filter user msg + info_messages = [] + for msg in self.messages: + if msg.role != MessageRoleEnum.USER.value: + continue + if len(msg.content) >= self.config.info_filter_msg_max_size: + continue + info_messages.append(msg) + + # gene prompt + user_query = "\n".join([f"{i + 1} 用户：{msg.content}" for i, msg in enumerate(info_messages)]) + info_filter_message = self.prompt_to_msg( + system_prompt=self.prompt_config.info_filter_system.format(batch_size=len(info_messages)), + few_shot=self.prompt_config.info_filter_few_shot, + user_query=self.prompt_config.info_filter_user_query.format(user_query=user_query)) + self.logger.info(f"info_filter_message={info_filter_message}") + + # call llm + response_text = self.gene_client.call(messages=info_filter_message, + model_name=self.config.info_filter_model, + max_token=self.config.info_filter_max_token, + temperature=self.config.info_filter_temperature, + top_k=self.config.info_filter_top_k) + + # return if empty + if not response_text: + self.add_run_info("info score call llm failed!", continue_run=False) + return + + # parse text + info_score_list = ResponseTextParser(response_text).parse_v1("info_filter") + if len(info_score_list) != len(info_messages): + self.add_run_info(f"info_score_size != info_messages_size, " + f"{len(info_score_list)} vs {len(info_messages)}", continue_run=False) + return + + # 过滤value=0的messages + filtered_messages = [] + for msg, info_score in zip(info_messages, info_score_list): + if not info_score: + continue + score = info_score[0] + # if score in ("1", "2",): + if score in ("3",): + msg.info_score = score + filtered_messages.append(msg) + + # 后续不会关注为0的msg，直接丢弃 + self.messages = filtered_messages diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/es_test2.py b/tests/es_test2.py new file mode 100644 index 00000000..25694c0d --- /dev/null +++ b/tests/es_test2.py @@ -0,0 +1,89 @@ +import sys + +sys.path.append("./") + +from common.elastic_search_client import ElasticSearchClient +from common.dash_embedding_client import DashEmbeddingClient +from enumeration.memory_type_enum import MemoryTypeEnum + +from config.bailian_memory_config import BailianMemoryConfig + +api_key: str = "sk-fc77951df1d94418bb5a6cd84da76b17" + +if __name__ == "__main__": + es_index_name: str = "memory_index" + es_user_name: str = "elastic" + es_password: str = "Beilianmemory_" + es_search_top_k = 50 + config = BailianMemoryConfig() + emb_client = DashEmbeddingClient( + request_id="123", + dash_scope_uid="123", + authorization=api_key, + workspace="") + + client = ElasticSearchClient(es_user_name=es_user_name, + es_password=es_password, + es_index_name=es_index_name, + embedding_client=emb_client) + + # result = client.exact_search(100, exact_filters={"status": [MemoryNodeStatus.ACTIVE.value, MemoryNodeStatus.EXPIRED.value]}) + # for k in result[:1]: + # print(k) + # # print(type(k)) + # # print(k["_index"]) + # # print(k["_id"]) + # # print(k["_score"]) + # # print(k["_source"]) + + # result = client.similar_search("可以帮忙准备一些菜吗？", size=100, exact_filters={ + # # "code": "jinli_0530_v2_TONGYI_MAIN_CHAT_profile_音乐偏好", + # # "memoryId": "jinli_0530_v2", + # # "status": MemoryNodeStatus.ACTIVE.value, + # # # "metaData.year": "2024", + # "scene": "TONGYI_MAIN_CHAT".lower(), + # # "memoryType": "profile", + # # "content_modified": True, + # }) + + # query = "我今天出差来深圳君悦酒店了，给张三发个邮件说一下事情" + # result = client.similar_search(text=query, + # size=100, + # exact_filters={ + # "memoryId": "jinli_0607_v11", + # "status": "active", + # "scene": "TONGYI_MAIN_CHAT".lower(), + # "memoryType": [MemoryTypeEnum.OBSERVATION.value, + # MemoryTypeEnum.INSIGHT.value, + # MemoryTypeEnum.OBS_CUSTOMIZED.value], + # + # }, + # wildcard_filters={ + # f"metaData.key_word": ["天气", "工作"], + # }) + + result = client.exact_search_v2(size=100, + term_filters={ + "memoryId": "jinli_0607_v11", + "status": "active", + "scene": "TONGYI_MAIN_CHAT".lower(), + "memoryType": [MemoryTypeEnum.OBSERVATION.value, + MemoryTypeEnum.INSIGHT.value, + MemoryTypeEnum.OBS_CUSTOMIZED.value], + + }, + match_filters={ + f"metaData.key_word": ["天气", "工作"], + } + ) + + for k in result: + # print(json.dumps(k, ensure_ascii=False)) + key_word = k["_source"]["metaData"].get("key_word", "") + content = k["_source"]["content"] + print(content, "||||", key_word) + # print(type(k)) + # print(k["_index"]) + # print(k["_id"]) + # print(k["_score"]) + # print(k["_source"]) diff --git a/tests/test_dash_api.py b/tests/test_dash_api.py new file mode 100644 index 00000000..263044d4 --- /dev/null +++ b/tests/test_dash_api.py @@ -0,0 +1,321 @@ +import sys + + +sys.path.append(".") +from common.dash_embedding_client import DashEmbeddingClient +from common.dash_generate_client import DashGenerateClient +from common.dash_rerank_client import DashReRankClient +from enumeration.env_type import EnvType + +KEY = "sk-fc77951df1d94418bb5a6cd84da76b17" + + +def test_emb(): + client = DashEmbeddingClient(authorization=KEY, request_id="", dash_scope_uid="", workspace="") + result = client.call(text="今天你吃饭了吗？") + print(len(result)) + print(result[:10]) + + +def test_gen(): + client = DashGenerateClient(authorization=KEY, request_id="", dash_scope_uid="", workspace="") + # result = client.call("今天你吃饭了吗？") + messages = [{'role': 'system', 'content': 'You are a helpful assistant.'}, + {'role': 'user', 'content': '今天你吃饭了吗？'}] + result = client.call(messages=messages) + print(result) + + +def test_rerank(): + query = "什么是文本排序模型" + documents = [ + "文本排序模型广泛用于搜索引擎和推荐系统中，它们根据文本相关性对候选文本进行排序", + "量子计算是计算科学的一个前沿领域", + "预训练语言模型的发展给文本排序模型带来了新的进展" + ] + client = DashReRankClient(authorization=KEY, request_id="", dash_scope_uid="", workspace="") + result = client.call(query=query, documents=documents) + print(result) + + +def test_rerank2(): + query = "工作地址" + documents = [ + "我在阿里工作", + ] + client = DashReRankClient(authorization=KEY, request_id="", dash_scope_uid="", workspace="") + result = client.call(query=query, documents=documents) + print(result) + result2 = client.call(query=documents[0], documents=[query]) + print(result2) + + +def test_gen2(): + client = DashGenerateClient(authorization=KEY, request_id="", dash_scope_uid="", workspace="") + # result = client.call("今天你吃饭了吗？") + messages = [{'role': 'system', 'content': 'You are a helpful assistant.'}, + {'role': 'user', 'content': '今天你吃饭了吗？'}] + # print(client.call(messages=messages)) + print(client.call(messages=messages, model_name="deepseek-7b-chat")) + print(client.call(messages=messages, model_name="qwen1.5-7b-chat")) + print(client.call(messages=messages, model_name="qwen1.5-4b-chat")) + print(client.call(messages=messages, model_name="baichuan2-7b-chat-v1")) + print(client.call(messages=messages, model_name="qwen-max")) + print(client.call(messages=messages, model_name="qwen-plus")) + + +def test_gen3(): + messages = [{'role': 'system', + 'content': '任务：从下面每一行的信息中提取出关于用户的可以挖掘的最多1个最重要的用户画像属性，每个<用户画像属性>最多4个字。\n注意：<用户画像属性>可能是一般的用户偏好，也可能是运动偏好，旅游偏好，饮食偏好等等，也可以是重要事件性质，比如最近重要的事情，也可以是一些高度概括的人生理想，价值观，人生观，性格等等。\n要求：根据<用户画像属性>，是可以从下面的信息中提取对应的信息的。\n一定要按如下格式输出，最后的结果一定要加<>:\n<序号> <用户画像属性>'}, + {'role': 'user', + 'content': '\n任务：从下面每一行的信息中提取出关于用户的可以挖掘的最多1个最重要的用户画像属性，每个<用户画像属性>最多4个字。\n注意：<用户画像属性>可能是一般的用户偏好，也可能是运动偏好，旅游偏好，饮食偏好等等，也可以是重要事件性质，比如最近重要的事情，也可以是一些高度概括的人生理想，价值观，人生观，性格等等。\n要求：根据<用户画像属性>，是可以从下面的信息中提取对应的信息的。\n一定要按如下格式输出，最后的结果一定要加<>:\n<序号> <用户画像属性>\n\n示例1\n信息：\n用户想知道明天上海的天气情况。\n用户可能在上海工作，并关心是否需要带伞上班。\n用户在阿里巴巴徐汇滨江园区附近工作。\n用户计划中午在公司附近用餐。\n用户对咖啡因过敏。\n用户喝了咖啡后晚上会出现失眠的情况。\n用户偏好口味较为清淡、不辣的中餐馆。\n用户刚开始了他们的第一份工作。\n用户的工作岗位是阿里巴巴的算法工程师。\n用户希望得到与该岗位相关的职场建议。\n用户面临的问题是在项目进展初期如何有效与上司沟通。\n用户的目标是及时同步项目状态给上司。\n用户希望了解image generation（图像生成）技术的发展概览和最新进展。\n用户对variational auto-encoder、GAN、Diffusion Model等技术及其相互关系感兴趣。\n问题：\n<1> <饮食偏好>\n<2> <技术方向>\n\n示例2\n信息：\n用户想要了解如何使用torchvision库来可视化深度学习任务的进度信息。\n用户希望了解如何将基于numpy和pytorch的并行计算方案迁移到CUDA支持的GPU上运行。\n用户询问是否需要依赖特定的包来完成这一任务。\n用户希望了解如何在Python中自定义进程和线程以实现并行计算。\n用户在编程中遇到了与并行计算相关的问题。\n用户希望学习如何使用Python（numpy，pytorch）在GPU上实现简单的并行计算。\n用户希望了解并行计算的基本概念，包括threads。\n用户询问有关世界各地著名菜系的信息。\n用户对全球各地的美食非常感兴趣。\n用户在寻求有关推拿按摩手法的教程或相关网站推荐。\n用户希望系统地学习正规的推拿按摩手法。\n用户对按摩感兴趣，并且经常去推拿按摩店。\n用户想了解自己在静息状态下一小时大概会消耗多少大卡热量。\n用户年龄为28岁。\n用户体重为70kg。\n用户是男性。\n用户关心其体重与运动消耗的额外热量及心率之间的关系。\n用户在询问为了实现这一目标，每天需要额外消耗多少大卡热量。\n用户希望每月减重1kg。\n用户希望得到类似战略类手机游戏的推荐。\n用户喜欢玩三国志系列、文明系列、全面战争、骑马与砍杀等战略类游戏。\n用户希望根据他们的喜好获得新的游戏推荐。\n用户列举了他们喜欢的具体游戏类型，包括：三国志系列、文明系列、全面战争、骑马与砍杀等。\n用户喜欢玩战略类游戏。\n问题：\n<1> <游戏偏好>\n<2> <运动计划>\n<3> <技术方向>\n\n示例3\n信息：\n用户寻求推荐一个相关课程或网址以进行学习。\n用户计划去青岛旅游。\n用户正为张三的女儿选购生日礼物。\n用户请求为一位名叫张三的人的女儿撰写一段温馨的祝福语。\n用户的同事名叫张三。\n用户与张三约定讨论阿里云百炼项目。\n用户与同事张三讨论了该项目的PRD（产品需求文档）。\n同事张三计划下周对PRD进行最终确定。\n张三还安排了在再下一周进行POC（Proof of Concept，概念验证）的讨论。\n用户希望获知该项目工程开发工作的负责团队信息，以了解项目执行的组织架构与分工情况。\n问题：\n<1> <张三关系>\n\n信息：\n用户对策略游戏感兴趣，希望寻找新的挑战。\n用户近期感到工作压力大，寻求放松方法。\n用户在上海有几位常聚的朋友。\n用户考虑更换工作，关注上海哪些区的工作机会较多。\n用户热衷于尝试新美食，求推荐美食应用。\n用户喜欢自己烹饪，需要海鲜菜谱推荐。\n用户想了解维持广泛社交关系的方法。\n问题：\n'}] + prompt = """ +任务：从下面每一行的信息中提取出关于用户的可以挖掘的最多1个最重要的用户画像属性，每个<用户画像属性>最多4个字。 +注意：<用户画像属性>可能是一般的用户偏好，也可能是运动偏好，旅游偏好，饮食偏好等等，也可以是重要事件性质，比如最近重要的事情，也可以是一些高度概括的人生理想，价值观，人生观，性格等等。 +要求：根据<用户画像属性>，是可以从下面的信息中提取对应的信息的。 +一定要按如下格式输出，最后的结果一定要加<>: +<序号> <用户画像属性> + +示例1 +信息： +用户想知道明天上海的天气情况。 +用户可能在上海工作，并关心是否需要带伞上班。 +用户在阿里巴巴徐汇滨江园区附近工作。 +用户计划中午在公司附近用餐。 +用户对咖啡因过敏。 +用户喝了咖啡后晚上会出现失眠的情况。 +用户偏好口味较为清淡、不辣的中餐馆。 +用户刚开始了他们的第一份工作。 +用户的工作岗位是阿里巴巴的算法工程师。 +用户希望得到与该岗位相关的职场建议。 +用户面临的问题是在项目进展初期如何有效与上司沟通。 +用户的目标是及时同步项目状态给上司。 +用户希望了解image generation（图像生成）技术的发展概览和最新进展。 +用户对variational auto-encoder、GAN、Diffusion Model等技术及其相互关系感兴趣。 +问题： +<1> <饮食偏好> +<2> <技术方向> +示例2 +信息： +用户想要了解如何使用torchvision库来可视化深度学习任务的进度信息。 +用户希望了解如何将基于numpy和pytorch的并行计算方案迁移到CUDA支持的GPU上运行。 +用户询问是否需要依赖特定的包来完成这一任务。 +用户希望了解如何在Python中自定义进程和线程以实现并行计算。 +用户在编程中遇到了与并行计算相关的问题。 +用户希望学习如何使用Python（numpy，pytorch）在GPU上实现简单的并行计算。 +用户希望了解并行计算的基本概念，包括threads。 +用户询问有关世界各地著名菜系的信息。 +用户对全球各地的美食非常感兴趣。 +用户在寻求有关推拿按摩手法的教程或相关网站推荐。 +用户希望系统地学习正规的推拿按摩手法。 +用户对按摩感兴趣，并且经常去推拿按摩店。 +用户想了解自己在静息状态下一小时大概会消耗多少大卡热量。 +用户年龄为28岁。 +用户体重为70kg。 +用户是男性。 +用户关心其体重与运动消耗的额外热量及心率之间的关系。 +用户在询问为了实现这一目标，每天需要额外消耗多少大卡热量。 +用户希望每月减重1kg。 +用户希望得到类似战略类手机游戏的推荐。 +用户喜欢玩三国志系列、文明系列、全面战争、骑马与砍杀等战略类游戏。 +用户希望根据他们的喜好获得新的游戏推荐。 +用户列举了他们喜欢的具体游戏类型，包括：三国志系列、文明系列、全面战争、骑马与砍杀等。 +用户喜欢玩战略类游戏。 +问题： +<1> <游戏偏好> +<2> <运动计划> +<3> <技术方向> +示例3 +信息： +用户寻求推荐一个相关课程或网址以进行学习。 +用户计划去青岛旅游。 +用户正为张三的女儿选购生日礼物。 +用户请求为一位名叫张三的人的女儿撰写一段温馨的祝福语。 +用户的同事名叫张三。 +用户与张三约定讨论阿里云百炼项目。 +用户与同事张三讨论了该项目的PRD（产品需求文档）。 +同事张三计划下周对PRD进行最终确定。 +张三还安排了在再下一周进行POC（Proof of Concept，概念验证）的讨论。 +用户希望获知该项目工程开发工作的负责团队信息，以了解项目执行的组织架构与分工情况。 +问题： +<1> <张三关系> +信息： +用户对策略游戏感兴趣，希望寻找新的挑战。 +用户近期感到工作压力大，寻求放松方法。 +用户在上海有几位常聚的朋友。 +用户考虑更换工作，关注上海哪些区的工作机会较多。 +用户热衷于尝试新美食，求推荐美食应用。 +用户喜欢自己烹饪，需要海鲜菜谱推荐。 +用户想了解维持广泛社交关系的方法。 +问题： + """ + client = DashGenerateClient(authorization=KEY, request_id="", dash_scope_uid="", workspace="") + result = client.call(prompt=prompt.strip(), model_name="qwen-max", seed=0, top_k=1, + repetition_penalty=10) # seed=10, repetition_penalty=0.001 + # print(client.call(prompt=content, model_name="qwen-plus")) + print(result) + + +def test_rerank3(): + """ +用户喜欢在家做饭，需海鲜菜谱推荐。 score=0.33822810090097916 +2024-05-31 16:49:19 INFO jinli_05 semantic_rerank_worker:67] content=用户喜欢在家做饭，求推荐海鲜菜谱。 score=0.3285956750439506 +2024-05-31 16:49:19 INFO jinli_05 semantic_rerank_worker:67] content=用户喜欢烹饪，特别是寻找海鲜菜谱。 score=0.2792495470520024 +2024-05-31 16:49:19 INFO jinli_05 semantic_rerank_worker:67] content=用户热爱烹饪新鲜海鲜，积极寻找美食应用和菜谱。 score=0.19372969292115966 +2024-05-31 16:49:19 INFO jinli_05 semantic_rerank_worker:67] content=用户喜爱探索新美食与在家烹饪海鲜。 score=0.1566147836382108 +2024-05-31 16:49:19 INFO jinli_05 semantic_rerank_worker:67] content=用户拥有几位要好朋友，常共同外出就餐。 score=0.14541493132600777 +2024-05-31 16:49:19 INFO jinli_05 semantic_rerank_worker:67] content=用户有几位常聚餐的好友。 score=0.11955426943198397 +2024-05-31 16:49:19 INFO jinli_05 semantic_rerank_worker:67] content=用户关注上海生活信息，包括寻找新鲜海鲜地点及询问工作机会多的区域。 score=0.08339651657739067 +2024-05-31 16:49:19 INFO jinli_05 semantic_rerank_worker:67] content=用户饮食偏好 (喜欢吃什么菜): 海鲜） score=0.08032847382954839 + + + """ + query = "用户喜欢吃什么菜" + # query = "运动" + documents = [ + "用户喜欢在家做饭，求推荐海鲜菜谱。", + "用户饮食偏好 (喜欢吃什么菜): 海鲜）。", + "用户好奇打篮球是否能促进身高增长。", + "用户感到在上海的工作压力大，寻求放松方法。", + ] + client = DashReRankClient(authorization=KEY, request_id="", dash_scope_uid="", workspace="") + result = client.call(query=query, documents=documents) + print(result) + # result1 = client.call(query=documents[0], documents=[query]) + # result2 = client.call(query=documents[1], documents=[query]) + # result3 = client.call(query=documents[2], documents=[query]) + # result4 = client.call(query=documents[3], documents=[query]) + # print(result1) + # print(result2) + # print(result3) + # print(result4) + + +def test_gen4(): + messages = [{'role': 'system', + 'content': """ +任务：从下面的信息中提取出关于用户的可以挖掘的最多3个重要的用户属性，每个用户属性最多4个字。 +注意：用户属性可能是一般的用户偏好，也可能是运动偏好，旅游偏好，饮食偏好等等，也可以是重要事件性质，比如最近重要的事情，也可以是一些高度概括的人生理想，价值观，人生观，性格等等。 +要求：根据用户属性，我们可以生成“用户的<用户属性>是什么？”的问题，以此从下面的信息中提取用户属性对应的值。 +每一行输出一个<用户属性>: +<用户属性> + """.strip() + }, + {'role': 'user', + 'content': """ +示例1 +信息： +用户想知道明天上海的天气情况。 +用户可能在上海工作，并关心是否需要带伞上班。 +用户在阿里巴巴徐汇滨江园区附近工作。 +用户计划中午在公司附近用餐。 +用户对咖啡因过敏。 +用户喝了咖啡后晚上会出现失眠的情况。 +用户偏好口味较为清淡、不辣的中餐馆。 +用户刚开始了他们的第一份工作。 +用户的工作岗位是阿里巴巴的算法工程师。 +用户希望得到与该岗位相关的职场建议。 +用户面临的问题是在项目进展初期如何有效与上司沟通。 +用户的目标是及时同步项目状态给上司。 +用户希望了解image generation（图像生成）技术的发展概览和最新进展。 +用户对variational auto-encoder、GAN、Diffusion Model等技术及其相互关系感兴趣。 +问题： +饮食偏好 +技术方向 + +示例2 +信息： +用户想要了解如何使用torchvision库来可视化深度学习任务的进度信息。 +用户希望了解如何将基于numpy和pytorch的并行计算方案迁移到CUDA支持的GPU上运行。 +用户询问是否需要依赖特定的包来完成这一任务。 +用户希望了解如何在Python中自定义进程和线程以实现并行计算。 +用户在编程中遇到了与并行计算相关的问题。 +用户希望学习如何使用Python（numpy，pytorch）在GPU上实现简单的并行计算。 +用户希望了解并行计算的基本概念，包括threads。 +用户询问有关世界各地著名菜系的信息。 +用户对全球各地的美食非常感兴趣。 +用户在寻求有关推拿按摩手法的教程或相关网站推荐。 +用户希望系统地学习正规的推拿按摩手法。 +用户对按摩感兴趣，并且经常去推拿按摩店。 +用户想了解自己在静息状态下一小时大概会消耗多少大卡热量。 +用户年龄为28岁。 +用户体重为70kg。 +用户是男性。 +用户关心其体重与运动消耗的额外热量及心率之间的关系。 +用户在询问为了实现这一目标，每天需要额外消耗多少大卡热量。 +用户希望每月减重1kg。 +用户希望得到类似战略类手机游戏的推荐。 +用户喜欢玩三国志系列、文明系列、全面战争、骑马与砍杀等战略类游戏。 +用户希望根据他们的喜好获得新的游戏推荐。 +用户列举了他们喜欢的具体游戏类型，包括：三国志系列、文明系列、全面战争、骑马与砍杀等。 +用户喜欢玩战略类游戏。 +问题： +游戏偏好 +运动计划 +技术方向 + +示例3 +信息： +用户寻求推荐一个相关课程或网址以进行学习。 +用户计划去青岛旅游。 +用户正为张三的女儿选购生日礼物。 +用户请求为一位名叫张三的人的女儿撰写一段温馨的祝福语。 +用户的同事名叫张三。 +用户与张三约定讨论阿里云百炼项目。 +用户与同事张三讨论了该项目的PRD（产品需求文档）。 +同事张三计划下周对PRD进行最终确定。 +张三还安排了在再下一周进行POC（Proof of Concept，概念验证）的讨论。 +用户希望获知该项目工程开发工作的负责团队信息，以了解项目执行的组织架构与分工情况。 +问题： +朋友关系 + +任务：从下面的信息中提取出关于用户的可以挖掘的最多3个重要的用户属性，每个用户属性最多4个字。 +注意：用户属性可能是一般的用户偏好，也可能是运动偏好，旅游偏好，饮食偏好等等，也可以是重要事件性质，比如最近重要的事情，也可以是一些高度概括的人生理想，价值观，人生观，性格等等。 +要求：根据用户属性，我们可以生成“用户的<用户属性>是什么？”的问题，以此从下面的信息中提取用户属性对应的值。 +每一行输出一个<用户属性>: +<用户属性> + +信息： +用户想知道上海哪里的海鲜最新鲜，表明用户在上海生活或访问，并对食物品质有要求。 +用户寻找策略游戏推荐，显示出对策略类游戏的兴趣和寻求新挑战的愿望。 +用户提到在上海的工作压力大，寻求放松建议，反映了其当前的生活压力状态和对减压方法的需求。 +用户拥有常一起吃饭的好友，强调了其社交活动和对友谊的重视。 +用户考虑更换工作，关注上海哪些区工作机会多，表明职业规划上的变动意向。 +用户喜欢尝试新美食并询问美食应用推荐，再次强调对美食的兴趣和探索欲。 +用户提到了自己烹饪的兴趣，特别是对海鲜菜谱的需求，细化了其个人爱好。 +用户询问维持广泛社交关系的方法，显示其对社交网络维护的关注。 +问题： + """.strip() + }] + client = DashGenerateClient(authorization=KEY, request_id="", dash_scope_uid="", workspace="") + # result = client.call(messages=messages, model_name="qwen-long", seed=0, top_k=1) + result = client.call(messages=messages, model_name="qwen-max", seed=0, top_k=1) + # result = client.call(messages=messages, model_name="qwen-max", seed=0) + + # seed=10, repetition_penalty=0.001 + print(result) + + +def test_gen_time(): + prompt = "任务指令：从语句与语句发生的时间，推断并提取语句内容中指向的时间段。回答尽可能完整的时间段。\n语句：好像是前天有人来找过你。\n时间：2074年12月7日，2074年第49周，周三，16时30分0秒。\n回答：" + # 1656375133437235, 197291 + client = DashGenerateClient(authorization="sk-AdrklI1sWM", request_id="", + dash_scope_uid="", workspace="", env_type=EnvType.DAILY) + # result = client.call(messages=messages, model_name="qwen-long", seed=0, top_k=1) + result = client.call(prompt=prompt, model_name="qwen_1_8_parse_time_service", seed=0, top_k=1) + # result = client.call(messages=messages, model_name="qwen-max", seed=0) + + # seed=10, repetition_penalty=0.001 + print(result) + + +if __name__ == "__main__": + # test_emb() + # test_gen() + # test_rerank() + # test_rerank2() + # test_gen2() + # test_gen3() + # test_gen4() + # test_rerank3() + test_gen_time() diff --git a/tests/test_memory.py b/tests/test_memory.py new file mode 100644 index 00000000..24d263a1 --- /dev/null +++ b/tests/test_memory.py @@ -0,0 +1,256 @@ +import datetime +import os +import time +from typing import List, Dict + +from common.logger import Logger +from constants.common_constants import NEW_USER_PROFILE, MODIFIED_MEMORIES, RELATED_MEMORIES +from enumeration.memory_method_enum import MemoryMethodEnum +from model.memory_node import MemoryNode +from model.user_attribute import UserAttribute +from request.memory import MemoryServiceRequestModel +from service.memory_service_bailian import MemoryServiceBailian + +""" +任务：随机生成一个用户的画像，随机种子0，并根据用户的画像虚拟一段用户和AI的对话。 +步骤： +1. 帮忙生成一个用户的画像：包括用户的姓名，性别，工作地点，朋友关系，饮食偏好，游戏偏好，运动偏好等等属性。 +2. 根据用户的画像虚拟一段用户和AI（比如通义千问）的对话，可以是用户的生活轨迹，生活事件，工作事件，也可以是用户的一些看法等等。要求对话中需要包含用户画像信息，并可以通过对话反推出部分用户画像。 +用户画像格式，最少10个用户画像： +<用户画像，例如性别>: <属性值，例如男性> +用户对话格式，最少20轮对话： +<轮次> <用户>：<用户问题> +<轮次> ：<回答> + """ + +os.environ["APP_ENV"] = "daily" +os.environ["memory_retrieve_pipeline"] = """ +parse_params,es.load_profile,[retrieve.extract_time|es.es_similar|es.es_keyword],retrieve.semantic_rank,retrieve.fuse_rerank +""".strip() +os.environ["memory_summary_short_pipeline"] = """ +parse_params,summary_short.info_filter,[es.es_today_obs|summary_short.get_observation|summary_short.get_observation_with_time],summary_short.contra_repeat,memory_store +""".strip() +os.environ["memory_summary_long_pipeline"] = """ +parse_params,[es.load_profile|es.es_new_obs|es.es_insight],[summary_long.update_insight|summary_long.get_reflection,summary_long.get_insight|summary_long.update_profile],summary_long.summary_collect,memory_store +""".strip() + +os.environ["memory_summary_long_reflect_obs_cnt_threshold"] = "5" +os.environ["memory_summary_long_max_workers"] = "5" + +# attrs = { +# "性别": ["男性或者女性", 1], +# "工作地点": ["工作所在城市", 1], +# "朋友关系": ["和谁是什么朋友", 0], +# "饮食偏好": ["喜欢吃什么菜", 0], +# "游戏偏好": ["喜欢玩什么游戏", 0], +# "运动偏好": ["喜欢什么运动", 0], +# "音乐偏好": ["喜欢听什么音乐", 0], +# "电影类型偏好": ["喜欢看什么类型的电影", 0], +# "阅读偏好": ["喜欢看什么书", 0], +# "购物习惯": ["喜欢买什么东西", 0], +# } +# +# os.environ["memory_summary_extra_user_attrs_TONGYI_MAIN_CHAT"] = ",".join( +# [f"{k}:{v[0]}:{v[1]}" for k, v in attrs.items()]) + +memory_id: str = "jinli_0607_v26" +workspace_id: str = "" +api_key: str = "sk-AdrklI1sWM" +scene: str = "TONGYI_MAIN_CHAT" +algo_version: str = "" +output_max_count: int = 3 +""" +'工作地点:工作所在城市:1', +'所在地:当前所在地点:1', +'饮食偏好:喜欢吃什么菜:0', +'游戏偏好:喜欢玩什么游戏:0', +'运动偏好:喜欢什么运动:0', +""" +user_profile: List[UserAttribute] = [ + UserAttribute(memory_key="运动偏好", value=["足球"], description="喜欢什么运动", is_unique=0), + UserAttribute(memory_key="工作地点", description="工作所在城市", is_unique=1), + UserAttribute(memory_key="所在地", description="当前所在地点", is_unique=1), + UserAttribute(memory_key="饮食偏好", description="喜欢吃什么菜", is_unique=0), + UserAttribute(memory_key="游戏偏好", description="喜欢玩什么游戏", is_unique=0), + UserAttribute(memory_key="运动偏好", description="喜欢什么运动", is_unique=0), +] +ext_info: Dict[str, str] = {} +trace_id: str = "jinli_0530_req_id" +request_id: str = "jinli_0530_req_id" +account_id: str = "jinli" +app_id: str = "jinli_id" +uid: str = "1656375133437235" + +messages1 = [ + {"role": "user", "content": "你知道北京哪里的海鲜最新鲜吗？", "time_created": "1717037394"}, + {"role": "user", "content": "有没有推荐的策略游戏？最近想找新的挑战。", "time_created": "1717037404"}, + {"role": "user", "content": "听说篮球运动对身体很好，是真的吗？", "time_created": "1717037414"}, + {"role": "user", "content": "最近在北京的工作压力太大，有什么放松的建议吗？", "time_created": "1717037424"}, + {"role": "user", "content": "说到朋友，我确实有几位很要好的朋友，我们经常一起出去吃饭。", + "time_created": "1717037434"}, + {"role": "user", "content": "对了，最近想换工作，你觉得北京的哪个区工作机会更多？", "time_created": "1717037444"}, + {"role": "user", "content": "听你这么说，我感觉挺有信心的，谢了！", "time_created": "1717037454"}, + {"role": "user", "content": "我很喜欢尝试新的美食，有没有推荐的美食应用？", "time_created": "1717037464"}, + {"role": "user", "content": "我有时也喜欢自己在家做饭，你有没有好的海鲜菜谱推荐？", "time_created": "1717037474"}, + {"role": "user", "content": "听说打篮球可以长高，这是真的吗？", "time_created": "1717037484"}, + {"role": "user", "content": "昨天是我的生日！", "time_created": "1717037494"}, + {"role": "user", "content": "昨天和同学一起在我家开了party，庆祝了我的生日！", "time_created": "1717037494"}, + {"role": "user", "content": "我在北京阿里云园区工作", "time_created": "1717037504"}, + {"role": "user", "content": "我是阿里云百炼的工程师", "time_created": "1717037504"}, + {"role": "user", "content": "最后一个问题，你知道怎么才能维持广泛的社交关系吗？", "time_created": "1717037504"}, + +] +dt_n = datetime.datetime(year=2024, month=6, day=1, hour=12) +ts = int(dt_n.timestamp()) +for i, msg in enumerate(messages1): + msg["time_created"] = str(ts + i * 10) + +messages2 = [ + # {"role": "user", "content": "今天我和客户团队的工程师张三讨论了技术方案，聊得很愉快", + # "time_created": "1717037394"}, + # {"role": "user", "content": "帮我记一下，我和他沟通约定3天后到杭州上门提供技术解决方案", + # "time_created": "1717037394"}, + {"role": "user", "content": "帮我记一下，我和客户团队的工程师张三沟通约定3天后到杭州上门提供技术解决方案", + "time_created": "1717037394"}, +] +dt_n = datetime.datetime(year=2024, month=6, day=3, hour=12) +ts = int(dt_n.timestamp()) +for i, msg in enumerate(messages2): + msg["time_created"] = str(ts + i * 10) + +messages3 = [ + {"role": "user", "content": "我今天出差来深圳君悦酒店了，给张三发个邮件说一下事情", "time_created": "1717037394"}, + {"role": "user", "content": "我最近肠胃不好，吃不了辣", "time_created": "1717037394"}, + {"role": "user", "content": "最近肠胃养好了，换一些川菜吧", "time_created": "1717037394"}, +] +dt_n = datetime.datetime(year=2024, month=6, day=6, hour=12) +ts = int(dt_n.timestamp()) +for i, msg in enumerate(messages3): + msg["time_created"] = str(ts + i * 10) + +messages4 = [ + {"role": "user", "content": "附近有什么好吃的", "time_created": "1717037394"}, + {"role": "user", "content": "今天天气怎么样？", "time_created": "1717037394"}, +] +dt_n = datetime.datetime(year=2024, month=6, day=6, hour=13) +ts = int(dt_n.timestamp()) +for i, msg in enumerate(messages3): + msg["time_created"] = str(ts + i * 10) + + +def summary_short(messages): + messages_pick_n = len(messages) + request: MemoryServiceRequestModel = MemoryServiceRequestModel( + messages=messages, + messages_pick_n=messages_pick_n, + memory_id=memory_id, + workspace_id=workspace_id, + api_key=api_key, + scene=scene, + algo_version=algo_version, + output_max_count=output_max_count, + user_profile=user_profile, + ext_info=ext_info, + trace_id=trace_id, + tenant_id=trace_id, + request_id=request_id, + account_id=account_id, + app_id=app_id, + uid=uid, + ) + + logger = Logger.get_memory_logger() + logger.set_trace_id(request.trace_id) + memory_service = MemoryServiceBailian(request, method=MemoryMethodEnum.SUMMARY_SHORT) + memory_service.run() + + modified_memories: List[MemoryNode] = memory_service.get_context(MODIFIED_MEMORIES) + return modified_memories + + +def summary_long(): + request: MemoryServiceRequestModel = MemoryServiceRequestModel( + messages=[], + messages_pick_n=0, + memory_id=memory_id, + workspace_id=workspace_id, + api_key=api_key, + scene=scene, + algo_version=algo_version, + output_max_count=output_max_count, + user_profile=user_profile, + ext_info=ext_info, + trace_id=trace_id, + tenant_id=trace_id, + request_id=request_id, + account_id=account_id, + app_id=app_id, + uid=uid, + ) + logger = Logger.get_memory_logger() + logger.set_trace_id(request.trace_id) + memory_service = MemoryServiceBailian(request, method=MemoryMethodEnum.SUMMARY_LONG) + memory_service.run() + + user_profiles: List[UserAttribute] = memory_service.get_context(NEW_USER_PROFILE) + modified_memories: List[MemoryNode] = memory_service.get_context(MODIFIED_MEMORIES) + # ext_infos = memory_service.get_context(RESPONSE_EXT_INFO) + # logger.info(f"user_profile=\n{json.dumps([x.model_dump() for x in user_profiles], ensure_ascii=False)}") + # logger.info(f"modified_memories=\n{json.dumps([x.model_dump() for x in modified_memories], ensure_ascii=False)}") + # logger.info(f"ext_info=\n{json.dumps(ext_infos, ensure_ascii=False)}") + return user_profiles, modified_memories + + +def retrieve(messages): + request: MemoryServiceRequestModel = MemoryServiceRequestModel( + messages=messages, + messages_pick_n=1, + memory_id=memory_id, + workspace_id=workspace_id, + api_key=api_key, + scene=scene, + algo_version=algo_version, + output_max_count=output_max_count, + user_profile=user_profile, + ext_info=ext_info, + trace_id=trace_id, + tenant_id=trace_id, + request_id=request_id, + account_id=account_id, + app_id=app_id, + uid=uid, + ) + + logger = Logger.get_memory_logger() + logger.set_trace_id(request.trace_id) + memory_service = MemoryServiceBailian(request, method=MemoryMethodEnum.RETRIEVE) + memory_service.run() + + modified_memories: List[str] = memory_service.get_context(RELATED_MEMORIES) + return modified_memories + + +if __name__ == "__main__": + logger = Logger.get_memory_logger() + + summary1 = summary_short(messages1) + logger.info(f"summary1={summary1}") + + summary2 = summary_short(messages2) + logger.info(f"summary2={summary2}") + + for i, msg in enumerate(messages3): + time.sleep(6) + retrieve_res = retrieve([msg]) + logger.info(f"index={i} retrieve_res={retrieve_res}") + + summary_res = summary_short([msg]) + logger.info(f"index={i} summary_res={summary_res}") + + summary3 = summary_long() + logger.info(f"summary3={summary3}") + + time.sleep(6) + for i, msg in enumerate(messages4): + retrieve_res = retrieve([msg]) + logger.info(f"index={i} retrieve_res={retrieve_res}")