planner.py

"""
OmniCoreX Planning Module

Implementation of high-level planning components central to OmniCoreX’s ultra high-tech AI.
This module generates multi-step reasoning plans, manages dynamic tool invocations,
and orchestrates execution flows with adaptive and context-aware strategies.

Features:
- Hierarchical multi-step plan generation based on input goals and context.
- Dynamic selection and invocation of external or internal tools.
- Real-time plan adaptation and interruption handling.
- Support for conditional branching and iterative reasoning.
- Integration hooks for memory and knowledge modules.
"""

from typing import Any, Dict, List, Optional, Callable
import logging

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.DEBUG)

class PlanStep:
    """
    Represents a single actionable step in a reasoning plan.
    """
    def __init__(self, step_type: str, content: Any = None, tool_name: Optional[str] = None,
                 tool_params: Optional[Dict[str, Any]] = None, conditions: Optional[Dict[str, Any]] = None):
        """
        Initializes a PlanStep.

        Args:
            step_type: Type of action ('tool', 'respond', 'compute', 'wait', etc.).
            content: Core content or instruction associated with the step.
            tool_name: Name of the tool to invoke, if applicable.
            tool_params: Parameters for the tool invocation.
            conditions: Optional dict specifying conditions for execution or branching.
        """
        self.step_type = step_type
        self.content = content
        self.tool_name = tool_name
        self.tool_params = tool_params or {}
        self.conditions = conditions or {}

    def __repr__(self):
        return (f"PlanStep(type={self.step_type}, content={self.content}, "
                f"tool={self.tool_name}, params={self.tool_params}, conditions={self.conditions})")


class Planner:
    """
    Core planner for OmniCoreX responsible for generating and managing multi-step reasoning plans.
    """

    def __init__(self,
                 tool_registry: Optional[Dict[str, Callable[..., Any]]] = None,
                 memory_provider: Optional[Callable[[str], Any]] = None):
        """
        Initializes the Planner.

        Args:
            tool_registry: Dict mapping tool names to callable implementations.
            memory_provider: Callable to access memory/context data by key or query.
        """
        self.tool_registry = tool_registry or {}
        self.memory_provider = memory_provider
        self.current_plan: List[PlanStep] = []
        self.current_step_index = 0

    def generate_plan(self, goal: str, context: Optional[str] = None) -> List[PlanStep]:
        """
        Generates a multi-step reasoning plan to achieve the specified goal.

        Args:
            goal: The high-level objective or query.
            context: Optional context information or memory.

        Returns:
            List of PlanStep objects defining the detailed plan.
        """
        logger.debug(f"Generating plan for goal: {goal} with context: {context}")

        # Placeholder plan generation that can be replaced with advanced AI methods
        plan = []

        # Step 1: Analyze Goal
        plan.append(PlanStep(step_type="compute", content=f"Analyze input goal: {goal}"))

        # Step 2: Look up context if available
        if context:
            plan.append(PlanStep(step_type="compute", content=f"Integrate contextual information."))

        # Step 3: Decide tool usage dynamically based on keywords (demo logic)
        if "search" in goal.lower():
            plan.append(PlanStep(step_type="tool",
                                 tool_name="web_search",
                                 tool_params={"query": goal}))
        elif "calculate" in goal.lower() or "compute" in goal.lower():
            plan.append(PlanStep(step_type="tool",
                                 tool_name="calculator",
                                 tool_params={"expression": goal}))
        else:
            plan.append(PlanStep(step_type="respond",
                                 content=f"Preliminary response: Processing {goal}"))

        # Step 4: Finalize response
        plan.append(PlanStep(step_type="respond", content="Present final response to user."))

        logger.debug(f"Generated plan: {plan}")
        self.current_plan = plan
        self.current_step_index = 0
        return plan

    def get_next_step(self) -> Optional[PlanStep]:
        """
        Gets the next step in the current plan, or None if plan is complete.

        Returns:
            The next PlanStep or None.
        """
        if self.current_step_index >= len(self.current_plan):
            logger.debug("Plan complete, no further steps.")
            return None
        step = self.current_plan[self.current_step_index]
        self.current_step_index += 1
        logger.debug(f"Advancing to next plan step: {step}")
        return step

    def execute_step(self, step: PlanStep) -> Any:
        """
        Executes the provided plan step, invoking tools or performing computations.

        Args:
            step: The PlanStep to execute.

        Returns:
            The result of the step's execution.
        """
        logger.debug(f"Executing plan step: {step}")

        if step.step_type == "tool":
            # Invoke registered tool
            tool_func = self.tool_registry.get(step.tool_name)
            if tool_func:
                try:
                    result = tool_func(**step.tool_params)
                    logger.debug(f"Tool '{step.tool_name}' returned: {result}")
                    return result
                except Exception as e:
                    logger.error(f"Error executing tool '{step.tool_name}': {e}")
                    return {"error": str(e)}
            else:
                error_msg = f"Tool '{step.tool_name}' not found."
                logger.error(error_msg)
                return {"error": error_msg}

        elif step.step_type == "compute":
            # Placeholder compute action: simulate or integrate with reasoning engine
            logger.debug(f"Computing: {step.content}")
            return f"Computed: {step.content}"

        elif step.step_type == "respond":
            # Prepare response action content
            logger.debug(f"Responding with content: {step.content}")
            return step.content

        elif step.step_type == "wait":
            # Wait or pause action
            duration = step.tool_params.get("duration", 1)
            logger.debug(f"Waiting for {duration} seconds.")
            import time
            time.sleep(duration)
            return f"Waited {duration} seconds"

        else:
            logger.warning(f"Unknown step_type '{step.step_type}'. No action performed.")
            return None

    def interrupt_and_replan(self, new_context: str, new_goal: Optional[str] = None) -> List[PlanStep]:
        """
        Interrupts current plan and generates a new plan based on updated context or goal.

        Args:
            new_context: New contextual data influencing replanning.
            new_goal: Optional new goal to re-plan for.

        Returns:
            New list of PlanStep objects forming the refreshed plan.
        """
        logger.info(f"Interrupting current plan for replanning with new context: {new_context} and goal: {new_goal}")
        goal = new_goal if new_goal else "Replan goal"
        return self.generate_plan(goal=goal, context=new_context)


# Example Tool Implementations (to be registered)
def web_search(query: str) -> Dict[str, Any]:
    # Simulate a mock web search tool
    logger.info(f"Web search tool searching for query: {query}")
    return {"results": [f"Result 1 for {query}", f"Result 2 for {query}"]}


def calculator(expression: str) -> Dict[str, Any]:
    # Simulate a calculator tool evaluating simple math expression
    logger.info(f"Calculator tool evaluating: {expression}")
    try:
        result = eval(expression, {"__builtins__": {}})
        return {"result": result}
    except Exception as e:
        return {"error": str(e)}


if __name__ == "__main__":
    # Minimal demonstration of planning and execution

    planner = Planner(tool_registry={
        "web_search": web_search,
        "calculator": calculator
    })

    goal = "Search latest AI research papers"
    context = "User interested in machine learning breakthroughs."

    plan = planner.generate_plan(goal, context)

    for step in plan:
        output = planner.execute_step(step)
        print(f"Step result: {output}")

    # Demonstrate interruption and replanning
    new_context = "Urgent request: calculate 2+2"
    new_plan = planner.interrupt_and_replan(new_context, new_goal="Calculate 2 + 2")
    print("\nNew plan after interruption:")
    for step in new_plan:
        output = planner.execute_step(step)
        print(f"Step result: {output}")