Fix credential

Author: pamelafox

evals/generate_ground_truth.py

@@ -3,7 +3,7 @@
 from collections.abc import Generator
 from pathlib import Path
 
-from azure.identity import DefaultAzureCredential
+from azure.identity import AzureDeveloperCliCredential
 from dotenv import load_dotenv
 from evaltools.gen.generate import generate_test_qa_data
 from sqlalchemy import create_engine, select
@@ -56,7 +56,7 @@ def get_openai_config_dict() -> dict:
             api_key = os.environ["AZURE_OPENAI_KEY"]
         else:
             logger.info("Using Azure OpenAI Service with Azure Developer CLI Credential")
-            azure_credential = DefaultAzureCredential(process_timeout=60)
+            azure_credential = AzureDeveloperCliCredential(process_timeout=60)
             api_key = azure_credential.get_token("https://cognitiveservices.azure.com/.default").token
         openai_config = {
             "api_type": "azure",

src/backend/fastapi_app/postgres_engine.py

@@ -1,7 +1,7 @@
 import logging
 import os
 
-from azure.identity import DefaultAzureCredential
+from azure.identity import AzureDeveloperCliCredential
 from sqlalchemy import event
 from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
 
@@ -60,7 +60,7 @@ async def create_postgres_engine_from_env(azure_credential=None) -> AsyncEngine:
 
 async def create_postgres_engine_from_args(args, azure_credential=None) -> AsyncEngine:
     if azure_credential is None and args.host.endswith(".database.azure.com"):
-        azure_credential = DefaultAzureCredential(process_timeout=60)
+        azure_credential = AzureDeveloperCliCredential(process_timeout=60)
 
     return await create_postgres_engine(
         host=args.host,

tests/conftest.py

@@ -274,22 +274,22 @@ async def mock_acreate(*args, **kwargs):
 
 
 @pytest.fixture(scope="function")
-def mock_default_azure_credential(mock_session_env):
+def mock_azure_credential(mock_session_env):
     """Mock the Azure credential for testing."""
-    with mock.patch("azure.identity.DefaultAzureCredential") as mock_default_azure_credential:
-        mock_default_azure_credential.return_value = MockAzureCredential()
-        yield mock_default_azure_credential
+    with mock.patch("azure.identity.AzureDeveloperCliCredential") as mock_azure_credential:
+        mock_azure_credential.return_value = MockAzureCredential()
+        yield mock_azure_credential
 
 
 @pytest_asyncio.fixture(scope="function")
-async def test_client(app, mock_default_azure_credential, mock_openai_embedding, mock_openai_chatcompletion):
+async def test_client(app, mock_azure_credential, mock_openai_embedding, mock_openai_chatcompletion):
     """Create a test client."""
     with TestClient(app) as test_client:
         yield test_client
 
 
 @pytest_asyncio.fixture(scope="function")
-async def db_session(mock_session_env, mock_default_azure_credential):
+async def db_session(mock_session_env, mock_azure_credential):
     """Create a new database session with a rollback at the end of the test."""
     engine = await create_postgres_engine_from_env()
     async_sesion = async_sessionmaker(autocommit=False, autoflush=False, bind=engine)
@@ -302,10 +302,10 @@ async def db_session(mock_session_env, mock_default_azure_credential):
 
 
 @pytest_asyncio.fixture(scope="function")
-async def postgres_searcher(mock_session_env, mock_default_azure_credential, db_session, mock_openai_embedding):
+async def postgres_searcher(mock_session_env, mock_azure_credential, db_session, mock_openai_embedding):
     from fastapi_app.postgres_searcher import PostgresSearcher
 
-    openai_embed_client = await create_openai_embed_client(mock_default_azure_credential)
+    openai_embed_client = await create_openai_embed_client(mock_azure_credential)
 
     yield PostgresSearcher(
         db_session=db_session,

tests/test_dependencies.py

@@ -34,7 +34,7 @@ async def test_get_common_parameters_openai(mock_session_env_openai):
 
 
 @pytest.mark.asyncio
-async def test_get_azure_credential(mock_session_env, mock_default_azure_credential):
+async def test_get_azure_credential(mock_session_env, mock_azure_credential):
     result = await get_azure_credential()
     token = result.get_token("https://vault.azure.net")
     assert token.expires_on == 9999999999

tests/test_embeddings.py

@@ -6,8 +6,8 @@
 
 
 @pytest.mark.asyncio
-async def test_compute_text_embedding(mock_default_azure_credential, mock_openai_embedding):
-    openai_embed_client = await create_openai_embed_client(mock_default_azure_credential)
+async def test_compute_text_embedding(mock_azure_credential, mock_openai_embedding):
+    openai_embed_client = await create_openai_embed_client(mock_azure_credential)
     result = await compute_text_embedding(
         q="test",
         openai_client=openai_embed_client,
@@ -18,8 +18,8 @@ async def test_compute_text_embedding(mock_default_azure_credential, mock_openai
 
 
 @pytest.mark.asyncio
-async def test_compute_text_embedding_dimensions(mock_default_azure_credential, mock_openai_embedding):
-    openai_embed_client = await create_openai_embed_client(mock_default_azure_credential)
+async def test_compute_text_embedding_dimensions(mock_azure_credential, mock_openai_embedding):
+    openai_embed_client = await create_openai_embed_client(mock_azure_credential)
     result = await compute_text_embedding(
         q="test",
         openai_client=openai_embed_client,

tests/test_openai_clients.py

@@ -5,8 +5,8 @@
 
 
 @pytest.mark.asyncio
-async def test_create_openai_embed_client(mock_default_azure_credential, mock_openai_embedding):
-    openai_embed_client = await create_openai_embed_client(mock_default_azure_credential)
+async def test_create_openai_embed_client(mock_azure_credential, mock_openai_embedding):
+    openai_embed_client = await create_openai_embed_client(mock_azure_credential)
     assert openai_embed_client.embeddings.create is not None
     embeddings = await openai_embed_client.embeddings.create(
         model="text-embedding-ada-002", input="test", dimensions=1536
@@ -15,8 +15,8 @@ async def test_create_openai_embed_client(mock_default_azure_credential, mock_op
 
 
 @pytest.mark.asyncio
-async def test_create_openai_chat_client(mock_default_azure_credential, mock_openai_chatcompletion):
-    openai_chat_client = await create_openai_chat_client(mock_default_azure_credential)
+async def test_create_openai_chat_client(mock_azure_credential, mock_openai_chatcompletion):
+    openai_chat_client = await create_openai_chat_client(mock_azure_credential)
     assert openai_chat_client.chat.completions.create is not None
     response = await openai_chat_client.chat.completions.create(
         model="gpt-4o-mini", messages=[{"content": "test", "role": "user"}]

tests/test_postgres_engine.py

@@ -11,14 +11,14 @@
 
 
 @pytest.mark.asyncio
-async def test_create_postgres_engine(mock_session_env, mock_default_azure_credential):
+async def test_create_postgres_engine(mock_session_env, mock_azure_credential):
     engine = await create_postgres_engine(
         host=os.environ["POSTGRES_HOST"],
         username=os.environ["POSTGRES_USERNAME"],
         database=os.environ["POSTGRES_DATABASE"],
         password=os.environ.get("POSTGRES_PASSWORD"),
         sslmode=os.environ.get("POSTGRES_SSL"),
-        azure_credential=mock_default_azure_credential,
+        azure_credential=mock_azure_credential,
     )
     assert engine.url.host == "localhost"
     assert engine.url.username == "admin"
@@ -28,9 +28,9 @@ async def test_create_postgres_engine(mock_session_env, mock_default_azure_crede
 
 
 @pytest.mark.asyncio
-async def test_create_postgres_engine_from_env(mock_session_env, mock_default_azure_credential):
+async def test_create_postgres_engine_from_env(mock_session_env, mock_azure_credential):
     engine = await create_postgres_engine_from_env(
-        azure_credential=mock_default_azure_credential,
+        azure_credential=mock_azure_credential,
     )
     assert engine.url.host == "localhost"
     assert engine.url.username == "admin"
@@ -40,7 +40,7 @@ async def test_create_postgres_engine_from_env(mock_session_env, mock_default_az
 
 
 @pytest.mark.asyncio
-async def test_create_postgres_engine_from_args(mock_default_azure_credential):
+async def test_create_postgres_engine_from_args(mock_azure_credential):
     args = type(
         "Args",
         (),
@@ -54,7 +54,7 @@ async def test_create_postgres_engine_from_args(mock_default_azure_credential):
     )
     engine = await create_postgres_engine_from_args(
         args=args,
-        azure_credential=mock_default_azure_credential,
+        azure_credential=mock_azure_credential,
     )
     assert engine.url.host == "localhost"
     assert engine.url.username == "admin"