AstrBot/astrbot/core/star/context.py

from asyncio import Queue
from typing import List, TypedDict, Union

from astrbot.core import sp
from astrbot.core.provider.provider import Provider
from astrbot.core.db import BaseDatabase
from astrbot.core.config.astrbot_config import AstrBotConfig
from astrbot.core.provider.func_tool_manager import FuncCall
from astrbot.core.platform.astr_message_event import MessageSesion
from astrbot.core.message.message_event_result import MessageChain
from astrbot.core.provider.manager import ProviderManager
from astrbot.core.platform.manager import PlatformManager
from .star import star_registry, StarMetadata
from .star_handler import star_handlers_registry, StarHandlerMetadata, EventType
from .filter.command import CommandFilter
from .filter.regex import RegexFilter
from typing import Awaitable
from astrbot.core.rag.knowledge_db_mgr import KnowledgeDBManager

class Context:
    '''
    暴露给插件的接口上下文。
    '''
    _event_queue: Queue = None
    '''事件队列。消息平台通过事件队列传递消息事件。'''
    
    _config: AstrBotConfig = None
    '''AstrBot 配置信息'''
    
    _db: BaseDatabase = None
    '''AstrBot 数据库'''
    
    provider_manager: ProviderManager = None
    
    platform_manager: PlatformManager = None
    
    # back compatibility
    _register_tasks: List[Awaitable] = []
    _star_manager = None

    def __init__(self, 
        event_queue: Queue, 
        config: AstrBotConfig, 
        db: BaseDatabase, 
        provider_manager: ProviderManager = None, 
        platform_manager: PlatformManager = None,
        knowledge_db_manager: KnowledgeDBManager = None
    ):
        self._event_queue = event_queue
        self._config = config
        self._db = db
        self.provider_manager = provider_manager
        self.platform_manager = platform_manager
        self.knowledge_db_manager = knowledge_db_manager

    def get_registered_star(self, star_name: str) -> StarMetadata:
        for star in star_registry:
            if star.name == star_name:
                return star

    def get_all_stars(self) -> List[StarMetadata]:
        return star_registry
    
    def get_llm_tool_manager(self) -> FuncCall:
        '''
        获取 LLM Tool Manager
        '''
        return self.provider_manager.llm_tools
        
    def register_llm_tool(self, name: str, func_args: list, desc: str, func_obj: Awaitable) -> None:
        '''
        为函数调用（function-calling / tools-use）添加工具。
        
        @param name: 函数名
        @param func_args: 函数参数列表，格式为 [{"type": "string", "name": "arg_name", "description": "arg_description"}, ...]
        @param desc: 函数描述
        @param func_obj: 异步处理函数。
        
        异步处理函数会接收到额外的的关键词参数：event: AstrMessageEvent, context: Context。
        '''
        md = StarHandlerMetadata(
            event_type=EventType.OnLLMRequestEvent,
            handler_full_name=func_obj.__module__ + "_" + func_obj.__name__,
            handler_name=func_obj.__name__,
            handler_module_path=func_obj.__module__,
            handler=func_obj,
            event_filters=[],
            desc=desc
        )
        star_handlers_registry.append(md)
        self.provider_manager.llm_tools.add_func(name, func_args, desc, func_obj, func_obj)
    
    def unregister_llm_tool(self, name: str) -> None:
        '''删除一个函数调用工具。如果再要启用，需要重新注册。'''
        self.provider_manager.llm_tools.remove_func(name)
        
    def activate_llm_tool(self, name: str) -> bool:
        '''激活一个已经注册的函数调用工具。注册的工具默认是激活状态。
        
        Returns:
            如果没找到，会返回 False
        '''
        func_tool = self.provider_manager.llm_tools.get_func(name)
        if func_tool is not None:
            func_tool.active = True
            
            inactivated_llm_tools: list = sp.get("inactivated_llm_tools", [])
            if name in inactivated_llm_tools:
                inactivated_llm_tools.remove(name)
                sp.put("inactivated_llm_tools", inactivated_llm_tools)
            
            return True
        return False
            
    def deactivate_llm_tool(self, name: str) -> bool:
        '''停用一个已经注册的函数调用工具。
        
        Returns:
            如果没找到，会返回 False'''
        func_tool = self.provider_manager.llm_tools.get_func(name)
        if func_tool is not None:
            func_tool.active = False
            
            inactivated_llm_tools: list = sp.get("inactivated_llm_tools", [])
            if name not in inactivated_llm_tools:
                inactivated_llm_tools.append(name)
                sp.put("inactivated_llm_tools", inactivated_llm_tools)
                
            return True
        return False
    
    def register_commands(self, star_name: str, command_name: str, desc: str, priority: int, awaitable: Awaitable, use_regex=False, ignore_prefix=False):
        '''
        注册一个命令。
        
        [Deprecated] 推荐使用装饰器注册指令。该方法将在未来的版本中被移除。
        
        @param star_name: 插件（Star）名称。
        @param command_name: 命令名称。
        @param desc: 命令描述。
        @param priority: 优先级。1-10。
        @param awaitable: 异步处理函数。
        
        '''
        md = StarHandlerMetadata(
            event_type=EventType.AdapterMessageEvent,
            handler_full_name=awaitable.__module__ + "_" + awaitable.__name__,
            handler_name=awaitable.__name__,
            handler_module_path=awaitable.__module__,
            handler=awaitable,
            event_filters=[],
            desc=desc
        )
        if use_regex:
            md.event_filters.append(RegexFilter(
                regex=command_name
            ))
        else:
            md.event_filters.append(CommandFilter(
                command_name=command_name,
                handler_md=md
            ))
        star_handlers_registry.append(md)
    
    def register_provider(self, provider: Provider):
        '''
        注册一个 LLM Provider(Chat_Completion 类型)。
        '''
        self.provider_manager.provider_insts.append(provider)
        
    def get_provider_by_id(self, provider_id: str) -> Provider:
        '''
        通过 ID 获取 LLM Provider(Chat_Completion 类型)。
        '''
        for provider in self.provider_manager.provider_insts:
            if provider.meta().id == provider_id:
                return provider
        return None
    
    def get_all_providers(self) -> List[Provider]:
        '''
        获取所有 LLM Provider(Chat_Completion 类型)。
        '''
        return self.provider_manager.provider_insts
    
    def get_using_provider(self) -> Provider:
        '''
        获取当前使用的 LLM Provider(Chat_Completion 类型)。
        
        通过 /provider 指令切换。
        '''
        return self.provider_manager.curr_provider_inst
    
    def get_config(self) -> AstrBotConfig:
        '''
        获取 AstrBot 配置信息。
        '''
        return self._config
    
    def get_db(self) -> BaseDatabase:
        '''
        获取 AstrBot 数据库。
        '''
        return self._db
    
    def get_event_queue(self) -> Queue:
        '''
        获取事件队列。
        '''
        return self._event_queue
    
    async def send_message(self, session: Union[str, MessageSesion], message_chain: MessageChain) -> bool:
        '''
        根据 session(unified_msg_origin) 发送消息。
        
        @param session: 消息会话。通过 event.session 或者 event.unified_msg_origin 获取。
        @param message_chain: 消息链。
        
        @return: 是否找到匹配的平台。
        
        当 session 为字符串时，会尝试解析为 MessageSesion 对象，如果解析失败，会抛出 ValueError 异常。
        '''
        
        if isinstance(session, str):
            try:
                session = MessageSesion.from_str(session)
            except BaseException as e:
                raise ValueError("不合法的 session 字符串: " + str(e))

        for platform in self.platform_manager.platform_insts:
            if platform.meta().name == session.platform_name:
                await platform.send_by_session(session, message_chain)
                return True
        return False

    def register_task(self, task: Awaitable, desc: str):
        '''
        注册一个异步任务。
        '''
        self._register_tasks.append(task)