[Automatic Import] Fix yarn draw-graphs command

x-pack/plugins/integration_assistant/scripts/draw_graphs_script.ts

@@ -20,6 +20,8 @@ import { getEcsGraph, getEcsSubGraph } from '../server/graphs/ecs/graph';
 import { getLogFormatDetectionGraph } from '../server/graphs/log_type_detection/graph';
 import { getRelatedGraph } from '../server/graphs/related/graph';
 import { getKVGraph } from '../server/graphs/kv/graph';
+import { getUnstructuredGraph } from '../server/graphs/unstructured';
+import { getCelGraph } from '../server/graphs/cel/graph';
 
 // Some mock elements just to get the graph to compile
 const model = new FakeLLM({
@@ -45,17 +47,20 @@ async function drawGraph(compiledGraph: RunnableGraph, graphName: string) {
   await saveFile(`${graphName}.png`, buffer);
 }
 
+const GRAPH_LIST = {
+  related_graph: getRelatedGraph,
+  log_detection_graph: getLogFormatDetectionGraph,
+  categorization_graph: getCategorizationGraph,
+  kv_graph: getKVGraph,
+  ecs_graph: getEcsGraph,
+  ecs_subgraph: getEcsSubGraph,
+  unstructured_graph: getUnstructuredGraph,
+  cel_graph: getCelGraph,
+};
+
 export async function drawGraphs() {
-  const relatedGraph = (await getRelatedGraph({ client, model })).getGraph();
-  const logFormatDetectionGraph = (await getLogFormatDetectionGraph({ client, model })).getGraph();
-  const categorizationGraph = (await getCategorizationGraph({ client, model })).getGraph();
-  const ecsSubGraph = (await getEcsSubGraph({ model })).getGraph();
-  const ecsGraph = (await getEcsGraph({ model })).getGraph();
-  const kvGraph = (await getKVGraph({ client, model })).getGraph();
-  drawGraph(relatedGraph, 'related_graph');
-  drawGraph(logFormatDetectionGraph, 'log_detection_graph');
-  drawGraph(categorizationGraph, 'categorization_graph');
-  drawGraph(ecsSubGraph, 'ecs_subgraph');
-  drawGraph(ecsGraph, 'ecs_graph');
-  drawGraph(kvGraph, 'kv_graph');
+  for (const [name, graph] of Object.entries(GRAPH_LIST)) {
+    const compiledGraph = (await graph({ client, model })).getGraph();
+    drawGraph(compiledGraph, name);
+  }
 }

x-pack/plugins/integration_assistant/server/graphs/categorization/graph.ts

@@ -233,6 +233,6 @@ export async function getCategorizationGraph({ client, model }: CategorizationGr
       }
     );
 
-  const compiledCategorizationGraph = workflow.compile().withConfig({ runName: 'Categorization' });
+  const compiledCategorizationGraph = workflow.compile();
   return compiledCategorizationGraph;
 }

x-pack/plugins/integration_assistant/server/graphs/cel/graph.ts

@@ -104,6 +104,6 @@ export async function getCelGraph({ model }: CelInputGraphParams) {
     .addEdge('handleGetStateVariables', 'handleGetStateDetails')
     .addEdge('handleGetStateDetails', 'modelOutput');
 
-  const compiledCelGraph = workflow.compile().withConfig({ runName: 'CEL' });
+  const compiledCelGraph = workflow.compile();
   return compiledCelGraph;
 }

x-pack/plugins/integration_assistant/server/graphs/ecs/graph.ts

@@ -78,7 +78,7 @@ export async function getEcsSubGraph({ model }: EcsGraphParams) {
     })
     .addEdge('modelSubOutput', END);
 
-  const compiledEcsSubGraph = workflow.compile().withConfig({ runName: 'ECS Mapping (Chunk)' });
+  const compiledEcsSubGraph = workflow.compile();
   return compiledEcsSubGraph;
 }
 
