amrita_core 0.1.0__py3-none-any.whl

amrita_core/logging.py

@@ -0,0 +1,71 @@
+# ref: https://github.com/NoneBot/NoneBot2/blob/main/nonebot/log.py
+import inspect
+import logging
+import os
+import sys
+from typing import TYPE_CHECKING, Protocol
+
+import loguru
+
+if TYPE_CHECKING:
+    from loguru import Logger, Record
+
+logger: "Logger" = loguru.logger
+
+debug: bool = False
+
+
+class ToStringAble(Protocol):
+    def __str__(self) -> str: ...
+
+
+def debug_log(message: ToStringAble) -> None:
+    global debug
+    if debug:
+        logger.debug(message)
+
+
+class LoguruHandler(logging.Handler):
+    def emit(self, record: logging.LogRecord):
+        try:
+            level = logger.level(record.levelname).name
+        except ValueError:
+            level = record.levelno
+
+        frame, depth = inspect.currentframe(), 0
+        while frame and (depth == 0 or frame.f_code.co_filename == logging.__file__):
+            frame = frame.f_back
+            depth += 1
+
+        logger.opt(depth=depth, exception=record.exc_info, colors=True).log(
+            level, record.getMessage()
+        )
+
+
+def default_filter(record: "Record"):
+    """Default filter for logging, change level from Environment"""
+    log_level = os.environ.get("LOG_LEVEL", "INFO")
+    levelno = logger.level(log_level).no if isinstance(log_level, str) else log_level
+    return record["level"].no >= levelno
+
+
+default_format: str = (
+    "<green>{time:YYYY-MM-DD HH:mm:ss}</green> | "
+    "<level>{level: <7}</level> | "
+    "<magenta>{name}:{function}:{line}</magenta> | "
+    "<level>{message}</level>"
+)
+"""Default log format"""
+
+logger.remove()
+logger_id = logger.add(
+    sys.stdout,
+    level=0,
+    diagnose=False,
+    filter=default_filter,
+    format=default_format,
+)
+"""Default log handler id"""
+
+
+__autodoc__ = {"logger_id": False}

amrita_core/preset.py

