sokuto_server/
state.rs

1//! # アプリケーション共有状態管理モジュール
2//!
3//! 推論エンジン、トークナイザー、較正設定、およびサーバー運用パラメータをスレッドセーフに保持する。
4
5use std::sync::Arc;
6
7use sokuto_core::contract::calibration::CalibrationConfig;
8use sokuto_core::gating::GatingConfig;
9use sokuto_runtime::dag::InProcessDagExecutor;
10use sokuto_runtime::engine::{
11    CoarseToFineConfig, DualSessionEngine, InferenceEngine, StateCacheManager,
12};
13use sokuto_runtime::hierarchical::HierarchicalRouterConfig;
14use sokuto_runtime::hierarchical::mapping::HierarchicalMapping;
15use sokuto_runtime::tokenizer::JevTokenizer;
16
17use crate::guardrails::{GuardrailConfig, GuardrailPipeline};
18
19/// デフォルトのリクエストあたり最大質問数。
20pub const DEFAULT_MAX_QUESTIONS_PER_REQUEST: usize = 128;
21
22/// デフォルトのバッチチャンクサイズ。
23pub const DEFAULT_CHUNK_SIZE: usize = 16;
24
25/// サーバー全体で共有されるアプリケーション状態。
26#[derive(Clone)]
27pub struct AppState {
28    /// ONNX Runtime 推論エンジン(内部でセッションプールと同期制御を保持)。
29    pub engine: Arc<InferenceEngine>,
30    /// Jev 高速トークナイザー。
31    pub tokenizer: Arc<JevTokenizer>,
32    /// 事後較正温度設定。
33    pub calib_config: Arc<CalibrationConfig>,
34    /// 大規模候補数向けの粗密 2 段階探索設定。
35    pub coarse_config: CoarseToFineConfig,
36    /// 1 リクエストで許容される最大質問数。
37    pub max_questions_per_request: usize,
38    /// マイクロバッチ分割時のチャンクサイズ。
39    pub chunk_size: usize,
40    /// 前処理ガードレールパイプライン。
41    pub guardrail_pipeline: GuardrailPipeline,
42    /// 確信度ゲーティング処理設定。
43    pub gating_config: GatingConfig,
44    /// インプロセス DAG 実行器(反事実・マイクロ決定グラフ向け)。
45    pub dag_executor: Arc<InProcessDagExecutor>,
46    /// 大分類・細分類オントロジー定義(存在する場合)。
47    pub hierarchical_mapping: Option<Arc<HierarchicalMapping>>,
48    /// 粗密二段階階層ルーティング設定。
49    pub hierarchical_config: HierarchicalRouterConfig,
50    /// 透過的粗密ルーティングモードの有効フラグ。
51    pub auto_hierarchical: bool,
52    /// 2 セッション分割型長文推論エンジン (存在する場合)。
53    pub dual_engine: Option<Arc<DualSessionEngine>>,
54    /// State 中間特徴量キャッシュマネージャー (存在する場合)。
55    pub cache_manager: Option<Arc<StateCacheManager>>,
56}
57
58impl AppState {
59    /// 新規 `AppState` を構築する。
60    ///
61    /// # 引数
62    /// - `engine`: 推論エンジン。
63    /// - `tokenizer`: 高速トークナイザー。
64    /// - `calib_config`: 較正設定。
65    pub fn new(
66        engine: Arc<InferenceEngine>,
67        tokenizer: Arc<JevTokenizer>,
68        calib_config: Arc<CalibrationConfig>,
69    ) -> Self {
70        let mut guardrail_config = GuardrailConfig::default();
71        guardrail_config.limits.max_questions = DEFAULT_MAX_QUESTIONS_PER_REQUEST;
72        let guardrail_pipeline = GuardrailPipeline::new(guardrail_config);
73
74        let auto_hierarchical = std::env::var("SOKUTO_AUTO_HIERARCHICAL")
75            .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
76            .unwrap_or(false);
77
78        let dag_executor = Arc::new(InProcessDagExecutor::new(
79            Arc::clone(&engine),
80            Arc::clone(&tokenizer),
81            (*calib_config).clone(),
82        ));
83
84        Self {
85            engine,
86            tokenizer,
87            calib_config,
88            coarse_config: CoarseToFineConfig::default(),
89            max_questions_per_request: DEFAULT_MAX_QUESTIONS_PER_REQUEST,
90            chunk_size: DEFAULT_CHUNK_SIZE,
91            guardrail_pipeline,
92            gating_config: GatingConfig {
93                enabled: false,
94                ..Default::default()
95            },
96            dag_executor,
97            hierarchical_mapping: None,
98            hierarchical_config: HierarchicalRouterConfig::default(),
99            auto_hierarchical,
100            dual_engine: None,
101            cache_manager: None,
102        }
103    }
104
105    /// 設定値をカスタマイズして `AppState` を構築する。
106    pub fn with_options(
107        engine: Arc<InferenceEngine>,
108        tokenizer: Arc<JevTokenizer>,
109        calib_config: Arc<CalibrationConfig>,
110        coarse_config: CoarseToFineConfig,
111        max_questions_per_request: usize,
112        chunk_size: usize,
113    ) -> Self {
114        let mut guardrail_config = GuardrailConfig::default();
115        guardrail_config.limits.max_questions = max_questions_per_request;
116        let guardrail_pipeline = GuardrailPipeline::new(guardrail_config);
117
118        let auto_hierarchical = std::env::var("SOKUTO_AUTO_HIERARCHICAL")
119            .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
120            .unwrap_or(false);
121
122        let dag_executor = Arc::new(InProcessDagExecutor::new(
123            Arc::clone(&engine),
124            Arc::clone(&tokenizer),
125            (*calib_config).clone(),
126        ));
127
128        Self {
129            engine,
130            tokenizer,
131            calib_config,
132            coarse_config,
133            max_questions_per_request,
134            chunk_size: chunk_size.max(1),
135            guardrail_pipeline,
136            gating_config: GatingConfig {
137                enabled: false,
138                ..Default::default()
139            },
140            dag_executor,
141            hierarchical_mapping: None,
142            hierarchical_config: HierarchicalRouterConfig::default(),
143            auto_hierarchical,
144            dual_engine: None,
145            cache_manager: None,
146        }
147    }
148
149    /// ガードレール設定も含めてフルカスタマイズした `AppState` を構築する。
150    pub fn with_guardrails(
151        engine: Arc<InferenceEngine>,
152        tokenizer: Arc<JevTokenizer>,
153        calib_config: Arc<CalibrationConfig>,
154        coarse_config: CoarseToFineConfig,
155        max_questions_per_request: usize,
156        chunk_size: usize,
157        guardrail_config: GuardrailConfig,
158    ) -> Self {
159        let auto_hierarchical = std::env::var("SOKUTO_AUTO_HIERARCHICAL")
160            .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
161            .unwrap_or(false);
162
163        let dag_executor = Arc::new(InProcessDagExecutor::new(
164            Arc::clone(&engine),
165            Arc::clone(&tokenizer),
166            (*calib_config).clone(),
167        ));
168
169        Self {
170            engine,
171            tokenizer,
172            calib_config,
173            coarse_config,
174            max_questions_per_request,
175            chunk_size: chunk_size.max(1),
176            guardrail_pipeline: GuardrailPipeline::new(guardrail_config),
177            gating_config: GatingConfig {
178                enabled: false,
179                ..Default::default()
180            },
181            dag_executor,
182            hierarchical_mapping: None,
183            hierarchical_config: HierarchicalRouterConfig::default(),
184            auto_hierarchical,
185            dual_engine: None,
186            cache_manager: None,
187        }
188    }
189
190    /// デュアルセッション推論エンジンおよびキャッシュマネージャーを設定する。
191    pub fn with_dual_engine(
192        mut self,
193        dual_engine: Arc<DualSessionEngine>,
194        cache_manager: Arc<StateCacheManager>,
195    ) -> Self {
196        self.dual_engine = Some(dual_engine);
197        self.cache_manager = Some(cache_manager);
198        self
199    }
200
201    /// ゲーティング設定をカスタマイズして設定する。
202    pub fn with_gating_config(mut self, gating_config: GatingConfig) -> Self {
203        self.gating_config = gating_config;
204        self
205    }
206
207    /// 階層オントロジーマッピングを設定する。
208    pub fn with_hierarchical_mapping(mut self, mapping: Arc<HierarchicalMapping>) -> Self {
209        self.hierarchical_mapping = Some(mapping);
210        self
211    }
212
213    /// 粗密階層ルーティング設定をカスタマイズして設定する。
214    pub fn with_hierarchical_config(mut self, config: HierarchicalRouterConfig) -> Self {
215        self.hierarchical_config = config;
216        self
217    }
218
219    /// 透過的粗密ルーティングモードの有効化を設定する。
220    pub fn with_auto_hierarchical(mut self, enabled: bool) -> Self {
221        self.auto_hierarchical = enabled;
222        self
223    }
224
225    /// カスタム DAG 実行器を設定する。
226    pub fn with_dag_executor(mut self, executor: Arc<InProcessDagExecutor>) -> Self {
227        self.dag_executor = executor;
228        self
229    }
230
231    /// 指定ディレクトリ内の `hierarchical_mapping.json` を読み込み、存在すれば登録する。
232    pub fn load_hierarchical_mapping_from_dir<P: AsRef<std::path::Path>>(mut self, dir: P) -> Self {
233        let mapping_path = dir.as_ref().join("hierarchical_mapping.json");
234        if mapping_path.exists()
235            && let Ok(content) = std::fs::read_to_string(&mapping_path)
236        {
237            if let Ok(mapping) = serde_json::from_str::<HierarchicalMapping>(&content) {
238                if mapping.validate().is_ok() {
239                    tracing::info!("粗密階層マッピングをロードしました: {}", mapping.name);
240                    self.hierarchical_mapping = Some(Arc::new(mapping));
241                } else {
242                    tracing::warn!("階層マッピングの検証に失敗しました: {:?}", mapping_path);
243                }
244            } else {
245                tracing::warn!(
246                    "階層マッピング JSON のパースに失敗しました: {:?}",
247                    mapping_path
248                );
249            }
250        }
251
252        self
253    }
254
255    /// サーバーがリクエストを安全に処理可能(Readiness)であるかを検査する。
256    pub fn is_ready(&self) -> bool {
257        self.engine.pool_size() > 0 && self.tokenizer.max_sequence_length() > 0
258    }
259}
sokuto_server/state.rs

sokuto_server/
state.rs
