evotraders/backend/agents/base/evo_agent.py

# -*- coding: utf-8 -*-
"""EvoAgent - Core agent implementation for EvoTraders.

This module provides the main EvoAgent class built on AgentScope's ReActAgent,
with integrated tools, skills, and memory management based on CoPaw design.

Key features:
- Workspace-driven configuration from Markdown files
- Dynamic skill loading from skills/active directories
- Tool-guard security interception
- Hook system for extensibility
- Runtime skill and prompt reloading
"""
from __future__ import annotations

import logging
from pathlib import Path
from typing import Any, Dict, List, Optional, Type, TYPE_CHECKING

from agentscope.agent import ReActAgent
from agentscope.memory import InMemoryMemory
from agentscope.message import Msg
from agentscope.tool import Toolkit

from .tool_guard import ToolGuardMixin
from .hooks import (
    HookManager,
    BootstrapHook,
    MemoryCompactionHook,
    HOOK_PRE_REASONING,
)
from ..prompts.builder import (
    PromptBuilder,
    build_system_prompt_from_workspace,
)
from ..agent_workspace import load_agent_workspace_config
from ..skills_manager import SkillsManager

if TYPE_CHECKING:
    from agentscope.formatter import FormatterBase
    from agentscope.model import ModelWrapperBase

logger = logging.getLogger(__name__)


