Compare commits
3 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 4d093548de | |||
| aaa593ac31 | |||
| 24a8dec151 |
@@ -1,100 +0,0 @@
|
||||
# NIT AI Studio — 第一阶段(聚合图像/数字人)设计
|
||||
|
||||
日期:2026-06-17
|
||||
分支:`nit_edu`
|
||||
状态:已通过 brainstorming,待落实现计划
|
||||
|
||||
## 背景
|
||||
|
||||
NIT 精英班专属的「AI 综合应用软件」是一套面向精英班学员的 AI 创作工具,完整愿景涵盖
|
||||
文本结构化创作(论文/小说/剧本/分镜/教材)、PPT 生成、图像(P图/制图)、视频
|
||||
(文生/图生/短视频)、数字人。体量过大,分阶段建设。
|
||||
|
||||
**本 spec 仅覆盖 AI Studio 第一阶段:聚合现有图像 + 数字人能力**,为精英班学员提供
|
||||
统一创作入口,复用 BoosAPI 已有的 `picture_agent`(图像)与 `digital_human`(数字人)。
|
||||
|
||||
## 已确认的关键决策
|
||||
|
||||
1. 定位:NIT 精英班专属模块,非通用套件。
|
||||
2. 第一阶段只聚合 `picture_agent`(图像)+ `digital_human`(数字人)。
|
||||
3. 视频(含 money_printer、工作流视频)留后续阶段,本阶段完全不碰 money_printer。
|
||||
4. 认证桥接:给精英班学员配后台下游 key + 扩展 surface 的 resolve_auth 接受 student JWT。
|
||||
5. 计费:学员创作消耗走积分轨(已建的 points track),管理员可给学员充积分。
|
||||
6. 文本创作 / PPT 留后续阶段。
|
||||
|
||||
## 前置依赖
|
||||
|
||||
本阶段 gate 在第一期(报考闭环)产物上,实现排在第一期之后:
|
||||
- student JWT(`audience='student'`,第一期发放)。
|
||||
- `enrollment.status=approved` 且 `tier=elite`。
|
||||
- 非精英班 / 未审核 → 403。
|
||||
|
||||
## 非目标(本阶段明确不做)
|
||||
|
||||
- 视频能力(文生/图生/短视频、money_printer、工作流视频)。
|
||||
- 文本结构化创作(论文/小说/剧本/分镜/教材)。
|
||||
- PPT 生成。
|
||||
- 通用套件(面向非 NIT 用户)。
|
||||
|
||||
## 学员下游 key 配置
|
||||
|
||||
- `downstream_api_keys` 加可空 `student_id` 列(关联 `students.id`,与现有 `user_id` 并列;
|
||||
二者互斥:一行要么属于 user,要么属于 student)。
|
||||
- 审核通过且 `tier=elite` 时,自动 provision 一个学员下游 key,沿用现有 `comfyui-agent`
|
||||
内置 key 的幂等 `ensure` 模式:缺失则建、被禁用则重新启用、保持唯一。
|
||||
- 学员从精英班降级或报考被驳回时,禁用该 key。
|
||||
- key 命名约定:`sk-nit-stu-{uuid_hex}`,`name='nit-student'`,便于识别。
|
||||
|
||||
## 认证扩展
|
||||
|
||||
仅改两个 surface:`src/routes/picture_agent.rs`、`src/routes/digital_human.rs`。
|
||||
**money_printer 不动。**
|
||||
|
||||
各自的 `resolve_auth` 加一个分支(现有 admin/proxy token、user JWT 逻辑不变):
|
||||
- 新增:`audience='student'` 的 JWT → 校验该学员存在 `approved` 且 `tier=elite` 的报考
|
||||
→ 查其 `student_id` 对应的下游 key → 作 `outbound_key`。
|
||||
- 抽共享 helper `resolve_student_outbound_key(student_id, state) -> AppResult<(key_id, key)>`,
|
||||
放在 `src/services/nit/`,两处 surface 复用,避免重复逻辑。
|
||||
- 校验失败(无 elite 报考 / key 缺失)→ 403。
|
||||
|
||||
## 计费(复用积分轨)
|
||||
|
||||
- 学员用图像/数字人,消耗扣在其下游 key 上,走已建的积分轨(`apply_points_delta`,track=points)。
|
||||
- 管理员给学员充积分:扩展现有 admin 积分调整端点支持 `student_id`(或新增 NIT 学员积分端点,
|
||||
实现时择一,优先扩展现有端点)。
|
||||
- 精英班开通时可送一笔初始积分,数额为配置项,默认 `0`。
|
||||
|
||||
## 统一入口(前端)
|
||||
|
||||
- `src/web/pages/nit/Studio.tsx`:精英班学员创作工作台首页,卡片式列出「图像 / 数字人」
|
||||
两个入口 + 显示积分余额。
|
||||
- 卡片跳进现有 `picture_agent` / `digital_human` 页面(带 student token)。
|
||||
- 非精英班学员:显示引导升级提示,不展示工具卡片。
|
||||
- `api.ts` 的 `nitApi` 扩展:查学员积分余额、查 Studio 可用工具。
|
||||
|
||||
## 错误处理
|
||||
|
||||
- 非精英班 / 未审核学员访问 → 403 + 友好提示。
|
||||
- 学员下游 key 缺失(provision 失败)→ 明确错误,提示联系管理员。
|
||||
- 积分余额不足 → 拦截创作,提示充值(复用现有 billing 引导文案范式)。
|
||||
|
||||
## 测试策略(TDD)
|
||||
|
||||
后端(rstest + 内存 SQLite,沿用现有 `recharge_test_pool` 范式):
|
||||
- 学员下游 key ensure 幂等性(缺失则建 / 禁用则启用 / 降级则禁用)。
|
||||
- `resolve_auth` 接受 elite student JWT、拒绝非精英班、拒绝未审核。
|
||||
- 学员创作扣积分、余额不足拦截。
|
||||
|
||||
前端(vitest):
|
||||
- `nitApi` 积分余额 / Studio 工具查询方法(沿用 `api.test.ts` 范式)。
|
||||
- Studio 页卡片渲染 + 余额展示 + 非精英班引导(沿用 `UsageRecords.test.tsx` 范式)。
|
||||
|
||||
## 验证
|
||||
|
||||
- `cargo build -j 2`、`cargo test -p boosapi --lib`(后端 nit / picture_agent / digital_human 相关)。
|
||||
- `npx vitest run`(相关前端测试)、`npm run typecheck`。
|
||||
|
||||
## 提交约定
|
||||
|
||||
全部在 `nit_edu` 分支,不混入 `refactor/rust-backend`。提交按逻辑拆分(迁移 / service /
|
||||
surface 认证 / 前端 分开),沿用项目中文 commit 风格。
|
||||
@@ -1,195 +0,0 @@
|
||||
# NIT 报考学习平台 — 第一期(报考闭环 A+B+C)设计
|
||||
|
||||
日期:2026-06-17
|
||||
分支:`nit_edu`
|
||||
状态:已通过 brainstorming,待落实现计划
|
||||
|
||||
## 背景
|
||||
|
||||
NIT 报考学习平台面向报考「NIT 人工智能与大模型应用」考试的学员,提供注册、报考、
|
||||
学习等功能。完整需求(见 `plan/NIT报考学习平台(初版需求备注解析).docx`)涵盖
|
||||
8 个子系统,体量过大,无法单 spec 完成,故拆分分期建设。
|
||||
|
||||
**本 spec 仅覆盖第一期:报考闭环 A+B+C** —— 学员注册 → 报考 → 缴费审核开课。
|
||||
学习中心(理论课/实操课/练习题)、OCR、校区负责人多级权限、内容保护等留后续期。
|
||||
|
||||
## 已确认的关键决策
|
||||
|
||||
1. 第一个子项目:A+B+C 报考闭环(注册 → 报考 → 缴费审核开课)。
|
||||
2. 独立 `students` 表,与现有 `users`(API 网关用户)完全隔离。
|
||||
3. 缴费:线下转账/扫码给校区 + 学员上传凭证截图 + 后台人工审核开课(非在线支付)。
|
||||
4. 第一期不做 OCR:学员手动填身份证号/姓名,二寸照手动命名上传。
|
||||
5. 开课权限只打状态标记(enrollment.status=approved + tier),学习中心留后续期再读此标记。
|
||||
6. 后台第一期只做平台超管(复用现有 admin 认证),校区负责人登录留后续期。
|
||||
7. 文件存储复用现有 workspace 机制(multipart + `record_bytes_asset` + `workspace_assets` 表)。
|
||||
8. 代码组织:NIT 作为 BoosAPI 内的新模块,复用现有 axum/sqlx/JWT/workspace 基建。
|
||||
9. 学员认证:发 `audience='student'` 的 JWT,与 admin/user token 互不通用。
|
||||
10. 报名信息表:固定核心字段 + 一个 JSON 扩展字段(应对考试中心模版变化),不做动态表单引擎。
|
||||
|
||||
## 非目标(第一期明确不做)
|
||||
|
||||
- 学习中心:理论课/实操课视频、二维码、学习进度、练习题、模拟题、错题集。
|
||||
- 实操课专属 AI 综合应用软件入口。
|
||||
- 身份证 OCR 自动识别与二寸照自动命名。
|
||||
- 校区负责人多级权限与数据隔离。
|
||||
- 内容保护(防截图/录屏/水印)。
|
||||
- 在线支付。
|
||||
|
||||
## 数据模型
|
||||
|
||||
三张新表,SQLite,迁移文件放 `migrations/`。
|
||||
|
||||
### `students` — 学员实名账户
|
||||
|
||||
| 字段 | 类型 | 说明 |
|
||||
|------|------|------|
|
||||
| `id` | INTEGER PK | |
|
||||
| `public_id` | TEXT UNIQUE | 对外 ID,格式 `stu_{uuid_hex}` |
|
||||
| `name` | TEXT NOT NULL | 姓名(实名,注册后不可改) |
|
||||
| `phone` | TEXT NOT NULL UNIQUE | 实名手机号,登录用 |
|
||||
| `id_card_no` | TEXT NOT NULL UNIQUE | 身份证号 |
|
||||
| `campus_id` | INTEGER NOT NULL | FK → `nit_campuses.id` |
|
||||
| `major` | TEXT | 所学专业(社会学员可空) |
|
||||
| `class_name` | TEXT | 班级(可空) |
|
||||
| `head_teacher` | TEXT | 班主任姓名(可空) |
|
||||
| `password_hash` | TEXT NOT NULL | 登录密码哈希(沿用现有 users 哈希方案) |
|
||||
| `privacy_agreed_at` | TEXT NOT NULL | 勾选隐私协议时间 |
|
||||
| `status` | TEXT NOT NULL DEFAULT 'active' | active / disabled |
|
||||
| `created_at` / `updated_at` | TEXT | |
|
||||
|
||||
### `nit_campuses` — 校区/机构库(超管录入,学员注册下拉选)
|
||||
|
||||
| 字段 | 类型 | 说明 |
|
||||
|------|------|------|
|
||||
| `id` | INTEGER PK | |
|
||||
| `public_id` | TEXT UNIQUE | `cmp_{uuid_hex}` |
|
||||
| `province` | TEXT NOT NULL | 省份 |
|
||||
| `name` | TEXT NOT NULL | 院校/机构名称 |
|
||||
| `manager_name` | TEXT | 负责人姓名 |
|
||||
| `manager_phone` | TEXT | 负责人联系方式 |
|
||||
| `enabled` | INTEGER NOT NULL DEFAULT 1 | 仅 enabled 的校区可被学员选择 |
|
||||
| `created_at` / `updated_at` | TEXT | |
|
||||
|
||||
### `nit_enrollments` — 报考记录
|
||||
|
||||
| 字段 | 类型 | 说明 |
|
||||
|------|------|------|
|
||||
| `id` | INTEGER PK | |
|
||||
| `public_id` | TEXT UNIQUE | `enr_{uuid_hex}` |
|
||||
| `student_id` | INTEGER NOT NULL | FK → `students.id` |
|
||||
| `subject` | TEXT NOT NULL | 第一期固定「NIT人工智能与大模型应用」 |
|
||||
| `tier` | TEXT NOT NULL | `basic`(980)/ `elite`(1980) |
|
||||
| `payment_time` | TEXT | 学员填写的缴费时间 |
|
||||
| `payment_proof_asset_id` | TEXT | 缴费凭证截图(workspace asset public_id) |
|
||||
| `id_card_front_asset_id` | TEXT | 身份证正面 |
|
||||
| `id_card_back_asset_id` | TEXT | 身份证反面 |
|
||||
| `photo_asset_id` | TEXT | 二寸蓝底照 |
|
||||
| `applicant_info` | TEXT (JSON) | 报名材料表可变字段 |
|
||||
| `status` | TEXT NOT NULL DEFAULT 'pending_review' | pending_review / approved / rejected |
|
||||
| `reviewed_by` | TEXT | 审核管理员标识 |
|
||||
| `reviewed_at` | TEXT | |
|
||||
| `review_note` | TEXT | 驳回原因/备注 |
|
||||
| `created_at` / `updated_at` | TEXT | |
|
||||
|
||||
**关系**:student → campus(多对一);enrollment → student(多对一)。一个学员允许在被驳回后
|
||||
重新提交报考(可有多条 rejected + 至多一条 active 的 pending_review/approved)。
|
||||
|
||||
**班型权限**:第一期权限即 `approved` 后 enrollment 上的 `tier` 标记;`basic` 对应理论课权限、
|
||||
`elite` 对应理论+实操权限,实际课程内容由后续期读取此标记实现。
|
||||
|
||||
**文件**:四类上传文件全部走现有 `workspace_assets`(`source='nit-enrollment'`,
|
||||
经 `record_bytes_asset` 落本地磁盘),enrollment 仅存其 `public_id` 引用。
|
||||
|
||||
## 认证与路由
|
||||
|
||||
### 学员认证
|
||||
|
||||
复用现有 `auth::jwt` 工具,给 students 发 `audience='student'` 的 JWT。新增
|
||||
`resolve_student(headers, state) -> AppResult<i64>`(返回 student_id),与 admin/user token
|
||||
互不通用。登录凭证为「手机号 + 密码」。
|
||||
|
||||
### 路由分布(挂进现有 axum router via `routes/mod.rs` merge)
|
||||
|
||||
- `src/routes/nit/mod.rs` — 装配 NIT 子路由,导出 `routes()`。
|
||||
- `src/routes/nit/student_auth.rs`
|
||||
- `POST /api/nit/register` — 学员注册
|
||||
- `POST /api/nit/login` — 学员登录
|
||||
- `src/routes/nit/enrollment.rs`(学员侧,需 student JWT)
|
||||
- `POST /api/nit/enrollment/uploads` — multipart 上传单个文件,返回 asset public_id
|
||||
- `POST /api/nit/enrollment` — 创建报考记录
|
||||
- `GET /api/nit/enrollment/mine` — 查本人报考状态
|
||||
- `src/routes/nit/admin.rs`(超管侧,复用现有 admin token 校验)
|
||||
- `GET/POST/PUT /api/admin/nit/campuses` — 校区 CRUD
|
||||
- `GET /api/admin/nit/enrollments` — 报考列表(按 省份/校区/班型/状态 筛选)
|
||||
- `POST /api/admin/nit/enrollments/{public_id}/approve`
|
||||
- `POST /api/admin/nit/enrollments/{public_id}/reject`
|
||||
- `GET /api/admin/nit/enrollments/export` — CSV 导出
|
||||
|
||||
业务逻辑(查询/状态机)抽到 `src/services/nit/`(如 `mod.rs` + `enrollment.rs`),
|
||||
路由层只做参数解析 + 鉴权 + 调用 service,保持文件聚焦。
|
||||
|
||||
## 数据流
|
||||
|
||||
### 注册
|
||||
1. 校验手机号、身份证号未被占用(唯一约束 + 友好报错)。
|
||||
2. 校验 `campus_id` 存在且 enabled。
|
||||
3. 必须 `privacy_agreed=true`。
|
||||
4. 写入 `students`(password 哈希),发 student JWT 返回。
|
||||
|
||||
### 报考(两步,避免半截事务)
|
||||
1. 学员逐个调 `POST /api/nit/enrollment/uploads`(multipart,带 `purpose` 区分四类文件)
|
||||
→ `record_bytes_asset` → 返回各 asset `public_id`。
|
||||
2. 学员调 `POST /api/nit/enrollment`,带各 asset id + tier + payment_time + applicant_info
|
||||
→ 校验本人无在审/已批的活动报考 → 创建 `pending_review` 记录。
|
||||
|
||||
### 审核(超管)
|
||||
- 列表筛选 `pending_review` → 查看凭证/证件图(经 workspace asset file 接口)。
|
||||
- `approve`:状态机 `pending_review → approved`,落 `tier` 权限标记,记 reviewed_by/at。
|
||||
- `reject`:`pending_review → rejected`,记 review_note;学员可据此重新提交。
|
||||
- 状态机只允许从 `pending_review` 流转,杜绝并发重复审核。
|
||||
|
||||
### 导出
|
||||
按 省份/校区/班型/状态 筛选,导出 CSV(实名信息 + 班型),用于核算报考人数/教材数量。
|
||||
|
||||
## 前端页面(`src/web/pages/nit/`)
|
||||
|
||||
- 学员注册页(实名信息 + 校区下拉 + 隐私协议勾选)
|
||||
- 学员登录页(手机号 + 密码)
|
||||
- 报考页(选班型 + 上传四类文件 + 填缴费时间 + 报名信息表 + 提交)
|
||||
- 「我的报考」状态页(pending/approved/rejected + 驳回原因 + 可重新提交)
|
||||
- 超管校区管理页(CRUD)
|
||||
- 超管报考审核页(列表 + 筛选 + 看图 + 通过/驳回 + 导出)
|
||||
|
||||
`api.ts` 新增 `nitApi`(学员侧,带 student token 注入)与 `adminApi` NIT 审核扩展。
|
||||
复用现有 React/Vite/api.ts 范式与 UI 组件。
|
||||
|
||||
## 错误处理
|
||||
|
||||
- 唯一约束冲突 → 友好提示「该手机号/身份证号已注册」。
|
||||
- 上传非图片或超大文件 → 拒绝并提示。
|
||||
- 报考时存在活动报考 → 拒绝重复提交。
|
||||
- 审核状态机:非 `pending_review` 状态的 approve/reject → 409/冲突提示。
|
||||
|
||||
## 测试策略(TDD)
|
||||
|
||||
后端(rstest + 内存 SQLite,沿用现有 `recharge_test_pool` 范式):
|
||||
- 注册手机号/身份证唯一性、校区校验、隐私协议必填。
|
||||
- 登录成功/失败。
|
||||
- 文件上传返回 asset id。
|
||||
- 报考创建 + 重复报考拦截。
|
||||
- 审核状态机(approve/reject/非法流转)。
|
||||
- 导出筛选正确性。
|
||||
|
||||
前端(vitest):
|
||||
- `nitApi` 各方法 URL/method/body 正确(沿用 `api.test.ts` 范式)。
|
||||
- 注册/报考/审核关键页面交互(沿用 `UsageRecords.test.tsx` 的 react-test-renderer 范式)。
|
||||
|
||||
## 验证
|
||||
|
||||
- `cargo build -j 2`(避免 OOM)、`cargo test -p boosapi --lib nit`(后端)。
|
||||
- `npx vitest run`(相关前端测试)、`npm run typecheck`。
|
||||
|
||||
## 提交约定
|
||||
|
||||
NIT 工作全部在 `nit_edu` 分支,不混入 `refactor/rust-backend`。提交按逻辑拆分
|
||||
(迁移 / service / 路由 / 前端 分开),沿用项目中文 commit 风格。
|
||||
@@ -1159,3 +1159,215 @@ pub async fn execute_image_edit(
|
||||
|
||||
Ok(json)
|
||||
}
|
||||
|
||||
// ── Gemini native passthrough ───────────────────────────────────────────
|
||||
//
|
||||
// Gemini's native protocol carries the model in the URL path
|
||||
// (`/v1beta/models/{model}:{op}`) and the API key in the `?key=` query, so the
|
||||
// request body needs no rewrite — we forward it verbatim. This mirrors AQBot's
|
||||
// `handle_native_request`, which deliberately skips `rewrite_body_model` for
|
||||
// Gemini, and BoosAPI's own `/v1/responses` passthrough.
|
||||
|
||||
/// Normalize an upstream base into a Gemini `v1beta` root. Accepts bases that
|
||||
/// already end in `/v1beta` (or contain `/v1beta/`) and leaves them as-is;
|
||||
/// otherwise appends `/v1beta`. Trailing slashes are trimmed.
|
||||
fn gemini_base_url(base_url: &str) -> String {
|
||||
let trimmed = base_url.trim().trim_end_matches('/');
|
||||
if trimmed.ends_with("/v1beta") || trimmed.contains("/v1beta/") {
|
||||
trimmed.to_string()
|
||||
} else {
|
||||
format!("{trimmed}/v1beta")
|
||||
}
|
||||
}
|
||||
|
||||
/// Build `{base}/v1beta/models/{model}:{op}?key={api_key}[&{extra_query}]`.
|
||||
/// `extra_query` carries client query params already filtered (e.g. `alt=sse`),
|
||||
/// without a leading `?`/`&`.
|
||||
fn gemini_operation_url(
|
||||
base_url: &str,
|
||||
model: &str,
|
||||
operation: &str,
|
||||
api_key: &str,
|
||||
extra_query: Option<&str>,
|
||||
) -> String {
|
||||
let base = gemini_base_url(base_url);
|
||||
let model = model.trim().trim_start_matches('/');
|
||||
let model = model.strip_prefix("models/").unwrap_or(model);
|
||||
let mut url = format!("{base}/models/{model}:{operation}");
|
||||
let mut sep = '?';
|
||||
if !api_key.is_empty() {
|
||||
url.push(sep);
|
||||
url.push_str("key=");
|
||||
url.push_str(api_key);
|
||||
sep = '&';
|
||||
}
|
||||
if let Some(query) = extra_query.filter(|q| !q.is_empty()) {
|
||||
url.push(sep);
|
||||
url.push_str(query);
|
||||
}
|
||||
url
|
||||
}
|
||||
|
||||
/// Execute a non-streaming Gemini operation (`generateContent` / `countTokens`)
|
||||
/// and return the upstream JSON verbatim.
|
||||
pub async fn execute_gemini_operation(
|
||||
state: AppState,
|
||||
route: UpstreamRoute,
|
||||
operation: &str,
|
||||
body: Value,
|
||||
extra_query: Option<&str>,
|
||||
) -> AppResult<Value> {
|
||||
let client = state.http.clone();
|
||||
let upstream_url =
|
||||
gemini_operation_url(&route.base_url, &route.model, operation, &route.api_key, extra_query);
|
||||
|
||||
let resp = client
|
||||
.post(&upstream_url)
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| AppError::UpstreamError(format!("Upstream request failed: {e}")))?;
|
||||
|
||||
let status = resp.status();
|
||||
let text = resp
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| AppError::UpstreamError(format!("Failed to read upstream response: {e}")))?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(AppError::UpstreamError(format!(
|
||||
"Upstream returned {status}: {text}"
|
||||
)));
|
||||
}
|
||||
|
||||
parse_upstream_json(&text, &upstream_url, status)
|
||||
}
|
||||
|
||||
/// Execute a streaming Gemini operation (`streamGenerateContent`) and return the
|
||||
/// live upstream response for SSE passthrough. `alt=sse` is enforced in the
|
||||
/// query so the upstream emits `data:`-framed events the accumulator can parse.
|
||||
pub async fn execute_gemini_stream_response(
|
||||
state: AppState,
|
||||
route: UpstreamRoute,
|
||||
operation: &str,
|
||||
body: Value,
|
||||
extra_query: Option<&str>,
|
||||
) -> AppResult<reqwest::Response> {
|
||||
let client = state.http.clone();
|
||||
let query = merge_alt_sse(extra_query);
|
||||
let upstream_url = gemini_operation_url(
|
||||
&route.base_url,
|
||||
&route.model,
|
||||
operation,
|
||||
&route.api_key,
|
||||
Some(&query),
|
||||
);
|
||||
|
||||
let resp = client
|
||||
.post(&upstream_url)
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| AppError::UpstreamError(format!("Upstream request failed: {e}")))?;
|
||||
|
||||
let status = resp.status();
|
||||
if !status.is_success() {
|
||||
let text = resp
|
||||
.text()
|
||||
.await
|
||||
.unwrap_or_else(|_| "unknown error".to_string());
|
||||
return Err(AppError::UpstreamError(format!(
|
||||
"Upstream returned {status}: {text}"
|
||||
)));
|
||||
}
|
||||
|
||||
Ok(resp)
|
||||
}
|
||||
|
||||
/// Ensure `alt=sse` is present in the extra query (Gemini streams plain JSON
|
||||
/// arrays without it). Preserves any other client-supplied params.
|
||||
fn merge_alt_sse(extra_query: Option<&str>) -> String {
|
||||
match extra_query.filter(|q| !q.is_empty()) {
|
||||
Some(q) if q.split('&').any(|p| p == "alt=sse" || p.starts_with("alt=")) => q.to_string(),
|
||||
Some(q) => format!("{q}&alt=sse"),
|
||||
None => "alt=sse".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod gemini_tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn base_url_appends_v1beta_when_missing() {
|
||||
assert_eq!(
|
||||
gemini_base_url("https://generativelanguage.googleapis.com"),
|
||||
"https://generativelanguage.googleapis.com/v1beta"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn base_url_preserves_existing_v1beta() {
|
||||
assert_eq!(
|
||||
gemini_base_url("https://generativelanguage.googleapis.com/v1beta/"),
|
||||
"https://generativelanguage.googleapis.com/v1beta"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn operation_url_embeds_model_op_and_key() {
|
||||
let url = gemini_operation_url(
|
||||
"https://generativelanguage.googleapis.com",
|
||||
"gemini-2.0-flash",
|
||||
"generateContent",
|
||||
"secret",
|
||||
None,
|
||||
);
|
||||
assert_eq!(
|
||||
url,
|
||||
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent?key=secret"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn operation_url_strips_models_prefix_and_appends_query() {
|
||||
let url = gemini_operation_url(
|
||||
"https://generativelanguage.googleapis.com/v1beta",
|
||||
"models/gemini-1.5-pro",
|
||||
"streamGenerateContent",
|
||||
"k",
|
||||
Some("alt=sse"),
|
||||
);
|
||||
assert_eq!(
|
||||
url,
|
||||
"https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-pro:streamGenerateContent?key=k&alt=sse"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn operation_url_without_key_uses_question_mark_for_query() {
|
||||
let url = gemini_operation_url(
|
||||
"https://example.test/v1beta",
|
||||
"gemini-x",
|
||||
"countTokens",
|
||||
"",
|
||||
Some("foo=bar"),
|
||||
);
|
||||
assert_eq!(
|
||||
url,
|
||||
"https://example.test/v1beta/models/gemini-x:countTokens?foo=bar"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_alt_sse_variants() {
|
||||
assert_eq!(merge_alt_sse(None), "alt=sse");
|
||||
assert_eq!(merge_alt_sse(Some("")), "alt=sse");
|
||||
assert_eq!(merge_alt_sse(Some("foo=bar")), "foo=bar&alt=sse");
|
||||
// already present → unchanged
|
||||
assert_eq!(merge_alt_sse(Some("alt=sse")), "alt=sse");
|
||||
assert_eq!(merge_alt_sse(Some("alt=json")), "alt=json");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2840,6 +2840,8 @@ mod tests {
|
||||
template_proxy_url: None,
|
||||
},
|
||||
comfyui: Default::default(),
|
||||
payment: Default::default(),
|
||||
pixelle_video: Default::default(),
|
||||
},
|
||||
pool,
|
||||
BuiltinKeys {
|
||||
|
||||
@@ -412,6 +412,8 @@ mod tests {
|
||||
template_proxy_url: None,
|
||||
},
|
||||
comfyui: Default::default(),
|
||||
payment: Default::default(),
|
||||
pixelle_video: Default::default(),
|
||||
},
|
||||
pool,
|
||||
BuiltinKeys {
|
||||
|
||||
@@ -19,6 +19,11 @@ use serde_json::{json, Value};
|
||||
use std::net::{IpAddr, SocketAddr};
|
||||
use tracing::{info, warn};
|
||||
|
||||
/// Default spending cap (USD) granted to a newly created user's auto-generated
|
||||
/// API key. Per the billing design (新用户默认 $1 额度), new users start with a
|
||||
/// $1 allowance on their default downstream key.
|
||||
pub const DEFAULT_NEW_USER_MAX_COST_USD: f64 = 1.0;
|
||||
|
||||
// ── Request / Response types ─────────────────────────────────────────
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
@@ -1019,6 +1024,7 @@ mod tests {
|
||||
},
|
||||
comfyui: Default::default(),
|
||||
payment: Default::default(),
|
||||
pixelle_video: Default::default(),
|
||||
},
|
||||
pool,
|
||||
BuiltinKeys {
|
||||
|
||||
@@ -127,6 +127,8 @@ mod tests {
|
||||
template_proxy_url: None,
|
||||
},
|
||||
comfyui: ComfyUIConfig::default(),
|
||||
payment: Default::default(),
|
||||
pixelle_video: Default::default(),
|
||||
},
|
||||
pool,
|
||||
BuiltinKeys {
|
||||
|
||||
@@ -1305,6 +1305,8 @@ mod tests {
|
||||
template_proxy_url: None,
|
||||
},
|
||||
comfyui: Default::default(),
|
||||
payment: Default::default(),
|
||||
pixelle_video: Default::default(),
|
||||
},
|
||||
pool,
|
||||
BuiltinKeys {
|
||||
|
||||
@@ -536,6 +536,8 @@ mod tests {
|
||||
template_proxy_url: None,
|
||||
},
|
||||
comfyui: Default::default(),
|
||||
payment: Default::default(),
|
||||
pixelle_video: Default::default(),
|
||||
},
|
||||
pool,
|
||||
BuiltinKeys {
|
||||
|
||||
@@ -2071,6 +2071,8 @@ mod tests {
|
||||
template_proxy_url: None,
|
||||
},
|
||||
comfyui,
|
||||
payment: Default::default(),
|
||||
pixelle_video: Default::default(),
|
||||
},
|
||||
pool,
|
||||
BuiltinKeys {
|
||||
|
||||
@@ -2550,6 +2550,8 @@ mod tests {
|
||||
template_proxy_url: None,
|
||||
},
|
||||
comfyui: ComfyUIConfig::default(),
|
||||
payment: Default::default(),
|
||||
pixelle_video: Default::default(),
|
||||
},
|
||||
pool,
|
||||
BuiltinKeys {
|
||||
|
||||
@@ -73,6 +73,16 @@ pub fn routes() -> Router<AppState> {
|
||||
.route("/v1/videos", post(videos_handler))
|
||||
.route("/v1/videos/{id}", get(video_task_handler))
|
||||
.route("/v1/models", get(list_models))
|
||||
// ── Gemini native inbound (v1beta) ────────────────────────────
|
||||
// Lets Gemini CLI / Google GenAI SDKs point their base URL straight
|
||||
// at this gateway. The model lives in the path and the key in the
|
||||
// `?key=` query (or x-goog-api-key header), so requests pass through
|
||||
// to a Gemini-shaped upstream without body rewriting.
|
||||
.route("/v1beta/models", get(gemini_list_models))
|
||||
.route(
|
||||
"/v1beta/models/{model_action}",
|
||||
post(gemini_model_operation),
|
||||
)
|
||||
// ── ComfyUI shim aliases ──────────────────────────────────────
|
||||
// ComfyUI's OpenAI nodes hit /proxy/openai/... paths (see
|
||||
// comfy_api_nodes/nodes_openai.py). Mounting them here as aliases
|
||||
@@ -157,6 +167,209 @@ async fn responses_handler(
|
||||
}
|
||||
}
|
||||
|
||||
// ── Gemini native inbound handlers (v1beta) ────────────────────────────
|
||||
//
|
||||
// Gemini clients (Gemini CLI / Google GenAI SDK) carry the model in the URL
|
||||
// path (`/v1beta/models/{model}:{op}`) and the gateway key in `?key=` or the
|
||||
// `x-goog-api-key` header. We authenticate against our downstream keys, route
|
||||
// the Gemini model through the same ProxyConductor as the OpenAI endpoints, and
|
||||
// forward the body verbatim to a Gemini-shaped upstream (no body rewrite).
|
||||
|
||||
/// Gemini operations carried in the `{model}:{op}` path segment.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum GeminiOperation {
|
||||
GenerateContent,
|
||||
StreamGenerateContent,
|
||||
CountTokens,
|
||||
}
|
||||
|
||||
impl GeminiOperation {
|
||||
fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
GeminiOperation::GenerateContent => "generateContent",
|
||||
GeminiOperation::StreamGenerateContent => "streamGenerateContent",
|
||||
GeminiOperation::CountTokens => "countTokens",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse the `{model}:{operation}` path segment (e.g.
|
||||
/// `gemini-2.5-pro:streamGenerateContent` or `models/gemini-x:countTokens`).
|
||||
/// Splits on the last `:` so model names that contain none are rejected.
|
||||
fn parse_gemini_model_action(model_action: &str) -> Option<(String, GeminiOperation)> {
|
||||
let (model, action) = model_action.rsplit_once(':')?;
|
||||
let operation = match action {
|
||||
"generateContent" => GeminiOperation::GenerateContent,
|
||||
"streamGenerateContent" => GeminiOperation::StreamGenerateContent,
|
||||
"countTokens" => GeminiOperation::CountTokens,
|
||||
_ => return None,
|
||||
};
|
||||
if model.trim().is_empty() {
|
||||
return None;
|
||||
}
|
||||
Some((model.to_string(), operation))
|
||||
}
|
||||
|
||||
/// Authenticate a Gemini-native request: the key arrives in `?key=` or the
|
||||
/// `x-goog-api-key` header (Gemini convention), not `Authorization: Bearer`.
|
||||
/// Reuses the same downstream-key authorization + limits as `check_auth`.
|
||||
async fn check_gemini_auth(
|
||||
state: &AppState,
|
||||
headers: &HeaderMap,
|
||||
query: &std::collections::HashMap<String, String>,
|
||||
) -> Result<Option<i64>, AppError> {
|
||||
let key = query
|
||||
.get("key")
|
||||
.map(|s| s.as_str())
|
||||
.or_else(|| headers.get("x-goog-api-key").and_then(|v| v.to_str().ok()))
|
||||
.map(str::trim)
|
||||
.filter(|s| !s.is_empty())
|
||||
.ok_or_else(|| {
|
||||
AppError::Unauthorized("Missing API key (use ?key= or x-goog-api-key)".into())
|
||||
})?;
|
||||
|
||||
let proxy_token = state.runtime_proxy_token();
|
||||
let result = authorize_downstream_token(key, &state.db, Some(&proxy_token))
|
||||
.await
|
||||
.map_err(AppError::Unauthorized)?;
|
||||
|
||||
match result {
|
||||
DownstreamAuthResult::Managed { key } => {
|
||||
check_downstream_key_limits(&key)?;
|
||||
Ok(Some(key.id))
|
||||
}
|
||||
DownstreamAuthResult::Global { .. } => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
/// Rebuild the upstream query string from client params, dropping the inbound
|
||||
/// gateway `key` (the upstream key is injected by the executor) and `alt` (the
|
||||
/// streaming path enforces its own `alt=sse`). Returns `None` when nothing
|
||||
/// remains. Values are forwarded as-is: real Gemini clients only put `key`/`alt`
|
||||
/// in the query and carry everything else (generationConfig, etc.) in the body.
|
||||
fn gemini_forward_query(
|
||||
query: &std::collections::HashMap<String, String>,
|
||||
drop_alt: bool,
|
||||
) -> Option<String> {
|
||||
let mut parts: Vec<(&String, &String)> = query
|
||||
.iter()
|
||||
.filter(|(k, _)| k.as_str() != "key" && !(drop_alt && k.as_str() == "alt"))
|
||||
.collect();
|
||||
if parts.is_empty() {
|
||||
return None;
|
||||
}
|
||||
// Stable order so the URL is deterministic (and testable).
|
||||
parts.sort_by(|a, b| a.0.cmp(b.0));
|
||||
let encoded = parts
|
||||
.iter()
|
||||
.map(|(k, v)| format!("{k}={v}"))
|
||||
.collect::<Vec<_>>()
|
||||
.join("&");
|
||||
Some(encoded)
|
||||
}
|
||||
|
||||
/// GET /v1beta/models — list models in Gemini's native shape.
|
||||
/// Like `/v1/models`, this is a directory listing from configured routes; no
|
||||
/// upstream call, no billing.
|
||||
async fn gemini_list_models(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Query(query): Query<std::collections::HashMap<String, String>>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
info!("GET /v1beta/models");
|
||||
|
||||
let _ = check_gemini_auth(&state, &headers, &query).await?;
|
||||
|
||||
let rows: Vec<(String, String)> = sqlx::query_as(
|
||||
"SELECT DISTINCT
|
||||
tr.model_pattern,
|
||||
COALESCE(NULLIF(TRIM(tr.model_mapping), ''), tr.model_pattern)
|
||||
FROM token_routes tr
|
||||
JOIN route_channels rc ON rc.route_id = tr.id
|
||||
JOIN accounts a ON a.id = rc.account_id
|
||||
JOIN sites s ON s.id = a.site_id
|
||||
LEFT JOIN account_tokens at ON at.id = rc.token_id AND at.enabled = 1
|
||||
WHERE tr.enabled = 1
|
||||
AND rc.enabled = 1
|
||||
AND a.status = 'active'
|
||||
AND s.status = 'active'
|
||||
AND TRIM(COALESCE(s.url, '')) <> ''
|
||||
AND (rc.cooldown_until IS NULL OR rc.cooldown_until <= datetime('now'))
|
||||
AND (rc.token_id IS NULL OR at.id IS NOT NULL)
|
||||
ORDER BY tr.id",
|
||||
)
|
||||
.fetch_all(&*state.db)
|
||||
.await?;
|
||||
|
||||
let models: Vec<Value> = rows
|
||||
.into_iter()
|
||||
.map(|(_pattern, model)| {
|
||||
let name = if model.starts_with("models/") {
|
||||
model.clone()
|
||||
} else {
|
||||
format!("models/{model}")
|
||||
};
|
||||
json!({
|
||||
"name": name,
|
||||
"supportedGenerationMethods": [
|
||||
"generateContent",
|
||||
"streamGenerateContent",
|
||||
"countTokens",
|
||||
],
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(Json(json!({ "models": models })))
|
||||
}
|
||||
|
||||
/// POST /v1beta/models/{model}:{op} — dispatch to the matching Gemini op.
|
||||
async fn gemini_model_operation(
|
||||
State(state): State<AppState>,
|
||||
axum::extract::Path(model_action): axum::extract::Path<String>,
|
||||
Query(query): Query<std::collections::HashMap<String, String>>,
|
||||
headers: HeaderMap,
|
||||
Json(body): Json<Value>,
|
||||
) -> Result<Response, AppError> {
|
||||
let Some((model, operation)) = parse_gemini_model_action(&model_action) else {
|
||||
return Err(AppError::BadRequest(format!(
|
||||
"Unsupported Gemini model action '{model_action}' \
|
||||
(expected {{model}}:generateContent|streamGenerateContent|countTokens)"
|
||||
)));
|
||||
};
|
||||
info!("POST /v1beta/models/{model}:{}", operation.as_str());
|
||||
|
||||
let key_id = check_gemini_auth(&state, &headers, &query).await?;
|
||||
|
||||
match operation {
|
||||
GeminiOperation::GenerateContent => {
|
||||
let extra_query = gemini_forward_query(&query, false);
|
||||
let result =
|
||||
execute_gemini_with_conductor(state, key_id, model, body, extra_query).await?;
|
||||
Ok(Json(result).into_response())
|
||||
}
|
||||
GeminiOperation::StreamGenerateContent => {
|
||||
let extra_query = gemini_forward_query(&query, true);
|
||||
execute_gemini_stream_with_conductor(state, key_id, model, body, extra_query).await
|
||||
}
|
||||
GeminiOperation::CountTokens => {
|
||||
// countTokens is a metadata call (no generation) — route for a
|
||||
// channel but skip billing, mirroring the lighter helper endpoints.
|
||||
let extra_query = gemini_forward_query(&query, false);
|
||||
let route = resolve_route(&state, &model).await?;
|
||||
let result = executors::execute_gemini_operation(
|
||||
state,
|
||||
route,
|
||||
"countTokens",
|
||||
body,
|
||||
extra_query.as_deref(),
|
||||
)
|
||||
.await?;
|
||||
Ok(Json(result).into_response())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// POST /v1/chat/completions
|
||||
async fn chat_completions(
|
||||
State(state): State<AppState>,
|
||||
@@ -1334,6 +1547,38 @@ fn responses_stream_usage_from_line(line: &str) -> Option<(i64, i64)> {
|
||||
(prompt != 0 || completion != 0).then_some((prompt, completion))
|
||||
}
|
||||
|
||||
/// Extract `(prompt, completion)` tokens from a Gemini response body's
|
||||
/// `usageMetadata` (`promptTokenCount` / `candidatesTokenCount`). Returns
|
||||
/// `(0, 0)` when the field is absent — mirrors `extract_usage` for OpenAI.
|
||||
fn gemini_extract_usage(response: &Value) -> (i64, i64) {
|
||||
let usage = response.get("usageMetadata");
|
||||
let prompt = usage
|
||||
.and_then(|u| u.get("promptTokenCount"))
|
||||
.and_then(|v| v.as_i64())
|
||||
.unwrap_or(0);
|
||||
let completion = usage
|
||||
.and_then(|u| u.get("candidatesTokenCount"))
|
||||
.and_then(|v| v.as_i64())
|
||||
.unwrap_or(0);
|
||||
(prompt, completion)
|
||||
}
|
||||
|
||||
/// Per-line usage parser for Gemini `streamGenerateContent?alt=sse` output.
|
||||
/// Gemini emits `data:`-framed JSON chunks; the final chunk carries the
|
||||
/// cumulative `usageMetadata`. Feeds the shared `StreamUsageAccumulator`.
|
||||
fn gemini_stream_usage_from_line(line: &str) -> Option<(i64, i64)> {
|
||||
let payload = line
|
||||
.strip_prefix("data: ")
|
||||
.or_else(|| line.strip_prefix("data:"))?
|
||||
.trim();
|
||||
if payload.is_empty() || payload == "[DONE]" {
|
||||
return None;
|
||||
}
|
||||
let value: Value = serde_json::from_str(payload).ok()?;
|
||||
let (prompt, completion) = gemini_extract_usage(&value);
|
||||
(prompt != 0 || completion != 0).then_some((prompt, completion))
|
||||
}
|
||||
|
||||
fn build_conductor(state: &AppState) -> ProxyConductor {
|
||||
let db = state.db.clone();
|
||||
let select_db = db.clone();
|
||||
@@ -1881,6 +2126,185 @@ async fn execute_responses_stream_with_conductor(
|
||||
}
|
||||
}
|
||||
|
||||
/// Non-streaming Gemini op (`generateContent`) through the ProxyConductor.
|
||||
/// Mirrors `execute_responses_with_conductor` but builds a Gemini upstream URL
|
||||
/// (model in path, key in query) and forwards the body verbatim — Gemini bodies
|
||||
/// carry no `model` field, so there is no body rewrite.
|
||||
async fn execute_gemini_with_conductor(
|
||||
state: AppState,
|
||||
key_id: Option<i64>,
|
||||
model: String,
|
||||
body: Value,
|
||||
extra_query: Option<String>,
|
||||
) -> AppResult<Value> {
|
||||
let started_at = Instant::now();
|
||||
let request_log_body = body.clone();
|
||||
let conductor = build_conductor(&state);
|
||||
let billing_state = state.clone();
|
||||
let billing_model = model.clone();
|
||||
let result = conductor
|
||||
.execute(ExecuteInput {
|
||||
requested_model: model,
|
||||
downstream_policy: None,
|
||||
attempt_fn: Box::new(move |ctx| {
|
||||
let state = state.clone();
|
||||
let body = body.clone();
|
||||
let extra_query = extra_query.clone();
|
||||
async move {
|
||||
match executors::execute_gemini_operation(
|
||||
state,
|
||||
selected_channel_to_route(ctx.selected),
|
||||
"generateContent",
|
||||
body,
|
||||
extra_query.as_deref(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(response) => AttemptResult::Success {
|
||||
response,
|
||||
latency_ms: None,
|
||||
cost: None,
|
||||
},
|
||||
Err(err) => AttemptResult::Failure {
|
||||
action: AttemptFailureAction::Failover,
|
||||
status: upstream_status_from_error(&err),
|
||||
raw_error_text: Some(err.to_string()),
|
||||
},
|
||||
}
|
||||
}
|
||||
.boxed()
|
||||
}),
|
||||
on_terminal_failure: None,
|
||||
})
|
||||
.await;
|
||||
|
||||
let (selected, response, attempts) = match result {
|
||||
ExecuteResult::Success {
|
||||
selected,
|
||||
response,
|
||||
attempts,
|
||||
} => (selected, response, attempts),
|
||||
ExecuteResult::Failure {
|
||||
reason,
|
||||
selected,
|
||||
status,
|
||||
raw_error_text,
|
||||
attempts,
|
||||
} => {
|
||||
record_proxy_log_failure(
|
||||
&billing_state,
|
||||
key_id,
|
||||
selected.as_ref(),
|
||||
&billing_model,
|
||||
status,
|
||||
Some(elapsed_ms_since(started_at)),
|
||||
raw_error_text.as_deref(),
|
||||
attempts,
|
||||
false,
|
||||
proxy_log_content_from_request(&request_log_body),
|
||||
)
|
||||
.await;
|
||||
return Err(AppError::UpstreamError(format!(
|
||||
"Proxy conductor failed after {attempts} attempts ({reason:?}, status: {:?}): {}",
|
||||
status,
|
||||
raw_error_text.unwrap_or_else(|| "no upstream error text".to_string())
|
||||
)));
|
||||
}
|
||||
};
|
||||
let (prompt, completion) = gemini_extract_usage(&response);
|
||||
let cost = estimate_request_cost(&selected, prompt, completion);
|
||||
record_proxy_log_success(
|
||||
&billing_state,
|
||||
key_id,
|
||||
&selected,
|
||||
&billing_model,
|
||||
Some(elapsed_ms_since(started_at)),
|
||||
prompt,
|
||||
completion,
|
||||
cost,
|
||||
attempts,
|
||||
false,
|
||||
proxy_log_content_from_request_response(&request_log_body, Some(&response)),
|
||||
)
|
||||
.await;
|
||||
record_billing(
|
||||
&billing_state,
|
||||
key_id,
|
||||
&billing_model,
|
||||
prompt,
|
||||
completion,
|
||||
cost,
|
||||
)
|
||||
.await;
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
/// Streaming Gemini op (`streamGenerateContent`) through the ProxyConductor.
|
||||
/// Mirrors `execute_responses_stream_with_conductor`, opening the upstream with
|
||||
/// failover and SSE-passing it through with the Gemini usage parser so streamed
|
||||
/// `usageMetadata` is billed.
|
||||
async fn execute_gemini_stream_with_conductor(
|
||||
state: AppState,
|
||||
key_id: Option<i64>,
|
||||
model: String,
|
||||
body: Value,
|
||||
extra_query: Option<String>,
|
||||
) -> Result<Response, AppError> {
|
||||
let started_at = Instant::now();
|
||||
let request_log_body = body.clone();
|
||||
let body_for_attempt = body.clone();
|
||||
let opened = open_stream_with_failover(&state, &model, move |state, selected| {
|
||||
let body = body_for_attempt.clone();
|
||||
let extra_query = extra_query.clone();
|
||||
async move {
|
||||
executors::execute_gemini_stream_response(
|
||||
state,
|
||||
selected_channel_to_route(selected),
|
||||
"streamGenerateContent",
|
||||
body,
|
||||
extra_query.as_deref(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
})
|
||||
.await;
|
||||
|
||||
match opened {
|
||||
Ok(opened) => sse_passthrough_response(
|
||||
state,
|
||||
opened,
|
||||
key_id,
|
||||
model,
|
||||
started_at,
|
||||
proxy_log_content_from_request(&request_log_body),
|
||||
gemini_stream_usage_from_line,
|
||||
),
|
||||
Err(failure) => {
|
||||
record_proxy_log_failure(
|
||||
&state,
|
||||
key_id,
|
||||
failure.selected.as_ref(),
|
||||
&model,
|
||||
failure.status,
|
||||
Some(elapsed_ms_since(started_at)),
|
||||
failure.raw_error_text.as_deref(),
|
||||
failure.attempts,
|
||||
true,
|
||||
proxy_log_content_from_request(&request_log_body),
|
||||
)
|
||||
.await;
|
||||
Err(AppError::UpstreamError(format!(
|
||||
"Proxy conductor failed after {} attempts (stream open, status: {:?}): {}",
|
||||
failure.attempts,
|
||||
failure.status,
|
||||
failure
|
||||
.raw_error_text
|
||||
.unwrap_or_else(|| "no upstream error text".to_string())
|
||||
)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -2121,6 +2545,8 @@ mod tests {
|
||||
template_proxy_url: None,
|
||||
},
|
||||
comfyui: Default::default(),
|
||||
payment: Default::default(),
|
||||
pixelle_video: Default::default(),
|
||||
},
|
||||
pool,
|
||||
crate::services::builtin_keys::BuiltinKeys {
|
||||
@@ -2155,4 +2581,87 @@ mod tests {
|
||||
.unwrap();
|
||||
assert!(fail_count > 0);
|
||||
}
|
||||
|
||||
// ── Gemini native inbound unit tests ───────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn parse_gemini_model_action_known_ops() {
|
||||
let (m, op) = parse_gemini_model_action("gemini-2.5-pro:generateContent").unwrap();
|
||||
assert_eq!(m, "gemini-2.5-pro");
|
||||
assert_eq!(op, GeminiOperation::GenerateContent);
|
||||
|
||||
let (m, op) = parse_gemini_model_action("gemini-2.5-pro:streamGenerateContent").unwrap();
|
||||
assert_eq!(m, "gemini-2.5-pro");
|
||||
assert_eq!(op, GeminiOperation::StreamGenerateContent);
|
||||
|
||||
// model names may carry the `models/` prefix
|
||||
let (m, op) = parse_gemini_model_action("models/gemini-1.5-pro:countTokens").unwrap();
|
||||
assert_eq!(m, "models/gemini-1.5-pro");
|
||||
assert_eq!(op, GeminiOperation::CountTokens);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_gemini_model_action_rejects_bad_input() {
|
||||
// no operation suffix
|
||||
assert!(parse_gemini_model_action("gemini-2.5-pro").is_none());
|
||||
// unknown operation
|
||||
assert!(parse_gemini_model_action("gemini-2.5-pro:embedContent").is_none());
|
||||
// empty model
|
||||
assert!(parse_gemini_model_action(":generateContent").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemini_extract_usage_reads_usage_metadata() {
|
||||
let body = json!({
|
||||
"candidates": [],
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 11,
|
||||
"candidatesTokenCount": 22,
|
||||
"totalTokenCount": 33
|
||||
}
|
||||
});
|
||||
assert_eq!(gemini_extract_usage(&body), (11, 22));
|
||||
// absent usageMetadata → zeros
|
||||
assert_eq!(gemini_extract_usage(&json!({})), (0, 0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemini_stream_usage_from_line_parses_data_frame() {
|
||||
let line = r#"data: {"candidates":[{"content":{"parts":[]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":4,"candidatesTokenCount":6,"totalTokenCount":10}}"#;
|
||||
assert_eq!(gemini_stream_usage_from_line(line), Some((4, 6)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemini_stream_usage_from_line_ignores_non_usage_and_done() {
|
||||
// a chunk with no usageMetadata yields None (don't overwrite accumulator)
|
||||
let no_usage = r#"data: {"candidates":[{"content":{"parts":[{"text":"hi"}]}}]}"#;
|
||||
assert_eq!(gemini_stream_usage_from_line(no_usage), None);
|
||||
// sentinels / blanks / non-data lines
|
||||
assert_eq!(gemini_stream_usage_from_line("data: [DONE]"), None);
|
||||
assert_eq!(gemini_stream_usage_from_line("data:"), None);
|
||||
assert_eq!(gemini_stream_usage_from_line(": comment"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemini_forward_query_drops_key_and_alt() {
|
||||
let mut q = std::collections::HashMap::new();
|
||||
q.insert("key".to_string(), "secret".to_string());
|
||||
q.insert("alt".to_string(), "sse".to_string());
|
||||
// stream path drops both key and alt → nothing left
|
||||
assert_eq!(gemini_forward_query(&q, true), None);
|
||||
// non-stream keeps alt but always drops key
|
||||
assert_eq!(gemini_forward_query(&q, false), Some("alt=sse".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemini_forward_query_preserves_other_params_sorted() {
|
||||
let mut q = std::collections::HashMap::new();
|
||||
q.insert("key".to_string(), "secret".to_string());
|
||||
q.insert("pageSize".to_string(), "50".to_string());
|
||||
q.insert("filter".to_string(), "x".to_string());
|
||||
assert_eq!(
|
||||
gemini_forward_query(&q, true),
|
||||
Some("filter=x&pageSize=50".to_string())
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -696,6 +696,8 @@ mod tests {
|
||||
template_proxy_url: None,
|
||||
},
|
||||
comfyui: ComfyUIConfig::default(),
|
||||
payment: Default::default(),
|
||||
pixelle_video: Default::default(),
|
||||
},
|
||||
pool,
|
||||
BuiltinKeys {
|
||||
|
||||
Reference in New Issue
Block a user