Requested changes

Author: asimurka

docs/providers.md

@@ -84,9 +84,9 @@ Example of LCS config section:
 
 ```yaml
 azure_entra_id:
-  tenant_id: ${env.AZURE_TENANT_ID}
-  client_id: ${env.AZURE_CLIENT_ID}
-  client_secret: ${env.AZURE_CLIENT_SECRET}
+  tenant_id: ${env.TENANT_ID}
+  client_id: ${env.CLIENT_ID}
+  client_secret: ${env.CLIENT_SECRET}
   # scope: "https://cognitiveservices.azure.com/.default"  # optional, this is the default
 ```
 

scripts/llama-stack-entrypoint.sh

@@ -12,24 +12,24 @@ ENV_FILE="/opt/app-root/.env"
 # Enrich config if lightspeed config exists
 if [ -f "$LIGHTSPEED_CONFIG" ]; then
     echo "Enriching llama-stack config..."
+    ENRICHMENT_FAILED=0
     python3 /opt/app-root/llama_stack_configuration.py \
         -c "$LIGHTSPEED_CONFIG" \
         -i "$INPUT_CONFIG" \
         -o "$ENRICHED_CONFIG" \
-        -e "$ENV_FILE" 2>&1 || true
+        -e "$ENV_FILE" 2>&1 || ENRICHMENT_FAILED=1
 
     # Source .env if generated (contains AZURE_API_KEY)
     if [ -f "$ENV_FILE" ]; then
         # shellcheck source=/dev/null
         set -a && . "$ENV_FILE" && set +a
     fi
 
-    if [ -f "$ENRICHED_CONFIG" ]; then
+    if [ -f "$ENRICHED_CONFIG" ] && [ "$ENRICHMENT_FAILED" -eq 0 ]; then
         echo "Using enriched config: $ENRICHED_CONFIG"
         exec llama stack run "$ENRICHED_CONFIG"
     fi
 fi
 
 echo "Using original config: $INPUT_CONFIG"
 exec llama stack run "$INPUT_CONFIG"
-

src/app/endpoints/query.py

