From 6747897a5398d37d6d220da1efbaf5e78e7b54d8 Mon Sep 17 00:00:00 2001 From: Abdul Basit Anees Date: Mon, 25 Aug 2025 12:11:18 +0000 Subject: [PATCH] add utility model support and add test --- .../team_agent_factory/inspector_factory.py | 2 +- .../team_agent/inspector_functional_test.py | 69 ++++++++++++++++++- 2 files changed, 69 insertions(+), 2 deletions(-) diff --git a/aixplain/factories/team_agent_factory/inspector_factory.py b/aixplain/factories/team_agent_factory/inspector_factory.py index 65b44dd5..52b5239c 100644 --- a/aixplain/factories/team_agent_factory/inspector_factory.py +++ b/aixplain/factories/team_agent_factory/inspector_factory.py @@ -31,7 +31,7 @@ from aixplain.utils.file_utils import _request_with_retry -INSPECTOR_SUPPORTED_FUNCTIONS = [Function.GUARDRAILS, Function.TEXT_GENERATION] +INSPECTOR_SUPPORTED_FUNCTIONS = [Function.GUARDRAILS, Function.TEXT_GENERATION, Function.UTILITIES] class InspectorFactory: diff --git a/tests/functional/team_agent/inspector_functional_test.py b/tests/functional/team_agent/inspector_functional_test.py index 97eecd8c..72525c5d 100644 --- a/tests/functional/team_agent/inspector_functional_test.py +++ b/tests/functional/team_agent/inspector_functional_test.py @@ -10,7 +10,7 @@ import pytest from aixplain import aixplain_v2 as v2 -from aixplain.factories import AgentFactory, TeamAgentFactory +from aixplain.factories import AgentFactory, TeamAgentFactory, ModelFactory from aixplain.enums.asset_status import AssetStatus from aixplain.modules.team_agent import InspectorTarget from aixplain.modules.team_agent.inspector import Inspector, InspectorPolicy, InspectorAction, InspectorOutput @@ -737,3 +737,70 @@ def process_response(model_response: ModelResponse, input_content: str) -> Inspe assert len(response_generator_steps) == 1, "Response generator should run exactly once after ABORT" team_agent.delete() + + +@pytest.mark.parametrize("TeamAgentFactory", [TeamAgentFactory, v2.TeamAgent]) +def test_team_agent_with_utility_inspector(run_input_map, delete_agents_and_team_agents, TeamAgentFactory): + """Test team agent with a Utility model as inspector""" + assert delete_agents_and_team_agents + + agents = create_agents_from_input_map(run_input_map) + + def lowercase_inspector(content: str) -> bool: + if content.islower(): + return "" + else: + return "Content is not all lowercase. There are uppercase characters in the content." + + utility_model = ModelFactory.create_utility_model( + name="Lowercase Inspector Test", + description="Inspect the content of the response. If the content is not all lowercase, provide feedback.'", + code=lowercase_inspector, + ) + utility_model.deploy() + + utility_model_id = utility_model.id + inspector = Inspector( + name="utility_inspector", + model_id=utility_model_id, + policy=InspectorPolicy.WARN, + ) + + # Create team agent with steps inspector + team_agent = create_team_agent( + TeamAgentFactory, + agents, + run_input_map, + use_mentalist=True, + inspectors=[inspector], + inspector_targets=[InspectorTarget.STEPS], + ) + + assert team_agent is not None + assert team_agent.status == AssetStatus.DRAFT + + # deploy team agent + team_agent.deploy() + team_agent = TeamAgentFactory.get(team_agent.id) + assert team_agent is not None + assert team_agent.status == AssetStatus.ONBOARDED + + # Run the team agent + response = team_agent.run(data=run_input_map["query"]) + + assert response is not None + assert response["completed"] is True + assert response["status"].lower() == "success" + + # Check for inspector steps + if "intermediate_steps" in response["data"]: + steps = response["data"]["intermediate_steps"] + verify_inspector_steps(steps, ["utility_inspector"], [InspectorTarget.STEPS]) + verify_response_generator(steps) + + # Verify inspector runs and execution continues + inspector_steps = [step for step in steps if "utility_inspector" in step.get("agent", "").lower()] + assert len(inspector_steps) > 0, "Utility inspector should run at least once" + + utility_model.delete() + team_agent.delete()