From 92099b277dc5b1448d14994d280674611ca9e261 Mon Sep 17 00:00:00 2001 From: Sergi Massaneda Date: Mon, 8 Jul 2024 20:14:57 +0200 Subject: [PATCH] [Security GenAI] Use AI setting to set langsmith tracing to the Integration Assistant (#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 Settings After one execution of the integration assistant: langsmith --------- Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Co-authored-by: Elastic Machine --- .../kbn-langchain/server}/tracers/README.mdx | 0 .../server/tracers/apm}/apm_tracer.ts | 0 .../kbn-langchain/server/tracers/apm/index.ts | 8 +++ .../server/tracers/langsmith/index.ts | 7 ++ .../tracers/langsmith/langsmith_tracer.ts | 64 +++++++++++++++++++ x-pack/packages/kbn-langchain/tsconfig.json | 3 +- .../execute_custom_llm_chain/index.ts | 4 +- .../executors/openai_functions_executor.ts | 4 +- .../graphs/default_assistant_graph/helpers.ts | 4 +- .../graphs/default_assistant_graph/index.ts | 2 +- .../tracers/{ => apm}/with_assistant_span.ts | 0 .../server/lib/model_evaluator/evaluation.ts | 3 +- .../server/routes/attack_discovery/helpers.ts | 2 +- .../server/routes/evaluate/post_evaluate.ts | 3 +- .../server/routes/evaluate/utils.ts | 55 +--------------- .../server/routes/helpers.ts | 2 +- .../plugins/elastic_assistant/tsconfig.json | 2 +- .../categorization_route.schema.yaml | 2 + .../categorization/categorization_route.ts | 2 + .../common/api/ecs/ecs_route.schema.yaml | 2 + .../common/api/ecs/ecs_route.ts | 2 + .../api/model/common_attributes.schema.yaml | 14 ++++ .../common/api/model/common_attributes.ts | 15 +++++ .../api/related/related_route.schema.yaml | 2 + .../common/api/related/related_route.ts | 2 + .../integration_assistant/kibana.jsonc | 3 + .../public/common/lib/lang_smith.ts | 37 +++++++++++ .../data_stream_step/generation_modal.tsx | 2 + .../server/routes/categorization_routes.ts | 20 ++++-- .../server/routes/ecs_routes.ts | 34 +++++----- .../server/routes/related_routes.ts | 19 ++++-- .../integration_assistant/tsconfig.json | 3 +- 32 files changed, 231 insertions(+), 91 deletions(-) rename x-pack/{plugins/elastic_assistant/server/lib/langchain => packages/kbn-langchain/server}/tracers/README.mdx (100%) rename x-pack/{plugins/elastic_assistant/server/lib/langchain/tracers => packages/kbn-langchain/server/tracers/apm}/apm_tracer.ts (100%) create mode 100644 x-pack/packages/kbn-langchain/server/tracers/apm/index.ts create mode 100644 x-pack/packages/kbn-langchain/server/tracers/langsmith/index.ts create mode 100644 x-pack/packages/kbn-langchain/server/tracers/langsmith/langsmith_tracer.ts rename x-pack/plugins/elastic_assistant/server/lib/langchain/tracers/{ => apm}/with_assistant_span.ts (100%) create mode 100644 x-pack/plugins/integration_assistant/public/common/lib/lang_smith.ts diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/tracers/README.mdx b/x-pack/packages/kbn-langchain/server/tracers/README.mdx similarity index 100% rename from x-pack/plugins/elastic_assistant/server/lib/langchain/tracers/README.mdx rename to x-pack/packages/kbn-langchain/server/tracers/README.mdx diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/tracers/apm_tracer.ts b/x-pack/packages/kbn-langchain/server/tracers/apm/apm_tracer.ts similarity index 100% rename from x-pack/plugins/elastic_assistant/server/lib/langchain/tracers/apm_tracer.ts rename to x-pack/packages/kbn-langchain/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 new file mode 100644 index 0000000000000..293db467d740b --- /dev/null +++ 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 new file mode 100644 index 0000000000000..179a5aa297c7e --- /dev/null +++ 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 new file mode 100644 index 0000000000000..15501d22dbe09 --- /dev/null +++ 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 index 949aca47794ec..e775b1a4ad50d 100644 --- 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 index 86e01a93ddcdc..cd1a032406326 100644 --- 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 index 6aa1aa3ce7890..edea22a888dff 100644 --- 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/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts index 482c89c10e969..fe46a5deae9fe 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/helpers.ts +++ b/x-pack/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/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts index 517ac10479461..87437a98898d4 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts +++ b/x-pack/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/x-pack/plugins/elastic_assistant/server/lib/langchain/tracers/with_assistant_span.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/tracers/apm/with_assistant_span.ts similarity index 100% rename from x-pack/plugins/elastic_assistant/server/lib/langchain/tracers/with_assistant_span.ts rename to x-pack/plugins/elastic_assistant/server/lib/langchain/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 index d6fb04f2a13dc..93f164835876a 100644 --- 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 index 5f5eb8d0d8659..1de3b86e74deb 100644 --- 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 index 990417b799234..79b601ed01073 100644 --- 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 index 17757b8778771..46909805510e2 100644 --- 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 index aa060e24bc5df..a8fe0d0c25d68 100644 --- 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 index 8f546d6e5fe01..a2d0ec6c1d68a 100644 --- 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/x-pack/plugins/integration_assistant/common/api/categorization/categorization_route.schema.yaml b/x-pack/plugins/integration_assistant/common/api/categorization/categorization_route.schema.yaml index 37b5750e01aca..e04a41b88e564 100644 --- a/x-pack/plugins/integration_assistant/common/api/categorization/categorization_route.schema.yaml +++ b/x-pack/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 index 067c3c88a07ed..c8e56a2af2f5e 100644 --- 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; 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 index d539ae5877da8..57505f64ebf2a 100644 --- 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 index b1973159f9304..fd1dca194ae59 100644 --- 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; 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 index 24cb71ed5274c..527899dc33727 100644 --- 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 index fdb20931ccce0..bde01d8bae245 100644 --- 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; +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 index 13990dc1b66d8..a588ac5852822 100644 --- 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 index 4c5242e48d2c3..961bc007adcd8 100644 --- 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; diff --git a/x-pack/plugins/integration_assistant/kibana.jsonc b/x-pack/plugins/integration_assistant/kibana.jsonc index b2ef3045e12be..d8d6e17026230 100644 --- a/x-pack/plugins/integration_assistant/kibana.jsonc +++ b/x-pack/plugins/integration_assistant/kibana.jsonc @@ -18,5 +18,8 @@ "actions", "stackConnectors", ], + "requiredBundles": [ + "kibanaUtils", + ] } } diff --git a/x-pack/plugins/integration_assistant/public/common/lib/lang_smith.ts b/x-pack/plugins/integration_assistant/public/common/lib/lang_smith.ts new file mode 100644 index 0000000000000..7234870439930 --- /dev/null +++ b/x-pack/plugins/integration_assistant/public/common/lib/lang_smith.ts @@ -0,0 +1,37 @@ +/* + * 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 { + DEFAULT_ASSISTANT_NAMESPACE, + TRACE_OPTIONS_SESSION_STORAGE_KEY, +} from '@kbn/elastic-assistant/impl/assistant_context/constants'; +import { Storage } from '@kbn/kibana-utils-plugin/public'; +import type { TraceOptions } from '@kbn/elastic-assistant/impl/assistant/types'; +import type { LangSmithOptions } from '../../../common/api/model/common_attributes'; + +const sessionStorage = new Storage(window.sessionStorage); + +/** + * Retrieves the LangSmith options from the AI Settings. + */ +export const getLangSmithOptions = ( + nameSpace: string = DEFAULT_ASSISTANT_NAMESPACE +): LangSmithOptions | undefined => { + // Get the LangSmith options stored by the AI Settings using the assistant context + // TODO: Encapsulate all AI Settings logic in a generic place. + const sessionStorageTraceOptions: TraceOptions = sessionStorage.get( + `${nameSpace}.${TRACE_OPTIONS_SESSION_STORAGE_KEY}` + ); + + if (!sessionStorageTraceOptions) { + return; + } + return { + projectName: sessionStorageTraceOptions.langSmithProject, + apiKey: sessionStorageTraceOptions.langSmithApiKey, + }; +}; diff --git a/x-pack/plugins/integration_assistant/public/components/create_integration/create_integration_assistant/steps/data_stream_step/generation_modal.tsx b/x-pack/plugins/integration_assistant/public/components/create_integration/create_integration_assistant/steps/data_stream_step/generation_modal.tsx index fe5a1dabb1374..237f4cd94ee33 100644 --- a/x-pack/plugins/integration_assistant/public/components/create_integration/create_integration_assistant/steps/data_stream_step/generation_modal.tsx +++ b/x-pack/plugins/integration_assistant/public/components/create_integration/create_integration_assistant/steps/data_stream_step/generation_modal.tsx @@ -24,6 +24,7 @@ import { import { isEmpty } from 'lodash/fp'; import React, { useCallback, useEffect, useMemo, useState } from 'react'; import { css } from '@emotion/react'; +import { getLangSmithOptions } from '../../../../../common/lib/lang_smith'; import type { CategorizationRequestBody, EcsMappingRequestBody, @@ -87,6 +88,7 @@ export const useGeneration = ({ dataStreamName: integrationSettings.dataStreamName ?? '', rawSamples: integrationSettings.logsSampleParsed ?? [], connectorId: connector.id, + langSmithOptions: getLangSmithOptions(), }; setProgress('ecs'); diff --git a/x-pack/plugins/integration_assistant/server/routes/categorization_routes.ts b/x-pack/plugins/integration_assistant/server/routes/categorization_routes.ts index 2dbdc63210a59..80ebe9eb65258 100644 --- a/x-pack/plugins/integration_assistant/server/routes/categorization_routes.ts +++ b/x-pack/plugins/integration_assistant/server/routes/categorization_routes.ts @@ -11,6 +11,8 @@ import { ActionsClientChatOpenAI, ActionsClientSimpleChatModel, } from '@kbn/langchain/server/language_models'; +import { APMTracer } from '@kbn/langchain/server/tracers/apm'; +import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith'; import { CATEGORIZATION_GRAPH_PATH, CategorizationRequestBody, @@ -46,7 +48,8 @@ export function registerCategorizationRoutes( }, withAvailability( async (context, req, res): Promise> => { - const { packageName, dataStreamName, rawSamples, currentPipeline } = req.body; + const { packageName, dataStreamName, rawSamples, currentPipeline, langSmithOptions } = + req.body; const services = await context.resolve(['core']); const { client } = services.core.elasticsearch; const { getStartServices, logger } = await context.integrationAssistant; @@ -76,13 +79,22 @@ export function registerCategorizationRoutes( streaming: false, }); - const graph = await getCategorizationGraph(client, model); - const results = await graph.invoke({ + const parameters = { packageName, dataStreamName, rawSamples, currentPipeline, - }); + }; + const options = { + callbacks: [ + new APMTracer({ projectName: langSmithOptions?.projectName ?? 'default' }, logger), + ...getLangSmithTracer({ ...langSmithOptions, logger }), + ], + }; + + const graph = await getCategorizationGraph(client, model); + const results = await graph.invoke(parameters, options); + return res.ok({ body: CategorizationResponse.parse(results) }); } catch (e) { return res.badRequest({ body: e }); diff --git a/x-pack/plugins/integration_assistant/server/routes/ecs_routes.ts b/x-pack/plugins/integration_assistant/server/routes/ecs_routes.ts index d177aeb4b2cdf..69b13a004e98e 100644 --- a/x-pack/plugins/integration_assistant/server/routes/ecs_routes.ts +++ b/x-pack/plugins/integration_assistant/server/routes/ecs_routes.ts @@ -11,6 +11,8 @@ import { ActionsClientChatOpenAI, ActionsClientSimpleChatModel, } from '@kbn/langchain/server/language_models'; +import { APMTracer } from '@kbn/langchain/server/tracers/apm'; +import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith'; import { ECS_GRAPH_PATH, EcsMappingRequestBody, EcsMappingResponse } from '../../common'; import { ROUTE_HANDLER_TIMEOUT } from '../constants'; import { getEcsGraph } from '../graphs/ecs'; @@ -39,7 +41,7 @@ export function registerEcsRoutes(router: IRouter> => { - const { packageName, dataStreamName, rawSamples, mapping } = req.body; + const { packageName, dataStreamName, rawSamples, mapping, langSmithOptions } = req.body; const { getStartServices, logger } = await context.integrationAssistant; const [, { actions: actionsPlugin }] = await getStartServices(); @@ -53,6 +55,7 @@ export function registerEcsRoutes(router: IRouter> => { - const { packageName, dataStreamName, rawSamples, currentPipeline } = req.body; + const { packageName, dataStreamName, rawSamples, currentPipeline, langSmithOptions } = + req.body; const services = await context.resolve(['core']); const { client } = services.core.elasticsearch; const { getStartServices, logger } = await context.integrationAssistant; @@ -68,13 +71,21 @@ export function registerRelatedRoutes(router: IRouter