Skip to main content

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::providers::{ProviderConfig, ProviderKind, create_provider};
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
77        let vendor = cfg.vendor.trim().to_owned();
78        if vendor.is_empty() {
79            return Err(anyhow::anyhow!(
80                "{}: vendor must be a non-empty string",
81                Self::MODULE_NAME
82            ));
83        }
84
85        // Register model-policy plugin schema in types-registry
86        let registry = ctx.client_hub().get::<dyn TypesRegistryClient>()?;
87        let schema_str = MiniChatModelPolicyPluginSpecV1::gts_schema_with_refs_as_string();
88        let mut schema_json: serde_json::Value = serde_json::from_str(&schema_str)?;
89        if let Some(obj) = schema_json.as_object_mut() {
90            obj.insert(
91                "additionalProperties".to_owned(),
92                serde_json::Value::Bool(false),
93            );
94        }
95        let results = registry.register(vec![schema_json]).await?;
96        RegisterResult::ensure_all_ok(&results)?;
97        info!(
98            schema_id = %MiniChatModelPolicyPluginSpecV1::gts_schema_id(),
99            "Registered model-policy plugin schema in types-registry"
100        );
101
102        self.url_prefix
103            .set(cfg.url_prefix)
104            .map_err(|_| anyhow::anyhow!("{} url_prefix already set", Self::MODULE_NAME))?;
105
106        let db = Arc::new(ctx.db_required()?);
107
108        let authz = ctx
109            .client_hub()
110            .get::<dyn AuthZResolverClient>()
111            .map_err(|e| anyhow::anyhow!("failed to get AuthZ resolver: {e}"))?;
112
113        let gateway = ctx
114            .client_hub()
115            .get::<dyn ServiceGatewayClientV1>()
116            .map_err(|e| anyhow::anyhow!("failed to get OAGW gateway: {e}"))?;
117
118        // TODO: provider kind and upstream alias should come from config in a
119        // follow-up — hardcoded to OpenAI Responses for initial P1 wiring.
120        let llm = create_provider(
121            gateway,
122            ProviderConfig {
123                kind: ProviderKind::OpenAiResponses,
124                upstream_alias: "openai".to_owned(),
125            },
126        );
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            llm,
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}