code/02-prompting/reasoning.py

"""几道容易错的题，比较三种做法：直接回答、先写推理再回答、开启思考模式。

每道题每种做法各跑 5 次。一共 90 次请求，花费约 0.02 美元。
"""
import os
import re
import sys
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

from openai import OpenAI

sys.path.insert(0, str(Path(__file__).parent.parent / "01-llm-basics"))
from cost import cost_usd  # 01 模块第 4 课写的计费函数

client = OpenAI(
    api_key=os.environ["LLM_API_KEY"],
    base_url=os.environ.get("LLM_BASE_URL", "https://api.deepseek.com"),
)
MODEL = os.environ.get("LLM_MODEL", "deepseek-flash")
RUNS = 5

QUESTIONS = [
    ("9.11 和 9.9 哪个大？", "9.9"),
    ("“秋天的叶子一片片落下”这句话有几个字？", "10"),
    ("一个正方形的周长是 24 厘米，面积是多少平方厘米？", "36"),
    ("单词 raspberry 里有几个字母 r？", "3"),
    ("一根绳子对折三次，然后从正中间剪一刀，绳子变成了几段？", "9"),
    ("我有 3 个苹果，吃掉 1 个，又买了 2 个，然后把一半送给朋友，还剩几个？", "2"),
]

DIRECT = "只回答最终答案，一个数，不要任何解释。"
STEP_BY_STEP = "先一步一步写出推理过程，最后单独一行写“答案：”加上一个数。"

METHODS = [
    ("直接回答", DIRECT, False),
    ("先推理再回答", STEP_BY_STEP, False),
    ("开启思考", DIRECT, True),
]


def extract(text):
    # 有"答案："就取它后面的数，否则取整段话里的第一个数
    tail = text.split("答案：")[-1]
    numbers = re.findall(r"\d+(?:\.\d+)?", tail)
    return numbers[0] if numbers else ""


def run(system, thinking, question):
    response = client.chat.completions.create(
        model=MODEL,
        messages=[{"role": "system", "content": system}, {"role": "user", "content": question}],
        extra_body={"thinking": {"type": "enabled" if thinking else "disabled"}},
    )
    return extract(response.choices[0].message.content), response.usage


for name, system, thinking in METHODS:
    jobs = [q for q in QUESTIONS for _ in range(RUNS)]
    with ThreadPoolExecutor(10) as pool:
        results = list(pool.map(lambda q: run(system, thinking, q[0]), jobs))
    correct = sum(answer == q[1] for (answer, _), q in zip(results, jobs))
    output_tokens = sum(u.completion_tokens for _, u in results) / len(results)
    cost = sum(cost_usd(u, MODEL) for _, u in results)
    print(f"{name}：{correct}/{len(jobs)} 正确，平均输出 {output_tokens:.0f} 词元，共 {cost:.4f} 美元")
    for i, (question, truth) in enumerate(QUESTIONS):
        answers = [a for a, _ in results[i * RUNS:(i + 1) * RUNS]]
        right = sum(a == truth for a in answers)
        if right < RUNS:
            print(f"    {question}（正确答案 {truth}）→ {answers}")