Tree-of-Thought

from openai import OpenAI
import json
from typing import List, Dict

client = OpenAI()

def generate_thoughts(problem: str, current_state: str, n_thoughts: int = 3) -> List[str]:
    """현재 상태에서 가능한 다음 생각들을 생성"""
    prompt = f"""문제: {problem}

현재 진행 상황: {current_state}

가능한 다음 단계를 {n_thoughts}가지 제안하세요.
각 단계는 한 줄로 간결하게 작성하세요.
형식: 1. [단계] 2. [단계] 3. [단계]"""

    response = client.chat.completions.create(
        model="gpt-4-turbo",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.7
    )

    # 생각들 파싱
    text = response.choices[0].message.content
    thoughts = [t.strip() for t in text.split("\n") if t.strip()]
    return thoughts[:n_thoughts]

def evaluate_thought(problem: str, thought: str) -> float:
    """각 생각의 유망도를 0-1 사이로 평가"""
    prompt = f"""문제: {problem}

제안된 접근법: {thought}

이 접근법이 문제 해결에 얼마나 유망한지 0에서 10 사이 점수로 평가하세요.
점수만 숫자로 답하세요."""

    response = client.chat.completions.create(
        model="gpt-4-turbo",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.1
    )

    try:
        score = float(response.choices[0].message.content.strip()) / 10
        return min(max(score, 0), 1)
    except:
        return 0.5

def tree_of_thought_solve(problem: str, max_depth: int = 3, beam_width: int = 2) -> str:
    """Tree-of-Thought로 문제 해결 (Beam Search 방식)"""

    # 초기 상태
    beams = [{"path": [], "score": 1.0}]

    for depth in range(max_depth):
        new_beams = []

        for beam in beams:
            current_state = " -> ".join(beam["path"]) if beam["path"] else "시작"

            # 다음 가능한 생각들 생성
            thoughts = generate_thoughts(problem, current_state)

            for thought in thoughts:
                # 각 생각 평가
                score = evaluate_thought(problem, thought)

                new_beams.append({
                    "path": beam["path"] + [thought],
                    "score": beam["score"] * score
                })

        # 상위 beam_width개만 유지 (가지치기)
        new_beams.sort(key=lambda x: x["score"], reverse=True)
        beams = new_beams[:beam_width]

        print(f"Depth {depth + 1}: Top beam score = {beams[0]['score']:.3f}")

    # 최종 답 생성
    best_path = " -> ".join(beams[0]["path"])
    final_prompt = f"""문제: {problem}

추론 과정: {best_path}

위 추론을 바탕으로 최종 답을 제시하세요."""

    response = client.chat.completions.create(
        model="gpt-4-turbo",
        messages=[{"role": "user", "content": final_prompt}]
    )

    return response.choices[0].message.content

# 사용 예시
problem = "8, 3, 1, 6 네 숫자를 사칙연산으로 조합하여 24를 만드세요."
solution = tree_of_thought_solve(problem)
print(f"\n최종 답:\n{solution}")