@@ -96,7 +96,7 @@ export async function getEcsGraph({ model }: EcsGraphParams) {
     .addNode('handleMergedSubGraphResponse', (state: EcsMappingState) =>
       modelMergedInputFromSubGraph({ state })
     )
-    .addNode('subGraph', subGraph)
+    .addNode('subGraph', subGraph.withConfig({ runName: 'ECS Mapping (Chunk)' }))
     .addEdge(START, 'modelInput')
     .addEdge('subGraph', 'handleMergedSubGraphResponse')
     .addEdge('handleDuplicates', 'handleValidation')
@@ -119,6 +119,6 @@ export async function getEcsGraph({ model }: EcsGraphParams) {
     })
     .addEdge('modelOutput', END);
 
-  const compiledEcsGraph = workflow.compile().withConfig({ runName: 'ECS Mapping' });
+  const compiledEcsGraph = workflow.compile();
   return compiledEcsGraph;
 }

x-pack/plugins/integration_assistant/server/graphs/kv/graph.ts

@@ -139,6 +139,6 @@ export async function getKVGraph({ model, client }: KVGraphParams) {
     })
     .addEdge('modelOutput', END);
 
-  const compiledKVGraph = workflow.compile().withConfig({ runName: 'Key-Value' });
+  const compiledKVGraph = workflow.compile();
   return compiledKVGraph;
 }

x-pack/plugins/integration_assistant/server/graphs/log_type_detection/graph.ts

@@ -118,8 +118,14 @@ export async function getLogFormatDetectionGraph({ model, client }: LogDetection
     .addNode('handleLogFormatDetection', (state: LogFormatDetectionState) =>
       handleLogFormatDetection({ state, model, client })
     )
-    .addNode('handleKVGraph', await getKVGraph({ model, client }))
-    .addNode('handleUnstructuredGraph', await getUnstructuredGraph({ model, client }))
+    .addNode(
+      'handleKVGraph',
+      (await getKVGraph({ model, client })).withConfig({ runName: 'Key-Value' })
+    )
+    .addNode(
+      'handleUnstructuredGraph',
+      (await getUnstructuredGraph({ model, client })).withConfig({ runName: 'Unstructured' })
+    )
     .addNode('handleCSV', (state: LogFormatDetectionState) => handleCSV({ state, model, client }))
     .addEdge(START, 'modelInput')
     .addEdge('modelInput', 'handleLogFormatDetection')
@@ -138,6 +144,6 @@ export async function getLogFormatDetectionGraph({ model, client }: LogDetection
       }
     );
 
-  const compiledLogFormatDetectionGraph = workflow.compile().withConfig({ runName: 'Log Format' });
+  const compiledLogFormatDetectionGraph = workflow.compile();
   return compiledLogFormatDetectionGraph;
 }

x-pack/plugins/integration_assistant/server/graphs/related/graph.ts

@@ -179,6 +179,6 @@ export async function getRelatedGraph({ client, model }: RelatedGraphParams) {
       }
     );
 
-  const compiledRelatedGraph = workflow.compile().withConfig({ runName: 'Related' });
+  const compiledRelatedGraph = workflow.compile();
   return compiledRelatedGraph;
 }

x-pack/plugins/integration_assistant/server/graphs/unstructured/graph.ts

@@ -107,6 +107,6 @@ export async function getUnstructuredGraph({ model, client }: UnstructuredGraphP
     .addEdge('handleUnstructuredError', 'handleUnstructuredValidate')
     .addEdge('modelOutput', END);
 
-  const compiledUnstructuredGraph = workflow.compile().withConfig({ runName: 'Unstructured' });
+  const compiledUnstructuredGraph = workflow.compile();
   return compiledUnstructuredGraph;
 }

x-pack/plugins/integration_assistant/server/routes/analyze_logs_routes.ts

@@ -91,7 +91,9 @@ export function registerAnalyzeLogsRoutes(
             logSamples,
           };
           const graph = await getLogFormatDetectionGraph({ model, client });
