Skip to content

Commit

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

## [Security Solution] [Attack discovery] Updates Gemini connector to ignore unknown properties in responses

### Summary

This PR updates the Gemini connector's `RunApiResponseSchema` to ignore unknown properties in responses.

It fixes an issue where the recent addition of a `modelVersion` in Gemini responses causes Attack discovery to fail with the following error:

```
Attack discovery generation failed

ActionsClientLlm: action result status is error: an error occurred while running the action - Response validation failed (Error: [modelVersion]: definition for this key is missing)
```

shown in the screenshot below:

![error](https://github.com/user-attachments/assets/8b706290-27ca-42d9-b972-3d8613af690c)

### Desk testing

To reproduce:

1) Navigate to Security > Attack discovery

2) Select a Gemini connector

3) Click Generate

**Expected result**

- Attack discoveries are generated

**Actual result**

- The following error is displayed:

```
Attack discovery generation failed

ActionsClientLlm: action result status is error: an error occurred while running the action - Response validation failed (Error: [modelVersion]: definition for this key is missing)
```

![error](https://github.com/user-attachments/assets/8b706290-27ca-42d9-b972-3d8613af690c)

(cherry picked from commit 699db81)
  • Loading branch information
andrew-goldstein committed Sep 17, 2024
1 parent 657f5f0 commit 8d7a0da
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 @@ -131,6 +131,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 8d7a0da

Please sign in to comment.