Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
159 changes: 139 additions & 20 deletions agentrun/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@
- Integration: 框架集成 / Framework integration
"""

from typing import TYPE_CHECKING

__version__ = "0.0.8"

# Agent Runtime
Expand Down Expand Up @@ -87,27 +89,50 @@
SandboxClient,
Template,
)
# Server
from agentrun.server import (
AgentRequest,
AgentResult,
AgentRunServer,
AsyncInvokeAgentHandler,
InvokeAgentHandler,
Message,
MessageRole,
OpenAIProtocolHandler,
ProtocolHandler,
SyncInvokeAgentHandler,
)
# ToolSet
from agentrun.toolset import ToolSet, ToolSetClient
from agentrun.utils.config import Config
from agentrun.utils.exception import (
ResourceAlreadyExistError,
ResourceNotExistError,
)
from agentrun.utils.model import Status

# Server - 延迟导入以避免可选依赖问题
# Type hints for IDE and type checkers
if TYPE_CHECKING:
from agentrun.server import (
AgentEvent,
AgentEventItem,
AgentRequest,
AgentResult,
AgentResultItem,
AgentReturnType,
AgentRunServer,
AguiEventNormalizer,
AGUIProtocolConfig,
AGUIProtocolHandler,
AsyncAgentEventGenerator,
AsyncAgentResultGenerator,
AsyncInvokeAgentHandler,
BaseProtocolHandler,
EventType,
InvokeAgentHandler,
MergeOptions,
Message,
MessageRole,
OpenAIProtocolConfig,
OpenAIProtocolHandler,
ProtocolConfig,
ProtocolHandler,
ServerConfig,
SyncAgentEventGenerator,
SyncAgentResultGenerator,
SyncInvokeAgentHandler,
Tool,
ToolCall,
)

__all__ = [
######## Agent Runtime ########
# base
Expand Down Expand Up @@ -185,22 +210,116 @@
######## ToolSet ########
"ToolSetClient",
"ToolSet",
######## Server ########
######## Server (延迟加载) ########
# Server
"AgentRunServer",
# Config
"ServerConfig",
"ProtocolConfig",
"OpenAIProtocolConfig",
"AGUIProtocolConfig",
# Request/Response Models
"AgentRequest",
"AgentResponse",
"AgentEvent",
"AgentResult",
"AgentStreamResponse",
"Message",
"MessageRole",
"Tool",
"ToolCall",
# Event Types
"EventType",
# Type Aliases
"AgentEventItem",
"AgentResultItem",
"AgentReturnType",
"SyncAgentEventGenerator",
"SyncAgentResultGenerator",
"AsyncAgentEventGenerator",
"AsyncAgentResultGenerator",
"InvokeAgentHandler",
"AsyncInvokeAgentHandler",
"SyncInvokeAgentHandler",
"Message",
"MessageRole",
# Protocol Base
"ProtocolHandler",
"BaseProtocolHandler",
# Protocol - OpenAI
"OpenAIProtocolHandler",
"AgentStreamIterator",
# Protocol - AG-UI
"AGUIProtocolHandler",
# Event Normalizer
"AguiEventNormalizer",
# Helpers
"MergeOptions",
######## Others ########
"Status",
"ResourceNotExistError",
"ResourceAlreadyExistError",
"Config",
Comment on lines 202 to +256
Copy link

Copilot AI Dec 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The Status export has been removed from __all__, but the import statement from agentrun.utils.model import Status still exists on line 99. This creates an inconsistency where Status is imported but not exported, making it unavailable to users who import from this module.

Copilot uses AI. Check for mistakes.
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@copilot open a new pull request to apply changes based on this feedback

]

# Server 模块的所有导出
_SERVER_EXPORTS = {
"AgentRunServer",
"ServerConfig",
"ProtocolConfig",
"OpenAIProtocolConfig",
"AGUIProtocolConfig",
"AgentRequest",
"AgentEvent",
"AgentResult",
"Message",
"MessageRole",
"Tool",
"ToolCall",
"EventType",
"AgentEventItem",
"AgentResultItem",
"AgentReturnType",
"SyncAgentEventGenerator",
"SyncAgentResultGenerator",
"AsyncAgentEventGenerator",
"AsyncAgentResultGenerator",
"InvokeAgentHandler",
"AsyncInvokeAgentHandler",
"SyncInvokeAgentHandler",
"ProtocolHandler",
"BaseProtocolHandler",
"OpenAIProtocolHandler",
"AGUIProtocolHandler",
"AguiEventNormalizer",
"MergeOptions",
}

# 可选依赖包映射:安装命令 -> 导入错误的包名列表
# Optional dependency mapping: installation command -> list of import error package names
# 将使用相同安装命令的包合并到一起 / Group packages with the same installation command
_OPTIONAL_PACKAGES = {
"agentrun-sdk[server]": ["fastapi", "uvicorn", "ag_ui"],
}


def __getattr__(name: str):
"""延迟加载 server 模块的导出,避免可选依赖导致导入失败

