code review feedback

Author: paulbatum

sdk/ai/azure-ai-projects/tests/agents/tools/multitool/test_agent_file_search_and_code_interpreter.py

@@ -19,6 +19,7 @@
 from devtools_testutils import is_live_and_not_recording
 from azure.ai.projects.models import PromptAgentDefinition, FileSearchTool, CodeInterpreterTool, CodeInterpreterToolAuto
 
+
 class TestAgentFileSearchAndCodeInterpreter(TestBase):
     """Tests for agents using File Search + Code Interpreter combination."""
 
@@ -40,7 +41,7 @@ def test_find_and_analyze_data(self, **kwargs):
 
         # Create data file
         txt_content = "Sample data: 10, 20, 30, 40, 50"
-        vector_store = openai_client.vector_stores.create(name="DataStore")        
+        vector_store = openai_client.vector_stores.create(name="DataStore")
 
         txt_file = BytesIO(txt_content.encode("utf-8"))
         txt_file.name = "data.txt"
@@ -109,8 +110,6 @@ def test_analyze_code_file(self, **kwargs):
 
         vector_store = openai_client.vector_stores.create(name="CodeAnalysisStore")
 
-        from io import BytesIO
-
         code_file = BytesIO(python_code.encode("utf-8"))
         code_file.name = "fibonacci.py"
 

sdk/ai/azure-ai-projects/tests/agents/tools/multitool/test_agent_file_search_and_function.py

@@ -62,8 +62,6 @@ def test_data_analysis_workflow(self, **kwargs):
         vector_store = openai_client.vector_stores.create(name="SalesDataStore")
         print(f"Vector store created (id: {vector_store.id})")
 
-        from io import BytesIO
-
         txt_file = BytesIO(txt_content.encode("utf-8"))
         txt_file.name = "sales_data.txt"
 
@@ -269,8 +267,6 @@ def calculate_sum(numbers):
         # Create vector store and upload
         vector_store = openai_client.vector_stores.create(name="CodeStore")
 
-        from io import BytesIO
-
         code_file = BytesIO(python_code.encode("utf-8"))
         code_file.name = "sample_code.py"
 

sdk/ai/azure-ai-projects/tests/agents/tools/multitool/test_agent_file_search_code_interpreter_function.py

@@ -15,6 +15,7 @@
 import os
 import json
 import pytest
