Skip to main content

sokuto_server/handlers/
state_batch.rs

1//! # 一括バッチ質問エンドポイントモジュール
2//!
3//! 単一の長文 `State` に対して複数の異種質問 (`Choice`/`Score`/`Noul`) を一括投入し、
4//! `State` エンコードを 1 度きりに抑えてミリ秒級で並列判定結果を返却する。
5
6use std::sync::Arc;
7use std::time::Instant;
8
9use axum::extract::State;
10use axum::http::HeaderMap;
11use axum::response::IntoResponse;
12use sokuto_core::schema::{StateBatchRequest, StateBatchResponse, SystemOneRequest};
13
14use crate::error::{ErrorResponse, ServerError};
15use crate::metrics::{record_http_duration, record_http_request};
16use crate::service::InferenceService;
17use crate::state::AppState;
18
19/// 単一 `State` に対する一括バッチ質問推論エンドポイント。
20///
21/// 共通 `State` を 1 回のみエンコード (またはキャッシュ取得) し、
22/// 全質問をプレフィックス共有バッチ化して一括評価する。
23#[utoipa::path(
24    post,
25    path = "/v1/systemone/state-batch",
26    tag = "Inference",
27    request_body = StateBatchRequest,
28    responses(
29        (status = 200, description = "バッチ推論および決定解決成功", body = StateBatchResponse),
30        (status = 400, description = "スキーマまたはパラメータ検証違反", body = ErrorResponse),
31        (status = 413, description = "質問数またはペイロード上限超過", body = ErrorResponse),
32        (status = 500, description = "推論エンジンまたはサーバー内部障害", body = ErrorResponse)
33    )
34)]
35pub async fn state_batch_handler(
36    State(state): State<Arc<AppState>>,
37    headers: HeaderMap,
38    axum::Json(req): axum::Json<StateBatchRequest>,
39) -> Result<impl IntoResponse, ServerError> {
40    let start_time = Instant::now();
41
42    req.validate()?;
43
44    let client_hash_str = headers
45        .get("x-state-hash")
46        .or_else(|| headers.get("X-State-Hash"))
47        .and_then(|v| v.to_str().ok())
48        .map(str::to_string)
49        .or_else(|| req.state_hash.clone());
50
51    let s1_req = SystemOneRequest {
52        model: req.model,
53        state: req.state,
54        questions: req.questions,
55        gating: req.gating,
56        auto_hierarchical: req.auto_hierarchical,
57        state_hash: req.state_hash,
58    };
59
60    let inference_service = InferenceService::new(state);
61    let output = inference_service.execute(s1_req, client_hash_str).await?;
62
63    let resp_headers = output.to_headers();
64    let hex_key_opt = output.state_key.map(|k| format!("{:016x}", k));
65
66    let mut response = StateBatchResponse::new(output.answers, output.usage, output.cache_hit);
67    if let Some(summary) = output.routing_summary {
68        response = response.with_routing(summary);
69    }
70    if let Some(ref hex_key) = hex_key_opt {
71        response = response.with_state_hash(hex_key);
72    }
73
74    let total_duration = start_time.elapsed().as_secs_f64();
75    record_http_duration("/v1/systemone/state-batch", total_duration);
76    record_http_request("/v1/systemone/state-batch", 200);
77
78    Ok((resp_headers, axum::Json(response)))
79}