[Security Assistant] Migrates to LangGraph and adds KB Tools

x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.gen.ts

@@ -33,11 +33,11 @@ export const KnowledgeBaseEntryErrorSchema = z
 export type Metadata = z.infer<typeof Metadata>;
 export const Metadata = z.object({
   /**
-   * Knowledge Base resource name
+   * Knowledge Base resource name for grouping entries, e.g. 'esql', 'lens-docs', etc
    */
   kbResource: z.string(),
   /**
-   * Original text content source
+   * Source document name or filepath
    */
   source: z.string(),
   /**

x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/common_attributes.schema.yaml

@@ -32,10 +32,10 @@ components:
       properties:
         kbResource:
           type: string
-          description: Knowledge Base resource name
+          description: Knowledge Base resource name for grouping entries, e.g. 'esql', 'lens-docs', etc
         source:
           type: string
-          description: Original text content source
+          description: Source document name or filepath
         required:
           type: boolean
           description: Whether or not this resource should always be included

x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.ts

@@ -18,8 +18,8 @@ import { CreateKnowledgeBaseEntrySchema } from './types';
 
 export interface CreateKnowledgeBaseEntryParams {
   esClient: ElasticsearchClient;
-  logger: Logger;
   knowledgeBaseIndex: string;
+  logger: Logger;
   spaceId: string;
   user: AuthenticatedUser;
   knowledgeBaseEntry: KnowledgeBaseEntryCreateProps;

x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts

@@ -6,6 +6,8 @@
  */
 
 import { errors } from '@elastic/elasticsearch';
+import { QueryDslQueryContainer } from '@elastic/elasticsearch/lib/api/types';
+import { AuthenticatedUser } from '@kbn/core-security-common';
 
 export const isModelAlreadyExistsError = (error: Error) => {
   return (
@@ -14,3 +16,87 @@ export const isModelAlreadyExistsError = (error: Error) => {
       error.body.error.type === 'status_exception')
   );
 };
+
+/**
+ * Returns an Elasticsearch query DSL that performs a vector search against the Knowledge Base for the given query/user/filter.
+ *
+ * @param filter - Optional filter to apply to the search
+ * @param kbResource - Specific resource tag to filter for, e.g. 'esql' or 'user'
+ * @param modelId - ID of the model to search with, e.g. `.elser_model_2`
+ * @param query - The search query provided by the user
+ * @param required - Whether to only include required entries
+ * @param user - The authenticated user
+ * @returns
+ */
+export const getKBVectorSearchQuery = ({
+  filter,
+  kbResource,
+  modelId,
+  query,
+  required,
+  user,
+}: {
+  filter?: QueryDslQueryContainer | undefined;
+  kbResource?: string | undefined;
+  modelId: string;
+  query: string;
+  required?: boolean | undefined;
+  user: AuthenticatedUser;
+}): QueryDslQueryContainer => {
+  const resourceFilter = kbResource
+    ? [
+        {
+          term: {
+            'metadata.kbResource': kbResource,
+          },
+        },
+      ]
+    : [];
+  const requiredFilter = required
+    ? [
+        {
+          term: {
+            'metadata.required': required,
+          },
+        },
+      ]
+    : [];
+
+  const userFilter = [
+    {
+      nested: {
+        path: 'users',
+        query: {
+          bool: {
+            must: [
+              {
+                match: user.profile_uid
+                  ? { 'users.id': user.profile_uid }
+                  : { 'users.name': user.username },
+              },
+            ],
+          },
+        },
+      },
+    },
+  ];
+
+  return {
+    bool: {
+      must: [
+        {
+          text_expansion: {
+            'vector.tokens': {
+              model_id: modelId,
+              model_text: query,
+            },
+          },
+        },
+        ...requiredFilter,
+        ...resourceFilter,
+        ...userFilter,
+      ],
+      filter,
+    },
+  };
+};

x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts

@@ -11,19 +11,23 @@ import {
 } from '@elastic/elasticsearch/lib/api/types';
 import type { MlPluginSetup } from '@kbn/ml-plugin/server';
 import type { KibanaRequest } from '@kbn/core-http-server';
-import type { Document } from 'langchain/document';
+import { Document } from 'langchain/document';
 import type { SavedObjectsClientContract } from '@kbn/core-saved-objects-api-server';
-import { KnowledgeBaseEntryResponse } from '@kbn/elastic-assistant-common';
+import {
+  KnowledgeBaseEntryCreateProps,
+  KnowledgeBaseEntryResponse,
+} from '@kbn/elastic-assistant-common';
 import pRetry from 'p-retry';
+import { QueryDslQueryContainer } from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
 import { AIAssistantDataClient, AIAssistantDataClientParams } from '..';
 import { ElasticsearchStore } from '../../lib/langchain/elasticsearch_store/elasticsearch_store';
 import { loadESQL } from '../../lib/langchain/content_loaders/esql_loader';
 import { GetElser } from '../../types';
-import { transformToCreateSchema } from './create_knowledge_base_entry';
+import { createKnowledgeBaseEntry, transformToCreateSchema } from './create_knowledge_base_entry';
 import { EsKnowledgeBaseEntrySchema } from './types';
 import { transformESSearchToKnowledgeBaseEntry } from './transforms';
 import { ESQL_DOCS_LOADED_QUERY } from '../../routes/knowledge_base/constants';
-import { isModelAlreadyExistsError } from './helpers';
+import { getKBVectorSearchQuery, isModelAlreadyExistsError } from './helpers';
 
 interface KnowledgeBaseDataClientParams extends AIAssistantDataClientParams {
   ml: MlPluginSetup;
@@ -217,8 +221,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
   /**
    * Adds LangChain Documents to the knowledge base
    *
-   * @param documents
-   * @param authenticatedUser
+   * @param documents LangChain Documents to add to the knowledge base
    */
   public addKnowledgeBaseDocuments = async ({
     documents,
@@ -261,4 +264,100 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
 
     return created?.data ? transformESSearchToKnowledgeBaseEntry(created?.data) : [];
   };
+
+  /**
+   * Performs similarity search to retrieve LangChain Documents from the knowledge base
+   */
+  public getKnowledgeBaseDocuments = async ({
+    filter,
+    kbResource,
+    query,
+    required,
+  }: {
+    filter?: QueryDslQueryContainer;
+    kbResource?: string;
+    query: string;
+    required?: boolean;
+  }): Promise<Document[]> => {
+    const user = this.options.currentUser;
+    if (user == null) {
+      throw new Error(
+        'Authenticated user not found! Ensure kbDataClient was initialized from a request.'
+      );
+    }
+
+    const esClient = await this.options.elasticsearchClientPromise;
+    const modelId = await this.options.getElserId();
+
+    const vectorSearchQuery = getKBVectorSearchQuery({
+      filter,
+      kbResource,
+      modelId,
+      query,
+      required,
+      user,
+    });
+
+    try {
+      const result = await esClient.search<EsKnowledgeBaseEntrySchema>({
+        index: this.indexTemplateAndPattern.alias,
+        size: 10,
+        query: vectorSearchQuery,
+      });
+
+      const results = result.hits.hits.map(
+        (hit) =>
+          new Document({
+            pageContent: hit?._source?.text ?? '',
+            metadata: hit?._source?.metadata ?? {},
+          })
+      );
+
+      this.options.logger.debug(
+        `getKnowledgeBaseDocuments() - Similarity Search Query:\n ${JSON.stringify(
+          vectorSearchQuery
+        )}`
+      );
+      this.options.logger.debug(
+        `getKnowledgeBaseDocuments() - Similarity Search Results:\n ${JSON.stringify(results)}`
+      );
+
+      return results;
+    } catch (e) {
+      this.options.logger.error(`Error performing KB Similarity Search: ${e.message}`);
+      return [];
+    }
+  };
+
+  /**
+   * Creates a new Knowledge Base Entry.
+   *
+   * @param knowledgeBaseEntry
+   */
+  public createKnowledgeBaseEntry = async ({
+    knowledgeBaseEntry,
+  }: {
+    knowledgeBaseEntry: KnowledgeBaseEntryCreateProps;
+  }): Promise<KnowledgeBaseEntryResponse | null> => {
+    const authenticatedUser = this.options.currentUser;
+    if (authenticatedUser == null) {
+      throw new Error(
+        'Authenticated user not found! Ensure kbDataClient was initialized from a request.'
+      );
+    }
+
+    this.options.logger.debug(
+      `Creating Knowledge Base Entry:\n ${JSON.stringify(knowledgeBaseEntry, null, 2)}`
+    );
+    this.options.logger.debug(`kbIndex: ${this.indexTemplateAndPattern.alias}`);
+    const esClient = await this.options.elasticsearchClientPromise;
+    return createKnowledgeBaseEntry({
+      esClient,
+      knowledgeBaseIndex: this.indexTemplateAndPattern.alias,
+      logger: this.options.logger,
+      spaceId: this.spaceId,
+      user: authenticatedUser,
+      knowledgeBaseEntry,
+    });
+  };
 }

x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts

@@ -89,12 +89,13 @@ export const callAgentExecutor: AgentExecutor<true | false> = async ({
 
   // Fetch any applicable tools that the source plugin may have registered
   const assistantToolParams: AssistantToolParams = {
-    anonymizationFields,
     alertsIndexPattern,
-    isEnabledKnowledgeBase,
+    anonymizationFields,
     chain,
-    llm,
     esClient,
+    isEnabledKnowledgeBase,
+    llm,
+    logger,
     modelExists,
     onNewReplacements,
     replacements,

x-pack/plugins/elastic_assistant/server/lib/langchain/executors/types.ts

@@ -9,14 +9,29 @@ import { PluginStartContract as ActionsPluginStart } from '@kbn/actions-plugin/s
 import { ElasticsearchClient } from '@kbn/core-elasticsearch-server';
 import { BaseMessage } from '@langchain/core/messages';
 import { Logger } from '@kbn/logging';
-import { KibanaRequest, ResponseHeaders } from '@kbn/core-http-server';
+import { KibanaRequest, KibanaResponseFactory, ResponseHeaders } from '@kbn/core-http-server';
 import type { LangChainTracer } from '@langchain/core/tracers/tracer_langchain';
 import { ExecuteConnectorRequestBody, Message, Replacements } from '@kbn/elastic-assistant-common';
 import { StreamResponseWithHeaders } from '@kbn/ml-response-stream/server';
 import { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen';
 import { ResponseBody } from '../types';
 import type { AssistantTool } from '../../../types';
 import { ElasticsearchStore } from '../elasticsearch_store/elasticsearch_store';
+import { AIAssistantKnowledgeBaseDataClient } from '../../../ai_assistant_data_clients/knowledge_base';
+import { AIAssistantConversationsDataClient } from '../../../ai_assistant_data_clients/conversations';
+import { AIAssistantDataClient } from '../../../ai_assistant_data_clients';
+
+export type OnLlmResponse = (
+  content: string,
+  traceData?: Message['traceData'],
+  isError?: boolean
+) => Promise<void>;
+
+export interface AssistantDataClients {
+  anonymizationFieldsDataClient?: AIAssistantDataClient;
+  conversationsDataClient?: AIAssistantConversationsDataClient;
+  kbDataClient?: AIAssistantKnowledgeBaseDataClient;
+}
 
 export interface AgentExecutorParams<T extends boolean> {
   abortSignal?: AbortSignal;
@@ -26,6 +41,8 @@ export interface AgentExecutorParams<T extends boolean> {
   isEnabledKnowledgeBase: boolean;
   assistantTools?: AssistantTool[];
   connectorId: string;
+  conversationId?: string;
+  dataClients?: AssistantDataClients;
   esClient: ElasticsearchClient;
   esStore: ElasticsearchStore;
   langChainMessages: BaseMessage[];
@@ -34,12 +51,9 @@ export interface AgentExecutorParams<T extends boolean> {
   onNewReplacements?: (newReplacements: Replacements) => void;
   replacements: Replacements;
   isStream?: T;
-  onLlmResponse?: (
-    content: string,
-    traceData?: Message['traceData'],
-    isError?: boolean
-  ) => Promise<void>;
+  onLlmResponse?: OnLlmResponse;
   request: KibanaRequest<unknown, unknown, ExecuteConnectorRequestBody>;
+  response?: KibanaResponseFactory;
   size?: number;
   traceOptions?: TraceOptions;
 }