Custom inspector policy

aixplain/factories/team_agent_factory/inspector_factory.py

@@ -19,7 +19,7 @@
 """
 
 import logging
-from typing import Dict, Optional, Text, Union
+from typing import Dict, Optional, Text, Union, Callable
 from urllib.parse import urljoin
 
 from aixplain.enums.asset_status import AssetStatus
@@ -48,7 +48,7 @@ def create_from_model(
         name: Text,
         model: Union[Text, Model],
         model_config: Optional[Dict] = None,
-        policy: InspectorPolicy = InspectorPolicy.ADAPTIVE,  # default: doing something dynamically
+        policy: Union[InspectorPolicy, Callable] = InspectorPolicy.ADAPTIVE,  # default: doing something dynamically
     ) -> Inspector:
         """Create a new inspector agent from an onboarded model.
 
@@ -62,11 +62,10 @@ def create_from_model(
                 to use for the inspector.
             model_config (Optional[Dict], optional): Configuration parameters for
                 the inspector model (e.g., prompts, thresholds). Defaults to None.
-            policy (InspectorPolicy, optional): Action to take upon negative feedback:
-                - WARN: Log warning but continue execution
-                - ABORT: Stop execution on negative feedback
-                - ADAPTIVE: Dynamically decide based on context
-                Defaults to InspectorPolicy.ADAPTIVE.
+            policy: Action to take upon negative feedback (WARN/ABORT/ADAPTIVE)
+                or a callable function. If callable, must have name "process_response",
+                arguments "model_response" and "input_content" (both strings), and
+                return InspectorAction. Defaults to ADAPTIVE.
 
         Returns:
             Inspector: Created and configured inspector agent.
@@ -124,7 +123,7 @@ def create_auto(
         cls,
         auto: InspectorAuto,
         name: Optional[Text] = None,
-        policy: InspectorPolicy = InspectorPolicy.ADAPTIVE,
+        policy: Union[InspectorPolicy, Callable] = InspectorPolicy.ADAPTIVE,
     ) -> Inspector:
         """Create a new inspector agent using automatic configuration.
 
@@ -136,11 +135,10 @@ def create_auto(
             auto (InspectorAuto): Pre-configured automatic inspector instance.
             name (Optional[Text], optional): Name for the inspector. If not provided,
                 uses the name from the auto configuration. Defaults to None.
-            policy (InspectorPolicy, optional): Action to take upon negative feedback:
-                - WARN: Log warning but continue execution
-                - ABORT: Stop execution on negative feedback
-                - ADAPTIVE: Dynamically decide based on context
-                Defaults to InspectorPolicy.ADAPTIVE.
+            policy: Action to take upon negative feedback (WARN/ABORT/ADAPTIVE)
+                or a callable function. If callable, must have name "process_response",
+                arguments "model_response" and "input_content" (both strings), and
+                return InspectorAction. Defaults to ADAPTIVE.
 
         Returns:
             Inspector: Created and configured inspector agent using automatic

aixplain/factories/team_agent_factory/utils.py

@@ -67,9 +67,21 @@ def build_team_agent(payload: Dict, agents: List[Agent] = None, api_key: Text =
                 continue
 
     # Ensure custom classes are instantiated: for compatibility with backend return format
-    inspectors = [
-        inspector if isinstance(inspector, Inspector) else Inspector(**inspector) for inspector in payload.get("inspectors", [])
-    ]
+    inspectors = []
+    for inspector_data in payload.get("inspectors", []):
+        try:
+            if isinstance(inspector_data, Inspector):
+                inspectors.append(inspector_data)
+            else:
+                # Handle both old format and new format with policy_type
+                if hasattr(Inspector, "model_validate"):
+                    inspectors.append(Inspector.model_validate(inspector_data))
+                else:
+                    inspectors.append(Inspector(**inspector_data))
+        except Exception as e:
+            logging.warning(f"Failed to create inspector from data: {e}")
+            continue
+
     inspector_targets = [InspectorTarget(target.lower()) for target in payload.get("inspectorTargets", [])]
 
     # Get LLMs from tools if present

aixplain/modules/team_agent/__init__.py

@@ -540,14 +540,17 @@ def from_dict(cls, data: Dict) -> "TeamAgent":
         if "inspectors" in data:
             for inspector_data in data["inspectors"]:
                 try:
-                    if hasattr(Inspector, "model_validate"):
+                    if isinstance(inspector_data, Inspector):
+                        inspectors.append(inspector_data)
+                    elif hasattr(Inspector, "model_validate"):
                         inspectors.append(Inspector.model_validate(inspector_data))
                     else:
                         inspectors.append(Inspector(**inspector_data))
                 except Exception as e:
                     import logging
 
                     logging.warning(f"Failed to create inspector from data: {e}")
+                    continue
 
         # Extract inspector targets
         inspector_targets = [InspectorTarget.STEPS]  # default