01. Planning Agents & Chain of Thought

# Standard Prompting
prompt_standard = f"Question: {problem}\nAnswer:"
 
# CoT Prompting
prompt_cot = f"Question: {problem}\nLet's think step by step."

problem = """
John had 5 apples. He gave 2 to Mary and ate 1.
Then he got 3 more from Mike, but dropped 1 on the way home.
How many apples does John have now?
"""
 
def get_completion(prompt):
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": prompt}
        ],
        temperature=0
    )
    return response.choices[0].message.content
 
# With CoT
response = get_completion(f"{problem}\nLet's think step by step.")
# Output:
# 1. John starts with 5 apples
# 2. Gives 2 to Mary: 5 - 2 = 3 apples
# 3. Eats 1: 3 - 1 = 2 apples
# 4. Gets 3 from Mike: 2 + 3 = 5 apples
# 5. Drops 1: 5 - 1 = 4 apples
# Answer: 4 apples

from pydantic import BaseModel, Field
from typing import List
 
class PlanStep(BaseModel):
    id: int = Field(description="Step number (starts from 1)")
    description: str = Field(description="What to do in this step")
    tool: str = Field(description="Tool to use (search or calculate)")
    args: str = Field(description="Arguments for the tool")
 
class Plan(BaseModel):
    steps: List[PlanStep] = Field(description="Ordered list of steps")
 
def create_plan(query: str) -> Plan:
    completion = client.beta.chat.completions.parse(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": "Create a step-by-step plan using available tools."},
            {"role": "user", "content": query}
        ],
        response_format=Plan
    )
    return completion.choices[0].message.parsed

def execute_plan(plan: Plan, tools: dict) -> str:
    results = {}
 
    for step in plan.steps:
        if step.tool in tools:
            result = tools[step.tool](step.args)
            results[step.id] = result
            print(f"Step {step.id}: {result}")
 
    # Synthesize final answer
    return synthesize_results(results)

def execute_with_replanning(plan: Plan, max_replans: int = 2):
    for attempt in range(max_replans):
        results, failed_step = execute_plan(plan)
 
        if failed_step is None:
            return results  # Success!
 
        # Replan from the failed step
        plan = replan(plan, failed_step, results)
 
    return results  # Best effort