use crate::{
Capability, CapabilityRegistry, ConnectionProvider, ConnectionProviderRegistry, DriverRegistry,
EgressService, EmailSender, UtilityLlmService,
traits::{DisabledSessionFileSystemFactory, SessionFileSystemFactory},
};
use serde_json::Value;
use std::sync::Arc;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum BuiltInHarnessRole {
Base,
Default,
Chat,
}
#[derive(Debug, Clone)]
pub struct BuiltInCapabilityDefinition {
pub id: String,
pub config: Value,
}
impl BuiltInCapabilityDefinition {
pub fn new(id: impl Into<String>) -> Self {
Self {
id: id.into(),
config: serde_json::json!({}),
}
}
pub fn with_config(id: impl Into<String>, config: Value) -> Self {
Self {
id: id.into(),
config,
}
}
}
#[derive(Debug, Clone)]
pub struct BuiltInHarnessDefinition {
pub name: String,
pub display_name: String,
pub description: String,
pub system_prompt: String,
pub parent_name: Option<String>,
pub tags: Vec<String>,
pub capabilities: Vec<BuiltInCapabilityDefinition>,
pub roles: Vec<BuiltInHarnessRole>,
}
impl BuiltInHarnessDefinition {
pub fn new(
name: impl Into<String>,
display_name: impl Into<String>,
description: impl Into<String>,
system_prompt: impl Into<String>,
) -> Self {
Self {
name: name.into(),
display_name: display_name.into(),
description: description.into(),
system_prompt: system_prompt.into(),
parent_name: None,
tags: Vec::new(),
capabilities: Vec::new(),
roles: Vec::new(),
}
}
pub fn with_tags<I, S>(mut self, tags: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.tags = tags.into_iter().map(Into::into).collect();
self
}
pub fn with_parent_name(mut self, parent_name: impl Into<String>) -> Self {
self.parent_name = Some(parent_name.into());
self
}
pub fn with_capabilities<I>(mut self, capabilities: I) -> Self
where
I: IntoIterator<Item = BuiltInCapabilityDefinition>,
{
self.capabilities = capabilities.into_iter().collect();
self
}
pub fn with_roles<I>(mut self, roles: I) -> Self
where
I: IntoIterator<Item = BuiltInHarnessRole>,
{
self.roles = roles.into_iter().collect();
self
}
pub fn has_role(&self, role: BuiltInHarnessRole) -> bool {
self.roles.contains(&role)
}
}
#[derive(Clone)]
pub struct PlatformDefinition {
capability_registry: CapabilityRegistry,
driver_registry: DriverRegistry,
connection_providers: ConnectionProviderRegistry,
built_in_harnesses: Vec<BuiltInHarnessDefinition>,
egress_service: Arc<dyn EgressService>,
email_sender: Arc<dyn EmailSender>,
utility_llm_service: Arc<dyn UtilityLlmService>,
session_file_system_factory: Arc<dyn SessionFileSystemFactory>,
}
impl PlatformDefinition {
pub fn new(capability_registry: CapabilityRegistry, driver_registry: DriverRegistry) -> Self {
Self {
capability_registry,
driver_registry,
connection_providers: ConnectionProviderRegistry::new(),
built_in_harnesses: Vec::new(),
egress_service: Arc::new(crate::DirectEgressService::default()),
email_sender: Arc::new(crate::DisabledEmailSender),
utility_llm_service: Arc::new(crate::DisabledUtilityLlmService),
session_file_system_factory: Arc::new(DisabledSessionFileSystemFactory),
}
}
pub fn builder() -> PlatformDefinitionBuilder {
PlatformDefinitionBuilder::new()
}
pub fn capability_registry(&self) -> &CapabilityRegistry {
&self.capability_registry
}
pub fn capability_registry_mut(&mut self) -> &mut CapabilityRegistry {
&mut self.capability_registry
}
pub fn driver_registry(&self) -> &DriverRegistry {
&self.driver_registry
}
pub fn driver_registry_mut(&mut self) -> &mut DriverRegistry {
&mut self.driver_registry
}
pub fn connection_providers(&self) -> &ConnectionProviderRegistry {
&self.connection_providers
}
pub fn connection_providers_mut(&mut self) -> &mut ConnectionProviderRegistry {
&mut self.connection_providers
}
pub fn built_in_harnesses(&self) -> &[BuiltInHarnessDefinition] {
&self.built_in_harnesses
}
pub fn built_in_harnesses_mut(&mut self) -> &mut Vec<BuiltInHarnessDefinition> {
&mut self.built_in_harnesses
}
pub fn egress_service(&self) -> Arc<dyn EgressService> {
self.egress_service.clone()
}
pub fn email_sender(&self) -> Arc<dyn EmailSender> {
self.email_sender.clone()
}
pub fn utility_llm_service(&self) -> Arc<dyn UtilityLlmService> {
self.utility_llm_service.clone()
}
pub fn session_file_system_factory(&self) -> Arc<dyn SessionFileSystemFactory> {
self.session_file_system_factory.clone()
}
pub fn add_built_in_harness(&mut self, harness: BuiltInHarnessDefinition) {
self.built_in_harnesses.push(harness);
}
pub fn harness_for_role(&self, role: BuiltInHarnessRole) -> Option<&BuiltInHarnessDefinition> {
self.built_in_harnesses.iter().find(|h| h.has_role(role))
}
}
impl Default for PlatformDefinition {
fn default() -> Self {
Self::new(CapabilityRegistry::new(), DriverRegistry::new())
}
}
impl std::fmt::Debug for PlatformDefinition {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let harness_keys: Vec<_> = self.built_in_harnesses.iter().map(|h| &h.name).collect();
f.debug_struct("PlatformDefinition")
.field("capabilities", &self.capability_registry)
.field("drivers", &self.driver_registry.registered_providers())
.field("connection_providers", &self.connection_providers)
.field("built_in_harnesses", &harness_keys)
.field("egress_service", &self.egress_service.name())
.field("email_sender", &self.email_sender.name())
.field("utility_llm_service", &self.utility_llm_service.name())
.field(
"session_file_system_factory",
&self.session_file_system_factory.name(),
)
.finish()
}
}
pub struct PlatformDefinitionBuilder {
platform: PlatformDefinition,
}
impl PlatformDefinitionBuilder {
pub fn new() -> Self {
Self {
platform: PlatformDefinition::default(),
}
}
pub fn capability_registry(mut self, registry: CapabilityRegistry) -> Self {
self.platform.capability_registry = registry;
self
}
pub fn capability(mut self, capability: impl Capability + 'static) -> Self {
self.platform.capability_registry.register(capability);
self
}
pub fn driver_registry(mut self, registry: DriverRegistry) -> Self {
self.platform.driver_registry = registry;
self
}
pub fn connection_providers(mut self, registry: ConnectionProviderRegistry) -> Self {
self.platform.connection_providers = registry;
self
}
pub fn connection_provider(mut self, provider: impl ConnectionProvider + 'static) -> Self {
self.platform.connection_providers.register(provider);
self
}
pub fn built_in_harnesses<I>(mut self, harnesses: I) -> Self
where
I: IntoIterator<Item = BuiltInHarnessDefinition>,
{
self.platform.built_in_harnesses = harnesses.into_iter().collect();
self
}
pub fn add_built_in_harness(mut self, harness: BuiltInHarnessDefinition) -> Self {
self.platform.built_in_harnesses.push(harness);
self
}
pub fn egress_service(mut self, service: Arc<dyn EgressService>) -> Self {
self.platform.egress_service = service;
self
}
pub fn email_sender(mut self, sender: Arc<dyn EmailSender>) -> Self {
self.platform.email_sender = sender;
self
}
pub fn utility_llm_service(mut self, service: Arc<dyn UtilityLlmService>) -> Self {
self.platform.utility_llm_service = service;
self
}
pub fn session_file_system_factory(
mut self,
factory: Arc<dyn SessionFileSystemFactory>,
) -> Self {
self.platform.session_file_system_factory = factory;
self
}
pub fn build(self) -> PlatformDefinition {
self.platform
}
}
impl Default for PlatformDefinitionBuilder {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::connection_provider::{
ConnectionFormSchema, ConnectionType, ConnectionValidation, FieldType, FormField,
};
use crate::{CapabilityStatus, CurrentTimeCapability};
use async_trait::async_trait;
struct TestProvider;
#[async_trait]
impl ConnectionProvider for TestProvider {
fn provider_id(&self) -> &str {
"test_provider"
}
fn display_name(&self) -> &str {
"Test Provider"
}
fn description(&self) -> &str {
"Test connection provider"
}
fn icon(&self) -> &str {
"plug"
}
fn connection_type(&self) -> ConnectionType {
ConnectionType::ApiKey
}
fn form_schema(&self) -> Option<ConnectionFormSchema> {
Some(ConnectionFormSchema {
fields: vec![FormField {
name: "api_key".to_string(),
label: "API Key".to_string(),
field_type: FieldType::Password,
required: true,
placeholder: None,
help_text: None,
}],
instructions_markdown: "Enter the API key.".to_string(),
})
}
async fn validate(&self, _credential: &str) -> Result<ConnectionValidation, String> {
Ok(ConnectionValidation {
provider_username: Some("test-user".to_string()),
provider_metadata: None,
})
}
}
#[test]
fn test_platform_definition_builder() {
let mut drivers = DriverRegistry::new();
crate::llmsim_driver::register_driver(&mut drivers);
let platform = PlatformDefinition::builder()
.driver_registry(drivers.clone())
.capability(CurrentTimeCapability)
.connection_provider(TestProvider)
.add_built_in_harness(
BuiltInHarnessDefinition::new(
"minimal",
"Minimal",
"Minimal harness",
"You are helpful.",
)
.with_roles([BuiltInHarnessRole::Base, BuiltInHarnessRole::Default]),
)
.build();
assert!(platform.capability_registry().has("current_time"));
assert!(platform.connection_providers().has("test_provider"));
assert_eq!(
platform
.harness_for_role(BuiltInHarnessRole::Base)
.unwrap()
.name,
"minimal"
);
assert!(
platform
.driver_registry()
.has_driver(&crate::ProviderType::LlmSim)
);
}
#[test]
fn test_platform_definition_mutation() {
let mut platform = PlatformDefinition::default();
platform
.capability_registry_mut()
.register(CurrentTimeCapability);
platform.connection_providers_mut().register(TestProvider);
platform.add_built_in_harness(
BuiltInHarnessDefinition::new("chat", "Chat", "Chat harness", "You are helpful.")
.with_roles([BuiltInHarnessRole::Chat]),
);
let info = crate::CapabilityInfo::from_core(
platform
.capability_registry()
.get("current_time")
.expect("current_time registered")
.as_ref(),
);
assert_eq!(info.status, CapabilityStatus::Available);
assert!(platform.connection_providers().has("test_provider"));
assert_eq!(
platform
.harness_for_role(BuiltInHarnessRole::Chat)
.unwrap()
.name,
"chat"
);
}
}