agnt5 0.2.8a3__cp310-abi3-manylinux_2_34_x86_64.whl → 0.2.8a5__cp310-abi3-manylinux_2_34_x86_64.whl

agnt5/__init__.py

@@ -7,7 +7,6 @@ with built-in durability guarantees and state management.
 
 from ._compat import _import_error, _rust_available
 from .agent import Agent, AgentContext, AgentRegistry, AgentResult, Handoff, agent, handoff
-from .agent_session import AgentSession
 from .client import Client, RunError
 from .context import Context
 from .function import FunctionContext
@@ -15,7 +14,7 @@ from .workflow import WorkflowContext
 from .entity import (
     Entity,
     EntityRegistry,
-    EntityStateManager,
+    EntityStateAdapter,
     EntityType,
     create_entity_context,
     with_entity_context,
@@ -27,9 +26,10 @@ from .exceptions import (
     ExecutionError,
     RetryError,
     StateError,
+    WaitingForUserInputException,
 )
 from .function import FunctionRegistry, function
-from .tool import Tool, ToolRegistry, tool
+from .tool import AskUserTool, RequestApprovalTool, Tool, ToolRegistry, tool
 from .types import BackoffPolicy, BackoffType, FunctionConfig, RetryPolicy, WorkflowConfig
 from .version import _get_version
 from .worker import Worker
@@ -55,7 +55,7 @@ __all__ = [
     "Entity",
     "EntityType",
     "EntityRegistry",
-    "EntityStateManager",
+    "EntityStateAdapter",
     "with_entity_context",
     "create_entity_context",
     "workflow",
@@ -63,11 +63,12 @@ __all__ = [
     "tool",
     "Tool",
     "ToolRegistry",
+    "AskUserTool",
+    "RequestApprovalTool",
     "agent",
     "Agent",
     "AgentRegistry",
     "AgentResult",
-    "AgentSession",
     "Handoff",
     "handoff",
     # Types
@@ -83,6 +84,7 @@ __all__ = [
     "RetryError",
     "StateError",
     "CheckpointError",
+    "WaitingForUserInputException",
     "RunError",
     # Language Model (Simplified API)
     "lm",

agnt5/agent.py

@@ -9,6 +9,7 @@ from __future__ import annotations
 import functools
 import json
 import logging
+import time
 from typing import Any, Callable, Dict, List, Optional
 
 from .context import Context
@@ -79,42 +80,46 @@ class AgentContext(Context):
         self._agent_name = agent_name
         self._session_id = session_id or run_id
 
-        # Determine state manager based on parent context
-        from .entity import EntityStateManager, _get_state_manager
+        # Determine state adapter based on parent context
+        from .entity import EntityStateAdapter, _get_state_adapter
 
         if state_manager:
-            # Explicit state manager provided
-            self._state_manager = state_manager
-            logger.debug(f"AgentContext using provided state manager")
+            # Explicit state adapter provided (parameter name kept for backward compat)
+            self._state_adapter = state_manager
+            logger.debug(f"AgentContext using provided state adapter")
         elif parent_context:
-            # Try to inherit state manager from parent
+            # Try to inherit state adapter from parent
             try:
                 # Check if parent is WorkflowContext or AgentContext
                 if hasattr(parent_context, '_workflow_entity'):
-                    # WorkflowContext - get state manager from worker context
-                    self._state_manager = _get_state_manager()
+                    # WorkflowContext - get state adapter from worker context
+                    self._state_adapter = _get_state_adapter()
                     logger.debug(f"AgentContext inheriting state from WorkflowContext")
-                elif hasattr(parent_context, '_state_manager'):
-                    # Parent AgentContext - share state manager
-                    self._state_manager = parent_context._state_manager
+                elif hasattr(parent_context, '_state_adapter'):
+                    # Parent AgentContext - share state adapter
+                    self._state_adapter = parent_context._state_adapter
                     logger.debug(f"AgentContext inheriting state from parent AgentContext")
+                elif hasattr(parent_context, '_state_manager'):
+                    # Backward compatibility: parent has old _state_manager
+                    self._state_adapter = parent_context._state_manager
+                    logger.debug(f"AgentContext inheriting state from parent (legacy)")
                 else:
-                    # FunctionContext or base Context - create new state manager
-                    self._state_manager = EntityStateManager()
-                    logger.debug(f"AgentContext created new state manager (parent has no state)")
+                    # FunctionContext or base Context - create new state adapter
+                    self._state_adapter = EntityStateAdapter()
+                    logger.debug(f"AgentContext created new state adapter (parent has no state)")
             except RuntimeError as e:
-                # _get_state_manager() failed (not in worker context) - create standalone
-                self._state_manager = EntityStateManager()
-                logger.debug(f"AgentContext created standalone state manager (not in worker context)")
+                # _get_state_adapter() failed (not in worker context) - create standalone
+                self._state_adapter = EntityStateAdapter()
+                logger.debug(f"AgentContext created standalone state adapter (not in worker context)")
         else:
             # Try to get from worker context first
             try:
-                self._state_manager = _get_state_manager()
-                logger.debug(f"AgentContext got state manager from worker context")
+                self._state_adapter = _get_state_adapter()
+                logger.debug(f"AgentContext got state adapter from worker context")
             except RuntimeError as e:
-                # Standalone - create new state manager
-                self._state_manager = EntityStateManager()
-                logger.debug(f"AgentContext created standalone state manager")
+                # Standalone - create new state adapter
+                self._state_adapter = EntityStateAdapter()
+                logger.debug(f"AgentContext created standalone state adapter")
 
         # Conversation key for state storage (used for in-memory state)
         self._conversation_key = f"agent:{agent_name}:{self._session_id}:messages"
@@ -127,24 +132,23 @@ class AgentContext(Context):
         """
         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:
-            EntityState instance for state operations
+            Dict-like object for state operations
 
         Example:
-            # Store conversation history
-            messages = ctx.state.get(f"agent:{agent_name}:{session_id}:messages", [])
-            messages.append({"role": "user", "content": "Hello"})
-            ctx.state.set(f"agent:{agent_name}:{session_id}:messages", messages)
-
-            # Store agent-specific data
-            ctx.state.set("research_results", data)
+            # Store agent-specific data (in-memory only)
+            ctx.state["research_results"] = data
+            ctx.state["iteration_count"] = 5
         """
-        from .entity import EntityState
-
-        # Use agent's conversation key as the state key
-        state_key = ("agent", self._conversation_key)
-        state_dict = self._state_manager.get_or_create_state(state_key)
-        return EntityState(state_dict)
+        # 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:
@@ -155,45 +159,30 @@ class AgentContext(Context):
         """
         Retrieve conversation history from state, loading from database if needed.
 
+        Uses the EntityStateAdapter which delegates to Rust core for cache-first loading.
+
         Returns:
             List of Message objects from conversation history
         """
-        # Try to load from database first if not in memory
         entity_type = "AgentSession"
-        entity_key = self._entity_key  # Use entity_key without :messages suffix
-        state_key = (entity_type, entity_key)
-
-        # Check if we need to load from platform
-        if state_key not in self._state_manager._states and self._state_manager._rust_manager:
-            try:
-                result = await self._state_manager._rust_manager.py_load_state(
-                    entity_type, entity_key
-                )
-                found, state_json_bytes, version = result
-                if found:
-                    import json
-                    state_json = state_json_bytes.decode('utf-8') if isinstance(state_json_bytes, bytes) else state_json_bytes
-                    session_data = json.loads(state_json)
-
-                    # Extract messages from session object (might be old format or new format)
-                    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:
-                        messages_data = []
-
-                    # Load messages into in-memory state (store just messages, not full session object)
-                    self._state_manager.load_state_from_platform(state_key, json.dumps(messages_data), version)
-                    logger.info(f"Loaded conversation history from database: {entity_key} (version {version})")
-            except Exception as e:
-                logger.warning(f"Failed to load conversation history from database: {e}")
-
-        # Get from in-memory state
-        messages_data = self.state.get(self._conversation_key, [])
-        logger.debug(f"Loaded {len(messages_data)} messages from conversation history")
+        entity_key = self._entity_key
+
+        # Load session data via adapter (Rust handles cache + platform load)
+        session_data = await self._state_adapter.load_state(entity_type, entity_key)
+
+        # 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"]
+            logger.debug(f"Loaded {len(messages_data)} messages from session {entity_key}")
+        elif isinstance(session_data, list):
+            # Old format - just messages array
+            messages_data = session_data
+            logger.debug(f"Loaded {len(messages_data)} messages (legacy format)")
+        else:
+            # No messages found
+            messages_data = []
+            logger.debug(f"No conversation history found for {entity_key}")
 
         # Convert dict representations back to Message objects
         messages = []
@@ -221,6 +210,8 @@ class AgentContext(Context):
         """
         Save conversation history to state and persist to database.
 
+        Uses the EntityStateAdapter which delegates to Rust core for version-checked saves.
+
         Args:
             messages: List of Message objects to persist
         """
@@ -231,49 +222,129 @@ class AgentContext(Context):
         for msg in messages:
             messages_data.append({
                 "role": msg.role.value if hasattr(msg.role, 'value') else str(msg.role),
-                "content": msg.content
+                "content": msg.content,
+                "timestamp": time.time()  # Add timestamp for each message
             })
 
-        # Save to in-memory state
-        self.state.set(self._conversation_key, messages_data)
+        entity_type = "AgentSession"
+        entity_key = self._entity_key
 
-        # Persist to database via Rust EntityStateManager
-        if self._state_manager._rust_manager:
-            try:
-                import json
-                import time
-                entity_type = "AgentSession"
-                entity_key = self._entity_key  # Use entity_key without :messages suffix
-                state_key = (entity_type, entity_key)
-
-                # Get current version
-                expected_version = self._state_manager._versions.get(state_key, 0)
-
-                # Build session object with metadata for the sessions API
-                now = time.time()
-                session_data = {
-                    "session_id": self._session_id,
-                    "agent_name": self._agent_name,
-                    "created_at": now if expected_version == 0 else None,  # Only set on first save
-                    "last_message_time": now,
-                    "message_count": len(messages_data),
-                    "messages": messages_data,
-                    "metadata": {}
-                }
+        # Load current state with version for optimistic locking
+        current_state, current_version = await self._state_adapter.load_with_version(
+            entity_type, entity_key
+        )
 
-                # Serialize to JSON bytes
-                state_json = json.dumps(session_data).encode('utf-8')
+        # Build session object with metadata
+        now = time.time()
 
-                # Save to platform
-                new_version = await self._state_manager._rust_manager.py_save_state(
-                    entity_type, entity_key, state_json, expected_version
-                )
+        # 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)
+        try:
+            new_version = await self._state_adapter.save_state(
+                entity_type,
+                entity_key,
+                session_data,
+                current_version
+            )
+            logger.info(
+                f"Persisted conversation history: {entity_key} (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
 
-                # Update version tracking
-                self._state_manager._versions[state_key] = new_version
-                logger.info(f"Persisted conversation history to database: {entity_key} (version {expected_version} -> {new_version})")
-            except Exception as e:
-                logger.error(f"Failed to persist conversation history to database: {e}")
+        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')}")
+            ```
+        """
+        entity_type = "AgentSession"
+        entity_key = self._entity_key
+
+        # Load session data
+        session_data = await self._state_adapter.load_state(entity_type, entity_key)
+
+        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)
 
 
 class Handoff: