mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-06 08:16:00 +00:00
219 lines
8.3 KiB
Python
219 lines
8.3 KiB
Python
import os
|
|
import time
|
|
from typing import List
|
|
|
|
import questionary
|
|
|
|
from memory_scope.chat.base_memory_chat import BaseMemoryChat
|
|
from memory_scope.constants.language_constants import DEFAULT_HUMAN_NAME
|
|
from memory_scope.enumeration.message_role_enum import MessageRoleEnum
|
|
from memory_scope.memory.service.base_memory_service import BaseMemoryService
|
|
from memory_scope.models.base_model import BaseModel
|
|
from memory_scope.scheme.message import Message
|
|
from memory_scope.scheme.model_response import ModelResponse, ModelResponseGen
|
|
from memory_scope.utils.global_context import G_CONTEXT
|
|
from memory_scope.utils.logger import Logger
|
|
from memory_scope.utils.prompt_handler import PromptHandler
|
|
from memory_scope.utils.tool_functions import char_logo
|
|
|
|
|
|
class CliMemoryChat(BaseMemoryChat):
|
|
USER_COMMANDS = {
|
|
"exit": "exit the CLI",
|
|
"clear": "clear commands",
|
|
"help": "get cli commands help",
|
|
"stream": "get stream response"
|
|
}
|
|
|
|
def __init__(self,
|
|
memory_service: str,
|
|
generation_model: str,
|
|
stream: bool = True,
|
|
human_name: str = DEFAULT_HUMAN_NAME[G_CONTEXT.language],
|
|
assistant_name: str = "AI",
|
|
**kwargs):
|
|
|
|
self._memory_service: BaseMemoryService | str = memory_service
|
|
self._generation_model: BaseModel | str = generation_model
|
|
self.stream: bool = stream
|
|
self.human_name: str = human_name
|
|
self.assistant_name: str = assistant_name
|
|
self.kwargs: dict = kwargs
|
|
|
|
self._logo = char_logo("MemoryScope")
|
|
self._prompt_handler: PromptHandler | None = None
|
|
G_CONTEXT.meta_data.update({
|
|
"human_name": human_name,
|
|
"assistant_name": assistant_name,
|
|
})
|
|
|
|
self.logger = Logger.get_logger()
|
|
|
|
@property
|
|
def prompt_handler(self) -> PromptHandler:
|
|
if self._prompt_handler is None:
|
|
self._prompt_handler = PromptHandler(__file__, **self.kwargs)
|
|
return self._prompt_handler
|
|
|
|
def print_logo(self):
|
|
for line in self._logo:
|
|
print(line)
|
|
|
|
@property
|
|
def memory_service(self) -> BaseMemoryService:
|
|
if isinstance(self._memory_service, str):
|
|
if self._memory_service not in G_CONTEXT.memory_service_dict:
|
|
raise ValueError("Missing declaration of memory_service in yaml configuration: " + self._memory_service)
|
|
self._memory_service = G_CONTEXT.memory_service_dict[self._memory_service]
|
|
self._memory_service.start_service()
|
|
return self._memory_service
|
|
|
|
@property
|
|
def generation_model(self) -> BaseModel:
|
|
if isinstance(self._generation_model, str):
|
|
if self._generation_model not in G_CONTEXT.model_dict:
|
|
raise ValueError(f"Missing declaration of generation model in yaml config: {self._generation_model}")
|
|
self._generation_model = G_CONTEXT.model_dict[self._generation_model]
|
|
return self._generation_model
|
|
|
|
def chat_with_memory(self, query: str, remember_response: bool = False) -> ModelResponse | ModelResponseGen:
|
|
new_message: Message = Message(role=MessageRoleEnum.USER.value, role_name=self.human_name, content=query)
|
|
self.memory_service.add_messages(new_message)
|
|
|
|
messages: List[Message] = []
|
|
|
|
# add memory to system prompt
|
|
system_prompt = self.prompt_handler.system_prompt
|
|
memories: str = self.memory_service.read_memory()
|
|
if memories:
|
|
memory_prompt = self.prompt_handler.memory_prompt
|
|
system_prompt = "\n".join([x.strip() for x in [system_prompt, memory_prompt, memories]])
|
|
messages.append(Message(role=MessageRoleEnum.SYSTEM, content=system_prompt))
|
|
|
|
# add history messages
|
|
history_messages = self.memory_service.read_message()
|
|
if history_messages:
|
|
messages.extend(history_messages)
|
|
|
|
# add new_message
|
|
messages.append(new_message)
|
|
self.logger.info(f"messages={messages}")
|
|
|
|
# call LLM. in stream mode, return generator. in non-stream mode, return response.
|
|
generated = self.generation_model.call(messages=messages, stream=self.stream)
|
|
|
|
# in non-stream mode, remember the response if user demand to do so.
|
|
if remember_response:
|
|
assert not self.stream
|
|
generated.message.role_name = self.assistant_name
|
|
self.memory_service.add_messages(generated.message)
|
|
|
|
# return response or generator
|
|
return generated
|
|
|
|
@staticmethod
|
|
def parse_query_command(query: str):
|
|
query_split = query.lstrip("/").lower().split(" ")
|
|
command = query_split[0]
|
|
args = query_split[1:]
|
|
kwargs = {}
|
|
for arg in args:
|
|
if not args:
|
|
continue
|
|
arg_split = arg.split("=")
|
|
if len(arg_split) >= 2:
|
|
k = arg_split[0]
|
|
v = arg_split[1]
|
|
if k and v:
|
|
kwargs[k] = v
|
|
return command, kwargs
|
|
|
|
def process_commands(self, query: str) -> bool:
|
|
continue_run = True
|
|
command, kwargs = self.parse_query_command(query)
|
|
|
|
if command == "exit":
|
|
self.memory_service.stop_service()
|
|
continue_run = False
|
|
|
|
elif command == "clear":
|
|
os.system("clear")
|
|
|
|
elif command == "help":
|
|
questionary.print("CLI commands", "bold")
|
|
for cmd, desc in self.USER_COMMANDS.items():
|
|
questionary.print(text=f" /{cmd}:", style="bold")
|
|
questionary.print(text=f" {desc}")
|
|
|
|
elif command == "stream":
|
|
self.stream = not self.stream
|
|
questionary.print(f"set stream: {self.stream}")
|
|
|
|
elif command in self.memory_service.op_description_dict:
|
|
refresh_time = kwargs.pop("refresh_time", "")
|
|
if refresh_time and refresh_time.isdigit():
|
|
refresh_time = int(refresh_time)
|
|
while True:
|
|
result = self.memory_service.do_operation(op_name=command, **kwargs)
|
|
os.system("clear")
|
|
self.print_logo()
|
|
questionary.print(result)
|
|
time.sleep(refresh_time)
|
|
|
|
else:
|
|
result = self.memory_service.do_operation(op_name=command, **kwargs)
|
|
questionary.print(result)
|
|
|
|
else:
|
|
questionary.print(f"Unknown command={command} received.")
|
|
|
|
return continue_run
|
|
|
|
def run(self):
|
|
self.print_logo()
|
|
self.USER_COMMANDS.update(self.memory_service.op_description_dict)
|
|
|
|
while True:
|
|
try:
|
|
query = questionary.text(message=f"{self.human_name}:", multiline=False, qmark=">").unsafe_ask()
|
|
if not query:
|
|
continue
|
|
|
|
query: str = query.strip()
|
|
|
|
# handle cli / commands with memory ops
|
|
if query.startswith("/"):
|
|
if self.process_commands(query=query):
|
|
continue
|
|
else:
|
|
break
|
|
|
|
questionary.print("> ", end="", style="fg:yellow")
|
|
questionary.print(f"{self.assistant_name}: ", end="", style="bold")
|
|
|
|
if self.stream:
|
|
model_response = None
|
|
for model_response in self.chat_with_memory(query=query):
|
|
questionary.print(model_response.delta, end="")
|
|
questionary.print("")
|
|
|
|
else:
|
|
model_response = self.chat_with_memory(query=query)
|
|
questionary.print(model_response.message.content)
|
|
|
|
# add response to memory
|
|
model_response.message.role_name = self.assistant_name
|
|
self.memory_service.add_messages(model_response.message)
|
|
|
|
except KeyboardInterrupt:
|
|
questionary.print("User interrupt occurred.")
|
|
is_exit = questionary.confirm("continue exit").unsafe_ask()
|
|
if is_exit:
|
|
self.memory_service.stop_service()
|
|
break
|
|
|
|
except Exception as e:
|
|
import traceback
|
|
traceback.print_exc()
|
|
self.logger.exception(f"An exception occurred when running cli memory chat. args={e.args}.")
|
|
continue
|