From fa53917c6378c47e84a168d6e53731e72695131d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Tue, 13 Jan 2026 14:01:44 +0100 Subject: [PATCH] feat: support custom url for openai embeddings & cohere (#150) * feat: support custom url for openai embeddings & cohere * feat: support custom url for openai embeddings & cohere --- hindsight-api/hindsight_api/config.py | 9 +++++ .../hindsight_api/engine/cross_encoder.py | 17 ++++++++-- .../hindsight_api/engine/embeddings.py | 34 +++++++++++++++---- hindsight-api/hindsight_api/main.py | 3 ++ .../docs/developer/configuration.md | 21 ++++++++++++ 5 files changed, 75 insertions(+), 9 deletions(-) diff --git a/hindsight-api/hindsight_api/config.py b/hindsight-api/hindsight_api/config.py index 18264998..d6c4f078 100644 --- a/hindsight-api/hindsight_api/config.py +++ b/hindsight-api/hindsight_api/config.py @@ -41,10 +41,13 @@ ENV_EMBEDDINGS_LOCAL_MODEL = "HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL" ENV_EMBEDDINGS_TEI_URL = "HINDSIGHT_API_EMBEDDINGS_TEI_URL" ENV_EMBEDDINGS_OPENAI_API_KEY = "HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY" ENV_EMBEDDINGS_OPENAI_MODEL = "HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL" +ENV_EMBEDDINGS_OPENAI_BASE_URL = "HINDSIGHT_API_EMBEDDINGS_OPENAI_BASE_URL" ENV_COHERE_API_KEY = "HINDSIGHT_API_COHERE_API_KEY" ENV_EMBEDDINGS_COHERE_MODEL = "HINDSIGHT_API_EMBEDDINGS_COHERE_MODEL" +ENV_EMBEDDINGS_COHERE_BASE_URL = "HINDSIGHT_API_EMBEDDINGS_COHERE_BASE_URL" ENV_RERANKER_COHERE_MODEL = "HINDSIGHT_API_RERANKER_COHERE_MODEL" +ENV_RERANKER_COHERE_BASE_URL = "HINDSIGHT_API_RERANKER_COHERE_BASE_URL" ENV_RERANKER_PROVIDER = "HINDSIGHT_API_RERANKER_PROVIDER" ENV_RERANKER_LOCAL_MODEL = "HINDSIGHT_API_RERANKER_LOCAL_MODEL" @@ -224,6 +227,8 @@ class HindsightConfig: embeddings_provider: str embeddings_local_model: str embeddings_tei_url: str | None + embeddings_openai_base_url: str | None + embeddings_cohere_base_url: str | None # Reranker reranker_provider: str @@ -232,6 +237,7 @@ class HindsightConfig: reranker_tei_batch_size: int reranker_tei_max_concurrent: int reranker_max_candidates: int + reranker_cohere_base_url: str | None # Server host: str @@ -300,6 +306,8 @@ class HindsightConfig: embeddings_provider=os.getenv(ENV_EMBEDDINGS_PROVIDER, DEFAULT_EMBEDDINGS_PROVIDER), embeddings_local_model=os.getenv(ENV_EMBEDDINGS_LOCAL_MODEL, DEFAULT_EMBEDDINGS_LOCAL_MODEL), embeddings_tei_url=os.getenv(ENV_EMBEDDINGS_TEI_URL), + embeddings_openai_base_url=os.getenv(ENV_EMBEDDINGS_OPENAI_BASE_URL) or None, + embeddings_cohere_base_url=os.getenv(ENV_EMBEDDINGS_COHERE_BASE_URL) or None, # Reranker reranker_provider=os.getenv(ENV_RERANKER_PROVIDER, DEFAULT_RERANKER_PROVIDER), reranker_local_model=os.getenv(ENV_RERANKER_LOCAL_MODEL, DEFAULT_RERANKER_LOCAL_MODEL), @@ -309,6 +317,7 @@ class HindsightConfig: os.getenv(ENV_RERANKER_TEI_MAX_CONCURRENT, str(DEFAULT_RERANKER_TEI_MAX_CONCURRENT)) ), reranker_max_candidates=int(os.getenv(ENV_RERANKER_MAX_CANDIDATES, str(DEFAULT_RERANKER_MAX_CANDIDATES))), + reranker_cohere_base_url=os.getenv(ENV_RERANKER_COHERE_BASE_URL) or None, # Server host=os.getenv(ENV_HOST, DEFAULT_HOST), port=int(os.getenv(ENV_PORT, DEFAULT_PORT)), diff --git a/hindsight-api/hindsight_api/engine/cross_encoder.py b/hindsight-api/hindsight_api/engine/cross_encoder.py index 1189d131..16ef2fc1 100644 --- a/hindsight-api/hindsight_api/engine/cross_encoder.py +++ b/hindsight-api/hindsight_api/engine/cross_encoder.py @@ -24,6 +24,7 @@ from ..config import ( DEFAULT_RERANKER_TEI_BATCH_SIZE, DEFAULT_RERANKER_TEI_MAX_CONCURRENT, ENV_COHERE_API_KEY, + ENV_RERANKER_COHERE_BASE_URL, ENV_RERANKER_COHERE_MODEL, ENV_RERANKER_FLASHRANK_CACHE_DIR, ENV_RERANKER_FLASHRANK_MODEL, @@ -392,6 +393,7 @@ class CohereCrossEncoder(CrossEncoderModel): self, api_key: str, model: str = DEFAULT_RERANKER_COHERE_MODEL, + base_url: str | None = None, timeout: float = 60.0, ): """ @@ -400,10 +402,12 @@ class CohereCrossEncoder(CrossEncoderModel): Args: api_key: Cohere API key model: Cohere rerank model name (default: rerank-english-v3.0) + base_url: Custom base URL for Cohere-compatible API (e.g., Azure-hosted endpoint) timeout: Request timeout in seconds (default: 60.0) """ self.api_key = api_key self.model = model + self.base_url = base_url self.timeout = timeout self._client = None @@ -421,8 +425,14 @@ class CohereCrossEncoder(CrossEncoderModel): except ImportError: raise ImportError("cohere is required for CohereCrossEncoder. Install it with: pip install cohere") - logger.info(f"Reranker: initializing Cohere provider with model {self.model}") - self._client = cohere.Client(api_key=self.api_key, timeout=self.timeout) + base_url_msg = f" at {self.base_url}" if self.base_url else "" + logger.info(f"Reranker: initializing Cohere provider with model {self.model}{base_url_msg}") + + # Build client kwargs, only including base_url if set (for Azure or custom endpoints) + client_kwargs = {"api_key": self.api_key, "timeout": self.timeout} + if self.base_url: + client_kwargs["base_url"] = self.base_url + self._client = cohere.Client(**client_kwargs) logger.info("Reranker: Cohere provider initialized") async def predict(self, pairs: list[tuple[str, str]]) -> list[float]: @@ -671,7 +681,8 @@ def create_cross_encoder_from_env() -> CrossEncoderModel: if not api_key: raise ValueError(f"{ENV_COHERE_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'cohere'") model = os.environ.get(ENV_RERANKER_COHERE_MODEL, DEFAULT_RERANKER_COHERE_MODEL) - return CohereCrossEncoder(api_key=api_key, model=model) + base_url = os.environ.get(ENV_RERANKER_COHERE_BASE_URL) or None + return CohereCrossEncoder(api_key=api_key, model=model, base_url=base_url) elif provider == "flashrank": model = os.environ.get(ENV_RERANKER_FLASHRANK_MODEL, DEFAULT_RERANKER_FLASHRANK_MODEL) cache_dir = os.environ.get(ENV_RERANKER_FLASHRANK_CACHE_DIR, DEFAULT_RERANKER_FLASHRANK_CACHE_DIR) diff --git a/hindsight-api/hindsight_api/engine/embeddings.py b/hindsight-api/hindsight_api/engine/embeddings.py index 3a4c953b..ee7a558c 100644 --- a/hindsight-api/hindsight_api/engine/embeddings.py +++ b/hindsight-api/hindsight_api/engine/embeddings.py @@ -21,9 +21,11 @@ from ..config import ( DEFAULT_EMBEDDINGS_OPENAI_MODEL, DEFAULT_EMBEDDINGS_PROVIDER, ENV_COHERE_API_KEY, + ENV_EMBEDDINGS_COHERE_BASE_URL, ENV_EMBEDDINGS_COHERE_MODEL, ENV_EMBEDDINGS_LOCAL_MODEL, ENV_EMBEDDINGS_OPENAI_API_KEY, + ENV_EMBEDDINGS_OPENAI_BASE_URL, ENV_EMBEDDINGS_OPENAI_MODEL, ENV_EMBEDDINGS_PROVIDER, ENV_EMBEDDINGS_TEI_URL, @@ -322,6 +324,7 @@ class OpenAIEmbeddings(Embeddings): self, api_key: str, model: str = DEFAULT_EMBEDDINGS_OPENAI_MODEL, + base_url: str | None = None, batch_size: int = 100, max_retries: int = 3, ): @@ -331,11 +334,13 @@ class OpenAIEmbeddings(Embeddings): Args: api_key: OpenAI API key model: OpenAI embedding model name (default: text-embedding-3-small) + base_url: Custom base URL for OpenAI-compatible API (e.g., Azure OpenAI endpoint) batch_size: Maximum batch size for embedding requests (default: 100) max_retries: Maximum number of retries for failed requests (default: 3) """ self.api_key = api_key self.model = model + self.base_url = base_url self.batch_size = batch_size self.max_retries = max_retries self._client = None @@ -361,8 +366,14 @@ class OpenAIEmbeddings(Embeddings): except ImportError: raise ImportError("openai is required for OpenAIEmbeddings. Install it with: pip install openai") - logger.info(f"Embeddings: initializing OpenAI provider with model {self.model}") - self._client = OpenAI(api_key=self.api_key, max_retries=self.max_retries) + base_url_msg = f" at {self.base_url}" if self.base_url else "" + logger.info(f"Embeddings: initializing OpenAI provider with model {self.model}{base_url_msg}") + + # Build client kwargs, only including base_url if set (for Azure or custom endpoints) + client_kwargs = {"api_key": self.api_key, "max_retries": self.max_retries} + if self.base_url: + client_kwargs["base_url"] = self.base_url + self._client = OpenAI(**client_kwargs) # Try to get dimension from known models, otherwise do a test embedding if self.model in self.MODEL_DIMENSIONS: @@ -435,6 +446,7 @@ class CohereEmbeddings(Embeddings): self, api_key: str, model: str = DEFAULT_EMBEDDINGS_COHERE_MODEL, + base_url: str | None = None, batch_size: int = 96, timeout: float = 60.0, input_type: str = "search_document", @@ -445,6 +457,7 @@ class CohereEmbeddings(Embeddings): Args: api_key: Cohere API key model: Cohere embedding model name (default: embed-english-v3.0) + base_url: Custom base URL for Cohere-compatible API (e.g., Azure-hosted endpoint) batch_size: Maximum batch size for embedding requests (default: 96, Cohere's limit) timeout: Request timeout in seconds (default: 60.0) input_type: Input type for embeddings (default: search_document). @@ -452,6 +465,7 @@ class CohereEmbeddings(Embeddings): """ self.api_key = api_key self.model = model + self.base_url = base_url self.batch_size = batch_size self.timeout = timeout self.input_type = input_type @@ -478,8 +492,14 @@ class CohereEmbeddings(Embeddings): except ImportError: raise ImportError("cohere is required for CohereEmbeddings. Install it with: pip install cohere") - logger.info(f"Embeddings: initializing Cohere provider with model {self.model}") - self._client = cohere.Client(api_key=self.api_key, timeout=self.timeout) + base_url_msg = f" at {self.base_url}" if self.base_url else "" + logger.info(f"Embeddings: initializing Cohere provider with model {self.model}{base_url_msg}") + + # Build client kwargs, only including base_url if set (for Azure or custom endpoints) + client_kwargs = {"api_key": self.api_key, "timeout": self.timeout} + if self.base_url: + client_kwargs["base_url"] = self.base_url + self._client = cohere.Client(**client_kwargs) # Try to get dimension from known models, otherwise do a test embedding if self.model in self.MODEL_DIMENSIONS: @@ -558,12 +578,14 @@ def create_embeddings_from_env() -> Embeddings: f"when {ENV_EMBEDDINGS_PROVIDER} is 'openai'" ) model = os.environ.get(ENV_EMBEDDINGS_OPENAI_MODEL, DEFAULT_EMBEDDINGS_OPENAI_MODEL) - return OpenAIEmbeddings(api_key=api_key, model=model) + base_url = os.environ.get(ENV_EMBEDDINGS_OPENAI_BASE_URL) or None + return OpenAIEmbeddings(api_key=api_key, model=model, base_url=base_url) elif provider == "cohere": api_key = os.environ.get(ENV_COHERE_API_KEY) if not api_key: raise ValueError(f"{ENV_COHERE_API_KEY} is required when {ENV_EMBEDDINGS_PROVIDER} is 'cohere'") model = os.environ.get(ENV_EMBEDDINGS_COHERE_MODEL, DEFAULT_EMBEDDINGS_COHERE_MODEL) - return CohereEmbeddings(api_key=api_key, model=model) + base_url = os.environ.get(ENV_EMBEDDINGS_COHERE_BASE_URL) or None + return CohereEmbeddings(api_key=api_key, model=model, base_url=base_url) else: raise ValueError(f"Unknown embeddings provider: {provider}. Supported: 'local', 'tei', 'openai', 'cohere'") diff --git a/hindsight-api/hindsight_api/main.py b/hindsight-api/hindsight_api/main.py index 30264cbe..f0dddc84 100644 --- a/hindsight-api/hindsight_api/main.py +++ b/hindsight-api/hindsight_api/main.py @@ -187,12 +187,15 @@ def main(): embeddings_provider=config.embeddings_provider, embeddings_local_model=config.embeddings_local_model, embeddings_tei_url=config.embeddings_tei_url, + embeddings_openai_base_url=config.embeddings_openai_base_url, + embeddings_cohere_base_url=config.embeddings_cohere_base_url, reranker_provider=config.reranker_provider, reranker_local_model=config.reranker_local_model, reranker_tei_url=config.reranker_tei_url, reranker_tei_batch_size=config.reranker_tei_batch_size, reranker_tei_max_concurrent=config.reranker_tei_max_concurrent, reranker_max_candidates=config.reranker_max_candidates, + reranker_cohere_base_url=config.reranker_cohere_base_url, host=args.host, port=args.port, log_level=args.log_level, diff --git a/hindsight-docs/docs/developer/configuration.md b/hindsight-docs/docs/developer/configuration.md index 02356daf..0347a5b4 100644 --- a/hindsight-docs/docs/developer/configuration.md +++ b/hindsight-docs/docs/developer/configuration.md @@ -144,8 +144,10 @@ export HINDSIGHT_API_REFLECT_LLM_MODEL=llama-3.3-70b-versatile | `HINDSIGHT_API_EMBEDDINGS_TEI_URL` | TEI server URL | - | | `HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY` | OpenAI API key (falls back to `HINDSIGHT_API_LLM_API_KEY`) | - | | `HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL` | OpenAI embedding model | `text-embedding-3-small` | +| `HINDSIGHT_API_EMBEDDINGS_OPENAI_BASE_URL` | Custom base URL for OpenAI-compatible API (e.g., Azure OpenAI) | - | | `HINDSIGHT_API_COHERE_API_KEY` | Cohere API key (shared for embeddings and reranker) | - | | `HINDSIGHT_API_EMBEDDINGS_COHERE_MODEL` | Cohere embedding model | `embed-english-v3.0` | +| `HINDSIGHT_API_EMBEDDINGS_COHERE_BASE_URL` | Custom base URL for Cohere-compatible API (e.g., Azure-hosted) | - | ```bash # Local (default) - uses SentenceTransformers @@ -157,6 +159,12 @@ export HINDSIGHT_API_EMBEDDINGS_PROVIDER=openai export HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY=sk-xxxxxxxxxxxx # or reuses HINDSIGHT_API_LLM_API_KEY export HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL=text-embedding-3-small # 1536 dimensions +# Azure OpenAI - embeddings via Azure endpoint +export HINDSIGHT_API_EMBEDDINGS_PROVIDER=openai +export HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY=your-azure-api-key +export HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL=text-embedding-3-small +export HINDSIGHT_API_EMBEDDINGS_OPENAI_BASE_URL=https://your-resource.openai.azure.com/openai/deployments/your-deployment + # TEI - HuggingFace Text Embeddings Inference (recommended for production) export HINDSIGHT_API_EMBEDDINGS_PROVIDER=tei export HINDSIGHT_API_EMBEDDINGS_TEI_URL=http://localhost:8080 @@ -165,6 +173,12 @@ export HINDSIGHT_API_EMBEDDINGS_TEI_URL=http://localhost:8080 export HINDSIGHT_API_EMBEDDINGS_PROVIDER=cohere export HINDSIGHT_API_COHERE_API_KEY=your-api-key export HINDSIGHT_API_EMBEDDINGS_COHERE_MODEL=embed-english-v3.0 # 1024 dimensions + +# Azure-hosted Cohere - embeddings via custom endpoint +export HINDSIGHT_API_EMBEDDINGS_PROVIDER=cohere +export HINDSIGHT_API_COHERE_API_KEY=your-azure-api-key +export HINDSIGHT_API_EMBEDDINGS_COHERE_MODEL=embed-english-v3.0 +export HINDSIGHT_API_EMBEDDINGS_COHERE_BASE_URL=https://your-azure-cohere-endpoint.com ``` #### Embedding Dimensions @@ -194,6 +208,7 @@ Supported OpenAI embedding dimensions: | `HINDSIGHT_API_RERANKER_TEI_BATCH_SIZE` | Batch size for TEI reranking | `128` | | `HINDSIGHT_API_RERANKER_TEI_MAX_CONCURRENT` | Max concurrent TEI reranking requests | `8` | | `HINDSIGHT_API_RERANKER_COHERE_MODEL` | Cohere rerank model | `rerank-english-v3.0` | +| `HINDSIGHT_API_RERANKER_COHERE_BASE_URL` | Custom base URL for Cohere-compatible API (e.g., Azure-hosted) | - | ```bash # Local (default) - uses SentenceTransformers CrossEncoder @@ -208,6 +223,12 @@ export HINDSIGHT_API_RERANKER_TEI_URL=http://localhost:8081 export HINDSIGHT_API_RERANKER_PROVIDER=cohere export HINDSIGHT_API_COHERE_API_KEY=your-api-key # shared with embeddings export HINDSIGHT_API_RERANKER_COHERE_MODEL=rerank-english-v3.0 + +# Azure-hosted Cohere - reranking via custom endpoint +export HINDSIGHT_API_RERANKER_PROVIDER=cohere +export HINDSIGHT_API_COHERE_API_KEY=your-azure-api-key +export HINDSIGHT_API_RERANKER_COHERE_MODEL=rerank-english-v3.0 +export HINDSIGHT_API_RERANKER_COHERE_BASE_URL=https://your-azure-cohere-endpoint.com ``` ### Authentication