FusionAGI/fusionagi/self_improvement/loop.py

"""AGI loop: wires self-correction, auto-recommend, and auto-training to events."""

from typing import Any, Callable

from fusionagi.schemas.task import TaskState
from fusionagi.schemas.recommendation import Recommendation, TrainingSuggestion
from fusionagi.core.event_bus import EventBus
from fusionagi._logger import logger

from fusionagi.self_improvement.correction import (
    SelfCorrectionLoop,
    StateManagerLike,
    OrchestratorLike,
    CriticLike,
)
from fusionagi.self_improvement.recommender import AutoRecommender
from fusionagi.self_improvement.training import AutoTrainer, ReflectiveMemoryLike


class FusionAGILoop:
    """
    High-level AGI loop: subscribes to task_state_changed and reflection_done,
    runs self-correction on failures, and runs auto-recommend + auto-training
    after reflection. Composes the world's most advanced agentic AGI self-improvement pipeline.
    """

    def __init__(
        self,
        event_bus: EventBus,
        state_manager: StateManagerLike,
        orchestrator: OrchestratorLike,
        critic_agent: CriticLike,
        reflective_memory: ReflectiveMemoryLike | None = None,
        *,
        auto_retry_on_failure: bool = False,
        max_retries_per_task: int = 2,
        on_recommendations: Callable[[list[Recommendation]], None] | None = None,
        on_training_suggestions: Callable[[list[TrainingSuggestion]], None] | None = None,
    ) -> None:
        """
        Initialize the FusionAGI loop.

        Args:
            event_bus: Event bus to subscribe to task_state_changed and reflection_done.
            state_manager: State manager for task state and traces.
            orchestrator: Orchestrator for plan and state transitions.
            critic_agent: Critic agent for evaluate_request -> evaluation_ready.
            reflective_memory: Optional reflective memory for lessons/heuristics.
            auto_retry_on_failure: If True, on FAILED transition prepare_retry automatically.
            max_retries_per_task: Max retries per task when auto_retry_on_failure is True.
            on_recommendations: Optional callback to receive recommendations (e.g. log or UI).
            on_training_suggestions: Optional callback to receive training suggestions.
        """
        self._event_bus = event_bus
        self._state = state_manager
        self._orchestrator = orchestrator
        self._critic = critic_agent
        self._memory = reflective_memory
        self._auto_retry = auto_retry_on_failure
        self._on_recs = on_recommendations
        self._on_training = on_training_suggestions

        self._correction = SelfCorrectionLoop(
            state_manager=state_manager,
            orchestrator=orchestrator,
            critic_agent=critic_agent,
            max_retries_per_task=max_retries_per_task,
        )
        self._recommender = AutoRecommender(reflective_memory=reflective_memory)
        self._trainer = AutoTrainer(reflective_memory=reflective_memory)

        self._event_bus.subscribe("task_state_changed", self._on_task_state_changed)
        self._event_bus.subscribe("reflection_done", self._on_reflection_done)
        logger.info("FusionAGILoop: subscribed to task_state_changed and reflection_done")

    def _on_task_state_changed(self, event_type: str, payload: dict[str, Any]) -> None:
        """On FAILED, optionally run self-correction and prepare retry."""
        try:
            to_state = payload.get("to_state")
            task_id = payload.get("task_id", "")
            if to_state != TaskState.FAILED.value or not task_id:
                return
            if self._auto_retry:
                ok, _ = self._correction.suggest_retry(task_id)
                if ok:
                    self._correction.prepare_retry(task_id)
            else:
                recs = self._correction.correction_recommendations(task_id)
                if recs and self._on_recs:
                    self._on_recs(recs)
        except Exception:
            logger.exception(
                "FusionAGILoop: _on_task_state_changed failed (best-effort)",
                extra={"event_type": event_type},
            )

    def _on_reflection_done(self, event_type: str, payload: dict[str, Any]) -> None:
        """After reflection, run auto-recommend and auto-training."""
        try:
            task_id = payload.get("task_id") or ""
            evaluation = payload.get("evaluation") or {}
            recs = self._recommender.recommend(
                task_id=task_id or None,
                evaluation=evaluation,
                include_lessons=True,
            )
            if self._on_recs:
                try:
                    self._on_recs(recs)
                except Exception:
                    logger.exception("FusionAGILoop: on_recommendations callback failed")
            suggestions = self._trainer.run_auto_training(
                task_id=task_id or None,
                evaluation=evaluation,
                apply_heuristics=True,
            )
            if self._on_training:
                try:
                    self._on_training(suggestions)
                except Exception:
                    logger.exception("FusionAGILoop: on_training_suggestions callback failed")
        except Exception:
            logger.exception(
                "FusionAGILoop: _on_reflection_done failed (best-effort)",
                extra={"event_type": event_type},
            )

    def run_after_reflection(
        self,
        task_id: str,
        evaluation: dict[str, Any],
    ) -> tuple[list[Recommendation], list[TrainingSuggestion]]:
        """
        Run auto-recommend and auto-training after a reflection (e.g. when
        not using reflection_done event). Returns (recommendations, training_suggestions).
        """
        recs = self._recommender.recommend(
            task_id=task_id,
            evaluation=evaluation,
            include_lessons=True,
        )
        suggestions = self._trainer.run_auto_training(
            task_id=task_id,
            evaluation=evaluation,
            apply_heuristics=True,
        )
        return recs, suggestions

    def unsubscribe(self) -> None:
        """Unsubscribe from event bus (for cleanup)."""
        self._event_bus.unsubscribe("task_state_changed", self._on_task_state_changed)
        self._event_bus.unsubscribe("reflection_done", self._on_reflection_done)
        logger.info("FusionAGILoop: unsubscribed from events")