amrita_core 0.1.0__py3-none-any.whl

amrita_core/config.py

@@ -0,0 +1,159 @@
+# TODO:分实例的配置
+from __future__ import annotations
+
+import random
+import string
+from typing import Literal
+
+from pydantic import Field
+from typing_extensions import LiteralString
+
+from amrita_core.types import BaseModel
+
+
+def random_alnum_string(length: int) -> str:
+    if length < 0:
+        raise ValueError("Length can't be smaller than zero!")
+
+    chars: LiteralString = string.ascii_letters + string.digits
+
+    return "".join(random.choice(chars) for _ in range(length))
+
+
+class CookieConfig(BaseModel):
+    """Amrita Core's cookie config"""
+
+    enable_cookie: bool = Field(
+        default=True, description="Whether to enable Cookie leak detection mechanism"
+    )
+    cookie: str = Field(
+        default_factory=lambda: random_alnum_string(16),
+        description="Cookie string for security detection",
+    )
+
+
+class FunctionConfig(BaseModel):
+    use_minimal_context: bool = Field(
+        default=True,
+        description="Whether to use minimal context, i.e. system prompt + user's last message (disabling this option will use all context from the message list, which may consume a large amount of Tokens during Agent workflow execution; enabling this option may effectively reduce token usage)",
+    )
+
+    tool_calling_mode: Literal["agent", "rag", "none"] = Field(
+        default="agent",
+        description="Tool calling mode, i.e. whether to use Agent or RAG to call tools",
+    )
+    agent_tool_call_limit: int = Field(
+        default=10, description="Tool call limit in agent mode"
+    )
+    agent_tool_call_notice: Literal["hide", "notify"] = Field(
+        default="hide",
+        description="Method of showing tool call status in agent mode, hide to conceal, notify to inform",
+    )
+    agent_thought_mode: Literal[
+        "reasoning", "chat", "reasoning-required", "reasoning-optional"
+    ] = Field(
+        default="chat",
+        description="Thinking mode in agent mode, reasoning mode will first perform reasoning process, then execute tasks; "
+        "reasoning-required requires task analysis for each Tool Calling; "
+        "reasoning-optional does not require reasoning but allows it; "
+        "chat mode executes tasks directly",
+    )
+    agent_reasoning_hide: bool = Field(
+        default=False, description="Whether to hide the thought process in agent mode"
+    )
+    agent_middle_message: bool = Field(
+        default=True,
+        description="Whether to allow Agent to send intermediate messages to users in agent mode",
+    )
+    agent_mcp_client_enable: bool = Field(
+        default=False, description="Whether to enable MCP client"
+    )
+    agent_mcp_server_scripts: list[str] = Field(
+        default=[], description="List of MCP server scripts"
+    )
+
+
+class LLMConfig(BaseModel):
+    require_tools: bool = Field(
+        default=False,
+        description="Whether to force at least one tool to be used per call",
+    )
+    memory_lenth_limit: int = Field(
+        default=50, description="Maximum number of messages in memory context"
+    )
+    max_tokens: int = Field(
+        default=100,
+        description="Maximum number of tokens generated in a single response",
+    )
+    tokens_count_mode: Literal["word", "bpe", "char"] = Field(
+        default="bpe",
+        description="Token counting mode: bpe(subwords)/word(words)/char(characters)",
+    )
+    enable_tokens_limit: bool = Field(
+        default=True, description="Whether to enable context length limits"
+    )
+    session_tokens_windows: int = Field(
+        default=5000, description="Session tokens window size"
+    )
+    llm_timeout: int = Field(
+        default=60, description="API request timeout duration (seconds)"
+    )
+    auto_retry: bool = Field(
+        default=True, description="Automatically retry on request failure"
+    )
+    max_retries: int = Field(default=3, description="Maximum number of retries")
+    enable_memory_abstract: bool = Field(
+        default=True,
+        description="Whether to enable context memory summarization (will delete context and insert a summary into system instruction)",
+    )
+    memory_abstract_proportion: float = Field(
+        default=15e-2, description="Context summarization proportion (0.15=15%)"
+    )
+    enable_multi_modal: bool = Field(
+        default=True,
+        description="Whether to enable multi-modal support (currently only supports image)",
+    )
+
+
+class AmritaConfig(BaseModel):
+    function_config: FunctionConfig = Field(
+        default_factory=FunctionConfig,
+        description="Function configuration",
+    )
+    llm: LLMConfig = Field(
+        default_factory=LLMConfig,
+        description="LLM configuration",
+    )
+    cookie: CookieConfig = Field(
+        default_factory=CookieConfig, description="Cookie configuration"
+    )
+
+
+__config = AmritaConfig()
+__inited: bool = False
+
+
+def get_config() -> AmritaConfig:
+    """Get amrita config
+
+    Raises:
+        RuntimeError: Raise it if amrita core is not initialized.
+
+    Returns:
+        AmritaConfig: Amrita core config
+    """
+    if not __inited:
+        raise RuntimeError("Config is not initialized. Please set config first.")
+    return __config
+
+
+def set_config(config: AmritaConfig):
+    """Override the config.
+
+    Args:
+        config (AmritaConfig): Configuration object to set
+    """
+    global __config, __inited
+    if not __inited:
+        __inited = True
+    __config = config