@@ -0,0 +1,166 @@
+import random
+import time
+import typing
+
+from typing_extensions import Self
+
+from .logging import debug_log
+from .protocol import AdapterManager
+from .tokenizer import hybrid_token_count
+from .types import BaseModel, Message, ModelPreset, TextContent, UniResponse
+
+TEST_MSG_PROMPT: Message[list[TextContent]] = Message(
+    role="system",
+    content=[TextContent(text="You are a helpful assistant.", type="text")],
+)
+TEST_MSG_USER: Message[list[TextContent]] = Message(
+    role="user",
+    content=[
+        TextContent(text="Hello, please briefly introduce yourself.", type="text")
+    ],
+)
+
+TEST_MSG_LIST: list[Message[list[TextContent]]] = [
+    TEST_MSG_PROMPT,
+    TEST_MSG_USER,
+]
+
+
+class PresetReport(BaseModel):
+    preset_name: str  # Name of the preset
+    preset_data: ModelPreset  # Preset data
+    test_input: tuple[Message, Message]  # Test input
+    test_output: Message | None  # Test output
+    token_prompt: int  # Token count of the prompt
+    token_completion: int  # Token count of the completion
+    status: bool  # Test result
+    message: str  # Test result message
+    time_used: float
+
+
+class PresetManager:
+    """
+    PresetManager is a singleton class that manages presets.
+    """
+
+    _default_preset: ModelPreset | None = None
+    _presets: dict[str, ModelPreset]
+    _instance = None
+
+    def __new__(cls) -> Self:
+        if cls._instance is None:
+            cls._presets = {}
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def set_default_preset(self, preset: ModelPreset | str) -> None:
+        """
+        Set the default preset.
+        """
+        if isinstance(preset, str):
+            preset = self.get_preset(preset)
+        if preset.name not in self._presets:
+            self.add_preset(preset)
+        self._default_preset = preset
+
+    def get_default_preset(self) -> ModelPreset:
+        """
+        Get the default preset.
+        """
+        if self._default_preset is None:
+            self._default_preset = random.choice(list(self._presets.values()))
+        return self._default_preset
+
+    def get_preset(self, name: str) -> ModelPreset:
+        """
+        Get a preset by name.
+        """
+        if name not in self._presets:
+            raise ValueError(f"Preset {name} not found")
+        return self._presets[name]
+
+    def add_preset(self, preset: ModelPreset) -> None:
+        """
+        Add a preset.
+        """
+        if preset.name in self._presets:
+            raise ValueError(f"Preset {preset.name} already exists")
+        self._presets[preset.name] = preset
+
+    def get_all_presets(self) -> list[ModelPreset]:
+        """
+        Get all presets.
+        """
+        return list(self._presets.values())
+
+    async def test_single_preset(self, preset: ModelPreset | str) -> PresetReport:
+        """Test a single preset for parallel execution"""
+        if isinstance(preset, str):
+            preset = self.get_preset(preset)
+        debug_log(f"Testing preset: {preset.name}...")
+        prompt_tokens = hybrid_token_count(
+            "".join(
+                [typing.cast(TextContent, msg.content[0]).text for msg in TEST_MSG_LIST]
+            )
+        )
+
+        adapter = AdapterManager().safe_get_adapter(preset.protocol)
+        if adapter is None:
+            return PresetReport(
+                preset_name=preset.name,
+                preset_data=preset,
+                test_input=(TEST_MSG_PROMPT, TEST_MSG_USER),
+                test_output=None,
+                token_prompt=prompt_tokens,
+                token_completion=0,
+                status=False,
+                message=f"Undefined protocol adapter: {preset.protocol}",
+                time_used=0,
+            )
+
+        try:
+            time_start = time.time()
+            debug_log(f"Calling preset: {preset.name}...")
+            data = [  # noqa: RUF015
+                i
+                async for i in adapter(preset).call_api(TEST_MSG_LIST)
+                if isinstance(i, UniResponse)
+            ][0]
+            time_end = time.time()
+            time_delta = time_end - time_start
+            debug_log(
+                f"Successfully called preset {preset.name}, took {time_delta:.2f} seconds"
+            )
+            return PresetReport(
+                preset_name=preset.name,
+                preset_data=preset,
+                test_input=(TEST_MSG_PROMPT, TEST_MSG_USER),
+                test_output=Message[list[TextContent]](
+                    role="assistant",
+                    content=[TextContent(type="text", text=data.content)],
+                ),
+                token_prompt=prompt_tokens,
+                token_completion=hybrid_token_count(data.content),
+                status=True,
+                message="",
+                time_used=time_delta,
+            )
+        except Exception as e:
+            debug_log(f"Error occurred while testing preset {preset.name}: {e}")
+            return PresetReport(
+                preset_name=preset.name,
+                preset_data=preset,
+                test_input=(TEST_MSG_PROMPT, TEST_MSG_USER),
+                test_output=None,
+                token_prompt=prompt_tokens,
+                token_completion=0,
+                status=False,
+                message=str(e),
+                time_used=0,
+            )
+
+    async def test_presets(self) -> typing.AsyncGenerator[PresetReport, None]:
+        presets: list[ModelPreset] = self.get_all_presets()
+        debug_log(f"Starting to test all presets ({len(presets)} total)...")
+        for preset in presets:
+            yield await self.test_single_preset(preset)

