Feature: Enable programmatically pass in api_key besides reading from env (#1134)

* Feature:Add tests for json_schema.rs

* Feature:Enable programmatically pass in api_key besides reading from env

* Moved api_key to Common LlmSpec Layer

* Fix formatting issues for GitHub Actions

- Fixed Python line length issue in llm.py by breaking long type annotation
- Fixed Rust function signature formatting in all LLM client files
- Fixed long function call formatting in embed_text.rs
- All formatting now complies with project standards

* Fix type mismatch error in embed_text.rs

- Fixed api_bail! usage in context expecting LlmApiConfig return type
- Replaced unwrap_or_else with proper if-let pattern matching
- Resolves compilation error in GitHub Actions build test

* Fix trailing whitespace formatting issues

- Removed trailing whitespace from all LLM client files
- Fixed formatting issues in gemini.rs, litellm.rs, openai.rs, openrouter.rs, vllm.rs
- Fixed trailing whitespace in embed_text.rs
- All files now comply with cargo fmt standards

* Feature: Enable programmatically pass in api_key besides reading from env

* Feature: Enable programmatically pass in api_key besides reading from env

* Feature: Enable programmatically pass in api_key besides reading from env

* Feature: Enable programmatically pass in api_key besides reading from env

* Feature: Enable programmatically pass in api_key besides reading from env

* Feature: Enable programmatically pass in api_key besides reading from env

Author: aryasoni98

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(),
@@ -253,11 +257,17 @@ 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);