enabled metadata support in ragulate (#580)

* enabled metadata support in ragulate

* fmt

* lint

Author: epinzur

libs/ragulate/ragstack_ragulate/datasets/__init__.py

@@ -1,4 +1,4 @@
-from .base_dataset import BaseDataset
+from .base_dataset import BaseDataset, QueryItem
 from .crag_dataset import CragDataset
 from .llama_dataset import LlamaDataset
 from .utils import find_dataset, get_dataset
@@ -7,6 +7,7 @@
     "BaseDataset",
     "CragDataset",
     "LlamaDataset",
+    "QueryItem",
     "find_dataset",
     "get_dataset",
 ]

libs/ragulate/ragstack_ragulate/datasets/base_dataset.py

@@ -3,24 +3,37 @@
 from abc import ABC, abstractmethod
 from os import makedirs, path
 from pathlib import Path
-from typing import Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional
 
 import aiofiles
 import aiohttp
 from tqdm.asyncio import tqdm
 
 
+class QueryItem():
+    query: str
+    metadata: Dict[str, Any]
+
+    def __init__(self, query:str, metadata: Dict[str, Any]):
+        self.query = query
+        self.metadata = metadata
+
+
 class BaseDataset(ABC):
 
     root_storage_path: str
     name: str
     _subsets: List[str] = []
+    _query_items: List[QueryItem]
+    _golden_set: List[Dict[str, str]]
 
     def __init__(
         self, dataset_name: str, root_storage_path: str = "datasets"
     ):
         self.name = dataset_name
         self.root_storage_path = root_storage_path
+        self._query_items = []
+        self._golden_set = []
 
     def storage_path(self) -> str:
         """returns the path where dataset files should be stored"""
@@ -55,8 +68,20 @@ def get_source_file_paths(self) -> List[str]:
         """gets a list of source file paths for for a dataset"""
 
     @abstractmethod
-    def get_queries_and_golden_set(self) -> Tuple[List[str], List[Dict[str, str]]]:
-        """gets a list of queries and golden_truth answers for a dataset"""
+    def _load_query_items_and_golden_set(self) -> None:
+        """loads query_items and golden_set"""
+
+    def get_query_items(self) -> List[QueryItem]:
+        """gets a list of query items for a dataset"""
+        if len(self._query_items) == 0:
+            self._load_query_items_and_golden_set()
+        return self._query_items
+
+    def get_golden_set(self) -> List[Dict[str, str]]:
+        """gets the set of ground_truth answers for a dataset"""
+        if len(self._golden_set) == 0:
+            self._load_query_items_and_golden_set()
+        return self._golden_set
 
     async def _download_file(
         self, session: aiohttp.ClientSession, url: str, temp_file_path: str

libs/ragulate/ragstack_ragulate/datasets/crag_dataset.py

@@ -1,9 +1,9 @@
 import asyncio
 import json
 from os import path
-from typing import Dict, List, Optional, Tuple
+from typing import List, Optional
 
-from .base_dataset import BaseDataset
+from .base_dataset import BaseDataset, QueryItem
 
 
 class CragDataset(BaseDataset):
@@ -52,10 +52,8 @@ def download_dataset(self) -> None:
     def get_source_file_paths(self) -> List[str]:
         raise NotImplementedError("Crag source files are not yet supported")
 
-    def get_queries_and_golden_set(self) -> Tuple[List[str], List[Dict[str, str]]]:
-        """gets a list of queries and golden_truth answers for a dataset"""
-        queries: List[str] = []
-        golden_set: List[Dict[str, str]] = []
+    def _load_query_items_and_golden_set(self) -> None:
+        """loads query_items and golden_set"""
 
         for subset in self.subsets:
             if subset not in self._subset_kinds:
@@ -74,10 +72,11 @@ def get_queries_and_golden_set(self) -> Tuple[List[str], List[Dict[str, str]]]:
 
                 query = data.get("query")
                 answer = data.get("answer")
+                del data["query"]
+                del data["answer"]
                 if query is not None and answer is not None:
-                    queries.append(query)
-                    golden_set.append({"query": query, "response": answer})
-
-        print(f"found {len(queries)} for subsets: {self.subsets}")
-
-        return queries, golden_set
+                    self._query_items.append(QueryItem(
+                        query=query,
+                        metadata=data,
+                    ))
+                    self._golden_set.append({"query": query, "response": answer})

libs/ragulate/ragstack_ragulate/datasets/llama_dataset.py

@@ -10,7 +10,7 @@
 )
 
 from ..logging_config import logger
-from .base_dataset import BaseDataset
+from .base_dataset import BaseDataset, QueryItem
 
 
 class LlamaDataset(BaseDataset):
@@ -69,17 +69,14 @@ def get_source_file_paths(self) -> List[str]:
         source_path = path.join(self._get_dataset_path(), "source_files")
         return self.list_files_at_path(path=source_path)
 