amrita_core/protocol.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+from abc import abstractmethod
+from collections.abc import AsyncGenerator, Iterable
+from dataclasses import dataclass
+
+from .logging import logger
+from .tools.models import ToolChoice, ToolFunctionSchema
+from .types import ModelPreset, ToolCall, UniResponse
+
+
+@dataclass
+class ModelAdapter:
+    """Base class for model adapter"""
+
+    preset: ModelPreset
+    __override__: bool = False  # Whether to allow overriding existing adapters
+
+    def __init_subclass__(cls) -> None:
+        super().__init_subclass__()
+        if not getattr(cls, "__abstract__", False):
+            AdapterManager().register_adapter(cls)
+
+    @abstractmethod
+    async def call_api(
+        self, messages: Iterable
+    ) -> AsyncGenerator[str | UniResponse[str, None], None]:
+        yield ""
+
+    async def call_tools(
+        self,
+        messages: Iterable,
+        tools: list[ToolFunctionSchema],
+        tool_choice: ToolChoice | None = None,
+    ) -> UniResponse[None, list[ToolCall] | None]:
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def get_adapter_protocol() -> str | tuple[str, ...]: ...
+
+    @property
+    def protocol(self):
+        """Get model protocol adapter"""
+        return self.get_adapter_protocol()
+
+
+class AdapterManager:
+    __instance = None
+    _adapter_class: dict[str, type[ModelAdapter]]
+
+    def __new__(cls):
+        if cls.__instance is None:
+            cls.__instance = super().__new__(cls)
+            cls.__instance._adapter_class = {}
+        return cls.__instance
+
+    def get_adapters(self) -> dict[str, type[ModelAdapter]]:
+        """Get all registered adapters"""
+        return self._adapter_class
+
+    def safe_get_adapter(self, protocol: str) -> type[ModelAdapter] | None:
+        """Get adapter"""
+        return self._adapter_class.get(protocol)
+
+    def get_adapter(self, protocol: str) -> type[ModelAdapter]:
+        """Get adapter"""
+        if protocol not in self._adapter_class:
+            raise ValueError(f"No adapter found for protocol {protocol}")
+        return self._adapter_class[protocol]
+
+    def register_adapter(self, adapter: type[ModelAdapter]):
+        """Register adapter"""
+        protocol = adapter.get_adapter_protocol()
+        override = adapter.__override__ if hasattr(adapter, "__override__") else False
+        if isinstance(protocol, str):
+            if protocol in self._adapter_class:
+                if not override:
+                    raise ValueError(
+                        f"Model protocol adapter {protocol} is already registered"
+                    )
+                logger.warning(
+                    f"Model protocol adapter {protocol} has been registered by {self._adapter_class[protocol].__name__}, overriding existing adapter"
+                )
+
+            self._adapter_class[protocol] = adapter
+        elif isinstance(protocol, tuple):
+            for p in protocol:
+                if not isinstance(p, str):
+                    raise TypeError(
+                        "Model protocol adapter must be a string or tuple of strings"
+                    )
+                if p in self._adapter_class:
+                    if not override:
+                        raise ValueError(
+                            f"Model protocol adapter {p} is already registered"
+                        )
+                    logger.warning(
+                        f"Model protocol adapter {p} has been registered by {self._adapter_class[p].__name__}, overriding existing adapter"
+                    )
+                self._adapter_class[p] = adapter

amrita_core/tokenizer.py

