feat: 添加内置提示词

This commit is contained in:
DongShengWu 2026-01-01 21:43:05 +08:00
parent 9ab0089a83
commit b68b30aaa9
2 changed files with 46 additions and 4 deletions

View File

@ -48,11 +48,19 @@ class StreamingLLM:
llm = StreamingLLM() llm = StreamingLLM()
for chunk in llm.chat("你好"): for chunk in llm.chat("你好"):
print(chunk.content, end="") print(chunk.content, end="")
# 使用预设角色
llm = StreamingLLM(role='voice') # 语音交互模式
llm = StreamingLLM(role='coder') # 编程助手
""" """
# 可用角色列表
ROLES = list(config.ROLES.keys())
def __init__( def __init__(
self, self,
model: str = None, model: str = None,
role: str = None,
enable_thinking: bool = None, enable_thinking: bool = None,
temperature: float = None, temperature: float = None,
top_p: float = None, top_p: float = None,
@ -64,20 +72,27 @@ class StreamingLLM:
Args: Args:
model: 模型名称 model: 模型名称
role: 预设角色 ('assistant', 'coder', 'translator', 'writer', 'voice', 'teacher')
enable_thinking: 是否启用思考模式 enable_thinking: 是否启用思考模式
temperature: 温度参数 temperature: 温度参数
top_p: Top P 采样 top_p: Top P 采样
max_tokens: 最大输出 Token max_tokens: 最大输出 Token
system_prompt: 系统提示词 system_prompt: 系统提示词 (优先级高于 role)
""" """
self._load_api_key() self._load_api_key()
self.model = model or config.MODEL self.model = model or config.MODEL
self.role = role or config.DEFAULT_ROLE
self.enable_thinking = enable_thinking if enable_thinking is not None else config.ENABLE_THINKING self.enable_thinking = enable_thinking if enable_thinking is not None else config.ENABLE_THINKING
self.temperature = temperature if temperature is not None else config.TEMPERATURE self.temperature = temperature if temperature is not None else config.TEMPERATURE
self.top_p = top_p if top_p is not None else config.TOP_P self.top_p = top_p if top_p is not None else config.TOP_P
self.max_tokens = max_tokens or config.MAX_TOKENS self.max_tokens = max_tokens or config.MAX_TOKENS
self.system_prompt = system_prompt or config.SYSTEM_PROMPT
# system_prompt 优先,否则使用角色配置
if system_prompt:
self.system_prompt = system_prompt
else:
self.system_prompt = config.ROLES.get(self.role, config.SYSTEM_PROMPT)
# 初始化 OpenAI 客户端 # 初始化 OpenAI 客户端
self._client = OpenAI( self._client = OpenAI(
@ -88,6 +103,19 @@ class StreamingLLM:
# 对话历史 # 对话历史
self._messages: List[Dict[str, str]] = [] self._messages: List[Dict[str, str]] = []
def set_role(self, role: str) -> None:
"""
切换角色
Args:
role: 角色名称
"""
if role in config.ROLES:
self.role = role
self.system_prompt = config.ROLES[role]
else:
raise ValueError(f"未知角色: {role}, 可用角色: {list(config.ROLES.keys())}")
def _load_api_key(self) -> None: def _load_api_key(self) -> None:
"""从 .env 加载 API Key""" """从 .env 加载 API Key"""
current_dir = Path(__file__).parent current_dir = Path(__file__).parent

View File

@ -46,12 +46,26 @@ TOP_P = 0.95
MAX_TOKENS = 4096 MAX_TOKENS = 4096
# ============================================================
# 角色配置
# ============================================================
# 默认角色
DEFAULT_ROLE = 'assistant'
# 预设角色 (角色名 -> 系统提示词)
ROLES = {
# 通用助手
'assistant': """ 你是心镜助手! """,
}
# ============================================================ # ============================================================
# 上下文配置 # 上下文配置
# ============================================================ # ============================================================
# 系统提示词 # 系统提示词 (默认使用角色配置)
SYSTEM_PROMPT = "You are a helpful assistant." SYSTEM_PROMPT = ROLES[DEFAULT_ROLE]
# 上下文长度限制 # 上下文长度限制
MAX_CONTEXT_LENGTH = 131072 MAX_CONTEXT_LENGTH = 131072