ruff, black

Author: mattgotteiner

tests/conftest.py

@@ -60,7 +60,7 @@ async def mock_search(self, *args, **kwargs):
 
 
 async def mock_retrieve(self, *args, **kwargs):
-    return 
+    return
 
 
 @pytest.fixture
@@ -250,6 +250,7 @@ def patch(openai_client):
 
     return patch
 
+
 @pytest.fixture
 def mock_acs_search(monkeypatch):
     monkeypatch.setattr(SearchClient, "search", mock_search)

tests/mocks.py

@@ -6,12 +6,19 @@
 import openai.types
 from azure.cognitiveservices.speech import ResultReason
 from azure.core.credentials_async import AsyncTokenCredential
+from azure.search.documents.agent.models import (
+    KnowledgeAgentAzureSearchDocReference,
+    KnowledgeAgentMessage,
+    KnowledgeAgentMessageTextContent,
+    KnowledgeAgentModelQueryPlanningActivityRecord,
+    KnowledgeAgentRetrievalResponse,
+    KnowledgeAgentSearchActivityRecord,
+    KnowledgeAgentSearchActivityRecordQuery,
+)
 from azure.search.documents.models import (
     VectorQuery,
 )
 from azure.storage.blob import BlobProperties
-from azure.search.documents.agent.models import KnowledgeAgentRetrievalResponse, KnowledgeAgentMessage, KnowledgeAgentMessageTextContent, KnowledgeAgentModelQueryPlanningActivityRecord, KnowledgeAgentSearchActivityRecord, KnowledgeAgentSearchActivityRecordQuery, KnowledgeAgentAzureSearchDocReference
-
 
 MOCK_EMBEDDING_DIMENSIONS = 1536
 MOCK_EMBEDDING_MODEL_NAME = "text-embedding-ada-002"
@@ -209,31 +216,40 @@ def mock_computervision_response():
         ),
     )
 
+
 def mock_retrieval_response():
     return KnowledgeAgentRetrievalResponse(
         response=[
             KnowledgeAgentMessage(
                 role="assistant",
-                content=[KnowledgeAgentMessageTextContent(text=r'[{"ref_id":0,"title":"Benefit_Options-2.pdf","content":"There is a whistleblower policy."}]')],
+                content=[
+                    KnowledgeAgentMessageTextContent(
+                        text=r'[{"ref_id":0,"title":"Benefit_Options-2.pdf","content":"There is a whistleblower policy."}]'
+                    )
+                ],
             )
         ],
         activity=[
             KnowledgeAgentModelQueryPlanningActivityRecord(id=0, input_tokens=10, output_tokens=20, elapsed_ms=200),
-            KnowledgeAgentSearchActivityRecord(id=1, target_index="index", query=KnowledgeAgentSearchActivityRecordQuery(search="whistleblower query"),count=10, elapsed_ms=50)
+            KnowledgeAgentSearchActivityRecord(
+                id=1,
+                target_index="index",
+                query=KnowledgeAgentSearchActivityRecordQuery(search="whistleblower query"),
+                count=10,
+                elapsed_ms=50,
+            ),
         ],
         references=[
             KnowledgeAgentAzureSearchDocReference(
                 id=0,
                 activity_source=1,
                 doc_key="Benefit_Options-2.pdf",
-                source_data={
-                    "content": "There is a whistleblower policy.",
-                    "sourcepage": "Benefit_Options-2.pdf"
-                }
+                source_data={"content": "There is a whistleblower policy.", "sourcepage": "Benefit_Options-2.pdf"},
             )
-        ]
+        ],
     )
 
+
 class MockAudio:
     def __init__(self, audio_data):
         self.audio_data = audio_data

tests/test_chatapproach.py

@@ -13,13 +13,14 @@
     MOCK_EMBEDDING_DIMENSIONS,
     MOCK_EMBEDDING_MODEL_NAME,
     MockAsyncSearchResultsIterator,
-    mock_retrieval_response
+    mock_retrieval_response,
 )
 
 
 async def mock_search(*args, **kwargs):
     return MockAsyncSearchResultsIterator(kwargs.get("search_text"), kwargs.get("vector_queries"))
 
+
 async def mock_retrieval(*args, **kwargs):
     return mock_retrieval_response()
 
@@ -264,6 +265,7 @@ async def validate_qr_and_mock_search(*args, **kwargs):
     assert len(results) == 1
     assert query_rewrites == "generative"
 
+
 @pytest.mark.asyncio
 async def test_agent_retrieval_results(monkeypatch):
     chat_approach = ChatReadRetrieveReadApproach(
@@ -284,18 +286,14 @@ async def test_agent_retrieval_results(monkeypatch):
         content_field="",
         query_language="en-us",
         query_speller="lexicon",
-        prompt_manager=PromptyManager()
+        prompt_manager=PromptyManager(),
     )
 
     agent_client = KnowledgeAgentRetrievalClient(endpoint="", agent_name="", credential=AzureKeyCredential(""))
 
     monkeypatch.setattr(KnowledgeAgentRetrievalClient, "retrieve", mock_retrieval)
 
-    _, results = await chat_approach.run_agentic_retrieval(
-        messages=[],
-        agent_client=agent_client,
-        search_index_name=""
-    )
+    _, results = await chat_approach.run_agentic_retrieval(messages=[], agent_client=agent_client, search_index_name="")
 
     assert len(results) == 1
     assert results[0].id == "Benefit_Options-2.pdf"