@@ -0,0 +1,115 @@
+import re
+from functools import lru_cache
+from typing import Literal
+
+import jieba
+
+
+@lru_cache(maxsize=2048)
+def hybrid_token_count(
+    text: str,
+    mode: Literal["word", "bpe", "char"] = "word",
+    truncate_mode: Literal["head", "tail", "middle"] = "head",
+) -> int:
+    """
+    Calculate token count for mixed Chinese-English text, supporting word, subword, and character modes
+
+    Args:
+        text: Input text
+        mode: Tokenization mode ['char'(character-level), 'word'(word-level), 'bpe'(mixed mode)], default bpe
+        truncate_mode: Truncation mode ['head'(head truncation), 'tail'(tail truncation), 'middle'(middle truncation)], default head
+
+    Returns:
+        int: Number of tokens
+    """
+    return Tokenizer(mode=mode, truncate_mode=truncate_mode).count_tokens(text=text)
+
+
+class Tokenizer:
+    """General purpose text tokenizer"""
+
+    def __init__(
+        self,
+        max_tokens: int = 2048,
+        mode: Literal["word", "bpe", "char"] = "bpe",
+        truncate_mode: Literal["head", "tail", "middle"] = "head",
+    ):
+        """
+        Initialize the tokenizer
+
+        :param max_tokens: Maximum token limit, default 2048 (only effective in Word mode)
+        :param mode: Tokenization mode ['char'(character-level), 'word'(word-level), 'bpe'(mixed mode)], default bpe
+        :param truncate_mode: Truncation mode ['head'(head truncation), 'tail'(tail truncation), 'middle'(middle truncation)], default head
+        """
+        self.max_tokens = max_tokens
+        self.mode = mode
+        self.truncate_mode = truncate_mode
+        self._word_pattern = re.compile(r"\w+|[^\w\s]")  # Match words or punctuation
+
+    def tokenize(self, text: str) -> list[str]:
+        """Perform tokenization operation, returning a list of tokens
+
+        Args:
+            text: Input text
+
+        Returns:
+            List[str]: List of tokens
+        """
+        if self.mode == "char":
+            return list(text)
+
+        # Mixed Chinese-English tokenization strategy
+        tokens = []
+        for chunk in re.findall(self._word_pattern, text):
+            if chunk.strip() == "":
+                continue
+
+            if self._is_english(chunk):
+                tokens.extend(chunk.split())
+            else:
+                tokens.extend(jieba.lcut(chunk))
+
+        return tokens[: self.max_tokens] if self.mode == "word" else tokens
+
+    def truncate(self, tokens: list[str]) -> list[str]:
+        """Perform token truncation operation
+
+        Args:
+            tokens: List of tokens
+
+        Returns:
+            List[str]: Truncated list of tokens
+        """
+        if len(tokens) <= self.max_tokens:
+            return tokens
+
+        if self.truncate_mode == "head":
+            return tokens[-self.max_tokens :]
+        elif self.truncate_mode == "tail":
+            return tokens[: self.max_tokens]
+        else:  # middle mode preserves head and tail
+            head_len = self.max_tokens // 2
+            tail_len = self.max_tokens - head_len
+            return tokens[:head_len] + tokens[-tail_len:]
+
+    def count_tokens(self, text: str) -> int:
+        """Count the number of tokens in text
+
+        Args:
+            text: Input text
+
+        Returns:
+            int: Number of tokens
+        """
+        return len(self.tokenize(text))
+
+    def _is_english(self, text: str) -> bool:
+        """Check if the text is English
+
+        Args:
+            text: Input text
+
+        Returns:
+            bool: Whether the text is English
+        """
+        return all(ord(c) < 128 for c in text)

amrita_core/tools/manager.py

