Skip to content

Commit ba44d2e

Browse files
feat(ai_guard): implement AI Guard SDK (#6486)
feat(ai_guard): implement AI Guard SDK
1 parent 5374f69 commit ba44d2e

File tree

22 files changed

+1176
-2
lines changed

22 files changed

+1176
-2
lines changed

.github/workflows/aiguard.yml

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
name: AI Guard
2+
3+
on:
4+
pull_request:
5+
push:
6+
branches: [master]
7+
schedule:
8+
- cron: 0 4 * * *
9+
10+
concurrency:
11+
group: ${{ github.workflow }}-${{ github.ref || github.run_id }}
12+
cancel-in-progress: true
13+
14+
env:
15+
MOCHA_OPTIONS: ${{ github.ref == 'refs/heads/master' && '--retries 1' || '' }}
16+
17+
jobs:
18+
19+
macos:
20+
runs-on: macos-latest
21+
steps:
22+
- uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
23+
- uses: ./.github/actions/node/latest
24+
- uses: ./.github/actions/install
25+
- run: yarn test:aiguard:ci
26+
- uses: codecov/codecov-action@5a1091511ad55cbe89839c7260b706298ca349f7 # v5.5.1
27+
28+
ubuntu:
29+
runs-on: ubuntu-latest
30+
steps:
31+
- uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
32+
- uses: ./.github/actions/node/oldest-maintenance-lts
33+
- uses: ./.github/actions/install
34+
- run: yarn test:aiguard:ci
35+
- uses: ./.github/actions/node/newest-maintenance-lts
36+
- run: yarn test:aiguard:ci
37+
- uses: ./.github/actions/node/active-lts
38+
- run: yarn test:aiguard:ci
39+
- uses: ./.github/actions/node/latest
40+
- run: yarn test:aiguard:ci
41+
- uses: codecov/codecov-action@5a1091511ad55cbe89839c7260b706298ca349f7 # v5.5.1
42+
43+
windows:
44+
runs-on: windows-latest
45+
steps:
46+
- uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
47+
- uses: ./.github/actions/node/latest
48+
- uses: ./.github/actions/install
49+
with:
50+
cache: 'true'
51+
- run: yarn test:aiguard:ci
52+
- uses: codecov/codecov-action@5a1091511ad55cbe89839c7260b706298ca349f7 # v5.5.1
53+
54+
integration:
55+
strategy:
56+
matrix:
57+
version: [maintenance, active, latest]
58+
runs-on: ubuntu-latest
59+
steps:
60+
- uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
61+
- uses: ./.github/actions/install
62+
- uses: ./.github/actions/node
63+
with:
64+
version: ${{ matrix.version }}
65+
- run: yarn test:integration:aiguard

CODEOWNERS

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,10 @@
33
/packages/dd-trace/src/appsec/ @DataDog/asm-js
44
/packages/dd-trace/test/appsec/ @DataDog/asm-js
55

6+
/integration-tests/aiguard/ @DataDog/asm-js
7+
/packages/dd-trace/src/aiguard/ @DataDog/asm-js
8+
/packages/dd-trace/test/aiguard/ @DataDog/asm-js
9+
610
/integration-tests/debugger/ @DataDog/dd-trace-js @DataDog/debugger
711
/packages/datadog-code-origin/ @DataDog/dd-trace-js @DataDog/debugger
812
/packages/datadog-plugin-*/**/code_origin.* @DataDog/dd-trace-js @DataDog/debugger
@@ -93,6 +97,7 @@
9397
# CI
9498
/.github/workflows/apm-capabilities.yml @DataDog/apm-sdk-capabilities-js
9599
/.github/workflows/apm-integrations.yml @Datadog/apm-idm-js
100+
/.github/workflows/aiguard.yml @DataDog/asm-js
96101
/.github/workflows/appsec.yml @DataDog/asm-js
97102
/.github/workflows/codeql-analysis.yml @DataDog/dd-trace-js
98103
/.github/workflows/debugger.yml @DataDog/debugger

docs/test.ts

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -686,3 +686,46 @@ llmobs.annotate(span, {
686686

687687
// flush
688688
llmobs.flush()
689+
690+
691+
// AI Guard typings tests
692+
693+
tracer.init({
694+
experimental: {
695+
aiguard: {
696+
enabled: true,
697+
endpoint: 'http://localhost',
698+
maxMessagesLength: 22,
699+
maxContentSize: 1024,
700+
timeout: 1000
701+
}
702+
}
703+
})
704+
705+
const aiguard = tracer.aiguard
706+
707+
aiguard.evaluate([
708+
{ role: 'user', content: 'What is 2 + 2' },
709+
]).then(result => {
710+
result.action && result.reason
711+
})
712+
713+
aiguard.evaluate([
714+
{
715+
role: 'assistant',
716+
tool_calls: [
717+
{
718+
id: 'call_1',
719+
function: { name: 'calc', arguments: '{ "operator": "+", "args": [2, 2] }' }
720+
},
721+
],
722+
}
723+
]).then(result => {
724+
result.action && result.reason
725+
})
726+
727+
aiguard.evaluate([
728+
{ role: 'tool', tool_call_id: 'call_1', content: '5' },
729+
]).then(result => {
730+
result.action && result.reason
731+
})

index.d.ts

Lines changed: 193 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,11 @@ interface Tracer extends opentracing.Tracer {
139139
*/
140140
llmobs: tracer.llmobs.LLMObs;
141141

142+
/**
143+
* AI Guard SDK
144+
*/
145+
aiguard: tracer.aiguard.AIGuard;
146+
142147
/**
143148
* @experimental
144149
* Provide same functionality as OpenTelemetry Baggage:
@@ -593,6 +598,29 @@ declare namespace tracer {
593598
*/
594599
enabled?: boolean
595600
}
601+
},
602+
603+
aiguard?: {
604+
/**
605+
* Set to `true` to enable the SDK.
606+
*/
607+
enabled?: boolean,
608+
/**
609+
* URL of the AI Guard REST API.
610+
*/
611+
endpoint?: string,
612+
/**
613+
* Timeout used in calls to the AI Guard REST API in milliseconds (default 5000)
614+
*/
615+
timeout?: number,
616+
/**
617+
* Maximum number of conversational messages allowed to be set in the meta-struct
618+
*/
619+
maxMessagesLength?: number,
620+
/**
621+
* Max size of the content property set in the meta-struct
622+
*/
623+
maxContentSize?: number
596624
}
597625
};
598626

@@ -907,7 +935,7 @@ declare namespace tracer {
907935
scope?: string,
908936

909937
/**
910-
* Custom fields to attach to the user (RBAC, Oauth, etc).
938+
* Custom fields to attach to the user (RBAC, Oauth, etc...).
911939
*/
912940
[key: string]: string | undefined
913941
}
@@ -1059,6 +1087,170 @@ declare namespace tracer {
10591087
eventTrackingV2: EventTrackingV2
10601088
}
10611089

1090+
export namespace aiguard {
1091+
1092+
/**
1093+
* Represents a tool call made by an AI assistant in an agentic workflow.
1094+
*/
1095+
export interface ToolCall {
1096+
/**
1097+
* Unique identifier for this specific tool call instance used to correlate the call with its response.
1098+
*/
1099+
id: string;
1100+
/**
1101+
* Details about the function being invoked.
1102+
*/
1103+
function: {
1104+
/**
1105+
* The name of the tool/function to be called.
1106+
*/
1107+
name: string;
1108+
/**
1109+
* String containing the arguments to pass to the tool.
1110+
*/
1111+
arguments: string;
1112+
};
1113+
}
1114+
1115+
/**
1116+
* A standard conversational message exchanged with a Large Language Model (LLM).
1117+
*/
1118+
export interface TextMessage {
1119+
/**
1120+
* The role of the message sender in the conversation (e.g.: 'system', 'user', 'assistant').
1121+
*/
1122+
role: string;
1123+
/**
1124+
* The textual content of the message.
1125+
*/
1126+
content: string;
1127+
}
1128+
1129+
/**
1130+
* A message from an AI assistant containing only textual content.
1131+
*/
1132+
export interface AssistantTextMessage {
1133+
/**
1134+
* The role identifier, always set to 'assistant'
1135+
*/
1136+
role: "assistant";
1137+
/**
1138+
* The textual response content from the assistant.
1139+
*/
1140+
content: string;
1141+
/**
1142+
* Explicitly excluded when content is present to maintain type safety.
1143+
*/
1144+
tool_calls?: never;
1145+
}
1146+
1147+
/**
1148+
* A message from an AI assistant that initiates one or more tool calls.
1149+
*/
1150+
export interface AssistantToolCallMessage {
1151+
/**
1152+
* The role identifier, always set to 'assistant'
1153+
*/
1154+
role: "assistant";
1155+
/**
1156+
* Array of tool calls that the assistant wants to execute.
1157+
*/
1158+
tool_calls: ToolCall[];
1159+
/**
1160+
* Explicitly excluded when tool calls are present to maintain type safety.
1161+
*/
1162+
content?: never;
1163+
}
1164+
1165+
/**
1166+
* A message containing the result of a tool invocation.
1167+
*/
1168+
export interface ToolMessage {
1169+
/**
1170+
* The role identifier, always set to 'tool' for tool execution results.
1171+
*/
1172+
role: "tool";
1173+
/**
1174+
* The unique identifier linking this result to the original tool call.
1175+
* Must correspond to a ToolCall.id from a previous AssistantToolCallMessage.
1176+
*/
1177+
tool_call_id: string;
1178+
/**
1179+
* The output returned by the tool execution.
1180+
*/
1181+
content: string;
1182+
}
1183+
1184+
export type Message =
1185+
| TextMessage
1186+
| AssistantTextMessage
1187+
| AssistantToolCallMessage
1188+
| ToolMessage;
1189+
1190+
/**
1191+
* The result returned by AI Guard after evaluating a conversation.
1192+
*/
1193+
export interface Evaluation {
1194+
/**
1195+
* The security action determined by AI Guard:
1196+
* - 'ALLOW': The conversation is safe to proceed
1197+
* - 'DENY': The current conversation exchange should be blocked
1198+
* - 'ABORT': The full workflow should be terminated immediately
1199+
*/
1200+
action: 'ALLOW' | 'DENY' | 'ABORT';
1201+
/**
1202+
* Human-readable explanation for why this action was chosen.
1203+
*/
1204+
reason: string;
1205+
}
1206+
1207+
/**
1208+
* Error thrown when AI Guard evaluation determines that a conversation should be blocked
1209+
* and the client is configured to enforce blocking mode.
1210+
*/
1211+
export interface AIGuardAbortError extends Error {
1212+
/**
1213+
* Human-readable explanation from AI Guard describing why the conversation was blocked.
1214+
*/
1215+
reason: string;
1216+
}
1217+
1218+
/**
1219+
* Error thrown when the AI Guard SDK encounters communication failures or API errors while attempting to
1220+
* evaluate conversations.
1221+
*/
1222+
export interface AIGuardClientError extends Error {
1223+
/**
1224+
* Detailed error information returned by the AI Guard API, formatted according to the JSON:API error
1225+
* specification.
1226+
*/
1227+
errors?: unknown[];
1228+
/**
1229+
* The underlying error that caused the communication failure, such as network timeouts, connection refused,
1230+
* or JSON parsing errors.
1231+
*/
1232+
cause?: Error;
1233+
}
1234+
1235+
/**
1236+
* AI Guard security client for evaluating AI conversations.
1237+
*/
1238+
export interface AIGuard {
1239+
/**
1240+
* Evaluates a conversation thread.
1241+
*
1242+
* @param messages - Array of conversation messages
1243+
* @param opts - Optional configuration object:
1244+
* - `block`: When true, throws an exception if evaluation result is not 'ALLOW'
1245+
* and the AI Guard service has blocking mode enabled (default: false).
1246+
* @returns Promise resolving to an Evaluation with the security decision and reasoning.
1247+
* The promise rejects with AIGuardAbortError when `opts.block` is true and the evaluation result would block the request.
1248+
* The promise rejects with AIGuardClientError when communication with the AI Guard service fails.
1249+
*/
1250+
evaluate (messages: Message[], opts?: { block?: boolean }): Promise<Evaluation>;
1251+
}
1252+
}
1253+
10621254
/** @hidden */
10631255
type anyObject = {
10641256
[key: string]: any;

0 commit comments

Comments
 (0)