add initial test

Author: mattgotteiner

tests/conftest.py

@@ -11,7 +11,6 @@
 import pytest
 import pytest_asyncio
 from azure.search.documents.agent.aio import KnowledgeAgentRetrievalClient
-from azure.search.documents.agent.models import KnowledgeAgentRetrievalResponse
 from azure.search.documents.aio import SearchClient
 from azure.search.documents.indexes.aio import SearchIndexClient
 from azure.search.documents.indexes.models import (
@@ -61,7 +60,7 @@ async def mock_search(self, *args, **kwargs):
 
 
 async def mock_retrieve(self, *args, **kwargs):
-    return KnowledgeAgentRetrievalResponse()
+    return 
 
 
 @pytest.fixture
@@ -251,7 +250,6 @@ def patch(openai_client):
 
     return patch
 
-
 @pytest.fixture
 def mock_acs_search(monkeypatch):
     monkeypatch.setattr(SearchClient, "search", mock_search)

tests/mocks.py

@@ -10,6 +10,8 @@
     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"
@@ -207,6 +209,30 @@ 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."}]')],
+            )
+        ],
+        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)
+        ],
+        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"
+                }
+            )
+        ]
+    )
 
 class MockAudio:
     def __init__(self, audio_data):

tests/test_chatapproach.py

@@ -13,12 +13,16 @@
     MOCK_EMBEDDING_DIMENSIONS,
     MOCK_EMBEDDING_MODEL_NAME,
     MockAsyncSearchResultsIterator,
+    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()
+
 
 @pytest.fixture
 def chat_approach():
@@ -27,7 +31,7 @@ def chat_approach():
         search_index_name=None,
         agent_model=None,
         agent_deployment=None,
-        agent_client=KnowledgeAgentRetrievalClient(endpoint="", agent_name="", credential=AzureKeyCredential("")),
+        agent_client=None,
         auth_helper=None,
         openai_client=None,
         chatgpt_model="gpt-4o-mini",
@@ -178,7 +182,7 @@ async def test_search_results_filtering_by_scores(
         search_index_name=None,
         agent_model=None,
         agent_deployment=None,
-        agent_client=KnowledgeAgentRetrievalClient(endpoint="", agent_name="", credential=AzureKeyCredential("")),
+        agent_client=None,
         auth_helper=None,
         openai_client=None,
         chatgpt_model="gpt-4o-mini",
@@ -221,7 +225,7 @@ async def test_search_results_query_rewriting(monkeypatch):
         search_index_name=None,
         agent_model=None,
         agent_deployment=None,
-        agent_client=KnowledgeAgentRetrievalClient(endpoint="", agent_name="", credential=AzureKeyCredential("")),
+        agent_client=None,
         auth_helper=None,
         openai_client=None,
         chatgpt_model="gpt-35-turbo",
@@ -259,3 +263,42 @@ 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(
+        search_client=None,
+        search_index_name=None,
+        agent_model=None,
+        agent_deployment=None,
+        agent_client=None,
+        auth_helper=None,
+        openai_client=None,
+        chatgpt_model="gpt-35-turbo",
+        chatgpt_deployment="chat",
+        embedding_deployment="embeddings",
+        embedding_model=MOCK_EMBEDDING_MODEL_NAME,
+        embedding_dimensions=MOCK_EMBEDDING_DIMENSIONS,
+        embedding_field="embedding3",
+        sourcepage_field="",
+        content_field="",
+        query_language="en-us",
+        query_speller="lexicon",
+        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=""
+    )
+
+    assert len(results) == 1
+    assert results[0].id == "Benefit_Options-2.pdf"
+    assert results[0].content == "There is a whistleblower policy."
+    assert results[0].sourcepage == "Benefit_Options-2.pdf"
+    assert results[0].search_agent_query == "whistleblower query"