llm-worker-rs/worker/examples/plugin_usage.rs

82 lines
2.6 KiB
Rust

use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use worker::{Worker, Role, plugin::{PluginRegistry, ProviderPlugin, example_provider::CustomProviderPlugin}};
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
// Initialize tracing for debugging
tracing_subscriber::fmt::init();
// Create a plugin registry
let plugin_registry = Arc::new(Mutex::new(PluginRegistry::new()));
// Create and initialize a custom provider plugin
let mut custom_plugin = CustomProviderPlugin::new();
let mut config = HashMap::new();
config.insert(
"base_url".to_string(),
serde_json::Value::String("https://api.custom-provider.com".to_string()),
);
config.insert(
"timeout".to_string(),
serde_json::Value::Number(serde_json::Number::from(60)),
);
custom_plugin.initialize(config).await?;
// Register the plugin
{
let mut registry = plugin_registry.lock().unwrap();
registry.register(Arc::new(custom_plugin))?;
}
// List available plugins
{
let registry = plugin_registry.lock().unwrap();
let plugins = registry.list();
println!("Available plugins:");
for plugin in plugins {
println!(" - {} ({}): {}", plugin.name, plugin.id, plugin.description);
println!(" Supported models: {:?}", plugin.supported_models);
}
}
// Create a Worker instance using the plugin
let role = Role::new(
"assistant",
"A helpful AI assistant",
"You are a helpful, harmless, and honest AI assistant powered by a custom LLM provider."
);
let worker = Worker::builder()
.plugin("custom-provider", plugin_registry.clone())
.model("custom-turbo")
.api_key("__plugin__", "custom-1234567890abcdefghijklmnop")
.role(role)
.build()?;
println!("\nWorker created successfully with custom provider plugin!");
// Example: List plugins from the worker
let plugin_list = worker.list_plugins()?;
println!("\nPlugins registered in worker:");
for metadata in plugin_list {
println!(" - {}: v{} by {}", metadata.name, metadata.version, metadata.author);
}
// Load plugins from directory (if dynamic loading is enabled)
#[cfg(feature = "dynamic-loading")]
{
use std::path::Path;
let plugin_dir = Path::new("./plugins");
if plugin_dir.exists() {
let mut worker = worker;
worker.load_plugins_from_directory(plugin_dir).await?;
println!("\nLoaded plugins from directory: {:?}", plugin_dir);
}
}
Ok(())
}