Add Azure OpenAI as a first-class LLM and embedding provider

Both factories advertised multi-provider support behind real ABCs but only ever
constructed public-OpenAI clients. This adds azure_openai to each.

Approach: target Azure's v1 OpenAI-compatible surface (<endpoint>/openai/v1) and reuse
the standard AsyncOpenAI client rather than AsyncAzureOpenAI, whose static types the
openai SDK's own README warns "can be incorrect". So the Azure branch is a different
base_url and a deployment name, not a second client implementation.

- config.py: azure_openai_endpoint / _api_key / _deployment / _embedding_deployment /
  _tokenizer_model, plus azure_openai_base_url() which normalises the endpoint to
  /openai/v1 idempotently (accepts bare host, trailing slash, or an already-complete
  URL).
- llm/factory.py + embeddings/factory.py: azure_openai branch, failing fast with a
  named variable when endpoint / key / deployment are missing.
- openai_llm.py + openai_embeddings.py: api_key widened to str | Callable[[], str] so
  an Entra token provider can be passed without either class knowing how the
  credential is obtained. Nothing here supplies one yet -- auth is API-key only.
- requirements.txt: openai>=1.106.0, the floor Microsoft documents for the v1 surface
  and callable token providers, and the first version exporting the error classes the
  health check now discriminates on. Was >=1.12.0.

Two Azure divergences that would otherwise be silent:

- health_check no longer treats a missing /models route as unhealthy. On Azure that
  route enumerates *deployments* and some configurations omit it entirely; a 404 means
  the endpoint answered, so credentials and networking are fine. 401/403 and unexpected
  statuses still fail. Previously /api/health would have reported a working Azure
  deployment as degraded.
- openai_embeddings.py takes tokenizer_model separately, because tiktoken resolves an
  encoding from a model id and on Azure `model` is a deployment name. Note honestly
  that the old bare `except KeyError: cl100k_base` was *accidentally* correct: every
  current OpenAI embedding model resolves to cl100k_base anyway. So this is a latent
  correctness fix plus a warning where there was silence, not a live bug fix. It would
  have mattered on an o200k_base embedding model or a non-OpenAI base_url.

Also, because Azure makes them reachable:

- openai_embedding_dimensions is now configurable and threaded through both factories.
  It was hardcoded to 1536 while the model was configurable, so a text-embedding-3-large
  deployment (3072) only failed when Chroma rejected the insert.
- embeddings/factory.py's unknown-provider branch now raises instead of falling back to
  OpenAI whenever a key happened to be set. That fallback dropped all seven rate-limit,
  batching and pacing arguments, so a typo in EMBEDDING_PROVIDER silently produced a
  differently-behaving client with no error. llm/factory.py already raised.

docker-compose.yml forwards the six new variables (55 total), and .env.example plus
docker/README.md document deployment-names-not-model-ids, the tokenizer requirement,
the dimensions match, and that Entra is not wired up.

Verified: 22 new unit tests covering URL normalisation, factory wiring, missing-config
errors, tokenizer resolution, dimensions and all five health-check branches; 111 tests
pass (was 89); ruff clean; both providers construct correctly from environment alone
and the default OpenAI path is unchanged.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Author: ShreeBohara

.env.example

@@ -24,12 +24,37 @@
 # For local LLM (optional, no key needed)
 # OLLAMA_BASE_URL=http://localhost:11434
 
