Skip to content
Merged
Show file tree
Hide file tree
Changes from 7 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions .github/workflows/system-tests.yml
Original file line number Diff line number Diff line change
Expand Up @@ -95,7 +95,7 @@ jobs:
persist-credentials: false
repository: 'DataDog/system-tests'
# Automatically managed, use scripts/update-system-tests-version to update
ref: '0526d7eb6ad321c14f7d4c7574cf8970089757c6'
ref: '24c2e96e7703a8128d5a29cae2e95776d86ef790'

- name: Download wheels to binaries directory
uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
Expand Down Expand Up @@ -148,7 +148,7 @@ jobs:
persist-credentials: false
repository: 'DataDog/system-tests'
# Automatically managed, use scripts/update-system-tests-version to update
ref: '0526d7eb6ad321c14f7d4c7574cf8970089757c6'
ref: '24c2e96e7703a8128d5a29cae2e95776d86ef790'

- name: Build runner
uses: ./.github/actions/install_runner
Expand Down Expand Up @@ -458,4 +458,4 @@ jobs:
needs.integration-frameworks-system-tests.result == 'cancelled'||
needs.tracer-release.result == 'failure' ||
needs.tracer-release.result == 'cancelled'
run: exit 1
run: exit 1
2 changes: 1 addition & 1 deletion .gitlab-ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ variables:
DD_VPA_TEMPLATE: "vpa-template-cpu-p70-10percent-2x-oom-min-cap"
# CI_DEBUG_SERVICES: "true"
# Automatically managed, use scripts/update-system-tests-version to update
SYSTEM_TESTS_REF: "0526d7eb6ad321c14f7d4c7574cf8970089757c6"
SYSTEM_TESTS_REF: "24c2e96e7703a8128d5a29cae2e95776d86ef790"

# Profiling native build image (built from dd/images/dd-trace-py/profiling_native)
PROFILING_NATIVE_IMAGE: "registry.ddbuild.io/dd-trace-py:v103334885-be1888c-profiling_native"
Expand Down
55 changes: 0 additions & 55 deletions .sg/rules/core-raising-dispatch.yml

This file was deleted.

108 changes: 0 additions & 108 deletions .sg/tests/__snapshots__/core-raising-dispatch-snapshot.yml

This file was deleted.

48 changes: 0 additions & 48 deletions .sg/tests/core-raising-dispatch-test.yml

This file was deleted.

18 changes: 9 additions & 9 deletions ddtrace/appsec/_ai_guard/_langchain.py
Original file line number Diff line number Diff line change
Expand Up @@ -224,8 +224,8 @@ def _langchain_generate_finally(*args, **kwargs):

Releases the AI Guard active counter that the matching ``.before``
listener bumped. Dispatched from the contrib's ``finally`` block so it
fires on every exit path — success, block (``_raising_dispatch`` raises
out of ``.before``), or exception inside the underlying LLM call. The
fires on every exit path — success, block (``core.dispatch(..., allow_raise=True)``
raises out of ``.before``), or exception inside the underlying LLM call. The
counter reset is a no-op when the counter is already zero, so listener
invocations that don't pair with a ``.before`` set are safe.
"""
Expand All @@ -247,21 +247,21 @@ def _langchain_llm_stream_before(client: AIGuardClient, instance, args, kwargs):


def _evaluate_langchain_messages(client: AIGuardClient, messages):
"""Evaluate the prompt and surface a block as a returned ``AIGuardAbortError``.
"""Evaluate the prompt and re-raise ``AIGuardAbortError`` on a block.

Returns the abort error so the contrib's dispatcher can re-raise it; the
``AIGuardClient`` already gates on ``ai_guard_config._ai_guard_block``,
so a returned error always represents a blocking decision.
Allow / skip paths return ``None``.
Re-raises so the contrib's ``core.dispatch(..., allow_raise=True)``
propagates the abort. The ``AIGuardClient`` already gates on
``ai_guard_config._ai_guard_block``, so a raised error always represents
a blocking decision. Allow / skip paths return ``None``.
"""
from langchain_core.messages import HumanMessage

