Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/error/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ pub enum LLMError {
BytesConversion(#[from] FromUtf8Error),
#[error("cli not found: {0}")]
CLINotFound(String),
#[error("empty credential: {0}")]
EmptyCredential(String),
#[error("missing env variable: {0}")]
Env(#[from] VarError),
#[error("failed to prompt: {0}")]
Expand Down
6 changes: 6 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,2 +1,8 @@
mod error;
pub mod llm;

// Re-export constants for convenient access
pub use llm::constants;

// Re-export Model trait and get_available_models function for convenient access
pub use llm::{get_available_models, Model};
11 changes: 8 additions & 3 deletions src/llm/anthropic.rs
Original file line number Diff line number Diff line change
@@ -1,24 +1,29 @@
use crate::error::LLMError;

use super::{Model, ModelFactory};
use super::{Model, ModelFactory, constants::names};

use std::env::var;

const ANTHROPIC_API_KEY: &'static str = "ANTHROPIC_API_KEY";

pub struct Anthropic {
api_key: String,
}

impl ModelFactory for Anthropic {
fn init() -> Result<Self, LLMError> {
let api_key = var("ANTHROPIC_API_KEY")?;
let api_key = var(ANTHROPIC_API_KEY)?.trim().to_string();
if api_key.is_empty() {
return Err(LLMError::EmptyCredential(ANTHROPIC_API_KEY.into()));
}

Ok(Self { api_key })
}
}

impl Model for Anthropic {
fn get_name(&self) -> String {
"Anthropic API".into()
names::ANTHROPIC.into()
}

fn prompt(&self, _: &str) -> Result<String, LLMError> {
Expand Down
6 changes: 3 additions & 3 deletions src/llm/claude.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,10 @@ use std::process::Command;

use crate::{
error::LLMError,
llm::{Model, ModelFactory},
llm::{constants::names, Model, ModelFactory},
};

const CLAUDE_CLI_NAME: &str = "claude";
const CLAUDE_CLI_NAME: &'static str = "claude";

pub struct Claude {}

Expand All @@ -22,7 +22,7 @@ impl ModelFactory for Claude {

impl Model for Claude {
fn get_name(&self) -> String {
"Claude CLI".into()
names::CLAUDE.into()
}

fn prompt(&self, input: &str) -> Result<String, LLMError> {
Expand Down
17 changes: 17 additions & 0 deletions src/llm/constants.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
/// Model name constants that can be used to check against
/// model names in external projects.
///
/// These constants match the values returned by `Model::get_name()`.
pub mod names {
/// The display name for the Cursor model: "Cursor CLI"
pub const CURSOR: &'static str = "Cursor CLI";

/// The display name for the Claude model: "Claude CLI"
pub const CLAUDE: &'static str = "Claude CLI";

/// The display name for the Anthropic model: "Anthropic API"
pub const ANTHROPIC: &'static str = "Anthropic API";

/// The display name for the OpenAI model: "OpenAI API"
pub const OPENAI: &'static str = "OpenAI API";
}
6 changes: 3 additions & 3 deletions src/llm/cursor.rs
Original file line number Diff line number Diff line change
@@ -1,11 +1,11 @@
use crate::{
error::LLMError,
llm::{Model, ModelFactory},
llm::{constants::names, Model, ModelFactory},
};

use std::process::Command;

const CURSOR_CLI_NAME: &str = "cursor-agent";
const CURSOR_CLI_NAME: &'static str = "cursor-agent";

pub struct CursorCLI {}

Expand All @@ -22,7 +22,7 @@ impl ModelFactory for CursorCLI {

impl Model for CursorCLI {
fn get_name(&self) -> String {
"Cursor CLI".into()
names::CURSOR.into()
}

fn prompt(&self, input: &str) -> Result<String, LLMError> {
Expand Down
7 changes: 7 additions & 0 deletions src/llm/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,10 @@ mod claude;
mod cursor;
mod openai;

/// Model name constants that can be used to check against
/// model names in external projects.
pub mod constants;

use crate::{
error::LLMError,
llm::{anthropic::Anthropic, claude::Claude, cursor::CursorCLI, openai::OpenAI},
Expand Down Expand Up @@ -35,6 +39,9 @@ pub trait Model: Send + Sync {

/// Returns the available models in the current
/// system context.
///
/// TODO: ideally this should also include the lifeness check for the found models (e.g. sending a
/// simple prompt).
pub fn get_available_models() -> Result<Vec<Box<dyn Model>>, LLMError> {
let mut models: Vec<Box<dyn Model>> = vec![];

Expand Down
12 changes: 9 additions & 3 deletions src/llm/openai.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,23 +2,29 @@ use std::env::var;

use crate::{
error::LLMError,
llm::{Model, ModelFactory},
llm::{Model, ModelFactory, constants::names},
};

const OPENAI_API_KEY: &'static str = "OPENAI_API_KEY";

pub struct OpenAI {
api_key: String,
}

impl ModelFactory for OpenAI {
fn init() -> Result<Self, LLMError> {
let api_key = var("OPENAI_API_KEY")?;
let api_key = var(OPENAI_API_KEY)?.trim().to_string();
if api_key.is_empty() {
return Err(LLMError::EmptyCredential(OPENAI_API_KEY.into()));
}

Ok(Self { api_key })
}
}

impl Model for OpenAI {
fn get_name(&self) -> String {
"OpenAI API".into()
names::OPENAI.into()
}

fn prompt(&self, _: &str) -> Result<String, LLMError> {
Expand Down
Loading