yingqi-z20-Agent-libOS

repository

loading code, commits, and activity

repositories

loading repo index

#1	from __future__ import annotations
#2
#3	import asyncio
#4	import json
#5	import time
#6	from abc import ABC, abstractmethod
#7	from collections.abc import Mapping
#8	from enum import Enum
#9	from typing import Any, ClassVar, Generic, TypeVar
#10
#11	from pydantic import BaseModel, ConfigDict, Field, ValidationError as PydanticValidationError
#12
#13	from agent_libos.config import DEFAULT_CONFIG
#14	from agent_libos.models.exceptions import (
#15	CapabilityDenied,
#16	HumanApprovalRequired,
#17	NotFound,
#18	ProcessError,
#19	ProcessMessageWaitRequired,
#20	ProcessWaitRequired,
#21	ValidationError as LibOSValidationError,
#22	)
#23	from agent_libos.models import ToolSpec
#24
#25	InputT = TypeVar("InputT", bound=BaseModel)
#26
#27	_TOOL_DEFAULTS = DEFAULT_CONFIG.tools
#28
#29
#30	class ToolErrorCode(str, Enum):
#31	VALIDATION_ERROR = "validation_error"
#32	PERMISSION_DENIED = "permission_denied"
#33	CONFIRMATION_REQUIRED = "confirmation_required"
#34	TIMEOUT = "timeout"
#35	RATE_LIMITED = "rate_limited"
#36	TRANSIENT_ERROR = "transient_error"
#37	EXECUTION_ERROR = "execution_error"
#38	UNSUPPORTED = "unsupported"
#39
#40
#41	class ToolPolicy(BaseModel):
#42	side_effects: bool = False
#43	idempotent: bool = True
#44	requires_confirmation: bool = False
#45	permissions: set[str] = Field(default_factory=set)
#46	timeout_s: float \| None = _TOOL_DEFAULTS.default_timeout_s
#47	max_retries: int = 0
#48
#49
#50	class ToolContext(BaseModel):
#51	trace_id: str
#52	call_id: str
#53	pid: str
#54	workspace_id: str \| None = None
#55	runtime: Any \| None = Field(default=None, exclude=True)
#56	granted_permissions: set[str] = Field(default_factory=set)
#57	metadata: dict[str, Any] = Field(default_factory=dict)
#58
#59	model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True)
#60
#61
#62	class ToolArtifact(BaseModel):
#63	kind: str
#64	uri: str
#65	name: str \| None = None
#66	mime_type: str \| None = None
#67	metadata: dict[str, Any] = Field(default_factory=dict)
#68
#69
#70	class ToolError(BaseModel):
#71	code: ToolErrorCode
#72	message: str
#73	retryable: bool = False
#74	details: dict[str, Any] = Field(default_factory=dict)
#75
#76
#77	class ToolResult(BaseModel):
#78	ok: bool
#79	content: str = ""
#80	data: Any \| None = None
#81	artifacts: list[ToolArtifact] = Field(default_factory=list)
#82	error: ToolError \| None = None
#83	metadata: dict[str, Any] = Field(default_factory=dict)
#84
#85	@classmethod
#86	def success(
#87	cls,
#88	*,
#89	content: str = "",
#90	data: Any \| None = None,
#91	artifacts: list[ToolArtifact] \| None = None,
#92	metadata: dict[str, Any] \| None = None,
#93	) -> "ToolResult":
#94	return cls(ok=True, content=content, data=data, artifacts=artifacts or [], metadata=metadata or {})
#95
#96	@classmethod
#97	def failure(
#98	cls,
#99	*,
#100	code: ToolErrorCode,
#101	message: str,
#102	retryable: bool = False,
#103	details: dict[str, Any] \| None = None,
#104	metadata: dict[str, Any] \| None = None,
#105	) -> "ToolResult":
#106	return cls(
#107	ok=False,
#108	content=message,
#109	error=ToolError(code=code, message=message, retryable=retryable, details=details or {}),
#110	metadata=metadata or {},
#111	)
#112
#113
#114	class ToolExecutionError(Exception):
#115	def __init__(
#116	self,
#117	message: str,
#118	*,
#119	code: ToolErrorCode = ToolErrorCode.EXECUTION_ERROR,
#120	retryable: bool = False,
#121	details: dict[str, Any] \| None = None,
#122	) -> None:
#123	super().__init__(message)
#124	self.code = code
#125	self.retryable = retryable
#126	self.details = details or {}
#127
#128
#129	class BaseAgentTool(ABC, Generic[InputT]):
#130	name: ClassVar[str]
#131	description: ClassVar[str]
#132	args_schema: ClassVar[type[InputT]]
#133
#134	output_schema: ClassVar[type[BaseModel] \| None] = None
#135	version: ClassVar[str] = _TOOL_DEFAULTS.version
#136	policy: ClassVar[ToolPolicy] = ToolPolicy()
#137	tags: ClassVar[list[str]] = []
#138	metadata: ClassVar[dict[str, Any]] = {}
#139	expose_internal_errors: ClassVar[bool] = False
#140
#141	def spec(self) -> ToolSpec:
#142	self._validate_contract()
#143	policy = self.policy.model_dump()
#144	return ToolSpec(
#145	name=self.name,
#146	description=self.description,
#147	version=self.version,
#148	input_schema=self.args_schema.model_json_schema(),
#149	output_schema=self.output_schema.model_json_schema() if self.output_schema is not None else {},
#150	policy=policy,
#151	tags=list(self.tags),
#152	metadata=dict(self.metadata),
#153	required_capabilities=[
#154	{"resource": f"permission:{permission}", "rights": ["execute"]}
#155	for permission in sorted(self.policy.permissions)
#156	],
#157	side_effects=sorted(self.policy.permissions) if self.policy.side_effects else [],
#158	)
#159
#160	def to_openai_chat_tool(self) -> dict[str, Any]:
#161	spec = self.spec()
#162	return {
#163	"type": "function",
#164	"function": {
#165	"name": spec.name,
#166	"description": spec.description,
#167	"parameters": spec.input_schema,
#168	},
#169	}
#170
#171	def to_mcp_tool(self) -> dict[str, Any]:
#172	spec = self.spec()
#173	return {
#174	"name": spec.name,
#175	"description": spec.description,
#176	"inputSchema": spec.input_schema,
#177	"_meta": {
#178	"version": spec.version,
#179	"tags": spec.tags,
#180	"policy": spec.policy,
#181	},
#182	}
#183
#184	async def ainvoke(self, raw_args: Mapping[str, Any] \| str \| InputT, ctx: ToolContext) -> ToolResult:
#185	started_at = time.perf_counter()
#186	try:
#187	args = self.parse_args(raw_args)
#188	except PydanticValidationError as exc:
#189	return ToolResult.failure(
#190	code=ToolErrorCode.VALIDATION_ERROR,
#191	message=f"Invalid arguments for tool `{self.name}`.",
#192	details={"errors": exc.errors()},
#193	metadata=self._base_metadata(ctx, started_at),
#194	)
#195	except Exception as exc:
#196	return ToolResult.failure(
#197	code=ToolErrorCode.VALIDATION_ERROR,
#198	message=f"Failed to parse arguments for tool `{self.name}`.",
#199	details={"error_type": type(exc).__name__},
#200	metadata=self._base_metadata(ctx, started_at),
#201	)
#202
#203	try:
#204	self._check_policy(ctx)
#205	if self.policy.timeout_s is None:
#206	raw_result = await self.execute(args, ctx)
#207	else:
#208	raw_result = await asyncio.wait_for(self.execute(args, ctx), timeout=self.policy.timeout_s)
#209	result = self._normalize_result(raw_result)
#210	result.metadata.update(self._base_metadata(ctx, started_at))
#211	return result
#212	except asyncio.TimeoutError:
#213	return ToolResult.failure(
#214	code=ToolErrorCode.TIMEOUT,
#215	message=f"Tool `{self.name}` timed out.",
#216	retryable=True,
#217	metadata=self._base_metadata(ctx, started_at),
#218	)
#219	except ToolExecutionError as exc:
#220	return ToolResult.failure(
#221	code=exc.code,
#222	message=str(exc),
#223	retryable=exc.retryable,
#224	details=exc.details,
#225	metadata=self._base_metadata(ctx, started_at),
#226	)
#227	except HumanApprovalRequired:
#228	raise
#229	except ProcessWaitRequired:
#230	raise
#231	except ProcessMessageWaitRequired:
#232	raise
#233	except CapabilityDenied as exc:
#234	return ToolResult.failure(
#235	code=ToolErrorCode.PERMISSION_DENIED,
#236	message=str(exc),
#237	details={"error_type": type(exc).__name__},
#238	metadata=self._base_metadata(ctx, started_at),
#239	)
#240	except NotFound as exc:
#241	return ToolResult.failure(
#242	code=ToolErrorCode.EXECUTION_ERROR,
#243	message=str(exc),
#244	details={"error_type": type(exc).__name__},
#245	metadata=self._base_metadata(ctx, started_at),
#246	)
#247	except ProcessError as exc:
#248	return ToolResult.failure(
#249	code=ToolErrorCode.EXECUTION_ERROR,
#250	message=str(exc),
#251	details={"error_type": type(exc).__name__},
#252	metadata=self._base_metadata(ctx, started_at),
#253	)
#254	except LibOSValidationError as exc:
#255	return ToolResult.failure(
#256	code=ToolErrorCode.VALIDATION_ERROR,
#257	message=str(exc),
#258	details={"error_type": type(exc).__name__},
#259	metadata=self._base_metadata(ctx, started_at),
#260	)
#261	except Exception as exc:
#262	details: dict[str, Any] = {"error_type": type(exc).__name__}
#263	if self.expose_internal_errors:
#264	details["message"] = str(exc)
#265	return ToolResult.failure(
#266	code=ToolErrorCode.EXECUTION_ERROR,
#267	message=f"Tool `{self.name}` failed during execution.",
#268	details=details,
#269	metadata=self._base_metadata(ctx, started_at),
#270	)
#271
#272	def invoke(self, raw_args: Mapping[str, Any] \| str \| InputT, ctx: ToolContext) -> ToolResult:
#273	try:
#274	asyncio.get_running_loop()
#275	except RuntimeError:
#276	return asyncio.run(self.ainvoke(raw_args, ctx))
#277	raise RuntimeError("Cannot call invoke() inside a running event loop. Use `await ainvoke(...)`.")
#278
#279	def parse_args(self, raw_args: Mapping[str, Any] \| str \| InputT) -> InputT:
#280	if isinstance(raw_args, self.args_schema):
#281	return raw_args
#282	if isinstance(raw_args, str):
#283	return self.args_schema.model_validate_json(raw_args)
#284	if isinstance(raw_args, Mapping):
#285	return self.args_schema.model_validate(dict(raw_args))
#286	raise TypeError(f"Tool arguments must be {self.args_schema.__name__}, dict, or JSON string.")
#287
#288	def _check_policy(self, ctx: ToolContext) -> None:
#289	missing_permissions = self.policy.permissions - ctx.granted_permissions
#290	if missing_permissions:
#291	raise ToolExecutionError(
#292	f"Permission denied for tool `{self.name}`.",
#293	code=ToolErrorCode.PERMISSION_DENIED,
#294	details={"missing_permissions": sorted(missing_permissions)},
#295	)
#296	if self.policy.requires_confirmation and not ctx.metadata.get("confirmed", False):
#297	raise ToolExecutionError(
#298	f"Confirmation required before executing tool `{self.name}`.",
#299	code=ToolErrorCode.CONFIRMATION_REQUIRED,
#300	)
#301
#302	def _normalize_result(self, raw_result: Any) -> ToolResult:
#303	if isinstance(raw_result, ToolResult):
#304	return raw_result
#305	if isinstance(raw_result, BaseModel):
#306	return ToolResult.success(content=raw_result.model_dump_json(), data=raw_result.model_dump())
#307	if isinstance(raw_result, (dict, list)):
#308	return ToolResult.success(content=json.dumps(raw_result, ensure_ascii=False, default=str), data=raw_result)
#309	if raw_result is None:
#310	return ToolResult.success()
#311	return ToolResult.success(content=str(raw_result), data=raw_result)
#312
#313	def _validate_contract(self) -> None:
#314	if not getattr(self, "name", None):
#315	raise TypeError(f"{self.__class__.__name__} must define non-empty `name`.")
#316	if not getattr(self, "description", None):
#317	raise TypeError(f"{self.__class__.__name__} must define non-empty `description`.")
#318	if not getattr(self, "args_schema", None):
#319	raise TypeError(f"{self.__class__.__name__} must define `args_schema`.")
#320	if not issubclass(self.args_schema, BaseModel):
#321	raise TypeError("`args_schema` must be a Pydantic BaseModel subclass.")
#322	if self.output_schema is not None and not issubclass(self.output_schema, BaseModel):
#323	raise TypeError("`output_schema` must be a Pydantic BaseModel subclass.")
#324
#325	def _base_metadata(self, ctx: ToolContext, started_at: float) -> dict[str, Any]:
#326	return {
#327	"tool_name": self.name,
#328	"tool_version": self.version,
#329	"trace_id": ctx.trace_id,
#330	"call_id": ctx.call_id,
#331	"duration_ms": round((time.perf_counter() - started_at) * 1000, 3),
#332	}
#333
#334	@abstractmethod
#335	async def execute(self, args: InputT, ctx: ToolContext) -> Any:
#336	raise NotImplementedError
#337
#338
#339	class SyncAgentTool(BaseAgentTool[InputT], ABC):
#340	async def execute(self, args: InputT, ctx: ToolContext) -> Any:
#341	return await asyncio.to_thread(self.run, args, ctx)
#342
#343	@abstractmethod
#344	def run(self, args: InputT, ctx: ToolContext) -> Any:
#345	raise NotImplementedError
#346

z6MkqRzA/yingqi-z20-Agent-libOS