python -m venv venv
source venv/bin/activate
pip install openai pydantic networkx rich
export OPENAI_API_KEY="sk-your-key"

# models.py
from pydantic import BaseModel
from typing import List, Optional
from enum import Enum
import logging

logger = logging.getLogger(__name__)


class StepStatus(str, Enum):
    """Status of a plan step."""
    PENDING = "pending"
    IN_PROGRESS = "in_progress"
    COMPLETED = "completed"
    FAILED = "failed"


class PlanStep(BaseModel):
    """A single step in a plan."""
    id: str
    description: str
    tool_needed: Optional[str] = None
    status: StepStatus = StepStatus.PENDING
    result: Optional[str] = None
    dependencies: List[str] = []


class Plan(BaseModel):
    """A complete plan for achieving a goal."""
    id: str
    goal: str
    steps: List[PlanStep]
    score: float = 0.0
    reasoning: str = ""


class PlanEvaluation(BaseModel):
    """Evaluation scores for a plan."""
    plan_id: str
    feasibility_score: float
    efficiency_score: float
    completeness_score: float
    overall_score: float
    feedback: str

# planner/cot_planner.py
from openai import OpenAI
from models import Plan, PlanStep
import json
import logging

logger = logging.getLogger(__name__)

COT_SYSTEM = """You are a planning expert. Break down goals into clear, 
executable steps. For each step, specify:
1. What needs to be done
2. What tool/information is needed
3. What dependencies exist

Always think step-by-step before generating the plan."""


class CoTPlanner:
    """
    Chain-of-Thought planner that generates a single plan.
    
    Uses step-by-step reasoning to create actionable plans.
    Fast but may miss better alternatives.
    """
    
    def __init__(self, model: str = "gpt-4-turbo-preview"):
        self.client = OpenAI()
        self.model = model

    def generate_plan(self, goal: str, available_tools: list[str]) -> Plan:
        """
        Generate a plan using chain-of-thought reasoning.
        
        Args:
            goal: The objective to achieve
            available_tools: List of available tool names
            
        Returns:
            Plan with ordered steps
        """
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": COT_SYSTEM},
                {"role": "user", "content": f"""Goal: {goal}
Available tools: {', '.join(available_tools)}

Think step-by-step, then generate a JSON plan with these fields:
- goal: the original goal
- steps: array of objects with id, description, tool_needed, dependencies
- reasoning: your chain-of-thought reasoning

Return ONLY valid JSON."""},
            ],
            temperature=0.0,
        )

        content = response.choices[0].message.content
        plan_data = self._parse_response(content)
        
        steps = [
            PlanStep(
                id=s["id"],
                description=s["description"],
                tool_needed=s.get("tool_needed"),
                dependencies=s.get("dependencies", []),
            )
            for s in plan_data["steps"]
        ]
        
        plan = Plan(
            id="plan_cot_001",
            goal=goal,
            steps=steps,
            reasoning=plan_data.get("reasoning", ""),
        )
        
        logger.info(f"Generated CoT plan with {len(steps)} steps")
        return plan

    def _parse_response(self, content: str) -> dict:
        """Parse LLM response as JSON."""
        try:
            return json.loads(content)
        except json.JSONDecodeError:
            import re
            json_match = re.search(r'\{[\s\S]*\}', content)
            if json_match:
                return json.loads(json_match.group())
            return {"steps": [], "reasoning": content}

# planner/tot_planner.py
from openai import OpenAI
from models import Plan, PlanStep, PlanEvaluation
from typing import List
import json
import logging

logger = logging.getLogger(__name__)

TOT_SYSTEM = """You are a strategic planner. For each goal, generate 
3 different candidate plans. For each plan:
1. Describe the approach
2. List the steps
3. Evaluate feasibility (0-1), efficiency (0-1), completeness (0-1)

Choose the BEST plan and explain why."""


class ToTPlanner:
    """
    Tree-of-Thought planner that explores multiple alternatives.
    
    Generates N candidate plans, evaluates each, and selects
    the best one. More thorough but costs more tokens.
    """
    
    def __init__(self, model: str = "gpt-4-turbo-preview", num_candidates: int = 3):
        self.client = OpenAI()
        self.model = model
        self.num_candidates = num_candidates

    def generate_plans(self, goal: str, available_tools: list[str]) -> List[Plan]:
        """
        Generate multiple candidate plans.
        
        Args:
            goal: The objective to achieve
            available_tools: List of available tool names
            
        Returns:
            List of candidate plans
        """
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": TOT_SYSTEM},
                {"role": "user", "content": f"""Goal: {goal}
Available tools: {', '.join(available_tools)}

Generate {self.num_candidates} different plans. For each plan provide:
- approach: description of the strategy
- steps: array of step objects (id, description, tool_needed, dependencies)
- scores: object with feasibility, efficiency, completeness (0-1 each)

Then select the best plan and explain your reasoning.

Return JSON: {{"plans": [...], "best_plan_id": "...", "reasoning": "..."}}"""},
            ],
            temperature=0.3,
            max_tokens=2000,
        )

        content = response.choices[0].message.content
        data = self._parse_response(content)
        
        plans = []
        for plan_data in data.get("plans", []):
            steps = [
                PlanStep(
                    id=s["id"],
                    description=s["description"],
                    tool_needed=s.get("tool_needed"),
                    dependencies=s.get("dependencies", []),
                )
                for s in plan_data.get("steps", [])
            ]
            scores = plan_data.get("scores", {})
            plan = Plan(
                id=plan_data.get("id", f"plan_{len(plans)}"),
                goal=goal,
                steps=steps,
                score=(
                    scores.get("feasibility", 0.5) +
                    scores.get("efficiency", 0.5) +
                    scores.get("completeness", 0.5)
                ) / 3,
                reasoning=plan_data.get("approach", ""),
            )
            plans.append(plan)

        logger.info(f"Generated {len(plans)} candidate plans")
        return plans

    def select_best(self, plans: List[Plan]) -> Plan:
        """Select the highest-scoring plan."""
        return max(plans, key=lambda p: p.score)

    def _parse_response(self, content: str) -> dict:
        """Parse LLM response as JSON."""
        try:
            return json.loads(content)
        except json.JSONDecodeError:
            import re
            match = re.search(r'\{[\s\S]*\}', content)
            if match:
                return json.loads(match.group())
            return {"plans": [], "reasoning": content}