# only call evaluator when the last message is an actual user prompt
if len(messages) > 0 and isinstance(messages[-1], HumanMessage):
try:
client.evaluate(_convert_messages(messages), Options(block=ai_guard_config._ai_guard_block))
except AIGuardAbortError as e:
return e
except AIGuardAbortError:
raise
except Exception:
logger.debug("Failed to evaluate chat model prompt", exc_info=True)
return None
19 changes: 16 additions & 3 deletions ddtrace/appsec/_ai_guard/_openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,19 @@ class OpenAIAIGuardAbortError(openai.UnprocessableEntityError, AIGuardAbortError
Catchable as either ``openai.APIError`` / ``openai.UnprocessableEntityError``
(idiomatic OpenAI error handling, no retry on 422) or
``AIGuardAbortError`` (Datadog-specific, exposes ``action`` / ``reason``).

AIDEV-NOTE: catchability asymmetry vs plain ``AIGuardAbortError``.
``AIGuardAbortError`` derives from ``DDBlockException(BaseException)``
so a generic ``except Exception:`` does NOT catch it (by design — a
block decision must not be silently swallowed). However,
``OpenAIAIGuardAbortError`` *also* inherits from
``openai.UnprocessableEntityError``, which is ``Exception``-derived,
so via MRO this subclass IS catchable by ``except Exception:``.
That asymmetry is intentional: the OpenAI contrib must keep
``except openai.APIError:`` blocks working unchanged for users
migrating from non-AI-Guard error handling. Code that wants
uniform block detection across providers should branch on
``isinstance(e, AIGuardAbortError)``.
"""

def __init__(self, action, reason, tags=None, sds=None, tag_probs=None):
Expand Down Expand Up @@ -285,7 +298,7 @@ def _openai_chat_completion_before(client, kwargs):
try:
client.evaluate(ai_guard_messages, Options(block=ai_guard_config._ai_guard_block))
except AIGuardAbortError as e:
return _wrap_abort_error(e)
raise _wrap_abort_error(e)
except Exception:
logger.debug("Failed to evaluate OpenAI chat completion request", exc_info=True)
return None
Expand All @@ -298,7 +311,7 @@ def _openai_chat_completion_after(client, kwargs, resp):
returns. Skips streaming responses (handled separately) and when a
framework evaluation is already active.

On block: returns an ``OpenAIAIGuardAbortError`` (or plain
On block: raises an ``OpenAIAIGuardAbortError`` (or plain
``AIGuardAbortError`` when the OpenAI SDK is not importable). Allow /
skip paths return ``None``.
"""
Expand All @@ -316,7 +329,7 @@ def _openai_chat_completion_after(client, kwargs, resp):
try:
client.evaluate(all_messages, Options(block=ai_guard_config._ai_guard_block))
except AIGuardAbortError as e:
return _wrap_abort_error(e)
raise _wrap_abort_error(e)
except Exception:
logger.debug("Failed to evaluate OpenAI chat completion response", exc_info=True)
return None
10 changes: 8 additions & 2 deletions ddtrace/appsec/ai_guard/_api_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
from ddtrace.ext import http
from ddtrace.internal import core
from ddtrace.internal import telemetry
from ddtrace.internal._exceptions import DDBlockException
import ddtrace.internal.logger as ddlogger
from ddtrace.internal.settings.asm import ai_guard_config
from ddtrace.internal.telemetry import TELEMETRY_NAMESPACE
Expand Down Expand Up @@ -94,8 +95,13 @@ def __init__(self, message: Optional[str], status: int = 0, errors: Optional[lis
super().__init__(message)


class AIGuardAbortError(Exception):
"""Exception to abort current execution due to security policy."""
class AIGuardAbortError(DDBlockException):
"""Exception to abort current execution due to security policy.

Inherits from ``DDBlockException`` (which is ``BaseException``-derived) so
that a generic ``except Exception:`` in user code does not accidentally
swallow an AI Guard block decision.
"""

def __init__(
self,
Expand Down
Loading
Loading