Skip to content

Commit

Permalink
[Rules Migration] Add yarn command to draw graphs (elastic#201753)
Browse files Browse the repository at this point in the history
## Summary

This PR implements a simple yarn command that is identical to the ones
used in Elastic Assistant and Automatic Import, which draws the relevant
Langgraph components used and places them in the docs folder under the
correct subfolder.

It also adds a tiny unit test to only ensure that the graph compiles
itself, which can be easy to miss during smaller fixes to the graph.

This is meant to be run manually and is not part of any CI process, but
provides benefits to navigate graphs when they become more complex.
  • Loading branch information
P1llus authored Nov 29, 2024
1 parent 2b7f584 commit 06bc18b
Show file tree
Hide file tree
Showing 7 changed files with 159 additions and 5 deletions.
17 changes: 17 additions & 0 deletions x-pack/plugins/security_solution/docs/siem_migration/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
# SIEM Migration Library

## Migration Process

The SIEM migration library defines a set of UI components and services that are used to migrate third party SIEM resources like detection rules and translate them into resources that can be used in the Elastic Security app.

## Graphs:

The below images are generated by running the following command from the security_solution directory:

```bash
yarn siem-migrations:graph:draw
```

Main agent graph:

![Agent Graph](./img/agent_graph.png)
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
3 changes: 2 additions & 1 deletion x-pack/plugins/security_solution/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
"test:generate:serverless-dev": "NODE_TLS_REJECT_UNAUTHORIZED=0 node --no-warnings scripts/endpoint/resolver_generator --node https://elastic_serverless:[email protected]:9200 --kibana http://elastic_serverless:[email protected]:5601",
"mappings:generate": "node scripts/mappings/mappings_generator",
"mappings:load": "node scripts/mappings/mappings_loader",
"siem-migrations:graph:draw": "node scripts/siem_migration/draw_graphs",
"junit:transform": "node scripts/junit_transformer --pathPattern '../../../target/kibana-security-solution/cypress/results/*.xml' --rootDirectory ../../../ --reportName 'Security Solution Cypress' --writeInPlace",
"openapi:generate": "node scripts/openapi/generate",
"openapi:generate:debug": "node --inspect-brk scripts/openapi/generate",
Expand All @@ -35,4 +36,4 @@
"openapi:bundle:entity-analytics": "node scripts/openapi/bundle_entity_analytics",
"openapi:bundle:endpoint-management": "node scripts/openapi/bundle_endpoint_management"
}
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
/*
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
* or more contributor license agreements. Licensed under the Elastic License
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/

require('../../../../../src/setup_node_env');
require('./draw_graphs_script').draw();
Original file line number Diff line number Diff line change
@@ -0,0 +1,80 @@
/*
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
* or more contributor license agreements. Licensed under the Elastic License
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/

import type { InferenceClient } from '@kbn/inference-plugin/server';
import type {
ActionsClientChatOpenAI,
ActionsClientSimpleChatModel,
} from '@kbn/langchain/server/language_models';
import type { Logger } from '@kbn/logging';
import { ToolingLog } from '@kbn/tooling-log';
import { FakeLLM } from '@langchain/core/utils/testing';
import fs from 'fs/promises';
import path from 'path';
import { getRuleMigrationAgent } from '../../server/lib/siem_migrations/rules/task/agent';
import type { IntegrationRetriever } from '../../server/lib/siem_migrations/rules/task/util/integration_retriever';
import type { PrebuiltRulesMapByName } from '../../server/lib/siem_migrations/rules/task/util/prebuilt_rules';
import type { RuleResourceRetriever } from '../../server/lib/siem_migrations/rules/task/util/rule_resource_retriever';

interface Drawable {
drawMermaidPng: () => Promise<Blob>;
}

const mockLlm = new FakeLLM({
response: JSON.stringify({}, null, 2),
}) as unknown as ActionsClientChatOpenAI | ActionsClientSimpleChatModel;

const inferenceClient = {} as InferenceClient;
const connectorId = 'draw_graphs';
const prebuiltRulesMap = {} as PrebuiltRulesMapByName;
const resourceRetriever = {} as RuleResourceRetriever;
const integrationRetriever = {} as IntegrationRetriever;

const createLlmInstance = () => {
return mockLlm;
};

async function getAgentGraph(logger: Logger): Promise<Drawable> {
const model = createLlmInstance();
const graph = getRuleMigrationAgent({
model,
inferenceClient,
prebuiltRulesMap,
resourceRetriever,
integrationRetriever,
connectorId,
logger,
});
return graph.getGraphAsync({ xray: true });
}

export const drawGraph = async ({
getGraphAsync,
outputFilename,
}: {
getGraphAsync: (logger: Logger) => Promise<Drawable>;
outputFilename: string;
}) => {
const logger = new ToolingLog({
level: 'info',
writeTo: process.stdout,
}) as unknown as Logger;
logger.info('Compiling graph');
const outputPath = path.join(__dirname, outputFilename);
const graph = await getGraphAsync(logger);
const output = await graph.drawMermaidPng();
const buffer = Buffer.from(await output.arrayBuffer());
logger.info(`Writing graph to ${outputPath}`);
await fs.writeFile(outputPath, buffer);
};

export const draw = async () => {
await drawGraph({
getGraphAsync: getAgentGraph,
outputFilename: '../../docs/siem_migration/img/agent_graph.png',
});
};
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
/*
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
* or more contributor license agreements. Licensed under the Elastic License
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/

import type { InferenceClient } from '@kbn/inference-plugin/server';
import type {
ActionsClientChatOpenAI,
ActionsClientSimpleChatModel,
} from '@kbn/langchain/server/language_models';
import { loggerMock } from '@kbn/logging-mocks';
import { FakeLLM } from '@langchain/core/utils/testing';
import type { IntegrationRetriever } from '../util/integration_retriever';
import type { PrebuiltRulesMapByName } from '../util/prebuilt_rules';
import type { RuleResourceRetriever } from '../util/rule_resource_retriever';
import { getRuleMigrationAgent } from './graph';

describe('getRuleMigrationAgent', () => {
const model = new FakeLLM({
response: JSON.stringify({}, null, 2),
}) as unknown as ActionsClientChatOpenAI | ActionsClientSimpleChatModel;

const inferenceClient = {} as InferenceClient;
const connectorId = 'draw_graphs';
const prebuiltRulesMap = {} as PrebuiltRulesMapByName;
const resourceRetriever = {} as RuleResourceRetriever;
const integrationRetriever = {} as IntegrationRetriever;
const logger = loggerMock.create();

it('Ensures that the graph compiles', async () => {
try {
await getRuleMigrationAgent({
model,
inferenceClient,
prebuiltRulesMap,
resourceRetriever,
integrationRetriever,
connectorId,
logger,
});
} catch (error) {
throw Error(`getRuleMigrationAgent threw an error: ${error}`);
}
});
});
Original file line number Diff line number Diff line change
Expand Up @@ -30,16 +30,16 @@ export function getRuleMigrationAgent({
logger,
});

const translateRuleGraph = new StateGraph(migrateRuleState)
const siemMigrationAgentGraph = new StateGraph(migrateRuleState)
// Nodes
.addNode('matchPrebuiltRule', matchPrebuiltRuleNode)
.addNode('translation', translationSubGraph)
.addNode('translationSubGraph', translationSubGraph)
// Edges
.addEdge(START, 'matchPrebuiltRule')
.addConditionalEdges('matchPrebuiltRule', matchedPrebuiltRuleConditional)
.addEdge('translation', END);
.addEdge('translationSubGraph', END);

const graph = translateRuleGraph.compile();
const graph = siemMigrationAgentGraph.compile();
graph.name = 'Rule Migration Graph'; // Customizes the name displayed in LangSmith
return graph;
}
Expand Down

0 comments on commit 06bc18b

Please sign in to comment.