# planner/plan_and_execute.py
from __future__ import annotations
import json
from openai import OpenAI
from models import Plan, PlanStep, StepStatus
from planner.cot_planner import CoTPlanner
from planner.tot_planner import ToTPlanner
from typing import Callable, Awaitable, Any
import logging

logger = logging.getLogger(__name__)


class PlanAndExecuteAgent:
    """
    Agent that plans before executing.
    
    Generates plans, executes steps, and re-plans
    on failures for robust completion.
    """
    
    def __init__(self, model: str = "gpt-4-turbo-preview"):
        self.client = OpenAI()
        self.model = model
        self.cot_planner = CoTPlanner(model)
        self.tot_planner = ToTPlanner(model)
        self.tools: dict[str, Callable] = {}

    def register_tool(self, name: str, func: Callable) -> None:
        """Register a tool for step execution."""
        self.tools[name] = func

    async def execute_goal(
        self,
        goal: str,
        strategy: str = "tot",
    ) -> dict:
        """
        Execute a goal using planning.
        
        Args:
            goal: The objective to achieve
            strategy: Planning strategy ("tot" or "cot")
            
        Returns:
            Execution results
        """
        # Generate plan(s)
        if strategy == "tot":
            plans = self.tot_planner.generate_plans(
                goal, list(self.tools.keys())
            )
            plan = self.tot_planner.select_best(plans)
        else:
            plan = self.cot_planner.generate_plan(
                goal, list(self.tools.keys())
            )

        # Execute steps
        execution_results = []
        for step in plan.steps:
            step.status = StepStatus.IN_PROGRESS
            result = await self._execute_step(step)
            execution_results.append(result)
            
            if result["success"]:
                step.status = StepStatus.COMPLETED
                step.result = result["output"]
            else:
                step.status = StepStatus.FAILED
                step.result = result["error"]
                
                # Re-plan on failure
                new_plan = await self._replan(goal, plan, step)
                if new_plan:
                    plan = new_plan
                    execution_results = []

        completed = all(
            s.status == StepStatus.COMPLETED for s in plan.steps
        )
        
        return {
            "goal": goal,
            "plan": plan.model_dump(),
            "results": execution_results,
            "completed": completed,
        }

    async def _execute_step(self, step: PlanStep) -> dict:
        """Execute a single plan step."""
        if not step.tool_needed or step.tool_needed not in self.tools:
            return await self._llm_step(step)

        tool_func = self.tools[step.tool_needed]
        try:
            if callable(tool_func):
                result = tool_func(step.description)
                return {"success": True, "output": str(result)}
            return {"success": False, "error": "Invalid tool"}
        except Exception as e:
            return {"success": False, "error": str(e)}

    async def _llm_step(self, step: PlanStep) -> dict:
        """Execute a step using LLM reasoning."""
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": "Execute the given step precisely."},
                {"role": "user", "content": step.description},
            ],
            temperature=0.0,
        )
        return {
            "success": True,
            "output": response.choices[0].message.content,
        }

    async def _replan(
        self, goal: str, current_plan: Plan, failed_step: PlanStep
    ) -> Plan | None:
        """Generate a new plan after a step failure."""
        logger.info(f"Re-planning after failure: {failed_step.description}")
        
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": """A step failed. Generate a new 
                plan that avoids the failed approach. Be concise."""},
                {"role": "user", "content": f"""Goal: {goal}
Failed step: {failed_step.description}
Error: {failed_step.result}

Generate a new plan as JSON with steps array."""},
            ],
            temperature=0.2,
        )
        
        content = response.choices[0].message.content
        try:
            data = json.loads(content)
            steps = [
                PlanStep(
                    id=s["id"],
                    description=s["description"],
                    tool_needed=s.get("tool_needed"),
                )
                for s in data.get("steps", [])
            ]
            return Plan(id="replan_001", goal=goal, steps=steps)
        except (json.JSONDecodeError, KeyError):
            return None

agent = PlanAndExecuteAgent()
result = await agent.execute_goal(
    "Extract data from CSV, transform to JSON, load into database",
    strategy="tot"
)
# Plans: sequential, parallel, or hybrid approaches
# Selects best based on data size and dependencies

result = await agent.execute_goal(
    "Research competitor pricing, analyze trends, create report",
    strategy="cot"
)
# Steps: search → analyze → visualize → write report

# Example: Safe planning with validation
class SafePlanAgent(PlanAndExecuteAgent):
    def validate_plan(self, plan: Plan) -> bool:
        for step in plan.steps:
            if step.tool_needed and step.tool_needed not in self.tools:
                return False
        return True