chore: add the tags returned by the service to the ai_guard span

Author: manuel-alvarez-alvarez

ddtrace/appsec/ai_guard/_api_client.py

@@ -224,6 +224,7 @@ def evaluate(self, messages: List[Message], options: Optional[Options] = None) -
                         attributes = result["data"]["attributes"]
                         action = attributes["action"]
                         reason = attributes.get("reason", None)
+                        tags = attributes.get("tags", [])
                         blocking_enabled = attributes.get("is_blocking_enabled", False)
                     except Exception as e:
                         value = json.dumps(result, indent=2)[:500]
@@ -239,6 +240,8 @@ def evaluate(self, messages: List[Message], options: Optional[Options] = None) -
                         )
 
                     span.set_tag(AI_GUARD.ACTION_TAG, action)
+                    for tag in tags:
+                        span.set_tag(AI_GUARD.TAG + ".tag." + tag, "true")
                     if reason:
                         span.set_tag(AI_GUARD.REASON_TAG, reason)
                 else:

tests/appsec/ai_guard/api/test_api_client.py

@@ -46,9 +46,9 @@
 
 def _build_test_params():
     actions = [
-        {"action": "ALLOW", "reason": "Go ahead"},
-        {"action": "DENY", "reason": "Nope"},
-        {"action": "ABORT", "reason": "Kill it with fire"},
+        {"action": "ALLOW", "reason": "Go ahead", "tags": []},
+        {"action": "DENY", "reason": "Nope", "tags": ["deny_everything", "test_deny"]},
+        {"action": "ABORT", "reason": "Kill it with fire", "tags": ["alarm_tag", "abort_everything"]},
     ]
     block = [True, False]
     suites = [
@@ -63,6 +63,7 @@ def _build_test_params():
             pytest.param(
                 action["action"],
                 action["reason"],
+                action["tags"],
                 block,
                 suite["suite"],
                 suite["target"],
@@ -78,14 +79,24 @@ def assert_telemetry(mocked, metric, tags):
     assert ("count", "appsec", metric, 1, tags) in metrics
 
 
-@pytest.mark.parametrize("action,reason,blocking,suite,target,messages", _build_test_params())
+@pytest.mark.parametrize("action,reason,tags,blocking,suite,target,messages", _build_test_params())
 @patch("ddtrace.internal.telemetry.telemetry_writer._namespace")
 @patch("ddtrace.appsec.ai_guard._api_client.AIGuardClient._execute_request")
 def test_evaluate_method(
-    mock_execute_request, telemetry_mock, ai_guard_client, tracer, action, reason, blocking, suite, target, messages
+    mock_execute_request,
+    telemetry_mock,
+    ai_guard_client,
+    tracer,
+    action,
+    reason,
+    tags,
+    blocking,
+    suite,
+    target,
+    messages,
 ):
     """Test different combinations of evaluations."""
-    mock_execute_request.return_value = mock_evaluate_response(action, reason, blocking)
+    mock_execute_request.return_value = mock_evaluate_response(action, reason, tags, blocking)
     should_block = blocking and action != "ALLOW"
 
     if should_block:
@@ -103,6 +114,8 @@ def test_evaluate_method(
         expected_tags.update({"ai_guard.tool_name": "calc"})
     if action != "ALLOW" and blocking:
         expected_tags.update({"ai_guard.blocked": "true"})
+    for tag in tags:
+        expected_tags.update({"ai_guard.tag." + tag: "true"})
     assert_ai_guard_span(
         tracer,
         messages,

tests/appsec/ai_guard/utils.py

@@ -38,11 +38,18 @@ def assert_ai_guard_span(tracer: DummyTracer, messages: List[Message], tags: Dic
     assert struct["messages"] == messages
 
 
-def mock_evaluate_response(action: str, reason: str = "", block: bool = True) -> Mock:
+def mock_evaluate_response(action: str, reason: str = "", tags: List[str] = None, block: bool = True) -> Mock:
     mock_response = Mock()
     mock_response.status = 200
     mock_response.get_json.return_value = {
-        "data": {"attributes": {"action": action, "reason": reason, "is_blocking_enabled": block}}
+        "data": {
+            "attributes": {
+                "action": action,
+                "reason": reason,
+                "tags": tags if tags else [],
+                "is_blocking_enabled": block,
+            }
+        }
     }
     return mock_response