[Elastic Assistant] Update default assistant graph

package.json

@@ -971,7 +971,7 @@
     "@langchain/community": "0.2.18",
     "@langchain/core": "^0.2.18",
     "@langchain/google-genai": "^0.0.23",
-    "@langchain/langgraph": "^0.0.31",
+    "@langchain/langgraph": "0.0.34",
     "@langchain/openai": "^0.1.3",
     "@langtrase/trace-attributes": "^3.0.8",
     "@launchdarkly/node-server-sdk": "^9.5.1",

x-pack/plugins/elastic_assistant/README.md

@@ -8,6 +8,16 @@ This plugin does NOT contain UI components. See `x-pack/packages/kbn-elastic-ass
 
 Maintained by the Security Solution team
 
+## Graph structure
+
+![DefaultAssistantGraph](./docs/img/default_assistant_graph.png)
+
+## Development
+
+### Generate graph structure
+
+To generate the graph structure, run `yarn draw-graph` from the plugin directory.
+The graph will be generated in the `docs/img` directory of the plugin.
 
 ### Testing
 

x-pack/plugins/elastic_assistant/package.json

@@ -5,6 +5,7 @@
   "private": true,
   "license": "Elastic License 2.0",
   "scripts": {
-    "evaluate-model": "node ./scripts/model_evaluator"
+    "evaluate-model": "node ./scripts/model_evaluator",
+    "draw-graph": "node ./scripts/draw_graph"
   }
-}
+}

x-pack/plugins/elastic_assistant/scripts/draw_graph.js

@@ -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_graph_script').draw();

x-pack/plugins/elastic_assistant/scripts/draw_graph_script.ts

@@ -0,0 +1,66 @@
+/*
+ * 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 { ToolingLog } from '@kbn/tooling-log';
+import fs from 'fs/promises';
+import path from 'path';
+import {
+  ActionsClientChatOpenAI,
+  ActionsClientSimpleChatModel,
+} from '@kbn/langchain/server/language_models';
+import type { Logger } from '@kbn/logging';
+import { ChatPromptTemplate } from '@langchain/core/prompts';
+import { FakeLLM } from '@langchain/core/utils/testing';
+import { createOpenAIFunctionsAgent } from 'langchain/agents';
+import { getDefaultAssistantGraph } from '../server/lib/langchain/graphs/default_assistant_graph/graph';
+
+// Just defining some test variables to get the graph to compile..
+const testPrompt = ChatPromptTemplate.fromMessages([
+  ['system', 'You are a helpful assistant'],
+  ['placeholder', '{chat_history}'],
+  ['human', '{input}'],
+  ['placeholder', '{agent_scratchpad}'],
+]);
+
+const mockLlm = new FakeLLM({
+  response: JSON.stringify({}, null, 2),
+}) as unknown as ActionsClientChatOpenAI | ActionsClientSimpleChatModel;
+
+const createLlmInstance = () => {
+  return mockLlm;
+};
+
+async function getGraph(logger: Logger) {
+  const agentRunnable = await createOpenAIFunctionsAgent({
+    llm: mockLlm,
+    tools: [],
+    prompt: testPrompt,
+    streamRunnable: false,
+  });
+  const graph = getDefaultAssistantGraph({
+    agentRunnable,
+    logger,
+    createLlmInstance,
+    tools: [],
+    replacements: {},
+  });
+  return graph.getGraph();
+}
+
+export const draw = async () => {
+  const logger = new ToolingLog({
+    level: 'info',
+    writeTo: process.stdout,
+  }) as unknown as Logger;
+  logger.info('Compiling graph');
+  const outputPath = path.join(__dirname, '../docs/img/default_assistant_graph.png');
+  const graph = await getGraph(logger);
+  const output = await graph.drawMermaidPng();
+  const buffer = Buffer.from(await output.arrayBuffer());
+  logger.info(`Writing graph to ${outputPath}`);
+  await fs.writeFile(outputPath, buffer);
+};

x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/constants.ts

@@ -0,0 +1,17 @@
+/*
+ * 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.
+ */
+export const NodeType = {
+  PERSIST_CONVERSATION_CHANGES: 'persistConversationChanges',
+  GET_PERSISTED_CONVERSATION: 'getPersistedConversation',
+  GENERATE_CHAT_TITLE: 'generateChatTitle',
+  AGENT: 'agent',
+  TOOLS: 'tools',
+  RESPOND: 'respond',
+  MODEL_INPUT: 'modelInput',
+  STEP_ROUTER: 'stepRouter',
+  END: 'end',
+} as const;