amrita_core/hook/event.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import TYPE_CHECKING
+
+from amrita_core.types import USER_INPUT, SendMessageWrap
+
+if TYPE_CHECKING:
+    from amrita_core.chatmanager import ChatObject
+
+
+class EventTypeEnum(str, Enum):
+    """
+    EventTypeEnum class is used to define and manage different event types.
+    It encapsulates the string identifiers of event types, providing a structured way
+    to handle and retrieve event types.
+
+    """
+
+    COMPLETION = "COMPLETION"
+    Nil = "Nil"
+    BEFORE_COMPLETION = "BEFORE_COMPLETION"
+
+    def validate(self, name: str) -> bool:
+        return name in self
+
+
+@dataclass
+class Event:
+    user_input: USER_INPUT
+    original_context: SendMessageWrap
+    chat_object: "ChatObject"
+
+    def __post_init__(self):
+        # Initialize event type as none
+        self._event_type = EventTypeEnum.Nil
+        # Validate and store messages using SendMessageWrap
+        self._context_messages = self.original_context
+
+    @property
+    def event_type(self) -> EventTypeEnum:
+        return self._event_type
+
+    @property
+    def message(self) -> SendMessageWrap:
+        return self._context_messages
+
+    def get_context_messages(self) -> SendMessageWrap:
+        return self._context_messages
+
+    def get_event_type(self) -> str:
+        raise NotImplementedError
+
+    def get_user_input(self) -> USER_INPUT:
+        return self.user_input
+
+
+@dataclass
+class CompletionEvent(Event):
+    model_response: str
+
+    def __post_init__(self):
+        super().__post_init__()
+        # Initialize event type as completion event
+        self._event_type = EventTypeEnum.COMPLETION
+
+    @property
+    def event_type(self):
+        return EventTypeEnum.COMPLETION
+
+    def get_event_type(self) -> str:
+        return EventTypeEnum.COMPLETION
+
+    def get_model_response(self) -> str:
+        return self.model_response
+
+
+@dataclass
+class PreCompletionEvent(Event):
+    def __post_init__(self):
+        super().__post_init__()
+        self._event_type = EventTypeEnum.BEFORE_COMPLETION
+
+    @property
+    def event_type(self) -> EventTypeEnum:
+        return self._event_type
+
+    def get_event_type(self) -> str:
+        return self._event_type

amrita_core/hook/exception.py

@@ -0,0 +1,14 @@
+class MatcherException(Exception):
+    """Base exception for Matcher."""
+
+
+class BlockException(MatcherException):
+    pass
+
+
+class CancelException(MatcherException):
+    pass
+
+
+class PassException(MatcherException):
+    pass

amrita_core/hook/matcher.py