+from io import BytesIO
 from test_base import TestBase, servicePreparer
 from devtools_testutils import is_live_and_not_recording
 from azure.ai.projects.models import (
@@ -50,8 +51,6 @@ def test_complete_analysis_workflow(self, **kwargs):
         txt_content = "Sample data for analysis"
         vector_store = openai_client.vector_stores.create(name="ThreeToolStore")
 
-        from io import BytesIO
-
         txt_file = BytesIO(txt_content.encode("utf-8"))
         txt_file.name = "data.txt"
 
@@ -126,8 +125,6 @@ def test_four_tools_combination(self, **kwargs):
         txt_content = "Test data"
         vector_store = openai_client.vector_stores.create(name="FourToolStore")
 
-        from io import BytesIO
-
         txt_file = BytesIO(txt_content.encode("utf-8"))
         txt_file.name = "data.txt"
 

sdk/ai/azure-ai-projects/tests/agents/tools/multitool/test_multitool_with_conversations.py

@@ -12,6 +12,7 @@
 
 import json
 import pytest
+from io import BytesIO
 from test_base import TestBase, servicePreparer
 from devtools_testutils import is_live_and_not_recording
 from azure.ai.projects.models import (
@@ -57,7 +58,6 @@ def test_file_search_and_function_with_conversation(self, **kwargs):
 """
 
         vector_store = openai_client.vector_stores.create(name="SalesDataStore")
-        from io import BytesIO
 
         file = BytesIO(doc_content.encode("utf-8"))
         file.name = "sales.txt"

sdk/ai/azure-ai-projects/tests/agents/tools/test_agent_ai_search.py

@@ -80,10 +80,6 @@ def test_agent_ai_search_question_answering(self, **kwargs):
 
         model = self.test_agents_params["model_deployment_name"]
 
-        # Setup
-        project_client = self.create_client(operation_group="agents", **kwargs)
-        openai_client = project_client.get_openai_client()
-
         # Get AI Search connection and index from environment
         ai_search_connection_id = kwargs.get("azure_ai_projects_tests_ai_search_project_connection_id")
         ai_search_index_name = kwargs.get("azure_ai_projects_tests_ai_search_index_name")
@@ -97,114 +93,118 @@ def test_agent_ai_search_question_answering(self, **kwargs):
         assert isinstance(ai_search_connection_id, str), "ai_search_connection_id must be a string"
         assert isinstance(ai_search_index_name, str), "ai_search_index_name must be a string"
 
-        # Create agent with Azure AI Search tool
-        agent = project_client.agents.create_version(
-            agent_name="ai-search-qa-agent",
-            definition=PromptAgentDefinition(
-                model=model,
-                instructions="""You are a helpful assistant that answers true/false questions based on the provided search results.
+        with (
+            self.create_client(operation_group="agents", **kwargs) as project_client,
+            project_client.get_openai_client() as openai_client,
+        ):
+            # Create agent with Azure AI Search tool
+            agent = project_client.agents.create_version(
+                agent_name="ai-search-qa-agent",
+                definition=PromptAgentDefinition(
+                    model=model,
+                    instructions="""You are a helpful assistant that answers true/false questions based on the provided search results.
                 Always use the Azure AI Search tool to find relevant information before answering.
                 Respond with only 'True' or 'False' based on what you find in the search results.
                 If you cannot find clear evidence in the search results, answer 'False'.""",
-                tools=[
-                    AzureAISearchAgentTool(
-                        azure_ai_search=AzureAISearchToolResource(
-                            indexes=[
-                                AISearchIndexResource(
-                                    project_connection_id=ai_search_connection_id,
-                                    index_name=ai_search_index_name,
-                                    query_type=AzureAISearchQueryType.SIMPLE,
-                                ),
-                            ]
+                    tools=[
+                        AzureAISearchAgentTool(
+                            azure_ai_search=AzureAISearchToolResource(
+                                indexes=[
+                                    AISearchIndexResource(
+                                        project_connection_id=ai_search_connection_id,
+                                        index_name=ai_search_index_name,
+                                        query_type=AzureAISearchQueryType.SIMPLE,
+                                    ),
+                                ]
+                            )
                         )
-                    )
-                ],
-            ),
-            description="Agent for testing AI Search question answering.",
-        )
-        print(f"Agent created (id: {agent.id}, name: {agent.name}, version: {agent.version})")
-        assert agent.id is not None
-        assert agent.name == "ai-search-qa-agent"
-        assert agent.version is not None
-
-        # Test each question
-        correct_answers = 0
-        total_questions = len(self.TEST_QUESTIONS)
-
-        for i, qa_pair in enumerate(self.TEST_QUESTIONS, 1):
-            question = qa_pair["question"]
-            expected_answer = qa_pair["answer"]
+                    ],
+                ),
+                description="Agent for testing AI Search question answering.",
+            )
+            print(f"Agent created (id: {agent.id}, name: {agent.name}, version: {agent.version})")
+            assert agent.id is not None
+            assert agent.name == "ai-search-qa-agent"
+            assert agent.version is not None
+
+            # Test each question
+            correct_answers = 0
+            total_questions = len(self.TEST_QUESTIONS)
+
+            for i, qa_pair in enumerate(self.TEST_QUESTIONS, 1):
+                question = qa_pair["question"]
+                expected_answer = qa_pair["answer"]
+
+                print(f"\n{'='*80}")
+                print(f"Question {i}/{total_questions}:")
+                print(f"Q: {question}")
+                print(f"Expected: {expected_answer}")
+
+                output_text = ""
+
+                stream_response = openai_client.responses.create(
+                    stream=True,
+                    tool_choice="required",
+                    input=f"Answer this question with only 'True' or 'False': {question}",
+                    extra_body={"agent": {"name": agent.name, "type": "agent_reference"}},
+                )
+
+                for event in stream_response:
+                    if event.type == "response.created":
+                        print(f"Response created with ID: {event.response.id}")
+                    elif event.type == "response.output_text.delta":
+                        pass  # Don't print deltas to reduce output
+                    elif event.type == "response.completed":
+                        output_text = event.response.output_text
+                        print(f"Agent's answer: {output_text}")
+
+                # Parse the answer from the output
+                # Look for "True" or "False" in the response
+                output_lower = output_text.lower()
+                agent_answer = None
+
+                # Try to extract boolean answer
+                if "true" in output_lower and "false" not in output_lower:
+                    agent_answer = True
+                elif "false" in output_lower and "true" not in output_lower:
+                    agent_answer = False
+                elif output_lower.strip() in ["true", "false"]:
+                    agent_answer = output_lower.strip() == "true"
+                else:
+                    # Try to determine based on more complex responses
+                    # Count occurrences
+                    true_count = output_lower.count("true")
+                    false_count = output_lower.count("false")
+                    if true_count > false_count:
+                        agent_answer = True
+                    elif false_count > true_count:
+                        agent_answer = False
+
+                if agent_answer is not None:
+                    is_correct = agent_answer == expected_answer
+                    if is_correct:
+                        correct_answers += 1
+                        print(f"✓ CORRECT (Agent: {agent_answer}, Expected: {expected_answer})")
+                    else:
+                        print(f"✗ INCORRECT (Agent: {agent_answer}, Expected: {expected_answer})")
+                else:
+                    print(f"✗ UNABLE TO PARSE ANSWER from: {output_text}")
 
+            # Print summary
             print(f"\n{'='*80}")
-            print(f"Question {i}/{total_questions}:")
-            print(f"Q: {question}")
-            print(f"Expected: {expected_answer}")
+            print(f"SUMMARY: {correct_answers}/{total_questions} questions answered correctly")
+            print(f"{'='*80}")
 
-            output_text = ""
-
-            stream_response = openai_client.responses.create(
-                stream=True,
-                tool_choice="required",
-                input=f"Answer this question with only 'True' or 'False': {question}",
-                extra_body={"agent": {"name": agent.name, "type": "agent_reference"}},
+            # Verify that at least 4 out of 5 questions were answered correctly
+            assert correct_answers >= 4, (
+                f"Expected at least 4 correct answers out of {total_questions}, "
+                f"but got {correct_answers}. The agent needs to answer at least 80% correctly."
             )
 
-            for event in stream_response:
-                if event.type == "response.created":
-                    print(f"Response created with ID: {event.response.id}")
-                elif event.type == "response.output_text.delta":
-                    pass  # Don't print deltas to reduce output
-                elif event.type == "response.completed":
-                    output_text = event.response.output_text
-                    print(f"Agent's answer: {output_text}")
-
-            # Parse the answer from the output
-            # Look for "True" or "False" in the response
-            output_lower = output_text.lower()
-            agent_answer = None
-
-            # Try to extract boolean answer
-            if "true" in output_lower and "false" not in output_lower:
-                agent_answer = True
-            elif "false" in output_lower and "true" not in output_lower:
-                agent_answer = False
-            elif output_lower.strip() in ["true", "false"]:
-                agent_answer = output_lower.strip() == "true"
-            else:
-                # Try to determine based on more complex responses
-                # Count occurrences
-                true_count = output_lower.count("true")
-                false_count = output_lower.count("false")
-                if true_count > false_count:
-                    agent_answer = True
-                elif false_count > true_count:
-                    agent_answer = False
+            print(
+                f"\n✓ Test passed! Agent answered {correct_answers}/{total_questions} questions correctly (>= 4 required)"
+            )
 
-            if agent_answer is not None:
-                is_correct = agent_answer == expected_answer
-                if is_correct:
-                    correct_answers += 1
-                    print(f"✓ CORRECT (Agent: {agent_answer}, Expected: {expected_answer})")
-                else:
-                    print(f"✗ INCORRECT (Agent: {agent_answer}, Expected: {expected_answer})")
-            else:
-                print(f"✗ UNABLE TO PARSE ANSWER from: {output_text}")
-
-        # Print summary
-        print(f"\n{'='*80}")
-        print(f"SUMMARY: {correct_answers}/{total_questions} questions answered correctly")
-        print(f"{'='*80}")
-
-        # Verify that at least 4 out of 5 questions were answered correctly
-        assert correct_answers >= 4, (
-            f"Expected at least 4 correct answers out of {total_questions}, "
-            f"but got {correct_answers}. The agent needs to answer at least 80% correctly."
-        )
-
-        print(
-            f"\n✓ Test passed! Agent answered {correct_answers}/{total_questions} questions correctly (>= 4 required)"
-        )
-
-        # Teardown
-        project_client.agents.delete_version(agent_name=agent.name, agent_version=agent.version)
-        print("Agent deleted")
+            # Teardown
+            project_client.agents.delete_version(agent_name=agent.name, agent_version=agent.version)
+            print("Agent deleted")