[Security GenAI] Use AI setting to set langsmith tracing to the Integ…

…ration Assistant (elastic#187466) ## Summary Enables tracing Langchain invocations in the integrations assistant using the Langsmith settings stored by the Security AI Settings. The evaluation settings tab is still under an experimental flag, to see it: ``` xpack.securitySolution.enableExperimental: ['assistantModelEvaluation'] ``` ### Screenshots <img width="1317" alt="Settings" src="https://github.com/elastic/kibana/assets/17747913/6aed1ef6-3750-4259-9fe2-b8bf1aed5504"> After one execution of the integration assistant: <img width="1240" alt="langsmith" src="https://github.com/elastic/kibana/assets/17747913/dd3dd99c-7c83-4a35-95b2-789e7a341031"> --------- Co-authored-by: kibanamachine <[email protected]> Co-authored-by: Elastic Machine <[email protected]>
kapral18 · Jul 8, 2024 · 92099b2 · 92099b2
1 parent 35ee0cc
commit 92099b2
Show file tree

Hide file tree

Showing 32 changed files with 231 additions and 91 deletions.
diff --git a/...t/server/lib/langchain/tracers/README.mdx → ...s/kbn-langchain/server/tracers/README.mdx b/...t/server/lib/langchain/tracers/README.mdx → ...s/kbn-langchain/server/tracers/README.mdx
diff --git a/...erver/lib/langchain/tracers/apm_tracer.ts → ...angchain/server/tracers/apm/apm_tracer.ts b/...erver/lib/langchain/tracers/apm_tracer.ts → ...angchain/server/tracers/apm/apm_tracer.ts
diff --git a/x-pack/packages/kbn-langchain/server/tracers/apm/index.ts b/x-pack/packages/kbn-langchain/server/tracers/apm/index.ts
@@ -0,0 +1,8 @@
+/*
+ * 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 { APMTracer } from './apm_tracer';
diff --git a/x-pack/packages/kbn-langchain/server/tracers/langsmith/index.ts b/x-pack/packages/kbn-langchain/server/tracers/langsmith/index.ts
@@ -0,0 +1,7 @@
+/*
+ * 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 { getLangSmithTracer, isLangSmithEnabled } from './langsmith_tracer';
diff --git a/x-pack/packages/kbn-langchain/server/tracers/langsmith/langsmith_tracer.ts b/x-pack/packages/kbn-langchain/server/tracers/langsmith/langsmith_tracer.ts
@@ -0,0 +1,64 @@
+/*
+ * 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 { Client } from 'langsmith';
+import type { Logger } from '@kbn/core/server';
+import { ToolingLog } from '@kbn/tooling-log';
+import { LangChainTracer } from '@langchain/core/tracers/tracer_langchain';
+
+/**
+ * Returns a custom LangChainTracer which adds the `exampleId` so Dataset 'Test' runs are written to LangSmith
+ * If `exampleId` is present (and a corresponding example exists in LangSmith) trace is written to the Dataset's `Tests`
+ * section, otherwise it is written to the `Project` provided
+ *
+ * @param apiKey API Key for LangSmith (will fetch from env vars if not provided)
+ * @param projectName Name of project to trace results to
+ * @param exampleId Dataset exampleId to associate trace with
+ * @param logger
+ */
+export const getLangSmithTracer = ({
+  apiKey,
+  projectName,
+  exampleId,
+  logger,
+}: {
+  apiKey?: string;
+  projectName?: string;
+  exampleId?: string;
+  logger: Logger | ToolingLog;
+}): LangChainTracer[] => {
+  try {
+    if (!isLangSmithEnabled() || apiKey == null) {
+      return [];
+    }
+    const lcTracer = new LangChainTracer({
+      projectName, // Shows as the 'test' run's 'name' in langsmith ui
+      exampleId,
+      client: new Client({ apiKey }),
+    });
+
+    return [lcTracer];
+  } catch (e) {
+    // Note: creating a tracer can fail if the LangSmith env vars are not set correctly
+    logger.error(`Error creating LangSmith tracer: ${e.message}`);
+  }
+
+  return [];
+};
+
+/**
+ * Returns true if LangSmith/tracing is enabled
+ */
+export const isLangSmithEnabled = (): boolean => {
+  try {
+    // Just checking if apiKey is available, if better way to check for enabled that is not env var please update
+    const config = Client.getDefaultClientConfig();
+    return config.apiKey != null;
+  } catch (e) {
+    return false;
+  }
+};
diff --git a/x-pack/packages/kbn-langchain/tsconfig.json b/x-pack/packages/kbn-langchain/tsconfig.json
@@ -18,6 +18,7 @@
     "@kbn/logging",
     "@kbn/actions-plugin",
     "@kbn/logging-mocks",
-    "@kbn/utility-types"
+    "@kbn/utility-types",
+    "@kbn/tooling-log"
   ]
 }
diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts
@@ -18,12 +18,12 @@ import {
   ActionsClientSimpleChatModel,
 } from '@kbn/langchain/server';
 import { MessagesPlaceholder } from '@langchain/core/prompts';
+import { APMTracer } from '@kbn/langchain/server/tracers/apm';
+import { withAssistantSpan } from '../tracers/apm/with_assistant_span';
 import { EsAnonymizationFieldsSchema } from '../../../ai_assistant_data_clients/anonymization_fields/types';
 import { transformESSearchToAnonymizationFields } from '../../../ai_assistant_data_clients/anonymization_fields/helpers';
 import { AgentExecutor } from '../executors/types';
-import { APMTracer } from '../tracers/apm_tracer';
 import { AssistantToolParams } from '../../../types';
-import { withAssistantSpan } from '../tracers/with_assistant_span';
 export const DEFAULT_AGENT_EXECUTOR_ID = 'Elastic AI Assistant Agent Executor';
 
 /**

diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/openai_functions_executor.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/executors/openai_functions_executor.ts
@@ -11,9 +11,9 @@ import { BufferMemory, ChatMessageHistory } from 'langchain/memory';
 import { ChainTool } from 'langchain/tools/chain';
 
 import { ActionsClientLlm } from '@kbn/langchain/server';
+import { APMTracer } from '@kbn/langchain/server/tracers/apm';
+import { withAssistantSpan } from '../tracers/apm/with_assistant_span';
 import { AgentExecutor } from './types';
-import { withAssistantSpan } from '../tracers/with_assistant_span';
-import { APMTracer } from '../tracers/apm_tracer';
 
 export const OPEN_AI_FUNCTIONS_AGENT_EXECUTOR_ID =
   'Elastic AI Assistant Agent Executor (OpenAI Functions)';

diff --git a/.../plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts b/.../plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts
@@ -11,11 +11,11 @@ import { streamFactory, StreamResponseWithHeaders } from '@kbn/ml-response-strea
 import { transformError } from '@kbn/securitysolution-es-utils';
 import type { KibanaRequest } from '@kbn/core-http-server';
 import type { ExecuteConnectorRequestBody, TraceData } from '@kbn/elastic-assistant-common';
+import { APMTracer } from '@kbn/langchain/server/tracers/apm';
+import { withAssistantSpan } from '../../tracers/apm/with_assistant_span';
 import { AGENT_NODE_TAG } from './nodes/run_agent';
 import { DEFAULT_ASSISTANT_GRAPH_ID, DefaultAssistantGraph } from './graph';
 import type { OnLlmResponse, TraceOptions } from '../../executors/types';
-import type { APMTracer } from '../../tracers/apm_tracer';
-import { withAssistantSpan } from '../../tracers/with_assistant_span';
 
 interface StreamGraphParams {
   apmTracer: APMTracer;

diff --git a/...ck/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts b/...ck/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts
@@ -13,11 +13,11 @@ import {
   ActionsClientSimpleChatModel,
 } from '@kbn/langchain/server';
 import { createOpenAIFunctionsAgent, createStructuredChatAgent } from 'langchain/agents';
+import { APMTracer } from '@kbn/langchain/server/tracers/apm';
 import { EsAnonymizationFieldsSchema } from '../../../../ai_assistant_data_clients/anonymization_fields/types';
 import { AssistantToolParams } from '../../../../types';
 import { AgentExecutor } from '../../executors/types';
 import { openAIFunctionAgentPrompt, structuredChatAgentPrompt } from './prompts';
-import { APMTracer } from '../../tracers/apm_tracer';
 import { getDefaultAssistantGraph } from './graph';
 import { invokeGraph, streamGraph } from './helpers';
 import { transformESSearchToAnonymizationFields } from '../../../../ai_assistant_data_clients/anonymization_fields/helpers';

diff --git a/.../langchain/tracers/with_assistant_span.ts → ...gchain/tracers/apm/with_assistant_span.ts b/.../langchain/tracers/with_assistant_span.ts → ...gchain/tracers/apm/with_assistant_span.ts
diff --git a/x-pack/plugins/elastic_assistant/server/lib/model_evaluator/evaluation.ts b/x-pack/plugins/elastic_assistant/server/lib/model_evaluator/evaluation.ts
@@ -15,9 +15,10 @@ import { ToolingLog } from '@kbn/tooling-log';
 import { LangChainTracer } from '@langchain/core/tracers/tracer_langchain';
 import { RunCollectorCallbackHandler } from '@langchain/core/tracers/run_collector';
 import { Dataset } from '@kbn/elastic-assistant-common';
+import { isLangSmithEnabled } from '@kbn/langchain/server/tracers/langsmith';
 import { AgentExecutorEvaluatorWithMetadata } from '../langchain/executors/types';
 import { callAgentWithRetry, getMessageFromLangChainResponse } from './utils';
-import { isLangSmithEnabled, writeLangSmithFeedback } from '../../routes/evaluate/utils';
+import { writeLangSmithFeedback } from '../../routes/evaluate/utils';
 import { ResponseBody } from '../langchain/types';
 
 export interface PerformEvaluationParams {

diff --git a/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/helpers.ts b/x-pack/plugins/elastic_assistant/server/routes/attack_discovery/helpers.ts
@@ -28,7 +28,7 @@ import type { ActionsClient } from '@kbn/actions-plugin/server';
 import moment from 'moment/moment';
 import { uniq } from 'lodash/fp';
 import { PublicMethodsOf } from '@kbn/utility-types';
-import { getLangSmithTracer } from '../evaluate/utils';
+import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith';
 import { getLlmType } from '../utils';
 import type { GetRegisteredTools } from '../../services/app_context';
 import {

diff --git a/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts b/x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts
@@ -18,6 +18,7 @@ import {
   ExecuteConnectorRequestBody,
 } from '@kbn/elastic-assistant-common';
 import { ActionsClientLlm } from '@kbn/langchain/server';
+import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith';
 import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/schemas/common';
 import { ESQL_RESOURCE, KNOWLEDGE_BASE_INDEX_PATTERN } from '../knowledge_base/constants';
 import { buildResponse } from '../../lib/build_response';
@@ -29,7 +30,7 @@ import {
   indexEvaluations,
   setupEvaluationIndex,
 } from '../../lib/model_evaluator/output_index/utils';
-import { fetchLangSmithDataset, getConnectorName, getLangSmithTracer } from './utils';
+import { fetchLangSmithDataset, getConnectorName } from './utils';
 import { DEFAULT_PLUGIN_NAME, getPluginNameFromRequest } from '../helpers';
 
 /**

diff --git a/x-pack/plugins/elastic_assistant/server/routes/evaluate/utils.ts b/x-pack/plugins/elastic_assistant/server/routes/evaluate/utils.ts
@@ -10,8 +10,8 @@ import type { ActionResult } from '@kbn/actions-plugin/server';
 import type { Logger } from '@kbn/core/server';
 import type { Run } from 'langsmith/schemas';
 import { ToolingLog } from '@kbn/tooling-log';
-import { LangChainTracer } from '@langchain/core/tracers/tracer_langchain';
 import { Dataset } from '@kbn/elastic-assistant-common';
+import { isLangSmithEnabled } from '@kbn/langchain/server/tracers/langsmith';
 
 /**
  * Return connector name for the given connectorId/connectors
@@ -97,56 +97,3 @@ export const writeLangSmithFeedback = async (
     return '';
   }
 };
-
-/**
- * Returns a custom LangChainTracer which adds the `exampleId` so Dataset 'Test' runs are written to LangSmith
- * If `exampleId` is present (and a corresponding example exists in LangSmith) trace is written to the Dataset's `Tests`
- * section, otherwise it is written to the `Project` provided
- *
- * @param apiKey API Key for LangSmith (will fetch from env vars if not provided)
- * @param projectName Name of project to trace results to
- * @param exampleId Dataset exampleId to associate trace with
- * @param logger
- */
-export const getLangSmithTracer = ({
-  apiKey,
-  projectName,
-  exampleId,
-  logger,
-}: {
-  apiKey?: string;
-  projectName?: string;
-  exampleId?: string;
-  logger: Logger | ToolingLog;
-}): LangChainTracer[] => {
-  try {
-    if (!isLangSmithEnabled() && apiKey == null) {
-      return [];
-    }
-    const lcTracer = new LangChainTracer({
-      projectName, // Shows as the 'test' run's 'name' in langsmith ui
-      exampleId,
-      client: new Client({ apiKey }),
-    });
-
-    return [lcTracer];
-  } catch (e) {
-    // Note: creating a tracer can fail if the LangSmith env vars are not set correctly
-    logger.error(`Error creating LangSmith tracer: ${e.message}`);
-  }
-
-  return [];
-};
-
-/**
- * Returns true if LangSmith/tracing is enabled
- */
-export const isLangSmithEnabled = (): boolean => {
-  try {
-    // Just checking if apiKey is available, if better way to check for enabled that is not env var please update
-    const config = Client.getDefaultClientConfig();
-    return config.apiKey != null;
-  } catch (e) {
-    return false;
-  }
-};
diff --git a/x-pack/plugins/elastic_assistant/server/routes/helpers.ts b/x-pack/plugins/elastic_assistant/server/routes/helpers.ts
@@ -27,6 +27,7 @@ import { i18n } from '@kbn/i18n';
 import { AwaitedProperties, PublicMethodsOf } from '@kbn/utility-types';
 import { ActionsClient } from '@kbn/actions-plugin/server';
 import { AssistantFeatureKey } from '@kbn/elastic-assistant-common/impl/capabilities';
