[HF Data Loader] Load onechat datasets with externally defined mappings. (elastic#230252)

## Summary
Closes: elastic#227993

This PR introduces OneChat datasets to the HuggingFace dataset loader
CLI tool.

### Key Changes

* OneChat Data Loading Support
- Added support for loading OneChat datasets (CSV-based) from
the `elastic/OneChatAgent` repository.
- OneChat datasets can now be specified using
the `onechat/<directory>/<dataset>` syntax.
* Refactored File Download and Parsing
    -   Extracted HuggingFace file download logic for reuse.
- Separated CSV and JSON parsing into dedicated functions for clarity
and maintainability. Previously only JSON download was supported. This
uses `papaparse` library for csv parsing which was very handy for
parsing multi-line csv files.
* Configurable Kibana URL
- The loader now accepts a `--kibana-url` flag to override
auto-discovery of the Kibana instance. This is needed locally for the
evaluation framework that uses scout server running on port `5620`.
* Documentation Updates:
- Expanded the README with clear instructions for loading both
HuggingFace and OneChat datasets.
- Added OneChat eval specific instructions in the corresponding README
* Fixed a bug with the existing dataset loader that was previously only
storing one document in ES.

### How to Test

Follow the instructions from the updated README:

1.  Prerequisites

    -   Ensure you have a running Kibana and Elasticsearch instance.
    -   Obtain a HuggingFace access token. 
- (For OneChat datasets) Ensure your HuggingFace account is part of the
Elastic organization.
2.  Run the Loader
```bash
HUGGING_FACE_ACCESS_TOKEN=<YOUR_TOKEN> \
node --require ./src/setup_node_env/index.js \
  x-pack/platform/packages/shared/kbn-ai-tools-cli/scripts/hf_dataset_loader.ts \
  --datasets onechat/knowledge-base/users,huffpost \
  --limit 100 \
  --clear \
  --kibana-url http://<username>:<password>@localhost:5601
```

3.  Verify

    -   Check the logs for successful ingestion and embedding.
    -   Confirm the new indices in your local ES instance.

### Checklist

Check the PR satisfies following conditions. 

Reviewers should verify this PR satisfies this list as well.

- [x]
[Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html)
was added for features that require explanation or tutorials
- [x] The PR description includes the appropriate Release Notes section,
and the correct `release_note:*` label is applied per the
[guidelines](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process)


---------

Co-authored-by: Quynh Nguyen (Quinn) <[email protected]>

Author: SrdjanLL

x-pack/platform/packages/shared/kbn-ai-tools-cli/index.ts

@@ -7,4 +7,4 @@
 
 export { loadHuggingFaceDatasets } from './src/hf_dataset_loader/load_hugging_face_datasets';
 export type { HuggingFaceDatasetSpec } from './src/hf_dataset_loader/types';
-export { ALL_HUGGING_FACE_DATASETS } from './src/hf_dataset_loader/config';
+export { PREDEFINED_HUGGING_FACE_DATASETS } from './src/hf_dataset_loader/datasets/config';

x-pack/platform/packages/shared/kbn-ai-tools-cli/scripts/hf_dataset_loader.ts

@@ -7,9 +7,14 @@
 
 import { run } from '@kbn/dev-cli-runner';
 import { createKibanaClient, toolingLogToLogger } from '@kbn/kibana-api-cli';
-import { castArray, keyBy } from 'lodash';
+import { castArray } from 'lodash';
 import { loadHuggingFaceDatasets } from '../src/hf_dataset_loader/load_hugging_face_datasets';
-import { ALL_HUGGING_FACE_DATASETS } from '../src/hf_dataset_loader/config';
+import {
+  PREDEFINED_HUGGING_FACE_DATASETS,
+  getDatasetSpecs,
+} from '../src/hf_dataset_loader/datasets/config';
+import { listAllOneChatDatasets } from '../src/hf_dataset_loader/datasets/onechat';
+import { HuggingFaceDatasetSpec } from '../src/hf_dataset_loader/types';
 
 interface Flags {
   // the number of rows per dataset to load into ES
@@ -18,11 +23,39 @@ interface Flags {
   datasets?: string | string[];
   // whether all specified dataset's indices should be cleared before loading
   clear?: boolean;
+  // the kibana URL to connect to
+  'kibana-url'?: string;
+}
+
+async function showAvailableDatasets(accessToken: string, logger: any) {
+  let output = 'No datasets specified. Here are the available datasets:\n\n';
+
+  output += 'Pre-defined HuggingFace datasets:\n';
+  output += PREDEFINED_HUGGING_FACE_DATASETS.map((d, index) => `  ${index + 1}. ${d.name}`).join(
+    '\n'
+  );
+  output += '\n\n';
+
+  const oneChatDatasets = await listAllOneChatDatasets(accessToken, logger);
+  output += 'OneChat datasets:\n';
+  if (oneChatDatasets.length > 0) {
+    output += oneChatDatasets.map((dataset, index) => `  ${index + 1}. ${dataset}`).join('\n');
+  } else {
+    output +=
+      '  (none available - you may need to join Elastic oranization on HuggingFace to access OneChat datasets)';
+  }
+
+  output += '\n\n';
+  output += 'Usage: Use --datasets to specify which datasets to load\n';
+  output += 'Example: --datasets onechat/knowledge-base/wix_knowledge_base';
+
+  logger.info(output);
 }
 
 run(
   async ({ log, flags }) => {
     const signal = new AbortController().signal;
+    const logger = toolingLogToLogger({ flags, log });
 
     const accessToken = process.env.HUGGING_FACE_ACCESS_TOKEN;
 
@@ -32,38 +65,40 @@ run(
       );
     }
 
+    // destructure and normalize CLI flags
+    const { limit, datasets, clear } = flags as Flags;
+    const kibanaUrl = typeof flags['kibana-url'] === 'string' ? flags['kibana-url'] : undefined;
+
     const kibanaClient = await createKibanaClient({
       log,
       signal,
+      baseUrl: kibanaUrl,
     });
 
-    // destructure and normalize CLI flags
-    const { limit, datasets, clear } = flags as Flags;
-
     const datasetNames = !!datasets
       ? castArray(datasets)
           .flatMap((set) => set.split(','))
           .map((set) => set.trim())
           .filter(Boolean)
       : undefined;
 
-    const specsByName = keyBy(ALL_HUGGING_FACE_DATASETS, (val) => val.name);
+    let specs: HuggingFaceDatasetSpec[];
 
-    const specs =
-      datasetNames?.map((name) => {
-        if (!specsByName[name]) {
-          throw new Error(`Dataset spec for ${name} not found`);
-        }
-        return specsByName[name];
-      }) ?? ALL_HUGGING_FACE_DATASETS;
+    if (datasetNames) {
+      specs = await getDatasetSpecs(accessToken, logger, datasetNames);
+    } else {
+      // Show available datasets and exit
+      await showAvailableDatasets(accessToken, logger);
+      return;
+    }
 
     if (!specs.length) {
       throw new Error(`No datasets to load`);
     }
 
     await loadHuggingFaceDatasets({
       esClient: kibanaClient.es,
-      logger: toolingLogToLogger({ flags, log }),
+      logger,
       clear: Boolean(clear),
       limit: !!limit ? Number(limit) : undefined,
       datasets: specs,
@@ -73,14 +108,17 @@ run(
   {
     description: `Loads HuggingFace datasets into an Elasticsearch cluster`,
     flags: {
-      string: ['limit', 'datasets'],
+      string: ['limit', 'datasets', 'kibana-url'],
       boolean: ['clear'],
       help: `
         Usage: node --require ./src/setup_node_env/index.js x-pack/platform/packages/shared/kbn-ai-tools-cli/scripts/hf_dataset_loader.ts [options]
 
-        --datasets          Comma-separated list of HuggingFace dataset names to load
+        --datasets          Comma-separated list of HuggingFace dataset names to load.
+                           For OneChat datasets, use format: onechat/<directory>/<dataset_name>
+                           Example: --datasets onechat/knowledge-base/wix_knowledge_base
         --limit             Number of rows per dataset to load into Elasticsearch
         --clear             Clear the existing indices for the specified datasets before loading
+        --kibana-url        Kibana URL to connect to (bypasses auto-discovery when provided)
       `,
       default: {
         clear: false,

x-pack/platform/packages/shared/kbn-ai-tools-cli/src/hf_dataset_loader/README.md

@@ -18,23 +18,63 @@ node --require ./src/setup_node_env/index.js \
   x-pack/platform/packages/shared/kbn-ai-tools-cli/scripts/hf_dataset_loader.ts \
   --datasets beir-trec-covid,beir-msmarco \
   --limit 1000 \
-  --clear
+  --clear \
+  --kibana-url http://<username>:<password>@localhost:5601
 ```
 
 ### CLI flags
 
-| Flag         | Type      | Description                                                                                           |
-| ------------ | --------- | ----------------------------------------------------------------------------------------------------- |
-| `--datasets` | `string`  | Comma-separated list of dataset **names** to load. Omit the flag to load **all** predefined datasets. |
-| `--limit`    | `number`  | Max docs per dataset (handy while testing). Defaults to 1k.                                           |
-| `--clear`    | `boolean` | Delete the target index **before** indexing. Defaults to `false`.                                     |
+| Flag           | Type      | Description                                                                                           |
+| -------------- | --------- | ----------------------------------------------------------------------------------------------------- |
+| `--datasets`   | `string`  | Comma-separated list of dataset **names** to load. Omit the flag to load **all** predefined datasets. |
+| `--limit`      | `number`  | Max docs per dataset (handy while testing). Defaults to 1k.                                           |
+| `--clear`      | `boolean` | Delete the target index **before** indexing. Defaults to `false`.                                     |
+| `--kibana-url` | `string`  | Kibana URL to connect to (bypasses auto-discovery when provided).                                     |
 
 ## Built-in dataset specs
 
 The script ships with ready-made specifications located in `config.ts`.
 
 Feel free to extend or tweak these specs in `src/hf_dataset_loader/config.ts`.
 
+## OneChat datasets
+
+The loader also supports **OneChat datasets** from the `elastic/OneChatAgent` repository. These are CSV-based datasets with predefined mappings stored in `index-mappings.jsonl` files.
+
+**Note**: To access OneChat datasets, you need to be a member of the Elastic organization on HuggingFace. Sign up with your `@elastic.co` email address to request access (automated process).
+
+### OneChat syntax
+
+Use the format `onechat/<directory>/<dataset>` to load OneChat datasets:
+
+```bash
+# Load a single OneChat dataset
+--datasets onechat/knowledge-base/wix_knowledge_base
+
+# Mix OneChat and regular datasets
+--datasets onechat/knowledge-base/wix_knowledge_base,beir-msmarco
+
+# Load multiple OneChat datasets
+--datasets onechat/knowledge-base/wix_knowledge_base,onechat/users/user_profiles
+```
+
+### How it works
+
+1. The loader fetches `<directory>/index-mappings.jsonl` from the OneChat repository
+2. Downloads the corresponding CSV file from `<directory>/datasets/`
+3. Creates Elasticsearch indices with the predefined mappings from `index-mappings.jsonl` file.
+4. Loads the CSV data into the index.
+
+### Available datasets
+
+Run the loader without `--datasets` to see all available OneChat and regular HuggingFace datasets.
+
+### Naming convention
+
+- Repository file: `knowledge-base/datasets/wix_knowledge_base.csv`
+- Loader dataset name: `onechat/knowledge-base/wix_knowledge_base`
+- Elasticsearch index: `wix_knowledge_base`
+
 ## Disabling local cache
 
 Set the environment variable `DISABLE_KBN_CLI_CACHE=1` to force fresh downloads instead of using the on-disk cache.

…ools-cli/src/hf_dataset_loader/config.ts …src/hf_dataset_loader/datasets/config.tsx-pack/platform/packages/shared/kbn-ai-tools-cli/src/hf_dataset_loader/config.ts renamed to x-pack/platform/packages/shared/kbn-ai-tools-cli/src/hf_dataset_loader/datasets/config.ts x-pack/platform/packages/shared/kbn-ai-tools-cli/src/hf_dataset_loader/config.ts renamed to x-pack/platform/packages/shared/kbn-ai-tools-cli/src/hf_dataset_loader/datasets/config.ts

@@ -5,7 +5,9 @@
  * 2.0.
  */
 
-import type { HuggingFaceDatasetSpec } from './types';
+import { Logger } from '@kbn/core/server';
+import type { HuggingFaceDatasetSpec } from '../types';
+import { createOneChatDatasetSpec, isOneChatDataset } from './onechat';
 
 const BEIR_NAMES = [
   'trec-covid',
@@ -81,7 +83,34 @@ const EXTRA_DATASETS: HuggingFaceDatasetSpec[] = [
   },
 ];
 
-export const ALL_HUGGING_FACE_DATASETS: HuggingFaceDatasetSpec[] = [
+export const PREDEFINED_HUGGING_FACE_DATASETS: HuggingFaceDatasetSpec[] = [
   ...BEIR_DATASETS,
   ...EXTRA_DATASETS,
 ];
+
+/**
+ * Get dataset specifications, including dynamically generated OneChat datasets
+ */
+export async function getDatasetSpecs(
+  accessToken: string,
+  logger: Logger,
+  datasetNames: string[]
+): Promise<HuggingFaceDatasetSpec[]> {
+  const specs: HuggingFaceDatasetSpec[] = [];
+  for (const name of datasetNames) {
+    if (isOneChatDataset(name)) {
+      const spec = await createOneChatDatasetSpec(name, accessToken, logger);
+      specs.push(spec);
+    } else {
+      // Look for static datasets
+      const staticSpec = PREDEFINED_HUGGING_FACE_DATASETS.find((spec) => spec.name === name);
+      if (staticSpec) {
+        specs.push(staticSpec);
+      } else {
+        throw new Error(`Dataset '${name}' not found`);
+      }
+    }
+  }
+
+  return specs;
+}