run make fix

Author: ks6088ts

template_langgraph/llms/azure_openais.py

@@ -1,6 +1,5 @@
-from functools import lru_cache
 import threading
-from typing import Optional
+from functools import lru_cache
 
 from azure.identity import DefaultAzureCredential
 from langchain_openai import AzureChatOpenAI, AzureOpenAIEmbeddings
@@ -37,38 +36,40 @@ class AzureOpenAiWrapper:
     _credentials: dict = {}
     _tokens: dict = {}
     _token_lock = threading.Lock()
-    
+
     def __init__(self, settings: Settings = None):
         if settings is None:
             settings = get_azure_openai_settings()
-        
+
         self.settings = settings
-        self._chat_model: Optional[AzureChatOpenAI] = None
-        self._reasoning_model: Optional[AzureChatOpenAI] = None
-        self._embedding_model: Optional[AzureOpenAIEmbeddings] = None
+        self._chat_model: AzureChatOpenAI | None = None
+        self._reasoning_model: AzureChatOpenAI | None = None
+        self._embedding_model: AzureOpenAIEmbeddings | None = None
 
     def _get_auth_key(self) -> str:
         """Generate a key for authentication caching based on settings."""
         return f"{self.settings.azure_openai_endpoint}_{self.settings.azure_openai_use_microsoft_entra_id}"
 
-    def _get_auth_token(self) -> Optional[str]:
+    def _get_auth_token(self) -> str | None:
         """Get authentication token with lazy initialization and caching."""
         if self.settings.azure_openai_use_microsoft_entra_id.lower() != "true":
             return None
-            
+
         auth_key = self._get_auth_key()
-        
+
         with self._token_lock:
             if auth_key not in self._credentials:
                 logger.info("Initializing Microsoft Entra ID authentication")
                 self._credentials[auth_key] = DefaultAzureCredential()
-            
+
             if auth_key not in self._tokens:
                 logger.info("Getting authentication token")
-                self._tokens[auth_key] = self._credentials[auth_key].get_token("https://cognitiveservices.azure.com/.default").token
-            
+                self._tokens[auth_key] = (
+                    self._credentials[auth_key].get_token("https://cognitiveservices.azure.com/.default").token
+                )
+
             return self._tokens[auth_key]
-    
+
     @property
     def chat_model(self) -> AzureChatOpenAI:
         """Lazily initialize and return chat model."""
@@ -92,7 +93,7 @@ def chat_model(self) -> AzureChatOpenAI:
                     streaming=True,
                 )
         return self._chat_model
-    
+
     @property
     def reasoning_model(self) -> AzureChatOpenAI:
         """Lazily initialize and return reasoning model."""
@@ -115,7 +116,7 @@ def reasoning_model(self) -> AzureChatOpenAI:
                     streaming=True,
                 )
         return self._reasoning_model
-    
+
     @property
     def embedding_model(self) -> AzureOpenAIEmbeddings:
         """Lazily initialize and return embedding model."""

tests/test_azure_openais.py

@@ -1,8 +1,6 @@
 import logging
 import threading
-import time
 from unittest.mock import Mock, patch
-import pytest
 
 from template_langgraph.llms.azure_openais import AzureOpenAiWrapper, Settings
 
@@ -26,8 +24,7 @@ def test_lazy_initialization_api_key(self, caplog):
         with caplog.at_level(logging.INFO):
             # Creating instances should not trigger authentication
             wrapper1 = AzureOpenAiWrapper(settings)
-            wrapper2 = AzureOpenAiWrapper(settings)
-            
+
             # No authentication logs yet
             assert "Using API key for authentication" not in caplog.text
 
@@ -49,7 +46,7 @@ def test_lazy_initialization_api_key(self, caplog):
             # Should still be only one authentication log per model type
             assert caplog.text.count("Using API key for authentication") >= 1
 
-    @patch('template_langgraph.llms.azure_openais.DefaultAzureCredential')
+    @patch("template_langgraph.llms.azure_openais.DefaultAzureCredential")
     def test_singleton_credential_entra_id(self, mock_credential_class, caplog):
         """Test that Microsoft Entra ID credentials are reused across instances."""
         # Mock the credential and token
@@ -85,7 +82,7 @@ def test_singleton_credential_entra_id(self, mock_credential_class, caplog):
             assert caplog.text.count("Initializing Microsoft Entra ID authentication") == 1
             assert caplog.text.count("Getting authentication token") == 1
 
-    @patch('template_langgraph.llms.azure_openais.DefaultAzureCredential')
+    @patch("template_langgraph.llms.azure_openais.DefaultAzureCredential")
     def test_thread_safety(self, mock_credential_class):
         """Test that authentication is thread-safe."""
         mock_credential = Mock()
@@ -174,10 +171,10 @@ def test_create_embedding_method_compatibility(self):
             pass  # Expected due to dummy credentials
 
         # Verify the method exists and is callable
-        assert hasattr(wrapper, 'create_embedding')
-        assert callable(getattr(wrapper, 'create_embedding'))
+        assert hasattr(wrapper, "create_embedding")
+        assert callable(getattr(wrapper, "create_embedding"))
 
-    @patch('template_langgraph.llms.azure_openais.DefaultAzureCredential')
+    @patch("template_langgraph.llms.azure_openais.DefaultAzureCredential")
     def test_mixed_authentication_methods(self, mock_credential_class, caplog):
         """Test using both authentication methods in different instances."""
         mock_credential = Mock()
@@ -212,4 +209,4 @@ def test_mixed_authentication_methods(self, mock_credential_class, caplog):
 
             # Should see both authentication methods being used
             assert "Using API key for authentication" in caplog.text
-            assert "Initializing Microsoft Entra ID authentication" in caplog.text
+            assert "Initializing Microsoft Entra ID authentication" in caplog.text