@@ -0,0 +1,213 @@
+import inspect
+from collections.abc import Awaitable, Callable
+from copy import deepcopy
+from types import FrameType
+from typing import (
+    Any,
+    ClassVar,
+    TypeAlias,
+)
+
+from pydantic import BaseModel, Field
+from typing_extensions import Self
+
+from amrita_core.logging import debug_log, logger
+
+from .event import Event
+from .exception import BlockException, CancelException, PassException
+
+ChatException: TypeAlias = BlockException | CancelException | PassException
+
+
+class FunctionData(BaseModel, arbitrary_types_allowed=True):
+    function: Callable[..., Awaitable[Any]] = Field(...)
+    signature: inspect.Signature = Field(...)
+    frame: FrameType = Field(...)
+    priority: int = Field(...)
+    block: bool = Field(...)
+    matcher: Any = Field(...)
+
+
+class EventRegistry:
+    _instance = None
+    __event_handlers: ClassVar[dict[str, list[FunctionData]]] = {}
+
+    def __new__(cls) -> Self:
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def register_handler(self, event_type: str, data: FunctionData):
+        self.__event_handlers.setdefault(event_type, []).append(data)
+
+    def get_handlers(self, event_type: str) -> list[FunctionData]:
+        self.__event_handlers.setdefault(event_type, [])
+        self.__event_handlers[event_type].sort(key=lambda x: x.priority, reverse=False)
+        return self.__event_handlers[event_type]
+
+    def _all(self) -> dict[str, list[FunctionData]]:
+        return self.__event_handlers
+
+
+class Matcher:
+    def __init__(self, event_type: str, priority: int = 10, block: bool = True):
+        """Constructor, initialize Matcher object.
+        Args:
+            event_type (str): Event type
+            priority (int, optional): Priority. Defaults to 10.
+            block (bool, optional): Whether to block subsequent events. Defaults to True.
+        """
+        if priority <= 0:
+            raise ValueError("Event priority cannot be zero or negative!")
+
+        self.event_type = event_type
+        self.priority = priority
+        self.block = block
+
+    def append_handler(self, func: Callable[..., Awaitable[Any]]):
+        frame = inspect.currentframe()
+        assert frame is not None, "Frame is None!!!"
+        func_data = FunctionData(
+            function=func,
+            signature=inspect.signature(func),
+            frame=frame,
+            priority=self.priority,
+            block=self.block,
+            matcher=self,
+        )
+        EventRegistry().register_handler(self.event_type, func_data)
+
+    def handle(self):
+        """
+        Event handler registration function
+        """
+
+        def wrapper(
+            func: Callable[..., Awaitable[Any]],
+        ):
+            self.append_handler(func)
+            return func
+
+        return wrapper
+
+    def stop_process(self):
+        """
+        Stop the event flow within the current chat plugin and immediately stop the current handler.
+        """
+        raise BlockException()
+
+    def cancel_matcher(self):
+        """
+        Stop event processing within the current chat plugin and cancel.
+        """
+        raise CancelException()
+
+    def pass_event(self):
+        """
+        Ignore the current handler and continue processing the next one.
+        """
+        raise PassException()
+
+
+class MatcherManager:
+    """
+    Event handling manager.
+    """
+
+    @staticmethod
+    async def trigger_event(*args, **kwargs) -> None:
+        """
+        Trigger a specific type of event and call all registered event handlers for that type.
+
+        Parameters:
+        - event: Event object containing event-related data.
+        - **kwargs: Keyword arguments passed to the dependency injection system.
+        - *args: Variable arguments passed to the dependency injection system.
+        """
+        event: Event | None = None
+        for i in args:
+            if isinstance(i, Event):
+                event = i
+                break
+        if not event:
+            raise RuntimeError("No event found in args")
+        event_type = event.get_event_type()  # Get event type
+        priority_tmp = 0
+        debug_log(f"Running matchers for event: {event_type}!")
+        # Check if there are handlers for this event type
+        if matcher_list := EventRegistry().get_handlers(event_type):
+            for matcher in matcher_list:
+                if matcher.priority != priority_tmp:
+                    priority_tmp = matcher.priority
+                    debug_log(f"Running matchers for priority {priority_tmp}......")
+
+                signature = matcher.signature
+                frame = matcher.frame
+                line_number = frame.f_lineno
+                file_name = frame.f_code.co_filename
+                handler = matcher.function
+                session_args = [matcher.matcher, *args]
+                session_kwargs = {**deepcopy(kwargs)}
+
+                args_types = {k: v.annotation for k, v in signature.parameters.items()}
+                filtered_args_types = {
+                    k: v for k, v in args_types.items() if v is not inspect._empty
+                }
+                if args_types != filtered_args_types:
+                    failed_args = list(args_types.keys() - filtered_args_types.keys())
+                    logger.warning(
+                        f"Matcher {matcher.function.__name__} (File: {file_name}: Line {frame.f_lineno!s}) has untyped parameters!"
+                        + f"(Args:{''.join(i + ',' for i in failed_args)}).Skipping......"
+                    )
+                    continue
+                new_args = []
+                used_indices = set()
+                for param_type in filtered_args_types.values():
+                    for i, arg in enumerate(session_args):
+                        if i in used_indices:
+                            continue
+                        if isinstance(arg, param_type):
+                            new_args.append(arg)
+                            used_indices.add(i)
+                            break
+
+                # Get keyword argument type annotations
+                kwparams = signature.parameters
+                f_kwargs = {  # TODO: kwparams dependent support
+                    param_name: session_kwargs[param.annotation]
+                    for param_name, param in kwparams.items()
+                    if param.annotation in session_kwargs
+                }
+                if len(new_args) != len(list(filtered_args_types)):
+                    continue
+
+                # Call the handler
+
+                try:
+                    logger.info(f"Starting to run Matcher: '{handler.__name__}'")
+
+                    await handler(*new_args, **f_kwargs)
+                except PassException:
+                    logger.info(
+                        f"Matcher '{handler.__name__}'(~{file_name}:{line_number}) was skipped"
+                    )
+                    continue
+                except CancelException:
+                    logger.info("Cancelled Matcher processing")
+                    return
+                except BlockException:
+                    break
+                except Exception as e:
+                    logger.opt(exception=e, colors=True).error(
+                        f"An error occurred while running '{handler.__name__}'({file_name}:{line_number}) "
+                    )
+
+                    continue
+                finally:
+                    logger.info(f"Handler {handler.__name__} finished")
+                    if matcher.block:
+                        break
+        else:
+            logger.warning(
+                f"No registered Matcher for {event_type} event, skipping processing."
+            )

