[o11y AI] Add get_hosts tool

src/platform/packages/shared/kbn-synthtrace-client/src/lib/infra/host.ts

@@ -50,16 +50,24 @@ export class Host extends Entity<HostDocument> {
     });
   }
 
-  memory() {
+  memory(fields?: {
+    'system.memory.actual.free'?: number;
+    'system.memory.actual.used.bytes'?: number;
+    'system.memory.actual.used.pct'?: number;
+    'system.memory.total'?: number;
+    'system.memory.used.bytes'?: number;
+    'system.memory.used.pct'?: number;
+    'process.memory.pct'?: number;
+  }) {
     return new HostMetrics({
       ...this.fields,
-      'system.memory.actual.free': 44704067584,
-      'system.memory.actual.used.bytes': 24015409152,
-      'system.memory.actual.used.pct': 0.35,
-      'system.memory.total': 68719476736,
-      'system.memory.used.bytes': 39964708864,
-      'system.memory.used.pct': 0.582,
-      'process.memory.pct': 0.1,
+      'system.memory.actual.free': fields?.['system.memory.actual.free'] ?? 44704067584,
+      'system.memory.actual.used.bytes': fields?.['system.memory.actual.used.bytes'] ?? 24015409152,
+      'system.memory.actual.used.pct': fields?.['system.memory.actual.used.pct'] ?? 0.35,
+      'system.memory.total': fields?.['system.memory.total'] ?? 68719476736,
+      'system.memory.used.bytes': fields?.['system.memory.used.bytes'] ?? 39964708864,
+      'system.memory.used.pct': fields?.['system.memory.used.pct'] ?? 0.582,
+      'process.memory.pct': fields?.['process.memory.pct'] ?? 0.1,
       'metricset.period': 10000,
       'metricset.name': 'memory',
     });

x-pack/platform/packages/shared/onechat/onechat-common/tools/tool_result.ts

@@ -75,7 +75,10 @@ export interface VisualizationResult {
   };
 }
 
-export type OtherResult = ToolResultMixin<ToolResultType.other, Record<string, unknown>>;
+export type OtherResult<T extends Object = Record<string, unknown>> = ToolResultMixin<
+  ToolResultType.other,
+  T
+>;
 
 export type ErrorResult = ToolResultMixin<
   ToolResultType.error,
@@ -86,13 +89,13 @@ export type ErrorResult = ToolResultMixin<
   }
 >;
 
-export type ToolResult =
+export type ToolResult<T extends Object = Record<string, unknown>> =
   | ResourceResult
   | TabularDataResult
   | QueryResult
   | VisualizationResult
   | DashboardResult
