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

class StepStatus(str, Enum):
    PENDING = "pending"
    IN_PROGRESS = "in_progress"
    COMPLETED = "completed"
    FAILED = "failed"

class PlanStep(BaseModel):
    id: str
    description: str
    tool_needed: Optional[str] = None
    status: StepStatus = StepStatus.PENDING
    result: Optional[str] = None
    dependencies: List[str] = []

class Plan(BaseModel):
    id: str
    goal: str
    steps: List[PlanStep]
    score: float = 0.0
    reasoning: str = ""

class PlanEvaluation(BaseModel):
    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

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:
    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:
        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"]
        ]
        return Plan(
            id="plan_cot_001",
            goal=goal,
            steps=steps,
            reasoning=plan_data.get("reasoning", ""),
        )

    def _parse_response(self, content: str) -> dict:
        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

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:
    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]:
        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)

        return plans

    def select_best(self, plans: List[Plan]) -> Plan:
        return max(plans, key=lambda p: p.score)

    def _parse_response(self, content: str) -> dict:
        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

class PlanAndExecuteAgent:
    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:
        self.tools[name] = func

    async def execute_goal(
        self,
        goal: str,
        strategy: str = "tot",
    ) -> dict:
        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())
            )

        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"]
                new_plan = await self._replan(goal, plan, step)
                if new_plan:
                    plan = new_plan
                    execution_results = []

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

    async def _execute_step(self, step: PlanStep) -> dict:
        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:
        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:
        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

# Parallel plan evaluation
import asyncio
from typing import List

class ParallelPlanEvaluator:
    def __init__(self, model: str = "gpt-4-turbo-preview"):
        from openai import AsyncOpenAI
        self.client = AsyncOpenAI()
        self.model = model

    async def evaluate_plans(self, plans: List[Plan]) -> List[dict]:
        tasks = [self._evaluate_one(plan) for plan in plans]
        return await asyncio.gather(*tasks)

    async def _evaluate_one(self, plan: Plan) -> dict:
        steps_desc = "\n".join(
            f"  {s.id}: {s.description}" for s in plan.steps
        )
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": "Evaluate this plan on feasibility (0-1), efficiency (0-1), completeness (0-1). Return JSON."},
                {"role": "user", "content": f"Goal: {plan.goal}\nSteps:\n{steps_desc}"},
            ],
            temperature=0.0,
        )
        import json
        try:
            return json.loads(response.choices[0].message.content)
        except:
            return {"feasibility": 0.5, "efficiency": 0.5, "completeness": 0.5}

import pytest
from planner.cot_planner import CoTPlanner
from planner.tot_planner import ToTPlanner

def test_cot_planner():
    planner = CoTPlanner()
    plan = planner.generate_plan(
        "Research AI agents and write a report",
        ["web_search", "write_file"],
    )
    assert len(plan.steps) > 0
    assert plan.goal == "Research AI agents and write a report"

def test_tot_planner():
    planner = ToTPlanner()
    plans = planner.generate_plans(
        "Analyze sales data and create visualization",
        ["read_csv", "create_chart"],
    )
    assert len(plans) >= 1
    best = planner.select_best(plans)
    assert best.score >= 0

# main.py
import asyncio
from planner.plan_and_execute import PlanAndExecuteAgent
from tools import calculator, web_search

async def main():
    agent = PlanAndExecuteAgent()
    agent.register_tool("calculator", calculator)
    agent.register_tool("web_search", web_search)

    result = await agent.execute_goal(
        "Research the top 3 Python web frameworks and compare their performance",
        strategy="tot",
    )
    print(f"Completed: {result['completed']}")
    print(f"Steps executed: {len(result['results'])}")

if __name__ == "__main__":
    asyncio.run(main())