From 9e6888201a869fb772238748c2adb836a5666bcf Mon Sep 17 00:00:00 2001 From: Soulter <905617992@qq.com> Date: Mon, 17 Nov 2025 11:01:31 +0800 Subject: [PATCH] feat: add conversation management features and enhance message handling in the API --- pyproject.toml | 3 + src/astr_agent_sdk/message.py | 168 ++++++++++ src/astr_agent_sdk/run_context.py | 17 ++ src/astr_agent_sdk/tool.py | 286 ++++++++++++++++++ src/astrbot_sdk/api/basic/conversation_mgr.py | 129 ++++---- src/astrbot_sdk/api/basic/entities.py | 23 ++ src/astrbot_sdk/api/provider/entities.py | 126 ++++++++ src/astrbot_sdk/api/star/context.py | 120 +++++++- .../runtime/api/conversation_mgr.py | 108 +++++++ .../runtime/stars/new_star_utils.py | 5 +- test_plugin/commands/hello.py | 3 + 11 files changed, 913 insertions(+), 75 deletions(-) create mode 100644 src/astr_agent_sdk/message.py create mode 100644 src/astr_agent_sdk/run_context.py create mode 100644 src/astr_agent_sdk/tool.py create mode 100644 src/astrbot_sdk/api/basic/entities.py create mode 100644 src/astrbot_sdk/api/provider/entities.py diff --git a/pyproject.toml b/pyproject.toml index daa085f82..3c173e398 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,10 +6,13 @@ readme = "README.md" requires-python = ">=3.12" dependencies = [ "aiohttp>=3.13.2", + "anthropic>=0.72.1", "certifi>=2025.10.5", "click>=8.3.0", "docstring-parser>=0.17.0", + "google-genai>=1.50.0", "loguru>=0.7.3", + "openai>=2.7.2", "pydantic>=2.12.3", "pyyaml>=6.0.3", ] diff --git a/src/astr_agent_sdk/message.py b/src/astr_agent_sdk/message.py new file mode 100644 index 000000000..11128c0f6 --- /dev/null +++ b/src/astr_agent_sdk/message.py @@ -0,0 +1,168 @@ +# Inspired by MoonshotAI/kosong, credits to MoonshotAI/kosong authors for the original implementation. +# License: Apache License 2.0 + +from typing import Any, ClassVar, Literal, cast + +from pydantic import BaseModel, GetCoreSchemaHandler +from pydantic_core import core_schema + + +class ContentPart(BaseModel): + """A part of the content in a message.""" + + __content_part_registry: ClassVar[dict[str, type["ContentPart"]]] = {} + + type: str + + def __init_subclass__(cls, **kwargs: Any) -> None: + super().__init_subclass__(**kwargs) + + invalid_subclass_error_msg = f"ContentPart subclass {cls.__name__} must have a `type` field of type `str`" + + type_value = getattr(cls, "type", None) + if type_value is None or not isinstance(type_value, str): + raise ValueError(invalid_subclass_error_msg) + + cls.__content_part_registry[type_value] = cls + + @classmethod + def __get_pydantic_core_schema__( + cls, source_type: Any, handler: GetCoreSchemaHandler + ) -> core_schema.CoreSchema: + # If we're dealing with the base ContentPart class, use custom validation + if cls.__name__ == "ContentPart": + + def validate_content_part(value: Any) -> Any: + # if it's already an instance of a ContentPart subclass, return it + if hasattr(value, "__class__") and issubclass(value.__class__, cls): + return value + + # if it's a dict with a type field, dispatch to the appropriate subclass + if isinstance(value, dict) and "type" in value: + type_value: Any | None = cast(dict[str, Any], value).get("type") + if not isinstance(type_value, str): + raise ValueError(f"Cannot validate {value} as ContentPart") + target_class = cls.__content_part_registry[type_value] + return target_class.model_validate(value) + + raise ValueError(f"Cannot validate {value} as ContentPart") + + return core_schema.no_info_plain_validator_function(validate_content_part) + + # for subclasses, use the default schema + return handler(source_type) + + +class TextPart(ContentPart): + """ + >>> TextPart(text="Hello, world!").model_dump() + {'type': 'text', 'text': 'Hello, world!'} + """ + + type: str = "text" + text: str + + +class ImageURLPart(ContentPart): + """ + >>> ImageURLPart(image_url="http://example.com/image.jpg").model_dump() + {'type': 'image_url', 'image_url': 'http://example.com/image.jpg'} + """ + + class ImageURL(BaseModel): + url: str + """The URL of the image, can be data URI scheme like `data:image/png;base64,...`.""" + id: str | None = None + """The ID of the image, to allow LLMs to distinguish different images.""" + + type: str = "image_url" + image_url: str + + +class AudioURLPart(ContentPart): + """ + >>> AudioURLPart(audio_url=AudioURLPart.AudioURL(url="https://example.com/audio.mp3")).model_dump() + {'type': 'audio_url', 'audio_url': {'url': 'https://example.com/audio.mp3', 'id': None}} + """ + + class AudioURL(BaseModel): + url: str + """The URL of the audio, can be data URI scheme like `data:audio/aac;base64,...`.""" + id: str | None = None + """The ID of the audio, to allow LLMs to distinguish different audios.""" + + type: str = "audio_url" + audio_url: AudioURL + + +class ToolCall(BaseModel): + """ + A tool call requested by the assistant. + + >>> ToolCall( + ... id="123", + ... function=ToolCall.FunctionBody( + ... name="function", + ... arguments="{}" + ... ), + ... ).model_dump() + {'type': 'function', 'id': '123', 'function': {'name': 'function', 'arguments': '{}'}} + """ + + class FunctionBody(BaseModel): + name: str + arguments: str | None + + type: Literal["function"] = "function" + + id: str + """The ID of the tool call.""" + function: FunctionBody + """The function body of the tool call.""" + + +class ToolCallPart(BaseModel): + """A part of the tool call.""" + + arguments_part: str | None = None + """A part of the arguments of the tool call.""" + + +class Message(BaseModel): + """A message in a conversation.""" + + role: Literal[ + "system", + "user", + "assistant", + "tool", + ] + + content: str | list[ContentPart] + """The content of the message.""" + + +class AssistantMessageSegment(Message): + """A message segment from the assistant.""" + + role: Literal["assistant"] = "assistant" + tool_calls: list[ToolCall] | list[dict] | None = None + + +class ToolCallMessageSegment(Message): + """A message segment representing a tool call.""" + + role: Literal["tool"] = "tool" + tool_call_id: str + + +class UserMessageSegment(Message): + """A message segment from the user.""" + + role: Literal["user"] = "user" + + +class SystemMessageSegment(Message): + """A message segment from the system.""" + + role: Literal["system"] = "system" diff --git a/src/astr_agent_sdk/run_context.py b/src/astr_agent_sdk/run_context.py new file mode 100644 index 000000000..395817679 --- /dev/null +++ b/src/astr_agent_sdk/run_context.py @@ -0,0 +1,17 @@ +from dataclasses import dataclass +from typing import Any, Generic + +from typing_extensions import TypeVar + +TContext = TypeVar("TContext", default=Any) + + +@dataclass +class ContextWrapper(Generic[TContext]): + """A context for running an agent, which can be used to pass additional data or state.""" + + context: TContext + tool_call_timeout: int = 60 # Default tool call timeout in seconds + + +NoContext = ContextWrapper[None] diff --git a/src/astr_agent_sdk/tool.py b/src/astr_agent_sdk/tool.py new file mode 100644 index 000000000..ae240d2e0 --- /dev/null +++ b/src/astr_agent_sdk/tool.py @@ -0,0 +1,286 @@ +from collections.abc import Awaitable, Callable +from typing import Any, Generic + +import jsonschema +import mcp +from deprecated import deprecated +from pydantic import model_validator +from pydantic.dataclasses import dataclass + +from .run_context import ContextWrapper, TContext + +ParametersType = dict[str, Any] + + +@dataclass +class ToolSchema: + """A class representing the schema of a tool for function calling.""" + + name: str + """The name of the tool.""" + + description: str + """The description of the tool.""" + + parameters: ParametersType + """The parameters of the tool, in JSON Schema format.""" + + @model_validator(mode="after") + def validate_parameters(self) -> "ToolSchema": + jsonschema.validate( + self.parameters, jsonschema.Draft202012Validator.META_SCHEMA + ) + return self + + +@dataclass +class FunctionTool(ToolSchema, Generic[TContext]): + """A callable tool, for function calling.""" + + handler: Callable[..., Awaitable[Any]] | None = None + """a callable that implements the tool's functionality. It should be an async function.""" + + handler_module_path: str | None = None + """ + The module path of the handler function. This is empty when the origin is mcp. + This field must be retained, as the handler will be wrapped in functools.partial during initialization, + causing the handler's __module__ to be functools + """ + active: bool = True + """ + Whether the tool is active. This field is a special field for AstrBot. + You can ignore it when integrating with other frameworks. + """ + + def __repr__(self): + return f"FuncTool(name={self.name}, parameters={self.parameters}, description={self.description})" + + async def call( + self, context: ContextWrapper[TContext], **kwargs + ) -> str | mcp.types.CallToolResult: + """Run the tool with the given arguments. The handler field has priority.""" + raise NotImplementedError( + "FunctionTool.call() must be implemented by subclasses or set a handler." + ) + + +class ToolSet: + """A set of function tools that can be used in function calling. + + This class provides methods to add, remove, and retrieve tools, as well as + convert the tools to different API formats (OpenAI, Anthropic, Google GenAI). + """ + + def __init__(self, tools: list[FunctionTool] | None = None): + self.tools: list[FunctionTool] = tools or [] + + def empty(self) -> bool: + """Check if the tool set is empty.""" + return len(self.tools) == 0 + + def add_tool(self, tool: FunctionTool): + """Add a tool to the set.""" + # 检查是否已存在同名工具 + for i, existing_tool in enumerate(self.tools): + if existing_tool.name == tool.name: + self.tools[i] = tool + return + self.tools.append(tool) + + def remove_tool(self, name: str): + """Remove a tool by its name.""" + self.tools = [tool for tool in self.tools if tool.name != name] + + def get_tool(self, name: str) -> FunctionTool | None: + """Get a tool by its name.""" + for tool in self.tools: + if tool.name == name: + return tool + return None + + @deprecated(reason="Use add_tool() instead", version="4.0.0") + def add_func( + self, + name: str, + func_args: list, + desc: str, + handler: Callable[..., Awaitable[Any]], + ): + """Add a function tool to the set.""" + params = { + "type": "object", # hard-coded here + "properties": {}, + } + for param in func_args: + params["properties"][param["name"]] = { + "type": param["type"], + "description": param["description"], + } + _func = FunctionTool( + name=name, + parameters=params, + description=desc, + handler=handler, + ) + self.add_tool(_func) + + @deprecated(reason="Use remove_tool() instead", version="4.0.0") + def remove_func(self, name: str): + """Remove a function tool by its name.""" + self.remove_tool(name) + + @deprecated(reason="Use get_tool() instead", version="4.0.0") + def get_func(self, name: str) -> FunctionTool | None: + """Get all function tools.""" + return self.get_tool(name) + + @property + def func_list(self) -> list[FunctionTool]: + """Get the list of function tools.""" + return self.tools + + def openai_schema(self, omit_empty_parameter_field: bool = False) -> list[dict]: + """Convert tools to OpenAI API function calling schema format.""" + result = [] + for tool in self.tools: + func_def = { + "type": "function", + "function": { + "name": tool.name, + "description": tool.description, + }, + } + + if ( + tool.parameters and tool.parameters.get("properties") + ) or not omit_empty_parameter_field: + func_def["function"]["parameters"] = tool.parameters + + result.append(func_def) + return result + + def anthropic_schema(self) -> list[dict]: + """Convert tools to Anthropic API format.""" + result = [] + for tool in self.tools: + input_schema = {"type": "object"} + if tool.parameters: + input_schema["properties"] = tool.parameters.get("properties", {}) + input_schema["required"] = tool.parameters.get("required", []) + tool_def = { + "name": tool.name, + "description": tool.description, + "input_schema": input_schema, + } + result.append(tool_def) + return result + + def google_schema(self) -> dict: + """Convert tools to Google GenAI API format.""" + + def convert_schema(schema: dict) -> dict: + """Convert schema to Gemini API format.""" + supported_types = { + "string", + "number", + "integer", + "boolean", + "array", + "object", + "null", + } + supported_formats = { + "string": {"enum", "date-time"}, + "integer": {"int32", "int64"}, + "number": {"float", "double"}, + } + + if "anyOf" in schema: + return {"anyOf": [convert_schema(s) for s in schema["anyOf"]]} + + result = {} + + if "type" in schema and schema["type"] in supported_types: + result["type"] = schema["type"] + if "format" in schema and schema["format"] in supported_formats.get( + result["type"], + set(), + ): + result["format"] = schema["format"] + else: + result["type"] = "null" + + support_fields = { + "title", + "description", + "enum", + "minimum", + "maximum", + "maxItems", + "minItems", + "nullable", + "required", + } + result.update({k: schema[k] for k in support_fields if k in schema}) + + if "properties" in schema: + properties = {} + for key, value in schema["properties"].items(): + prop_value = convert_schema(value) + if "default" in prop_value: + del prop_value["default"] + properties[key] = prop_value + + if properties: + result["properties"] = properties + + if "items" in schema: + result["items"] = convert_schema(schema["items"]) + + return result + + tools = [] + for tool in self.tools: + d: dict[str, Any] = { + "name": tool.name, + "description": tool.description, + } + if tool.parameters: + d["parameters"] = convert_schema(tool.parameters) + tools.append(d) + + declarations = {} + if tools: + declarations["function_declarations"] = tools + return declarations + + @deprecated(reason="Use openai_schema() instead", version="4.0.0") + def get_func_desc_openai_style(self, omit_empty_parameter_field: bool = False): + return self.openai_schema(omit_empty_parameter_field) + + @deprecated(reason="Use anthropic_schema() instead", version="4.0.0") + def get_func_desc_anthropic_style(self): + return self.anthropic_schema() + + @deprecated(reason="Use google_schema() instead", version="4.0.0") + def get_func_desc_google_genai_style(self): + return self.google_schema() + + def names(self) -> list[str]: + """获取所有工具的名称列表""" + return [tool.name for tool in self.tools] + + def __len__(self): + return len(self.tools) + + def __bool__(self): + return len(self.tools) > 0 + + def __iter__(self): + return iter(self.tools) + + def __repr__(self): + return f"ToolSet(tools={self.tools})" + + def __str__(self): + return f"ToolSet(tools={self.tools})" diff --git a/src/astrbot_sdk/api/basic/conversation_mgr.py b/src/astrbot_sdk/api/basic/conversation_mgr.py index 103838036..4d775ceb2 100644 --- a/src/astrbot_sdk/api/basic/conversation_mgr.py +++ b/src/astrbot_sdk/api/basic/conversation_mgr.py @@ -1,4 +1,5 @@ -# from astr_agent_sdk.message import AssistantMessageSegment, UserMessageSegment +from astr_agent_sdk.message import AssistantMessageSegment, UserMessageSegment +from ...api.basic.entities import Conversation class BaseConversationManager: @@ -31,7 +32,9 @@ class BaseConversationManager: """ ... - async def switch_conversation(self, unified_msg_origin: str, conversation_id: str): + async def switch_conversation( + self, unified_msg_origin: str, conversation_id: str + ) -> None: """切换会话的对话 Args: @@ -75,60 +78,60 @@ class BaseConversationManager: """ ... - # async def get_conversation( - # self, - # unified_msg_origin: str, - # conversation_id: str, - # create_if_not_exists: bool = False, - # ) -> Conversation | None: - # """获取会话的对话. + async def get_conversation( + self, + unified_msg_origin: str, + conversation_id: str, + create_if_not_exists: bool = False, + ) -> Conversation | None: + """获取会话的对话. - # Args: - # unified_msg_origin (str): 统一的消息来源字符串。格式为 platform_name:message_type:session_id - # conversation_id (str): 对话 ID, 是 uuid 格式的字符串 - # create_if_not_exists (bool): 如果对话不存在,是否创建一个新的对话 - # Returns: - # conversation (Conversation): 对话对象 + Args: + unified_msg_origin (str): 统一的消息来源字符串。格式为 platform_name:message_type:session_id + conversation_id (str): 对话 ID, 是 uuid 格式的字符串 + create_if_not_exists (bool): 如果对话不存在,是否创建一个新的对话 + Returns: + conversation (Conversation): 对话对象 - # """ - # ... + """ + ... - # async def get_conversations( - # self, - # unified_msg_origin: str | None = None, - # platform_id: str | None = None, - # ) -> list[Conversation]: - # """获取对话列表. + async def get_conversations( + self, + unified_msg_origin: str | None = None, + platform_id: str | None = None, + ) -> list[Conversation]: + """获取对话列表. - # Args: - # unified_msg_origin (str): 统一的消息来源字符串。格式为 platform_name:message_type:session_id，可选 - # platform_id (str): 平台 ID, 可选参数, 用于过滤对话 - # Returns: - # conversations (List[Conversation]): 对话对象列表 + Args: + unified_msg_origin (str): 统一的消息来源字符串。格式为 platform_name:message_type:session_id，可选 + platform_id (str): 平台 ID, 可选参数, 用于过滤对话 + Returns: + conversations (List[Conversation]): 对话对象列表 - # """ - # ... + """ + ... - # async def get_filtered_conversations( - # self, - # page: int = 1, - # page_size: int = 20, - # platform_ids: list[str] | None = None, - # search_query: str = "", - # **kwargs, - # ) -> tuple[list[Conversation], int]: - # """获取过滤后的对话列表. + async def get_filtered_conversations( + self, + page: int = 1, + page_size: int = 20, + platform_ids: list[str] | None = None, + search_query: str = "", + **kwargs, + ) -> tuple[list[Conversation], int]: + """获取过滤后的对话列表. - # Args: - # page (int): 页码, 默认为 1 - # page_size (int): 每页大小, 默认为 20 - # platform_ids (list[str]): 平台 ID 列表, 可选 - # search_query (str): 搜索查询字符串, 可选 - # Returns: - # conversations (list[Conversation]): 对话对象列表 + Args: + page (int): 页码, 默认为 1 + page_size (int): 每页大小, 默认为 20 + platform_ids (list[str]): 平台 ID 列表, 可选 + search_query (str): 搜索查询字符串, 可选 + Returns: + conversations (list[Conversation]): 对话对象列表 - # """ - # ... + """ + ... async def update_conversation( self, @@ -184,23 +187,23 @@ class BaseConversationManager: """ ... - # async def add_message_pair( - # self, - # cid: str, - # user_message: UserMessageSegment | dict, - # assistant_message: AssistantMessageSegment | dict, - # ) -> None: - # """Add a user-assistant message pair to the conversation history. + async def add_message_pair( + self, + cid: str, + user_message: UserMessageSegment | dict, + assistant_message: AssistantMessageSegment | dict, + ) -> None: + """Add a user-assistant message pair to the conversation history. - # Args: - # cid (str): Conversation ID - # user_message (UserMessageSegment | dict): OpenAI-format user message object or dict - # assistant_message (AssistantMessageSegment | dict): OpenAI-format assistant message object or dict + Args: + cid (str): Conversation ID + user_message (UserMessageSegment | dict): OpenAI-format user message object or dict + assistant_message (AssistantMessageSegment | dict): OpenAI-format assistant message object or dict - # Raises: - # Exception: If the conversation with the given ID is not found - # """ - # ... + Raises: + Exception: If the conversation with the given ID is not found + """ + ... async def get_human_readable_context( self, diff --git a/src/astrbot_sdk/api/basic/entities.py b/src/astrbot_sdk/api/basic/entities.py new file mode 100644 index 000000000..05c272f31 --- /dev/null +++ b/src/astrbot_sdk/api/basic/entities.py @@ -0,0 +1,23 @@ +from dataclasses import dataclass + + +@dataclass +class Conversation: + """The conversation entity representing a chat session.""" + + platform_id: str + """The platform ID in AstrBot""" + user_id: str + """The user ID associated with the conversation.""" + cid: str + """The conversation ID, in UUID format.""" + history: str = "" + """The conversation history as a string.""" + title: str | None = "" + """The title of the conversation. For now, it's only used in WebChat.""" + persona_id: str | None = "" + """The persona ID associated with the conversation.""" + created_at: int = 0 + """The timestamp when the conversation was created.""" + updated_at: int = 0 + """The timestamp when the conversation was last updated.""" diff --git a/src/astrbot_sdk/api/provider/entities.py b/src/astrbot_sdk/api/provider/entities.py new file mode 100644 index 000000000..d28d4c030 --- /dev/null +++ b/src/astrbot_sdk/api/provider/entities.py @@ -0,0 +1,126 @@ +from __future__ import annotations +import json +from anthropic.types import Message as AnthropicMessage +from google.genai.types import GenerateContentResponse +from openai.types.chat.chat_completion import ChatCompletion +from dataclasses import dataclass, field +from ..message.chain import MessageChain +from ..message import components as Comp +from typing import Any +from astr_agent_sdk.message import ToolCall + + +@dataclass +class LLMResponse: + role: str + """角色, assistant, tool, err""" + result_chain: MessageChain | None = None + """返回的消息链""" + tools_call_args: list[dict[str, Any]] = field(default_factory=list) + """工具调用参数""" + tools_call_name: list[str] = field(default_factory=list) + """工具调用名称""" + tools_call_ids: list[str] = field(default_factory=list) + """工具调用 ID""" + + raw_completion: ( + ChatCompletion | GenerateContentResponse | AnthropicMessage | None + ) = None + _new_record: dict[str, Any] | None = None + + _completion_text: str = "" + + is_chunk: bool = False + """是否是流式输出的单个 Chunk""" + + def __init__( + self, + role: str, + completion_text: str = "", + result_chain: MessageChain | None = None, + tools_call_args: list[dict[str, Any]] | None = None, + tools_call_name: list[str] | None = None, + tools_call_ids: list[str] | None = None, + raw_completion: ChatCompletion + | GenerateContentResponse + | AnthropicMessage + | None = None, + _new_record: dict[str, Any] | None = None, + is_chunk: bool = False, + ): + """初始化 LLMResponse + + Args: + role (str): 角色, assistant, tool, err + completion_text (str, optional): 返回的结果文本，已经过时，推荐使用 result_chain. Defaults to "". + result_chain (MessageChain, optional): 返回的消息链. Defaults to None. + tools_call_args (List[Dict[str, any]], optional): 工具调用参数. Defaults to None. + tools_call_name (List[str], optional): 工具调用名称. Defaults to None. + raw_completion (ChatCompletion, optional): 原始响应, OpenAI 格式. Defaults to None. + + """ + if tools_call_args is None: + tools_call_args = [] + if tools_call_name is None: + tools_call_name = [] + if tools_call_ids is None: + tools_call_ids = [] + + self.role = role + self.completion_text = completion_text + self.result_chain = result_chain + self.tools_call_args = tools_call_args + self.tools_call_name = tools_call_name + self.tools_call_ids = tools_call_ids + self.raw_completion = raw_completion + self._new_record = _new_record + self.is_chunk = is_chunk + + @property + def completion_text(self): + if self.result_chain: + return self.result_chain.get_plain_text() + return self._completion_text + + @completion_text.setter + def completion_text(self, value): + if self.result_chain: + self.result_chain.chain = [ + comp + for comp in self.result_chain.chain + if not isinstance(comp, Comp.Plain) + ] # 清空 Plain 组件 + self.result_chain.chain.insert(0, Comp.Plain(text=value)) + else: + self._completion_text = value + + def to_openai_tool_calls(self) -> list[dict]: + """Convert to OpenAI tool calls format. Deprecated, use to_openai_to_calls_model instead.""" + ret = [] + for idx, tool_call_arg in enumerate(self.tools_call_args): + ret.append( + { + "id": self.tools_call_ids[idx], + "function": { + "name": self.tools_call_name[idx], + "arguments": json.dumps(tool_call_arg), + }, + "type": "function", + }, + ) + return ret + + def to_openai_to_calls_model(self) -> list[ToolCall]: + """The same as to_openai_tool_calls but return pydantic model.""" + ret = [] + for idx, tool_call_arg in enumerate(self.tools_call_args): + ret.append( + ToolCall( + id=self.tools_call_ids[idx], + function=ToolCall.FunctionBody( + name=self.tools_call_name[idx], + arguments=json.dumps(tool_call_arg), + ), + ), + ) + return ret diff --git a/src/astrbot_sdk/api/star/context.py b/src/astrbot_sdk/api/star/context.py index 46c198140..ed6c8e0a7 100644 --- a/src/astrbot_sdk/api/star/context.py +++ b/src/astrbot_sdk/api/star/context.py @@ -1,16 +1,21 @@ from abc import ABC from typing import Any, Callable from ..basic.conversation_mgr import BaseConversationManager +from astr_agent_sdk.tool import ToolSet, FunctionTool +from astr_agent_sdk.message import Message +from ..provider.entities import LLMResponse +from ..message.chain import MessageChain class Context(ABC): conversation_manager: BaseConversationManager + persona_manager: Any def __init__(self): self._registered_managers: dict[str, Any] = {} self._registered_functions: dict[str, Callable] = {} - def register_component(self, *components: Any) -> None: + def _register_component(self, *components: Any) -> None: """Register a components instance and its public methods. This allows the components's methods to be called via RPC using the pattern: @@ -31,22 +36,117 @@ class Context(ABC): full_name = f"{class_name}.{attr_name}" self._registered_functions[full_name] = attr - def get_registered_function(self, full_name: str) -> Callable | None: - """Get a registered function by its full name. + async def llm_generate( + self, + chat_provider_id: str, + prompt: str | None = None, + image_urls: list[str] | None = None, + tools: ToolSet | None = None, + system_prompt: str | None = None, + contexts: list[Message] | list[dict] | None = None, + **kwargs: Any, + ) -> LLMResponse: + """Call the LLM to generate a response. The method will not automatically execute tool calls. If you want to use tool calls, please use `tool_loop_agent()`. Args: - full_name: Full name in format "ComponentClassName.method_name" + chat_provider_id: The chat provider ID to use. + prompt: The prompt to send to the LLM, if `contexts` and `prompt` are both provided, `prompt` will be appended as the last user message + image_urls: List of image URLs to include in the prompt, if `contexts` and `prompt` are both provided, `image_urls` will be appended to the last user message + tools: ToolSet of tools available to the LLM + system_prompt: System prompt to guide the LLM's behavior, if provided, it will always insert as the first system message in the context + contexts: context messages for the LLM + **kwargs: Additional keyword arguments for LLM generation, OpenAI compatible + + Raises: + ChatProviderNotFoundError: If the specified chat provider ID is not found + Exception: For other errors during LLM generation + """ + ... + + async def tool_loop_agent( + self, + chat_provider_id: str, + prompt: str | None = None, + image_urls: list[str] | None = None, + tools: ToolSet | None = None, + system_prompt: str | None = None, + contexts: list[Message] | list[dict] | None = None, + max_steps: int = 30, + **kwargs: Any, + ) -> LLMResponse: + """Run an agent loop that allows the LLM to call tools iteratively until a final answer is produced. + + Args: + chat_provider_id: The chat provider ID to use. + prompt: The prompt to send to the LLM, if `contexts` and `prompt` are both provided, `prompt` will be appended as the last user message + image_urls: List of image URLs to include in the prompt, if `contexts` and `prompt` are both provided, `image_urls` will be appended to the last user message + tools: ToolSet of tools available to the LLM + system_prompt: System prompt to guide the LLM's behavior, if provided, it will always insert as the first system message in the context + contexts: context messages for the LLM + max_steps: Maximum number of tool calls before stopping the loop + **kwargs: Additional keyword arguments for LLM generation, OpenAI compatible Returns: - The callable function or None if not found + The final LLMResponse after tool calls are completed. + + Raises: + ChatProviderNotFoundError: If the specified chat provider ID is not found + Exception: For other errors during LLM generation """ + ... - return self._registered_functions.get(full_name) + async def send_message( + self, + session: str, + message_chain: MessageChain, + ) -> None: + """Send a message to a user or group. - def list_registered_functions(self) -> list[str]: - """List all registered function names. + Args: + session: unified message origin(umo), this can represent a user or group in a specific platform instance + message_chain: The MessageChain to send + + Raises: + Exception: If sending the message fails + """ + ... + + async def add_llm_tools(self, *tools: FunctionTool) -> None: + """Add tools to the LLM's toolset. + + Args: + tools: The FunctionTool instances to add + """ + ... + + async def put_kv_data( + self, + key: str, + value: str, + ) -> None: + """Insert a key-value pair data. The data will permanently stored in AstrBot unless user explicitly deleted. + + Args: + key: The key to insert + value: The value to insert + """ + ... + + async def get_kv_data(self, key: str) -> str | None: + """Get a value by key from the key-value store. + + Args: + key: The key to retrieve Returns: - List of full function names in format "ComponentClassName.method_name" + The value associated with the key, or None if not found """ - return list(self._registered_functions.keys()) + ... + + async def delete_kv_data(self, key: str) -> None: + """Delete a key-value pair by key. + + Args: + key: The key to delete + """ + ... diff --git a/src/astrbot_sdk/runtime/api/conversation_mgr.py b/src/astrbot_sdk/runtime/api/conversation_mgr.py index 1b36db3fb..9d490ab1e 100644 --- a/src/astrbot_sdk/runtime/api/conversation_mgr.py +++ b/src/astrbot_sdk/runtime/api/conversation_mgr.py @@ -1,3 +1,5 @@ +from astr_agent_sdk.message import AssistantMessageSegment, UserMessageSegment +from astrbot_sdk.api.basic.entities import Conversation from ...api.basic.conversation_mgr import BaseConversationManager from ..star_runner import StarRunner @@ -30,3 +32,109 @@ class ConversationManager(BaseConversationManager): }, ) return result["data"] + + async def switch_conversation( + self, unified_msg_origin: str, conversation_id: str + ) -> None: + await self.runner.call_context_function( + f"{self.__class__.__name__}.{self.switch_conversation.__name__}", + { + "unified_msg_origin": unified_msg_origin, + "conversation_id": conversation_id, + }, + ) + + async def delete_conversation( + self, + unified_msg_origin: str, + conversation_id: str | None = None, + ) -> None: + await self.runner.call_context_function( + f"{self.__class__.__name__}.{self.delete_conversation.__name__}", + { + "unified_msg_origin": unified_msg_origin, + "conversation_id": conversation_id, + }, + ) + + async def delete_conversations_by_user_id(self, unified_msg_origin: str) -> None: + await self.runner.call_context_function( + f"{self.__class__.__name__}.{self.delete_conversations_by_user_id.__name__}", + { + "unified_msg_origin": unified_msg_origin, + }, + ) + + async def get_curr_conversation_id(self, unified_msg_origin: str) -> str | None: + result = await self.runner.call_context_function( + f"{self.__class__.__name__}.{self.get_curr_conversation_id.__name__}", + { + "unified_msg_origin": unified_msg_origin, + }, + ) + return result["data"] + + async def get_conversation( + self, + unified_msg_origin: str, + conversation_id: str, + create_if_not_exists: bool = False, + ) -> Conversation | None: + result = await self.runner.call_context_function( + f"{self.__class__.__name__}.{self.get_conversation.__name__}", + { + "unified_msg_origin": unified_msg_origin, + "conversation_id": conversation_id, + "create_if_not_exists": create_if_not_exists, + }, + ) + return Conversation(**result["data"]) if result["data"] else None + + async def get_conversations( + self, unified_msg_origin: str | None = None, platform_id: str | None = None + ) -> list[Conversation]: + result = await self.runner.call_context_function( + f"{self.__class__.__name__}.{self.get_conversations.__name__}", + { + "unified_msg_origin": unified_msg_origin, + "platform_id": platform_id, + }, + ) + return [Conversation(**conv) for conv in result["data"]] + + async def update_conversation( + self, + unified_msg_origin: str, + conversation_id: str | None = None, + history: list[dict] | None = None, + title: str | None = None, + persona_id: str | None = None, + ) -> None: + await self.runner.call_context_function( + f"{self.__class__.__name__}.{self.update_conversation.__name__}", + { + "unified_msg_origin": unified_msg_origin, + "conversation_id": conversation_id, + "history": history, + "title": title, + "persona_id": persona_id, + }, + ) + + async def add_message_pair( + self, + cid: str, + user_message: UserMessageSegment | dict, + assistant_message: AssistantMessageSegment | dict, + ) -> None: + """Add a user-assistant message pair to the conversation history. + + Args: + cid (str): Conversation ID + user_message (UserMessageSegment | dict): OpenAI-format user message object or dict + assistant_message (AssistantMessageSegment | dict): OpenAI-format assistant message object or dict + + Raises: + Exception: If the conversation with the given ID is not found + """ + ... diff --git a/src/astrbot_sdk/runtime/stars/new_star_utils.py b/src/astrbot_sdk/runtime/stars/new_star_utils.py index cf75137a1..06ce7d481 100644 --- a/src/astrbot_sdk/runtime/stars/new_star_utils.py +++ b/src/astrbot_sdk/runtime/stars/new_star_utils.py @@ -5,6 +5,7 @@ from loguru import logger from ...api.event.astr_message_event import AstrMessageEvent, AstrMessageEventModel from ...api.star.star import StarMetadata +from ...api.star.context import Context from ..rpc.client import JSONRPCClient from ..rpc.request_helper import RPCRequestHelper from ..rpc.jsonrpc import ( @@ -181,7 +182,7 @@ class ClientHandshakeHandler: class PluginRequestHandler: """Handles JSON-RPC requests from plugins calling core methods.""" - def __init__(self, context: Any): + def __init__(self, context: Context): """Initialize the plugin request handler. Args: @@ -251,7 +252,7 @@ class PluginRequestHandler: ) # Get the registered function from context - func = self._context.get_registered_function(func_full_name) + func = self._context._registered_functions.get(func_full_name) if func is None: raise ValueError(f"Function not found: {func_full_name}") diff --git a/test_plugin/commands/hello.py b/test_plugin/commands/hello.py index 1d7dd152e..67c25f517 100644 --- a/test_plugin/commands/hello.py +++ b/test_plugin/commands/hello.py @@ -12,3 +12,6 @@ class HelloCommand(CommandComponent): ret = await self.context.conversation_manager.new_conversation("hello") print(f"New conversation created: {ret}") yield event.plain_result(f"Hello, Astrbot! Created conversation ID: {ret}") + yield event.plain_result("Hello, Astrbot!") + yield event.plain_result("Hello again, Astrbot!") + yield event.plain_result("Goodbye, Astrbot!")