agent-runtime-definition-builder 0.1.0

//! Ergonomic builders for Runtime Agent packages.

use std::sync::Arc;

use runtime_extension_api::{
    AgentPackage, CompletionMode, ContextHook, ContextSource, Guard, Observer, PackageManifest,
    Tool,
};
use runtime_types::PackageId;
use sha2::{Digest, Sha256};
use thiserror::Error;

#[derive(Debug, Error)]
pub enum PackageBuildError {
    #[error("invalid package id: {0}")]
    InvalidId(String),
    #[error("package version must not be empty")]
    EmptyVersion,
    #[error("runtime type must not be empty")]
    EmptyRuntimeType,
    #[error("system prompt must not be empty")]
    EmptySystemPrompt,
    #[error("duplicate tool: {0}")]
    DuplicateTool(String),
    #[error("required completion tool is not registered: {0}")]
    MissingCompletionTool(String),
}

pub struct AgentPackageBuilder {
    id: String,
    version: String,
    runtime_type: String,
    system_prompt: String,
    tools: Vec<Arc<dyn Tool>>,
    context_sources: Vec<Arc<dyn ContextSource>>,
    context_hooks: Vec<Arc<dyn ContextHook>>,
    guards: Vec<Arc<dyn Guard>>,
    observers: Vec<Arc<dyn Observer>>,
    completion_mode: CompletionMode,
}

impl AgentPackageBuilder {
    pub fn new(id: impl Into<String>, version: impl Into<String>) -> Self {
        Self {
            id: id.into(),
            version: version.into(),
            runtime_type: "generic".into(),
            system_prompt: String::new(),
            tools: Vec::new(),
            context_sources: Vec::new(),
            context_hooks: Vec::new(),
            guards: Vec::new(),
            observers: Vec::new(),
            completion_mode: CompletionMode::ModelOrTool,
        }
    }

    pub fn runtime_type(mut self, value: impl Into<String>) -> Self {
        self.runtime_type = value.into();
        self
    }
    pub fn system_prompt(mut self, value: impl Into<String>) -> Self {
        self.system_prompt = value.into();
        self
    }
    pub fn tool(mut self, value: Arc<dyn Tool>) -> Self {
        self.tools.push(value);
        self
    }
    pub fn context_source(mut self, value: Arc<dyn ContextSource>) -> Self {
        self.context_sources.push(value);
        self
    }
    pub fn context_hook(mut self, value: Arc<dyn ContextHook>) -> Self {
        self.context_hooks.push(value);
        self
    }
    pub fn guard(mut self, value: Arc<dyn Guard>) -> Self {
        self.guards.push(value);
        self
    }
    pub fn observer(mut self, value: Arc<dyn Observer>) -> Self {
        self.observers.push(value);
        self
    }
    pub fn completion_mode(mut self, value: CompletionMode) -> Self {
        self.completion_mode = value;
        self
    }

    pub fn build(self) -> Result<Arc<AgentPackage>, PackageBuildError> {
        let id = PackageId::new(self.id)
            .map_err(|error| PackageBuildError::InvalidId(error.to_string()))?;
        if self.version.trim().is_empty() {
            return Err(PackageBuildError::EmptyVersion);
        }
        if self.runtime_type.trim().is_empty() {
            return Err(PackageBuildError::EmptyRuntimeType);
        }
        if self.system_prompt.trim().is_empty() {
            return Err(PackageBuildError::EmptySystemPrompt);
        }
        let mut names = std::collections::BTreeSet::new();
        for tool in &self.tools {
            let name = tool.descriptor().name;
            if !names.insert(name.clone()) {
                return Err(PackageBuildError::DuplicateTool(name));
            }
        }
        if let CompletionMode::RequiredTool(name) = &self.completion_mode
            && !names.contains(name)
        {
            return Err(PackageBuildError::MissingCompletionTool(name.clone()));
        }
        let digest = package_digest(
            id.as_str(),
            &self.version,
            &self.runtime_type,
            &self.system_prompt,
            names.iter().map(String::as_str),
        );
        Ok(Arc::new(AgentPackage {
            manifest: PackageManifest {
                id,
                version: self.version,
                runtime_type: self.runtime_type,
                schema_version: 1,
                digest,
            },
            system_prompt: self.system_prompt,
            tools: self.tools,
            context_sources: self.context_sources,
            context_hooks: self.context_hooks,
            guards: self.guards,
            observers: self.observers,
            completion_mode: self.completion_mode,
        }))
    }
}

fn package_digest<'a>(
    id: &str,
    version: &str,
    runtime_type: &str,
    prompt: &str,
    tools: impl Iterator<Item = &'a str>,
) -> String {
    let mut digest = Sha256::new();
    for value in [id, version, runtime_type, prompt] {
        digest.update(value.len().to_be_bytes());
        digest.update(value.as_bytes());
    }
    for tool in tools {
        digest.update(tool.len().to_be_bytes());
        digest.update(tool.as_bytes());
    }
    format!("sha256:{:x}", digest.finalize())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn package_digest_is_reproducible() {
        let one = package_digest("a", "1", "generic", "prompt", ["x", "y"].into_iter());
        let two = package_digest("a", "1", "generic", "prompt", ["x", "y"].into_iter());
        assert_eq!(one, two);
    }
}