Parcourir la source

Merge pull request 'codex/lq-token-test-init' (#1) from codex/lq-token-test-init into main

Reviewed-on: https://git.malls.iformall.com/server/new-api-check/pulls/1
main
liushen il y a 1 mois
Parent
révision
91e863b448
16 fichiers modifiés avec 2020 ajouts et 145 suppressions
  1. +75
    -3
      README.md
  2. +12
    -0
      config.example.yaml
  3. +50
    -4
      docs/USAGE.zh-CN.md
  4. +88
    -5
      docs/testing-guide.md
  5. +77
    -0
      py_demo/gemini_api_test.py
  6. +54
    -0
      py_demo/open_api_test.py
  7. +3
    -2
      src/benchmarks/aime.rs
  8. +9
    -7
      src/benchmarks/gpqa.rs
  9. +41
    -46
      src/benchmarks/judge.rs
  10. +73
    -50
      src/cli.rs
  11. +28
    -0
      src/config.rs
  12. +404
    -17
      src/protocols/anthropic.rs
  13. +649
    -0
      src/protocols/google.rs
  14. +1
    -0
      src/protocols/mod.rs
  15. +333
    -11
      src/protocols/openai.rs
  16. +123
    -0
      src/runner.rs

+ 75
- 3
README.md Voir le fichier

@@ -1,6 +1,6 @@
# lq_token_test

`lq_token_test` is a Rust CLI for checking LLM relay compatibility, running small benchmark probes, and testing request pacing against RPM targets. It supports OpenAI-compatible and Anthropic-compatible relay protocols, fetches local benchmark datasets, prints terminal summaries, and writes JSON reports for later comparison.
`lq_token_test` is a Rust CLI for checking LLM relay compatibility, running small benchmark probes, and testing request pacing against RPM targets. It supports OpenAI-compatible, Anthropic-compatible, and Google Gemini relay protocols, fetches local benchmark datasets, prints terminal summaries, and writes JSON reports for later comparison.

## Build And Test

@@ -46,6 +46,18 @@ providers:
budget_tokens: 10000
display: "omitted"

google:
protocol: google
base_url: "https://generativelanguage.googleapis.com/v1beta"
api_token: "${GOOGLE_API_KEY}"
default_model: "gemini-3-pro-preview"
stream: true
thinking:
enabled: true
budget_tokens: 5000
effort: "high"
display: "summarized"

benchmarks:
data_dir: "data/benchmarks"
aime2026:
@@ -61,9 +73,53 @@ Values written as `${ENV_NAME}` are expanded when the config is loaded. For exam
```bash
export OPENAI_RELAY_TOKEN="..."
export ANTHROPIC_RELAY_TOKEN="..."
export GOOGLE_API_KEY="..."
```

Thinking can also be enabled per run with CLI overrides such as `--thinking true`, `--thinking-type enabled`, `--thinking-budget-tokens 10000`, `--thinking-display omitted`, `--reasoning-effort high`, and `--reasoning-summary auto`. For Anthropic, enabling thinking omits `temperature` from the upstream request.
Thinking can also be enabled per run with CLI overrides such as `--thinking true`, `--thinking-type enabled`, `--thinking-budget-tokens 10000`, `--thinking-effort high`, `--thinking-display omitted`, `--reasoning-effort high`, and `--reasoning-summary auto`. For Anthropic, enabling thinking omits `temperature` from the upstream request.

### Google Gemini Thinking

For Google Gemini, thinking settings are sent under `generationConfig.thinkingConfig`:

- `budget_tokens` maps to `thinkingBudget`.
- `effort` maps to `thinkingLevel`.
- `display: summarized` sends `includeThoughts: true`.
- `display: omitted` sends `includeThoughts: false`.

If both `budget_tokens` and `effort` are configured, both fields are sent. Some Gemini backends or relays may reject a mixed Gemini 2.5/Gemini 3 style request, so prefer model-specific configs.

Recommended Gemini 3 config:

```yaml
google:
protocol: google
base_url: "https://generativelanguage.googleapis.com/v1beta"
api_token: "${GOOGLE_API_KEY}"
default_model: "gemini-3-pro-preview"
stream: true
thinking:
enabled: true
effort: "high"
display: "summarized"
```

Recommended Gemini 2.5 config:

```yaml
google:
protocol: google
base_url: "https://generativelanguage.googleapis.com/v1beta"
api_token: "${GOOGLE_API_KEY}"
default_model: "gemini-2.5-pro"
stream: true
thinking:
enabled: true
budget_tokens: 5000
display: "summarized"
```

If `enabled: true` is set without `budget_tokens`, `effort`, or `display`, the Google adapter does not send `thinkingConfig`; add `display` when you want to explicitly request or suppress thought summaries.

## Dataset Fetching

@@ -111,6 +167,16 @@ cargo run -- check \
--prompt "Reply with the word ready."
```

Run a simple Google Gemini relay check:

```bash
cargo run -- check \
--config config.yaml \
--provider google \
--model gemini-3-pro-preview \
--prompt "Reply with the word ready."
```

The check command prints the HTTP status, elapsed milliseconds, and model text.

## Benchmarks
@@ -126,6 +192,8 @@ cargo run -- bench aime2026 \
--limit 10
```

AIME prompts ask the model to put the final answer in `\boxed{}`. Scoring extracts the last boxed integer first, then falls back only to explicit final-answer forms such as `Final answer: 393`, `Answer: 393`, or `The answer is 393`; unconstrained trailing integers are treated as `no_answer`.

Run a GPQA-Diamond benchmark:

```bash
@@ -139,6 +207,8 @@ cargo run -- bench gpqa-diamond \

Omit `--limit` to run all locally available cases.

GPQA-Diamond prompts and scoring follow the OpenAI `simple-evals` style: the model is instructed to put `Answer: $LETTER` on the last line, where `LETTER` is one of `A`, `B`, `C`, or `D`. Scoring extracts answers with the same strict `Answer:` pattern, so a bare final `C` or prose such as `I choose C` is treated as `no_answer`.

## RPM Testing

Run a sustained RPM test. This is the default mode and starts requests at a stable interval:
@@ -237,7 +307,9 @@ Benchmark and RPM commands print a terminal summary with success counts, failure

Benchmark reports include `wrong_cases`, with each wrong case containing the case id, question, expected answer, extracted actual answer, and raw model output. RPM reports include request counts, mode, target RPM, observed RPM, latency, error counts, and mode-specific details such as burst summaries, probe summaries, window-boundary summaries, and optional limiter inference.

Use `--debug-raw` with `check`, `bench`, or `rpm` to write upstream raw responses under `outputs/debug/`. Non-streaming requests save the raw JSON body, and streaming requests save the raw SSE lines. The directory is ignored by git and can help diagnose relay-side response rewriting.
When an upstream request returns a non-success HTTP status such as 400, 429, or 504, `check`, `bench`, and `rpm` automatically write a request/response debug JSON file under `outputs/debug/`. The debug file includes the full request URL, redacted request headers, full request body including the prompt, response status, response headers, and full response body. If the request fails before an HTTP response is available, for example a connect timeout, read failure, or streaming interruption counted as `request_error`, the same directory gets a `*-request-error` debug JSON with `response.status: null`, `response.error_kind: "request_error"`, and the local error message. API tokens are redacted, but prompts and model outputs are preserved for troubleshooting.

Use `--debug-raw` with `check`, `bench`, or `rpm` when you also want to save successful upstream raw responses. Non-streaming success responses save the raw JSON body, and streaming success responses save the raw SSE lines. The directory is ignored by git and can help diagnose relay-side response rewriting.

## Comparing Scores



+ 12
- 0
config.example.yaml Voir le fichier

@@ -25,6 +25,18 @@ providers:
effort: high # Anthropic adaptive 模式可用
display: summarized # summarized | omitted

google:
protocol: google
base_url: "https://generativelanguage.googleapis.com/v1beta"
api_token: "${GOOGLE_API_KEY}"
default_model: "gemini-3-pro-preview"
stream: true
thinking:
enabled: true
budget_tokens: 5000 # Gemini 2.5: generationConfig.thinkingConfig.thinkingBudget
effort: high # Gemini 3: generationConfig.thinkingConfig.thinkingLevel
display: summarized # summarized -> includeThoughts: true, omitted -> false

benchmarks:
data_dir: "data/benchmarks"
aime2026:


+ 50
- 4
docs/USAGE.zh-CN.md Voir le fichier

@@ -6,6 +6,7 @@

- OpenAI-compatible 协议
- Anthropic-compatible 协议
- Google Gemini 协议
- YAML 配置文件
- AIME 2026 数据集
- GPQA-Diamond 数据集
@@ -69,6 +70,18 @@ providers:
budget_tokens: 10000
display: "omitted"

google:
protocol: google
base_url: "https://generativelanguage.googleapis.com/v1beta"
api_token: "${GOOGLE_API_KEY}"
default_model: "gemini-3-pro-preview"
stream: true
thinking:
enabled: true
budget_tokens: 5000
effort: "high"
display: "summarized"

benchmarks:
data_dir: "data/benchmarks"
aime2026:
@@ -84,6 +97,7 @@ benchmarks:
```bash
export OPENAI_RELAY_TOKEN="..."
export ANTHROPIC_RELAY_TOKEN="..."
export GOOGLE_API_KEY="..."
```

注意:工具只会解析当前使用的 provider token。比如只运行 `--provider anthropic` 时,不需要设置 `OPENAI_RELAY_TOKEN`。
@@ -105,6 +119,8 @@ cargo run -- check --provider openai --prompt "hello" \

Anthropic 开启 thinking 时,请求体不会再发送 `temperature`。

Google Gemini 开启 thinking 时,会写入 `generationConfig.thinkingConfig`:`budget_tokens` 对应 `thinkingBudget`,`effort` 对应 Gemini 3 的 `thinkingLevel`,`display: summarized` 会发送 `includeThoughts: true`,`display: omitted` 会发送 `includeThoughts: false`。

## 3. 协议连通性测试

OpenAI-compatible:
@@ -127,6 +143,16 @@ cargo run -- check \
--prompt "Reply with the word ready."
```

Google Gemini:

```bash
cargo run -- check \
--config config.yaml \
--provider google \
--model gemini-3-pro-preview \
--prompt "Reply with the word ready."
```

如果不传 `--model`,会使用配置里的 `default_model`。

## 4. 下载数据集
@@ -179,6 +205,8 @@ cargo run -- bench aime2026 \
--concurrency 4
```

AIME prompt 会要求模型把最终答案写进 `\boxed{}`。评分时优先提取最后一个 boxed integer,其次只接受 `Final answer: 393`、`Answer: 393`、`The answer is 393` 这类明确 final-answer 格式;解释文本末尾的普通数字不会再被当成答案,会记为 `no_answer`。

GPQA-Diamond 小样本:

```bash
@@ -198,6 +226,8 @@ cargo run -- bench gpqa-diamond \
--concurrency 4
```

GPQA-Diamond 的 prompt 和评分按 OpenAI `simple-evals` 风格处理:要求模型最后一行输出 `Answer: $LETTER`,其中 `LETTER` 是 `A`、`B`、`C` 或 `D`。评分只按这个严格 `Answer:` 格式提取答案;如果模型只在最后单独输出 `C`,或输出 `I choose C`,会记为 `no_answer`。

benchmark 会输出:

- accuracy
@@ -382,9 +412,25 @@ ls -lt reports | head
jq . reports/<报告文件>.json
```

### 原始响应 Debug
### 错误请求/响应 Debug

当上游返回 400、429、504 等非成功 HTTP 状态时,`check`、`bench`、`rpm` 会自动在 `outputs/debug/` 写入请求/响应 debug JSON,不需要额外参数。请求在拿到 HTTP 响应前失败时,例如连接超时、读取响应失败、流式响应中断并被统计成 `request_error`,也会写入 `*-request-error` debug JSON。

debug JSON 包含:

- 完整请求 URL
- 脱敏后的请求 headers
- 完整请求 body,包括 prompt
- 响应 status;`request_error` 场景为 `null`
- 完整响应 headers;`request_error` 场景为空对象
- 完整响应 body;`request_error` 场景为空字符串
- `request_error` 场景额外包含 `response.error_kind: "request_error"` 和本地错误信息 `response.error`

API token 会被脱敏,但 prompt、题目内容和模型输出会原样保存。排查完后可以按需清理 `outputs/debug/`。

### 原始成功响应 Debug

如果需要排查中转站是否改写了模型响应,可以开启 `--debug-raw`:
如果需要排查中转站是否改写了成功响应,可以开启 `--debug-raw`:

```bash
cargo run -- check --provider anthropic --stream --debug-raw --prompt "hello"
@@ -392,13 +438,13 @@ cargo run -- bench aime2026 --provider anthropic --stream --debug-raw --limit 3
cargo run -- rpm --provider anthropic --rpm 60 --duration 30s --stream --debug-raw --prompt "hello"
```

开启后,原始响应会写到:
开启后,成功请求的原始响应会写到:

```text
outputs/debug/
```

非流式请求保存完整 JSON body;流式请求保存原始 SSE 行,包括 `event:` 和 `data:`。文件不包含 API token,但可能包含模型输出内容,所以 `outputs/` 不会提交到 git。
非流式成功请求保存完整 JSON body;流式成功请求保存原始 SSE 行,包括 `event:` 和 `data:`。文件不包含 API token,但可能包含模型输出内容,所以 `outputs/` 不会提交到 git。

benchmark report 包含:



+ 88
- 5
docs/testing-guide.md Voir le fichier

@@ -6,6 +6,7 @@
# 设置环境变量
export ANTHROPIC_RELAY_TOKEN="your-token-here"
export OPENAI_RELAY_TOKEN="your-token-here"
export GOOGLE_API_KEY="your-token-here"

# 构建
cargo build --release
@@ -13,6 +14,24 @@ cargo build --release

确认 `config.yaml` 中 provider 配置正确(base_url、protocol、default_model)。

Google Gemini provider 推荐配置:

```yaml
providers:
google:
protocol: google
base_url: "https://generativelanguage.googleapis.com/v1beta"
api_token: "${GOOGLE_API_KEY}"
default_model: "gemini-3-pro-preview"
stream: true
thinking:
enabled: true
effort: high
display: summarized
```

如果走中转站,把 `base_url` 改成中转站提供的 Gemini base url。

---

## 阶段 1:RPM 限流测试(Anthropic 协议)
@@ -32,6 +51,15 @@ cargo run -- check --provider anthropic --stream --prompt "Reply with pong."
- 非流式:返回文本 + elapsed_ms
- 流式:额外输出 first_token_ms,且 first_token_ms < elapsed_ms

Google Gemini 连通性验证:

```bash
cargo run -- check --provider google --prompt "Reply with pong."
cargo run -- check --provider google --stream --prompt "Reply with pong."
```

验证要点相同。Google 流式请求会走 `streamGenerateContent`,非流式请求会走 `generateContent`。

### 1.2 Sustained 模式 — 持续发送

以目标 RPM 匀速发送,观察成功率。
@@ -145,8 +173,13 @@ cargo run -- bench aime2026 --provider anthropic

# 指定更强模型
cargo run -- bench aime2026 --provider anthropic --model anthropic/claude-sonnet-4-20250514

# Google Gemini 快速验证
cargo run -- bench aime2026 --provider google --model gemini-3-pro-preview --limit 10
```

AIME prompt 会要求模型把最终答案写进 `\boxed{}`。评分时优先提取最后一个 boxed integer,其次只接受 `Final answer: 393`、`Answer: 393`、`The answer is 393` 这类明确 final-answer 格式;解释文本末尾的普通数字不会再被当成答案,会记为 `no_answer`。

### 2.3 GPQA Diamond(研究生级别多选题)

```bash
@@ -158,8 +191,13 @@ cargo run -- bench gpqa-diamond --provider anthropic

# 指定模型
cargo run -- bench gpqa-diamond --provider anthropic --model anthropic/claude-sonnet-4-20250514

# Google Gemini 快速验证
cargo run -- bench gpqa-diamond --provider google --model gemini-3-pro-preview --limit 10
```

GPQA-Diamond 的 prompt 和评分按 OpenAI `simple-evals` 风格处理:要求模型最后一行输出 `Answer: $LETTER`,其中 `LETTER` 是 `A`、`B`、`C` 或 `D`。评分只按这个严格 `Answer:` 格式提取答案;如果模型只在最后单独输出 `C`,或输出 `I choose C`,会记为 `no_answer`。

### 2.4 精度参考基线

| 数据集 | claude-haiku-4.5 | claude-sonnet-4 | gpt-4o-mini |
@@ -217,13 +255,58 @@ cargo run -- bench aime2026 \
--thinking true \
--reasoning-effort high \
--reasoning-summary auto

# Google Gemini 3 thinkingLevel
cargo run -- bench aime2026 \
--provider google \
--model gemini-3-pro-preview \
--limit 5 \
--thinking true \
--thinking-effort high \
--thinking-display summarized

# Google Gemini 2.5 thinkingBudget
cargo run -- bench aime2026 \
--provider google \
--model gemini-2.5-pro \
--limit 5 \
--thinking true \
--thinking-budget-tokens 5000 \
--thinking-display summarized
```

Anthropic 开启 thinking 后,请求体不会发送 `temperature`。benchmark 和 RPM 的 JSON report 会记录本次 thinking 参数,方便复现实验。

## 原始响应 Debug 模式
Google Gemini 开启 thinking 后,请求体会按以下规则写入 `generationConfig.thinkingConfig`:

- `budget_tokens` 对应 `thinkingBudget`,推荐给 Gemini 2.5 使用。
- `effort` 对应 `thinkingLevel`,推荐给 Gemini 3 使用。
- `display: summarized` 会发送 `includeThoughts: true`。
- `display: omitted` 会发送 `includeThoughts: false`。

如果同时配置 `budget_tokens` 和 `effort`,两个字段都会发送。为了减少后端或中转站兼容性问题,建议 Gemini 3 只配 `effort`,Gemini 2.5 只配 `budget_tokens`。

如果只配置 `enabled: true`,但不配置 `budget_tokens`、`effort`、`display`,当前实现不会发送 `thinkingConfig`。如果想显式请求或关闭 thought summary,至少加上 `display: summarized` 或 `display: omitted`。

## 错误请求/响应 Debug

当上游返回 400、429、504 等非成功 HTTP 状态时,`check`、`bench`、`rpm` 会自动在 `outputs/debug/` 写入一份请求/响应 debug JSON,不需要额外加参数。请求在拿到 HTTP 响应前失败时,例如连接超时、读取响应失败、流式响应中断并被统计成 `request_error`,也会写入 `*-request-error` debug JSON。

debug JSON 包含:

- 完整请求 URL
- 脱敏后的请求 headers
- 完整请求 body,包括 prompt
- 响应 status;`request_error` 场景为 `null`
- 完整响应 headers;`request_error` 场景为空对象
- 完整响应 body;`request_error` 场景为空字符串
- `request_error` 场景额外包含 `response.error_kind: "request_error"` 和本地错误信息 `response.error`

API token 会被脱敏,但 prompt、题目内容和模型输出会原样保存,排查完后按需清理 `outputs/debug/`。

## 原始成功响应 Debug 模式

排查中转站是否改写响应时,可以加 `--debug-raw`:
排查中转站是否改写成功响应时,可以加 `--debug-raw`:

```bash
cargo run -- check --provider anthropic --stream --debug-raw --prompt "Reply with pong."
@@ -231,10 +314,10 @@ cargo run -- bench aime2026 --provider anthropic --stream --debug-raw --limit 3
cargo run -- rpm --provider anthropic --rpm 60 --duration 30s --stream --debug-raw --prompt "Hi"
```

开启后,程序会把上游原始响应写到 `outputs/debug/`:
开启后,程序会把成功请求的上游原始响应写到 `outputs/debug/`:

- 非流式请求保存完整 JSON body
- 流式请求保存原始 SSE 行,包括 `event:` 和 `data:`
- 非流式成功请求保存完整 JSON body
- 流式成功请求保存原始 SSE 行,包括 `event:` 和 `data:`
- 文件不包含 API token
- 文件可能包含模型输出内容,`outputs/` 不会提交到 git



+ 77
- 0
py_demo/gemini_api_test.py Voir le fichier

@@ -0,0 +1,77 @@
from google import genai
from google.genai import types
import os

GEMINI_API_KEY = os.environ["GEMINI_API_KEY"]
GEMINI_BASE_URL = os.environ.get("GEMINI_BASE_URL", "https://generativelanguage.googleapis.com/v1beta")
MODEL_ID = os.environ.get("GEMINI_MODEL", "gemini-3.1-pro-preview")


def _make_client() -> genai.Client:
return genai.Client(
api_key=GEMINI_API_KEY,
http_options=types.HttpOptions(base_url=GEMINI_BASE_URL),
)


def list_models():
client = _make_client()
models = client.models.list()
model_ids = sorted([m.name for m in models])
print(f"共找到 {len(model_ids)} 个模型:")
for model_id in model_ids:
print(f" - {model_id}")


def test_model(
stream: bool = False,
thinking: bool = True,
prompt: str = "Hello, how are you?",
):
client = _make_client()
thinking_config = (
types.ThinkingConfig(include_thoughts=True, thinking_budget=5000)
if thinking
else types.ThinkingConfig(include_thoughts=False)
)
config = types.GenerateContentConfig(thinking_config=thinking_config)

if stream:
in_thinking = False
for chunk in client.models.generate_content_stream(
model=MODEL_ID, contents=prompt, config=config
):
if not chunk.candidates:
continue
for part in chunk.candidates[0].content.parts or []:
if part.thought:
if not in_thinking:
print("<thinking>", flush=True)
in_thinking = True
print(part.text, end="", flush=True)
else:
if in_thinking:
print("\n</thinking>\n", flush=True)
in_thinking = False
print(part.text or "", end="", flush=True)
if in_thinking:
print("\n</thinking>", flush=True)
print()
else:
response = client.models.generate_content(
model=MODEL_ID, contents=prompt, config=config
)
for part in response.candidates[0].content.parts:
if part.thought:
print(f"<thinking>\n{part.text}\n</thinking>\n")
else:
print(part.text or "", end="")
print()


if __name__ == "__main__":
test_model(
stream=True,
thinking=False,
prompt="解释什么是 MVCC,并举一个 PostgreSQL 中的应用例子,控制在 150 字内。",
)

+ 54
- 0
py_demo/open_api_test.py Voir le fichier

@@ -0,0 +1,54 @@
import os

import openai

OPENAI_API_KEY = os.environ["OPENAI_API_KEY"]
OPENAI_BASE_URL = os.environ.get("OPENAI_BASE_URL", "https://lancerouter.ai/v1")
MODEL_ID = os.environ.get("OPENAI_MODEL", "google/gemini-3.1-pro-preview")

def list_models():
client = openai.OpenAI(api_key=OPENAI_API_KEY, base_url=OPENAI_BASE_URL)
models = client.models.list()
model_ids = sorted([m.id for m in models.data])
print(f"共找到 {len(model_ids)} 个模型:")
for model_id in model_ids:
print(f" - {model_id}")


def test_model(stream: bool = False, thinking: bool = True, prompt: str = "Hello, how are you?"):
client = openai.OpenAI(api_key=OPENAI_API_KEY, base_url=OPENAI_BASE_URL)

extra_body = {"thinking": {"type": "enabled", "budget_tokens": 5000}} if thinking else {}

if stream:
response = client.chat.completions.create(
model=MODEL_ID,
messages=[{"role": "user", "content": prompt}],
stream=True,
extra_body=extra_body,
)
for chunk in response:
if not chunk.choices:
continue
delta = chunk.choices[0].delta
# 输出思考内容(thinking block)
if hasattr(delta, "thinking") and delta.thinking:
print(delta.thinking, end="", flush=True)
elif delta.content:
print(delta.content, end="", flush=True)
else:
response = client.chat.completions.create(
model=MODEL_ID,
messages=[{"role": "user", "content": prompt}],
extra_body=extra_body,
)
message = response.choices[0].message
# 输出思考内容(thinking block)
if hasattr(message, "thinking") and message.thinking:
print(f"<thinking>\n{message.thinking}\n</thinking>\n")
print(message.content)


if __name__ == "__main__":
# list_models()
test_model(stream=True, thinking=False, prompt="解释什么是 MVCC,并举一个 PostgreSQL 中的应用例子,控制在 150 字内。")

+ 3
- 2
src/benchmarks/aime.rs Voir le fichier

@@ -74,7 +74,7 @@ pub fn load_cases(data_dir: &Path) -> Result<LoadedAimeCases> {
impl AimeCase {
pub fn prompt(&self) -> String {
format!(
"Solve the following AIME problem.\n\n{}\n\nRespond with the final integer answer.",
"Solve the following AIME problem. Think step by step before answering.\n\n{}\n\nPut your final answer in \\boxed{{}}.",
self.problem
)
}
@@ -95,7 +95,8 @@ mod tests {
let prompt = case.prompt();

assert!(prompt.contains("What is 20 + 22?"));
assert!(prompt.contains("final integer answer"));
assert!(prompt.contains("Think step by step before answering."));
assert!(prompt.contains(r"\boxed{}"));
}

#[test]


+ 9
- 7
src/benchmarks/gpqa.rs Voir le fichier

@@ -160,12 +160,12 @@ impl GpqaCase {
let choices = self
.choices
.iter()
.map(|(label, text)| format!("{label}. {text}"))
.map(|(label, text)| format!("{label}) {text}"))
.collect::<Vec<_>>()
.join("\n");

format!(
"Answer the following multiple-choice question.\n\n{}\n\n{}\n\nPlease answer with exactly one letter: A, B, C, or D.",
"Answer the following multiple choice question.\nThe last line of your response should be of the following format: 'Answer: $LETTER' (without quotes) where LETTER is one of ABCD. Think step by step before answering.\n{}\n\n{}",
self.question, choices
)
}
@@ -192,11 +192,13 @@ mod tests {
let prompt = case.prompt();

assert!(prompt.contains("Which option is correct?"));
assert!(prompt.contains("A. Alpha"));
assert!(prompt.contains("B. Beta"));
assert!(prompt.contains("C. Gamma"));
assert!(prompt.contains("D. Delta"));
assert!(prompt.contains("answer with exactly one letter"));
assert!(prompt.contains("A) Alpha"));
assert!(prompt.contains("B) Beta"));
assert!(prompt.contains("C) Gamma"));
assert!(prompt.contains("D) Delta"));
assert!(prompt.contains("The last line of your response"));
assert!(prompt.contains("'Answer: $LETTER'"));
assert!(prompt.contains("Think step by step before answering."));
}

#[test]


+ 41
- 46
src/benchmarks/judge.rs Voir le fichier

@@ -2,31 +2,26 @@ use regex::Regex;

pub fn extract_final_integer(text: &str) -> Option<String> {
let boxed = Regex::new(r"\\boxed\{\s*(-?\d+)\s*\}").expect("valid boxed integer regex");
if let Some(captures) = boxed.captures(text) {
if let Some(captures) = boxed.captures_iter(text).last() {
return captures.get(1).map(|value| value.as_str().to_string());
}

let integer = Regex::new(r"-?\d+").expect("valid integer regex");
integer
.find_iter(text)
let explicit =
Regex::new(r"(?i)\b(?:final\s+answer|answer|the\s+answer\s+is)\s*(?:is|:|-)?\s*(-?\d+)\b")
.expect("valid explicit integer answer regex");
explicit
.captures_iter(text)
.last()
.map(|value| value.as_str().to_string())
.and_then(|captures| captures.get(1).map(|value| value.as_str().to_string()))
}

pub fn extract_choice(text: &str) -> Option<char> {
let explicit = Regex::new(
r"(?i)\b(?:final\s+answer|answer|choose|choice)\s*(?:is|:|-)?\s*\(?([A-D])\)?\b",
)
.expect("valid explicit choice regex");
if let Some(captures) = explicit.captures_iter(text).last() {
return choice_from_capture(&captures, 1);
}

let whole_output =
Regex::new(r"(?i)^(?:([A-D])|\(([A-D])\))$").expect("valid whole output choice regex");
whole_output.captures(text.trim()).and_then(|captures| {
choice_from_capture(&captures, 1).or_else(|| choice_from_capture(&captures, 2))
})
let explicit =
Regex::new(r"(?i)Answer[ \t]*:[ \t]*\$?([A-D])\$?").expect("valid GPQA answer regex");
explicit
.captures_iter(text)
.last()
.and_then(|captures| choice_from_capture(&captures, 1))
}

fn choice_from_capture(captures: &regex::Captures<'_>, index: usize) -> Option<char> {
@@ -49,9 +44,9 @@ mod tests {
use super::*;

#[test]
fn extracts_final_integer_from_plain_answer() {
fn extracts_final_integer_from_answer_prefix() {
assert_eq!(
extract_final_integer("The answer is 42."),
extract_final_integer("Final answer: 42."),
Some("42".to_string())
);
}
@@ -65,58 +60,58 @@ mod tests {
}

#[test]
fn extracts_choice_from_answer_prefix() {
assert_eq!(extract_choice("Answer: C"), Some('C'));
fn extracts_last_boxed_integer() {
assert_eq!(
extract_final_integer(r"First \boxed{12}, final \boxed{34}"),
Some("34".to_string())
);
}

#[test]
fn extracts_choice_from_parenthesized_lowercase() {
assert_eq!(extract_choice("I choose (b)."), Some('B'));
fn ignores_trailing_integer_without_explicit_final_answer() {
assert_eq!(
extract_final_integer("The answer should be 393. I verified 30 cases."),
None
);
}

#[test]
fn extracts_choice_from_final_answer_instead_of_earlier_prose_option() {
assert_eq!(
extract_choice(
"Option A is tempting, but after checking the calculation the answer is D"
),
Some('D')
);
fn extracts_choice_from_answer_prefix() {
assert_eq!(extract_choice("Answer: C"), Some('C'));
}

#[test]
fn extracts_choice_from_final_answer_after_considered_options() {
assert_eq!(
extract_choice("I considered A and B before deciding. Final answer: C"),
Some('C')
);
fn extracts_choice_from_dollar_wrapped_answer() {
assert_eq!(extract_choice("Answer: $B$"), Some('B'));
}

#[test]
fn extracts_last_explicit_choice() {
fn extracts_last_answer_prefix() {
assert_eq!(
extract_choice("Answer: A was my first thought. Final answer: D"),
extract_choice("Answer: A was my first thought.\nAnswer: D"),
Some('D')
);
}

#[test]
fn ignores_prose_without_final_choice() {
assert_eq!(
extract_choice("This explanation has no final option."),
None
);
fn ignores_choice_without_answer_prefix() {
assert_eq!(extract_choice("I choose (b)."), None);
}

#[test]
fn ignores_bare_final_line_choice() {
assert_eq!(extract_choice("<think>...</think>\n\nC"), None);
}

#[test]
fn judges_integer_by_extracted_value() {
assert!(judge_integer("Final: 42.", "42"));
assert!(!judge_integer("Final: 43.", "42"));
assert!(judge_integer("Final answer: 42.", "42"));
assert!(!judge_integer("Final answer: 43.", "42"));
}

#[test]
fn judges_choice_by_extracted_value() {
assert!(judge_choice("I choose (b).", 'B'));
assert!(judge_choice("Answer: B", 'B'));
assert!(!judge_choice("Answer: C", 'B'));
}
}

+ 73
- 50
src/cli.rs Voir le fichier

@@ -138,7 +138,7 @@ pub enum BenchCommand {
limit: Option<usize>,
#[arg(long, num_args = 0..=1, default_missing_value = "true")]
stream: Option<bool>,
#[arg(long, default_value_t = 32768)]
#[arg(long, default_value_t = 32_768)]
max_tokens: u32,
#[arg(long)]
debug_raw: bool,
@@ -170,7 +170,7 @@ pub enum BenchCommand {
limit: Option<usize>,
#[arg(long, num_args = 0..=1, default_missing_value = "true")]
stream: Option<bool>,
#[arg(long, default_value_t = 32768)]
#[arg(long, default_value_t = 32_768)]
max_tokens: u32,
#[arg(long)]
debug_raw: bool,
@@ -324,7 +324,7 @@ async fn dispatch_bench(command: BenchCommand) -> Result<()> {
reasoning_effort,
reasoning_summary,
} => {
run_aime_benchmark(
run_aime_benchmark(BenchmarkCommandOptions {
config,
provider,
model,
@@ -333,7 +333,7 @@ async fn dispatch_bench(command: BenchCommand) -> Result<()> {
stream,
max_tokens,
debug_raw,
ThinkingOverrides {
thinking_overrides: ThinkingOverrides {
thinking,
thinking_type,
thinking_budget_tokens,
@@ -342,7 +342,7 @@ async fn dispatch_bench(command: BenchCommand) -> Result<()> {
reasoning_effort,
reasoning_summary,
},
)
})
.await
}
BenchCommand::GpqaDiamond {
@@ -362,7 +362,7 @@ async fn dispatch_bench(command: BenchCommand) -> Result<()> {
reasoning_effort,
reasoning_summary,
} => {
run_gpqa_benchmark(
run_gpqa_benchmark(BenchmarkCommandOptions {
config,
provider,
model,
@@ -371,7 +371,7 @@ async fn dispatch_bench(command: BenchCommand) -> Result<()> {
stream,
max_tokens,
debug_raw,
ThinkingOverrides {
thinking_overrides: ThinkingOverrides {
thinking,
thinking_type,
thinking_budget_tokens,
@@ -380,14 +380,14 @@ async fn dispatch_bench(command: BenchCommand) -> Result<()> {
reasoning_effort,
reasoning_summary,
},
)
})
.await
}
}
}

async fn run_aime_benchmark(
config_path: PathBuf,
struct BenchmarkCommandOptions {
config: PathBuf,
provider: Option<String>,
model: Option<String>,
concurrency: usize,
@@ -396,11 +396,15 @@ async fn run_aime_benchmark(
max_tokens: u32,
debug_raw: bool,
thinking_overrides: ThinkingOverrides,
) -> Result<()> {
let config = AppConfig::load(&config_path)?;
let provider_name = provider_name(&config, provider.as_deref())?;
}

async fn run_aime_benchmark(options: BenchmarkCommandOptions) -> Result<()> {
let config = AppConfig::load(&options.config)?;
let provider_name = provider_name(&config, options.provider.as_deref())?;
let provider_config = config.resolved_provider(Some(&provider_name))?;
let model = model.unwrap_or_else(|| provider_config.default_model.clone());
let model = options
.model
.unwrap_or_else(|| provider_config.default_model.clone());
let loaded = benchmarks::aime::load_cases(Path::new(&config.benchmarks.data_dir))?;
let dataset = dataset_report(
config
@@ -410,15 +414,17 @@ async fn run_aime_benchmark(
.map(|dataset| (dataset.source.as_str(), dataset.split.as_str())),
&loaded.local_path,
);
let cases = apply_limit(loaded.cases, limit);
let cases = apply_limit(loaded.cases, options.limit);
let total = cases.len() as u64;
let started_at = Utc::now();
let started = Instant::now();
let mut base_request = request_template(&provider_config, &model, 0.0, max_tokens);
base_request.stream = stream.unwrap_or(provider_config.stream);
base_request.raw_debug = raw_debug_config(debug_raw, &provider_name, &model);
base_request.thinking =
merged_thinking_config(provider_config.thinking.as_ref(), thinking_overrides);
let mut base_request = request_template(&provider_config, &model, 0.0, options.max_tokens);
base_request.stream = options.stream.unwrap_or(provider_config.stream);
base_request.raw_debug = raw_debug_config(options.debug_raw, &provider_name, &model);
base_request.thinking = merged_thinking_config(
provider_config.thinking.as_ref(),
options.thinking_overrides,
);
let protocol = provider_config.protocol;

let pb = ProgressBar::new(total);
@@ -437,7 +443,7 @@ async fn run_aime_benchmark(
(case, result)
}
})
.buffer_unordered(nonzero_concurrency(concurrency));
.buffer_unordered(nonzero_concurrency(options.concurrency));

let mut metrics = Metrics::new();
let mut wrong_cases = Vec::new();
@@ -489,9 +495,9 @@ async fn run_aime_benchmark(
dataset,
started_at,
duration_ms: started.elapsed().as_millis(),
concurrency,
limit,
max_tokens,
concurrency: options.concurrency,
limit: options.limit,
max_tokens: options.max_tokens,
summary,
correct_samples,
wrong_cases,
@@ -501,21 +507,13 @@ async fn run_aime_benchmark(
Ok(())
}

async fn run_gpqa_benchmark(
config_path: PathBuf,
provider: Option<String>,
model: Option<String>,
concurrency: usize,
limit: Option<usize>,
stream: Option<bool>,
max_tokens: u32,
debug_raw: bool,
thinking_overrides: ThinkingOverrides,
) -> Result<()> {
let config = AppConfig::load(&config_path)?;
let provider_name = provider_name(&config, provider.as_deref())?;
async fn run_gpqa_benchmark(options: BenchmarkCommandOptions) -> Result<()> {
let config = AppConfig::load(&options.config)?;
let provider_name = provider_name(&config, options.provider.as_deref())?;
let provider_config = config.resolved_provider(Some(&provider_name))?;
let model = model.unwrap_or_else(|| provider_config.default_model.clone());
let model = options
.model
.unwrap_or_else(|| provider_config.default_model.clone());
let loaded = benchmarks::gpqa::load_cases(Path::new(&config.benchmarks.data_dir))?;
let dataset = dataset_report(
config
@@ -525,15 +523,17 @@ async fn run_gpqa_benchmark(
.map(|dataset| (dataset.source.as_str(), dataset.split.as_str())),
&loaded.local_path,
);
let cases = apply_limit(loaded.cases, limit);
let cases = apply_limit(loaded.cases, options.limit);
let total = cases.len() as u64;
let started_at = Utc::now();
let started = Instant::now();
let mut base_request = request_template(&provider_config, &model, 0.0, max_tokens);
base_request.stream = stream.unwrap_or(provider_config.stream);
base_request.raw_debug = raw_debug_config(debug_raw, &provider_name, &model);
base_request.thinking =
merged_thinking_config(provider_config.thinking.as_ref(), thinking_overrides);
let mut base_request = request_template(&provider_config, &model, 0.0, options.max_tokens);
base_request.stream = options.stream.unwrap_or(provider_config.stream);
base_request.raw_debug = raw_debug_config(options.debug_raw, &provider_name, &model);
base_request.thinking = merged_thinking_config(
provider_config.thinking.as_ref(),
options.thinking_overrides,
);
let protocol = provider_config.protocol;

let pb = ProgressBar::new(total);
@@ -552,7 +552,7 @@ async fn run_gpqa_benchmark(
(case, result)
}
})
.buffer_unordered(nonzero_concurrency(concurrency));
.buffer_unordered(nonzero_concurrency(options.concurrency));

let mut metrics = Metrics::new();
let mut wrong_cases = Vec::new();
@@ -606,9 +606,9 @@ async fn run_gpqa_benchmark(
dataset,
started_at,
duration_ms: started.elapsed().as_millis(),
concurrency,
limit,
max_tokens,
concurrency: options.concurrency,
limit: options.limit,
max_tokens: options.max_tokens,
summary,
correct_samples,
wrong_cases,
@@ -1096,12 +1096,13 @@ fn provider_name(config: &AppConfig, provider: Option<&str>) -> Result<String> {
}

fn raw_debug_config(enabled: bool, provider: &str, model: &str) -> Option<RawDebugConfig> {
enabled.then(|| {
Some(
RawDebugConfig::new(
PathBuf::from("outputs/debug"),
format!("{provider}-{model}"),
)
})
.with_success_raw(enabled),
)
}

#[derive(Debug, Default)]
@@ -1574,6 +1575,21 @@ mod tests {
assert_eq!(stream, Some(false));
}

#[test]
fn bench_command_defaults_max_tokens_to_32768() {
let cli = Cli::try_parse_from(["lq_token_test", "bench", "gpqa-diamond"])
.expect("parse gpqa bench");

let Command::Bench {
command: BenchCommand::GpqaDiamond { max_tokens, .. },
} = cli.command
else {
panic!("expected gpqa-diamond bench command");
};

assert_eq!(max_tokens, 32_768);
}

#[test]
fn check_command_parses_debug_raw() {
let cli = Cli::try_parse_from([
@@ -1594,6 +1610,13 @@ mod tests {
assert!(debug_raw);
}

#[test]
fn raw_debug_config_is_created_even_when_success_raw_is_disabled() {
let raw_debug = raw_debug_config(false, "google", "gemini-test");

assert!(raw_debug.is_some());
}

#[test]
fn check_command_parses_thinking_overrides() {
let cli = Cli::try_parse_from([


+ 28
- 0
src/config.rs Voir le fichier

@@ -32,6 +32,7 @@ pub enum ConfigError {
pub enum ProtocolKind {
Openai,
Anthropic,
Google,
}

#[derive(Clone, Deserialize)]
@@ -357,6 +358,33 @@ providers:
assert_eq!(provider.api_token, "anthropic-secret");
}

#[test]
fn loads_google_protocol_provider() {
let config = AppConfig::load_from_str_with_env(
r#"
default_provider: google
providers:
google:
protocol: google
base_url: https://generativelanguage.googleapis.com/v1beta
api_token: ${GOOGLE_API_KEY}
default_model: gemini-3-pro-preview
"#,
|_| Err(()),
)
.expect("load config should not require provider tokens");

let provider = config
.resolved_provider_with_env(None, |name| match name {
"GOOGLE_API_KEY" => Ok("google-secret".to_string()),
_ => Err(()),
})
.expect("google provider");

assert_eq!(provider.protocol, ProtocolKind::Google);
assert_eq!(provider.api_token, "google-secret");
}

#[test]
fn provider_stream_defaults_false_and_can_be_enabled() {
let config = AppConfig::load_from_str_with_env(


+ 404
- 17
src/protocols/anthropic.rs Voir le fichier

@@ -1,4 +1,7 @@
use crate::runner::{ModelRequest, ModelResponse};
use crate::runner::{
HttpDebugRequest, HttpDebugResponse, ModelRequest, ModelResponse, request_headers_for_debug,
response_headers_for_debug,
};
use anyhow::{Context, Result, bail};
use futures::StreamExt;
use reqwest::Client;
@@ -8,22 +11,48 @@ use std::time::Instant;

pub async fn send(client: &Client, request: &ModelRequest) -> Result<ModelResponse> {
let url = super::endpoint_url(&request.base_url, "/v1/messages")?;
let debug_url = url.to_string();
let request_body = request_body(request, false);
let response = client
.post(url)
.header("x-api-key", &request.api_token)
.header("anthropic-version", "2023-06-01")
.json(&request_body(request, false))
.json(&request_body)
.send()
.await
.map_err(|error| {
write_request_error_debug_blocking(
request,
&debug_url,
request_body.clone(),
format!("failed to send Anthropic messages request: {error}"),
"anthropic-request-error",
);
error
})
.context("failed to send Anthropic messages request")?;

let status = response.status();
let status_code = status.as_u16();
let response_headers = response.headers().clone();
let body = response
.text()
.await
.map_err(|error| {
write_request_error_debug_blocking(
request,
&debug_url,
request_body.clone(),
format!("failed to read Anthropic response body: {error}"),
"anthropic-request-error",
);
error
})
.context("failed to read Anthropic response body")?;
if let Some(raw_debug) = &request.raw_debug {
if status.is_success()
&& let Some(raw_debug) = &request.raw_debug
&& raw_debug.write_success_raw()
{
raw_debug
.write_response("anthropic-json", &body)
.await
@@ -31,22 +60,54 @@ pub async fn send(client: &Client, request: &ModelRequest) -> Result<ModelRespon
}

if !status.is_success() {
write_error_debug(
request,
&debug_url,
request_body,
status_code,
&response_headers,
&body,
"anthropic-error-http",
)
.await?;
bail!(
"{}",
super::upstream_error_message("Anthropic", status_code, &body)
);
}

let parsed: AnthropicResponse =
serde_json::from_str(&body).context("failed to parse Anthropic response JSON")?;
let text = parsed
.content
.into_iter()
.find_map(|block| match block {
AnthropicContentBlock::Text { text } if !text.is_empty() => Some(text),
_ => None,
})
.context("Anthropic response missing first text content block")?;
let parsed: AnthropicResponse = match serde_json::from_str(&body) {
Ok(parsed) => parsed,
Err(error) => {
write_response_request_error_debug(
request,
&debug_url,
request_body,
status_code,
&response_headers,
&body,
format!("failed to parse Anthropic response JSON: {error}"),
)
.await?;
return Err(error).context("failed to parse Anthropic response JSON");
}
};
let Some(text) = parsed.content.into_iter().find_map(|block| match block {
AnthropicContentBlock::Text { text } if !text.is_empty() => Some(text),
_ => None,
}) else {
write_response_request_error_debug(
request,
&debug_url,
request_body,
status_code,
&response_headers,
&body,
"Anthropic response missing first text content block".to_string(),
)
.await?;
bail!("Anthropic response missing first text content block");
};

Ok(ModelResponse {
text,
@@ -58,27 +119,60 @@ pub async fn send(client: &Client, request: &ModelRequest) -> Result<ModelRespon

pub async fn send_stream(client: &Client, request: &ModelRequest) -> Result<ModelResponse> {
let url = super::endpoint_url(&request.base_url, "/v1/messages")?;
let debug_url = url.to_string();
let started = Instant::now();
let request_body = request_body(request, true);
let response = client
.post(url)
.header("x-api-key", &request.api_token)
.header("anthropic-version", "2023-06-01")
.json(&request_body(request, true))
.json(&request_body)
.send()
.await
.map_err(|error| {
write_request_error_debug_blocking(
request,
&debug_url,
request_body.clone(),
format!("failed to send Anthropic streaming request: {error}"),
"anthropic-request-error",
);
error
})
.context("failed to send Anthropic streaming request")?;

let status = response.status();
let status_code = status.as_u16();
let response_headers = response.headers().clone();

if !status.is_success() {
let body = response
.text()
.await
.map_err(|error| {
write_request_error_debug_blocking(
request,
&debug_url,
request_body.clone(),
format!("failed to read Anthropic error response body: {error}"),
"anthropic-request-error",
);
error
})
.context("failed to read Anthropic error response body")?;
if let Some(raw_debug) = &request.raw_debug {
raw_debug
.write_response("anthropic-error", &body)
.write_http_error(
"anthropic-error-http",
debug_request(request, &debug_url, request_body),
HttpDebugResponse {
status: Some(status_code),
headers: response_headers_for_debug(&response_headers),
body: body.clone(),
error_kind: None,
error: None,
},
)
.await
.context("failed to write Anthropic raw debug error response")?;
}
@@ -97,7 +191,18 @@ pub async fn send_stream(client: &Client, request: &ModelRequest) -> Result<Mode
let mut done = false;

while let Some(chunk) = stream.next().await {
let chunk = chunk.context("Anthropic stream interrupted")?;
let chunk = chunk
.map_err(|error| {
write_request_error_debug_blocking(
request,
&debug_url,
request_body.clone(),
format!("Anthropic stream interrupted: {error}"),
"anthropic-request-error",
);
error
})
.context("Anthropic stream interrupted")?;
for line in buffer.feed(&chunk) {
raw_stream.push_str(&line);
raw_stream.push('\n');
@@ -131,7 +236,9 @@ pub async fn send_stream(client: &Client, request: &ModelRequest) -> Result<Mode
}
}

if let Some(raw_debug) = &request.raw_debug {
if let Some(raw_debug) = &request.raw_debug
&& raw_debug.write_success_raw()
{
raw_debug
.write_response("anthropic-sse", &raw_stream)
.await
@@ -139,6 +246,16 @@ pub async fn send_stream(client: &Client, request: &ModelRequest) -> Result<Mode
}

if text.is_empty() {
write_response_request_error_debug(
request,
&debug_url,
request_body,
status_code,
&response_headers,
&raw_stream,
"Anthropic stream completed without producing any content".to_string(),
)
.await?;
bail!("Anthropic stream completed without producing any content");
}

@@ -209,6 +326,90 @@ fn request_body(request: &ModelRequest, stream: bool) -> Value {
body
}

async fn write_error_debug(
request: &ModelRequest,
url: &str,
request_body: Value,
status: u16,
response_headers: &reqwest::header::HeaderMap,
response_body: &str,
response_kind: &str,
) -> Result<()> {
if let Some(raw_debug) = &request.raw_debug {
raw_debug
.write_http_error(
response_kind,
debug_request(request, url, request_body),
HttpDebugResponse {
status: Some(status),
headers: response_headers_for_debug(response_headers),
body: response_body.to_string(),
error_kind: None,
error: None,
},
)
.await
.context("failed to write Anthropic raw debug error response")?;
}
Ok(())
}

async fn write_response_request_error_debug(
request: &ModelRequest,
url: &str,
request_body: Value,
status: u16,
response_headers: &reqwest::header::HeaderMap,
response_body: &str,
error: String,
) -> Result<()> {
if let Some(raw_debug) = &request.raw_debug {
raw_debug
.write_http_error(
"anthropic-request-error",
debug_request(request, url, request_body),
HttpDebugResponse {
status: Some(status),
headers: response_headers_for_debug(response_headers),
body: response_body.to_string(),
error_kind: Some("request_error".to_string()),
error: Some(error),
},
)
.await
.context("failed to write Anthropic raw debug request error response")?;
}
Ok(())
}

fn debug_request(request: &ModelRequest, url: &str, body: Value) -> HttpDebugRequest {
HttpDebugRequest {
method: "POST".to_string(),
url: url.to_string(),
headers: request_headers_for_debug(&[
("x-api-key", request.api_token.clone()),
("anthropic-version", "2023-06-01".to_string()),
]),
body,
}
}

fn write_request_error_debug_blocking(
request: &ModelRequest,
url: &str,
request_body: Value,
error: String,
response_kind: &str,
) {
if let Some(raw_debug) = &request.raw_debug {
let _ = futures::executor::block_on(raw_debug.write_request_error(
response_kind,
debug_request(request, url, request_body),
error,
));
}
}

#[cfg(test)]
mod tests {
use crate::runner::{ModelRequest, RawDebugConfig, ThinkingConfig};
@@ -409,6 +610,123 @@ mod tests {
assert!(!message.contains("sk-leaked-token"));
}

#[tokio::test]
async fn non_success_error_debug_records_request_and_response_without_token() {
let server = MockServer::start().await;
let temp_dir = tempfile::tempdir().expect("create temp dir");
Mock::given(method("POST"))
.and(path("/v1/messages"))
.respond_with(
ResponseTemplate::new(504)
.insert_header("x-request-id", "anthropic-req")
.set_body_json(serde_json::json!({
"error": {
"message": "gateway timeout"
}
})),
)
.mount(&server)
.await;

let request = ModelRequest {
base_url: server.uri(),
api_token: "anthropic-secret-token".to_string(),
model: "claude-test".to_string(),
prompt: "full anthropic prompt".to_string(),
temperature: 0.0,
max_tokens: 1024,
stream: false,
raw_debug: Some(RawDebugConfig::new(
temp_dir.path().to_path_buf(),
"anthropic-claude-test".to_string(),
)),
thinking: None,
};

let _ = super::send(&Client::new(), &request)
.await
.expect_err("non-success should fail");

let debug_files = std::fs::read_dir(temp_dir.path())
.expect("read debug dir")
.collect::<Result<Vec<_>, _>>()
.expect("debug entries");
assert_eq!(debug_files.len(), 1);
let raw = std::fs::read_to_string(debug_files[0].path()).expect("read raw debug file");
let debug: serde_json::Value = serde_json::from_str(&raw).expect("debug json");

assert_eq!(debug["request"]["headers"]["x-api-key"], "[REDACTED]");
assert_eq!(
debug["request"]["body"]["messages"][0]["content"],
"full anthropic prompt"
);
assert_eq!(debug["response"]["status"], 504);
assert_eq!(
debug["response"]["headers"]["x-request-id"],
"anthropic-req"
);
assert!(
debug["response"]["body"]
.as_str()
.expect("response body")
.contains("gateway timeout")
);
assert!(!raw.contains("anthropic-secret-token"));
}

#[tokio::test]
async fn request_error_debug_records_request_and_local_error_without_token() {
let temp_dir = tempfile::tempdir().expect("create temp dir");
let request = ModelRequest {
base_url: "http://127.0.0.1:9".to_string(),
api_token: "anthropic-secret-token".to_string(),
model: "claude-test".to_string(),
prompt: "prompt before anthropic connect error".to_string(),
temperature: 0.0,
max_tokens: 1024,
stream: false,
raw_debug: Some(RawDebugConfig::new(
temp_dir.path().to_path_buf(),
"anthropic-claude-test".to_string(),
)),
thinking: None,
};

let _ = super::send(&Client::new(), &request)
.await
.expect_err("connection should fail");

let debug_files = std::fs::read_dir(temp_dir.path())
.expect("read debug dir")
.collect::<Result<Vec<_>, _>>()
.expect("debug entries");
assert_eq!(debug_files.len(), 1);
let raw = std::fs::read_to_string(debug_files[0].path()).expect("read raw debug file");
let debug: serde_json::Value = serde_json::from_str(&raw).expect("debug json");

assert_eq!(debug["request"]["method"], "POST");
assert!(
debug["request"]["url"]
.as_str()
.expect("request url")
.ends_with("/v1/messages")
);
assert_eq!(debug["request"]["headers"]["x-api-key"], "[REDACTED]");
assert_eq!(
debug["request"]["body"]["messages"][0]["content"],
"prompt before anthropic connect error"
);
assert_eq!(debug["response"]["status"], serde_json::Value::Null);
assert_eq!(debug["response"]["error_kind"], "request_error");
assert!(
debug["response"]["error"]
.as_str()
.expect("local error")
.contains("failed to send Anthropic messages request")
);
assert!(!raw.contains("anthropic-secret-token"));
}

#[tokio::test]
async fn base_url_with_v1_prefix_does_not_duplicate_messages_path() {
let server = MockServer::start().await;
@@ -504,6 +822,75 @@ mod tests {
assert!(raw.contains("\"text\":\"hi \""));
}

#[tokio::test]
async fn empty_stream_request_error_debug_records_request_and_raw_stream() {
let server = MockServer::start().await;
let temp_dir = tempfile::tempdir().expect("create temp dir");
Mock::given(method("POST"))
.and(path("/v1/messages"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(
"event: message_start\n\
data: {\"type\":\"message_start\"}\n\n\
event: message_stop\n\
data: {\"type\":\"message_stop\"}\n\n",
),
)
.mount(&server)
.await;

let request = ModelRequest {
base_url: server.uri(),
api_token: "anthropic-secret-token".to_string(),
model: "claude-test".to_string(),
prompt: "prompt before empty stream".to_string(),
temperature: 0.0,
max_tokens: 1024,
stream: true,
raw_debug: Some(
RawDebugConfig::new(
temp_dir.path().to_path_buf(),
"anthropic-claude-test".to_string(),
)
.with_success_raw(false),
),
thinking: None,
};

let error = super::send_stream(&Client::new(), &request)
.await
.expect_err("empty stream should fail");
assert!(
error
.to_string()
.contains("completed without producing any content")
);

let debug_files = std::fs::read_dir(temp_dir.path())
.expect("read debug dir")
.collect::<Result<Vec<_>, _>>()
.expect("debug entries");
assert_eq!(debug_files.len(), 1);
let raw = std::fs::read_to_string(debug_files[0].path()).expect("read raw debug file");
let debug: serde_json::Value = serde_json::from_str(&raw).expect("debug json");

assert_eq!(
debug["request"]["body"]["messages"][0]["content"],
"prompt before empty stream"
);
assert_eq!(debug["response"]["status"], 200);
assert_eq!(debug["response"]["error_kind"], "request_error");
assert!(
debug["response"]["body"]
.as_str()
.expect("response body")
.contains("event: message_stop")
);
assert!(!raw.contains("anthropic-secret-token"));
}

fn anthropic_request(base_url: String) -> ModelRequest {
ModelRequest {
base_url,


+ 649
- 0
src/protocols/google.rs Voir le fichier

@@ -0,0 +1,649 @@
use crate::runner::{
HttpDebugRequest, HttpDebugResponse, ModelRequest, ModelResponse, request_headers_for_debug,
response_headers_for_debug,
};
use anyhow::{Context, Result, bail};
use futures::StreamExt;
use reqwest::Client;
use serde::Deserialize;
use serde_json::{Value, json};
use std::time::Instant;

pub async fn send(client: &Client, request: &ModelRequest) -> Result<ModelResponse> {
let url = google_endpoint(&request.base_url, &request.model, "generateContent")?;
let debug_url = url.to_string();
let request_body = request_body(request);
let response = client
.post(url)
.header("x-goog-api-key", &request.api_token)
.json(&request_body)
.send()
.await
.map_err(|error| {
write_request_error_debug_blocking(
request,
&debug_url,
request_body.clone(),
format!("failed to send Google generateContent request: {error}"),
"google-request-error",
);
error
})
.context("failed to send Google generateContent request")?;

let status = response.status();
let status_code = status.as_u16();
let response_headers = response.headers().clone();
let body = response
.text()
.await
.map_err(|error| {
write_request_error_debug_blocking(
request,
&debug_url,
request_body.clone(),
format!("failed to read Google response body: {error}"),
"google-request-error",
);
error
})
.context("failed to read Google response body")?;
if status.is_success()
&& let Some(raw_debug) = &request.raw_debug
&& raw_debug.write_success_raw()
{
raw_debug
.write_response("google-json", &body)
.await
.context("failed to write Google raw debug response")?;
}

if !status.is_success() {
write_error_debug(
request,
&debug_url,
request_body,
status_code,
&response_headers,
&body,
"google-error-http",
)
.await?;
bail!(
"{}",
super::upstream_error_message("Google", status_code, &body)
);
}

let parsed: GoogleResponse = match serde_json::from_str(&body) {
Ok(parsed) => parsed,
Err(error) => {
write_response_request_error_debug(
request,
&debug_url,
request_body,
status_code,
&response_headers,
&body,
format!("failed to parse Google response JSON: {error}"),
)
.await?;
return Err(error).context("failed to parse Google response JSON");
}
};
let Some(text) = response_text(parsed) else {
write_response_request_error_debug(
request,
&debug_url,
request_body,
status_code,
&response_headers,
&body,
"Google response missing non-thought text part".to_string(),
)
.await?;
bail!("Google response missing non-thought text part");
};

Ok(ModelResponse {
text,
status: status_code,
elapsed_ms: 0,
first_token_ms: None,
})
}

pub async fn send_stream(client: &Client, request: &ModelRequest) -> Result<ModelResponse> {
let url = google_endpoint(&request.base_url, &request.model, "streamGenerateContent")?;
let debug_url = url.to_string();
let started = Instant::now();
let request_body = request_body(request);
let response = client
.post(url)
.header("x-goog-api-key", &request.api_token)
.json(&request_body)
.send()
.await
.map_err(|error| {
write_request_error_debug_blocking(
request,
&debug_url,
request_body.clone(),
format!("failed to send Google streamGenerateContent request: {error}"),
"google-request-error",
);
error
})
.context("failed to send Google streamGenerateContent request")?;

let status = response.status();
let status_code = status.as_u16();
let response_headers = response.headers().clone();

if !status.is_success() {
let body = response
.text()
.await
.map_err(|error| {
write_request_error_debug_blocking(
request,
&debug_url,
request_body.clone(),
format!("failed to read Google error response body: {error}"),
"google-request-error",
);
error
})
.context("failed to read Google error response body")?;
if let Some(raw_debug) = &request.raw_debug {
raw_debug
.write_http_error(
"google-error-http",
debug_request(request, &debug_url, request_body),
HttpDebugResponse {
status: Some(status_code),
headers: response_headers_for_debug(&response_headers),
body: body.clone(),
error_kind: None,
error: None,
},
)
.await
.context("failed to write Google raw debug error response")?;
}
bail!(
"{}",
super::upstream_error_message("Google", status_code, &body)
);
}

let mut stream = response.bytes_stream();
let mut buffer = super::SseLineBuffer::new();
let mut text = String::new();
let mut raw_stream = String::new();
let mut first_token_ms: Option<u128> = None;

while let Some(chunk) = stream.next().await {
let chunk = chunk
.map_err(|error| {
write_request_error_debug_blocking(
request,
&debug_url,
request_body.clone(),
format!("Google stream interrupted: {error}"),
"google-request-error",
);
error
})
.context("Google stream interrupted")?;
for line in buffer.feed(&chunk) {
raw_stream.push_str(&line);
raw_stream.push('\n');
let Some(data) = line.strip_prefix("data: ") else {
continue;
};
let Some(content) = serde_json::from_str::<GoogleResponse>(data)
.ok()
.and_then(response_text)
.filter(|content| !content.is_empty())
else {
continue;
};
if first_token_ms.is_none() {
first_token_ms = Some(started.elapsed().as_millis());
}
text.push_str(&content);
}
}

if let Some(raw_debug) = &request.raw_debug
&& raw_debug.write_success_raw()
{
raw_debug
.write_response("google-sse", &raw_stream)
.await
.context("failed to write Google raw debug stream")?;
}

if text.is_empty() {
write_response_request_error_debug(
request,
&debug_url,
request_body,
status_code,
&response_headers,
&raw_stream,
"Google stream completed without producing any content".to_string(),
)
.await?;
bail!("Google stream completed without producing any content");
}

Ok(ModelResponse {
text,
status: status_code,
elapsed_ms: 0,
first_token_ms,
})
}

fn google_endpoint(base_url: &str, model: &str, action: &str) -> Result<reqwest::Url> {
let model_path = if model.starts_with("models/") {
model.to_string()
} else {
format!("models/{model}")
};
super::endpoint_url(base_url, &format!("/{model_path}:{action}"))
}

fn request_body(request: &ModelRequest) -> Value {
let mut body = json!({
"contents": [{
"role": "user",
"parts": [{"text": request.prompt}]
}],
"generationConfig": {
"temperature": request.temperature,
"maxOutputTokens": request.max_tokens
}
});

if let Some(thinking) = &request.thinking
&& thinking.enabled
{
let mut thinking_config = serde_json::Map::new();
if let Some(display) = &thinking.display {
thinking_config.insert(
"includeThoughts".to_string(),
json!(display != "omitted" && display != "false"),
);
}
if let Some(budget_tokens) = thinking.budget_tokens {
thinking_config.insert("thinkingBudget".to_string(), json!(budget_tokens));
}
if let Some(effort) = &thinking.effort {
thinking_config.insert("thinkingLevel".to_string(), json!(effort));
}
if !thinking_config.is_empty() {
body["generationConfig"]["thinkingConfig"] = Value::Object(thinking_config);
}
}

body
}

fn response_text(response: GoogleResponse) -> Option<String> {
let text = response
.candidates
.into_iter()
.flat_map(|candidate| candidate.content.parts)
.filter(|part| !part.thought.unwrap_or(false))
.filter_map(|part| part.text)
.filter(|text| !text.is_empty())
.collect::<String>();
(!text.is_empty()).then_some(text)
}

async fn write_error_debug(
request: &ModelRequest,
url: &str,
request_body: Value,
status: u16,
response_headers: &reqwest::header::HeaderMap,
response_body: &str,
response_kind: &str,
) -> Result<()> {
if let Some(raw_debug) = &request.raw_debug {
raw_debug
.write_http_error(
response_kind,
debug_request(request, url, request_body),
HttpDebugResponse {
status: Some(status),
headers: response_headers_for_debug(response_headers),
body: response_body.to_string(),
error_kind: None,
error: None,
},
)
.await
.context("failed to write Google raw debug error response")?;
}
Ok(())
}

async fn write_response_request_error_debug(
request: &ModelRequest,
url: &str,
request_body: Value,
status: u16,
response_headers: &reqwest::header::HeaderMap,
response_body: &str,
error: String,
) -> Result<()> {
if let Some(raw_debug) = &request.raw_debug {
raw_debug
.write_http_error(
"google-request-error",
debug_request(request, url, request_body),
HttpDebugResponse {
status: Some(status),
headers: response_headers_for_debug(response_headers),
body: response_body.to_string(),
error_kind: Some("request_error".to_string()),
error: Some(error),
},
)
.await
.context("failed to write Google raw debug request error response")?;
}
Ok(())
}

fn debug_request(request: &ModelRequest, url: &str, body: Value) -> HttpDebugRequest {
HttpDebugRequest {
method: "POST".to_string(),
url: url.to_string(),
headers: request_headers_for_debug(&[("x-goog-api-key", request.api_token.clone())]),
body,
}
}

fn write_request_error_debug_blocking(
request: &ModelRequest,
url: &str,
request_body: Value,
error: String,
response_kind: &str,
) {
if let Some(raw_debug) = &request.raw_debug {
let _ = futures::executor::block_on(raw_debug.write_request_error(
response_kind,
debug_request(request, url, request_body),
error,
));
}
}

#[derive(Debug, Deserialize)]
struct GoogleResponse {
candidates: Vec<GoogleCandidate>,
}

#[derive(Debug, Deserialize)]
struct GoogleCandidate {
content: GoogleContent,
}

#[derive(Debug, Deserialize)]
struct GoogleContent {
parts: Vec<GooglePart>,
}

#[derive(Debug, Deserialize)]
struct GooglePart {
text: Option<String>,
#[serde(default)]
thought: Option<bool>,
}

#[cfg(test)]
mod tests {
use crate::runner::{ModelRequest, RawDebugConfig, ThinkingConfig};
use reqwest::Client;
use wiremock::matchers::{body_json, header, method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};

#[tokio::test]
async fn sends_generate_content_with_thinking_config_and_extracts_text() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/models/gemini-test:generateContent"))
.and(header("x-goog-api-key", "test-token"))
.and(body_json(serde_json::json!({
"contents": [{
"role": "user",
"parts": [{"text": "hello"}]
}],
"generationConfig": {
"temperature": 0.0,
"maxOutputTokens": 1024,
"thinkingConfig": {
"includeThoughts": true,
"thinkingBudget": 5000,
"thinkingLevel": "high"
}
}
})))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"candidates": [{
"content": {
"parts": [
{"text": "hidden thought", "thought": true},
{"text": "hi from gemini"}
]
}
}]
})))
.mount(&server)
.await;

let mut request = google_request(server.uri());
request.thinking = Some(ThinkingConfig {
enabled: true,
kind: None,
budget_tokens: Some(5000),
effort: Some("high".to_string()),
display: Some("summarized".to_string()),
reasoning_effort: None,
reasoning_summary: None,
});

let response = super::send(&Client::new(), &request)
.await
.expect("response");

assert_eq!(response.status, 200);
assert_eq!(response.text, "hi from gemini");
}

#[tokio::test]
async fn sends_stream_generate_content_and_extracts_non_thought_parts() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/models/gemini-test:streamGenerateContent"))
.and(header("x-goog-api-key", "test-token"))
.and(body_json(serde_json::json!({
"contents": [{
"role": "user",
"parts": [{"text": "hello"}]
}],
"generationConfig": {
"temperature": 0.0,
"maxOutputTokens": 1024,
"thinkingConfig": {
"includeThoughts": false,
"thinkingBudget": 0
}
}
})))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(
"data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"hi \"}]}}]}\n\n\
data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"thinking\",\"thought\":true},{\"text\":\"there\"}]}}]}\n\n",
),
)
.mount(&server)
.await;

let mut request = google_request(server.uri());
request.stream = true;
request.thinking = Some(ThinkingConfig {
enabled: true,
kind: None,
budget_tokens: Some(0),
effort: None,
display: Some("omitted".to_string()),
reasoning_effort: None,
reasoning_summary: None,
});

let response = super::send_stream(&Client::new(), &request)
.await
.expect("stream response");

assert_eq!(response.status, 200);
assert_eq!(response.text, "hi there");
assert!(response.first_token_ms.is_some());
}

#[tokio::test]
async fn non_success_error_debug_records_request_and_response_without_token() {
let server = MockServer::start().await;
let temp_dir = tempfile::tempdir().expect("create temp dir");
Mock::given(method("POST"))
.and(path("/models/gemini-test:generateContent"))
.respond_with(
ResponseTemplate::new(400)
.insert_header("x-request-id", "google-req")
.set_body_json(serde_json::json!({
"error": {
"message": "bad google request"
}
})),
)
.mount(&server)
.await;

let request = ModelRequest {
base_url: server.uri(),
api_token: "google-secret-token".to_string(),
model: "gemini-test".to_string(),
prompt: "full google prompt".to_string(),
temperature: 0.0,
max_tokens: 1024,
stream: false,
raw_debug: Some(RawDebugConfig::new(
temp_dir.path().to_path_buf(),
"google-gemini-test".to_string(),
)),
thinking: None,
};

let _ = super::send(&Client::new(), &request)
.await
.expect_err("non-success should fail");

let debug_files = std::fs::read_dir(temp_dir.path())
.expect("read debug dir")
.collect::<Result<Vec<_>, _>>()
.expect("debug entries");
assert_eq!(debug_files.len(), 1);
let raw = std::fs::read_to_string(debug_files[0].path()).expect("read raw debug file");
let debug: serde_json::Value = serde_json::from_str(&raw).expect("debug json");

assert_eq!(debug["request"]["headers"]["x-goog-api-key"], "[REDACTED]");
assert_eq!(
debug["request"]["body"]["contents"][0]["parts"][0]["text"],
"full google prompt"
);
assert_eq!(debug["response"]["status"], 400);
assert_eq!(debug["response"]["headers"]["x-request-id"], "google-req");
assert!(
debug["response"]["body"]
.as_str()
.expect("response body")
.contains("bad google request")
);
assert!(!raw.contains("google-secret-token"));
}

#[tokio::test]
async fn request_error_debug_records_request_and_local_error_without_token() {
let temp_dir = tempfile::tempdir().expect("create temp dir");
let request = ModelRequest {
base_url: "http://127.0.0.1:9".to_string(),
api_token: "google-secret-token".to_string(),
model: "gemini-test".to_string(),
prompt: "prompt before google connect error".to_string(),
temperature: 0.0,
max_tokens: 1024,
stream: false,
raw_debug: Some(RawDebugConfig::new(
temp_dir.path().to_path_buf(),
"google-gemini-test".to_string(),
)),
thinking: None,
};

let _ = super::send(&Client::new(), &request)
.await
.expect_err("connection should fail");

let debug_files = std::fs::read_dir(temp_dir.path())
.expect("read debug dir")
.collect::<Result<Vec<_>, _>>()
.expect("debug entries");
assert_eq!(debug_files.len(), 1);
let raw = std::fs::read_to_string(debug_files[0].path()).expect("read raw debug file");
let debug: serde_json::Value = serde_json::from_str(&raw).expect("debug json");

assert_eq!(debug["request"]["method"], "POST");
assert!(
debug["request"]["url"]
.as_str()
.expect("request url")
.ends_with("/models/gemini-test:generateContent")
);
assert_eq!(debug["request"]["headers"]["x-goog-api-key"], "[REDACTED]");
assert_eq!(
debug["request"]["body"]["contents"][0]["parts"][0]["text"],
"prompt before google connect error"
);
assert_eq!(debug["response"]["status"], serde_json::Value::Null);
assert_eq!(debug["response"]["error_kind"], "request_error");
assert!(
debug["response"]["error"]
.as_str()
.expect("local error")
.contains("failed to send Google generateContent request")
);
assert!(!raw.contains("google-secret-token"));
}

fn google_request(base_url: String) -> ModelRequest {
ModelRequest {
base_url,
api_token: "test-token".to_string(),
model: "gemini-test".to_string(),
prompt: "hello".to_string(),
temperature: 0.0,
max_tokens: 1024,
stream: false,
raw_debug: None,
thinking: None,
}
}
}

+ 1
- 0
src/protocols/mod.rs Voir le fichier

@@ -1,4 +1,5 @@
pub mod anthropic;
pub mod google;
pub mod openai;

use anyhow::{Context, Result};


+ 333
- 11
src/protocols/openai.rs Voir le fichier

@@ -1,4 +1,7 @@
use crate::runner::{ModelRequest, ModelResponse};
use crate::runner::{
HttpDebugRequest, HttpDebugResponse, ModelRequest, ModelResponse, request_headers_for_debug,
response_headers_for_debug,
};
use anyhow::{Context, Result, bail};
use futures::StreamExt;
use reqwest::Client;
@@ -8,21 +11,47 @@ use std::time::Instant;

pub async fn send(client: &Client, request: &ModelRequest) -> Result<ModelResponse> {
let url = super::endpoint_url(&request.base_url, "/chat/completions")?;
let debug_url = url.to_string();
let request_body = request_body(request, false);
let response = client
.post(url)
.bearer_auth(&request.api_token)
.json(&request_body(request, false))
.json(&request_body)
.send()
.await
.map_err(|error| {
write_request_error_debug_blocking(
request,
&debug_url,
request_body.clone(),
format!("failed to send OpenAI chat completion request: {error}"),
"openai-request-error",
);
error
})
.context("failed to send OpenAI chat completion request")?;

let status = response.status();
let status_code = status.as_u16();
let response_headers = response.headers().clone();
let body = response
.text()
.await
.map_err(|error| {
write_request_error_debug_blocking(
request,
&debug_url,
request_body.clone(),
format!("failed to read OpenAI response body: {error}"),
"openai-request-error",
);
error
})
.context("failed to read OpenAI response body")?;
if let Some(raw_debug) = &request.raw_debug {
if status.is_success()
&& let Some(raw_debug) = &request.raw_debug
&& raw_debug.write_success_raw()
{
raw_debug
.write_response("openai-json", &body)
.await
@@ -30,21 +59,57 @@ pub async fn send(client: &Client, request: &ModelRequest) -> Result<ModelRespon
}

if !status.is_success() {
write_error_debug(
request,
&debug_url,
request_body,
status_code,
&response_headers,
&body,
"openai-error-http",
)
.await?;
bail!(
"{}",
super::upstream_error_message("OpenAI", status_code, &body)
);
}

let parsed: OpenAiResponse =
serde_json::from_str(&body).context("failed to parse OpenAI response JSON")?;
let text = parsed
let parsed: OpenAiResponse = match serde_json::from_str(&body) {
Ok(parsed) => parsed,
Err(error) => {
write_response_request_error_debug(
request,
&debug_url,
request_body,
status_code,
&response_headers,
&body,
format!("failed to parse OpenAI response JSON: {error}"),
)
.await?;
return Err(error).context("failed to parse OpenAI response JSON");
}
};
let Some(text) = parsed
.choices
.into_iter()
.next()
.and_then(|choice| choice.message.content)
.filter(|content| !content.is_empty())
.context("OpenAI response missing choices[0].message.content")?;
else {
write_response_request_error_debug(
request,
&debug_url,
request_body,
status_code,
&response_headers,
&body,
"OpenAI response missing choices[0].message.content".to_string(),
)
.await?;
bail!("OpenAI response missing choices[0].message.content");
};

Ok(ModelResponse {
text,
@@ -56,17 +121,30 @@ pub async fn send(client: &Client, request: &ModelRequest) -> Result<ModelRespon

pub async fn send_stream(client: &Client, request: &ModelRequest) -> Result<ModelResponse> {
let url = super::endpoint_url(&request.base_url, "/chat/completions")?;
let debug_url = url.to_string();
let started = Instant::now();
let request_body = request_body(request, true);
let response = client
.post(url)
.bearer_auth(&request.api_token)
.json(&request_body(request, true))
.json(&request_body)
.send()
.await
.map_err(|error| {
write_request_error_debug_blocking(
request,
&debug_url,
request_body.clone(),
format!("failed to send OpenAI streaming request: {error}"),
"openai-request-error",
);
error
})
.context("failed to send OpenAI streaming request")?;

let status = response.status();
let status_code = status.as_u16();
let response_headers = response.headers().clone();

if !status.is_success() {
let body = response
@@ -75,7 +153,17 @@ pub async fn send_stream(client: &Client, request: &ModelRequest) -> Result<Mode
.context("failed to read OpenAI error response body")?;
if let Some(raw_debug) = &request.raw_debug {
raw_debug
.write_response("openai-error", &body)
.write_http_error(
"openai-error-http",
debug_request(request, &debug_url, request_body),
HttpDebugResponse {
status: Some(status_code),
headers: response_headers_for_debug(&response_headers),
body: body.clone(),
error_kind: None,
error: None,
},
)
.await
.context("failed to write OpenAI raw debug error response")?;
}
@@ -93,7 +181,18 @@ pub async fn send_stream(client: &Client, request: &ModelRequest) -> Result<Mode
let mut done = false;

while let Some(chunk) = stream.next().await {
let chunk = chunk.context("OpenAI stream interrupted")?;
let chunk = chunk
.map_err(|error| {
write_request_error_debug_blocking(
request,
&debug_url,
request_body.clone(),
format!("OpenAI stream interrupted: {error}"),
"openai-request-error",
);
error
})
.context("OpenAI stream interrupted")?;
for line in buffer.feed(&chunk) {
raw_stream.push_str(&line);
raw_stream.push('\n');
@@ -121,7 +220,9 @@ pub async fn send_stream(client: &Client, request: &ModelRequest) -> Result<Mode
}
}

if let Some(raw_debug) = &request.raw_debug {
if let Some(raw_debug) = &request.raw_debug
&& raw_debug.write_success_raw()
{
raw_debug
.write_response("openai-sse", &raw_stream)
.await
@@ -129,6 +230,16 @@ pub async fn send_stream(client: &Client, request: &ModelRequest) -> Result<Mode
}

if text.is_empty() {
write_response_request_error_debug(
request,
&debug_url,
request_body,
status_code,
&response_headers,
&raw_stream,
"OpenAI stream completed without producing any content".to_string(),
)
.await?;
bail!("OpenAI stream completed without producing any content");
}

@@ -200,6 +311,90 @@ fn request_body(request: &ModelRequest, stream: bool) -> Value {
body
}

async fn write_error_debug(
request: &ModelRequest,
url: &str,
request_body: Value,
status: u16,
response_headers: &reqwest::header::HeaderMap,
response_body: &str,
response_kind: &str,
) -> Result<()> {
if let Some(raw_debug) = &request.raw_debug {
raw_debug
.write_http_error(
response_kind,
debug_request(request, url, request_body),
HttpDebugResponse {
status: Some(status),
headers: response_headers_for_debug(response_headers),
body: response_body.to_string(),
error_kind: None,
error: None,
},
)
.await
.context("failed to write OpenAI raw debug error response")?;
}
Ok(())
}

async fn write_response_request_error_debug(
request: &ModelRequest,
url: &str,
request_body: Value,
status: u16,
response_headers: &reqwest::header::HeaderMap,
response_body: &str,
error: String,
) -> Result<()> {
if let Some(raw_debug) = &request.raw_debug {
raw_debug
.write_http_error(
"openai-request-error",
debug_request(request, url, request_body),
HttpDebugResponse {
status: Some(status),
headers: response_headers_for_debug(response_headers),
body: response_body.to_string(),
error_kind: Some("request_error".to_string()),
error: Some(error),
},
)
.await
.context("failed to write OpenAI raw debug request error response")?;
}
Ok(())
}

fn debug_request(request: &ModelRequest, url: &str, body: Value) -> HttpDebugRequest {
HttpDebugRequest {
method: "POST".to_string(),
url: url.to_string(),
headers: request_headers_for_debug(&[(
"authorization",
format!("Bearer {}", request.api_token),
)]),
body,
}
}

fn write_request_error_debug_blocking(
request: &ModelRequest,
url: &str,
request_body: Value,
error: String,
response_kind: &str,
) {
if let Some(raw_debug) = &request.raw_debug {
let _ = futures::executor::block_on(raw_debug.write_request_error(
response_kind,
debug_request(request, url, request_body),
error,
));
}
}

#[cfg(test)]
mod tests {
use crate::runner::{ModelRequest, RawDebugConfig, ThinkingConfig};
@@ -359,6 +554,133 @@ mod tests {
assert!(!message.contains("sk-leaked-token"));
}

#[tokio::test]
async fn non_success_error_debug_records_request_and_response_without_token() {
let server = MockServer::start().await;
let temp_dir = tempfile::tempdir().expect("create temp dir");
Mock::given(method("POST"))
.and(path("/chat/completions"))
.respond_with(
ResponseTemplate::new(429)
.insert_header("x-request-id", "req-test")
.set_body_json(serde_json::json!({
"error": {
"message": "rate limited with full response"
}
})),
)
.mount(&server)
.await;

let request = ModelRequest {
base_url: server.uri(),
api_token: "sk-real-token".to_string(),
model: "gpt-test".to_string(),
prompt: "full prompt should be recorded".to_string(),
temperature: 0.0,
max_tokens: 1024,
stream: false,
raw_debug: Some(RawDebugConfig::new(
temp_dir.path().to_path_buf(),
"openai-gpt-test".to_string(),
)),
thinking: None,
};

let _ = super::send(&Client::new(), &request)
.await
.expect_err("non-success should fail");

let debug_files = std::fs::read_dir(temp_dir.path())
.expect("read debug dir")
.collect::<Result<Vec<_>, _>>()
.expect("debug entries");
assert_eq!(debug_files.len(), 1);
let raw = std::fs::read_to_string(debug_files[0].path()).expect("read raw debug file");
let debug: serde_json::Value = serde_json::from_str(&raw).expect("debug json");

assert_eq!(debug["request"]["method"], "POST");
assert!(
debug["request"]["url"]
.as_str()
.expect("request url")
.ends_with("/chat/completions")
);
assert_eq!(
debug["request"]["headers"]["authorization"],
"Bearer [REDACTED]"
);
assert_eq!(
debug["request"]["body"]["messages"][0]["content"],
"full prompt should be recorded"
);
assert_eq!(debug["response"]["status"], 429);
assert_eq!(debug["response"]["headers"]["x-request-id"], "req-test");
assert!(
debug["response"]["body"]
.as_str()
.expect("response body")
.contains("rate limited with full response")
);
assert!(!raw.contains("sk-real-token"));
}

#[tokio::test]
async fn request_error_debug_records_request_and_local_error_without_token() {
let temp_dir = tempfile::tempdir().expect("create temp dir");
let request = ModelRequest {
base_url: "http://127.0.0.1:9".to_string(),
api_token: "sk-real-token".to_string(),
model: "gpt-test".to_string(),
prompt: "prompt before connect error".to_string(),
temperature: 0.0,
max_tokens: 1024,
stream: false,
raw_debug: Some(RawDebugConfig::new(
temp_dir.path().to_path_buf(),
"openai-gpt-test".to_string(),
)),
thinking: None,
};

let _ = super::send(&Client::new(), &request)
.await
.expect_err("connection should fail");

let debug_files = std::fs::read_dir(temp_dir.path())
.expect("read debug dir")
.collect::<Result<Vec<_>, _>>()
.expect("debug entries");
assert_eq!(debug_files.len(), 1);
let raw = std::fs::read_to_string(debug_files[0].path()).expect("read raw debug file");
let debug: serde_json::Value = serde_json::from_str(&raw).expect("debug json");

assert_eq!(debug["request"]["method"], "POST");
assert!(
debug["request"]["url"]
.as_str()
.expect("request url")
.ends_with("/chat/completions")
);
assert_eq!(
debug["request"]["headers"]["authorization"],
"Bearer [REDACTED]"
);
assert_eq!(
debug["request"]["body"]["messages"][0]["content"],
"prompt before connect error"
);
assert_eq!(debug["response"]["status"], serde_json::Value::Null);
assert_eq!(debug["response"]["error_kind"], "request_error");
assert!(
debug["response"]["error"]
.as_str()
.expect("local error")
.contains("failed to send OpenAI chat completion request")
);
assert!(!raw.contains("sk-real-token"));
}

#[tokio::test]
async fn base_url_with_v1_prefix_keeps_chat_completion_path() {
let server = MockServer::start().await;


+ 123
- 0
src/runner.rs Voir le fichier

@@ -3,6 +3,10 @@ use crate::protocols;
use anyhow::{Context, Result};
use chrono::Utc;
use reqwest::Client;
use reqwest::header::HeaderMap;
use serde::Serialize;
use serde_json::Value;
use std::collections::BTreeMap;
use std::fmt;
use std::path::PathBuf;
use std::sync::Arc;
@@ -38,6 +42,7 @@ pub struct RawDebugConfig {
output_dir: PathBuf,
prefix: String,
counter: Arc<AtomicU64>,
write_success_raw: bool,
}

impl RawDebugConfig {
@@ -46,10 +51,57 @@ impl RawDebugConfig {
output_dir,
prefix: sanitize_filename_component(&prefix),
counter: Arc::new(AtomicU64::new(0)),
write_success_raw: true,
}
}

pub fn with_success_raw(mut self, write_success_raw: bool) -> Self {
self.write_success_raw = write_success_raw;
self
}

pub fn write_success_raw(&self) -> bool {
self.write_success_raw
}

pub async fn write_response(&self, response_kind: &str, contents: &str) -> Result<PathBuf> {
self.write_debug_file(response_kind, contents).await
}

pub async fn write_http_error(
&self,
response_kind: &str,
request: HttpDebugRequest,
response: HttpDebugResponse,
) -> Result<PathBuf> {
let envelope = HttpDebugEnvelope { request, response };
let contents = serde_json::to_string_pretty(&envelope)
.context("failed to serialize raw debug HTTP error")?;
self.write_debug_file(response_kind, &contents).await
}

pub async fn write_request_error(
&self,
response_kind: &str,
request: HttpDebugRequest,
error: String,
) -> Result<PathBuf> {
let envelope = HttpDebugEnvelope {
request,
response: HttpDebugResponse {
status: None,
headers: BTreeMap::new(),
body: String::new(),
error_kind: Some("request_error".to_string()),
error: Some(error),
},
};
let contents = serde_json::to_string_pretty(&envelope)
.context("failed to serialize raw debug request error")?;
self.write_debug_file(response_kind, &contents).await
}

async fn write_debug_file(&self, response_kind: &str, contents: &str) -> Result<PathBuf> {
tokio::fs::create_dir_all(&self.output_dir)
.await
.with_context(|| {
@@ -74,6 +126,75 @@ impl RawDebugConfig {
}
}

#[derive(Debug, Serialize)]
pub struct HttpDebugEnvelope {
pub request: HttpDebugRequest,
pub response: HttpDebugResponse,
}

#[derive(Debug, Serialize)]
pub struct HttpDebugRequest {
pub method: String,
pub url: String,
pub headers: BTreeMap<String, String>,
pub body: Value,
}

#[derive(Debug, Serialize)]
pub struct HttpDebugResponse {
pub status: Option<u16>,
pub headers: BTreeMap<String, String>,
pub body: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub error_kind: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
}

pub fn request_headers_for_debug(headers: &[(&str, String)]) -> BTreeMap<String, String> {
headers
.iter()
.map(|(name, value)| {
let name = name.to_ascii_lowercase();
let value = if is_sensitive_header(&name) {
redact_header_value(&name, value)
} else {
value.clone()
};
(name, value)
})
.collect()
}

pub fn response_headers_for_debug(headers: &HeaderMap) -> BTreeMap<String, String> {
headers
.iter()
.map(|(name, value)| {
(
name.as_str().to_ascii_lowercase(),
value.to_str().unwrap_or("<non-utf8>").to_string(),
)
})
.collect()
}

fn is_sensitive_header(name: &str) -> bool {
matches!(
name,
"authorization" | "x-api-key" | "x-goog-api-key" | "api-key"
)
}

fn redact_header_value(name: &str, value: &str) -> String {
if name == "authorization"
&& let Some((scheme, _)) = value.split_once(' ')
{
format!("{scheme} [REDACTED]")
} else {
"[REDACTED]".to_string()
}
}

impl fmt::Debug for ModelRequest {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
@@ -140,11 +261,13 @@ pub async fn run_model_request_with_client(
match protocol {
ProtocolKind::Openai => protocols::openai::send_stream(client, request).await?,
ProtocolKind::Anthropic => protocols::anthropic::send_stream(client, request).await?,
ProtocolKind::Google => protocols::google::send_stream(client, request).await?,
}
} else {
match protocol {
ProtocolKind::Openai => protocols::openai::send(client, request).await?,
ProtocolKind::Anthropic => protocols::anthropic::send(client, request).await?,
ProtocolKind::Google => protocols::google::send(client, request).await?,
}
};
response.elapsed_ms = started.elapsed().as_millis();


Chargement…
Annuler
Enregistrer