cortex-agent 0.2.1

mod agent;
mod config;
mod memory;
mod messages;
mod plugin;
mod provider;
mod providers;
mod session;
mod tool;
mod tools;
mod ui;

use std::collections::HashMap;
use std::path::Path;

use clap::Parser;

use agent::Agent;
use config::Config;
use memory::store::MemoryStore;

/// Template config written on first run — guides the user through setup.
const FIRST_RUN_CONFIG_TEMPLATE: &str = r##"# Cortex Agent — configuration
# ${ENV_VAR} references are resolved from the environment at runtime.

# --- Active provider & model ---
active_provider: openai     # Change this to your preferred provider
active_model: gpt-4o        # Set the model you want to use

# --- Providers (OpenAI-compatible API) ---
# Each entry needs an api_key and base_url. Keys can be inline or ${ENV_VAR}.
providers:
  openai:
    api_key: ${OPENAI_API_KEY}
    base_url: https://api.openai.com/v1
  openrouter:
    api_key: ${OPENROUTER_API_KEY}
    base_url: https://openrouter.ai/api/v1
  anthropic:
    api_key: ${ANTHROPIC_API_KEY}
    base_url: https://api.anthropic.com/v1
  deepseek:
    api_key: ${DEEPSEEK_API_KEY}
    base_url: https://api.deepseek.com/v1
  groq:
    api_key: ${GROQ_API_KEY}
    base_url: https://api.groq.com/openai/v1
  together:
    api_key: ${TOGETHER_API_KEY}
    base_url: https://api.together.xyz/v1
  mistral:
    api_key: ${MISTRAL_API_KEY}
    base_url: https://api.mistral.ai/v1

# --- Agent Behaviour ---
system_prompt: >
  You are Cortex, an autonomous agent. Answer from your own knowledge first.
  Only use tools when genuinely needed.

  CONCISENESS:
  - Be BRIEF. Short answers, no fluff.
  - Don't introduce yourself unless asked.
  - Let the user drive the conversation.

max_tokens: 4096
max_iterations: 10
temperature: 0.7

# --- Memory ---
memory_enabled: true
memory_dir: ~/.cortex/memory
memory_db: cortex.db

# --- Diagnostics ---
verbose: false
"##;

/// On first run (no config file found), create `~/.cortex/config.yaml` with a
/// commented template and print setup instructions, then exit cleanly.
fn first_run_setup() {
    let home = std::env::var("HOME").unwrap_or_else(|_| ".".into());
    let config_dir = Path::new(&home).join(".cortex");
    let config_path = config_dir.join("config.yaml");

    // Create the .cortex directory
    if let Err(e) = std::fs::create_dir_all(&config_dir) {
        eprintln!("Error: Could not create config directory ~/.cortex: {}", e);
        std::process::exit(1);
    }

    // Write the template config
    if let Err(e) = std::fs::write(&config_path, FIRST_RUN_CONFIG_TEMPLATE) {
        eprintln!("Error: Could not write config file: {}", e);
        std::process::exit(1);
    }

    // Show a friendly setup message on stderr so it doesn't interfere with piping
    eprintln!(
        "\
  ╭─────────────────────────────────────────────────────────────╮
  │  Welcome to Cortex!                                          │
  │                                                              │
  │  A template config has been created at:                      │
  │  ~/.cortex/config.yaml                                       │
  │                                                              │
  │  Next steps:                                                 │
  │  1. Set your API keys via environment variables or edit      │
  │     the config file directly.                                │
  │  2. Run 'cortex' again to start chatting.                    │
  │                                                              │
  │  Supported providers (set the corresponding env var):        │
  │    openai       →  OPENAI_API_KEY                            │
  │    openrouter   →  OPENROUTER_API_KEY                        │
  │    anthropic    →  ANTHROPIC_API_KEY                         │
  │    deepseek     →  DEEPSEEK_API_KEY                          │
  │    groq         →  GROQ_API_KEY                              │
  │    together     →  TOGETHER_API_KEY                          │
  │    mistral      →  MISTRAL_API_KEY                           │
  │                                                              │
  │  Example:                                                    │
  │    export OPENAI_API_KEY=\"sk-...\"                            │
  │    cortex                                                     │
  ╰─────────────────────────────────────────────────────────────╯"
    );
}

/// Cortex Agent — self-learning agent framework with persistent memory.
#[derive(Parser)]
#[command(
    name = "cortex",
    version = "0.2.1",
    about = "Self-learning AI agent with persistent memory, tools, and a beautiful terminal UI",
    after_help = "Examples:\n  cortex  # interactive with memory\n  cortex --one-shot \"Remember I use Vim\"\n  cortex --config ./config.yaml\n  cortex --no-memory"
)]
struct Cli {
    /// Run a single question and exit
    #[arg(long, value_name = "QUESTION")]
    one_shot: Option<String>,

