Skip to content

Commit

Permalink
[8.x] [Security Solution] [Attack discovery] Updates Gemini connector…
Browse files Browse the repository at this point in the history
… to ignore unknown properties in responses (#192915) (#193222)

# Backport

This will backport the following commits from `main` to `8.x`:
- [[Security Solution] [Attack discovery] Updates Gemini connector to
ignore unknown properties in responses
(#192915)](#192915)

<!--- Backport version: 9.4.3 -->

### Questions ?
Please refer to the [Backport tool
documentation](https://github.com/sqren/backport)

<!--BACKPORT [{"author":{"name":"Andrew
Macri","email":"andrew.macri@elastic.co"},"sourceCommit":{"committedDate":"2024-09-17T19:46:48Z","message":"[Security
Solution] [Attack discovery] Updates Gemini connector to ignore unknown
properties in responses (#192915)\n\n## [Security Solution] [Attack
discovery] Updates Gemini connector to ignore unknown properties in
responses\r\n\r\n### Summary\r\n\r\nThis PR updates the Gemini
connector's `RunApiResponseSchema` to ignore unknown properties in
responses.\r\n\r\nIt fixes an issue where the recent addition of a
`modelVersion` in Gemini responses causes Attack discovery to fail with
the following error:\r\n\r\n```\r\nAttack discovery generation
failed\r\n\r\nActionsClientLlm: action result status is error: an error
occurred while running the action - Response validation failed (Error:
[modelVersion]: definition for this key is missing)\r\n```\r\n\r\nshown
in the screenshot
below:\r\n\r\n![error](https://github.com/user-attachments/assets/8b706290-27ca-42d9-b972-3d8613af690c)\r\n\r\n###
Desk testing\r\n\r\nTo reproduce:\r\n\r\n1) Navigate to Security >
Attack discovery\r\n\r\n2) Select a Gemini connector\r\n\r\n3) Click
Generate\r\n\r\n**Expected result**\r\n\r\n- Attack discoveries are
generated\r\n\r\n**Actual result**\r\n\r\n- The following error is
displayed:\r\n\r\n```\r\nAttack discovery generation
failed\r\n\r\nActionsClientLlm: action result status is error: an error
occurred while running the action - Response validation failed (Error:
[modelVersion]: definition for this key is
missing)\r\n```\r\n\r\n![error](https://github.com/user-attachments/assets/8b706290-27ca-42d9-b972-3d8613af690c)","sha":"699db81856eaf9a3666f4f46fa58ee4619b3b17c","branchLabelMapping":{"^v9.0.0$":"main","^v8.16.0$":"8.x","^v(\\d+).(\\d+).\\d+$":"$1.$2"}},"sourcePullRequest":{"labels":["bug","release_note:fix","v9.0.0","Team:
SecuritySolution","Team:Security Generative
AI","v8.16.0","v8.15.2"],"title":"[Security Solution] [Attack discovery]
Updates Gemini connector to ignore unknown properties in
responses","number":192915,"url":"#192915
Solution] [Attack discovery] Updates Gemini connector to ignore unknown
properties in responses (#192915)\n\n## [Security Solution] [Attack
discovery] Updates Gemini connector to ignore unknown properties in
responses\r\n\r\n### Summary\r\n\r\nThis PR updates the Gemini
connector's `RunApiResponseSchema` to ignore unknown properties in
responses.\r\n\r\nIt fixes an issue where the recent addition of a
`modelVersion` in Gemini responses causes Attack discovery to fail with
the following error:\r\n\r\n```\r\nAttack discovery generation
failed\r\n\r\nActionsClientLlm: action result status is error: an error
occurred while running the action - Response validation failed (Error:
[modelVersion]: definition for this key is missing)\r\n```\r\n\r\nshown
in the screenshot
below:\r\n\r\n![error](https://github.com/user-attachments/assets/8b706290-27ca-42d9-b972-3d8613af690c)\r\n\r\n###
Desk testing\r\n\r\nTo reproduce:\r\n\r\n1) Navigate to Security >
Attack discovery\r\n\r\n2) Select a Gemini connector\r\n\r\n3) Click
Generate\r\n\r\n**Expected result**\r\n\r\n- Attack discoveries are
generated\r\n\r\n**Actual result**\r\n\r\n- The following error is
displayed:\r\n\r\n```\r\nAttack discovery generation
failed\r\n\r\nActionsClientLlm: action result status is error: an error
occurred while running the action - Response validation failed (Error:
[modelVersion]: definition for this key is
missing)\r\n```\r\n\r\n![error](#192915
Solution] [Attack discovery] Updates Gemini connector to ignore unknown
properties in responses (#192915)\n\n## [Security Solution] [Attack
discovery] Updates Gemini connector to ignore unknown properties in
responses\r\n\r\n### Summary\r\n\r\nThis PR updates the Gemini
connector's `RunApiResponseSchema` to ignore unknown properties in
responses.\r\n\r\nIt fixes an issue where the recent addition of a
`modelVersion` in Gemini responses causes Attack discovery to fail with
the following error:\r\n\r\n```\r\nAttack discovery generation
failed\r\n\r\nActionsClientLlm: action result status is error: an error
occurred while running the action - Response validation failed (Error:
[modelVersion]: definition for this key is missing)\r\n```\r\n\r\nshown
in the screenshot
below:\r\n\r\n![error](https://github.com/user-attachments/assets/8b706290-27ca-42d9-b972-3d8613af690c)\r\n\r\n###
Desk testing\r\n\r\nTo reproduce:\r\n\r\n1) Navigate to Security >
Attack discovery\r\n\r\n2) Select a Gemini connector\r\n\r\n3) Click
Generate\r\n\r\n**Expected result**\r\n\r\n- Attack discoveries are
generated\r\n\r\n**Actual result**\r\n\r\n- The following error is
displayed:\r\n\r\n```\r\nAttack discovery generation
failed\r\n\r\nActionsClientLlm: action result status is error: an error
occurred while running the action - Response validation failed (Error:
[modelVersion]: definition for this key is
missing)\r\n```\r\n\r\n![error](https://github.com/user-attachments/assets/8b706290-27ca-42d9-b972-3d8613af690c)","sha":"699db81856eaf9a3666f4f46fa58ee4619b3b17c"}},{"branch":"8.x","label":"v8.16.0","branchLabelMappingKey":"^v8.16.0$","isSourceBranch":false,"state":"NOT_CREATED"},{"branch":"8.15","label":"v8.15.2","branchLabelMappingKey":"^v(\\d+).(\\d+).\\d+$","isSourceBranch":false,"state":"NOT_CREATED"}]}]
BACKPORT-->

Co-authored-by: Andrew Macri <andrew.macri@elastic.co>
  • Loading branch information
kibanamachine and andrew-goldstein committed Sep 17, 2024
1 parent dcd1cb6 commit 664d067
Show file tree
Hide file tree
Showing 2 changed files with 40 additions and 8 deletions.
19 changes: 11 additions & 8 deletions x-pack/plugins/stack_connectors/common/gemini/schema.ts
Original file line number Diff line number Diff line change
Expand Up @@ -29,14 +29,17 @@ export const RunActionParamsSchema = schema.object({
raw: schema.maybe(schema.boolean()),
});

export const RunApiResponseSchema = schema.object({
candidates: schema.any(),
usageMetadata: schema.object({
promptTokenCount: schema.number(),
candidatesTokenCount: schema.number(),
totalTokenCount: schema.number(),
}),
});
export const RunApiResponseSchema = schema.object(
{
candidates: schema.any(),
usageMetadata: schema.object({
promptTokenCount: schema.number(),
candidatesTokenCount: schema.number(),
totalTokenCount: schema.number(),
}),
},
{ unknowns: 'ignore' } // unknown keys will NOT fail validation, but will be removed
);

export const RunActionResponseSchema = schema.object(
{
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,35 @@ describe('GeminiConnector', () => {

expect(response).toEqual(connectorResponse);
});

describe('RunApiResponseSchema', () => {
it('successfully validates a response that only has known properties', () => {
const onlyKnownProperties = {
...defaultResponse.data,
};

expect(RunApiResponseSchema.validate(onlyKnownProperties)).toEqual(onlyKnownProperties);
});

it('fails validation when the response does NOT conform to the schema', () => {
const missingRequiredFields = {
// missing candidates and usageMetadata
};

expect(() => RunApiResponseSchema.validate(missingRequiredFields)).toThrowError();
});

it('removes unknown properties, but does NOT fail validation when they are present', () => {
const hasUnknownProperties = {
...defaultResponse.data,
modelVersion: '1.0.0', // <-- an unknown property
};

expect(RunApiResponseSchema.validate(hasUnknownProperties)).toEqual({
...defaultResponse.data,
});
});
});
});

describe('invokeAI', () => {
Expand Down

0 comments on commit 664d067

Please sign in to comment.