[DAGE-20] alignment with Naz update on opensearch query template

Author: nicolo-rinaldi

rre-dataset-generator/dataset_generator.py

@@ -38,10 +38,11 @@ def add_user_queries(config: Config, data_store: DataStore):
     if config.queries is not None:
         with open(config.queries, 'r', encoding='utf-8') as file:
             for line in file:
-                if line.strip():
-                    data_store.add_query(line)
+                cleaned = line.strip()
+                if cleaned:
+                    data_store.add_query(cleaned)
 
-def generate_and_add_queries(config: Config, data_store: DataStore, llm_service: LLMService, docs_to_generate_queries: List[Document]) -> None:
+def generate_and_add_queries(llm_service: LLMService, config: Config, data_store: DataStore) -> None:
     num_queries_per_doc: int = int(((config.num_queries_needed - len(data_store.get_queries())) // config.doc_number) * 1.5)
 
     for doc in docs_to_generate_queries:
@@ -54,7 +55,7 @@ def generate_and_add_queries(config: Config, data_store: DataStore, llm_service:
             data_store.add_rating_score(query_id, doc.id, max(config.relevance_label_set))
 
 
-def retrieve_and_add_documents(config: Config, data_store: DataStore, search_engine: BaseSearchEngine) -> None:
+def retrieve_and_add_documents(config: Config, data_store: DataStore) -> None:
     for query_rating_context in data_store.get_queries():
         docs_eval: List[Document] = search_engine.fetch_for_evaluation(keyword=query_rating_context.get_query(),
                                                                        query_template=config.query_template,
@@ -63,7 +64,7 @@ def retrieve_and_add_documents(config: Config, data_store: DataStore, search_eng
             if not data_store.has_document(doc.id):
                 data_store.add_document(doc.id, doc)
 
-def add_cartesian_product_scores(config: Config, data_store: DataStore, llm_service: LLMService) -> None:
+def add_cartesian_product_scores(llm_service: LLMService, config: Config, data_store: DataStore) -> None:
     for query_rating_context in data_store.get_queries():
         for doc in data_store.get_documents():
             if not data_store.has_rating_score(query_rating_context.get_query_id(), doc.id):
@@ -97,11 +98,11 @@ def add_cartesian_product_scores(config: Config, data_store: DataStore, llm_serv
                                                                                         doc_fields=config.doc_fields)
     log.debug(f"Number of documents retrieved for generation: {len(docs_to_generate_queries)}")
 
-    generate_and_add_queries(config, data_store, service, docs_to_generate_queries)
+    generate_and_add_queries(service, config, data_store)
 
-    retrieve_and_add_documents(config, data_store, search_engine)
+    retrieve_and_add_documents(config, data_store)
 
-    add_cartesian_product_scores(config, data_store, service)
+    add_cartesian_product_scores(service, config, data_store)
 
     writer.write(config.output_destination)
 

rre-dataset-generator/src/search_engine/elasticsearch_search_engine.py

@@ -1,4 +1,6 @@
 import json
+from json import JSONDecodeError
+
 import requests
 from urllib.parse import urljoin
 from pydantic import HttpUrl
@@ -86,13 +88,17 @@ def fetch_for_evaluation(self, query_template: str, doc_fields: List[str], keywo
         Returns:
             List[Document]: A list of documents matching the query.
         """
-        if keyword:
-             payload = json.loads(query_template.replace(self.PLACEHOLDER, keyword))
-        else:
-            payload = {
-            "query": {"match_all": {}}
-            }
-        payload["_source"] = doc_fields
+        try:
+            payload: Dict[str, Any] = json.loads(query_template)
+        except JSONDecodeError as e:
+            raise ValueError(f"Invalid JSON query_template: {e}")
+
+        query_string_obj = payload.get("query", {}).get("query_string", {})
+        if "query" in query_string_obj:
+            query_string_obj["query"] = query_string_obj["query"].replace(self.PLACEHOLDER, keyword)
+
+        fields = doc_fields if self.UNIQUE_KEY in doc_fields else doc_fields + [self.UNIQUE_KEY]
+        payload["_source"] = fields
         return self._search(payload)
 
     def _search(self, payload: Dict[str, Any]) -> List[Document]:

rre-dataset-generator/tests/unit/resources/elasticsearch_good_config.yaml

@@ -1,4 +1,11 @@
-query_template: '{ "query": { "match": { "message": "#$query##" } } }'
+query_template: > 
+    {
+        "query": { 
+            "match": { 
+                "message": "#$query##" 
+            } 
+        } 
+    }
 search_engine_type: "elasticsearch"
 search_engine_collection_endpoint: 'http://localhost:9200/testcore/'
 doc_number: 100