# -*- coding: utf-8 -*- """The agent state class.""" from typing import Any, Type from pydantic import BaseModel, Field, field_serializer, model_validator import aiofiles.os from .._utils._common import _generate_id from ._task import Task from ..message import ( TextBlock, DataBlock, Msg, ToolCallBlock, ToolCallState, ToolResultBlock, HintBlock, ) from ..permission import PermissionContext class ReadCacheEntry(BaseModel): """The read file cache.""" lines: list[str] updated_at: float bytes: float file_path: str class ToolContext(BaseModel): """The tool context, e.g. tool cache""" max_cache_files: int = Field(default=100, gt=1) """The maximum number of cached files.""" max_cache_bytes: float = Field(default=25000, gt=10000) """The maximum size of the accumulated read file cache.""" read_file_cache: list[ReadCacheEntry] = Field(default_factory=list) """The cache for Read/Write/Edit file tools.""" activated_groups: list[str] = Field(default_factory=list) """The names of the activated tool groups, each group contains a set of tools.""" async def get_cache(self, file_path: str) -> ReadCacheEntry | None: """Get cached file content if still valid. Args: file_path: The absolute path of the file. Returns: The cached entry if valid, otherwise None. """ # Find the cache entry for entry in self.read_file_cache: if entry.file_path == file_path: # Check if cache is still valid try: updated_at = await aiofiles.os.path.getmtime(file_path) if updated_at == entry.updated_at: return entry else: # Cache is outdated, remove it self.read_file_cache.remove(entry) return None except Exception: # File might not exist anymore self.read_file_cache.remove(entry) return None return None async def cache_file(self, file_path: str, lines: list[str]) -> None: """Cache file content with LRU eviction. Args: file_path: The absolute path of the file. lines: The lines of the file content. """ try: updated_at = await aiofiles.os.path.getmtime(file_path) except Exception: # Cannot get mtime, skip caching return # Calculate size in KB new_entry_bytes = ( sum(len(line.encode("utf-8")) for line in lines) / 1024 ) # Remove existing cache for this file if present self.read_file_cache = [ entry for entry in self.read_file_cache if entry.file_path != file_path ] # Evict the oldest entries if exceeding max_cache_files while len(self.read_file_cache) >= self.max_cache_files: self.read_file_cache.pop(0) # Evict the oldest entries if exceeding max_cache_bytes current_size = sum(entry.bytes for entry in self.read_file_cache) while ( self.read_file_cache and current_size + new_entry_bytes > self.max_cache_bytes ): removed = self.read_file_cache.pop(0) current_size -= removed.bytes # Add new entry to the end (most recent) self.read_file_cache.append( ReadCacheEntry( lines=lines, updated_at=updated_at, bytes=new_entry_bytes, file_path=file_path, ), ) async def clean_file_cache( self, reserved_file_paths: set[str] | None = None, ) -> None: """Drop read caches whose paths are not in ``reserved_file_paths``. Args: reserved_file_paths: File paths from Read calls that remain in the context. Caches for these files are kept; all others are evicted. """ reserved_file_paths = reserved_file_paths or set() self.read_file_cache = [ entry for entry in self.read_file_cache if entry.file_path in reserved_file_paths ] class TaskContext(BaseModel): """The task context.""" tasks: list[Task] = Field(default_factory=lambda: []) """The task context.""" class ReplyContext(BaseModel): """The context of the current agent reply.""" reply_id: str = Field(default_factory=_generate_id) """The id of the current reply, which is also used as the id of the final message of the reply.""" cur_iter: int = 0 """The current iteration of the agent's reasoning-acting loop in this reply.""" structured_schema: Type[BaseModel] | dict | None = None """The reply's structured output requirement, a pydantic model class in process and serialized as its JSON schema dict.""" @field_serializer("structured_schema") def _serialize_structured_schema( self, value: Type[BaseModel] | dict | None, ) -> dict | None: """Serialize a schema class into its JSON schema dict.""" return value.model_json_schema() if isinstance(value, type) else value structured_output: dict | None = None """The structured output generated within this reply.""" class AgentState(BaseModel): """The agent state that should be saved and loaded from storage.""" session_id: str = Field(default_factory=_generate_id) """The session id of the agent. Normally, each session will maintain one independent agent state for each agent.""" summary: str | list[TextBlock | DataBlock] = "" """The compressed summary of the context, which will be prepended to the context when fed into the LLM.""" context: list[Msg] = Field(default_factory=list) """The uncompressed conversation context, which will be fed into the LLM""" # ================================================================= # For backword compatibility # ================================================================= @model_validator(mode="before") @classmethod def _migrate_legacy_reply_fields(cls, data: Any) -> Any: """Migrate the top-level ``reply_id``/``cur_iter`` fields (the pre-``reply_context`` storage format) into ``reply_context``, so that the states saved by previous versions load correctly.""" if isinstance(data, dict) and ( "reply_id" in data or "cur_iter" in data ): data = dict(data) reply_context = dict(data.get("reply_context") or {}) for key in ("reply_id", "cur_iter"): if key in data and key not in reply_context: reply_context[key] = data.pop(key) data["reply_context"] = reply_context return data @property def reply_id(self) -> str: """The reply id of the current reply.""" return self.reply_context.reply_id @reply_id.setter def reply_id(self, value: str) -> None: """Set the reply id of the current reply.""" self.reply_context.reply_id = value @property def cur_iter(self) -> int: """The current iteration of the agent's reasoning-acting loop in this reply.""" return self.reply_context.cur_iter @cur_iter.setter def cur_iter(self, value: int) -> None: """Set the current iteration of the agent's reasoning-acting loop in this reply.""" self.reply_context.cur_iter = value # ================================================================= # The reply context # ================================================================= reply_context: ReplyContext = Field(default_factory=ReplyContext) """The reply related context.""" # ================================================================= # The permission context # ================================================================= permission_context: PermissionContext = Field( default_factory=PermissionContext, ) """The permission context that will be passed to the toolkit to determine the tool permissions.""" # ================================================================= # The tool context # ================================================================= tool_context: ToolContext = Field(default_factory=ToolContext) # ================================================================= # The tasks context # ================================================================= tasks_context: TaskContext = Field(default_factory=TaskContext) """The task context that records the agent tasks.""" # ================================================================= # The middleware context # ================================================================= middle_context: dict[str, Any] = Field(default_factory=dict) """The context that allow the middlewares to store/get data across different replies.""" def append_context( self, name: str, blocks: list[ TextBlock | DataBlock | HintBlock | ToolCallBlock | ToolResultBlock ], ) -> None: """Append the given blocks to the agent's own message with the current `reply_id`. If such message doesn't exist, a new assistant message with agent's name and current reply ID will be created. """ # If append to the latest message if ( self.context and self.context[-1].role == "assistant" and self.context[-1].name == name and self.context[-1].id == self.reply_id ): self.context[-1].content.extend(blocks) else: # Create a new assistant message with the current reply ID self.context.append( Msg( id=self.reply_id, role="assistant", name=name, content=blocks, ), ) def has_awaiting_tool_calls(self, name: str) -> bool: """Whether the tail assistant message written by ``name`` has any tool call still awaiting an outside response — an ``ASKING`` user confirmation or a ``SUBMITTED`` external execution with no matching tool result yet. Args: name (`str`): Only messages authored by this agent name are inspected; observed messages from other agents are ignored. Returns: `bool`: ``True`` when at least one such tool call is pending. """ return bool(self.get_awaiting_tool_calls(name)) def get_awaiting_tool_calls(self, name: str) -> list[ToolCallBlock]: """Get the tail assistant message's tool calls still awaiting an outside response — an ``ASKING`` user confirmation or a ``SUBMITTED`` external execution with no matching tool result yet. Args: name (`str`): Only messages authored by this agent name are inspected; observed messages from other agents are ignored. Returns: `list[ToolCallBlock]`: The awaiting tool call blocks, empty if none. """ if not self.context: return [] last_msg = self.context[-1] if last_msg.role != "assistant" or last_msg.name != name: return [] result_ids = {b.id for b in last_msg.get_content_blocks("tool_result")} return [ tc for tc in last_msg.get_content_blocks("tool_call") if tc.state == ToolCallState.ASKING or ( tc.state == ToolCallState.SUBMITTED and tc.id not in result_ids ) ] def get_unfinished_tool_calls(self, name: str) -> list[ToolCallBlock]: """Get the current reply's tool calls with no matching tool result yet, whatever their state. A reply accumulates into a single message, so the calls of the finished reasoning-acting rounds are excluded while the ones of the ongoing round are returned. Args: name (`str`): Only messages authored by this agent name are inspected; observed messages from other agents are ignored. Returns: `list[ToolCallBlock]`: The unfinished tool call blocks, empty if none. """ if not self.context: return [] last_msg = self.context[-1] if ( last_msg.role != "assistant" or last_msg.name != name or last_msg.id != self.reply_id ): return [] result_ids = {b.id for b in last_msg.get_content_blocks("tool_result")} return [ tc for tc in last_msg.get_content_blocks("tool_call") if tc.id not in result_ids ]