mini_chat/infra/plugins/static_model_policy/
gear.rs1use std::sync::{Arc, OnceLock};
2
3use async_trait::async_trait;
4use mini_chat_sdk::{MiniChatModelPolicyPluginClientV1, MiniChatModelPolicyPluginSpecV1};
5use toolkit::Gear;
6use toolkit::client_hub::ClientScope;
7use toolkit::context::GearCtx;
8use toolkit::gts::PluginV1;
9use tracing::info;
10use types_registry_sdk::{RegisterResult, TypesRegistryClient};
11
12use super::config::StaticMiniChatPolicyPluginConfig;
13use super::service::Service;
14
15#[toolkit::gear(
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 Gear for StaticMiniChatModelPolicyPlugin {
36 async fn init(&self, ctx: &GearCtx) -> anyhow::Result<()> {
37 let cfg: StaticMiniChatPolicyPluginConfig = ctx.config_or_default()?;
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!("{} gear already initialized", Self::MODULE_NAME))?;
56
57 let (instance_id, instance_json) =
59 PluginV1::<MiniChatModelPolicyPluginSpecV1>::build_registration(
60 "cf.core._.static_mini_chat_model_policy.v1",
61 cfg.vendor.clone(),
62 cfg.priority,
63 )?;
64
65 let registry = ctx.client_hub().get::<dyn TypesRegistryClient>()?;
67 let results = registry.register(vec![instance_json]).await?;
68 RegisterResult::ensure_all_ok(&results)?;
69
70 let api: Arc<dyn MiniChatModelPolicyPluginClientV1> = service;
72 ctx.client_hub()
73 .register_scoped::<dyn MiniChatModelPolicyPluginClientV1>(
74 ClientScope::gts_id(&instance_id),
75 api,
76 );
77
78 info!(instance_id = %instance_id, "Static mini-chat model policy plugin registered");
79 Ok(())
80 }
81}