diff --git a/Cargo.toml b/Cargo.toml index f8526a6..8150e66 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,6 +3,13 @@ name = "arms" version = "0.0.1" edition = "2024" +description = "🧬 The adaptive model routing system for exploration and exploitation." +license = "MIT" +repository = "https://github.com/InftyAI/AMRS" +readme = "README.md" +keywords = ["llmops", "ai-gateway", "ai-proxy", "sdk", "llm"] +categories = ["router"] + [dependencies] async-openai = { version = "0.31.1", features = ["_api", "response-types", "responses", "chat-completion"] } async-trait = "0.1.89" diff --git a/README.md b/README.md index e663b8d..52ac385 100644 --- a/README.md +++ b/README.md @@ -23,7 +23,17 @@ AMRS builds on top of [async-openai](https://github.com/64bit/async-openai) to p - OpenAI compatible providers (OpenAI, DeepInfra, etc.) - More on the way -## How to use +## How to Install + +Run the following Cargo command in your project directory: + +`cargo add arms` + +Or add the following line to your Cargo.toml: + +`arms = "0.0.1"` + +## How to Use Here's a simple example with the Weighted Round Robin (WRR) routing mode. Before running the code, make sure to set your provider API key in the environment variable by running `export _API_KEY="your_openai_api_key"`. Here we use OpenAI as an example. @@ -32,54 +42,60 @@ Here we use OpenAI as an example. ```rust # Make sure OPENAI_API_KEY is set in your environment variables before running this code. -use tokio::runtime::Runtime; use arms::client; use arms::types::chat; +use tokio::runtime::Runtime; -let config = client::Config::builder() - .provider("openai") - .routing_mode(client::RoutingMode::WRR) - .model( - client::ModelConfig::builder() - .name("gpt-3.5-turbo") - .weight(2) - .build() - .unwrap(), - ) - .model( - client::ModelConfig::builder() - .name("gpt-4") - .weight(1) - .build() - .unwrap(), - ) - .build() - .unwrap(); - -let mut client = client::Client::new(config); -let request = chat::CreateChatCompletionRequestArgs::default() - .messages([ - chat::ChatCompletionRequestSystemMessage::from("You are a helpful assistant.").into(), - chat::ChatCompletionRequestUserMessage::from("Who won the FIFA World Cup in 2025?").into(), - ]) - .build() - .unwrap(); - -let result = Runtime::new().unwrap().block_on(client.create_completion(request)); -match result { - Ok(response) => { - for choice in response.choices { - println!("Response: {:?}", choice.message.content); +fn main() { + let config = client::Config::builder() + .provider("deepinfra") + .routing_mode(client::RoutingMode::WRR) + .model( + client::ModelConfig::builder() + .name("deepseek-ai/DeepSeek-V3.2") + .weight(2) + .build() + .unwrap(), + ) + .model( + client::ModelConfig::builder() + .name("nvidia/Nemotron-3-Nano-30B-A3B") + .weight(1) + .build() + .unwrap(), + ) + .build() + .unwrap(); + + let mut client = client::Client::new(config); + let request = chat::CreateChatCompletionRequestArgs::default() + .messages([ + chat::ChatCompletionRequestSystemMessage::from("You are a helpful assistant.").into(), + chat::ChatCompletionRequestUserMessage::from("How long it takes to learn Rust?").into(), + ]) + .build() + .unwrap(); + + let result = Runtime::new() + .unwrap() + .block_on(client.create_completion(request)); + match result { + Ok(response) => { + for choice in response.choices { + println!("Response: {:?}", choice.message.content); + } + } + Err(e) => { + eprintln!("Error: {}", e); } - } - Err(e) => { - eprintln!("Error: {}", e); } } ``` +See more examples [here](/examples) folder. + ## Contributing 🚀 All kinds of contributions are welcomed ! Please follow [Contributing](/CONTRIBUTING.md). -[![Star History Chart](https://api.star-history.com/svg?repos=inftyai/amrs&type=Date)](https://www.star-history.com/#inftyai/amrs&Date) \ No newline at end of file +[![Star History Chart](https://api.star-history.com/svg?repos=inftyai/amrs&type=Date)](https://www.star-history.com/#inftyai/amrs&Date) diff --git a/examples/wrr.rs b/examples/wrr.rs new file mode 100644 index 0000000..6050b8f --- /dev/null +++ b/examples/wrr.rs @@ -0,0 +1,48 @@ +use arms::client; +use arms::types::chat; +use tokio::runtime::Runtime; + +fn main() { + let config = client::Config::builder() + .provider("deepinfra") + .routing_mode(client::RoutingMode::WRR) + .model( + client::ModelConfig::builder() + .name("deepseek-ai/DeepSeek-V3.2") + .weight(2) + .build() + .unwrap(), + ) + .model( + client::ModelConfig::builder() + .name("nvidia/Nemotron-3-Nano-30B-A3B") + .weight(1) + .build() + .unwrap(), + ) + .build() + .unwrap(); + + let mut client = client::Client::new(config); + let request = chat::CreateChatCompletionRequestArgs::default() + .messages([ + chat::ChatCompletionRequestSystemMessage::from("You are a helpful assistant.").into(), + chat::ChatCompletionRequestUserMessage::from("How long it takes to learn Rust?").into(), + ]) + .build() + .unwrap(); + + let result = Runtime::new() + .unwrap() + .block_on(client.create_completion(request)); + match result { + Ok(response) => { + for choice in response.choices { + println!("Response: {:?}", choice.message.content); + } + } + Err(e) => { + eprintln!("Error: {}", e); + } + } +} diff --git a/src/main.rs b/src/main.rs deleted file mode 100644 index 8ab0e53..0000000 --- a/src/main.rs +++ /dev/null @@ -1,105 +0,0 @@ -use tokio::runtime::Runtime; - -use arms::client; -use arms::types::chat; - -fn main() { - // case 1: completion with DeepInfra provider. - let config = client::Config::builder() - .provider("deepinfra") - .routing_mode(client::RoutingMode::WRR) - .model( - client::ModelConfig::builder() - .name("nvidia/Nemotron-3-Nano-30B-A3B") - .weight(1) - .build() - .unwrap(), - ) - .model( - client::ModelConfig::builder() - .name("deepseek-ai/DeepSeek-V3.2") - .weight(2) - .build() - .unwrap(), - ) - .build() - .unwrap(); - - let mut client = client::Client::new(config); - - let request = chat::CreateChatCompletionRequestArgs::default() - .messages([ - chat::ChatCompletionRequestSystemMessage::from("You are a helpful assistant.").into(), - chat::ChatCompletionRequestUserMessage::from("Who won the world series in 2020?") - .into(), - chat::ChatCompletionRequestAssistantMessage::from( - "The Los Angeles Dodgers won the World Series in 2020.", - ) - .into(), - chat::ChatCompletionRequestUserMessage::from("Where was it played?").into(), - ]) - .build() - .unwrap(); - - let result = Runtime::new() - .unwrap() - .block_on(client.create_completion(request)); - - match result { - Ok(response) => { - for choice in response.choices { - println!("Response: {:?}", choice.message.content); - } - } - Err(e) => { - eprintln!("Error: {}", e); - } - } - - // case 2: response with DeepInfra provider. - // let config = client::Config::builder() - // .provider("deepinfra") - // .routing_mode(client::RoutingMode::WRR) - // .model( - // client::ModelConfig::builder() - // .name("nvidia/Nemotron-3-Nano-30B-A3B") - // .weight(1) - // .build() - // .unwrap(), - // ) - // .model( - // client::ModelConfig::builder() - // .name("deepseek-ai/DeepSeek-V3.2") - // .weight(2) - // .build() - // .unwrap(), - // ) - // .build() - // .unwrap(); - - // let mut client = client::Client::new(config); - - // let request = responses::CreateResponseArgs::default() - // .input(responses::InputParam::Items(vec![ - // responses::InputItem::EasyMessage(responses::EasyInputMessage { - // r#type: responses::MessageType::Message, - // role: responses::Role::User, - // content: responses::EasyInputContent::Text("What is AGI?".to_string()), - // }), - // ])) - // .build() - // .unwrap(); - - // let result = Runtime::new() - // .unwrap() - // .block_on(client.create_response(request)); - - // match result { - // Ok(response) => { - // println!("Response ID: {}", response.id); - // } - // Err(e) => { - // eprintln!("Error: {}", e); - // } - // } -}