class EvoAgent(ToolGuardMixin, ReActAgent):
    """EvoAgent with integrated tools, skills, and memory management.

    This agent extends ReActAgent with:
    - Workspace-driven configuration from AGENTS.md/SOUL.md/PROFILE.md/etc.
    - Dynamic skill loading from skills/active directories
    - Tool-guard security interception (via ToolGuardMixin)
    - Hook system for extensibility (bootstrap, memory compaction)
    - Runtime skill and prompt reloading

    MRO note
    ~~~~~~~~
    ``ToolGuardMixin`` overrides ``_acting`` and ``_reasoning`` via
    Python's MRO: EvoAgent → ToolGuardMixin → ReActAgent.

    Example:
        agent = EvoAgent(
            agent_id="fundamentals_analyst",
            config_name="smoke_fullstack",
            workspace_dir=Path("runs/smoke_fullstack/agents/fundamentals_analyst"),
            model=model_instance,
            formatter=formatter_instance,
        )
    """

    def __init__(
        self,
        agent_id: str,
        config_name: str,
        workspace_dir: Path,
        model: "ModelWrapperBase",
        formatter: "FormatterBase",
        skills_manager: Optional[SkillsManager] = None,
        sys_prompt: Optional[str] = None,
        max_iters: int = 10,
        memory: Optional[Any] = None,
        enable_tool_guard: bool = True,
        enable_bootstrap_hook: bool = True,
        enable_memory_compaction: bool = False,
        memory_manager: Optional[Any] = None,
        memory_compact_threshold: Optional[int] = None,
        env_context: Optional[str] = None,
        prompt_files: Optional[List[str]] = None,
    ):
        """Initialize EvoAgent.

        Args:
            agent_id: Unique identifier for this agent
            config_name: Run configuration name (e.g., "smoke_fullstack")
            workspace_dir: Agent workspace directory containing markdown files
            model: LLM model instance
            formatter: Message formatter instance
            skills_manager: Optional SkillsManager instance
            sys_prompt: Optional override for system prompt
            max_iters: Maximum reasoning-acting iterations
            memory: Optional memory instance (defaults to InMemoryMemory)
            enable_tool_guard: Enable tool-guard security interception
            enable_bootstrap_hook: Enable bootstrap guidance on first interaction
            enable_memory_compaction: Enable automatic memory compaction
            memory_manager: Optional memory manager for compaction
            memory_compact_threshold: Token threshold for memory compaction
            env_context: Optional environment context to prepend to system prompt
            prompt_files: List of markdown files to load (defaults to standard set)
        """
        self.agent_id = agent_id
        self.config_name = config_name
        self.workspace_dir = Path(workspace_dir)
        self._skills_manager = skills_manager or SkillsManager()
        self._env_context = env_context
        self._prompt_files = prompt_files

        # Initialize tool guard
        if enable_tool_guard:
            self._init_tool_guard()

        # Load agent configuration from workspace
        self._agent_config = self._load_agent_config()

        # Build or use provided system prompt
        if sys_prompt is not None:
            self._sys_prompt = sys_prompt
        else:
            self._sys_prompt = self._build_system_prompt()

        # Create toolkit with skills
        toolkit = self._create_toolkit()

        # Initialize hook manager
        self._hook_manager = HookManager()

        # Initialize parent ReActAgent
        super().__init__(
            name=agent_id,
            model=model,
            sys_prompt=self._sys_prompt,
            toolkit=toolkit,
            memory=memory or InMemoryMemory(),
            formatter=formatter,
            max_iters=max_iters,
        )

        # Register hooks
        self._register_hooks(
            enable_bootstrap=enable_bootstrap_hook,
            enable_memory_compaction=enable_memory_compaction,
            memory_manager=memory_manager,
            memory_compact_threshold=memory_compact_threshold,
        )

        logger.info(
            "EvoAgent initialized: %s (workspace: %s)",
            agent_id,
            workspace_dir,
        )

    def _load_agent_config(self) -> Dict[str, Any]:
        """Load agent configuration from workspace.

        Returns:
            Agent configuration dictionary
        """
        config_path = self.workspace_dir / "agent.yaml"
        if config_path.exists():
            loaded = load_agent_workspace_config(config_path)
            return dict(loaded.values)
        return {}

    def _build_system_prompt(self) -> str:
        """Build system prompt from workspace markdown files.

        Uses PromptBuilder to load and combine AGENTS.md, SOUL.md,
        PROFILE.md, and other configured files.

        Returns:
            Complete system prompt string
        """
        prompt = build_system_prompt_from_workspace(
            workspace_dir=self.workspace_dir,
            enabled_files=self._prompt_files,
            agent_id=self.agent_id,
            extra_context=self._env_context,
        )
        return prompt

    def _create_toolkit(self) -> Toolkit:
        """Create and populate toolkit with agent skills.

        Loads skills from the agent's active skills directory and
        registers them with the toolkit.

        Returns:
            Configured Toolkit instance
        """
        toolkit = Toolkit(
            agent_skill_instruction=(
                "<system-info>You have access to specialized skills. "
                "Each skill lives in a directory and is described by SKILL.md. "
                "Follow the skill instructions when they are relevant to the current task."
                "</system-info>"
            ),
            agent_skill_template="- {name} (dir: {dir}): {description}",
        )

        # Register skills from active directory
        active_skills_dir = self._skills_manager.get_agent_active_root(
            self.config_name,
            self.agent_id,
        )

        if active_skills_dir.exists():
            for skill_dir in sorted(active_skills_dir.iterdir()):
                if skill_dir.is_dir() and (skill_dir / "SKILL.md").exists():
                    try:
                        toolkit.register_agent_skill(str(skill_dir))
                        logger.debug("Registered skill: %s", skill_dir.name)
                    except Exception as e:
                        logger.error(
                            "Failed to register skill '%s': %s",
                            skill_dir.name,
                            e,
                        )

        return toolkit

    def _register_hooks(
        self,
        enable_bootstrap: bool,
        enable_memory_compaction: bool,
        memory_manager: Optional[Any],
        memory_compact_threshold: Optional[int],
    ) -> None:
        """Register agent hooks.

        Args:
            enable_bootstrap: Enable bootstrap hook
            enable_memory_compaction: Enable memory compaction hook
            memory_manager: Memory manager instance
            memory_compact_threshold: Token threshold for compaction
        """
        # Bootstrap hook - checks BOOTSTRAP.md on first interaction
        if enable_bootstrap:
            bootstrap_hook = BootstrapHook(
                workspace_dir=self.workspace_dir,
                language="zh",
            )
            self._hook_manager.register(
                hook_type=HOOK_PRE_REASONING,
                hook_name="bootstrap",
                hook=bootstrap_hook,
            )
            logger.debug("Registered bootstrap hook")

        # Memory compaction hook
        if enable_memory_compaction and memory_manager is not None:
            compaction_hook = MemoryCompactionHook(
                memory_manager=memory_manager,
                memory_compact_threshold=memory_compact_threshold,
            )
            self._hook_manager.register(
                hook_type=HOOK_PRE_REASONING,
                hook_name="memory_compaction",
                hook=compaction_hook,
            )
            logger.debug("Registered memory compaction hook")

    async def _reasoning(self, **kwargs) -> Msg:
        """Override reasoning to execute pre-reasoning hooks.

        Args:
            **kwargs: Arguments for reasoning

        Returns:
            Response message
        """
        # Execute pre-reasoning hooks
        kwargs = await self._hook_manager.execute(
            hook_type=HOOK_PRE_REASONING,
            agent=self,
            kwargs=kwargs,
        )

        # Call parent (which may be ToolGuardMixin's _reasoning)
        return await super()._reasoning(**kwargs)

    def reload_skills(self, active_skill_dirs: Optional[List[Path]] = None) -> None:
        """Reload skills at runtime.

        Rebuilds the toolkit with current skills from the active directory.

        Args:
            active_skill_dirs: Optional list of specific skill directories to load
        """
        logger.info("Reloading skills for agent: %s", self.agent_id)

        # Create new toolkit
        new_toolkit = Toolkit(
            agent_skill_instruction=(
                "<system-info>You have access to specialized skills. "
                "Each skill lives in a directory and is described by SKILL.md. "
                "Follow the skill instructions when they are relevant to the current task."
                "</system-info>"
            ),
            agent_skill_template="- {name} (dir: {dir}): {description}",
        )

        # Register skills
        if active_skill_dirs is None:
            active_skills_dir = self._skills_manager.get_agent_active_root(
                self.config_name,
                self.agent_id,
            )
            if active_skills_dir.exists():
                active_skill_dirs = [
                    d for d in active_skills_dir.iterdir()
                    if d.is_dir() and (d / "SKILL.md").exists()
                ]
            else:
                active_skill_dirs = []

        for skill_dir in active_skill_dirs:
            if skill_dir.exists() and (skill_dir / "SKILL.md").exists():
                try:
                    new_toolkit.register_agent_skill(str(skill_dir))
                    logger.debug("Reloaded skill: %s", skill_dir.name)
                except Exception as e:
                    logger.error(
                        "Failed to reload skill '%s': %s",
                        skill_dir.name,
                        e,
                    )

        # Replace toolkit
        self.toolkit = new_toolkit
        logger.info("Skills reloaded for agent: %s", self.agent_id)

    def rebuild_sys_prompt(self) -> None:
        """Rebuild and replace the system prompt at runtime.

        Useful after updating AGENTS.md, SOUL.md, PROFILE.md, etc.
        to ensure the prompt reflects the latest configuration.

        Updates both self._sys_prompt and the first system-role
        message stored in self.memory.content.
        """
        logger.info("Rebuilding system prompt for agent: %s", self.agent_id)

        # Reload agent config in case it changed
        self._agent_config = self._load_agent_config()

        # Rebuild prompt
        self._sys_prompt = self._build_system_prompt()

        # Update memory if system message exists
        if hasattr(self, "memory") and self.memory.content:
            for msg, _marks in self.memory.content:
                if getattr(msg, "role", None) == "system":
                    msg.content = self._sys_prompt
                    logger.debug("Updated system message in memory")
                    break

        logger.info("System prompt rebuilt for agent: %s", self.agent_id)

    async def reply(
        self,
        msg: Msg | List[Msg] | None = None,
        structured_model: Optional[Type[Any]] = None,
    ) -> Msg:
        """Process a message and return a response.

        Args:
            msg: Input message(s) from user
            structured_model: Optional pydantic model for structured output

        Returns:
            Response message
        """
        # Handle list of messages
        if isinstance(msg, list):
            # Process each message in sequence
            for m in msg[:-1]:
                await self.memory.add(m)
            msg = msg[-1] if msg else None

        return await super().reply(msg=msg, structured_model=structured_model)

    def get_agent_info(self) -> Dict[str, Any]:
        """Get agent information.

        Returns:
            Dictionary with agent metadata
        """
        return {
            "agent_id": self.agent_id,
            "config_name": self.config_name,
            "workspace_dir": str(self.workspace_dir),
            "skills_count": len([
                s for s in self._skills_manager.list_active_skill_metadata(
                    self.config_name,
                    self.agent_id,
                )
            ]),
            "registered_hooks": self._hook_manager.list_hooks(),
        }


__all__ = ["EvoAgent"]