agnt5 0.3.0a8__cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

agnt5/agent/__init__.py

@@ -0,0 +1,48 @@
+"""Agent module - AI agents with streaming execution.
+
+This module provides the core agent primitives for building AI-powered
+applications with tool orchestration and multi-agent collaboration.
+
+Example:
+    ```python
+    from agnt5.agent import Agent, AgentResult, handoff
+
+    # Create an agent
+    agent = Agent(
+        name="researcher",
+        model="openai/gpt-4o",
+        instructions="You are a research assistant.",
+    )
+
+    # Streaming execution (recommended)
+    async for event in agent.run("Find recent AI papers"):
+        if event.event_type == EventType.LM_MESSAGE_DELTA:
+            print(event.data, end="")  # data is raw content string for deltas
+
+    # Non-streaming execution
+    result = await agent.run_sync("Find recent AI papers")
+    print(result.output)
+    ```
+"""
+
+# Import from split modules
+from .context import AgentContext
+from .result import AgentResult
+from .handoff import Handoff, handoff
+from .registry import AgentRegistry
+from .core import Agent
+from .decorator import agent
+
+__all__ = [
+    # Core classes
+    "Agent",
+    "AgentContext",
+    "AgentResult",
+    # Handoff support
+    "Handoff",
+    "handoff",
+    # Registry
+    "AgentRegistry",
+    # Decorator
+    "agent",
+]

agnt5/agent/context.py