+# -----------------------
+# Azure OpenAI (optional)
+# -----------------------
+# Uses Azure's v1 OpenAI-compatible surface, so the standard OpenAI client talks to
+# it directly -- the endpoint below is normalised to <endpoint>/openai/v1 for you.
+#
+# The two things that differ from public OpenAI:
+#   1. You pass DEPLOYMENT NAMES, not model ids. Azure sends the deployment name
+#      where a model id normally goes.
+#   2. Because a deployment name is not a model id, tiktoken cannot derive an
+#      encoding from it. Name the underlying model in AZURE_OPENAI_TOKENIZER_MODEL
+#      so token counting (used for truncation and batch splitting) stays exact.
+#
+# LLM_PROVIDER=azure_openai
+# EMBEDDING_PROVIDER=azure_openai
+# AZURE_OPENAI_ENDPOINT=https://my-resource.openai.azure.com
+# AZURE_OPENAI_API_KEY=...
+# AZURE_OPENAI_DEPLOYMENT=my-gpt4o-deployment
+# AZURE_OPENAI_EMBEDDING_DEPLOYMENT=my-embedding-deployment
+# AZURE_OPENAI_TOKENIZER_MODEL=text-embedding-3-small
+#
+# Set this to your deployed embedding model's output size, or Chroma will reject the
+# insert: text-embedding-3-small is 1536, text-embedding-3-large is 3072.
+# OPENAI_EMBEDDING_DIMENSIONS=1536
+
 # -----------------------
 # Embedding Providers
 # -----------------------
 # Uses OpenAI by default. Uncomment to use alternatives:
-# EMBEDDING_PROVIDER=openai  # openai or ollama
-# LLM_PROVIDER=openai        # openai, anthropic, or ollama
+# EMBEDDING_PROVIDER=openai  # openai, azure_openai, or ollama
+# LLM_PROVIDER=openai        # openai, azure_openai, anthropic, or ollama
 # OPENAI_EMBEDDING_MAX_TOKENS_PER_REQUEST=250000
 # OPENAI_EMBEDDING_MAX_TEXTS_PER_REQUEST=128
 # OPENAI_EMBEDDING_REQUEST_CONCURRENCY=1

apps/api/requirements.txt

@@ -20,7 +20,10 @@ aiosqlite>=0.19.0
 chromadb>=1.0.0
 
 # LLM Providers
-openai>=1.12.0
+# >=1.106.0 is the floor Microsoft documents for Azure OpenAI's v1 surface and for
+# passing a callable token provider as api_key. NotFoundError/PermissionDeniedError,
+# used by the provider-aware health check, are also only exported on modern versions.
+openai>=1.106.0
 anthropic>=0.18.0
 tiktoken>=0.6.0
 

apps/api/src/config.py

@@ -29,6 +29,7 @@
     "chroma_persist_dir",
     "repos_dir",
     "vector_db_type",
+    "azure_openai_tokenizer_model",
 )
 
 
@@ -64,18 +65,40 @@ class Settings(BaseSettings):
     qdrant_api_key: Optional[str] = None
 
     # LLM Providers
-    llm_provider: str = "openai"  # "openai", "anthropic", "ollama"
+    llm_provider: str = "openai"  # "openai", "azure_openai", "anthropic", "ollama"
     openai_api_key: Optional[str] = None
     openai_model: str = "gpt-4o"
     anthropic_api_key: Optional[str] = None
     anthropic_model: str = "claude-sonnet-4-20250514"
     ollama_base_url: str = "http://localhost:11434"
     ollama_model: str = "llama3.1"
 
+    # Azure OpenAI
+    #
+    # Targets Azure's v1 OpenAI-compatible surface, so the standard OpenAI client is
+    # used rather than AzureOpenAI (whose static types the openai SDK README warns
+    # "can be incorrect"). azure_openai_base_url() below appends /openai/v1.
+    #
+    # Note that on Azure the *deployment name* takes the place of the model name in
+    # API calls. It is frequently not a model id, which is why the tokenizer must be
+    # named separately -- see azure_openai_tokenizer_model.
+    azure_openai_endpoint: Optional[str] = None  # e.g. https://my-resource.openai.azure.com
+    azure_openai_api_key: Optional[str] = None
+    azure_openai_deployment: Optional[str] = None  # chat deployment name
+    azure_openai_embedding_deployment: Optional[str] = None  # embedding deployment name
+    # tiktoken cannot resolve an encoding from a deployment name; without this it
+    # silently falls back to cl100k_base, which is wrong for o200k_base models and
+    # makes every token count (and therefore every truncation) quietly inaccurate.
+    azure_openai_tokenizer_model: str = "text-embedding-3-small"
+
     # Embedding Providers
