Needle In A Haystack — Prompt Phase Shifting Test

评估「前置相位偏移」（Prompt Phase Shifting, PPS）对 LLM 超长上下文关键信息检索能力的影响。

本仓库由大海捞针（Needle-in-a-Haystack）测试项目改造而来，已改为针对 PPS 的专项性能评测：对同一个模型、同一套 sweep 网格，分别在开启与关闭 PPS 的 两种条件下运行，逐格对比命中率差异。

当前仓库已包含支持 DeepSeek 模型的测试配置文件。 仓库内其他模型的配置文件（Anthropic / OpenAI / Cohere）不属于当前测试支持范畴，仅供改造参考——请以 configs/models/deepseek-flash.yaml 为 唯一标准模板，把它们改写成对应的 PPS 测试配置（见下文「接入其他模型」）。

---

快速开始

uv sync --extra dev --extra pps     # dev = 测试与 lint；pps = 相位偏移 SDK

echo "DEEPSEEK_API_KEY=sk-..." > .env    # niah 启动时自动加载 .env

# 零配置冒烟：不调用真实模型，验证整条流水线
uv run niah run configs/runs/smoke.fake.yaml

# 真实 A/B 测试（A 组基线 / B 组实验，网格完全一致）
uv run niah run configs/runs/deepseek-a-test.yaml
uv run niah run configs/runs/deepseek-b-test.yaml

前置相位偏移 SDK

前置相位偏移 SDK prompt-phase-shifting 位于另一公开仓库 FactLab01/Prompt_Phase_Shifting，不在本仓库内，也暂未发布到公开的 Python index 上。pyproject.toml 的 [tool.uv.sources] 已声明该 git 来源，使用 uv sync --extra pps 会自动拉取安装，不需要手动 clone。

不装 SDK 也能跑通除 PPS 改写之外的全部流程：未安装时 PPS 桥记一次告警后按原文发送， 即 A/B 测试的基线行为。

两组结果分别落在 results/deepseek-a-baseline.jsonl 与 results/deepseek-b-phaseshift.jsonl，可逐格对比。resume: true 已开启， 中断后重跑只会补齐缺失的 cell。

渲染热力图（Context Length × Depth Percent → Score）：

uv run python plot_headmap.py --out-dir out     # 产出 heatmap-a.png / heatmap-b.png

---

接入前置相位偏移 SDK 的实现

PPS 的思想：把 prompt 中的「关键 token」用零宽字符物理推进到目标模型的注意力 偏置位置（token_index % slice_length == 0），让针落在模型最容易注意到的相位上。

接入点在 needlehaystack/providers/openai.py 的 complete() 内：messages 构建完成、即将交给 SDK 之前。这是唯一能同时覆盖 runner、CLI、库调用三条路径的位置。只改写消息的 content 文本，model / temperature / 工具调用 / 流式等参数一律不碰。

实现见 needlehaystack/providers/phaseshift.py， 要点：

• 默认关闭。未配置 phase_shift_* 时 provider 行为与关闭时完全一致，即 A/B 测试的基线组。
• 软依赖。prompt-phase-shifting 未安装时记一次告警后按原文发送，不会让整轮 测试报错。
• 旋钮隔离。phase_shift_* 由 NIAH 自己消费，不会转发给 API。

---

A/B 测试的唯一变量

两组配置只有一个差别：模型 YAML 里的 phase_shift_enabled。

	A 组（基线）	B 组（实验）
模型配置	configs/models/deepseek-flash.yaml	configs/models/deepseek-flash-phaseshift.yaml
运行配置	configs/runs/deepseek-a-test.yaml	configs/runs/deepseek-b-test.yaml
phase_shift_enabled	false	true

model / max_tokens / temperature / client 完全一致，sweep 网格完全一致。 改任一组时请同步另一组，否则对比失效。

---

PPS 旋钮

全部可选，写在模型 YAML 的 request: 或 client: 块（request 优先）：

旋钮	说明
phase_shift_enabled	总开关。也可用环境变量 NIAH_PHASE_SHIFT=1 打开
phase_shift_keywords	静态关键词列表，针内容固定时用
phase_shift_capture_pattern	动态捕获正则，针内容每次随机时用
phase_shift_slice_length	目标模型的切片长度，默认 4
phase_shift_min_chars	短于该长度的 prompt 不改写，0 = 一律处理
phase_shift_max_chars	phase_shifting 的输入长度上限，0 = 按 prompt 自动放宽

UUID 场景为什么要动态捕获

UUID 单针测试的针内容每次随机生成，写死关键词命中不了：UUID 本体的 token 切分 完全随机（同一个 UUID 会被切成 b/dd/640/-…，换个 UUID 又不一样）。

但它前面的引导短语 The secret code is 由 UuidSingle.template 固定，永远紧邻 UUID 且稳定存在。因此捕获「引导短语 + UUID 整体」作为关键词，即可把整颗针锚定到 相位 0。这就是 phase_shift_capture_pattern 的用途。

---

自定义运行

抛弃 demo 式的默认值，用两个 YAML 驱动自己的 sweep：一个 run 配置引用一个模型 配置。完整示例在 configs/。

uv run niah validate configs/runs/deepseek-b-test.yaml   # 只解析，不调用模型
uv run niah run      configs/runs/deepseek-b-test.yaml   # 实际运行

Run 配置（configs/runs/deepseek-a-test.yaml）

run_name: "deepseek-a-baseline"
model: "deepseek-flash"              # 按 id 解析到 configs/models/

task:
  type: "uuid"                       # UUID 单针测试
  needle_template: "The secret code is {uuid}."
  question_template: "What is the secret code?"

haystack:
  type: "files"
  path: "PaulGrahamEssays"