@@ -0,0 +1,458 @@
+"""Agent execution context with conversation state management."""
+
+import logging
+import time
+from typing import Any, Dict, List, Optional, TYPE_CHECKING
+
+from ..context import Context
+from ..lm import Message
+
+if TYPE_CHECKING:
+    pass
+
+logger = logging.getLogger(__name__)
+
+
+class AgentContext(Context):
+    """
+    Context for agent execution with conversation state management.
+
+    Extends base Context with:
+    - State management via EntityStateManager
+    - Conversation history persistence
+    - Context inheritance (child agents share parent's state)
+
+    Three initialization modes:
+    1. Standalone: Creates own state manager (playground testing)
+    2. Inherit WorkflowContext: Shares parent's state manager
+    3. Inherit parent AgentContext: Shares parent's state manager
+
+    Example:
+        ```python
+        # Standalone agent with conversation history
+        ctx = AgentContext(run_id="session-1", agent_name="tutor")
+        result = await agent.run_sync("Hello", context=ctx)
+        result = await agent.run_sync("Continue", context=ctx)  # Remembers previous message
+
+        # Agent in workflow - shares workflow state
+        @workflow
+        async def research_workflow(ctx: WorkflowContext):
+            agent_result = await research_agent.run_sync("Find AI trends", context=ctx)
+            # Agent has access to workflow state via inherited context
+        ```
+    """
+
+    def __init__(
+        self,
+        run_id: str,
+        agent_name: str,
+        session_id: Optional[str] = None,
+        state_manager: Optional[Any] = None,
+        parent_context: Optional[Context] = None,
+        attempt: int = 0,
+        runtime_context: Optional[Any] = None,
+        is_streaming: bool = False,
+        tenant_id: Optional[str] = None,
+    ):
+        """
+        Initialize agent context.
+
+        Args:
+            run_id: Unique execution identifier
+            agent_name: Name of the agent
+            session_id: Session identifier for conversation history (default: run_id)
+            state_manager: Optional state manager (for context inheritance)
+            parent_context: Parent context to inherit state from
+            attempt: Retry attempt number
+            runtime_context: RuntimeContext for trace correlation
+            is_streaming: Whether this is a streaming request (for real-time SSE log delivery)
+            tenant_id: Tenant identifier for multi-tenant deployments
+        """
+        # Inherit is_streaming and tenant_id from parent context if not explicitly provided
+        if parent_context and not is_streaming:
+            is_streaming = getattr(parent_context, '_is_streaming', False)
+        if parent_context and not tenant_id:
+            tenant_id = getattr(parent_context, '_tenant_id', None)
+
+        super().__init__(run_id, attempt, runtime_context, is_streaming, tenant_id)
+
+        self._agent_name = agent_name
+        self._session_id = session_id or run_id
+        self.parent_context = parent_context  # Store for context chain traversal
+
+        # Determine state adapter based on parent context
+        from ..entity import EntityStateAdapter, _get_state_adapter
+
+        if state_manager:
+            # Explicit state adapter provided (parameter name kept for backward compat)
+            self._state_adapter = state_manager
+        elif parent_context:
+            # Try to inherit state adapter from parent
+            try:
+                # Check if parent is WorkflowContext or AgentContext
+                if hasattr(parent_context, '_workflow_entity'):
+                    # WorkflowContext - get state adapter from worker context
+                    self._state_adapter = _get_state_adapter()
+                elif hasattr(parent_context, '_state_adapter'):
+                    # Parent AgentContext - share state adapter
+                    self._state_adapter = parent_context._state_adapter
+                elif hasattr(parent_context, '_state_manager'):
+                    # Backward compatibility: parent has old _state_manager
+                    self._state_adapter = parent_context._state_manager
+                else:
+                    # FunctionContext or base Context - create new state adapter
+                    self._state_adapter = EntityStateAdapter()
+            except RuntimeError:
+                # _get_state_adapter() failed (not in worker context) - create standalone
+                self._state_adapter = EntityStateAdapter()
+        else:
+            # Try to get from worker context first
+            try:
+                self._state_adapter = _get_state_adapter()
+            except RuntimeError:
+                # Standalone - create new state adapter
+                self._state_adapter = EntityStateAdapter()
+
+        # Conversation key for state storage (used for in-memory state)
+        self._conversation_key = f"agent:{agent_name}:{self._session_id}:messages"
+        # Entity key for database persistence (without :messages suffix to match API expectations)
+        self._entity_key = f"agent:{agent_name}:{self._session_id}"
+
+        # Determine storage mode: "workflow" if parent is WorkflowContext, else "standalone"
+        self._storage_mode = "standalone"  # Default mode
+        self._workflow_entity = None
+
+        if parent_context and hasattr(parent_context, '_workflow_entity'):
+            # Agent is running within a workflow - store conversation in workflow state
+            self._storage_mode = "workflow"
+            self._workflow_entity = parent_context._workflow_entity
+            logger.debug(
+                f"Agent '{agent_name}' using workflow storage mode "
+                f"(workflow entity: {self._workflow_entity.key})"
+            )
+
+    @property
+    def state(self):
+        """
+        Get state interface for agent state management.
+
+        Note: This is a simplified in-memory state interface for agent-specific data.
+        Conversation history is managed separately via get_conversation_history() and
+        save_conversation_history() which use the Rust-backed persistence layer.
+
+        Returns:
+            Dict-like object for state operations
+
+        Example:
+            # Store agent-specific data (in-memory only)
+            ctx.state["research_results"] = data
+            ctx.state["iteration_count"] = 5
+        """
+        # Simple dict-based state for agent-specific data
+        # This is in-memory only and not persisted to platform
+        if not hasattr(self, '_agent_state'):
+            self._agent_state = {}
+        return self._agent_state
+
+    @property
+    def session_id(self) -> str:
+        """Get session identifier for this agent context."""
+        return self._session_id
+
+    async def get_conversation_history(self) -> List[Message]:
+        """
+        Retrieve conversation history from state, loading from database if needed.
+
+        Uses the EntityStateAdapter which delegates to Rust core for cache-first loading.
+        If running within a workflow, loads from workflow entity state instead.
+
+        Returns:
+            List of Message objects from conversation history
+        """
+        if self._storage_mode == "workflow":
+            return await self._load_from_workflow_state()
+        else:
+            return await self._load_from_entity_storage()
+
+    async def _load_from_workflow_state(self) -> List[Message]:
+        """Load conversation history from workflow entity state."""
+        key = f"agent.{self._agent_name}"
+        agent_data = self._workflow_entity.state.get(key, {})
+        messages_data = agent_data.get("messages", [])
+
+        # Convert dict representations back to Message objects
+        return self._convert_dicts_to_messages(messages_data)
+
+    async def _load_from_entity_storage(self) -> List[Message]:
+        """Load conversation history from AgentSession entity (standalone mode)."""
+        entity_type = "AgentSession"
+        entity_key = self._entity_key
+
+        # Load session data via adapter (Rust handles cache + platform load)
+        # Use session scope with session_id for proper entity isolation
+        session_data = await self._state_adapter.load_state(
+            entity_type,
+            entity_key,
+            scope="session",
+            scope_id=self._session_id,
+        )
+
+        # Extract messages from session object
+        if isinstance(session_data, dict) and "messages" in session_data:
+            # New format with session metadata
+            messages_data = session_data["messages"]
+        elif isinstance(session_data, list):
+            # Old format - just messages array
+            messages_data = session_data
+        else:
+            # No messages found
+            messages_data = []
+
+        # Convert dict representations back to Message objects
+        return self._convert_dicts_to_messages(messages_data)
+
+    def _convert_dicts_to_messages(self, messages_data: list) -> List[Message]:
+        """Convert list of message dicts to Message objects."""
+        from ..lm import MessageRole
+
+        messages = []
+        for msg_dict in messages_data:
+            if isinstance(msg_dict, dict):
+                role = msg_dict.get("role", "user")
+                content = msg_dict.get("content", "")
+                if role == "user":
+                    messages.append(Message.user(content))
+                elif role == "assistant":
+                    messages.append(Message.assistant(content))
+                else:
+                    # Generic message - create with MessageRole enum
+                    msg_role = (
+                        MessageRole(role)
+                        if role in ("user", "assistant", "system")
+                        else MessageRole.USER
+                    )
+                    msg = Message(role=msg_role, content=content)
+                    messages.append(msg)
+            else:
+                # Already a Message object
+                messages.append(msg_dict)
+
+        return messages
+
+    async def save_conversation_history(self, messages: List[Message]) -> None:
+        """
+        Save conversation history to state and persist to database.
+
+        Uses the EntityStateAdapter which delegates to Rust core for version-checked saves.
+        If running within a workflow, saves to workflow entity state instead.
+
+        Args:
+            messages: List of Message objects to persist
+        """
+        if self._storage_mode == "workflow":
+            await self._save_to_workflow_state(messages)
+        else:
+            await self._save_to_entity_storage(messages)
+
+    async def _save_to_workflow_state(self, messages: List[Message]) -> None:
+        """Save conversation history to workflow entity state."""
+        # Convert Message objects to dict for JSON serialization
+        messages_data = []
+        for msg in messages:
+            messages_data.append({
+                "role": msg.role.value if hasattr(msg.role, 'value') else str(msg.role),
+                "content": msg.content,
+                "timestamp": time.time()
+            })
+
+        # Build agent data structure
+        key = f"agent.{self._agent_name}"
+        current_data = self._workflow_entity.state.get(key, {})
+        now = time.time()
+
+        agent_data = {
+            "session_id": self._session_id,
+            "agent_name": self._agent_name,
+            "created_at": current_data.get("created_at", now),
+            "last_message_time": now,
+            "message_count": len(messages_data),
+            "messages": messages_data,
+            "metadata": getattr(self, '_custom_metadata', {})
+        }
+
+        # Store in workflow state (WorkflowEntity handles persistence)
+        self._workflow_entity.state.set(key, agent_data)
+        logger.info(f"Saved conversation to workflow state: {key} ({len(messages_data)} messages)")
+
+    async def _save_to_entity_storage(self, messages: List[Message]) -> None:
+        """Save conversation history to AgentSession entity (standalone mode)."""
+        # Convert Message objects to dict for JSON serialization
+        messages_data = []
+        for msg in messages:
+            messages_data.append({
+                "role": msg.role.value if hasattr(msg.role, 'value') else str(msg.role),
+                "content": msg.content,
+                "timestamp": time.time()  # Add timestamp for each message
+            })
+
+        entity_type = "AgentSession"
+        entity_key = self._entity_key
+
+        # Load current state with version for optimistic locking
+        # Use session scope with session_id for proper entity isolation
+        current_state, current_version = await self._state_adapter.load_with_version(
+            entity_type,
+            entity_key,
+            scope="session",
+            scope_id=self._session_id,
+        )
+
+        # Build session object with metadata
+        now = time.time()
+
+        # Get custom metadata from instance variable or preserve from loaded state
+        custom_metadata = getattr(self, '_custom_metadata', current_state.get("metadata", {}))
+
+        session_data = {
+            "session_id": self._session_id,
+            "agent_name": self._agent_name,
+            "created_at": current_state.get("created_at", now),  # Preserve existing or set new
+            "last_message_time": now,
+            "message_count": len(messages_data),
+            "messages": messages_data,
+            "metadata": custom_metadata  # Save custom metadata
+        }
+
+        # Save to platform via adapter (Rust handles optimistic locking)
+        # Use session scope with session_id for proper entity isolation
+        try:
+            new_version = await self._state_adapter.save_state(
+                entity_type,
+                entity_key,
+                session_data,
+                current_version,
+                scope="session",
+                scope_id=self._session_id,
+            )
+            logger.info(
+                f"Persisted conversation history: {entity_key} "
+                f"(version {current_version} -> {new_version})"
+            )
+        except Exception as e:
+            logger.error(f"Failed to persist conversation history to database: {e}")
+            # Don't fail - conversation is still in memory for this execution
+
+    async def get_metadata(self) -> Dict[str, Any]:
+        """
+        Get conversation session metadata.
+
+        Returns session metadata including:
+        - created_at: Timestamp of first message (float, Unix timestamp)
+        - last_activity: Timestamp of last message (float, Unix timestamp)
+        - message_count: Number of messages in conversation (int)
+        - custom: Dict of user-provided custom metadata
+
+        Returns:
+            Dictionary with metadata. If no conversation exists yet, returns defaults.
+
+        Example:
+            ```python
+            metadata = await context.get_metadata()
+            print(f"Session created: {metadata['created_at']}")
+            print(f"User ID: {metadata['custom'].get('user_id')}")
+            ```
+        """
+        if self._storage_mode == "workflow":
+            return await self._get_metadata_from_workflow()
+        else:
+            return await self._get_metadata_from_entity()
+
+    async def _get_metadata_from_workflow(self) -> Dict[str, Any]:
+        """Get metadata from workflow entity state."""
+        key = f"agent.{self._agent_name}"
+        agent_data = self._workflow_entity.state.get(key, {})
+
+        if not agent_data:
+            # No conversation exists yet - return defaults
+            return {
+                "created_at": None,
+                "last_activity": None,
+                "message_count": 0,
+                "custom": getattr(self, '_custom_metadata', {})
+            }
+
+        messages = agent_data.get("messages", [])
+        return {
+            "created_at": agent_data.get("created_at"),
+            "last_activity": agent_data.get("last_message_time"),
+            "message_count": len(messages),
+            "custom": agent_data.get("metadata", {})
+        }
+
+    async def _get_metadata_from_entity(self) -> Dict[str, Any]:
+        """Get metadata from AgentSession entity (standalone mode)."""
+        entity_type = "AgentSession"
+        entity_key = self._entity_key
+
+        # Load session data with session scope
+        session_data = await self._state_adapter.load_state(
+            entity_type,
+            entity_key,
+            scope="session",
+            scope_id=self._session_id,
+        )
+
+        if not session_data:
+            # No conversation exists yet - return defaults
+            return {
+                "created_at": None,
+                "last_activity": None,
+                "message_count": 0,
+                "custom": getattr(self, '_custom_metadata', {})
+            }
+
+        messages = session_data.get("messages", [])
+
+        # Derive timestamps from messages if available
+        created_at = session_data.get("created_at")
+        last_activity = session_data.get("last_message_time")
+
+        return {
+            "created_at": created_at,
+            "last_activity": last_activity,
+            "message_count": len(messages),
+            "custom": session_data.get("metadata", {})
+        }
+
+    def update_metadata(self, **kwargs) -> None:
+        """
+        Update custom session metadata.
+
+        Metadata will be persisted alongside conversation history on next save.
+        Use this to store application-specific data like user_id, preferences, etc.
+
+        Args:
+            **kwargs: Key-value pairs to store as metadata
+
+        Example:
+            ```python
+            # Store user identification and preferences
+            context.update_metadata(
+                user_id="user-123",
+                subscription_tier="premium",
+                preferences={"theme": "dark", "language": "en"}
+            )
+
+            # Later retrieve it
+            metadata = await context.get_metadata()
+            user_id = metadata["custom"]["user_id"]
+            ```
+
+        Note:
+            - Metadata is merged with existing metadata (doesn't replace)
+            - Changes persist on next save_conversation_history() call
+            - Use simple JSON-serializable types (str, int, float, dict, list)
+        """
+        if not hasattr(self, '_custom_metadata'):
+            self._custom_metadata = {}
+        self._custom_metadata.update(kwargs)