-          const graphResults = await graph.invoke(logFormatParameters, options);
+          const graphResults = await graph
+            .withConfig({ runName: 'Log Format' })
+            .invoke(logFormatParameters, options);
           const graphLogFormat = graphResults.results.samplesFormat.name;
           if (graphLogFormat === 'unsupported') {
             throw new UnsupportedLogFormatError(GenerationErrorCode.UNSUPPORTED_LOG_SAMPLES_FORMAT);

x-pack/plugins/integration_assistant/server/routes/categorization_routes.test.ts

@@ -23,7 +23,9 @@ const mockResult = jest.fn().mockResolvedValue({
 jest.mock('../graphs/categorization', () => {
   return {
     getCategorizationGraph: jest.fn().mockResolvedValue({
-      invoke: () => mockResult(),
+      withConfig: () => ({
+        invoke: () => mockResult(),
+      }),
     }),
   };
 });

x-pack/plugins/integration_assistant/server/routes/categorization_routes.ts

@@ -99,7 +99,9 @@ export function registerCategorizationRoutes(
             };
 
             const graph = await getCategorizationGraph({ client, model });
-            const results = await graph.invoke(parameters, options);
+            const results = await graph
+              .withConfig({ runName: 'Categorization' })
+              .invoke(parameters, options);
 
             return res.ok({ body: CategorizationResponse.parse(results) });
           } catch (err) {

x-pack/plugins/integration_assistant/server/routes/cel_route.test.ts

@@ -22,7 +22,9 @@ const mockResult = jest.fn().mockResolvedValue({
 jest.mock('../graphs/cel', () => {
   return {
     getCelGraph: jest.fn().mockResolvedValue({
-      invoke: () => mockResult(),
+      withConfig: () => ({
+        invoke: () => mockResult(),
+      }),
     }),
   };
 });

x-pack/plugins/integration_assistant/server/routes/cel_routes.ts

@@ -78,7 +78,7 @@ export function registerCelInputRoutes(router: IRouter<IntegrationAssistantRoute
           };
 
           const graph = await getCelGraph({ model });
-          const results = await graph.invoke(parameters, options);
+          const results = await graph.withConfig({ runName: 'CEL' }).invoke(parameters, options);
 
           return res.ok({ body: CelInputResponse.parse(results) });
         } catch (e) {

x-pack/plugins/integration_assistant/server/routes/ecs_routes.test.ts

@@ -29,7 +29,9 @@ const mockResult = jest.fn().mockResolvedValue({
 jest.mock('../graphs/ecs', () => {
   return {
     getEcsGraph: jest.fn().mockResolvedValue({
-      invoke: () => mockResult(),
+      withConfig: () => ({
+        invoke: () => mockResult(),
+      }),
     }),
   };
 });

x-pack/plugins/integration_assistant/server/routes/ecs_routes.ts

@@ -92,7 +92,9 @@ export function registerEcsRoutes(router: IRouter<IntegrationAssistantRouteHandl
           };
 
           const graph = await getEcsGraph({ model });
-          const results = await graph.invoke(parameters, options);
+          const results = await graph
+            .withConfig({ runName: 'ECS Mapping' })
+            .invoke(parameters, options);
 
           return res.ok({ body: EcsMappingResponse.parse(results) });
         } catch (err) {

x-pack/plugins/integration_assistant/server/routes/related_routes.test.ts

@@ -23,7 +23,9 @@ const mockResult = jest.fn().mockResolvedValue({
 jest.mock('../graphs/related', () => {
   return {
     getRelatedGraph: jest.fn().mockResolvedValue({
-      invoke: () => mockResult(),
+      withConfig: () => ({
+        invoke: () => mockResult(),
+      }),
     }),
   };
 });

x-pack/plugins/integration_assistant/server/routes/related_routes.ts

@@ -90,7 +90,9 @@ export function registerRelatedRoutes(router: IRouter<IntegrationAssistantRouteH
           };
 
           const graph = await getRelatedGraph({ client, model });
-          const results = await graph.invoke(parameters, options);
+          const results = await graph
+            .withConfig({ runName: 'Related' })
+            .invoke(parameters, options);
           return res.ok({ body: RelatedResponse.parse(results) });
         } catch (err) {
           try {