-    embedding_provider: str = "openai"  # "openai" or "ollama"
+    embedding_provider: str = "openai"  # "openai", "azure_openai" or "ollama"
     openai_embedding_model: str = "text-embedding-3-small"
     openai_base_url: Optional[str] = None  # Optional: OpenAI-compatible endpoint (e.g., LM Studio)
+    # Must match the deployed model's output size. text-embedding-3-small is 1536,
+    # text-embedding-3-large is 3072; a mismatch is only discovered when Chroma
+    # rejects the insert, so it is configurable rather than hardcoded.
+    openai_embedding_dimensions: int = 1536
     openai_embedding_max_tokens_per_request: int = 250000
     openai_embedding_max_texts_per_request: int = 128
     openai_embedding_request_concurrency: int = 1
@@ -223,6 +246,26 @@ def _blank_falls_back_to_default(cls, value, info: ValidationInfo):
                 return field.default
         return value
 
+    def azure_openai_base_url(self) -> str:
+        """
+        Base URL for Azure's v1 OpenAI-compatible surface.
+
+        Azure exposes an OpenAI-compatible API at <endpoint>/openai/v1, which lets the
+        standard OpenAI client talk to it directly. Accepts an endpoint with or without
+        a trailing slash, and is idempotent if the caller already included /openai/v1.
+        """
+        endpoint = (self.azure_openai_endpoint or "").strip().rstrip("/")
+        if not endpoint:
+            raise ValueError(
+                "AZURE_OPENAI_ENDPOINT is required when using the azure_openai provider "
+                "(e.g. https://my-resource.openai.azure.com)"
+            )
+        if endpoint.endswith("/openai/v1"):
+            return endpoint
+        if endpoint.endswith("/openai"):
+            return f"{endpoint}/v1"
+        return f"{endpoint}/openai/v1"
+
     @property
     def cors_origins(self) -> List[str]:
         """

apps/api/src/core/embeddings/factory.py

@@ -8,11 +8,37 @@ def create_embedding_service() -> BaseEmbeddings:
     """Factory function to create embedding service based on configuration."""
     provider = settings.embedding_provider.lower()
 
-    if provider == "openai":
+    if provider in ("azure_openai", "azure"):
+        # Same client, Azure v1 base_url, deployment name in place of the model id.
+        # tokenizer_model is passed separately because tiktoken cannot resolve an
+        # encoding from a deployment name.
+        if not settings.azure_openai_api_key:
+            raise ValueError("AZURE_OPENAI_API_KEY required for the azure_openai provider")
+        if not settings.azure_openai_embedding_deployment:
+            raise ValueError(
+                "AZURE_OPENAI_EMBEDDING_DEPLOYMENT required for the azure_openai "
+                "embedding provider (the embedding deployment name)"
+            )
+        return OpenAIEmbeddings(
+            api_key=settings.azure_openai_api_key,
+            model=settings.azure_openai_embedding_deployment,
+            base_url=settings.azure_openai_base_url(),
+            dimensions=settings.openai_embedding_dimensions,
+            tokenizer_model=settings.azure_openai_tokenizer_model,
+            max_tokens_per_request=settings.openai_embedding_max_tokens_per_request,
+            max_texts_per_request=settings.openai_embedding_max_texts_per_request,
+            request_concurrency=settings.openai_embedding_request_concurrency,
+            min_seconds_between_requests=settings.openai_embedding_min_seconds_between_requests,
+            rate_limit_max_retries=settings.openai_embedding_rate_limit_max_retries,
+            rate_limit_base_backoff_seconds=settings.openai_embedding_rate_limit_base_backoff_seconds,
+            rate_limit_max_backoff_seconds=settings.openai_embedding_rate_limit_max_backoff_seconds,
+        )
+    elif provider == "openai":
         return OpenAIEmbeddings(
             api_key=settings.openai_api_key,
             model=settings.openai_embedding_model,
             base_url=settings.openai_base_url,
+            dimensions=settings.openai_embedding_dimensions,
             max_tokens_per_request=settings.openai_embedding_max_tokens_per_request,
             max_texts_per_request=settings.openai_embedding_max_texts_per_request,
             request_concurrency=settings.openai_embedding_request_concurrency,
@@ -31,11 +57,12 @@ def create_embedding_service() -> BaseEmbeddings:
             max_failure_ratio=settings.ollama_embedding_max_failure_ratio,
         )
     else:
-        # Fallback/Default or Raise
-        # For now, if unknown, default to OpenAI if key exists, else error
-        if settings.openai_api_key:
-             return OpenAIEmbeddings(
-                api_key=settings.openai_api_key,
-                model=settings.openai_embedding_model
-            )
-        raise ValueError(f"Unknown embedding provider: {provider}")
+        # Fail fast, matching src/core/llm/factory.py. This previously fell back to
+        # OpenAI whenever a key happened to be present, which meant a typo in
+        # EMBEDDING_PROVIDER silently produced an OpenAI client with *none* of the
+        # rate-limit, batching or pacing settings applied -- so indexing behaved
+        # differently from the configured provider with no error anywhere.
+        raise ValueError(
+            f"Unknown embedding provider: {provider!r}. "
+            "Expected one of: openai, azure_openai, ollama."
+        )

apps/api/src/core/embeddings/openai_embeddings.py

@@ -8,7 +8,7 @@
 import random
 import threading
 import time
-from typing import List, Sequence
+from typing import Callable, List, Sequence
 
 import tiktoken
 from openai import AsyncOpenAI, RateLimitError
@@ -23,9 +23,11 @@ class OpenAIEmbeddings(BaseEmbeddings):
 
     def __init__(
         self,
-        api_key: str = None,
+        api_key: str | Callable[[], str] | None = None,
         model: str = "text-embedding-3-small",
         base_url: str | None = None,
+        dimensions: int = 1536,
+        tokenizer_model: str | None = None,
         max_tokens_per_request: int = 250000,
         max_texts_per_request: int = 128,
         request_concurrency: int = 1,
@@ -39,7 +41,7 @@ def __init__(
             client_kwargs["base_url"] = base_url
         self._client = AsyncOpenAI(**client_kwargs)
         self._model = model
-        self._dimensions = 1536
+        self._dimensions = int(dimensions)
         self._max_tokens = 8000  # Leave some buffer from 8192 limit
         self._max_tokens_per_request = max(1, max_tokens_per_request)
         self._max_texts_per_request = max(1, max_texts_per_request)
@@ -54,10 +56,23 @@ def __init__(
         )
         self._request_pacing_lock = threading.Lock()
         self._next_request_time = 0.0
+        # tiktoken resolves an encoding from a *model id*. On Azure `model` is a
+        # deployment name, which will not resolve -- and the bare fallback below is
+        # silent, so every token count (and therefore every truncation in
+        # _truncate_text and every batch split in _split_batches) would be computed
+        # with the wrong encoding without any signal. tokenizer_model lets the caller
+        # name the real model; the fallback now warns instead of hiding it.
+        resolve_from = tokenizer_model or model
         try:
-            self._tokenizer = tiktoken.encoding_for_model(model)
+            self._tokenizer = tiktoken.encoding_for_model(resolve_from)
         except KeyError:
             self._tokenizer = tiktoken.get_encoding("cl100k_base")
+            logger.warning(
+                "tiktoken has no encoding for %r; falling back to cl100k_base. Token "
+                "counts will be approximate. Set AZURE_OPENAI_TOKENIZER_MODEL (or pass "
+                "tokenizer_model) to the underlying model id to fix this.",
+                resolve_from,
+            )
 
     @property
     def dimensions(self) -> int:

apps/api/src/core/llm/factory.py

@@ -15,6 +15,22 @@ def create_llm() -> BaseLLM:
             model=settings.openai_model,
             base_url=settings.openai_base_url,
         )
+    elif provider in ("azure_openai", "azure"):
+        # Azure's v1 surface is OpenAI-compatible, so the same client is reused with a
+        # different base_url. The deployment name takes the place of the model name.
+        if not settings.azure_openai_api_key:
+            raise ValueError("AZURE_OPENAI_API_KEY required for the azure_openai provider")
+        if not settings.azure_openai_deployment:
+            raise ValueError(
+                "AZURE_OPENAI_DEPLOYMENT required for the azure_openai provider "
+                "(the chat deployment name, which Azure uses in place of a model id)"
+            )
+        return OpenAILLM(
+            api_key=settings.azure_openai_api_key,
+            model=settings.azure_openai_deployment,
+            base_url=settings.azure_openai_base_url(),
+            provider_label="azure_openai",
+        )
     elif provider == "anthropic":
         if not settings.anthropic_api_key:
             # Don't raise immediately, allow app to start but fail on use if key missing

apps/api/src/core/llm/openai_llm.py

@@ -4,9 +4,15 @@
 
 import asyncio
 import logging
-from typing import AsyncGenerator, Dict, List
+from typing import AsyncGenerator, Callable, Dict, List
 
-from openai import AsyncOpenAI
+from openai import (
+    APIStatusError,
+    AsyncOpenAI,
+    AuthenticationError,
+    NotFoundError,
+    PermissionDeniedError,
+)
 
 from src.core.llm.base import BaseLLM, stream_error_text
 
@@ -16,13 +22,23 @@
 class OpenAILLM(BaseLLM):
     """OpenAI LLM service with retry logic."""
 
-    def __init__(self, api_key: str = None, model: str = "gpt-4o", base_url: str | None = None):
+    def __init__(
+        self,
+        api_key: str | Callable[[], str] | None = None,
+        model: str = "gpt-4o",
+        base_url: str | None = None,
+        provider_label: str = "openai",
+    ):
+        # api_key accepts a callable so a token provider (e.g. Entra ID) can be passed
+        # without this class needing to know how the credential is obtained.
         client_kwargs = {"api_key": api_key}
         if base_url:
             client_kwargs["base_url"] = base_url
         self._client = AsyncOpenAI(**client_kwargs)
         self._model = model
         self._max_retries = 3
+        # Only used for log messages; behaviour is identical across OpenAI-compatible hosts.
+        self._provider_label = provider_label
 
     async def _retry_with_backoff(self, func, *args, **kwargs):
         """Retry with exponential backoff."""
@@ -122,11 +138,36 @@ async def generate_stream(
                 return
 
     async def health_check(self) -> bool:
-        """Check OpenAI API availability."""
+        """
+        Check provider availability.
+
+        Distinguishes "cannot reach the provider" from "provider does not implement
+        /models". Azure serves an OpenAI-compatible surface but /models enumerates
+        *deployments*, and some configurations do not expose it at all -- a 404 there
+        means the endpoint answered, so credentials and networking are fine and the
+        service is usable. Treating that as unhealthy would report a working Azure
+        deployment as down.
+        """
         try:
-            # Simple models list call to verify API key
             await self._client.models.list()
             return True
+        except NotFoundError:
+            logger.info(
+                "%s does not expose /models; treating as reachable (endpoint responded)",
+                self._provider_label,
+            )
+            return True
+        except (AuthenticationError, PermissionDeniedError) as e:
+            logger.warning("%s health check failed: bad credentials: %s", self._provider_label, e)
+            return False
+        except APIStatusError as e:
+            # Any other HTTP status still proves the endpoint is reachable, but an
+            # unexpected status is worth surfacing rather than silently passing.
+            logger.warning(
+                "%s health check got unexpected status %s: %s",
+                self._provider_label, e.status_code, e,
+            )
+            return False
         except Exception as e:
-            logger.warning(f"OpenAI health check failed: {e}")
+            logger.warning("%s health check failed: %s", self._provider_label, e)
             return False