reflex_server/gateway/
state.rs

1use genai::Client;
2use std::path::PathBuf;
3use std::sync::Arc;
4
5use reflex::cache::{BqSearchBackend, StorageLoader, TieredCache};
6use reflex::scoring::CrossEncoderScorer;
7
8#[derive(Clone)]
9pub struct HandlerState<
10    B: BqSearchBackend + Clone + Send + Sync + 'static,
11    S: StorageLoader + Clone + Send + Sync + 'static,
12> {
13    pub tiered_cache: Arc<TieredCache<B, S>>,
14
15    pub scorer: Arc<CrossEncoderScorer>,
16
17    pub storage_path: PathBuf,
18
19    pub bq_client: B,
20
21    pub collection_name: String,
22
23    pub genai_client: Client,
24
25    pub mock_provider: bool,
26}
27
28impl<B, S> HandlerState<B, S>
29where
30    B: BqSearchBackend + Clone + Send + Sync + 'static,
31    S: StorageLoader + Clone + Send + Sync + 'static,
32{
33    pub fn new(
34        tiered_cache: Arc<TieredCache<B, S>>,
35        scorer: Arc<CrossEncoderScorer>,
36        storage_path: PathBuf,
37        bq_client: B,
38        collection_name: String,
39    ) -> Self {
40        let mock_provider = std::env::var_os("REFLEX_MOCK_PROVIDER").is_some_and(|v| !v.is_empty());
41        Self {
42            tiered_cache,
43            scorer,
44            storage_path,
45            bq_client,
46            collection_name,
47            genai_client: Client::default(),
48            mock_provider,
49        }
50    }
51
52    pub fn new_with_mock_provider(
53        tiered_cache: Arc<TieredCache<B, S>>,
54        scorer: Arc<CrossEncoderScorer>,
55        storage_path: PathBuf,
56        bq_client: B,
57        collection_name: String,
58        mock_provider: bool,
59    ) -> Self {
60        Self {
61            tiered_cache,
62            scorer,
63            storage_path,
64            bq_client,
65            collection_name,
66            genai_client: Client::default(),
67            mock_provider,
68        }
69    }
70}