polish API

Signed-off-by: kerthcet <kerthcet@gmail.com>

Author: kerthcet

.env.integration-test

@@ -1,3 +1,3 @@
 AMRS_API_KEY=your_amrs_api_key_here
 OPENAI_API_KEY=your_openai_api_key_here
-FAKE_API_KEY=your_fake_api_key_here
+FAKER_API_KEY=your_faker_api_key_here

Cargo.lock

@@ -12,4 +12,4 @@ lazy_static = "1.5.0"
 rand = "0.9.2"
 reqwest = "0.12.26"
 serde = "1.0.228"
-tokio = "1.48.0"
+tokio = { version = "1.48.0", features = ["full"] }

Cargo.toml

@@ -11,7 +11,7 @@ Thanks to [async-openai](https://github.com/64bit/async-openai), AMRS builds on
 - Flexible routing strategies, including:
   - **Random**: Randomly selects a model from the available models.
   - **WRR**: Weighted Round Robin selects models based on predefined weights.
-  - **UCB**: Upper Confidence Bound based model selection (coming soon).
+  - **UCB1**: Upper Confidence Bound based model selection (coming soon).
   - **Adaptive**: Dynamically selects models based on performance metrics (coming soon).
 
 - Broad provider support:
@@ -27,20 +27,21 @@ Here's a simple example with the Weighted Round Robin (WRR) routing mode:
 // Before running the code, make sure to set your OpenAI API key in the environment variable:
 // export OPENAI_API_KEY="your_openai_api_key"
 
-use arms::{Client, Config, ModelConfig, CreateResponseArgs, RoutingMode};
+use arms::client;
+use arms::types::responses;
 
-let config = Config::builder()
+let config = client::Config::builder()
     .provider("openai")
-    .routing_mode(RoutingMode::WRR)
+    .routing_mode(client::RoutingMode::WRR)
     .model(
-        ModelConfig::builder()
+        client::ModelConfig::builder()
             .id("gpt-3.5-turbo")
             .weight(2)
             .build()
             .unwrap(),
     )
     .model(
-        ModelConfig::builder()
+        client::ModelConfig::builder()
             .id("gpt-4")
             .weight(1)
             .build()
@@ -49,8 +50,8 @@ let config = Config::builder()
     .build()
     .unwrap();
 
-let mut client = Client::new(config);
-let request = CreateResponseArgs::default()
+let mut client = client::Client::new(config);
+let request = responses::CreateResponseArgs::default()
     .input("give me a poem about nature")
     .build()
     .unwrap();

README.md

@@ -5,7 +5,7 @@ use derive_builder::Builder;
 use lazy_static::lazy_static;
 
 // ------------------ Provider ------------------
-const OPENAI_PROVIDER: &str = "OPENAI";
+pub const DEFAULT_PROVIDER: &str = "OPENAI";
 
 lazy_static! {
     pub static ref PROVIDER_BASE_URLS: HashMap<&'static str, &'static str> = {
@@ -14,7 +14,7 @@ lazy_static! {
         m.insert("DEEPINFRA", "https://api.deepinfra.com/v1/openai");
         m.insert("OPENROUTER", "https://openrouter.ai/api/v1");
 
-        m.insert("FAKE", "http://localhost:8080"); // test only
+        m.insert("FAKER", "http://localhost:8080"); // test only
         // TODO: support more providers here...
         m
     };
@@ -82,9 +82,9 @@ impl ModelConfig {
 #[builder(build_fn(validate = "Self::validate"), pattern = "mutable")]
 pub struct Config {
     // global configs for models, will be overridden by model-specific configs
-    #[builder(default = "https://api.openai.com/v1".to_string())]
-    pub(crate) base_url: String,
-    #[builder(default = "OPENAI_PROVIDER.to_string()", setter(custom))]
+    #[builder(default=None, setter(custom))]
+    pub(crate) base_url: Option<String>,
+    #[builder(default = "DEFAULT_PROVIDER.to_string()", setter(custom))]
     pub(crate) provider: String,
     #[builder(default = "0.8")]
     pub(crate) temperature: f32,
@@ -104,22 +104,34 @@ impl Config {
 
     // populate will fill in the missing model-specific configs with global configs.
     pub fn populate(&mut self) -> &mut Self {
+        let global_base_url = match self.base_url.is_some() {
+            true => self.base_url.clone(),
+            false => Some(
+                PROVIDER_BASE_URLS
+                    .get(self.provider.as_str())
+                    .unwrap()
+                    .to_string(),
+            ),
+        };
+
         for model in &mut self.models {
-            let model_url_exist = model.base_url.is_some();
+            if model.base_url.is_none() {
+                if model.provider.is_some() {
+                    model.base_url = Some(
+                        PROVIDER_BASE_URLS
+                            .get(model.provider.as_ref().unwrap().as_str())
+                            .unwrap()
+                            .to_string(),
+                    );
+                } else {
+                    model.base_url = global_base_url.clone();
+                }
+            }
 
             if model.provider.is_none() {
                 model.provider = Some(self.provider.clone());
             }
 
-            if !model_url_exist
-                && PROVIDER_BASE_URLS.contains_key(model.provider.as_ref().unwrap().as_str())
-            {
-                model.base_url =
-                    Some(PROVIDER_BASE_URLS[model.provider.as_ref().unwrap().as_str()].to_string());
-            }
-            if !model_url_exist {
-                model.base_url = Some(self.base_url.clone());
-            }
             if model.temperature.is_none() {
                 model.temperature = Some(self.temperature);
             }
@@ -132,6 +144,10 @@ impl Config {
 }
 
 impl ConfigBuilder {
+    pub fn base_url<S: AsRef<str>>(&mut self, url: S) -> &mut Self {
+        self.base_url = Some(Some(url.as_ref().to_string()));
+        self
+    }
     pub fn model(&mut self, model: ModelConfig) -> &mut Self {
         let mut models = self.models.clone().unwrap_or_default();
         models.push(model);
@@ -180,7 +196,7 @@ impl ConfigBuilder {
 
             // check the existence of API key in environment variables
             if let Some(provider) = &model.provider {
-                let env_var = format!("{}_API_KEY", provider.to_uppercase());
+                let env_var = format!("{}_API_KEY", provider);
                 if env::var(&env_var).is_err() {
                     return Err(format!(
                         "API key for provider '{}' not found in environment variable '{}'",
@@ -194,16 +210,15 @@ impl ConfigBuilder {
                     "{}_API_KEY",
                     self.provider
                         .as_ref()
-                        .unwrap_or(&OPENAI_PROVIDER.to_string())
+                        .unwrap_or(&DEFAULT_PROVIDER.to_string())
                         .to_uppercase()
                 );
                 if env::var(&env_var).is_err() {
                     return Err(format!(
                         "API key for provider '{}' not found in environment variable '{}'",
                         self.provider
                             .as_ref()
-                            .unwrap_or(&OPENAI_PROVIDER.to_string())
-                            .to_uppercase(),
+                            .unwrap_or(&DEFAULT_PROVIDER.to_string()),
                         env_var
                     ));
                 }
@@ -234,10 +249,8 @@ mod tests {
             )
             .build();
         assert!(valid_simplest_models_cfg.is_ok());
-        assert!(valid_simplest_models_cfg.as_ref().unwrap().provider == OPENAI_PROVIDER);
-        assert!(
-            valid_simplest_models_cfg.as_ref().unwrap().base_url == "https://api.openai.com/v1"
-        );
+        assert!(valid_simplest_models_cfg.as_ref().unwrap().provider == DEFAULT_PROVIDER);
+        assert!(valid_simplest_models_cfg.as_ref().unwrap().base_url == None);
         assert!(valid_simplest_models_cfg.as_ref().unwrap().temperature == 0.8);
         assert!(
             valid_simplest_models_cfg
@@ -285,7 +298,7 @@ mod tests {
         // case 4:
         // AMRS_API_KEY is set in .env.test already.
         let valid_cfg_with_customized_provider = Config::builder()
-            .base_url("http://example.ai".to_string())
+            .base_url("http://example.ai")
             .max_output_tokens(2048)
             .model(
                 ModelConfig::builder()

src/client/config.rs

@@ -8,7 +8,7 @@ mod router {
 }
 
 mod provider {
-    mod fake;
+    mod faker;
     mod openai;
     pub mod provider;
 }

src/lib.rs

@@ -1,3 +1,52 @@
+use tokio::runtime::Runtime;
+
+use arms::client;
+use arms::types::responses;
+
 fn main() {
-    println!("Hello, world!");
+    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);
+        }
+    }
 }

src/main.rs

@@ -1,18 +1,18 @@
 use async_trait::async_trait;
 
 use crate::client::config::{ModelConfig, ModelName};
-use crate::provider::provider::{Provider, validate_request};
+use crate::provider::provider;
 use crate::types::error::OpenAIError;
 use crate::types::responses::{
     AssistantRole, CreateResponse, OutputItem, OutputMessage, OutputMessageContent, OutputStatus,
     OutputTextContent, Response, Status,
 };
 
-pub struct FakeProvider {
+pub struct FakerProvider {
     model: ModelName,
 }
 
-impl FakeProvider {
+impl FakerProvider {
     pub fn new(config: ModelConfig) -> Self {
         Self {
             model: config.name.clone(),
@@ -21,13 +21,13 @@ impl FakeProvider {
 }
 
 #[async_trait]
-impl Provider for FakeProvider {
+impl provider::Provider for FakerProvider {
     fn name(&self) -> &'static str {
         "FakeProvider"
     }
 
     async fn create_response(&self, request: CreateResponse) -> Result<Response, OpenAIError> {
-        validate_request(&request)?;
+        provider::validate_responses_request(&request)?;
 
         Ok(Response {
             id: "fake-response-id".to_string(),