01. Planning Agents & Chain of Thought

# 표준 프롬프팅
prompt_standard = f"질문: {problem}\n답변:"
 
# CoT 프롬프팅
prompt_cot = f"질문: {problem}\n단계별로 생각해봅시다."

problem = """
존이 사과 5개를 가지고 있었습니다. 메리에게 2개를 주고 1개를 먹었습니다.
그 후 마이크에게서 3개를 더 받았지만, 집에 오는 길에 1개를 떨어뜨렸습니다.
존은 지금 사과가 몇 개 있을까요?
"""
 
def get_completion(prompt):
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": "당신은 도움이 되는 어시스턴트입니다."},
            {"role": "user", "content": prompt}
        ],
        temperature=0
    )
    return response.choices[0].message.content
 
# CoT 적용
response = get_completion(f"{problem}\n단계별로 생각해봅시다.")
# 출력:
# 1. 존은 사과 5개로 시작
# 2. 메리에게 2개 줌: 5 - 2 = 3개
# 3. 1개 먹음: 3 - 1 = 2개
# 4. 마이크에게서 3개 받음: 2 + 3 = 5개
# 5. 1개 떨어뜨림: 5 - 1 = 4개
# 답: 4개

from pydantic import BaseModel, Field
from typing import List
 
class PlanStep(BaseModel):
    id: int = Field(description="단계 번호 (1부터 시작)")
    description: str = Field(description="이 단계에서 할 일")
    tool: str = Field(description="사용할 도구 (search 또는 calculate)")
    args: str = Field(description="도구에 전달할 인자")
 
class Plan(BaseModel):
    steps: List[PlanStep] = Field(description="순서가 있는 단계 목록")
 
def create_plan(query: str) -> Plan:
    completion = client.beta.chat.completions.parse(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": "사용 가능한 도구로 단계별 계획을 생성하세요."},
            {"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}")
 
    # 최종 답변 종합
    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  # 성공!
 
        # 실패한 단계부터 재계획
        plan = replan(plan, failed_step, results)
 
    return results  # 최선의 결과