@@ -307,36 +307,34 @@ async def query_endpoint_handler_base(  # pylint: disable=R0914
     try:
         check_tokens_available(configuration.quota_limiters, user_id)
         # try to get Llama Stack client
-        client_holder = AsyncLlamaStackClientHolder()
-        client = client_holder.get_client()
+        client = AsyncLlamaStackClientHolder().get_client()
         llama_stack_model_id, model_id, provider_id = select_model_and_provider_id(
             await client.models.list(),
             *evaluate_model_hints(
                 user_conversation=user_conversation, query_request=query_request
             ),
         )
 
-        if provider_id == "azure":
-            if (
-                AzureEntraIDManager().is_entra_id_configured
-                and AzureEntraIDManager().is_token_expired
-            ):
-                await AzureEntraIDManager().refresh_token()
-
-                if client_holder.is_library_client:
-                    client = await client_holder.reload_library_client()
-                else:
-                    azure_config = next(
-                        p.config
-                        for p in await client.providers.list()
-                        if p.provider_type == "remote::azure"
-                    )
-                    client = client_holder.update_provider_data(
-                        {
-                            "azure_api_key": AzureEntraIDManager().access_token.get_secret_value(),
-                            "azure_api_base": str(azure_config.get("api_base")),
-                        }
-                    )
+        if (
+            provider_id == "azure"
+            and AzureEntraIDManager().is_entra_id_configured
+            and AzureEntraIDManager().is_token_expired
+            and AzureEntraIDManager().refresh_token()
+        ):
+            if AsyncLlamaStackClientHolder().is_library_client:
+                client = await AsyncLlamaStackClientHolder().reload_library_client()
+            else:
+                azure_config = next(
+                    p.config
+                    for p in await client.providers.list()
+                    if p.provider_type == "remote::azure"
+                )
+                client = AsyncLlamaStackClientHolder().update_provider_data(
+                    {
+                        "azure_api_key": AzureEntraIDManager().access_token.get_secret_value(),
+                        "azure_api_base": str(azure_config.get("api_base")),
+                    }
+                )
 
         summary, conversation_id, referenced_documents, token_usage = (
             await retrieve_response_func(

src/app/endpoints/streaming_query.py

@@ -884,36 +884,34 @@ async def streaming_query_endpoint_handler_base(  # pylint: disable=too-many-loc
 
     try:
         # try to get Llama Stack client
-        client_holder = AsyncLlamaStackClientHolder()
-        client = client_holder.get_client()
+        client = AsyncLlamaStackClientHolder().get_client()
         llama_stack_model_id, model_id, provider_id = select_model_and_provider_id(
             await client.models.list(),
             *evaluate_model_hints(
                 user_conversation=user_conversation, query_request=query_request
             ),
         )
 
-        if provider_id == "azure":
-            if (
-                AzureEntraIDManager().is_entra_id_configured
-                and AzureEntraIDManager().is_token_expired
-            ):
-                await AzureEntraIDManager().refresh_token()
-
-                if client_holder.is_library_client:
-                    client = await client_holder.reload_library_client()
-                else:
-                    azure_config = next(
-                        p.config
-                        for p in await client.providers.list()
-                        if p.provider_type == "remote::azure"
-                    )
-                    client = client_holder.update_provider_data(
-                        {
-                            "azure_api_key": AzureEntraIDManager().access_token.get_secret_value(),
-                            "azure_api_base": str(azure_config.get("api_base")),
-                        }
-                    )
+        if (
+            provider_id == "azure"
+            and AzureEntraIDManager().is_entra_id_configured
+            and AzureEntraIDManager().is_token_expired
+            and AzureEntraIDManager().refresh_token()
+        ):
+            if AsyncLlamaStackClientHolder().is_library_client:
+                client = await AsyncLlamaStackClientHolder().reload_library_client()
+            else:
+                azure_config = next(
+                    p.config
+                    for p in await client.providers.list()
+                    if p.provider_type == "remote::azure"
+                )
+                client = AsyncLlamaStackClientHolder().update_provider_data(
+                    {
+                        "azure_api_key": AzureEntraIDManager().access_token.get_secret_value(),
+                        "azure_api_base": str(azure_config.get("api_base")),
+                    }
+                )
 
         response, conversation_id = await retrieve_response_func(
             client,

src/app/main.py

@@ -44,14 +44,11 @@ async def lifespan(_app: FastAPI) -> AsyncIterator[None]:
     azure_config = configuration.configuration.azure_entra_id
     if azure_config is not None:
         AzureEntraIDManager().set_config(azure_config)
-        try:
-            await AzureEntraIDManager().refresh_token()
-            os.environ["AZURE_API_KEY"] = (
-                AzureEntraIDManager().access_token.get_secret_value()
+        if not AzureEntraIDManager().refresh_token():
+            logger.warning(
+                "Failed to refresh Azure token at startup. "
+                "Token refresh will be retried on next Azure request."
             )
-            logger.info("Azure Entra ID token set in environment")
-        except ValueError as e:
-            logger.error("Failed to refresh Azure token: %s", e)
 
     await AsyncLlamaStackClientHolder().load(configuration.configuration.llama_stack)
     client = AsyncLlamaStackClientHolder().get_client()

src/authorization/azure_token_manager.py

@@ -56,11 +56,11 @@ def access_token(self) -> SecretStr:
         """Return the access token from environment variable as SecretStr."""
         return SecretStr(os.environ.get("AZURE_API_KEY", ""))
 
-    async def refresh_token(self) -> None:
+    def refresh_token(self) -> bool:
         """Refresh the cached Azure access token.
 
-        This is async to enforce proper ordering in the event loop -
-        callers must await this before using the refreshed token.
+        Returns:
+            bool: True if token was successfully refreshed, False otherwise.
 
         Raises:
             ValueError: If Entra ID configuration has not been set.
@@ -72,6 +72,8 @@ async def refresh_token(self) -> None:
         token_obj = self._retrieve_access_token()
         if token_obj:
             self._update_access_token(token_obj.token, token_obj.expires_on)
+            return True
+        return False
 
     def _update_access_token(self, token: str, expires_on: int) -> None:
         """Update the token in env var and track expiration time."""
@@ -96,5 +98,5 @@ def _retrieve_access_token(self) -> Optional[AccessToken]:
             return credential.get_token(self._entra_id_config.scope)
 
         except (ClientAuthenticationError, CredentialUnavailableError):
-            logger.error("Failed to retrieve Azure access token")
+            logger.warning("Failed to retrieve Azure access token")
             return None

src/llama_stack_configuration.py

@@ -8,6 +8,7 @@
 import logging
 import os
 from argparse import ArgumentParser
+from pathlib import Path
 from typing import Any
 
 from azure.core.exceptions import ClientAuthenticationError
@@ -62,14 +63,14 @@ def setup_azure_entra_id_token(
     tenant_id = azure_config.get("tenant_id")
     client_id = azure_config.get("client_id")
     client_secret = azure_config.get("client_secret")
+    scope = azure_config.get("scope", "https://cognitiveservices.azure.com/.default")
 
     if not all([tenant_id, client_id, client_secret]):
         logger.warning(
             "Azure Entra ID: Missing required fields (tenant_id, client_id, client_secret)"
         )
         return
 
-    scope = "https://cognitiveservices.azure.com/.default"
     try:
         credential = ClientSecretCredential(
             tenant_id=str(tenant_id),
@@ -80,10 +81,12 @@ def setup_azure_entra_id_token(
         token = credential.get_token(scope)
 
         # Write to .env file
+        # Create file if it doesn't exist
+        Path(env_file).touch()
+
         lines = []
-        if os.path.exists(env_file):
-            with open(env_file, "r", encoding="utf-8") as f:
-                lines = f.readlines()
+        with open(env_file, "r", encoding="utf-8") as f:
+            lines = f.readlines()
 
         # Update or add AZURE_API_KEY
         key_found = False
@@ -123,7 +126,7 @@ def construct_vector_dbs_section(
         ls_config (dict[str, Any]): Existing Llama Stack configuration mapping
         used as the base; existing `vector_dbs` entries are preserved if
         present.
-        byok_rag (list[ByokRag]): List of BYOK RAG definitions to be added to
+        byok_rag (list[dict[str, Any]]): List of BYOK RAG definitions to be added to
         the `vector_dbs` section.
 
     Returns:
@@ -143,10 +146,10 @@ def construct_vector_dbs_section(
     for brag in byok_rag:
         output.append(
             {
-                "vector_db_id": brag.get("vector_db_id"),
+                "vector_db_id": brag.get("vector_db_id", ""),
                 "provider_id": "byok_" + brag.get("vector_db_id", ""),
-                "embedding_model": brag.get("embedding_model", "all-MiniLM-L6-v2"),
-                "embedding_dimension": brag.get("embedding_dimension", 384),
+                "embedding_model": brag.get("embedding_model", ""),
+                "embedding_dimension": brag.get("embedding_dimension"),
             }
         )
     logger.info(
@@ -170,7 +173,7 @@ def construct_vector_io_providers_section(
         ls_config (dict[str, Any]): Existing Llama Stack configuration
         dictionary; if it contains providers.vector_io, those entries are used
         as the starting list.
-        byok_rag (list[ByokRag]): List of BYOK RAG specifications to convert
+        byok_rag (list[dict[str, Any]]): List of BYOK RAG specifications to convert
         into provider entries.
 
     Returns:
@@ -303,13 +306,9 @@ def main() -> None:
     )
     args = parser.parse_args()
 
-    try:
-        with open(args.config, "r", encoding="utf-8") as f:
-            config = yaml.safe_load(f)
-            config = replace_env_vars(config)
-    except FileNotFoundError:
-        logger.error("Config not found: %s", args.config)
-        return
+    with open(args.config, "r", encoding="utf-8") as f:
+        config = yaml.safe_load(f)
+        config = replace_env_vars(config)
 
     generate_configuration(args.input, args.output, config, args.env_file)
 

test.containerfile

@@ -13,7 +13,7 @@ RUN pip install faiss-cpu==1.11.0 azure-identity && \
 # Copy enrichment scripts for runtime config enrichment
 COPY src/llama_stack_configuration.py /opt/app-root/llama_stack_configuration.py
 COPY scripts/llama-stack-entrypoint.sh /opt/app-root/enrich-entrypoint.sh
-RUN chmod +x /opt/app-root/enrich-entrypoint.sh /opt/app-root/llama_stack_configuration.py && \
+RUN chmod +x /opt/app-root/enrich-entrypoint.sh && \
     chown 1001:0 /opt/app-root/enrich-entrypoint.sh /opt/app-root/llama_stack_configuration.py
 
 # Switch back to the original user

tests/unit/authorization/test_azure_token_manager.py

@@ -75,13 +75,12 @@ def test_token_expiration_logic(self, token_manager: AzureEntraIDManager) -> Non
         token_manager._expires_on = 0
         assert token_manager.is_token_expired
 
-    @pytest.mark.asyncio
-    async def test_refresh_token_raises_without_config(
+    def test_refresh_token_raises_without_config(
         self, token_manager: AzureEntraIDManager
     ) -> None:
         """Raise ValueError when refresh_token is called without config."""
         with pytest.raises(ValueError, match="Azure Entra ID configuration not set"):
-            await token_manager.refresh_token()
+            token_manager.refresh_token()
 
     def test_update_access_token_sets_token_and_expiration(
         self, token_manager: AzureEntraIDManager
@@ -92,8 +91,7 @@ def test_update_access_token_sets_token_and_expiration(
         assert token_manager.access_token.get_secret_value() == "test-token"
         assert token_manager._expires_on == expires_on - TOKEN_EXPIRATION_LEEWAY
 
-    @pytest.mark.asyncio
-    async def test_refresh_token_success(
+    def test_refresh_token_success(
         self,
         token_manager: AzureEntraIDManager,
         dummy_config: AzureEntraIdConfiguration,
@@ -111,16 +109,14 @@ async def test_refresh_token_success(
             return_value=mock_credential_instance,
         )
 
-        await token_manager.refresh_token()
+        result = token_manager.refresh_token()
 
+        assert result is True
         assert token_manager.access_token.get_secret_value() == "token_value"
         assert not token_manager.is_token_expired
-        mock_credential_instance.get_token.assert_called_once_with(
-            "https://cognitiveservices.azure.com/.default"
-        )
+        mock_credential_instance.get_token.assert_called_once_with(dummy_config.scope)
 
-    @pytest.mark.asyncio
-    async def test_refresh_token_failure_logs_error(
+    def test_refresh_token_failure_logs_error(
         self,
         token_manager: AzureEntraIDManager,
         dummy_config: AzureEntraIdConfiguration,
@@ -138,8 +134,9 @@ async def test_refresh_token_failure_logs_error(
             return_value=mock_credential_instance,
         )
 
-        with caplog.at_level("ERROR"):
-            await token_manager.refresh_token()
+        with caplog.at_level("WARNING"):
+            result = token_manager.refresh_token()
+            assert result is False
             assert "Failed to retrieve Azure access token" in caplog.text
 
     def test_token_expired_property_dynamic(

tests/unit/test_llama_stack_configuration.py

@@ -65,15 +65,6 @@ def test_construct_vector_dbs_section_adds_new() -> None:
     assert output[0]["embedding_dimension"] == 512
 
 
-def test_construct_vector_dbs_section_uses_defaults() -> None:
-    """Test uses default values when not specified."""
-    ls_config: dict[str, Any] = {}
-    byok_rag = [{"vector_db_id": "db1"}]
-    output = construct_vector_dbs_section(ls_config, byok_rag)
-    assert output[0]["embedding_model"] == "all-MiniLM-L6-v2"
-    assert output[0]["embedding_dimension"] == 384
-
-
 def test_construct_vector_dbs_section_merge() -> None:
     """Test merges existing and new entries."""
     ls_config = {"vector_dbs": [{"vector_db_id": "existing"}]}
@@ -119,14 +110,6 @@ def test_construct_vector_io_providers_section_adds_new() -> None:
     assert output[0]["provider_type"] == "inline::faiss"
 
 
-def test_construct_vector_io_providers_section_uses_defaults() -> None:
-    """Test uses default values when not specified."""
-    ls_config: dict[str, Any] = {"providers": {}}
-    byok_rag = [{"vector_db_id": "db1"}]
-    output = construct_vector_io_providers_section(ls_config, byok_rag)
-    assert output[0]["provider_type"] == "inline::faiss"
-
-
 # =============================================================================
 # Test generate_configuration
 # =============================================================================