[rag evals] refactor & add ability to eval retrieval + generation in agentic eval pipeline (#664)

# What does this PR do?

- See https://github.com/meta-llama/llama-stack/pull/666 &
https://github.com/meta-llama/llama-stack/pull/668

- Refactor BaseScoringFn to be just a minimal interface, add new
RegistrableBaseScoring
- Refactor data schema check
- To separately evaluate retrieval component in RAG, we will have
scoring functions needing "context" column additionally.
- Refactor braintrust eval (more scoring fn added & tested in following
PR)

## Test Plan

```
pytest -v -s -m llm_as_judge_scoring_together_inference scoring/test_scoring.py --judge-model meta-llama/Llama-3.2-3B-Instruct
pytest -v -s -m basic_scoring_together_inference scoring/test_scoring.py
pytest -v -s -m braintrust_scoring_together_inference scoring/test_scoring.py
```

<img width="847" alt="image"
src="https://github.com/user-attachments/assets/d099cb2d-6f9c-4bdf-9d0d-f388cf758c0f"
/>

```
pytest -v -s -m meta_reference_eval_together_inference eval/test_eval.py
pytest -v -s -m meta_reference_eval_together_inference_huggingface_datasetio eval/test_eval.py
```
<img width="850" alt="image"
src="https://github.com/user-attachments/assets/dce28fc3-0493-4d34-820a-567260873cc8"
/>



## Before submitting

- [ ] This PR fixes a typo or improves the docs (you can dismiss the
other checks if that's the case).
- [ ] Ran pre-commit to handle lint / formatting issues.
- [ ] Read the [contributor
guideline](https://github.com/meta-llama/llama-stack/blob/main/CONTRIBUTING.md),
      Pull Request section?
- [ ] Updated relevant documentation.
- [ ] Wrote necessary unit or integration tests.

llama_stack/providers/inline/eval/meta_reference/eval.py

@@ -3,23 +3,24 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
-from enum import Enum
 from typing import Any, Dict, List, Optional
 
 from tqdm import tqdm
 
-from llama_stack.apis.agents import Agents
-from llama_stack.apis.common.type_system import (
-    ChatCompletionInputType,
-    CompletionInputType,
-    StringType,
-)
+from llama_stack.apis.agents import Agents, StepType
 from llama_stack.apis.datasetio import DatasetIO
 from llama_stack.apis.datasets import Datasets
 from llama_stack.apis.eval_tasks import EvalTask
 from llama_stack.apis.inference import Inference, UserMessage
 from llama_stack.apis.scoring import Scoring
+from llama_stack.distribution.datatypes import Api
 from llama_stack.providers.datatypes import EvalTasksProtocolPrivate
+
+from llama_stack.providers.utils.common.data_schema_validator import (
+    ColumnName,
+    DataSchemaValidatorMixin,
+    get_valid_schemas,
+)
 from llama_stack.providers.utils.kvstore import kvstore_impl
 
 from .....apis.common.job_types import Job
@@ -30,15 +31,7 @@ from .config import MetaReferenceEvalConfig
 EVAL_TASKS_PREFIX = "eval_tasks:"
 
 
-class ColumnName(Enum):
-    input_query = "input_query"
-    expected_answer = "expected_answer"
-    chat_completion_input = "chat_completion_input"
-    completion_input = "completion_input"
-    generated_answer = "generated_answer"
-
-
-class MetaReferenceEvalImpl(Eval, EvalTasksProtocolPrivate):
+class MetaReferenceEvalImpl(Eval, EvalTasksProtocolPrivate, DataSchemaValidatorMixin):
     def __init__(
         self,
         config: MetaReferenceEvalConfig,
@@ -82,29 +75,6 @@ class MetaReferenceEvalImpl(Eval, EvalTasksProtocolPrivate):
         )
         self.eval_tasks[task_def.identifier] = task_def
 
-    async def validate_eval_input_dataset_schema(self, dataset_id: str) -> None:
-        dataset_def = await self.datasets_api.get_dataset(dataset_id=dataset_id)
-        if not dataset_def.dataset_schema or len(dataset_def.dataset_schema) == 0:
-            raise ValueError(f"Dataset {dataset_id} does not have a schema defined.")
-
-        expected_schemas = [
-            {
-                ColumnName.input_query.value: StringType(),
-                ColumnName.expected_answer.value: StringType(),
-                ColumnName.chat_completion_input.value: ChatCompletionInputType(),
-            },
-            {
-                ColumnName.input_query.value: StringType(),
-                ColumnName.expected_answer.value: StringType(),
-                ColumnName.completion_input.value: CompletionInputType(),
-            },
-        ]
-
-        if dataset_def.dataset_schema not in expected_schemas:
-            raise ValueError(
-                f"Dataset {dataset_id} does not have a correct input schema in {expected_schemas}"
-            )
-
     async def run_eval(
         self,
         task_id: str,
@@ -114,8 +84,10 @@ class MetaReferenceEvalImpl(Eval, EvalTasksProtocolPrivate):
         dataset_id = task_def.dataset_id
         candidate = task_config.eval_candidate
         scoring_functions = task_def.scoring_functions
-
-        await self.validate_eval_input_dataset_schema(dataset_id=dataset_id)
+        dataset_def = await self.datasets_api.get_dataset(dataset_id=dataset_id)
+        self.validate_dataset_schema(
+            dataset_def.dataset_schema, get_valid_schemas(Api.eval.value)
+        )
         all_rows = await self.datasetio_api.get_rows_paginated(
             dataset_id=dataset_id,
             rows_in_page=(
@@ -167,11 +139,21 @@ class MetaReferenceEvalImpl(Eval, EvalTasksProtocolPrivate):
                 )
             ]
             final_event = turn_response[-1].event.payload
-            generations.append(
-                {
-                    ColumnName.generated_answer.value: final_event.turn.output_message.content
-                }
+
+            # check if there's a memory retrieval step and extract the context
+            memory_rag_context = None
+            for step in final_event.turn.steps:
+                if step.step_type == StepType.memory_retrieval.value:
+                    memory_rag_context = " ".join(x.text for x in step.inserted_context)
+
+            agent_generation = {}
+            agent_generation[ColumnName.generated_answer.value] = (
+                final_event.turn.output_message.content
             )
+            if memory_rag_context:
+                agent_generation[ColumnName.context.value] = memory_rag_context
+
+            generations.append(agent_generation)
 
         return generations