-    def get_queries_and_golden_set(self) -> Tuple[List[str], List[Dict[str, str]]]:
-        """gets a list of queries and golden_truth answers for a dataset"""
+    def _load_query_items_and_golden_set(self) -> None:
+        """loads query_items and golden_set"""
         json_path = path.join(self._get_dataset_path(), "rag_dataset.json")
         with open(json_path, "r") as f:
             examples = json.load(f)["examples"]
-            queries = [e["query"] for e in examples]
-            golden_set = [
-                {
-                    "query": e["query"],
-                    "response": e["reference_answer"],
-                }
-                for e in examples
-            ]
-            return queries, golden_set
+            for example in examples:
+                self._query_items.append(QueryItem(query=example["query"], metadata={}))
+                self._golden_set.append({
+                    "query": example["query"],
+                    "response": example["reference_answer"],
+                })

libs/ragulate/ragstack_ragulate/pipelines/query_pipeline.py

@@ -19,7 +19,7 @@
 from .feedbacks import Feedbacks
 
 if TYPE_CHECKING:
-    from ragstack_ragulate.datasets import BaseDataset
+    from ragstack_ragulate.datasets import BaseDataset, QueryItem
 
 
 class QueryPipeline(BasePipeline):
@@ -30,7 +30,7 @@ class QueryPipeline(BasePipeline):
     _tru: Tru
     _name: str
     _progress: tqdm[Never]
-    _queries: dict[str, list[str]]
+    _query_items: dict[str, list[QueryItem]]
     _golden_sets: dict[str, list[dict[str, str]]]
     _total_queries: int = 0
     _total_feedbacks: int = 0
@@ -61,7 +61,7 @@ def __init__(
         llm_provider: str = "OpenAI",
         model_name: str | None = None,
     ):
-        self._queries = {}
+        self._query_items = {}
         self._golden_sets = {}
         super().__init__(
             recipe_name=recipe_name,
@@ -89,14 +89,14 @@ def __init__(
 
         total_existing_queries = 0
         for dataset in datasets:
-            queries, golden_set = dataset.get_queries_and_golden_set()
+            query_items = dataset.get_query_items()
             if self.sample_percent < 1.0:
                 if self.random_seed is not None:
                     random.seed(self.random_seed)
                 sampled_indices = random.sample(
-                    range(len(queries)), int(self.sample_percent * len(queries))
+                    range(len(query_items)), int(self.sample_percent * len(query_items))
                 )
-                queries = [queries[i] for i in sampled_indices]
+                query_items = [query_items[i] for i in sampled_indices]
 
             # Check for existing records and filter queries
             existing_records, _feedbacks = self._tru.get_records_and_feedback(
@@ -105,11 +105,15 @@ def __init__(
             existing_queries = existing_records["input"].dropna().tolist()
             total_existing_queries += len(existing_queries)
 
-            queries = [query for query in queries if query not in existing_queries]
+            query_items = [
+                query_item
+                for query_item in query_items
+                if query_item.query not in existing_queries
+            ]
 
-            self._queries[dataset.name] = queries
-            self._golden_sets[dataset.name] = golden_set
-            self._total_queries += len(self._queries[dataset.name])
+            self._query_items[dataset.name] = query_items
+            self._golden_sets[dataset.name] = dataset.get_golden_set()
+            self._total_queries += len(self._query_items[dataset.name])
 
         metric_count = 4
         self._total_feedbacks = self._total_queries * metric_count
@@ -129,7 +133,7 @@ def start_evaluation(self) -> None:
 
     def export_results(self) -> None:
         """Export results."""
-        for dataset_name in self._queries:
+        for dataset_name in self._query_items:
             records, _feedback_names = self._tru.get_records_and_feedback(
                 app_ids=[dataset_name]
             )
@@ -217,7 +221,7 @@ def query(self) -> None:
             initial=self._finished_queries,
         )
 
-        for dataset_name in self._queries:
+        for dataset_name in self._query_items:
             feedback_functions = [
                 feedbacks.answer_correctness(
                     golden_set=self._golden_sets[dataset_name]
@@ -234,14 +238,15 @@ def query(self) -> None:
                 feedback_mode=FeedbackMode.DEFERRED,
             )
 
-            for query in self._queries[dataset_name]:
+            for query_item in self._query_items[dataset_name]:
                 if self._sigint_received:
                     break
                 try:
-                    with recorder:
-                        pipeline.invoke(query)
+                    with recorder as recording:
+                        recording.record_metadata = query_item.metadata
+                        pipeline.invoke(query_item.query)
                 except Exception as e:  # noqa: BLE001
-                    err = f"Query: '{query}' caused exception, skipping."
+                    err = f"Query: '{query_item.query}' caused exception, skipping."
                     logger.exception(err)
                     # TODO: figure out why the logger isn't working after tru-lens starts. For now use print().  # noqa: E501
                     print(f"{err} Exception {e}")