mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-06 08:16:00 +00:00
72 lines
2.1 KiB
Python
72 lines
2.1 KiB
Python
import json
|
|
import os.path
|
|
from typing import Dict
|
|
|
|
import yaml
|
|
|
|
from memory_scope.utils.global_context import G_CONTEXT
|
|
|
|
|
|
class PromptHandler(object):
|
|
|
|
def __init__(self, class_path: str, prompt_file: str = "", prompt_dict: dict = None, **kwargs):
|
|
self._class_path: str = class_path
|
|
self._prompt_dict: Dict[str, str] = {}
|
|
self.kwargs = kwargs
|
|
|
|
file_path = self._class_path.strip(".py")
|
|
self.add_prompt_file(file_path)
|
|
|
|
if prompt_file:
|
|
self.add_prompt_file(prompt_file)
|
|
|
|
if prompt_dict:
|
|
self.add_prompt_dict(prompt_dict)
|
|
|
|
@staticmethod
|
|
def file_path_completion(file_path: str) -> str:
|
|
if file_path.endswith(".yaml") or file_path.endswith(".json"):
|
|
return file_path
|
|
|
|
if os.path.exists(f"{file_path}.yaml"):
|
|
return f"{file_path}.yaml"
|
|
|
|
if os.path.exists(f"{file_path}.json"):
|
|
return f"{file_path}.json"
|
|
|
|
raise RuntimeError(f"{file_path}/yaml/json is not exists!")
|
|
|
|
def add_prompt_file(self, file_path: str):
|
|
file_path = self.file_path_completion(file_path)
|
|
|
|
prompt_dict = {}
|
|
|
|
if file_path.endswith(".yaml"):
|
|
with open(file_path) as f:
|
|
prompt_dict = yaml.load(f, yaml.FullLoader)
|
|
|
|
elif file_path.endswith(".json"):
|
|
with open(f"{file_path}.json") as f:
|
|
prompt_dict = json.load(f)
|
|
|
|
self.add_prompt_dict(prompt_dict)
|
|
|
|
def add_prompt_dict(self, prompt_dict: dict):
|
|
for key, language_dict in prompt_dict.items():
|
|
prompts = language_dict.get(G_CONTEXT.language)
|
|
if not prompts:
|
|
raise RuntimeError(f"{key}.prompt.{G_CONTEXT.language} is empty!")
|
|
self._prompt_dict[key] = prompts
|
|
|
|
@property
|
|
def prompt_dict(self):
|
|
return self._prompt_dict
|
|
|
|
def __getitem__(self, key: str):
|
|
return self._prompt_dict[key]
|
|
|
|
def __setitem__(self, key: str, value: str):
|
|
self._prompt_dict[key] = value
|
|
|
|
def __getattr__(self, key: str):
|
|
return self._prompt_dict[key]
|