当用户访问 server 相关的类时,才尝试导入 server 模块。
如果 server 可选依赖未安装,会抛出清晰的错误提示。
"""
if name in _SERVER_EXPORTS:
try:
from agentrun import server

return getattr(server, name)
except ImportError as e:
# 检查是否是缺少可选依赖导致的错误
error_str = str(e)
for install_cmd, package_names in _OPTIONAL_PACKAGES.items():
for package_name in package_names:
if package_name in error_str:
raise ImportError(
f"'{name}' requires the 'server' optional dependencies. "
f"Install with: pip install {install_cmd}\n"
f"Original error: {e}"
) from e
# 其他导入错误继续抛出
raise

raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
4 changes: 2 additions & 2 deletions agentrun/credential/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -134,7 +134,7 @@ def outbound_tool_ak_sk(
cls,
provider: str,
access_key_id: str,
access_key_secred: str,
access_key_secret: str,
account_id: str,
):
"""配置访问第三方工具的 ak/sk 凭证"""
Expand All @@ -148,7 +148,7 @@ def outbound_tool_ak_sk(
"accountId": account_id,
},
},
credential_secret=access_key_secred,
credential_secret=access_key_secret,
)

@classmethod
Expand Down
1 change: 1 addition & 0 deletions agentrun/model/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -131,6 +131,7 @@ class ProxyConfigTokenRateLimiter(BaseModel):

class ProxyConfigAIGuardrailConfig(BaseModel):
"""AI 防护配置"""

check_request: Optional[bool] = None
check_response: Optional[bool] = None

Expand Down
72 changes: 38 additions & 34 deletions agentrun/server/agui_protocol.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,36 +19,12 @@
)
import uuid

from ag_ui.core import AssistantMessage
from ag_ui.core import CustomEvent as AguiCustomEvent
from ag_ui.core import EventType as AguiEventType
from ag_ui.core import Message as AguiMessage
from ag_ui.core import MessagesSnapshotEvent
from ag_ui.core import RawEvent as AguiRawEvent
from ag_ui.core import (
RunErrorEvent,
RunFinishedEvent,
RunStartedEvent,
StateDeltaEvent,
StateSnapshotEvent,
StepFinishedEvent,
StepStartedEvent,
SystemMessage,
TextMessageContentEvent,
TextMessageEndEvent,
TextMessageStartEvent,
)
from ag_ui.core import Tool as AguiTool
from ag_ui.core import ToolCall as AguiToolCall
from ag_ui.core import (
ToolCallArgsEvent,
ToolCallEndEvent,
ToolCallResultEvent,
ToolCallStartEvent,
)
from ag_ui.core import ToolMessage as AguiToolMessage
from ag_ui.core import UserMessage
from ag_ui.encoder import EventEncoder
if TYPE_CHECKING:
from ag_ui.core import (
Message as AguiMessage,
)
from ag_ui.encoder import EventEncoder

from fastapi import APIRouter, Request
from fastapi.responses import StreamingResponse
import pydash
Expand Down Expand Up @@ -101,16 +77,20 @@ class StreamStateMachine:
run_errored: bool = False

def end_all_tools(
self, encoder: EventEncoder, exclude: Optional[str] = None
self, encoder: "EventEncoder", exclude: Optional[str] = None
) -> Iterator[str]:
from ag_ui.core import ToolCallEndEvent

for tool_id, state in self.tool_call_states.items():
if exclude and tool_id == exclude:
continue
if state.started and not state.ended:
yield encoder.encode(ToolCallEndEvent(tool_call_id=tool_id))
state.ended = True

def ensure_text_started(self, encoder: EventEncoder) -> Iterator[str]:
def ensure_text_started(self, encoder: "EventEncoder") -> Iterator[str]:
from ag_ui.core import TextMessageStartEvent

if not self.text.started or self.text.ended:
if self.text.ended:
self.text = TextState()
Expand All @@ -123,7 +103,9 @@ def ensure_text_started(self, encoder: EventEncoder) -> Iterator[str]:
self.text.started = True
self.text.ended = False

def end_text_if_open(self, encoder: EventEncoder) -> Iterator[str]:
def end_text_if_open(self, encoder: "EventEncoder") -> Iterator[str]:
from ag_ui.core import TextMessageEndEvent

if self.text.started and not self.text.ended:
yield encoder.encode(
TextMessageEndEvent(message_id=self.text.message_id)
Expand Down Expand Up @@ -168,6 +150,8 @@ class AGUIProtocolHandler(BaseProtocolHandler):
name = "ag-ui"

def __init__(self, config: Optional[ServerConfig] = None):
from ag_ui.encoder import EventEncoder

self._config = config.agui if config else None
self._encoder = EventEncoder()

Expand Down Expand Up @@ -363,6 +347,8 @@ async def _format_stream(
Yields:
SSE 格式的字符串
"""
from ag_ui.core import RunFinishedEvent, RunStartedEvent

