134 lines
4.6 KiB
Python
134 lines
4.6 KiB
Python
import abc
|
|
import json
|
|
from collections import defaultdict
|
|
from typing import List
|
|
from astrbot.core.db import BaseDatabase
|
|
from astrbot.core import logger
|
|
from typing import TypedDict
|
|
from astrbot.core.provider.func_tool_manager import FuncCall
|
|
from astrbot.core.provider.entites import LLMResponse
|
|
from dataclasses import dataclass
|
|
class Personality(TypedDict):
|
|
prompt: str = ""
|
|
name: str = ""
|
|
|
|
@dataclass
|
|
class ProviderMeta():
|
|
id: str
|
|
model: str
|
|
type: str
|
|
|
|
|
|
class Provider(abc.ABC):
|
|
def __init__(
|
|
self,
|
|
provider_config: dict,
|
|
provider_settings: dict,
|
|
persistant_history: bool = True,
|
|
db_helper: BaseDatabase = None
|
|
) -> None:
|
|
self.model_name = ""
|
|
'''当前使用的模型名称'''
|
|
|
|
self.session_memory = defaultdict(list)
|
|
'''维护了 session_id 的上下文,**不包含 system 指令**。'''
|
|
|
|
self.provider_config = provider_config
|
|
|
|
self.provider_settings = provider_settings
|
|
|
|
self.curr_personality = Personality(prompt=provider_settings['default_personality'])
|
|
'''维护了当前的使用的 persona,即人格。'''
|
|
|
|
self.db_helper = db_helper
|
|
'''用于持久化的数据库操作对象。'''
|
|
|
|
if persistant_history:
|
|
# 读取历史记录
|
|
try:
|
|
for history in db_helper.get_llm_history(provider_type=provider_config['type']):
|
|
self.session_memory[history.session_id] = json.loads(history.content)
|
|
except BaseException as e:
|
|
logger.warning(f"读取 LLM 对话历史记录 失败:{e}。仍可正常使用。")
|
|
|
|
def set_model(self, model_name: str):
|
|
'''设置当前使用的模型名称'''
|
|
self.model_name = model_name
|
|
|
|
def get_model(self) -> str:
|
|
'''获得当前使用的模型名称'''
|
|
return self.model_name
|
|
|
|
@abc.abstractmethod
|
|
def get_current_key(self) -> str:
|
|
raise NotImplementedError()
|
|
|
|
def get_keys(self) -> List[str]:
|
|
'''获得提供商 Key'''
|
|
return self.provider_config.get("key", [])
|
|
|
|
@abc.abstractmethod
|
|
def set_key(self, key: str):
|
|
raise NotImplementedError()
|
|
|
|
@abc.abstractmethod
|
|
def get_models(self) -> List[str]:
|
|
'''获得支持的模型列表'''
|
|
raise NotImplementedError()
|
|
|
|
@abc.abstractmethod
|
|
async def get_human_readable_context(self, session_id: str, page: int, page_size: int):
|
|
'''获取人类可读的上下文
|
|
|
|
page 从 1 开始
|
|
|
|
Example:
|
|
|
|
["User: 你好", "Assistant: 你好!"]
|
|
|
|
Return:
|
|
contexts: List[str]: 上下文列表
|
|
total_pages: int: 总页数
|
|
'''
|
|
raise NotImplementedError()
|
|
|
|
@abc.abstractmethod
|
|
async def text_chat(self,
|
|
prompt: str,
|
|
session_id: str=None,
|
|
image_urls: List[str]=None,
|
|
func_tool: FuncCall=None,
|
|
contexts: List=None,
|
|
system_prompt: str=None,
|
|
**kwargs) -> LLMResponse:
|
|
'''获得 LLM 的文本对话结果。会使用当前的模型进行对话。
|
|
|
|
Args:
|
|
prompt: 提示词
|
|
session_id: 会话 ID
|
|
image_urls: 图片 URL 列表
|
|
tools: Function-calling 工具
|
|
contexts: 上下文
|
|
kwargs: 其他参数
|
|
|
|
Notes:
|
|
- 如果传入了 contexts,将会提前加上上下文。否则使用 session_memory 中的上下文。
|
|
- 可以选择性地传入 session_id,如果传入了 session_id,将会使用 session_id 对应的上下文进行对话,
|
|
并且也会记录相应的对话上下文,实现多轮对话。如果不传入则不会记录上下文。
|
|
- 如果传入了 image_urls,将会在对话时附上图片。如果模型不支持图片输入,将会抛出错误。
|
|
- 如果传入了 tools,将会使用 tools 进行 Function-calling。如果模型不支持 Function-calling,将会抛出错误。
|
|
'''
|
|
raise NotImplementedError()
|
|
|
|
@abc.abstractmethod
|
|
async def forget(self, session_id: str) -> bool:
|
|
'''重置某一个 session_id 的上下文'''
|
|
raise NotImplementedError()
|
|
|
|
def meta(self) -> ProviderMeta:
|
|
'''获取 Provider 的元数据'''
|
|
return ProviderMeta(
|
|
id=self.provider_config['id'],
|
|
model=self.get_model(),
|
|
type=self.provider_config['type']
|
|
) |