oagi-core 0.9.0__py3-none-any.whl

oagi/agent/tasker/taskee_agent.py

@@ -0,0 +1,395 @@
+# -----------------------------------------------------------------------------
+#  Copyright (c) OpenAGI Foundation
+#  All rights reserved.
+#
+#  This file is part of the official API project.
+#  Licensed under the MIT License.
+# -----------------------------------------------------------------------------
+
+import logging
+from datetime import datetime
+from typing import Any
+
+from oagi import AsyncTask
+from oagi.types import AsyncActionHandler, AsyncImageProvider
+
+from ..protocol import AsyncAgent
+from .memory import PlannerMemory
+from .models import Action, ExecutionResult
+from .planner import Planner
+
+logger = logging.getLogger(__name__)
+
+
+class TaskeeAgent(AsyncAgent):
+    """Executes a single todo with planning and reflection capabilities.
+
+    This agent uses a Planner to:
+    1. Convert a todo into a clear actionable instruction
+    2. Execute the instruction using OAGI API
+    3. Periodically reflect on progress and adjust approach
+    4. Generate execution summaries
+    """
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        base_url: str | None = None,
+        model: str = "lux-v1",
+        max_steps_per_subtask: int = 10,
+        reflection_interval: int = 20,
+        temperature: float = 0.0,
+        planner: Planner | None = None,
+        external_memory: PlannerMemory | None = None,
+        todo_index: int | None = None,
+    ):
+        """Initialize the taskee agent.
+
+        Args:
+            api_key: OAGI API key
+            base_url: OAGI API base URL
+            model: Model to use for vision tasks
+            max_steps_per_subtask: Maximum steps before reinitializing task
+            reflection_interval: Number of actions before triggering reflection
+            temperature: Sampling temperature
+            planner: Planner for planning and reflection
+            external_memory: External memory from parent agent
+            todo_index: Index of the todo being executed
+        """
+        self.api_key = api_key
+        self.base_url = base_url
+        self.model = model
+        self.max_steps_per_subtask = max_steps_per_subtask
+        self.reflection_interval = reflection_interval
+        self.temperature = temperature
+        self.planner = planner or Planner()
+        self.external_memory = external_memory
+        self.todo_index = todo_index
+
+        # Internal state
+        self.task: AsyncTask | None = None
+        self.current_todo: str = ""
+        self.current_instruction: str = ""
+        self.actions: list[Action] = []
+        self.total_actions = 0
+        self.since_reflection = 0
+        self.success = False
+
+    async def execute(
+        self,
+        instruction: str,
+        action_handler: AsyncActionHandler,
+        image_provider: AsyncImageProvider,
+    ) -> bool:
+        """Execute the todo using planning and reflection.
+
+        Args:
+            instruction: The todo description to execute
+            action_handler: Handler for executing actions
+            image_provider: Provider for capturing screenshots
+
+        Returns:
+            True if successful, False otherwise
+        """
+        self.current_todo = instruction
+        self.actions = []
+        self.total_actions = 0
+        self.since_reflection = 0
+        self.success = False
+
+        try:
+            # Initial planning
+            await self._initial_plan(image_provider)
+
+            # Main execution loop with reinitializations
+            max_total_steps = self.max_steps_per_subtask * 3  # Allow up to 3 reinits
+            remaining_steps = max_total_steps
+
+            while remaining_steps > 0 and not self.success:
+                # Execute subtask
+                steps_taken = await self._execute_subtask(
+                    min(self.max_steps_per_subtask, remaining_steps),
+                    action_handler,
+                    image_provider,
+                )
+                remaining_steps -= steps_taken
+
+                # Check if we should continue
+                if not self.success and remaining_steps > 0:
+                    # Reflect and potentially get new instruction
+                    should_continue = await self._reflect_and_decide(image_provider)
+                    if not should_continue:
+                        break
+
+            # Generate final summary
+            await self._generate_summary()
+
+            return self.success
+
+        except Exception as e:
+            logger.error(f"Error executing todo: {e}")
+            self._record_action(
+                action_type="error",
+                target=None,
+                reasoning=str(e),
+            )
+            return False
+        finally:
+            # Clean up task
+            if self.task:
+                await self.task.close()
+                self.task = None
+
+    async def _initial_plan(self, image_provider: AsyncImageProvider) -> None:
+        """Generate initial plan for the todo.
+
+        Args:
+            image_provider: Provider for capturing screenshots
+        """
+        logger.info("Generating initial plan for todo")
+
+        # Capture initial screenshot
+        screenshot = await image_provider()
+
+        # Get context from external memory if available
+        context = self._get_context()
+
+        # Generate plan using LLM planner
+        plan_output = await self.planner.initial_plan(
+            self.current_todo,
+            context,
+            screenshot,
+            memory=self.external_memory,
+            todo_index=self.todo_index,
+        )
+
+        # Record planning action
+        self._record_action(
+            action_type="plan",
+            target="initial",
+            reasoning=plan_output.reasoning,
+            result=plan_output.instruction,
+        )
+
+        # Set current instruction
+        self.current_instruction = plan_output.instruction
+        logger.info(f"Initial instruction: {self.current_instruction}")
+
+        # Handle subtodos if any
+        if plan_output.subtodos:
+            logger.info(f"Planner created {len(plan_output.subtodos)} subtodos")
+            # Could potentially add these to memory for tracking
+
+    async def _execute_subtask(
+        self,
+        max_steps: int,
+        action_handler: AsyncActionHandler,
+        image_provider: AsyncImageProvider,
+    ) -> int:
+        """Execute a subtask with the current instruction.
+
+        Args:
+            max_steps: Maximum steps for this subtask
+            action_handler: Handler for executing actions
+            image_provider: Provider for capturing screenshots
+
+        Returns:
+            Number of steps taken
+        """
+        logger.info(f"Executing subtask with max {max_steps} steps")
+
+        # Use async with for automatic resource management
+        async with AsyncTask(
+            api_key=self.api_key,
+            base_url=self.base_url,
+            model=self.model,
+            temperature=self.temperature,
+        ) as task:
+            # Store reference for potential cleanup in execute's finally block
+            self.task = task
+
+            # Initialize task with current instruction
+            await task.init_task(self.current_instruction)
+
+            steps_taken = 0
+            for step_num in range(max_steps):
+                # Capture screenshot
+                screenshot = await image_provider()
+
+                # Get next step from OAGI
+                try:
+                    step = await task.step(screenshot, instruction=None)
+                except Exception as e:
+                    logger.error(f"Error getting step from OAGI: {e}")
+                    self._record_action(
+                        action_type="error",
+                        target="oagi_step",
+                        reasoning=str(e),
+                    )
+                    break
+
+                # Record OAGI actions
+                if step.actions:
+                    for action in step.actions:
+                        self._record_action(
+                            action_type=action.type.lower(),
+                            target=action.argument,
+                            reasoning=step.reason,
+                        )
+
+                    # Execute actions
+                    await action_handler(step.actions)
+                    self.total_actions += len(step.actions)
+                    self.since_reflection += len(step.actions)
+
+                steps_taken += 1
+
+                # Check if task is complete
+                if step.stop:
+                    logger.info("OAGI signaled task completion")
+                    self.success = True
+                    break
+
+                # Check if reflection is needed
+                if self.since_reflection >= self.reflection_interval:
+                    logger.info("Reflection interval reached")
+                    break
+
+            # Task will be automatically closed by async with context manager
+            # Clear reference after context manager closes it
+            self.task = None
+            return steps_taken
+
+    async def _reflect_and_decide(self, image_provider: AsyncImageProvider) -> bool:
+        """Reflect on progress and decide whether to continue.
+
+        Args:
+            image_provider: Provider for capturing screenshots
+
+        Returns:
+            True to continue, False to stop
+        """
+        logger.info("Reflecting on progress")
+
+        # Capture current screenshot
+        screenshot = await image_provider()
+
+        # Get context
+        context = self._get_context()
+        context["current_todo"] = self.current_todo
+
+        # Get recent actions for reflection
+        recent_actions = self.actions[-self.since_reflection :]
+
+        # Reflect using planner
+        reflection = await self.planner.reflect(
+            recent_actions,
+            context,
+            screenshot,
+            memory=self.external_memory,
+            todo_index=self.todo_index,
+            current_instruction=self.current_instruction,
+        )
+
+        # Record reflection
+        self._record_action(
+            action_type="reflect",
+            target=None,
+            reasoning=reflection.reasoning,
+            result=("continue" if reflection.continue_current else "pivot"),
+        )
+
+        # Update success assessment
+        if reflection.success_assessment:
+            self.success = True
+            logger.info("Reflection indicates task is successful")
+            return False
+
+        # Reset reflection counter
+        self.since_reflection = 0
+
+        # Update instruction if needed
+        if not reflection.continue_current and reflection.new_instruction:
+            logger.info(f"Pivoting to new instruction: {reflection.new_instruction}")
+            self.current_instruction = reflection.new_instruction
+            return True
+
+        return reflection.continue_current
+
+    async def _generate_summary(self) -> None:
+        """Generate execution summary."""
+        logger.info("Generating execution summary")
+
+        context = self._get_context()
+        context["current_todo"] = self.current_todo
+
+        summary = await self.planner.summarize(
+            self.actions,
+            context,
+            memory=self.external_memory,
+            todo_index=self.todo_index,
+        )
+
+        # Record summary
+        self._record_action(
+            action_type="summary",
+            target=None,
+            reasoning=summary,
+        )
+
+        logger.info(f"Execution summary: {summary}")
+
+    def _record_action(
+        self,
+        action_type: str,
+        target: str | None,
+        reasoning: str | None = None,
+        result: str | None = None,
+    ) -> None:
+        """Record an action to the history.
+
+        Args:
+            action_type: Type of action
+            target: Target of the action
+            reasoning: Reasoning for the action
+            result: Result of the action
+        """
+        action = Action(
+            timestamp=datetime.now().isoformat(),
+            action_type=action_type,
+            target=target,
+            reasoning=reasoning,
+            result=result,
+            details={},
+        )
+        self.actions.append(action)
+
+    def _get_context(self) -> dict[str, Any]:
+        """Get execution context.
+
+        Returns:
+            Dictionary with context information
+        """
+        if self.external_memory:
+            return self.external_memory.get_context()
+        return {}
+
+    def return_execution_results(self) -> ExecutionResult:
+        """Return the execution results.
+
+        Returns:
+            ExecutionResult with success status, actions, and summary
+        """
+        # Find summary in actions
+        summary = ""
+        for action in reversed(self.actions):
+            if action.action_type == "summary":
+                summary = action.reasoning or ""
+                break
+
+        return ExecutionResult(
+            success=self.success,
+            actions=self.actions,
+            summary=summary,
+            total_steps=self.total_actions,
+        )