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

ddtrace/appsec/ai_guard/_api_client.py

@@ -211,7 +211,8 @@ def evaluate(self, messages: List[Message], options: Optional[Options] = None) -
                     span.set_tag(AI_GUARD.TOOL_NAME_TAG, tool_name)
                 else:
                     span.set_tag(AI_GUARD.TARGET_TAG, "prompt")
-                span.set_struct_tag(AI_GUARD.STRUCT, {"messages": self._messages_for_meta_struct(messages)})
+
+                span._set_struct_tag(AI_GUARD.STRUCT, {"messages": self._messages_for_meta_struct(messages)})
 
                 try:
                     response = self._execute_request(f"{self._endpoint}/evaluate", payload)
@@ -224,6 +225,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 +241,9 @@ def evaluate(self, messages: List[Message], options: Optional[Options] = None) -
                         )
 
                     span.set_tag(AI_GUARD.ACTION_TAG, action)
+                    if len(tags) > 0:
+                        meta_struct = span._get_struct_tag(AI_GUARD.STRUCT)
+                        meta_struct.update({"attack_categories": tags})
                     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,10 +114,13 @@ def test_evaluate_method(
         expected_tags.update({"ai_guard.tool_name": "calc"})
     if action != "ALLOW" and blocking:
         expected_tags.update({"ai_guard.blocked": "true"})
+    expected_meta_struct = {"messages": messages}
+    if len(tags) > 0:
+        expected_meta_struct.update({"attack_categories": tags})
     assert_ai_guard_span(
         tracer,
-        messages,
         expected_tags,
+        expected_meta_struct,
     )
     assert_telemetry(
         telemetry_mock,

tests/appsec/ai_guard/utils.py

@@ -29,20 +29,33 @@ def find_ai_guard_span(tracer: DummyTracer) -> Span:
     return span
 
 
-def assert_ai_guard_span(tracer: DummyTracer, messages: List[Message], tags: Dict[str, Any]) -> None:
+def assert_ai_guard_span(
+    tracer: DummyTracer,
+    tags: Dict[str, Any],
+    meta_struct: Dict[str, Any],
+) -> None:
     span = find_ai_guard_span(tracer)
     for tag, value in tags.items():
         assert tag in span.get_tags(), f"Missing {tag} from spans tags"
         assert span.get_tag(tag) == value, f"Wrong value {span.get_tag(tag)}, expected {value}"
     struct = span._get_struct_tag(AI_GUARD.TAG)
-    assert struct["messages"] == messages
+    for meta, value in meta_struct.items():
+        assert meta in struct.keys(), f"Missing {meta} from meta_struct keys"
+        assert struct[meta] == value, f"Wrong value {struct[meta]}, expected {value}"
 
 
-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