@@ -0,0 +1,163 @@
+import typing
+from collections.abc import Awaitable, Callable
+from typing import Any, ClassVar, overload
+
+from typing_extensions import Self
+
+from .models import FunctionDefinitionSchema, ToolContext, ToolData, ToolFunctionSchema
+
+T = typing.TypeVar("T")
+
+
+class ToolsManager:
+    _instance = None
+    _models: ClassVar[dict[str, ToolData]] = {}
+    _disabled_tools: ClassVar[set[str]] = (
+        set()
+    )  # Disabled tools, has_tool and get_tool will not return disabled tools
+
+    def __new__(cls) -> Self:
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def has_tool(self, name: str) -> bool:
+        return False if name in self._disabled_tools else name in self._models
+
+    @overload
+    def get_tool(self, name: str) -> ToolData | None: ...
+    @overload
+    def get_tool(self, name: str, default: T) -> ToolData | T: ...
+    def get_tool(self, name: str, default: T = None) -> ToolData | T | None:
+        if not self.has_tool(name):
+            return default
+        tool: ToolData = self._models[name]
+        return tool if tool.enable_if() else default
+
+    @overload
+    def get_tool_meta(self, name: str) -> ToolFunctionSchema | None: ...
+    @overload
+    def get_tool_meta(self, name: str, default: T) -> ToolFunctionSchema | T: ...
+    def get_tool_meta(
+        self, name: str, default: T | None = None
+    ) -> ToolFunctionSchema | None | T:
+        func_data = self.get_tool(name)
+        if func_data is None:
+            return default
+        if isinstance(func_data, ToolData):
+            return func_data.data
+        return default
+
+    @overload
+    def get_tool_func(
+        self, name: str, default: T
+    ) -> (
+        Callable[[dict[str, Any]], Awaitable[str]]
+        | Callable[[ToolContext], Awaitable[str | None]]
+        | T
+    ): ...
+    @overload
+    def get_tool_func(
+        self,
+        name: str,
+    ) -> (
+        Callable[[dict[str, Any]], Awaitable[str]]
+        | Callable[[ToolContext], Awaitable[str | None]]
+        | None
+    ): ...
+    def get_tool_func(
+        self, name: str, default: T | None = None
+    ) -> (
+        Callable[[dict[str, Any]], Awaitable[str]]
+        | Callable[[ToolContext], Awaitable[str | None]]
+        | None
+        | T
+    ):
+        func_data = self.get_tool(name)
+        if func_data is None:
+            return default
+        if isinstance(func_data, ToolData):
+            return func_data.func
+        return default
+
+    def get_tools(self) -> dict[str, ToolData]:
+        return {
+            name: data
+            for name, data in self._models.items()
+            if (name not in self._disabled_tools and data.enable_if())
+        }
+
+    def tools_meta(self) -> dict[str, ToolFunctionSchema]:
+        return {
+            k: v.data
+            for k, v in self._models.items()
+            if (k not in self._disabled_tools and v.enable_if())
+        }
+
+    def tools_meta_dict(self, **kwargs) -> dict[str, dict[str, Any]]:
+        return {
+            k: v.data.model_dump(**kwargs)
+            for k, v in self._models.items()
+            if (k not in self._disabled_tools and v.enable_if())
+        }
+
+    def register_tool(self, tool: ToolData) -> None:
+        if tool.data.function.name not in self._models:
+            self._models[tool.data.function.name] = tool
+        else:
+            raise ValueError(f"Tool {tool.data.function.name} already exists")
+
+    def remove_tool(self, name: str) -> None:
+        self._models.pop(name, None)
+        if name in self._disabled_tools:
+            self._disabled_tools.remove(name)
+
+    def enable_tool(self, name: str) -> None:
+        if name in self._disabled_tools:
+            self._disabled_tools.remove(name)
+        else:
+            raise ValueError(f"Tool {name} is not disabled")
+
+    def disable_tool(self, name: str) -> None:
+        if self.has_tool(name):
+            self._disabled_tools.add(name)
+        else:
+            raise ValueError(f"Tool {name} does not exist or has been disabled")
+
+    def get_disabled_tools(self) -> list[str]:
+        return list(self._disabled_tools)
+
+
+def on_tools(
+    data: FunctionDefinitionSchema,
+    custom_run: bool = False,
+    strict: bool = False,
+    enable_if: Callable[[], bool] = lambda: True,
+) -> Callable[
+    ...,
+    Callable[[dict[str, Any]], Awaitable[str]]
+    | Callable[[ToolContext], Awaitable[str | None]],
+]:
+    """Tool registration decorator
+
+    Args:
+        data (FunctionDefinitionSchema): Function metadata
+        custom_run (bool, optional): Whether to enable custom run mode. Defaults to False.
+        strict (bool, optional): Whether to enable strict mode. Defaults to False.
+        show_call (bool, optional): Whether to show tool call. Defaults to True.
+    """
+
+    def decorator(
+        func: Callable[[dict[str, Any]], Awaitable[str]]
+        | Callable[[ToolContext], Awaitable[str | None]],
+    ):
+        tool_data = ToolData(
+            func=func,
+            data=ToolFunctionSchema(function=data, type="function", strict=strict),
+            custom_run=custom_run,
+            enable_if=enable_if,
+        )
+        ToolsManager().register_tool(tool_data)
+        return func
+
+    return decorator