    /// Path to YAML config file
    #[arg(long)]
    config: Option<String>,

    /// Print internal reasoning and tool calls
    #[arg(short, long)]
    verbose: bool,

    /// Disable persistent memory for this session
    #[arg(long)]
    no_memory: bool,
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let args = Cli::parse();

    // ── Load config ──
    let config_path = if let Some(ref p) = args.config {
        Some(p.clone())
    } else {
        let candidates = [
            Path::new("config.yaml").to_path_buf(),
            dirs::home_dir().map(|h| h.join(".cortex/config.yaml")).unwrap_or_default(),
        ];
        candidates.iter().find(|p| p.exists()).map(|p| p.to_string_lossy().to_string())
    };

    // ── First-run setup ──
    // If no config file was found anywhere, create a template and guide the user.
    if config_path.is_none() {
        first_run_setup();
        return Ok(());
    }

    let mut cfg = match config_path {
        Some(ref path) => Config::from_yaml(path).unwrap_or_default(),
        None => Config::default(),
    };

    if args.verbose { cfg.verbose = true; }
    if args.no_memory { cfg.memory_enabled = false; }

    // ── Load session state (last used provider/model) ──
    let session_state = session::load_session_state();
    if !session_state.last_provider.is_empty() && cfg.providers.contains_key(&session_state.last_provider) {
        cfg.active_provider = session_state.last_provider.clone();
        cfg.active_model = session_state.last_model.clone();
        if cfg.verbose {
            eprintln!("[Session: resuming {} / {}]", cfg.active_provider, cfg.active_model);
        }
    }

    // ── Health check ──
    let (provider_name, api_key, base_url) = cfg.get_active_provider_config();
    let model = if cfg.active_model.is_empty() { &cfg.model } else { &cfg.active_model };

    if api_key.is_empty() {
        eprintln!(
            "Error: No API key configured for provider '{}'.\n\
             Set it in config.yaml (~/.cortex/config.yaml) or set the corresponding\n\
             environment variable (e.g. OPENAI_API_KEY for the 'openai' provider).\n\
             \n\
             Available providers: {}",
            provider_name,
            cfg.get_provider_names().join(", ")
        );
        std::process::exit(1);
    }

    if cfg.verbose {
        let (ok, msg) = session::health_check(&base_url, &api_key).await;
        if ok {
            eprintln!("[Health: {} - {}]", provider_name, msg);
        } else {
            eprintln!("[Health: {} - {}]", provider_name, msg);
        }
    }

    let provider = providers::openai_compat::create_provider("openai", model, &api_key, Some(&base_url))
        .map_err(|e| anyhow::anyhow!("Failed to create provider: {}", e))?;

    // ── Build memory store ──
    let memory_store = if cfg.memory_enabled {
        let db_dir = cfg.memory_dir_resolved();
        let db_name = cfg.memory_db_resolved();
        match MemoryStore::new(&db_dir, &db_name) {
            Ok(store) => {
                if cfg.verbose { eprintln!("[Memory store: {}]", store.db_path); }
                Some(store)
            }
            Err(e) => { eprintln!("Warning: Failed to initialize memory store: {}", e); None }
        }
    } else { None };

    // ── Build tools ──
    let memory_arc = memory_store.as_ref().map(|_| {
        std::sync::Arc::new(std::sync::Mutex::new(
            MemoryStore::new(&cfg.memory_dir_resolved(), &cfg.memory_db_resolved()).unwrap()
        ))
    });
    let tools = tools::default_tools(memory_arc);

    if cfg.verbose {
        let tool_names: Vec<&str> = tools.iter().map(|t| t.name.as_str()).collect();
        eprintln!("Config: active_provider={}, model={}, memory_enabled={}", cfg.active_provider, cfg.active_model, cfg.memory_enabled);
        eprintln!("Tools ({}): {:?}", tool_names.len(), tool_names);
    }

    // Build provider map for switching
    let provider_map: HashMap<String, serde_json::Value> = cfg.providers.iter().map(|(k, v)| {
        (k.clone(), serde_json::json!({ "api_key": v.api_key, "base_url": v.base_url }))
    }).collect();
    let provider_names = cfg.get_provider_names();

    // ── Build agent ──
    let mut agent = Agent::new(
        provider,
        tools,
        memory_store,
        &cfg.system_prompt,
        cfg.max_iterations,
        cfg.max_tokens,
        cfg.temperature,
        cfg.verbose,
        provider_map,
        provider_names,
        cfg.active_provider.clone(),
        cfg.active_model.clone(),
    );

    // ── Run ──
    if let Some(question) = args.one_shot {
        let response = agent.run(&question, true).await;
        println!("{}", response);
    } else {
        agent.chat().await;
    }

    Ok(())
}