Feature: Enable programmatically pass in api_key besides reading from env

docs/docs/ai/llm.mdx

@@ -234,6 +234,12 @@ To use the Vertex AI API:
     gcloud auth application-default login
     ```
 
+:::note API Key Not Supported
+
+Vertex AI does **not** support the `api_key` parameter. Vertex AI uses Application Default Credentials (ADC) for authentication instead of API keys. If you provide an `api_key` parameter when using `LlmApiType.VERTEX_AI`, an error will be raised.
+
+:::
+
 Spec for Vertex AI takes additional `api_config` field, in type `cocoindex.llm.VertexAiConfig` with the following fields:
 -   `project` (type: `str`, required): The project ID of the Google Cloud project.
 -   `region` (type: `str`, optional): The region of the Google Cloud project. Use `global` if not specified.

python/cocoindex/functions/_engine_builtin_specs.py

@@ -4,6 +4,7 @@
 from typing import Literal
 
 from .. import llm, op
+from ..auth_registry import TransientAuthEntryReference
 
 
 class ParseJson(op.FunctionSpec):
@@ -56,6 +57,7 @@ class EmbedText(op.FunctionSpec):
     output_dimension: int | None = None
     task_type: str | None = None
     api_config: llm.VertexAiConfig | None = None
+    api_key: TransientAuthEntryReference[str] | None = None
 
 
 class ExtractByLlm(op.FunctionSpec):

python/cocoindex/llm.py

@@ -1,6 +1,8 @@
 from dataclasses import dataclass
 from enum import Enum
 
+from .auth_registry import TransientAuthEntryReference
+
 
 class LlmApiType(Enum):
     """The type of LLM API to use."""
@@ -44,4 +46,5 @@ class LlmSpec:
     api_type: LlmApiType
     model: str
     address: str | None = None
+    api_key: TransientAuthEntryReference[str] | None = None
     api_config: VertexAiConfig | OpenAiConfig | None = None

rust/cocoindex/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(),

rust/cocoindex/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,11 +275,15 @@ 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() {
             api_bail!("VertexAi API address is not supported for VertexAi API type");
         }
+        if api_key.is_some() {
+            api_bail!("VertexAi API key is not supported for VertexAi API type. Vertex AI uses Application Default Credentials (ADC) for authentication. Please set up ADC using 'gcloud auth application-default login' instead.");
+        }
         let Some(super::LlmApiConfig::VertexAi(config)) = api_config else {
             api_bail!("VertexAi API config is required for VertexAi API type");
         };

rust/cocoindex/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);

rust/cocoindex/src/llm/mod.rs

@@ -45,6 +45,7 @@ pub struct LlmSpec {
     pub api_type: LlmApiType,
     pub address: Option<String>,
     pub model: String,
+    pub api_key: Option<spec::AuthEntryReference<String>>,
     pub api_config: Option<LlmApiConfig>,
 }
 
@@ -119,61 +120,64 @@ 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::OpenAi => Box::new(openai::Client::new(address, api_key, api_config)?)
+            as Box<dyn LlmGenerationClient>,
         LlmApiType::Gemini => {
-            Box::new(gemini::AiStudioClient::new(address)?) as Box<dyn LlmGenerationClient>
+            Box::new(gemini::AiStudioClient::new(address, api_key)?) as Box<dyn LlmGenerationClient>
         }
-        LlmApiType::VertexAi => Box::new(gemini::VertexAiClient::new(address, api_config).await?)
-            as Box<dyn LlmGenerationClient>,
-        LlmApiType::Anthropic => {
-            Box::new(anthropic::Client::new(address).await?) as Box<dyn LlmGenerationClient>
+        LlmApiType::VertexAi => {
+            Box::new(gemini::VertexAiClient::new(address, api_key, api_config).await?)
+                as Box<dyn LlmGenerationClient>
         }
+        LlmApiType::Anthropic => Box::new(anthropic::Client::new(address, api_key).await?)
+            as Box<dyn LlmGenerationClient>,
         LlmApiType::Bedrock => {
             Box::new(bedrock::Client::new(address).await?) as Box<dyn LlmGenerationClient>
         }
-        LlmApiType::LiteLlm => {
-            Box::new(litellm::Client::new_litellm(address).await?) as Box<dyn LlmGenerationClient>
-        }
-        LlmApiType::OpenRouter => Box::new(openrouter::Client::new_openrouter(address).await?)
+        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

rust/cocoindex/src/llm/openai.rs

@@ -31,7 +31,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"),
@@ -48,13 +52,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),
         })
     }

rust/cocoindex/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);

rust/cocoindex/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);

rust/cocoindex/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(),

rust/cocoindex/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<AuthEntryReference<String>>,
 }
 
 struct Args {
@@ -112,15 +113,26 @@ impl SimpleFunctionFactoryBase for Factory {
         &'a self,
         spec: &'a Spec,
         args_resolver: &mut OpArgsResolver<'a>,
-        _context: &FlowInstanceContext,
+        context: &FlowInstanceContext,
     ) -> Result<(Self::ResolvedArgs, EnrichedValueType)> {
         let text = args_resolver
             .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 api_key = spec
+            .api_key
+            .as_ref()
+            .map(|key_ref| context.auth_registry.get(key_ref))
+            .transpose()?;
+
+        let client = new_llm_embedding_client(
+            spec.api_type,
+            spec.address.clone(),
+            api_key,
+            spec.api_config.clone(),
+        )
+        .await?;
         let output_dimension = match spec.output_dimension {
             Some(output_dimension) => output_dimension,
             None => {
@@ -171,6 +183,7 @@ mod tests {
             api_config: None,
             output_dimension: None,
             task_type: None,
+            api_key: None,
         };
 
         let factory = Arc::new(Factory);