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

mini_chat/
module.rs
