nonebot-plugin-amrita 0.1.0__py3-none-any.whl

nonebot_plugin_amrita/__init__.py

@@ -0,0 +1,111 @@
+import asyncio
+import contextlib
+import sys
+
+import amrita_core
+from amrita_core import ChatManager, ChatObject, minimal_init
+from amrita_core.config import (
+    AmritaConfig,
+    BuiltinAgentConfig,
+    CookieConfig,
+    FunctionConfig,
+    LLMConfig,
+)
+from nonebot import get_driver, logger
+from nonebot import log as nb_log
+from nonebot.plugin import PluginMetadata
+
+from . import agent, database, dirty, memory
+from . import config as conf_module
+from .config import Config
+from .database import InsightsModel, UserDataExecutor
+from .memory import CachedUserDataRepository, MemorySchema
+
+__plugin_meta__ = PluginMetadata(
+    name="LibAmritaCore",
+    description="Add AmritaCore (a high performance agent core) support to nonebot2",
+    usage="View `https://amrita-core.suggar.top/zh` for details.",
+    type="library",
+    config=Config,
+    supported_adapters=None,
+)
+
+
+def _patch_logger():
+    logger.remove()
+    nb_log.logger_id = logger.add(
+        sys.stdout,
+        level=0,
+        diagnose=False,
+        filter=nb_log.default_filter,
+        format=nb_log.default_format,
+    )
+    amrita_core.logging.logger_id = nb_log.logger_id
+
+
+_patch_logger()  # AmritaCore会修改Loguru的配置，这里重置为NoneBot2的默认配置
+
+
+def replace_config(config: Config):
+
+    if not isinstance(config, Config):
+        raise TypeError("config must be Config")
+    conf_module._config = config
+
+
+@get_driver().on_startup
+async def init():
+    _config = conf_module._config
+    am_cookie_conf = CookieConfig(
+        enable_cookie=_config.amrita_cookie_enable, cookie=_config.amrita_cookie
+    )
+    am_function_conf = FunctionConfig(
+        agent_mcp_client_enable=_config.amrita_mcp_enable,
+        agent_mcp_server_scripts=_config.amrita_mcp_clients,
+        agent_tool_call_limit=_config.amrita_tool_call_limit,
+    )
+    am_builtin_conf = BuiltinAgentConfig(
+        tool_calling_mode=_config.amrita_agent_mode,
+        agent_thought_mode=_config.amrita_agent_thought_mode,
+    )
+    am_llm_conf = LLMConfig(
+        memory_length_limit=_config.amrita_memory_length,
+        session_tokens_windows=_config.amrita_memory_token_limit,
+        max_tokens=_config.amrita_prompt_token_limit,
+    )
+    am_conf = AmritaConfig(
+        cookie=am_cookie_conf,
+        function_config=am_function_conf,
+        builtin=am_builtin_conf,
+        llm=am_llm_conf,
+    )
+    await minimal_init(am_conf)
+
+
+@get_driver().on_shutdown
+async def shutdown():
+    logger.info("Shutting down AmritaCore...")
+
+    async def kill_all(objs: list[ChatObject]):
+        for obj in objs:
+            with contextlib.suppress(Exception):
+                obj.terminate()
+
+    await asyncio.gather(
+        *[kill_all(objs) for objs in ChatManager().running_chat_object.values()],
+        return_exceptions=True,
+    )
+
+
+__all__ = [
+    "CachedUserDataRepository",
+    "ChatManager",
+    "ChatObject",
+    "InsightsModel",
+    "MemorySchema",
+    "UserDataExecutor",
+    "agent",
+    "database",
+    "dirty",
+    "memory",
+]

nonebot_plugin_amrita/agent.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+import typing
+
+import typing_extensions
+from amrita_core import (
+    AgentRuntime as AmRuntime,
+)
+from amrita_core import (
+    AgentStrategy,
+    AmritaConfig,
+    ChatObject,
+    ModelPreset,
+    PresetManager,
+    SessionsManager,
+    get_config,
+)
+from amrita_core.builtins.agent import AmritaAgentStrategy
+from amrita_core.consts import DEFAULT_TEMPLATE
+from amrita_core.sessions import SessionData
+from amrita_core.types import Content, Message
+from jinja2 import Template
+from nonebot.adapters import Event
+from nonebot.params import Depends
+
+from nonebot_plugin_amrita.database import InsightsModel, make_id
+from nonebot_plugin_amrita.lock import lock_by_session
+from nonebot_plugin_amrita.memory import (
+    AwaredMemory,
+    CachedUserDataRepository,
+    add_usage,
+)
+
+
+class AgentSession(AmRuntime):
+    chat_objs: list[ChatObject]
+
+    def __init__(
+        self,
+        config: AmritaConfig,
+        preset: ModelPreset,
+        train: dict[str, str] | Message[str],
+        strategy: type[AgentStrategy] = AmritaAgentStrategy,
+        template: Template | str = DEFAULT_TEMPLATE,
+        session: SessionData | str | None = None,
+        no_session: bool = False,
+    ):
+        super().__init__(config, preset, train, strategy, template, session, no_session)
+        self.chat_objs = []
+
+    async def __aenter__(self) -> typing_extensions.Self:
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        del exc_type, exc, tb
+        uni_id = self.session_id
+        if not self.chat_objs:
+            return
+        async with lock_by_session(uni_id):  # Thread safe
+            dm = CachedUserDataRepository()
+            metadata = await dm.get_metadata(uni_id)
+            insight = await InsightsModel.get()
+            for chat_object in self.chat_objs:
+                if chat_object.response.usage:
+                    add_usage(metadata, chat_object.response.usage)
+                    add_usage(insight, chat_object.response.usage)
+            await insight.save()
+            await dm.update_metadata(metadata)
+
+    @classmethod
+    async def load_from(
+        cls,
+        id_or_event: Event | str,
+        train: Message[str] | dict[str, str],
+        config: AmritaConfig | None = None,
+        preset: ModelPreset | None = None,
+        **kwargs,
+    ) -> AgentSession:
+        uni_id = make_id(id_or_event)
+        dm = CachedUserDataRepository()
+        memory = await dm.get_memory(uni_id)
+        config = config or get_config()
+        preset = preset or PresetManager().get_default_preset()
+        SessionsManager().init_session(uni_id)
+        session = SessionsManager().get_session_data(uni_id)
+        session.memory = memory.memory_json
+        return cls(config, preset, train, session=session, **kwargs)
+
+    async def save_context(self):
+        session_id = self.session_id
+        dm = CachedUserDataRepository()
+        context: AwaredMemory = typing.cast(AwaredMemory, self.context)
+        mem = await dm.get_memory(session_id)
+        mem.memory_json = context
+        await dm.update_memory_data(mem)
+
+    @typing_extensions.override
+    def get_chatobject(
+        self, input: typing.Sequence[Content] | str | None, **kwargs
+    ) -> ChatObject:
+        obj = super().get_chatobject(input, **kwargs)
+        self.chat_objs.append(obj)
+        return obj
+
+
+def SessionDepends(
+    train: Message[str] | dict[str, str],
+    config: AmritaConfig | None = None,
+    preset: ModelPreset | None = None,
+    **kwargs,
+):
+    async def constructor(event: Event) -> AgentSession:
+        return await AgentSession.load_from(event, train, config, preset, **kwargs)
+
+    return Depends(constructor)