code/01-llm-basics/long_context.py

49 行 · 2.1 KB

コードと実行結果は実際に動かしたときのまま載せているため、コメントと出力は中国語です。

"""两个实验:把整份 httpx 文档放进一次请求,看缓存命中;再在文档里藏一句话,看模型能不能找到。

在 AI-Course 目录下运行:python code/01-llm-basics/long_context.py
一共发出 7 次请求,按 2026 年 9 月的价格花费约 0.05 美元。
"""
import os
import pathlib

from openai import OpenAI

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")
NO_THINKING = {"thinking": {"type": "disabled"}}

docs = pathlib.Path("data/httpx-docs")
text = "\n\n".join(p.read_text() for p in sorted(docs.rglob("*.md")) if p.name != "LICENSE.md")
print(f"文档共 {len(text)} 个字符")

# 实验一:同样的请求连发两次
system = "你是 httpx 的答疑助手。下面是 httpx 的全部文档:\n\n" + text
for i in range(2):
    r = client.chat.completions.create(
        model=MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": "httpx 默认的超时时间是多少秒?一句话回答。"},
        ],
        extra_body=NO_THINKING,
    )
    u = r.usage
    print(f"第 {i + 1} 次:输入 {u.prompt_tokens},其中缓存命中 {u.prompt_cache_hit_tokens},"
          f"输出 {u.completion_tokens} | {r.choices[0].message.content}")

# 实验二:把一句无关的话插到文档的不同位置
needle = "(备注:RepoBot 项目的内部口令是“蓝鲸四十二”。)"
paragraphs = text.split("\n\n")
for fraction in [0, 0.25, 0.5, 0.75, 1.0]:
    k = int(len(paragraphs) * fraction)
    haystack = "\n\n".join(paragraphs[:k] + [needle] + paragraphs[k:])
    r = client.chat.completions.create(
        model=MODEL,
        messages=[{"role": "user", "content": haystack + "\n\n问题:上面的文字里提到的 RepoBot 内部口令是什么?只回答口令本身。"}],
        extra_body=NO_THINKING,
    )
    print(f"口令放在 {fraction:>4.0%} 处:输入 {r.usage.prompt_tokens} 词元 → {r.choices[0].message.content}")