amrita_core/hook/on.py

@@ -0,0 +1,14 @@
+from .event import EventTypeEnum
+from .matcher import Matcher
+
+
+def on_completion(priority: int = 10, block: bool = True):
+    return on_event(EventTypeEnum.COMPLETION, priority, block)
+
+
+def on_precompletion(priority: int = 10, block: bool = True):
+    return on_event(EventTypeEnum.BEFORE_COMPLETION, priority, block)
+
+
+def on_event(event_type: EventTypeEnum, priority: int = 10, block: bool = True):
+    return Matcher(event_type, priority, block)

amrita_core/libchat.py

@@ -0,0 +1,189 @@
+from __future__ import annotations
+
+import typing
+from collections.abc import AsyncGenerator, Generator
+
+from pydantic import ValidationError
+
+from amrita_core.preset import PresetManager
+
+from .config import get_config
+from .logging import debug_log
+from .protocol import (
+    AdapterManager,
+    ModelAdapter,
+)
+from .tokenizer import hybrid_token_count
+from .tools.models import ToolChoice
+from .types import (
+    CONTENT_LIST_TYPE,
+    Message,
+    ModelPreset,
+    ToolCall,
+    ToolResult,
+    UniResponse,
+    UniResponseUsage,
+)
+
+T = typing.TypeVar("T")
+
+
+def text_generator(
+    memory: CONTENT_LIST_TYPE, split_role: bool = False
+) -> Generator[str, None, str]:
+    memory_l = [(i.model_dump() if hasattr(i, "model_dump") else i) for i in memory]
+    role_map = {
+        "assistant": "<BOT's response>",
+        "user": "<User's query>",
+        "tool": "<Tool call>",
+    }
+    for st in memory_l:
+        if st["content"] is None:
+            continue
+        if isinstance(st["content"], str):
+            yield (
+                st["content"]
+                if not split_role
+                else role_map.get(st["role"], "") + st["content"]
+            )
+        else:
+            for s in st["content"]:
+                if s["type"] == "text" and s.get("text") is not None:
+                    yield (
+                        s["text"]
+                        if not split_role
+                        else role_map.get(st["role"], "") + s["text"]
+                    )
+    return ""
+
+
+async def get_tokens(
+    memory: CONTENT_LIST_TYPE, response: UniResponse[str, None]
+) -> UniResponseUsage[int]:
+    """Calculate token counts for messages and response
+
+    Args:
+        memory: Message history list
+        response: Model response
+
+    Returns:
+        Object containing token usage information
+    """
+    if (
+        response.usage is not None
+        and response.usage.total_tokens is not None
+        and response.usage.completion_tokens is not None
+        and response.usage.prompt_tokens is not None
+    ):
+        return response.usage
+
+    it = hybrid_token_count(
+        "".join(list(text_generator(memory))),
+        get_config().llm.tokens_count_mode,
+    )
+
+    ot = hybrid_token_count(response.content)
+    return UniResponseUsage(
+        prompt_tokens=it, total_tokens=it + ot, completion_tokens=ot
+    )
+
+
+def _validate_msg_list(
+    messages: CONTENT_LIST_TYPE,
+) -> CONTENT_LIST_TYPE:
+    validated_messages = []
+    for msg in messages:
+        if isinstance(msg, dict):
+            # Ensure message has role field
+            if "role" not in msg:
+                raise ValueError("Message dictionary is missing 'role' field")
+            try:
+                validated_msg = (
+                    Message.model_validate(msg)
+                    if msg["role"] != "tool"
+                    else ToolResult.model_validate(msg)
+                )
+            except ValidationError as e:
+                raise ValueError(f"Invalid message format: {e}")
+            validated_messages.append(validated_msg)
+        else:
+            validated_messages.append(msg)
+    return validated_messages
+
+
+async def _call_with_reflection(
+    preset: ModelPreset,
+    call_func: typing.Callable[..., typing.Awaitable[T]],
+    *args,
+    **kwargs,
+) -> T:
+    """Call the specified function with the list of presets"""
+    adapter_class = AdapterManager().safe_get_adapter(preset.protocol)
+    if adapter_class:
+        debug_log(
+            f"Using adapter {adapter_class.__name__} to handle protocol {preset.protocol}"
+        )
+    else:
+        raise ValueError(f"Undefined protocol adapter: {preset.protocol}")
+
+    debug_log(f"Getting chat for {preset.model}")
+    debug_log(f"Preset: {preset.name}")
+    debug_log(f"Key: {preset.api_key[:7]}...")
+    debug_log(f"Protocol: {preset.protocol}")
+    debug_log(f"API URL: {preset.base_url}")
+    debug_log(f"Model: {preset.model}")
+    adapter = adapter_class(preset)
+    return await call_func(adapter, *args, **kwargs)
+
+
+async def tools_caller(
+    messages: CONTENT_LIST_TYPE,
+    tools: list,
+    preset: ModelPreset | None = None,
+    tool_choice: ToolChoice | None = None,
+) -> UniResponse[None, list[ToolCall] | None]:
+    async def _call_tools(
+        adapter: ModelAdapter,
+        messages: CONTENT_LIST_TYPE,
+        tools,
+        tool_choice,
+    ):
+        return await adapter.call_tools(messages, tools, tool_choice)
+
+    preset = preset or PresetManager().get_default_preset()
+    return await _call_with_reflection(
+        preset, _call_tools, messages, tools, tool_choice
+    )
+
+
+async def call_completion(
+    messages: CONTENT_LIST_TYPE,
+    preset: ModelPreset | None = None,
+) -> AsyncGenerator[str | UniResponse[str, None], None]:
+    """Get chat response"""
+    messages = _validate_msg_list(messages)
+    preset = preset or PresetManager().get_default_preset()
+
+    async def _call_api(adapter: ModelAdapter, messages: CONTENT_LIST_TYPE):
+        async def inner():
+            async for i in adapter.call_api([(i.model_dump()) for i in messages]):
+                yield i
+
+        return inner
+
+    # Call adapter to get chat response
+    response = await _call_with_reflection(preset, _call_api, messages)
+
+    async for i in response():
+        yield i
+
+
+async def get_last_response(
+    generator: AsyncGenerator[str | UniResponse[str, None], None],
+) -> UniResponse[str, None]:
+    ls: list[UniResponse[str, None]] = [
+        i async for i in generator if isinstance(i, UniResponse)
+    ]
+    if len(ls) == 0:
+        raise RuntimeError("No response found in generator.")
+    return ls[-1]