+import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith';
 import { MINIMUM_AI_ASSISTANT_LICENSE } from '../../common/constants';
 import { ESQL_RESOURCE, KNOWLEDGE_BASE_INDEX_PATTERN } from './knowledge_base/constants';
 import { callAgentExecutor } from '../lib/langchain/execute_custom_llm_chain';
@@ -39,7 +40,6 @@ import {
 import { executeAction, StaticResponse } from '../lib/executor';
 import { getLangChainMessages } from '../lib/langchain/helpers';
 
-import { getLangSmithTracer } from './evaluate/utils';
 import { ElasticsearchStore } from '../lib/langchain/elasticsearch_store/elasticsearch_store';
 import { AIAssistantConversationsDataClient } from '../ai_assistant_data_clients/conversations';
 import { INVOKE_ASSISTANT_SUCCESS_EVENT } from '../lib/telemetry/event_based_telemetry';

diff --git a/x-pack/plugins/elastic_assistant/tsconfig.json b/x-pack/plugins/elastic_assistant/tsconfig.json
@@ -27,7 +27,6 @@
     "@kbn/core-elasticsearch-server",
     "@kbn/logging",
     "@kbn/ml-plugin",
-    "@kbn/apm-utils",
     "@kbn/elastic-assistant-common",
     "@kbn/core-http-router-server-mocks",
     "@kbn/data-stream-adapter",
@@ -46,6 +45,7 @@
     "@kbn/langchain",
     "@kbn/stack-connectors-plugin",
     "@kbn/security-plugin",
+    "@kbn/apm-utils",
   ],
   "exclude": [
     "target/**/*",

diff --git a/.../plugins/integration_assistant/common/api/categorization/categorization_route.schema.yaml b/.../plugins/integration_assistant/common/api/categorization/categorization_route.schema.yaml
@@ -34,6 +34,8 @@ paths:
                   $ref: "../model/common_attributes.schema.yaml#/components/schemas/Pipeline"
                 connectorId:
                   $ref: "../model/common_attributes.schema.yaml#/components/schemas/Connector"
+                langSmithOptions:
+                  $ref: "../model/common_attributes.schema.yaml#/components/schemas/LangSmithOptions"
       responses:
         200:
           description: Indicates a successful call.

diff --git a/x-pack/plugins/integration_assistant/common/api/categorization/categorization_route.ts b/x-pack/plugins/integration_assistant/common/api/categorization/categorization_route.ts
@@ -10,6 +10,7 @@ import { z } from 'zod';
 import {
   Connector,
   DataStreamName,
+  LangSmithOptions,
   PackageName,
   Pipeline,
   RawSamples,
@@ -23,6 +24,7 @@ export const CategorizationRequestBody = z.object({
   rawSamples: RawSamples,
   currentPipeline: Pipeline,
   connectorId: Connector,
+  langSmithOptions: LangSmithOptions.optional(),
 });
 export type CategorizationRequestBodyInput = z.input<typeof CategorizationRequestBody>;
 

diff --git a/x-pack/plugins/integration_assistant/common/api/ecs/ecs_route.schema.yaml b/x-pack/plugins/integration_assistant/common/api/ecs/ecs_route.schema.yaml
@@ -33,6 +33,8 @@ paths:
                   $ref: "../model/common_attributes.schema.yaml#/components/schemas/Mapping"
                 connectorId:
                   $ref: "../model/common_attributes.schema.yaml#/components/schemas/Connector"
+                langSmithOptions:
+                  $ref: "../model/common_attributes.schema.yaml#/components/schemas/LangSmithOptions"
       responses:
         200:
           description: Indicates a successful call.

diff --git a/x-pack/plugins/integration_assistant/common/api/ecs/ecs_route.ts b/x-pack/plugins/integration_assistant/common/api/ecs/ecs_route.ts
@@ -10,6 +10,7 @@ import { z } from 'zod';
 import {
   Connector,
   DataStreamName,
+  LangSmithOptions,
   Mapping,
   PackageName,
   RawSamples,
@@ -23,6 +24,7 @@ export const EcsMappingRequestBody = z.object({
   rawSamples: RawSamples,
   mapping: Mapping.optional(),
   connectorId: Connector,
+  langSmithOptions: LangSmithOptions.optional(),
 });
 export type EcsMappingRequestBodyInput = z.input<typeof EcsMappingRequestBody>;
 

diff --git a/x-pack/plugins/integration_assistant/common/api/model/common_attributes.schema.yaml b/x-pack/plugins/integration_assistant/common/api/model/common_attributes.schema.yaml
@@ -143,3 +143,17 @@ components:
         logo:
           type: string
           description: The logo of the integration.
+
+    LangSmithOptions:
+      type: object
+      description: The LangSmith options object.
+      required:
+        - projectName
+        - apiKey
+      properties:
+        projectName:
+          type: string
+          description: The project name.
+        apiKey:
+          type: string
+          description: The apiKey to use for tracing.
diff --git a/x-pack/plugins/integration_assistant/common/api/model/common_attributes.ts b/x-pack/plugins/integration_assistant/common/api/model/common_attributes.ts
@@ -156,3 +156,18 @@ export const Integration = z.object({
    */
   logo: z.string().optional(),
 });
+
+/**
+ * The LangSmith options object.
+ */
+export type LangSmithOptions = z.infer<typeof LangSmithOptions>;
+export const LangSmithOptions = z.object({
+  /**
+   * The project name to use with tracing.
+   */
+  projectName: z.string(),
+  /**
+   * The api key for the project
+   */
+  apiKey: z.string(),
+});
diff --git a/x-pack/plugins/integration_assistant/common/api/related/related_route.schema.yaml b/x-pack/plugins/integration_assistant/common/api/related/related_route.schema.yaml
@@ -34,6 +34,8 @@ paths:
                   $ref: "../model/common_attributes.schema.yaml#/components/schemas/Pipeline"
                 connectorId:
                   $ref: "../model/common_attributes.schema.yaml#/components/schemas/Connector"
+                langSmithOptions:
+                  $ref: "../model/common_attributes.schema.yaml#/components/schemas/LangSmithOptions"
       responses:
         200:
           description: Indicates a successful call.

diff --git a/x-pack/plugins/integration_assistant/common/api/related/related_route.ts b/x-pack/plugins/integration_assistant/common/api/related/related_route.ts
@@ -10,6 +10,7 @@ import { z } from 'zod';
 import {
   Connector,
   DataStreamName,
+  LangSmithOptions,
   PackageName,
   Pipeline,
   RawSamples,
@@ -23,6 +24,7 @@ export const RelatedRequestBody = z.object({
   rawSamples: RawSamples,
   currentPipeline: Pipeline,
   connectorId: Connector,
+  langSmithOptions: LangSmithOptions.optional(),
 });
 export type RelatedRequestBodyInput = z.input<typeof RelatedRequestBody>;