95 lines
2.9 KiB
Rust
95 lines
2.9 KiB
Rust
use std::collections::HashMap;
|
|
use std::sync::{Arc, Mutex};
|
|
use worker::{
|
|
PromptError, SystemPromptContext, Worker,
|
|
plugin::{PluginRegistry, ProviderPlugin, example_provider::CustomProviderPlugin},
|
|
};
|
|
use worker_types::Message;
|
|
|
|
#[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
|
|
fn plugin_prompt(
|
|
_ctx: &SystemPromptContext,
|
|
_messages: &[Message],
|
|
) -> Result<String, PromptError> {
|
|
Ok("You are a helpful assistant powered by a custom provider.".to_string())
|
|
}
|
|
|
|
let mut blueprint = Worker::blueprint();
|
|
blueprint
|
|
.plugin_provider("custom-provider", plugin_registry.clone())
|
|
.model("custom-turbo")
|
|
.api_key("__plugin__", "custom-1234567890abcdefghijklmnop")
|
|
.system_prompt_fn(plugin_prompt);
|
|
let worker = blueprint.instantiate()?;
|
|
|
|
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(())
|
|
}
|