From f45154734884ff248cc0c88b2d54707a2e70f1e0 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Wed, 26 Jun 2024 13:04:32 -0600 Subject: [PATCH] Initial package creation with readme and placeholder interface --- .../kbn-elastic-assistant-agents/README.md | 160 ++++++++++++++++++ .../kbn-elastic-assistant-agents/index.ts | 8 + .../jest.config.js | 12 ++ .../kbn-elastic-assistant-agents/kibana.jsonc | 5 + .../kbn-elastic-assistant-agents/package.json | 6 + .../src/interfaces.ts | 26 +++ .../tsconfig.json | 19 +++ 7 files changed, 236 insertions(+) create mode 100644 x-pack/packages/kbn-elastic-assistant-agents/README.md create mode 100644 x-pack/packages/kbn-elastic-assistant-agents/index.ts create mode 100644 x-pack/packages/kbn-elastic-assistant-agents/jest.config.js create mode 100644 x-pack/packages/kbn-elastic-assistant-agents/kibana.jsonc create mode 100644 x-pack/packages/kbn-elastic-assistant-agents/package.json create mode 100644 x-pack/packages/kbn-elastic-assistant-agents/src/interfaces.ts create mode 100644 x-pack/packages/kbn-elastic-assistant-agents/tsconfig.json diff --git a/x-pack/packages/kbn-elastic-assistant-agents/README.md b/x-pack/packages/kbn-elastic-assistant-agents/README.md new file mode 100644 index 0000000000000..171ef0fec7fdd --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-agents/README.md @@ -0,0 +1,160 @@ +# @kbn/elastic-assistant-agents + +Package for defining a common 'agent interface' for use in the Observability/Security AI Assistants, within the Search Playground, or anywhere within Kibana server code where common LLM-backed tasks are needed. + +In the context of building conversational UI's, or LLM-backed features like Contextual Insights and Attack Discovery, an 'agent' is defined as a function that takes an input query (or set of messages), some default dependencies, and returns a string response. This interface closely resembles function or tool definitions as used when making requests to LLMs. + +The reference agent presented in this package is the `esql-query-gen-agent`, which can be used to generate ESQL queries from a natural language input. + +The first revision of this interface is an extraction of common dependencies used by the Observability Assistant's [`query`](https://github.com/elastic/kibana/blob/4fc13a4a586db2b3cc351e6a3d87f2c6471c5a55/x-pack/plugins/observability_solution/observability_ai_assistant_app/server/functions/query/index.ts#L73) function, and the Security Assistant's [`esql-knowledge-base-tool`](https://github.com/elastic/kibana/blob/e2e1fb350464b94a0ebb691db04d16d5badd4ef1/x-pack/plugins/security_solution/server/assistant/tools/esql_language_knowledge_base/esql_language_knowledge_base_tool.ts#L37). + + +### Observability `registerQueryFunction` dependencies: + +Top level in-plugin registration via: + +```ts +export const registerFunctions = async (registrationParameters: FunctionRegistrationParameters) => {} +``` +with dependency chain: + +``` ts +export type FunctionRegistrationParameters = Omit< + Parameters[0], + 'registerContext' | 'hasFunction' +> & { pluginsStart: ObservabilityAIAssistantAppPluginStartDependencies }; + +///// + +// Main entry point interface: see `observability_ai_assistant_app/server/functions/index.ts` +export type RegistrationCallback = ({}: { + signal: AbortSignal; // unused + resources: RespondFunctionResources; + client: ObservabilityAIAssistantClient; // unused + functions: ChatFunctionClient; +}) => Promise; + +///// + +export type RespondFunctionResources = Pick< + ObservabilityAIAssistantRouteHandlerResources, + 'context' | 'logger' | 'plugins' | 'request' +>; + +///// + +export class ChatFunctionClient { + registerFunction: RegisterFunction = (definition, respond) => {; + registerInstruction: RegisterInstruction = (instruction) => {}; + validate(name: string, parameters: unknown) {} + getInstructions(): RegisteredInstruction[] {} + hasAction(name: string) {} + getFunctions({ filter }: { filter?: string; } = {}): FunctionHandler[] {} + getActions() {} + hasFunction(name: string): boolean {} + async executeFunction({ chat, name, args, messages, signal }: { + chat: FunctionCallChatFunction; + name: string; + args: string | undefined; + messages: Message[]; + signal: AbortSignal; + }): Promise {} +} +``` + +With main public interface for registration through `ChatFunctionClient`, by means of `registerFunction`, with input interface: + +```ts +export interface FunctionDefinition { + name: string; + description: string; + visibility?: FunctionVisibility; + descriptionForUser?: string; + parameters?: TParameters; +} +``` + +### Security `esql-knowledge-base-tool` tool dependencies: + +Main public interface for registration is through `pluginStart` via +``` +registerTools: (pluginName: string, tools: AssistantTool[]) => void; +``` + +With top-level `AssistantTool` interface defined as: + +```ts +export interface AssistantTool { + id: string; + name: string; + description: string; + sourceRegister: string; + isSupported: (params: AssistantToolParams) => boolean; + getTool: (params: AssistantToolParams) => Tool | DynamicStructuredTool | null; +} +``` + +and params provided to tools for execution: + +```ts +export interface AssistantToolParams { + alertsIndexPattern?: string; + anonymizationFields?: AnonymizationFieldResponse[]; + isEnabledKnowledgeBase: boolean; + chain?: RetrievalQAChain; + esClient: ElasticsearchClient; + kbDataClient?: AIAssistantKnowledgeBaseDataClient; + langChainTimeout?: number; + llm?: ActionsClientLlm | ActionsClientChatOpenAI | ActionsClientSimpleChatModel; + logger: Logger; + modelExists: boolean; + onNewReplacements?: (newReplacements: Replacements) => void; + replacements?: Replacements; + request: KibanaRequest< + unknown, + unknown, + ExecuteConnectorRequestBody | AttackDiscoveryPostRequestBody + >; + size?: number; +} +``` + + +## Interface Overview + +The Observability in-plugin `registerFunctions()` function (which is responsible for registering the esql query generation behavior), provides a swathe of dependencies, however actual usage falls mostly in line with the public API of the `ChatFunctionsClient`, with `resources` only providing access to a scoped `esClient` and `logger`. The public `registerFunction()` method of `ChatFunctionsClient` handles the function definition, json schema, and any additional arguments. + +On the Security side, the same publicly exposed `registerTools()` function is used for registering the `esql-knowledge-base-tool`, and makes available a hodgepodge of dependencies (mostly optional) via the `AssistantToolParams` defined above, with the `getTool` callback encapsulating the zod schema definition by means of the LangChain `Tool | DynamicStructuredTool` abstractions. + +### Interface Commonalities + +Between the two, we have common function/tool declaration fields, like `id`, `name`, and `description`, some differing descriptor fields like `visibility`, `isSupported` and `sourceRegister` (registering plugin), and then the generic params needed by the function/tool, like `esClients`, the `Kibana Request`, and abstractions for making requests to the LLM. + +Note: Generic params are mostly made available via each implementation's provided callbacks, either the `respond` callback when registering the function with the Observability Assistant, or the `getTool` callback when registering with Security Assistant. + +--- + +### Proposed common interface: + +```ts +export interface AgentInterface { + id: string; + name: string; + description: string; + visibility?: FunctionVisibility; + isSupported?: (params: AssistantToolParams) => boolean; + sourceRegister?: string; + execute: (params: AgentParams) => Promise; +} + +export interface AgentParams { + esClient: ElasticsearchClient; + logger: Logger; + request: KibanaRequest; + // Additional parameters can be added as needed +} + +export interface AgentResponse { + // Define the structure of the function response +} +``` \ No newline at end of file diff --git a/x-pack/packages/kbn-elastic-assistant-agents/index.ts b/x-pack/packages/kbn-elastic-assistant-agents/index.ts new file mode 100644 index 0000000000000..f16a23a804aaa --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-agents/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 { AgentInterface, AgentParams, AgentResponse } from './src/interfaces'; diff --git a/x-pack/packages/kbn-elastic-assistant-agents/jest.config.js b/x-pack/packages/kbn-elastic-assistant-agents/jest.config.js new file mode 100644 index 0000000000000..0bbbd34b5967c --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-agents/jest.config.js @@ -0,0 +1,12 @@ +/* + * 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. + */ + +module.exports = { + preset: '@kbn/test/jest_node', + rootDir: '../../..', + roots: ['/x-pack/packages/kbn-elastic-assistant-agents'], +}; diff --git a/x-pack/packages/kbn-elastic-assistant-agents/kibana.jsonc b/x-pack/packages/kbn-elastic-assistant-agents/kibana.jsonc new file mode 100644 index 0000000000000..2b29fc4aea388 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-agents/kibana.jsonc @@ -0,0 +1,5 @@ +{ + "type": "shared-common", + "id": "@kbn/elastic-assistant-agents", + "owner": ["@elastic/obs-knowledge-team", "@elastic/search-kibana", "@elastic/security-generative-ai"] +} diff --git a/x-pack/packages/kbn-elastic-assistant-agents/package.json b/x-pack/packages/kbn-elastic-assistant-agents/package.json new file mode 100644 index 0000000000000..180e22b873613 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-agents/package.json @@ -0,0 +1,6 @@ +{ + "name": "@kbn/elastic-assistant-agents", + "private": true, + "version": "1.0.0", + "license": "Elastic License 2.0" +} \ No newline at end of file diff --git a/x-pack/packages/kbn-elastic-assistant-agents/src/interfaces.ts b/x-pack/packages/kbn-elastic-assistant-agents/src/interfaces.ts new file mode 100644 index 0000000000000..4a663ed023b2c --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-agents/src/interfaces.ts @@ -0,0 +1,26 @@ +/* + * 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 interface AgentInterface { + id: string; + name: string; + description: string; + visibility?: FunctionVisibility; + isSupported?: (params: AssistantToolParams) => boolean; + sourceRegister?: string; + execute: (params: AgentParams) => Promise; +} + +export interface AgentParams { + esClient: ElasticsearchClient; + logger: Logger; + request: KibanaRequest; + // Additional parameters can be added as needed +} + +// eslint-disable-next-line @typescript-eslint/no-empty-interface +export interface AgentResponse {} diff --git a/x-pack/packages/kbn-elastic-assistant-agents/tsconfig.json b/x-pack/packages/kbn-elastic-assistant-agents/tsconfig.json new file mode 100644 index 0000000000000..b270f5a1b6c67 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-agents/tsconfig.json @@ -0,0 +1,19 @@ +{ + "extends": "../../../tsconfig.base.json", + "compilerOptions": { + "outDir": "target/types", + "types": [ + "jest", + "node" + ] + }, + "include": [ + "**/*.ts", + ], + "exclude": [ + "target/**/*" + ], + "kbn_references": [ + "@kbn/tooling-log" + ] +}