mini_chat/infra/plugins/static_model_policy/
module.rs1use std::sync::{Arc, OnceLock};
2
3use async_trait::async_trait;
4use mini_chat_sdk::{MiniChatModelPolicyPluginClientV1, MiniChatModelPolicyPluginSpecV1};
5use modkit::Module;
6use modkit::client_hub::ClientScope;
7use modkit::context::ModuleCtx;
8use modkit::gts::BaseModkitPluginV1;
9use tracing::info;
10use types_registry_sdk::{RegisterResult, TypesRegistryClient};
11
12use super::config::StaticMiniChatPolicyPluginConfig;
13use super::service::Service;
14
15#[modkit::module(
19 name = "static-mini-chat-model-policy-plugin",
20 deps = ["types-registry"]
21)]
22pub struct StaticMiniChatModelPolicyPlugin {
23 service: OnceLock<Arc<Service>>,
24}
25
26impl Default for StaticMiniChatModelPolicyPlugin {
27 fn default() -> Self {
28 Self {
29 service: OnceLock::new(),
30 }
31 }
32}
33
34#[async_trait]
35impl Module for StaticMiniChatModelPolicyPlugin {
36 async fn init(&self, ctx: &ModuleCtx) -> anyhow::Result<()> {
37 let cfg: StaticMiniChatPolicyPluginConfig = ctx.config()?;
38 info!(
39 vendor = %cfg.vendor,
40 priority = cfg.priority,
41 models = cfg.model_catalog.len(),
42 "Loaded static mini-chat model policy plugin configuration"
43 );
44
45 let service = Arc::new(Service::new(
48 cfg.model_catalog,
49 cfg.kill_switches,
50 cfg.default_standard_limits,
51 cfg.default_premium_limits,
52 ));
53 self.service
54 .set(service.clone())
55 .map_err(|_| anyhow::anyhow!("{} module already initialized", Self::MODULE_NAME))?;
56
57 let instance_id = MiniChatModelPolicyPluginSpecV1::gts_make_instance_id(
59 "x.core._.static_mini_chat_model_policy.v1",
60 );
61
62 let registry = ctx.client_hub().get::<dyn TypesRegistryClient>()?;
64 let instance = BaseModkitPluginV1::<MiniChatModelPolicyPluginSpecV1> {
65 id: instance_id.clone(),
66 vendor: cfg.vendor.clone(),
67 priority: cfg.priority,
68 properties: MiniChatModelPolicyPluginSpecV1,
69 };
70 let instance_json = serde_json::to_value(&instance)?;
71
72 let results = registry.register(vec![instance_json]).await?;
73 RegisterResult::ensure_all_ok(&results)?;
74
75 let api: Arc<dyn MiniChatModelPolicyPluginClientV1> = service;
77 ctx.client_hub()
78 .register_scoped::<dyn MiniChatModelPolicyPluginClientV1>(
79 ClientScope::gts_id(&instance_id),
80 api,
81 );
82
83 info!(instance_id = %instance_id, "Static mini-chat model policy plugin registered");
84 Ok(())
85 }
86}