sweep:
  context_lengths: {min: 4000, max: 256000, num: 5, scale: "linear"}
  depth_percents:  {min: 0,    max: 100,   num: 5, scale: "linear"}
  seeds: [1, 2]

runner:
  concurrency: 4
  retries: 2
  resume: true

store:
  type: "jsonl"
  path: "results/deepseek-a-baseline.jsonl"

模型配置（configs/models/deepseek-flash-phaseshift.yaml）

id: "deepseek-flash-phaseshift"
runtime:
  sdk: "openai-python"               # 必须与 openai.py 里 register_provider 的键一致
  api: "chat_completions"            # 是下划线，不是点
client:
  api_key_env: "DEEPSEEK_API_KEY"    # 存的是环境变量名，provider 会 os.environ.get()
request:
  model: "deepseek-flash"
  max_tokens: 10240
  temperature: 0.0
  phase_shift_enabled: true
  phase_shift_capture_pattern: 'The\s+secret\s+code\s+is\s+[0-9a-fA-F-]{36}'
  phase_shift_slice_length: 4

request: 下除 model / stream 与 phase_shift_* 外的内容，原样转发给 SDK。 新增 provider 特有参数（thinking、reasoning_effort、top_p…）不需要改代码。

---

接入其他模型

仓库内 anthropic-opus-4-medium.yaml、openai-gpt-4o.yaml、 openai-gpt-4o-mini.yaml、cohere 等配置不是当前测试支持范畴内的文件， 仅作为字段写法参考。

要把某个模型纳入 PPS 测试，以 configs/models/deepseek-flash.yaml 为唯一标准 模板改造：

1. 复制 configs/models/deepseek-flash.yaml 为 <your-model>.yaml（基线组）与 <your-model>-phaseshift.yaml（实验组），两组只差 phase_shift_enabled。
2. 按目标 SDK 改 runtime.sdk / runtime.api（OpenAI 兼容服务可继续用 openai-python / chat_completions，并配 client.base_url_env）。
3. 改 client.api_key_env 指向该模型的凭据环境变量名。
4. 填 request.model / max_tokens / temperature，必要时加 pricing。
5. 复制一份 run 配置，把 model 指向新 id，并同步 sweep 网格。

若目标 SDK 不是 OpenAI 兼容接口，还需在 needlehaystack/providers/ 写一个插件并 register_provider(sdk, api, factory)，参考 anthropic.py。

---

CLI

niah run        <run.yaml>            运行 sweep，追加写入 JSONL
niah run        <run.yaml> --dry-run  解析、校验、打印计划，不调用模型
niah validate   <run.yaml>            解析 + 校验，不调用模型
niah reconstruct <results.jsonl> --row N [--out file]
                                      重建该行模型实际看到的上下文

--model-dir DIR（可重复）为裸模型 id 追加搜索路径。

---

结果行与重建

JSONL 每行只有几 KB，与上下文大小无关：不存储渲染后的长上下文。

{
  "schema_version": 2,
  "run_name": "deepseek-a-baseline",
  "model_id": "deepseek-flash",
  "model_request_name": "deepseek-flash",
  "task_type": "uuid",
  "context_length": 32000,
  "target_depth_percent": 50.0,
  "recipe": {
    "haystack": {"type": "files", "path": "PaulGrahamEssays"},
    "inserter": "single_depth",
    "needle_placements": [
      {"text": "The secret code is abc-...", "insertion_token_index": 15876, "actual_depth_percent": 49.6}
    ],
    "final_context_token_count": 32082
  },
  "expected_answer": "the-final-uuid",
  "prompt_question": "What is the secret code?",
  "response": "...",
  "score": {"value": 1.0, "details": {}},
  "usage": {"input_tokens": 32100, "output_tokens": 412},
  "cost_usd": 0.171,
  "duration_seconds": 12.4,
  "status": "ok",
  "seed": 1,
  "timestamp_utc": "2026-..."
}

niah reconstruct 按配方重放，产出与模型实际收到的内容逐字节一致的上下文。

---

内置任务

task.type	说明
single	一条静态事实置于一个深度，精确匹配计分
multi	N 条事实均匀散布于上下文，按命中比例计分
uuid	一个随机 UUID 置于一个深度，要求模型复述。PPS 测试默认使用此任务
uuid_chain	A → B → C → … 链式跳转，问题不透露链结构，模型需自行发现跳转路径

任务是 Protocol，见 needlehaystack/tasks/。新增任务只需 一个文件加一次注册调用，runner 不需要改动。

from needlehaystack.tasks import register_task

class MyCustomTask:
    name = "my_task"
    inserter_name = "single_depth"
    def generate_needle(self, seed): ...
    def insert(self, ctx, needle, depth): ...
    def question(self, needle): ...
    def score(self, response, needle): ...

register_task("my_task", MyCustomTask)

---

贡献

本项目基于 Needle-in-a-Haystack 项目进行开发，你可前往原始项目进行共享，亦可参与本项目共享贡献。

git clone https://github.com/FactLab01/Prompt_Phase_Shifting_Test_NeedleInAHaystack.git
cd Prompt_Phase_Shifting_Test_NeedleInAHaystack
uv sync --extra dev --extra pps

uv run ruff check .          # lint
uv run ruff format --check . # 格式检查（CI 同款）
uv run mypy needlehaystack   # 类型检查
uv run pytest                # 测试

测试不得访问网络；provider 调用一律走 FakeProvider （见 needlehaystack/providers/fake.py）。 详见 CONTRIBUTING.md。

---

历史结果与许可

2023 年的原始运行数据（OpenAI GPT-4 / Anthropic Claude 2.1）保留在 original_results/。其 schema 与 v2 不兼容，无法用现有工具加载，仅作存档。

MIT — 见 LICENSE.txt。使用本软件需保留对原作者与项目的署名。