-  | OtherResult
+  | OtherResult<T>
   | ErrorResult;
 
 export const isResourceResult = (result: ToolResult): result is ResourceResult => {

x-pack/platform/packages/shared/onechat/onechat-server/allow_lists.ts

@@ -25,6 +25,7 @@ export const AGENT_BUILDER_BUILTIN_TOOLS: string[] = [
   `${internalNamespaces.observability}.get_services`,
   `${internalNamespaces.observability}.get_downstream_dependencies`,
   `${internalNamespaces.observability}.get_correlated_logs`,
+  `${internalNamespaces.observability}.get_hosts`,
 
   // Dashboards
   'platform.dashboard.create_dashboard',

x-pack/platform/packages/shared/onechat/onechat-server/tools/handler.ts

@@ -21,13 +21,16 @@ import type {
 /**
  * Tool result as returned by the tool handler.
  */
-export type ToolHandlerResult = Omit<ToolResult, 'tool_result_id'> & { tool_result_id?: string };
+export type ToolHandlerResult<TResult extends ToolResult = ToolResult> = Omit<
+  TResult,
+  'tool_result_id'
+> & { tool_result_id?: string };
 
 /**
  * Return value for {@link ToolHandlerFn} / {@link BuiltinToolDefinition}
  */
-export interface ToolHandlerReturn {
-  results: ToolHandlerResult[];
+export interface ToolHandlerReturn<TResult extends ToolResult = ToolResult> {
+  results: Array<ToolHandlerResult<TResult>>;
 }
 
 /**

x-pack/solutions/observability/plugins/infra/kibana.jsonc

@@ -50,6 +50,7 @@
       "cloud",
       "profilingDataAccess",
       "observabilityAIAssistant",
+      "observabilityAgentBuilder",
       "licenseManagement",
       "serverless"
     ],

x-pack/solutions/observability/plugins/infra/moon.yml

@@ -132,6 +132,7 @@ dependsOn:
   - '@kbn/deeplinks-analytics'
   - '@kbn/esql-composer'
   - '@kbn/fleet-plugin'
+  - '@kbn/observability-agent-builder-plugin'
 tags:
   - plugin
   - prod

x-pack/solutions/observability/plugins/infra/server/agent_builder/register_data_providers.ts

@@ -0,0 +1,131 @@
+/*
+ * 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 type { CoreSetup, Logger, KibanaRequest, CoreRequestHandlerContext } from '@kbn/core/server';
+import type {
+  InfraServerPluginSetupDeps,
+  InfraServerPluginStartDeps,
+} from '../lib/adapters/framework';
+import { getHosts } from '../routes/infra/lib/host/get_hosts';
+import type { InfraBackendLibs } from '../lib/infra_types';
+import { getInfraMetricsClient } from '../lib/helpers/get_infra_metrics_client';
+import { getInfraAlertsClient } from '../lib/helpers/get_infra_alerts_client';
+import { getInfraRequestHandlerContext } from '../utils/get_infra_request_handler_context';
+import type { InfraPluginRequestHandlerContext } from '../types';
+import { getApmDataAccessClient } from '../lib/helpers/get_apm_data_access_client';
+
+// Default metrics to retrieve - same as HOST_TABLE_METRICS in infra plugin
+const DEFAULT_HOST_METRICS = [
+  'cpuV2',
+  'memory',
+  'memoryFree',
+  'diskSpaceUsage',
+  'normalizedLoad1m',
+  'rxV2',
+  'txV2',
+] as const;
+
+export function registerDataProviders({
+  core,
+  plugins,
+  libs,
+  logger,
+}: {
+  core: CoreSetup<InfraServerPluginStartDeps>;
+  plugins: InfraServerPluginSetupDeps;
+  libs: InfraBackendLibs;
+  logger: Logger;
+}) {
+  const { observabilityAgentBuilder } = plugins;
+  if (!observabilityAgentBuilder) {
+    return;
+  }
+
+  observabilityAgentBuilder.registerDataProvider(
+    'infraHosts',
+    async ({ request, from, to, limit, kqlFilter, hostNames }) => {
+      const infraToolResources = await buildInfraToolResources({
+        core,
+        plugins,
+        libs,
+        request,
+      });
+
+      // Build query filter
+      const mustFilters: unknown[] = [];
+
+      if (kqlFilter) {
+        mustFilters.push({
+          query_string: {
+            query: kqlFilter,
+            analyze_wildcard: true,
+          },
+        });
+      }
+
+      if (hostNames && hostNames.length > 0) {
+        mustFilters.push({
+          terms: {
+            'host.name': hostNames,
+          },
+        });
+      }
+
+      const query = mustFilters.length > 0 ? { bool: { must: mustFilters } } : undefined;
+
+      const result = await getHosts({
+        from: new Date(from).getTime(),
+        to: new Date(to).getTime(),
+        metrics: [...DEFAULT_HOST_METRICS],
+        limit,
+        query,
+        alertsClient: infraToolResources.alertsClient,
+        infraMetricsClient: infraToolResources.infraMetricsClient,
+        apmDataAccessServices: infraToolResources.apmDataAccessServices,
+        schema: 'ecs',
+      });
+
+      return result;
+    }
+  );
+}
+
+async function buildInfraToolResources({
+  core,
+  plugins,
+  libs,
+  request,
+}: {
+  core: CoreSetup<InfraServerPluginStartDeps>;
+  plugins: InfraServerPluginSetupDeps;
+  libs: InfraBackendLibs;
+  request: KibanaRequest;
+}) {
+  const [coreStart] = await core.getStartServices();
+  const soClient = coreStart.savedObjects.getScopedClient(request, { includedHiddenTypes: [] });
+  const uiSettingsClient = coreStart.uiSettings.asScopedToClient(soClient);
+  const esClient = coreStart.elasticsearch.client.asScoped(request);
+
+  const coreContext = {
+    savedObjects: { client: soClient },
+    uiSettings: { client: uiSettingsClient },
+    elasticsearch: { client: esClient },
+  } as unknown as CoreRequestHandlerContext;
+
+  const infraContext = await getInfraRequestHandlerContext({ coreContext, request, plugins });
+  const context = {
+    core: Promise.resolve(coreContext),
+    infra: Promise.resolve(infraContext),
+  } as unknown as InfraPluginRequestHandlerContext;
+
+  const infraMetricsClient = await getInfraMetricsClient({ libs, context, request });
+  const alertsClient = await getInfraAlertsClient({ libs, request });
+  const apmDataAccessClient = getApmDataAccessClient({ request, libs, context });
+  const apmDataAccessServices = await apmDataAccessClient.getServices();
+
+  return { infraMetricsClient, alertsClient, apmDataAccessServices };
+}

x-pack/solutions/observability/plugins/infra/server/lib/adapters/framework/adapter_types.ts

@@ -40,6 +40,7 @@ import type {
 } from '@kbn/apm-data-access-plugin/server';
 import type { LogsDataAccessPluginStart } from '@kbn/logs-data-access-plugin/server';
 import type { ServerlessPluginStart } from '@kbn/serverless/server';
+import type { ObservabilityAgentBuilderPluginSetup } from '@kbn/observability-agent-builder-plugin/server';
 
 export interface InfraServerPluginSetupDeps {
   alerting: AlertingServerSetup;
@@ -58,6 +59,7 @@ export interface InfraServerPluginSetupDeps {
   profilingDataAccess?: ProfilingDataAccessPluginSetup;
   apmDataAccess: ApmDataAccessPluginSetup;
   serverless?: ServerlessPluginStart;
+  observabilityAgentBuilder?: ObservabilityAgentBuilderPluginSetup;
 }
 
 export interface InfraServerPluginStartDeps {

x-pack/solutions/observability/plugins/infra/server/plugin.ts

@@ -16,7 +16,6 @@ import {
   InventoryLocatorDefinition,
   MetricsExplorerLocatorDefinition,
 } from '@kbn/observability-shared-plugin/common';
-import { DEFAULT_SPACE_ID } from '@kbn/spaces-plugin/common';
 import { mapValues } from 'lodash';
 import { LOGS_FEATURE_ID, METRICS_FEATURE_ID } from '../common/constants';
 import { getMetricsFeature } from './features';
@@ -54,6 +53,8 @@ import type {
 } from './types';
 import { UsageCollector } from './usage/usage_collector';
 import { mapSourceToLogView } from './utils/map_source_to_log_view';
+import { registerDataProviders } from './agent_builder/register_data_providers';
+import { getInfraRequestHandlerContext } from './utils/get_infra_request_handler_context';
 
 export interface KbnServer extends Server {
   usage: any;
@@ -200,36 +201,15 @@ export class InfraServerPlugin
       'infra',
       async (context, request) => {
         const coreContext = await context.core;
-        const savedObjectsClient = coreContext.savedObjects.client;
-        const uiSettingsClient = coreContext.uiSettings.client;
-
-        const mlSystem = plugins.ml?.mlSystemProvider(request, savedObjectsClient);
-        const mlAnomalyDetectors = plugins.ml?.anomalyDetectorsProvider(
-          request,
-          savedObjectsClient
-        );
-        const spaceId = plugins.spaces?.spacesService.getSpaceId(request) ?? DEFAULT_SPACE_ID;
-
-        const getMetricsIndices = async () => {
-          return metricsClient.getMetricIndices({
-            savedObjectsClient,
-          });
-        };
-
-        return {
-          mlAnomalyDetectors,
-          mlSystem,
-          spaceId,
-          savedObjectsClient,
-          uiSettingsClient,
-          getMetricsIndices,
-        };
+        return getInfraRequestHandlerContext({ coreContext, request, plugins });
       }
     );
 
     // Telemetry
     UsageCollector.registerUsageCollector(plugins.usageCollection);
 
+    registerDataProviders({ core, plugins, libs: this.libs, logger: this.logger });
+
     return {
       inventoryViews,
       metricsExplorerViews,

x-pack/solutions/observability/plugins/infra/server/utils/get_infra_request_handler_context.ts

@@ -0,0 +1,44 @@
+/*
+ * 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 type { CoreRequestHandlerContext, KibanaRequest } from '@kbn/core/server';
+import { DEFAULT_SPACE_ID } from '@kbn/spaces-plugin/common';
+import type { InfraPluginRequestHandlerContext } from '../types';
+import type { InfraServerPluginSetupDeps } from '../lib/adapters/framework';
+
+export async function getInfraRequestHandlerContext({
+  coreContext,
+  request,
+  plugins,
+}: {
+  coreContext: CoreRequestHandlerContext;
+  request: KibanaRequest;
+  plugins: InfraServerPluginSetupDeps;
+}): Promise<InfraPluginRequestHandlerContext['infra']> {
+  const savedObjectsClient = coreContext.savedObjects.client;
+  const uiSettingsClient = coreContext.uiSettings.client;
+  const metricsClient = plugins.metricsDataAccess.client;
+
+  const mlSystem = plugins.ml?.mlSystemProvider(request, savedObjectsClient);
+  const mlAnomalyDetectors = plugins.ml?.anomalyDetectorsProvider(request, savedObjectsClient);
+  const spaceId = plugins.spaces?.spacesService.getSpaceId(request) ?? DEFAULT_SPACE_ID;
+
+  const getMetricsIndices = async () => {
+    return metricsClient.getMetricIndices({
+      savedObjectsClient,
+    });
+  };
+
+  return {
+    mlAnomalyDetectors,
+    mlSystem,
+    spaceId,
+    savedObjectsClient,
+    uiSettingsClient,
+    getMetricsIndices,
+  };
+}

x-pack/solutions/observability/plugins/infra/tsconfig.json

@@ -8,7 +8,7 @@
     "common/**/*",
     "public/**/*",
     "server/**/*",
-    "types/**/*",
+    "types/**/*"
   ],
   "kbn_references": [
     "@kbn/core",
@@ -125,6 +125,7 @@
     "@kbn/deeplinks-analytics",
     "@kbn/esql-composer",
     "@kbn/fleet-plugin",
+    "@kbn/observability-agent-builder-plugin"
   ],
   "exclude": ["target/**/*"]
 }