Improves data ingestion by pulling the data either suing GCP or by using hugging face

Author: mafaldasv-dev

examples/opik_vs_elastic/Instructions.md

@@ -1,35 +1,26 @@
 # Instructions
 
-Setup and run the Opik vs elastic-evals-sdk-python PoC.
+This PoC demonstrates how Agent Builder can use the `kbn/evals` Python SDK for dataset
+management, experiment tracing, and evaluator score ingestion—capabilities previously
+handled through Opik—and how to build custom evaluators.
 
-## 1. Python environment
+## 1. Dependencies
 
-From the SDK repo root:
-
-```bash
-cd elastic-evals-sdk-python
-uv sync --extra dev --extra runner --extra poc
-source .venv/bin/activate
-```
-
-## 2. Dependencies
-
-```bash
-uv add --optional poc datasets pandas python-dotenv ipykernel elasticsearch opik loguru
-uv sync --extra poc
-```
-
-Orca must be importable. Clone the `orca` repo as a sibling of `elastic-evals-sdk-python/` so the layout is:
+`run2.py` uses external Orca evaluators. Clone the `orca` repo as a sibling of
+`elastic-evals-sdk-python/` so the layout is:
 
 ```
 ── elastic-evals-sdk-python/
 ── orca/
 ```
 
-Then install it editable and register it under the `poc` extra:
+## 2. Python environment
+
+From the SDK repo root:
 
 ```bash
-uv add --optional poc --editable ../orca
+uv sync --group dev --extra runner --extra poc
+source .venv/bin/activate
 ```
 
 Register the venv as a Jupyter kernel. Only needed if your notebook/IDE doesn't pick up `.venv` automatically:
@@ -40,21 +31,50 @@ uv run python -m ipykernel install --user --name elastic-evals-poc
 
 ## 3. Secrets
 
-Retrieve the API keys (Opik, OpenRouter, HuggingFace) from Vault, then fill them into `.env`:
+Create `.env` next to `.env.example`:
+
+```bash
+cp examples/opik_vs_elastic/.env.example examples/opik_vs_elastic/.env
+```
+
+Set the local URLs, `ELASTICSEARCH_API_KEY`, `KIBANA_API_KEY`, `CONNECTOR_ID`,
+and `EVALUATION_CONNECTOR_ID`. The Opik variables are used when `run2.py` runs
+the tracked external Orca evaluators. Retrieve internal credentials from Vault
+when needed:
 
 ```bash
 VAULT_ADDR=https://secrets.elastic.co:8200 vault login --method oidc
 ```
 
-## 4. GCP access
+The public Hugging Face dataset does not require an API key.
+
+## 4. Data source and sample size
+
+Set these values near the top of the script before running it:
+
+```python
+USE_ENTIRE_DATASET = False
+DATASET_SAMPLE_SIZE = 10
+USE_GCP = False
+```
+
+With `USE_GCP = False`, the scripts load the public `Wix/WixQA` dataset and
+knowledge base from Hugging Face. Set `USE_GCP = True` to use the internal GCS
+files. `DATASET_SAMPLE_SIZE` is ignored when `USE_ENTIRE_DATASET` is `True`.
+The entire knowledge base is always indexed.
+
+The available examples and their order may differ between Hugging Face and GCS.
+`run.py` defaults to 10 examples and `run2.py` defaults to 3.
+
+## 5. GCP access
 
-Only needed if you import data from a GCP bucket (e.g. `gs://agent-builder-data-science-datasets/...`). Authenticate with your `@elastic.co` account:
+Only needed when `USE_GCP = True`. Authenticate with your `@elastic.co` account:
 
 ```bash
 gcloud auth application-default login
 ```
 
-## 5. Local stack
+## 6. Local stack
 
 Use a separate terminal for each service and leave it running.
 
