ENG 2333: Add Utility Model Support

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:

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()