feat(config): more configurable services

src/config.rs

@@ -2,7 +2,13 @@ use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 use std::fs;
 use std::path::PathBuf;
-use toml::Value;
+
+#[derive(Debug, Deserialize)]
+pub struct ServiceConfig {
+    #[serde(skip_serializing)] // internal use only
+    pub api_key: String,
+    pub url: String,
+}
 
 #[derive(Debug, Deserialize, Serialize)]
 pub struct Prompt {
@@ -32,20 +38,20 @@ fn resolve_config_path() -> PathBuf {
     }
 }
 
-pub fn get_api_key(service: &str) -> String {
+pub fn get_service_config(service: &str) -> ServiceConfig {
     let api_keys_path = resolve_config_path().join(API_KEYS_FILE);
     let content = fs::read_to_string(&api_keys_path)
         .unwrap_or_else(|error| panic!("Could not read file {:?}, {:?}", api_keys_path, error));
-    let value: Value = content.parse().expect("Failed to parse TOML");
 
-    // Extract the API key from the TOML table.
-    let api_key = value
-        .get("API_KEYS")
-        .expect("API_KEYS section not found")
-        .get(service)
-        .unwrap_or_else(|| panic!("No api key found for service {}.", &service));
+    let mut service_configs: HashMap<String, ServiceConfig> = toml::from_str(&content).unwrap();
 
-    api_key.to_string()
+    service_configs.remove(service).unwrap_or_else(|| {
+        panic!(
+            "Prompt {} not found, availables ones are: {:?}",
+            service,
+            service_configs.keys().collect::<Vec<_>>()
+        )
+    })
 }
 
 pub fn get_prompts() -> HashMap<String, Prompt> {

src/input_processing.rs

@@ -1,4 +1,4 @@
-use crate::config::{get_api_key, Prompt, PLACEHOLDER_TOKEN};
+use crate::config::{get_service_config, Prompt, PLACEHOLDER_TOKEN};
 use crate::request::{make_authenticated_request, OpenAiResponse};
 use std::io::{Read, Result, Write};
 
@@ -51,8 +51,8 @@ pub fn process_input_with_request<R: Read, W: Write>(
     for message in prompt.messages.iter_mut() {
         message.content = message.content.replace(PLACEHOLDER_TOKEN, &input)
     }
-    let api_key = get_api_key(&prompt.service);
-    let response: OpenAiResponse = make_authenticated_request(&api_key, prompt)
+    let service_config = get_service_config(&prompt.service);
+    let response: OpenAiResponse = make_authenticated_request(service_config, prompt)
         .unwrap()
         .into_json()?;
 

src/request.rs

@@ -1,5 +1,7 @@
 use serde::{Deserialize, Serialize};
 
+use crate::config::ServiceConfig;
+
 #[derive(Debug, Deserialize)]
 pub struct Message {
     pub role: String,
@@ -32,13 +34,16 @@ pub struct OpenAiResponse {
 }
 
 pub fn make_authenticated_request(
-    api_key: &str,
+    service_config: ServiceConfig,
     data: impl Serialize,
 ) -> Result<ureq::Response, ureq::Error> {
-    println!("Trying to reach openai with {}", &api_key);
-    ureq::post("https://api.openai.com/v1/chat/completions")
+    println!("Trying to reach openai with {}", service_config.api_key);
+    ureq::post(&service_config.url)
         .set("Content-Type", "application/json")
-        .set("Authorization", &format!("Bearer {}", api_key))
+        .set(
+            "Authorization",
+            &format!("Bearer {}", service_config.api_key),
+        )
         .send_json(data)
     //     .send_json(ureq::json!(
     //         {