sekuire 0.1.0

The official SDK for the Sekuire Agent Identity Protocol
Documentation
use anyhow::{Context, Result};
use std::collections::HashMap;
use std::fs;
use std::path::Path;
use std::sync::Mutex;

use crate::config::{get_agent_config, load_config, load_system_prompt, load_tools, LLMConfig};
use crate::llm::{create_llm_provider, ChatOptions, LLMProvider, Message};
use crate::policy::{ActivePolicy, PolicyEnforcer};

// Simple global cache for policy files
static POLICY_CACHE: Mutex<Vec<(String, ActivePolicy)>> = Mutex::new(Vec::new());

/// Sekuire Agent - Config-first AI agent
pub struct SekuireAgent {
    pub name: String,
    llm: Box<dyn LLMProvider>,
    system_prompt: String,
    tools: Vec<String>,
    conversation_history: Vec<Message>,
    max_history_messages: usize,
    policy_enforcer: Option<PolicyEnforcer>,
}

impl SekuireAgent {
    /// Create a new agent
    pub fn new(
        name: String,
        llm: Box<dyn LLMProvider>,
        system_prompt: String,
        tools: Vec<String>,
        max_history_messages: Option<usize>,
        policy_enforcer: Option<PolicyEnforcer>,
    ) -> Self {
        Self {
            name,
            llm,
            system_prompt,
            tools,
            conversation_history: Vec::new(),
            max_history_messages: max_history_messages.unwrap_or(10),
            policy_enforcer,
        }
    }

    /// Chat with the agent
    pub async fn chat(
        &mut self,
        user_message: &str,
        options: Option<ChatOptions>,
    ) -> Result<String> {
        // Enforce Policy
        if let Some(enforcer) = &self.policy_enforcer {
            enforcer.enforce_model(self.get_model_name())?;
            enforcer.enforce_rate_limit("request", 1)?;
        }

        // Build messages array
        let mut messages = vec![Message {
            role: "system".to_string(),
            content: self.system_prompt.clone(),
        }];

        // Add conversation history (last N messages)
        let history_start = self
            .conversation_history
            .len()
            .saturating_sub(self.max_history_messages);
        messages.extend_from_slice(&self.conversation_history[history_start..]);

        // Add current user message
        messages.push(Message {
            role: "user".to_string(),
            content: user_message.to_string(),
        });

        // Get response from LLM
        let response = self.llm.chat(&messages, options).await?;

        // Store in conversation history
        self.conversation_history.push(Message {
            role: "user".to_string(),
            content: user_message.to_string(),
        });
        self.conversation_history.push(Message {
            role: "assistant".to_string(),
            content: response.content.clone(),
        });

        Ok(response.content)
    }

    /// Clear conversation history
    pub fn clear_history(&mut self) {
        self.conversation_history.clear();
    }

    /// Get conversation history
    pub fn get_history(&self) -> &[Message] {
        &self.conversation_history
    }

    /// Get LLM provider name
    pub fn get_llm_provider(&self) -> &str {
        self.llm.get_provider_name()
    }

    /// Get LLM model name
    pub fn get_model_name(&self) -> &str {
        self.llm.get_model_name()
    }

    /// Get available tools
    pub fn get_tools(&self) -> Vec<String> {
        if let Some(enforcer) = &self.policy_enforcer {
            self.tools
                .iter()
                .filter(|t| enforcer.enforce_tool(t).is_ok())
                .cloned()
                .collect()
        } else {
            self.tools.clone()
        }
    }

    /// Get the active policy enforcer
    pub fn get_policy_enforcer(&self) -> Option<&PolicyEnforcer> {
        self.policy_enforcer.as_ref()
    }
}

