Skip to content

Commit beeb176

Browse files
Update to use meta struct for the rules
1 parent 9e53810 commit beeb176

File tree

3 files changed

+11
-6
lines changed

3 files changed

+11
-6
lines changed

ddtrace/appsec/ai_guard/_api_client.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -211,7 +211,8 @@ def evaluate(self, messages: List[Message], options: Optional[Options] = None) -
211211
span.set_tag(AI_GUARD.TOOL_NAME_TAG, tool_name)
212212
else:
213213
span.set_tag(AI_GUARD.TARGET_TAG, "prompt")
214-
span.set_struct_tag(AI_GUARD.STRUCT, {"messages": self._messages_for_meta_struct(messages)})
214+
215+
span._set_struct_tag(AI_GUARD.STRUCT, {"messages": self._messages_for_meta_struct(messages)})
215216

216217
try:
217218
response = self._execute_request(f"{self._endpoint}/evaluate", payload)
@@ -241,7 +242,8 @@ def evaluate(self, messages: List[Message], options: Optional[Options] = None) -
241242

242243
span.set_tag(AI_GUARD.ACTION_TAG, action)
243244
if len(tags) > 0:
244-
span.set_tag(AI_GUARD.TAG + ".matching_rules", json.dumps(tags))
245+
meta_struct = span._get_struct_tag(AI_GUARD.STRUCT)
246+
meta_struct.update({"matching_rules": tags})
245247
if reason:
246248
span.set_tag(AI_GUARD.REASON_TAG, reason)
247249
else:

tests/appsec/ai_guard/api/test_api_client.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -115,12 +115,13 @@ def test_evaluate_method(
115115
expected_tags.update({"ai_guard.tool_name": "calc"})
116116
if action != "ALLOW" and blocking:
117117
expected_tags.update({"ai_guard.blocked": "true"})
118+
expected_meta_struct = {"messages": messages}
118119
if len(tags) > 0:
119-
expected_tags.update({"ai_guard.matching_rules": json.dumps(tags)})
120+
expected_meta_struct.update({"matching_rules": tags})
120121
assert_ai_guard_span(
121122
tracer,
122-
messages,
123123
expected_tags,
124+
expected_meta_struct,
124125
)
125126
assert_telemetry(
126127
telemetry_mock,

tests/appsec/ai_guard/utils.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,13 +29,15 @@ def find_ai_guard_span(tracer: DummyTracer) -> Span:
2929
return span
3030

3131

32-
def assert_ai_guard_span(tracer: DummyTracer, messages: List[Message], tags: Dict[str, Any]) -> None:
32+
def assert_ai_guard_span(tracer: DummyTracer, tags: Dict[str, Any], meta_struct: Dict[str, Any], ) -> None:
3333
span = find_ai_guard_span(tracer)
3434
for tag, value in tags.items():
3535
assert tag in span.get_tags(), f"Missing {tag} from spans tags"
3636
assert span.get_tag(tag) == value, f"Wrong value {span.get_tag(tag)}, expected {value}"
3737
struct = span._get_struct_tag(AI_GUARD.TAG)
38-
assert struct["messages"] == messages
38+
for meta, value in meta_struct.items():
39+
assert meta in struct.keys(), f"Missing {meta} from meta_struct keys"
40+
assert struct[meta] == value, f"Wrong value {struct[meta]}, expected {value}"
3941

4042

4143
def mock_evaluate_response(action: str, reason: str = "", tags: List[str] = None, block: bool = True) -> Mock:

0 commit comments

Comments
 (0)