prime 0.2.10

A CLI tool for interacting with LLMs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118

//! Entry point for Prime CLI
//! v0.2.5: Enhanced with streaming responses and rich display

mod commands;
mod config;
mod console;
mod memory;
mod session;
mod parser;
mod streaming;
mod display;

use std::env;
use std::process;

use anyhow::{Context as AnyhowContext, Result};
use crossterm::style::Stylize;
use llm::builder::{LLMBackend, LLMBuilder};
use session::PrimeSession;
use crate::config::Config;

pub const APP_NAME: &str = "prime";

#[tokio::main]
async fn main() -> Result<()> {
    console::display_banner();

    let config = match config::load_config() {
        Ok(cfg) => cfg,
        Err(e) => {
            eprintln!("{}", format!("[ERROR] Failed to load configuration: {}", e).red());
            process::exit(1);
        }
    };

    let session = match init_session(config).await {
        Ok(session) => session,
        Err(e) => {
            eprintln!("{}", format!("[ERROR] Initialization error: {}", e).red());
            process::exit(1);
        }
    };

    if let Err(e) = console::run_repl(session).await {
        eprintln!("{}", format!("[ERROR] Session ended with an error: {}", e).red());
        process::exit(1);
    }

    Ok(())
}

async fn init_session(config: Config) -> Result<PrimeSession> {
    let provider = env::var("LLM_PROVIDER").unwrap_or(config.provider);
    let model_from_env = env::var("LLM_MODEL").ok();
    
    let model = model_from_env.or(config.model).unwrap_or_else(|| {
        match provider.as_str() {
            "google" => "gemini-2.5-flash-lite".to_string(),
            "ollama" => "gemma2".to_string(),
            _ => "gemma2".to_string(),
        }
    });

    let temperature = env::var("LLM_TEMPERATURE")
        .ok()
        .and_then(|s| s.parse::<f32>().ok())
        .unwrap_or(config.temperature);
        
    let max_tokens = env::var("LLM_MAX_TOKENS")
        .ok()
        .and_then(|s| s.parse::<u32>().ok())
        .unwrap_or(config.max_tokens);

    let prime_config_base_dir = dirs::home_dir()
        .ok_or_else(|| anyhow::anyhow!("Could not determine home directory"))?
        .join(".prime");

    let workspace_dir = env::current_dir().context("Failed to get current working directory")?;

    let (llm, provider_name) = match provider.as_str() {
        "google" => {
            let api_key = env::var("GEMINI_API_KEY").unwrap_or(config.gemini_api_key);
            if api_key.is_empty() {
                return Err(anyhow::anyhow!("GEMINI_API_KEY not set in environment or config.toml. Please get a key from Google AI Studio."));
            }
            let llm = LLMBuilder::new()
                .backend(LLMBackend::Google)
                .api_key(api_key)
                .model(model.clone())
                .max_tokens(max_tokens)
                .temperature(temperature)
                .build()
                .context("Failed to build LLM provider (Google)")?;
            (llm, "Google AI Platform")
        },
        "ollama" => {
            let api_key = env::var("OLLAMA_API_KEY").unwrap_or(config.ollama_api_key);
            let llm = LLMBuilder::new()
                .backend(LLMBackend::Ollama)
                .api_key(api_key)
                .model(model.clone())
                .max_tokens(max_tokens)
                .temperature(temperature)
                .build()
                .context("Failed to build LLM provider (Ollama)")?;
            (llm, "Ollama")
        },
        _ => {
            return Err(anyhow::anyhow!("Unsupported LLM provider: {}", provider));
        }
    };

    console::display_init_info(&model, provider_name, &prime_config_base_dir, &workspace_dir);

    let session = PrimeSession::new(prime_config_base_dir, llm)?;

    Ok(session)
}