batuta/serve/banco/
state.rs

1//! Banco application state shared across all handlers via `Arc`.
2
3use crate::serve::backends::{BackendSelector, PrivacyTier, ServingBackend};
4use crate::serve::circuit_breaker::{CircuitState, CostCircuitBreaker};
5use crate::serve::context::ContextManager;
6use crate::serve::router::SpilloverRouter;
7use crate::serve::templates::ChatTemplateEngine;
8use std::sync::Arc;
9use std::time::Instant;
10
11use super::audit::AuditLog;
12use super::auth::AuthStore;
13use super::batch::BatchStore;
14use super::config::BancoConfig;
15use super::conversations::ConversationStore;
16use super::eval::EvalStore;
17use super::events::EventBus;
18use super::experiment::ExperimentStore;
19use super::metrics::MetricsCollector;
20use super::model_slot::ModelSlot;
21use super::prompts::PromptStore;
22use super::rag::RagIndex;
23use super::recipes::RecipeStore;
24use super::storage::FileStore;
25use super::tools::ToolRegistry;
26use super::training::TrainingStore;
27use super::types::{HealthResponse, InferenceParams, ModelInfo, ModelsResponse, SystemResponse};
28use std::sync::RwLock;
29
30// ============================================================================
31// BANCO-STA-001: State
32// ============================================================================
33
34/// Inner state — not `Clone` because of atomics in router/circuit breaker.
35pub struct BancoStateInner {
36    pub backend_selector: BackendSelector,
37    pub router: SpilloverRouter,
38    pub circuit_breaker: CostCircuitBreaker,
39    pub context_manager: ContextManager,
40    pub template_engine: ChatTemplateEngine,
41    pub privacy_tier: PrivacyTier,
42    pub start_time: Instant,
43    pub conversations: Arc<ConversationStore>,
44    pub prompts: PromptStore,
45    pub auth: AuthStore,
46    pub model: ModelSlot,
47    pub inference_params: RwLock<InferenceParams>,
48    pub files: Arc<FileStore>,
49    pub recipes: Arc<RecipeStore>,
50    pub rag: RagIndex,
51    pub evals: Arc<EvalStore>,
52    pub training: Arc<TrainingStore>,
53    pub experiments: Arc<ExperimentStore>,
54    pub batches: Arc<BatchStore>,
55    pub tools: ToolRegistry,
56    pub metrics: MetricsCollector,
57    pub audit_log: AuditLog,
58    pub events: EventBus,
59}
60
61/// Shared handle passed to axum handlers.
62pub type BancoState = Arc<BancoStateInner>;
63
64impl BancoStateInner {
65    /// Create default state (Standard privacy, default backends).
66    #[must_use]
67    pub fn with_defaults() -> BancoState {
68        Self::with_privacy(PrivacyTier::Standard)
69    }
70
71    /// Create state from `~/.banco/config.toml` (loads on disk, falls back to defaults).
72    #[must_use]
73    pub fn from_config() -> BancoState {
74        let config = BancoConfig::load();
75        let tier: PrivacyTier = config.server.privacy_tier.into();
76        let cb_config = crate::serve::circuit_breaker::CircuitBreakerConfig {
77            daily_budget_usd: config.budget.daily_limit_usd,
78            max_request_cost_usd: config.budget.max_request_usd,
79            ..Default::default()
80        };
81
82        // Use disk-backed stores when ~/.banco/ is available
83        let data_dir = BancoConfig::config_dir();
84        let conversations = match &data_dir {
85            Some(dir) => ConversationStore::with_data_dir(dir.join("conversations")),
86            None => ConversationStore::in_memory(),
87        };
88        let files = match &data_dir {
89            Some(dir) => FileStore::with_data_dir(dir.clone()),
90            None => FileStore::in_memory(),
91        };
92
93        let state = Arc::new(Self {
94            backend_selector: BackendSelector::new().with_privacy(tier),
95            router: SpilloverRouter::with_defaults(),
96            circuit_breaker: CostCircuitBreaker::new(cb_config),
97            context_manager: ContextManager::default(),
98            template_engine: ChatTemplateEngine::default(),
99            privacy_tier: tier,
100            start_time: Instant::now(),
101            conversations,
102            prompts: PromptStore::new(),
103            auth: AuthStore::local(),
104            model: ModelSlot::empty(),
105            inference_params: RwLock::new(InferenceParams::default()),
106            files,
107            recipes: RecipeStore::new(),
108            rag: RagIndex::new(),
109            evals: EvalStore::new(),
110            training: TrainingStore::new(),
111            experiments: match &data_dir {
112                Some(dir) => ExperimentStore::with_data_dir(dir.join("experiments")),
113                None => ExperimentStore::new(),
114            },
115            batches: BatchStore::new(),
116            tools: ToolRegistry::default(),
117            metrics: MetricsCollector::default(),
118            audit_log: match &data_dir {
119                Some(dir) => AuditLog::with_file(dir.join("audit.jsonl")),
120                None => AuditLog::new(),
121            },
122            events: EventBus::default(),
123        });
124
125        // Re-index loaded files into RAG
126        let loaded_files = state.files.list();
127        for file_info in &loaded_files {
128            if let Some(content) = state.files.read_content(&file_info.id) {
129                let text = String::from_utf8_lossy(&content);
130                state.rag.index_document(&file_info.id, &file_info.name, &text);
131            }
132        }
133        if !loaded_files.is_empty() {
134            eprintln!("[banco] Indexed {} files for RAG", loaded_files.len());
135        }
136
137        state
138    }
139
140    /// Create state with a specific privacy tier.
141    #[must_use]
142    pub fn with_privacy(tier: PrivacyTier) -> BancoState {
143        Arc::new(Self {
144            backend_selector: BackendSelector::new().with_privacy(tier),
145            router: SpilloverRouter::with_defaults(),
146            circuit_breaker: CostCircuitBreaker::with_defaults(),
147            context_manager: ContextManager::default(),
148            template_engine: ChatTemplateEngine::default(),
149            privacy_tier: tier,
150            start_time: Instant::now(),
151            conversations: ConversationStore::in_memory(),
152            prompts: PromptStore::new(),
153            auth: AuthStore::local(),
154            model: ModelSlot::empty(),
155            inference_params: RwLock::new(InferenceParams::default()),
156            files: FileStore::in_memory(),
157            recipes: RecipeStore::new(),
158            rag: RagIndex::new(),
159            evals: EvalStore::new(),
160            training: TrainingStore::new(),
161            experiments: ExperimentStore::new(),
162            batches: BatchStore::new(),
163            tools: ToolRegistry::default(),
164            metrics: MetricsCollector::default(),
165            audit_log: AuditLog::new(),
166            events: EventBus::default(),
167        })
168    }
169
170    /// Build a `HealthResponse` snapshot.
171    #[must_use]
172    pub fn health_status(&self) -> HealthResponse {
173        let cb_state = match self.circuit_breaker.state() {
174            CircuitState::Closed => "closed",
175            CircuitState::Open => "open",
176            CircuitState::HalfOpen => "half_open",
177        };
178        HealthResponse {
179            status: "ok".to_string(),
180            circuit_breaker_state: cb_state.to_string(),
181            uptime_secs: self.start_time.elapsed().as_secs(),
182        }
183    }
184
185    /// Build a `ModelsResponse` from recommended backends.
186    #[must_use]
187    pub fn list_models(&self) -> ModelsResponse {
188        let backends = self.backend_selector.recommend();
189        let data = backends
190            .iter()
191            .map(|b| ModelInfo {
192                id: format!("{b:?}").to_lowercase(),
193                object: "model".to_string(),
194                owned_by: "batuta".to_string(),
195                local: b.is_local(),
196            })
197            .collect();
198        ModelsResponse { object: "list".to_string(), data }
199    }
200
201    /// Build a `SystemResponse`.
202    #[must_use]
203    pub fn system_info(&self) -> SystemResponse {
204        let backends = self.backend_selector.recommend();
205        let rag_status = self.rag.status();
206        let tokenizer = if self.model.is_loaded() {
207            #[cfg(feature = "aprender")]
208            {
209                Some(if self.model.has_bpe_tokenizer() { "bpe" } else { "greedy" }.to_string())
210            }
211            #[cfg(not(feature = "aprender"))]
212            {
213                Some("greedy".to_string())
214            }
215        } else {
216            None
217        };
218        SystemResponse {
219            privacy_tier: format!("{:?}", self.privacy_tier),
220            backends: backends.iter().map(|b| format!("{b:?}")).collect(),
221            gpu_available: backends.contains(&ServingBackend::Realizar),
222            version: env!("CARGO_PKG_VERSION").to_string(),
223            telemetry: false,
224            model_loaded: self.model.is_loaded(),
225            model_id: self.model.info().map(|m| m.model_id),
226            hint: if self.model.is_loaded() {
227                None
228            } else {
229                Some(
230                    "Load a model: POST /api/v1/models/load {\"model\": \"./model.gguf\"}"
231                        .to_string(),
232                )
233            },
234            tokenizer,
235            endpoints: 82,
236            files: self.files.len(),
237            conversations: self.conversations.len(),
238            rag_indexed: rag_status.indexed,
239            rag_chunks: rag_status.chunk_count,
240            training_runs: self.training.list().len(),
241            audit_entries: self.audit_log.len(),
242        }
243    }
244}
batuta/serve/banco/state.rs

batuta/serve/banco/
state.rs
