Feature: Enable programmatically pass in api_key besides reading from env

python/cocoindex/functions.py

@@ -67,6 +67,7 @@ class EmbedText(op.FunctionSpec):
     output_dimension: int | None = None
     task_type: str | None = None
     api_config: llm.VertexAiConfig | None = None
+    api_key: str | None = None
 
 
 class ExtractByLlm(op.FunctionSpec):

python/cocoindex/functions/_engine_builtin_specs.py

@@ -56,6 +56,7 @@ class EmbedText(op.FunctionSpec):
     output_dimension: int | None = None
     task_type: str | None = None
     api_config: llm.VertexAiConfig | None = None
+    api_key: str | None = None
 
 
 class ExtractByLlm(op.FunctionSpec):

python/cocoindex/llm.py

@@ -43,4 +43,5 @@ class LlmSpec:
     api_type: LlmApiType
     model: str
     address: str | None = None
+    api_key: str | None = None
     api_config: VertexAiConfig | OpenAiConfig | None = None

src/llm/anthropic.rs

@@ -14,14 +14,19 @@ pub struct Client {
 }
 
 impl Client {
-    pub async fn new(address: Option<String>) -> Result<Self> {
+    pub async fn new(address: Option<String>, api_key: Option<String>) -> Result<Self> {
         if address.is_some() {
             api_bail!("Anthropic doesn't support custom API address");
         }
-        let api_key = match std::env::var("ANTHROPIC_API_KEY") {
-            Ok(val) => val,
-            Err(_) => api_bail!("ANTHROPIC_API_KEY environment variable must be set"),
+
+        let api_key = if let Some(key) = api_key {
+            key
+        } else {
+            std::env::var("ANTHROPIC_API_KEY").map_err(|_| {
+                anyhow::anyhow!("ANTHROPIC_API_KEY environment variable must be set")
+            })?
         };
+
         Ok(Self {
             api_key,
             client: reqwest::Client::new(),

src/llm/gemini.rs

@@ -34,14 +34,18 @@ pub struct AiStudioClient {
 }
 
 impl AiStudioClient {
-    pub fn new(address: Option<String>) -> Result<Self> {
+    pub fn new(address: Option<String>, api_key: Option<String>) -> Result<Self> {
         if address.is_some() {
             api_bail!("Gemini doesn't support custom API address");
         }
-        let api_key = match std::env::var("GEMINI_API_KEY") {
-            Ok(val) => val,
-            Err(_) => api_bail!("GEMINI_API_KEY environment variable must be set"),
+
+        let api_key = if let Some(key) = api_key {
+            key
+        } else {
+            std::env::var("GEMINI_API_KEY")
+                .map_err(|_| anyhow::anyhow!("GEMINI_API_KEY environment variable must be set"))?
         };
+
         Ok(Self {
             api_key,
             client: reqwest::Client::new(),
@@ -271,6 +275,7 @@ static SHARED_RETRY_THROTTLER: LazyLock<SharedRetryThrottler> =
 impl VertexAiClient {
     pub async fn new(
         address: Option<String>,
+        _api_key: Option<String>,
         api_config: Option<super::LlmApiConfig>,
     ) -> Result<Self> {
         if address.is_some() {

src/llm/litellm.rs

@@ -4,9 +4,14 @@ use async_openai::config::OpenAIConfig;
 pub use super::openai::Client;
 
 impl Client {
-    pub async fn new_litellm(address: Option<String>) -> anyhow::Result<Self> {
+    pub async fn new_litellm(
+        address: Option<String>,
+        api_key: Option<String>,
+    ) -> anyhow::Result<Self> {
         let address = address.unwrap_or_else(|| "http://127.0.0.1:4000".to_string());
-        let api_key = std::env::var("LITELLM_API_KEY").ok();
+
+        let api_key = api_key.or_else(|| std::env::var("LITELLM_API_KEY").ok());
+
         let mut config = OpenAIConfig::new().with_api_base(address);
         if let Some(api_key) = api_key {
             config = config.with_api_key(api_key);

src/llm/mod.rs

@@ -44,6 +44,7 @@ pub struct LlmSpec {
     pub api_type: LlmApiType,
     pub address: Option<String>,
     pub model: String,
+    pub api_key: Option<String>,
     pub api_config: Option<LlmApiConfig>,
 }
 
@@ -117,58 +118,61 @@ mod voyage;
 pub async fn new_llm_generation_client(
     api_type: LlmApiType,
     address: Option<String>,
+    api_key: Option<String>,
     api_config: Option<LlmApiConfig>,
 ) -> Result<Box<dyn LlmGenerationClient>> {
     let client = match api_type {
         LlmApiType::Ollama => {
             Box::new(ollama::Client::new(address).await?) as Box<dyn LlmGenerationClient>
         }
-        LlmApiType::OpenAi => {
-            Box::new(openai::Client::new(address, api_config)?) as Box<dyn LlmGenerationClient>
-        }
-        LlmApiType::Gemini => {
-            Box::new(gemini::AiStudioClient::new(address)?) as Box<dyn LlmGenerationClient>
-        }
-        LlmApiType::VertexAi => Box::new(gemini::VertexAiClient::new(address, api_config).await?)
+        LlmApiType::OpenAi => Box::new(openai::Client::new(address, api_key, api_config)?)
             as Box<dyn LlmGenerationClient>,
-        LlmApiType::Anthropic => {
-            Box::new(anthropic::Client::new(address).await?) as Box<dyn LlmGenerationClient>
+        LlmApiType::Gemini => {
+            Box::new(gemini::AiStudioClient::new(address, api_key)?) as Box<dyn LlmGenerationClient>
         }
-        LlmApiType::LiteLlm => {
-            Box::new(litellm::Client::new_litellm(address).await?) as Box<dyn LlmGenerationClient>
+        LlmApiType::VertexAi => {
+            Box::new(gemini::VertexAiClient::new(address, api_key, api_config).await?)
+                as Box<dyn LlmGenerationClient>
         }
-        LlmApiType::OpenRouter => Box::new(openrouter::Client::new_openrouter(address).await?)
+        LlmApiType::Anthropic => Box::new(anthropic::Client::new(address, api_key).await?)
+            as Box<dyn LlmGenerationClient>,
+        LlmApiType::LiteLlm => Box::new(litellm::Client::new_litellm(address, api_key).await?)
             as Box<dyn LlmGenerationClient>,
+        LlmApiType::OpenRouter => {
+            Box::new(openrouter::Client::new_openrouter(address, api_key).await?)
+                as Box<dyn LlmGenerationClient>
+        }
         LlmApiType::Voyage => {
             api_bail!("Voyage is not supported for generation")
         }
-        LlmApiType::Vllm => {
-            Box::new(vllm::Client::new_vllm(address).await?) as Box<dyn LlmGenerationClient>
-        }
+        LlmApiType::Vllm => Box::new(vllm::Client::new_vllm(address, api_key).await?)
+            as Box<dyn LlmGenerationClient>,
     };
     Ok(client)
 }
 
 pub async fn new_llm_embedding_client(
     api_type: LlmApiType,
     address: Option<String>,
+    api_key: Option<String>,
     api_config: Option<LlmApiConfig>,
 ) -> Result<Box<dyn LlmEmbeddingClient>> {
     let client = match api_type {
         LlmApiType::Ollama => {
             Box::new(ollama::Client::new(address).await?) as Box<dyn LlmEmbeddingClient>
         }
         LlmApiType::Gemini => {
-            Box::new(gemini::AiStudioClient::new(address)?) as Box<dyn LlmEmbeddingClient>
-        }
-        LlmApiType::OpenAi => {
-            Box::new(openai::Client::new(address, api_config)?) as Box<dyn LlmEmbeddingClient>
+            Box::new(gemini::AiStudioClient::new(address, api_key)?) as Box<dyn LlmEmbeddingClient>
         }
+        LlmApiType::OpenAi => Box::new(openai::Client::new(address, api_key, api_config)?)
+            as Box<dyn LlmEmbeddingClient>,
         LlmApiType::Voyage => {
-            Box::new(voyage::Client::new(address)?) as Box<dyn LlmEmbeddingClient>
+            Box::new(voyage::Client::new(address, api_key)?) as Box<dyn LlmEmbeddingClient>
+        }
+        LlmApiType::VertexAi => {
+            Box::new(gemini::VertexAiClient::new(address, api_key, api_config).await?)
+                as Box<dyn LlmEmbeddingClient>
         }
-        LlmApiType::VertexAi => Box::new(gemini::VertexAiClient::new(address, api_config).await?)
-            as Box<dyn LlmEmbeddingClient>,
         LlmApiType::OpenRouter | LlmApiType::LiteLlm | LlmApiType::Vllm | LlmApiType::Anthropic => {
             api_bail!("Embedding is not supported for API type {:?}", api_type)
         }

src/llm/openai.rs

@@ -32,7 +32,11 @@ impl Client {
         Self { client }
     }
 
-    pub fn new(address: Option<String>, api_config: Option<super::LlmApiConfig>) -> Result<Self> {
+    pub fn new(
+        address: Option<String>,
+        api_key: Option<String>,
+        api_config: Option<super::LlmApiConfig>,
+    ) -> Result<Self> {
         let config = match api_config {
             Some(super::LlmApiConfig::OpenAi(config)) => config,
             Some(_) => api_bail!("unexpected config type, expected OpenAiConfig"),
@@ -49,13 +53,16 @@ impl Client {
         if let Some(project_id) = config.project_id {
             openai_config = openai_config.with_project_id(project_id);
         }
-
-        // Verify API key is set
-        if std::env::var("OPENAI_API_KEY").is_err() {
-            api_bail!("OPENAI_API_KEY environment variable must be set");
+        if let Some(key) = api_key {
+            openai_config = openai_config.with_api_key(key);
+        } else {
+            // Verify API key is set in environment if not provided in config
+            if std::env::var("OPENAI_API_KEY").is_err() {
+                api_bail!("OPENAI_API_KEY environment variable must be set");
+            }
         }
+
         Ok(Self {
-            // OpenAI client will use OPENAI_API_KEY and OPENAI_API_BASE env variables by default
             client: OpenAIClient::with_config(openai_config),
         })
     }

src/llm/openrouter.rs

@@ -4,9 +4,14 @@ use async_openai::config::OpenAIConfig;
 pub use super::openai::Client;
 
 impl Client {
-    pub async fn new_openrouter(address: Option<String>) -> anyhow::Result<Self> {
+    pub async fn new_openrouter(
+        address: Option<String>,
+        api_key: Option<String>,
+    ) -> anyhow::Result<Self> {
         let address = address.unwrap_or_else(|| "https://openrouter.ai/api/v1".to_string());
-        let api_key = std::env::var("OPENROUTER_API_KEY").ok();
+
+        let api_key = api_key.or_else(|| std::env::var("OPENROUTER_API_KEY").ok());
+
         let mut config = OpenAIConfig::new().with_api_base(address);
         if let Some(api_key) = api_key {
             config = config.with_api_key(api_key);

src/llm/vllm.rs

@@ -4,9 +4,14 @@ use async_openai::config::OpenAIConfig;
 pub use super::openai::Client;
 
 impl Client {
-    pub async fn new_vllm(address: Option<String>) -> anyhow::Result<Self> {
+    pub async fn new_vllm(
+        address: Option<String>,
+        api_key: Option<String>,
+    ) -> anyhow::Result<Self> {
         let address = address.unwrap_or_else(|| "http://127.0.0.1:8000/v1".to_string());
-        let api_key = std::env::var("VLLM_API_KEY").ok();
+
+        let api_key = api_key.or_else(|| std::env::var("VLLM_API_KEY").ok());
+
         let mut config = OpenAIConfig::new().with_api_base(address);
         if let Some(api_key) = api_key {
             config = config.with_api_key(api_key);

src/llm/voyage.rs

@@ -33,14 +33,18 @@ pub struct Client {
 }
 
 impl Client {
-    pub fn new(address: Option<String>) -> Result<Self> {
+    pub fn new(address: Option<String>, api_key: Option<String>) -> Result<Self> {
         if address.is_some() {
             api_bail!("Voyage AI doesn't support custom API address");
         }
-        let api_key = match std::env::var("VOYAGE_API_KEY") {
-            Ok(val) => val,
-            Err(_) => api_bail!("VOYAGE_API_KEY environment variable must be set"),
+
+        let api_key = if let Some(key) = api_key {
+            key
+        } else {
+            std::env::var("VOYAGE_API_KEY")
+                .map_err(|_| anyhow::anyhow!("VOYAGE_API_KEY environment variable must be set"))?
         };
+
         Ok(Self {
             api_key,
             client: reqwest::Client::new(),

src/ops/functions/embed_text.rs

@@ -13,6 +13,7 @@ struct Spec {
     api_config: Option<LlmApiConfig>,
     output_dimension: Option<u32>,
     task_type: Option<String>,
+    api_key: Option<String>,
 }
 
 struct Args {
@@ -91,9 +92,14 @@ impl SimpleFunctionFactoryBase for Factory {
             .next_arg("text")?
             .expect_type(&ValueType::Basic(BasicValueType::Str))?
             .required()?;
-        let client =
-            new_llm_embedding_client(spec.api_type, spec.address.clone(), spec.api_config.clone())
-                .await?;
+
+        let client = new_llm_embedding_client(
+            spec.api_type,
+            spec.address.clone(),
+            spec.api_key.clone(),
+            spec.api_config.clone(),
+        )
+        .await?;
         let output_dimension = match spec.output_dimension {
             Some(output_dimension) => output_dimension,
             None => {
@@ -144,6 +150,7 @@ mod tests {
             api_config: None,
             output_dimension: None,
             task_type: None,
+            api_key: None,
         };
 
         let factory = Arc::new(Factory);

src/ops/functions/extract_by_llm.rs

@@ -55,6 +55,7 @@ impl Executor {
         let client = new_llm_generation_client(
             spec.llm_spec.api_type,
             spec.llm_spec.address,
+            spec.llm_spec.api_key,
             spec.llm_spec.api_config,
         )
         .await?;
@@ -204,6 +205,7 @@ mod tests {
                 api_type: crate::llm::LlmApiType::OpenAi,
                 model: "gpt-4o".to_string(),
                 address: None,
+                api_key: None,
                 api_config: None,
             },
             output_type: output_type_spec,
@@ -274,6 +276,7 @@ mod tests {
                 api_type: crate::llm::LlmApiType::OpenAi,
                 model: "gpt-4o".to_string(),
                 address: None,
+                api_key: None,
                 api_config: None,
             },
             output_type: make_output_type(BasicValueType::Str),