@@ -128,25 +148,40 @@ curl --silent --show-error \
   http://localhost:5601/dev/api/status
 ```
 
-## 6. Elasticsearch API key
+## 7. Elasticsearch API key
 
-`run.py` authenticates to Elasticsearch with `ELASTICSEARCH_API_KEY` (and reuses the
-same key for Kibana, which validates Elasticsearch API keys). Create one against your
-cluster and paste the `encoded` field into `ELASTICSEARCH_API_KEY` in `.env`:
+Both scripts authenticate to Elasticsearch and Kibana. Create an API key against
+the local cluster and paste the `encoded` field into both
+`ELASTICSEARCH_API_KEY` and `KIBANA_API_KEY` in `.env`:
 
 ```bash
 curl -u elastic:changeme -XPOST http://localhost:9200/_security/api_key \
   -H 'Content-Type: application/json' -d '{"name":"evals-poc"}'
 ```
 
-This step is **optional** — skip it if you already have a valid `ELASTICSEARCH_API_KEY`
-for the cluster in `ES_URL`. It's required when the key is missing or invalid, e.g.
-after starting a fresh local Elasticsearch (`yarn es snapshot`): API keys are
-cluster-specific, so a key from a previous cluster returns `401`.
+This step is optional if both variables already contain a valid key for the
+current cluster. API keys are cluster-specific, so a key from a previous local
+Elasticsearch snapshot returns `401`.
 
-## 7. Run the PoC
+## 8. Run the PoC
+
+### `run.py`: managed workflow
+
+Demonstrates the higher-level workflow. It uses
+`ElasticEvalsClient.run_experiment()` to execute the selected WixQA examples,
+run SDK-side and custom evaluators, and ingest their scores.
 
 ```bash
 cd /Users/mafaldasavelho/Documents/work-repos/kibana-fork/evals-python-sdk/elastic-evals-sdk-python
 uv run --extra poc python examples/opik_vs_elastic/run.py
 ```
+
+### `run2.py`: granular workflow
+
+Demonstrates the lower-level workflow without `run_experiment()`. It directly
+coordinates the Dataset, Evaluators, and Score Ingestion APIs, runs the custom
+Document Recall evaluator, and attaches external Orca scores.
+
+```bash
+uv run --extra poc python examples/opik_vs_elastic/run2.py
+```

examples/opik_vs_elastic/helpers/data.py

@@ -0,0 +1,114 @@
+# 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.
+
+"""Load and normalize WixQA data from GCS or Hugging Face."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping
+from typing import Any
+
+import pandas as pd
+from datasets import load_dataset
+
+from examples.opik_vs_elastic.helpers.helpers import (
+    GROUND_TRUTH_COLUMN,
+    WIX_KNOWLEDGE_BASE_PATH,
+    WIX_QA_DATASET_PATH,
+    _parse_relevant_doc_ids,  # noqa: PLC2701
+)
+
+HUGGING_FACE_DATASET = "Wix/WixQA"
+HUGGING_FACE_QA_CONFIG = "wixqa_expertwritten"
+HUGGING_FACE_KB_CONFIG = "wix_kb_corpus"
+
+
+def _require_columns(dataframe: pd.DataFrame, columns: set[str], *, source: str) -> None:
+    missing = sorted(columns - set(dataframe.columns))
+    if missing:
+        raise ValueError(f"{source} is missing required columns: {', '.join(missing)}")
+
+
+def _to_document_ids(value: Any) -> list[str]:
+    if isinstance(value, str) or not isinstance(value, Iterable) or isinstance(value, Mapping):
+        return []
+    return list(dict.fromkeys(str(item) for item in value if item is not None))
+
+
+def _load_hugging_face_config(config: str) -> pd.DataFrame:
+    dataset = load_dataset(HUGGING_FACE_DATASET, config, split="train")
+    return dataset.to_pandas()
+
+
+def _normalize_gcs_qa(dataframe: pd.DataFrame) -> pd.DataFrame:
+    required = {
+        "meta_query_id",
+        "input_question",
+        "output_expected",
+        GROUND_TRUTH_COLUMN,
+    }
+    _require_columns(dataframe, required, source="GCS WixQA dataset")
+
+    normalized = dataframe.copy()
+    normalized["relevant_doc_ids"] = normalized[GROUND_TRUTH_COLUMN].apply(_parse_relevant_doc_ids)
+    return normalized
+
+
+def _normalize_hugging_face_qa(dataframe: pd.DataFrame) -> pd.DataFrame:
+    dataframe = dataframe.reset_index(drop=True)
+    _require_columns(
+        dataframe,
+        {"question", "answer", "article_ids"},
+        source="Hugging Face WixQA dataset",
+    )
+
+    document_ids = dataframe["article_ids"].apply(_to_document_ids)
+    return pd.DataFrame(
+        {
+            "meta_query_id": [f"wixqa_expertwritten_{index + 1}" for index in range(len(dataframe))],
+            "input_question": dataframe["question"].astype(str),
+            "output_expected": dataframe["answer"].astype(str),
+            GROUND_TRUTH_COLUMN: document_ids.apply(lambda ids: {document_id: True for document_id in ids}),
+            "relevant_doc_ids": document_ids,
+        }
+    )
+
+
+def _normalize_knowledge_base(dataframe: pd.DataFrame, *, source: str) -> pd.DataFrame:
+    _require_columns(dataframe, {"id", "contents"}, source=source)
+    normalized = dataframe.copy()
+    normalized["id"] = normalized["id"].astype(str)
+    return normalized
+
+
+def load_wix_data(*, use_gcp: bool) -> tuple[pd.DataFrame, pd.DataFrame]:
+    """Load WixQA examples and their knowledge-base corpus."""
+    if use_gcp:
+        qa = _normalize_gcs_qa(pd.read_csv(WIX_QA_DATASET_PATH))
+        knowledge_base = _normalize_knowledge_base(
+            pd.read_csv(WIX_KNOWLEDGE_BASE_PATH),
+            source="GCS Wix knowledge base",
+        )
+    else:
+        qa = _normalize_hugging_face_qa(_load_hugging_face_config(HUGGING_FACE_QA_CONFIG))
+        knowledge_base = _normalize_knowledge_base(
+            _load_hugging_face_config(HUGGING_FACE_KB_CONFIG),
+            source="Hugging Face Wix knowledge base",
+        )
+
+    return qa, knowledge_base
+
+
+def select_qa_examples(
+    dataframe: pd.DataFrame,
+    *,
+    use_entire_dataset: bool,
+    sample_size: int,
+) -> pd.DataFrame:
+    """Return the full QA dataset or its first configured examples."""
+    if use_entire_dataset:
+        return dataframe.reset_index(drop=True).copy()
+    if sample_size < 1:
+        raise ValueError("DATASET_SAMPLE_SIZE must be at least 1")
+    return dataframe.head(sample_size).reset_index(drop=True).copy()

examples/opik_vs_elastic/run.py

@@ -9,7 +9,6 @@
 from typing import Any
 
 import httpx
-import pandas as pd
 from dotenv import load_dotenv
 
 from elastic_evals.agent_builder import (
@@ -46,6 +45,7 @@
     EvaluatorParams,
     Example,
 )
+from examples.opik_vs_elastic.helpers.data import load_wix_data, select_qa_examples
 from examples.opik_vs_elastic.helpers.helpers import (
     AGENT_ID,
     AGENT_INSTRUCTIONS,
@@ -55,14 +55,18 @@
     INDEX_NAME,
     SEARCH_TOOL_DESCRIPTION,
     SEARCH_TOOL_ID,
-    WIX_KNOWLEDGE_BASE_PATH,
-    WIX_QA_DATASET_PATH,
     _extract_retrieved_doc_ids,  # noqa: PLC2701
-    _parse_relevant_doc_ids,  # noqa: PLC2701
     _to_string_list,  # noqa: PLC2701
 )
 from examples.opik_vs_elastic.helpers.indexing import get_elasticsearch_client
 
+USE_ENTIRE_DATASET = False
+DATASET_SAMPLE_SIZE = 10
+USE_GCP = False
+
+DATASET_NAME = "wix_qa_managed_workflow"
+DATASET_DESCRIPTION = "WixQA golden Q&A pairs for the managed evaluation workflow."
+
 WIX_RESPONSE_CRITERIA = [
     "The response directly addresses the user's Wix support question.",
     "The response provides clear and actionable guidance.",
@@ -152,16 +156,15 @@ async def main() -> None:
     load_dotenv(ENV_PATH)
 
     # (1) Prepare data:
-    print("\nLoading QA pairs from GCS (wix_qa dataset)...")
-    qa_wix = pd.read_csv(WIX_QA_DATASET_PATH)
-    qa_wix["relevant_doc_ids"] = qa_wix[GROUND_TRUTH_COLUMN].apply(_parse_relevant_doc_ids)
-    print(f"Dataset shape: {qa_wix.shape}")  # output: (52, 5)
-    # print(qa_wix.head())
-
-    print("\nLoading knowledge base from GCS (wix_knowledge_base dataset)...")
-    kb_wix = pd.read_csv(WIX_KNOWLEDGE_BASE_PATH)
-    print(f"Dataset shape: {kb_wix.shape}")  # output: (6222, 4)
-    # print(kb_wix.head())
+    source = "GCS" if USE_GCP else "Hugging Face"
+    print(f"\nLoading WixQA data from {source}...")
+    all_qa_wix, kb_wix = load_wix_data(use_gcp=USE_GCP)
+    qa_wix = select_qa_examples(
+        all_qa_wix,
+        use_entire_dataset=USE_ENTIRE_DATASET,
+        sample_size=DATASET_SAMPLE_SIZE,
+    )
+    print(f"Selected {len(qa_wix)} of {len(all_qa_wix)} QA examples; knowledge base contains {len(kb_wix)} documents")
 
     # (2) Indexing the knowledge base into elasticsearch (if not indexed already):
     print("Indexing knowledge base into Elasticsearch...")
@@ -200,8 +203,8 @@ async def main() -> None:
     ]
 
     examples_dataset = EvaluationDataset(
-        name="wix_qa_smoke",
-        description="Small sample of WixQA golden Q&A pairs for smoke-testing the pipeline.",
+        name=DATASET_NAME,
+        description=DATASET_DESCRIPTION,
         examples=[
             Example(
                 input={"question": row["input_question"]},
@@ -213,8 +216,7 @@ async def main() -> None:
                 },
             )
             for _, row in qa_wix.iterrows()
-            if _ < 10
-        ],  # take only the first 10 examples for a quick iteration
+        ],
     )
 
     # [IMPORTANT] NOTE: if the dataset doesn't exist, then it creates a new one with the given name. If the dataset already exists,
@@ -223,8 +225,8 @@ async def main() -> None:
     # UpsertDatasetResponse(dataset_id='0b5ee7b6-9f4a-5c66-b196-6b8cc5154eec', added=0, removed=47, unchanged=5)
 
     upsert_dataset_response = await elastic_evals_client._datasets_client.upsert(
-        name="wix_qa_smoke",
-        description="Small sample of WixQA golden Q&A pairs for smoke-testing the pipeline.",
+        name=DATASET_NAME,
+        description=DATASET_DESCRIPTION,
         examples=examples,
     )
     print(