state = StreamStateMachine()

# 发送 RUN_STARTED
Expand Down Expand Up @@ -420,6 +406,18 @@ def _process_event_with_boundaries(
"""处理事件并注入边界事件"""
import json

from ag_ui.core import CustomEvent as AguiCustomEvent
from ag_ui.core import (
RunErrorEvent,
StateDeltaEvent,
StateSnapshotEvent,
TextMessageContentEvent,
ToolCallArgsEvent,
ToolCallEndEvent,
ToolCallResultEvent,
ToolCallStartEvent,
)

# RAW 事件直接透传
if event.event == EventType.RAW:
raw_data = event.data.get("raw", "")
Expand Down Expand Up @@ -703,7 +701,7 @@ def _process_event_with_boundaries(

def _convert_messages_for_snapshot(
self, messages: List[Dict[str, Any]]
) -> List[AguiMessage]:
) -> List["AguiMessage"]:
"""将消息列表转换为 ag-ui-protocol 格式

Args:
Expand All @@ -712,6 +710,10 @@ def _convert_messages_for_snapshot(
Returns:
ag-ui-protocol 消息列表
"""
from ag_ui.core import AssistantMessage, SystemMessage
from ag_ui.core import ToolMessage as AguiToolMessage
from ag_ui.core import UserMessage

result = []
for msg in messages:
if not isinstance(msg, dict):
Expand Down Expand Up @@ -779,6 +781,8 @@ async def _error_stream(self, message: str) -> AsyncIterator[str]:
Yields:
SSE 格式的错误事件
"""
from ag_ui.core import RunErrorEvent, RunStartedEvent

thread_id = str(uuid.uuid4())
run_id = str(uuid.uuid4())

Expand Down
Loading