feat(openai-config): allow custom address, org_id and project_id (#852)

Author: georgeh0

docs/docs/ai/llm.mdx

@@ -75,7 +75,11 @@ CocoIndex integrates with various LLM APIs for these functions.
 To use the OpenAI LLM API, you need to set the environment variable `OPENAI_API_KEY`.
 You can generate the API key from [OpenAI Dashboard](https://platform.openai.com/api-keys).
 
-Currently we don't support custom address for OpenAI API.
+If you want to use a custom address, you can either provide the `address` parameter in `LlmSpec` / `EmbedText`, or set the environment variable `OPENAI_API_BASE`. The `address` parameter takes precedence over the environment variable.
+
+Spec for OpenAI takes additional `api_config` field, in type `cocoindex.llm.OpenAiConfig` with the following fields:
+-   `org_id` (type: `str`, optional): The organization ID of the OpenAI account.
+-   `project_id` (type: `str`, optional): The project ID of the OpenAI account.
 
 You can find the full list of models supported by OpenAI [here](https://platform.openai.com/docs/models).
 

python/cocoindex/llm.py

@@ -26,11 +26,21 @@ class VertexAiConfig:
     region: str | None = None
 
 
+@dataclass
+class OpenAiConfig:
+    """A specification for a OpenAI LLM."""
+
+    kind = "OpenAi"
+
+    org_id: str | None = None
+    project_id: str | None = None
+
+
 @dataclass
 class LlmSpec:
     """A specification for a LLM."""
 
     api_type: LlmApiType
     model: str
     address: str | None = None
-    api_config: VertexAiConfig | None = None
+    api_config: VertexAiConfig | OpenAiConfig | None = None

src/llm/mod.rs

@@ -26,10 +26,17 @@ pub struct VertexAiConfig {
     pub region: Option<String>,
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize, Default)]
+pub struct OpenAiConfig {
+    pub org_id: Option<String>,
+    pub project_id: Option<String>,
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(tag = "kind")]
 pub enum LlmApiConfig {
     VertexAi(VertexAiConfig),
+    OpenAi(OpenAiConfig),
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -113,7 +120,7 @@ pub async fn new_llm_generation_client(
             Box::new(ollama::Client::new(address).await?) as Box<dyn LlmGenerationClient>
         }
         LlmApiType::OpenAi => {
-            Box::new(openai::Client::new(address)?) as Box<dyn LlmGenerationClient>
+            Box::new(openai::Client::new(address, api_config)?) as Box<dyn LlmGenerationClient>
         }
         LlmApiType::Gemini => {
             Box::new(gemini::AiStudioClient::new(address)?) as Box<dyn LlmGenerationClient>
@@ -151,7 +158,7 @@ pub async fn new_llm_embedding_client(
             Box::new(gemini::AiStudioClient::new(address)?) as Box<dyn LlmEmbeddingClient>
         }
         LlmApiType::OpenAi => {
-            Box::new(openai::Client::new(address)?) as Box<dyn LlmEmbeddingClient>
+            Box::new(openai::Client::new(address, api_config)?) as Box<dyn LlmEmbeddingClient>
         }
         LlmApiType::Voyage => {
             Box::new(voyage::Client::new(address)?) as Box<dyn LlmEmbeddingClient>

src/llm/openai.rs

@@ -33,17 +33,31 @@ impl Client {
         Self { client }
     }
 
-    pub fn new(address: Option<String>) -> Result<Self> {
+    pub fn new(address: 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"),
+            None => super::OpenAiConfig::default(),
+        };
+
+        let mut openai_config = OpenAIConfig::new();
         if let Some(address) = address {
-            api_bail!("OpenAI doesn't support custom API address: {address}");
+            openai_config = openai_config.with_api_base(address);
+        }
+        if let Some(org_id) = config.org_id {
+            openai_config = openai_config.with_org_id(org_id);
         }
+        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");
         }
         Ok(Self {
-            // OpenAI client will use OPENAI_API_KEY env variable by default
-            client: OpenAIClient::new(),
+            // OpenAI client will use OPENAI_API_KEY and OPENAI_API_BASE env variables by default
+            client: OpenAIClient::with_config(openai_config),
         })
     }
 }