gitdataai/lib/ai/agent/config.rs

use crate::error::{AiError, AiResult};

pub const DEFAULT_SYSTEM_PROMPT: &str = r#"You are a precise autonomous agent that executes tasks through tool calls.

## Core Principles
- Use tools when they can materially improve correctness or efficiency
- After each action, verify results and adjust approach if needed
- Keep reasoning concise and focus on actionable outcomes
- Return only the final useful answer to the user

## Workflow
1. Analyze the request and plan your approach
2. Execute actions using appropriate tools
3. Review observations and verify assumptions
4. Iterate until the task is complete
5. Provide a clear, concise final response

## Title Generation
If this is the first user message in a new conversation with a default title, you SHOULD call `set_conversation_title` as your first action to create a short, descriptive title (max 100 chars). This helps keep the conversation organized. Only do this once at the very beginning."#;

#[derive(Clone, Debug)]
pub struct AgentConfig {
    pub model: String,
    pub provider: String,
    pub api_mode: String,
    pub system_prompt: String,
    pub max_iterations: usize,
    pub iteration_budget: usize,
    pub temperature: Option<f64>,
    pub max_completion_tokens: Option<u64>,
    pub max_total_tokens_per_run: Option<i64>,
    pub enabled_toolsets: Vec<String>,
    pub disabled_toolsets: Vec<String>,
    pub allowed_tools: Vec<String>,
    pub denied_tools: Vec<String>,
    pub retry_max_attempts: usize,
    pub retry_base_delay_ms: u64,
    pub retry_jitter: bool,
    pub fallback_model: Option<String>,
    pub skip_memory: bool,
    pub skip_context_files: bool,
    pub skip_compression: bool,
    pub quiet_mode: bool,
    pub save_trajectories: bool,
    pub reasoning_effort: Option<String>,
    pub service_tier: Option<String>,
    pub platform: Option<String>,
    pub session_id: Option<uuid::Uuid>,
}

impl AgentConfig {
    pub fn new(model: impl Into<String>) -> AiResult<Self> {
        let config = Self {
            model: model.into(),
            provider: String::new(),
            api_mode: String::from("chat_completions"),
            system_prompt: DEFAULT_SYSTEM_PROMPT.to_string(),
            max_iterations: 64,
            iteration_budget: 90,
            temperature: Some(0.2),
            max_completion_tokens: None,
            max_total_tokens_per_run: Some(128_000),
            enabled_toolsets: Vec::new(),
            disabled_toolsets: Vec::new(),
            allowed_tools: Vec::new(),
            denied_tools: Vec::new(),
            retry_max_attempts: 3,
            retry_base_delay_ms: 1_000,
            retry_jitter: true,
            fallback_model: None,
            skip_memory: false,
            skip_context_files: false,
            skip_compression: false,
            quiet_mode: false,
            save_trajectories: false,
            reasoning_effort: None,
            service_tier: None,
            platform: None,
            session_id: None,
        };
        config.validate()?;
        Ok(config)
    }

    pub fn validate(&self) -> AiResult<()> {
        if self.model.trim().is_empty() {
            return Err(AiError::Config("agent model is required".to_string()));
        }
        if self.max_iterations == 0 {
            return Err(AiError::Config(
                "agent max_iterations must be greater than 0".to_string(),
            ));
        }
        if let Some(tokens) = self.max_total_tokens_per_run
            && tokens <= 0
        {
            return Err(AiError::Config(
                "agent max_total_tokens_per_run must be > 0".to_string(),
            ));
        }
        Ok(())
    }

    pub fn with_provider(mut self, provider: impl Into<String>) -> Self {
        self.provider = provider.into();
        self
    }

    pub fn with_api_mode(mut self, mode: impl Into<String>) -> Self {
        self.api_mode = mode.into();
        self
    }

    pub fn with_max_iterations(mut self, max: usize) -> Self {
        self.max_iterations = max;
        self.iteration_budget = self.iteration_budget.max(max);
        self
    }

    pub fn with_iteration_budget(mut self, budget: usize) -> Self {
        self.iteration_budget = budget;
        self
    }

    pub fn with_system_prompt(mut self, prompt: impl Into<String>) -> Self {
        self.system_prompt = prompt.into();
        self
    }

    pub fn with_temperature(mut self, temperature: Option<f64>) -> Self {
        self.temperature = temperature;
        self
    }

    pub fn with_max_completion_tokens(
        mut self,
        max_completion_tokens: Option<u64>,
    ) -> Self {
        self.max_completion_tokens = max_completion_tokens;
        self
    }

    pub fn with_max_total_tokens(mut self, limit: Option<i64>) -> Self {
        self.max_total_tokens_per_run = limit;
        self
    }

    pub fn with_toolset_policy(
        mut self,
        enabled: Vec<String>,
        disabled: Vec<String>,
    ) -> Self {
        self.enabled_toolsets = enabled;
        self.disabled_toolsets = disabled;
        self
    }

    pub fn with_tool_policy(
        mut self,
        allowed_tools: Vec<String>,
        denied_tools: Vec<String>,
    ) -> Self {
        self.allowed_tools = allowed_tools;
        self.denied_tools = denied_tools;
        self
    }

    pub fn with_retry(
        mut self,
        max_attempts: usize,
        base_delay_ms: u64,
    ) -> Self {
        self.retry_max_attempts = max_attempts;
        self.retry_base_delay_ms = base_delay_ms;
        self
    }

    pub fn with_retry_jitter(mut self, jitter: bool) -> Self {
        self.retry_jitter = jitter;
        self
    }

    pub fn with_fallback_model(
        mut self,
        fallback_model: impl Into<String>,
    ) -> Self {
        self.fallback_model = Some(fallback_model.into());
        self
    }

    pub fn with_skip_memory(mut self, skip: bool) -> Self {
        self.skip_memory = skip;
        self
    }

    pub fn with_skip_compression(mut self, skip: bool) -> Self {
        self.skip_compression = skip;
        self
    }

    pub fn with_quiet_mode(mut self, quiet: bool) -> Self {
        self.quiet_mode = quiet;
        self
    }

    pub fn with_platform(mut self, platform: impl Into<String>) -> Self {
        self.platform = Some(platform.into());
        self
    }

    pub fn with_session_id(mut self, session_id: uuid::Uuid) -> Self {
        self.session_id = Some(session_id);
        self
    }

    pub fn with_reasoning_effort(mut self, effort: impl Into<String>) -> Self {
        self.reasoning_effort = Some(effort.into());
        self
    }

    pub fn is_tool_exposed(&self, name: &str) -> bool {
        let denied = self.denied_tools.iter().any(|tool| tool == name);
        if denied {
            return false;
        }
        if self.allowed_tools.is_empty() {
            return true;
        }
        self.allowed_tools.iter().any(|tool| tool == name)
    }
}

pub fn default_system_prompt() -> &'static str {
    DEFAULT_SYSTEM_PROMPT
}