/// Load an agent from sekuire.yml configuration
///
/// # Examples
///
/// ```no_run
/// use sekuire_sdk::get_agent;
///
/// #[tokio::main]
/// async fn main() -> anyhow::Result<()> {
///     // Load agent from config
///     let mut agent = get_agent(Some("my-agent-openai"), None).await?;
///     
///     // Chat with the agent
///     let response = agent.chat("Hello!", None).await?;
///     println!("{}", response);
///     
///     Ok(())
/// }
/// ```
pub async fn get_agent(
    agent_name: Option<&str>,
    config_path: Option<&str>,
) -> Result<SekuireAgent> {
    // Load configuration
    let config_file = config_path.unwrap_or("./sekuire.yml");
    let config = load_config(config_file)?;

    // Get agent configuration
    let agent_config = get_agent_config(&config, agent_name)?;

    // Create LLM provider
    let llm_provider = create_llm_provider_from_config(&agent_config.llm).await?;

    // Load system prompt
    let system_prompt = load_system_prompt(&agent_config.system_prompt, Some(&".".to_string()))?;

    // Load tools schema
    let tools_schema = load_tools(&agent_config.tools, Some(&".".to_string()))?;
    let tool_names: Vec<String> = tools_schema.tools.iter().map(|t| t.name.clone()).collect();

    // Get memory config
    let max_history_messages = agent_config.memory.as_ref().and_then(|m| m.max_messages);

    // Load Policy (with caching)
    let policy_path = "policy.json";
    let mut policy_enforcer = None;

    // Check cache
    let mut cached_policy: Option<ActivePolicy> = None;
    if let Ok(cache) = POLICY_CACHE.lock() {
        if let Some((_, policy)) = cache.iter().find(|(k, _)| k == policy_path) {
            cached_policy = Some(policy.clone());
        }
    }

    if let Some(policy) = cached_policy {
        policy_enforcer = Some(PolicyEnforcer::new(policy));
    } else {
        // Load from disk
        if Path::new(policy_path).exists() {
            if let Ok(content) = fs::read_to_string(policy_path) {
                if let Ok(policy) = serde_json::from_str::<ActivePolicy>(&content) {
                    // Update cache
                    if let Ok(mut cache) = POLICY_CACHE.lock() {
                        cache.push((policy_path.to_string(), policy.clone()));
                    }
                    policy_enforcer = Some(PolicyEnforcer::new(policy));
                }
            }
        }
    }

    // Create and return agent
    Ok(SekuireAgent::new(
        agent_config.name,
        llm_provider,
        system_prompt,
        tool_names,
        max_history_messages,
        policy_enforcer,
    ))
}

/// Create LLM provider from configuration
async fn create_llm_provider_from_config(llm_config: &LLMConfig) -> Result<Box<dyn LLMProvider>> {
    // Get API key from environment
    let api_key = std::env::var(&llm_config.api_key_env).with_context(|| {
        format!(
            "API key not found: {}. Please set this environment variable.",
            llm_config.api_key_env
        )
    })?;

    // Create provider
    create_llm_provider(
        &llm_config.provider,
        api_key,
        llm_config.model.clone(),
        llm_config.base_url.clone(),
        llm_config.temperature,
        llm_config.max_tokens,
    )
    .await
}

/// Load all agents from configuration
///
/// # Examples
///
/// ```no_run
/// use sekuire_sdk::get_agents;
///
/// #[tokio::main]
/// async fn main() -> anyhow::Result<()> {
///     let agents = get_agents(None).await?;
///     
///     if let Some(openai) = agents.get("my-agent-openai") {
///         println!("OpenAI agent loaded");
///     }
///     
///     Ok(())
/// }
/// ```
pub async fn get_agents(config_path: Option<&str>) -> Result<HashMap<String, SekuireAgent>> {
    let config_file = config_path.unwrap_or("./sekuire.yml");
    let config = load_config(config_file)?;
    let mut agents_map = HashMap::new();

    if let Some(agents) = &config.agents {
        // Multi-agent configuration
        for agent_name in agents.keys() {
            let agent = get_agent(Some(agent_name), Some(config_file)).await?;
            agents_map.insert(agent_name.clone(), agent);
        }
    } else if config.agent.is_some() {
        // Single agent configuration
        let agent = get_agent(None, Some(config_file)).await?;
        agents_map.insert(config.project.name.clone(), agent);
    }

    Ok(agents_map)
}

/// Get system prompt for a specific agent from configuration
pub async fn get_system_prompt(
    agent_name: Option<&str>,
    config_path: Option<&str>,
) -> Result<String> {
    let config_file = config_path.unwrap_or("./sekuire.yml");
    let config = load_config(config_file)?;
    let agent_config = get_agent_config(&config, agent_name)?;
    let base_path = ".".to_string();
    load_system_prompt(&agent_config.system_prompt, Some(&base_path))
}

/// Get tools for a specific agent from configuration
pub async fn get_tools(agent_name: Option<&str>, config_path: Option<&str>) -> Result<Vec<String>> {
    let config_file = config_path.unwrap_or("./sekuire.yml");
    let config = load_config(config_file)?;
    let agent_config = get_agent_config(&config, agent_name)?;
    let base_path = ".".to_string();
    let tools_schema = load_tools(&agent_config.tools, Some(&base_path))?;
    Ok(tools_schema.tools.iter().map(|t| t.name.clone()).collect())
}