mini_chat/
module.rs

1use std::sync::{Arc, OnceLock};
2
3use async_trait::async_trait;
4use authz_resolver_sdk::AuthZResolverClient;
5use mini_chat_sdk::MiniChatModelPolicyPluginSpecV1;
6use modkit::api::OpenApiRegistry;
7use modkit::{DatabaseCapability, Module, ModuleCtx, RestApiCapability};
8use oagw_sdk::ServiceGatewayClientV1;
9use sea_orm_migration::MigrationTrait;
10use tracing::info;
11use types_registry_sdk::{RegisterResult, TypesRegistryClient};
12
13use crate::api::rest::routes;
14use crate::domain::service::{AppServices as GenericAppServices, Repositories};
15
16pub(crate) type AppServices = GenericAppServices<
17    TurnRepository,
18    MessageRepository,
19    QuotaUsageRepository,
20    ReactionRepository,
21    ChatRepository,
22>;
23use crate::infra::db::repo::attachment_repo::AttachmentRepository;
24use crate::infra::db::repo::chat_repo::ChatRepository;
25use crate::infra::db::repo::message_repo::MessageRepository;
26use crate::infra::db::repo::model_pref_repo::ModelPrefRepository;
27use crate::infra::db::repo::quota_usage_repo::QuotaUsageRepository;
28use crate::infra::db::repo::reaction_repo::ReactionRepository;
29use crate::infra::db::repo::thread_summary_repo::ThreadSummaryRepository;
30use crate::infra::db::repo::turn_repo::TurnRepository;
31use crate::infra::db::repo::vector_store_repo::VectorStoreRepository;
32use crate::infra::llm::provider_resolver::ProviderResolver;
33use crate::infra::model_policy::ModelPolicyGateway;
34
35/// Default URL prefix for all mini-chat REST routes.
36pub const DEFAULT_URL_PREFIX: &str = "/mini-chat";
37
38/// The mini-chat module: multi-tenant AI chat with SSE streaming.
39#[modkit::module(
40    name = "mini-chat",
41    deps = ["types-registry", "authz-resolver", "oagw"],
42    capabilities = [db, rest],
43)]
44pub struct MiniChatModule {
45    service: OnceLock<Arc<AppServices>>,
46    url_prefix: OnceLock<String>,
47}
48
49impl Default for MiniChatModule {
50    fn default() -> Self {
51        Self {
52            service: OnceLock::new(),
53            url_prefix: OnceLock::new(),
54        }
55    }
56}
57
58#[async_trait]
59impl Module for MiniChatModule {
60    async fn init(&self, ctx: &ModuleCtx) -> anyhow::Result<()> {
61        info!("Initializing {} module", Self::MODULE_NAME);
62
63        let cfg: crate::config::MiniChatConfig = ctx.config()?;
64        cfg.streaming
65            .validate()
66            .map_err(|e| anyhow::anyhow!("streaming config: {e}"))?;
67        cfg.estimation_budgets
68            .validate()
69            .map_err(|e| anyhow::anyhow!("estimation_budgets config: {e}"))?;
70        cfg.quota
71            .validate()
72            .map_err(|e| anyhow::anyhow!("quota config: {e}"))?;
73        cfg.outbox
74            .validate()
75            .map_err(|e| anyhow::anyhow!("outbox config: {e}"))?;
76        for (id, entry) in &cfg.providers {
77            entry
78                .validate(id)
79                .map_err(|e| anyhow::anyhow!("providers config: {e}"))?;
80        }
81
82        let vendor = cfg.vendor.trim().to_owned();
83        if vendor.is_empty() {
84            return Err(anyhow::anyhow!(
85                "{}: vendor must be a non-empty string",
86                Self::MODULE_NAME
87            ));
88        }
89
90        // Register model-policy plugin schema in types-registry
91        let registry = ctx.client_hub().get::<dyn TypesRegistryClient>()?;
92        let schema_str = MiniChatModelPolicyPluginSpecV1::gts_schema_with_refs_as_string();
93        let mut schema_json: serde_json::Value = serde_json::from_str(&schema_str)?;
94        if let Some(obj) = schema_json.as_object_mut() {
95            obj.insert(
96                "additionalProperties".to_owned(),
97                serde_json::Value::Bool(false),
98            );
99        }
100        let results = registry.register(vec![schema_json]).await?;
101        RegisterResult::ensure_all_ok(&results)?;
102        info!(
103            schema_id = %MiniChatModelPolicyPluginSpecV1::gts_schema_id(),
104            "Registered model-policy plugin schema in types-registry"
105        );
106
107        self.url_prefix
108            .set(cfg.url_prefix)
109            .map_err(|_| anyhow::anyhow!("{} url_prefix already set", Self::MODULE_NAME))?;
110
111        let db = Arc::new(ctx.db_required()?);
112
113        let authz = ctx
114            .client_hub()
115            .get::<dyn AuthZResolverClient>()
116            .map_err(|e| anyhow::anyhow!("failed to get AuthZ resolver: {e}"))?;
117
118        let gateway = ctx
119            .client_hub()
120            .get::<dyn ServiceGatewayClientV1>()
121            .map_err(|e| anyhow::anyhow!("failed to get OAGW gateway: {e}"))?;
122
123        // Register OAGW upstreams for each configured provider.
124        crate::infra::oagw_provisioning::register_oagw_upstreams(&gateway, &cfg.providers).await?;
125
126        let provider_resolver = Arc::new(ProviderResolver::new(&gateway, cfg.providers));
127
128        let repos = Repositories {
129            chat: Arc::new(ChatRepository::new(modkit_db::odata::LimitCfg {
130                default: 20,
131                max: 100,
132            })),
133            attachment: Arc::new(AttachmentRepository),
134            message: Arc::new(MessageRepository::new(modkit_db::odata::LimitCfg {
135                default: 20,
136                max: 100,
137            })),
138            quota: Arc::new(QuotaUsageRepository),
139            turn: Arc::new(TurnRepository),
140            reaction: Arc::new(ReactionRepository),
141            model_pref: Arc::new(ModelPrefRepository),
142            thread_summary: Arc::new(ThreadSummaryRepository),
143            vector_store: Arc::new(VectorStoreRepository),
144        };
145
146        let model_policy_gw = Arc::new(ModelPolicyGateway::new(ctx.client_hub(), vendor));
147        let services = Arc::new(AppServices::new(
148            &repos,
149            db,
150            authz,
151            &(model_policy_gw.clone() as Arc<dyn crate::domain::repos::ModelResolver>),
152            provider_resolver,
153            cfg.streaming,
154            model_policy_gw.clone() as Arc<dyn crate::domain::repos::PolicySnapshotProvider>,
155            model_policy_gw as Arc<dyn crate::domain::repos::UserLimitsProvider>,
156            cfg.estimation_budgets,
157            cfg.quota,
158        ));
159
160        self.service
161            .set(services)
162            .map_err(|_| anyhow::anyhow!("{} module already initialized", Self::MODULE_NAME))?;
163
164        info!("{} module initialized successfully", Self::MODULE_NAME);
165        Ok(())
166    }
167}
168
169impl DatabaseCapability for MiniChatModule {
170    fn migrations(&self) -> Vec<Box<dyn MigrationTrait>> {
171        use sea_orm_migration::MigratorTrait;
172        info!("Providing mini-chat database migrations");
173        crate::infra::db::migrations::Migrator::migrations()
174    }
175}
176
177impl RestApiCapability for MiniChatModule {
178    fn register_rest(
179        &self,
180        _ctx: &ModuleCtx,
181        router: axum::Router,
182        openapi: &dyn OpenApiRegistry,
183    ) -> anyhow::Result<axum::Router> {
184        let services = self
185            .service
186            .get()
187            .ok_or_else(|| anyhow::anyhow!("{} not initialized", Self::MODULE_NAME))?;
188
189        info!("Registering mini-chat REST routes");
190        let prefix = self
191            .url_prefix
192            .get()
193            .ok_or_else(|| anyhow::anyhow!("{} not initialized (url_prefix)", Self::MODULE_NAME))?;
194
195        let router = routes::register_routes(router, openapi, Arc::clone(services), prefix);
196        info!("Mini-chat REST routes registered successfully");
197        Ok(router)
198    }
199}
mini_chat/module.rs

mini_chat/
module.rs
