DAGE-20: added ElasticSearch support

rre-dataset-generator/README.md

@@ -92,11 +92,11 @@ Depending on your Docker version, you may need to use `docker compose` instead o
 If you have Docker Compose v1 installed, use:
 
 ```bash
-docker-compose up --build
+docker-compose -f docker-compose.solr.yml up --build
 ```
 If you have Docker Compose v2 installed, use:
 ```bash
-docker compose up --build
+docker compose -f docker-compose.solr.yml up --build
 ```
 
 This will start 2 services:
@@ -126,3 +126,26 @@ This will start 2 services:
  - `opensearch`, available at http://localhost:9200/
  - `opensearch-init`, loads documents (`bulk indexing`) from opensearch-init/data/dataset.jsonl.
 
+
+##### Running Elasticsearch (Single Node)
+
+Similarly to Solr, to run a local Elasticsearch test environment using docker-compose:
+```bash
+cd tests/integration/
+```
+
+Depending on your Docker version, you may need to use `docker compose` instead of `docker-compose`.
+If you have Docker Compose v1 installed, use:
+
+```bash
+docker-compose -f docker-compose.elasticsearch.yml up --build 
+```
+If you have Docker Compose v2 installed, use:
+```bash
+docker compose -f docker-compose.elasticsearch.yml up --build
+```
+
+This will start 2 services:
+ - `elasticsearch`, available at http://localhost:9200
+ - `elasticsearch-init`, loads documents from elasticsearch-init/data/dataset.jsonl only if Elasticsearch doesn't have 
+any documents in the index.

rre-dataset-generator/dataset_generator.py

@@ -44,7 +44,7 @@ def add_user_queries(config: Config, data_store: DataStore):
                     data_store.add_query(cleaned)
 
 
-def generate_and_add_queries(llm_service: LLMService, config: Config, data_store: DataStore) -> None:
+def generate_and_add_queries(config: Config, data_store: DataStore, llm_service: LLMService, docs_to_generate_queries: List[Document]) -> None:
     num_queries_per_doc: int = int(
         ((config.num_queries_needed - len(data_store.get_queries())) // config.doc_number) * 1.5)
 
@@ -59,7 +59,7 @@ def generate_and_add_queries(llm_service: LLMService, config: Config, data_store
                                         "Default max rating is assigned because the query is generated by the document")
 
 
-def retrieve_and_add_documents(config: Config, data_store: DataStore) -> None:
+def retrieve_and_add_documents(config: Config, data_store: DataStore, search_engine: BaseSearchEngine) -> 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_text(),
                                                                        query_template=config.query_template,
@@ -68,8 +68,7 @@ def retrieve_and_add_documents(config: Config, data_store: DataStore) -> None:
             if not data_store.has_document(doc.id):
                 data_store.add_document(doc.id, doc)
 
-
-def add_cartesian_product_scores(llm_service: LLMService, config: Config, data_store: DataStore) -> None:
+def add_cartesian_product_scores(config: Config, data_store: DataStore, llm_service: LLMService) -> 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):
@@ -79,7 +78,8 @@ def add_cartesian_product_scores(llm_service: LLMService, config: Config, data_s
                                                                               config.save_llm_explanation)
                 data_store.add_rating_score(query_rating_context.get_query_id(),
                                             doc.id,
-                                            score_response.get_score(), score_response.get_explanation())
+                                            score_response.get_score(),
+                                            score_response.get_explanation())
 
 
 if __name__ == "__main__":
@@ -105,11 +105,11 @@ def add_cartesian_product_scores(llm_service: LLMService, config: Config, data_s
         doc_fields=config.doc_fields)
     log.debug(f"Number of documents retrieved for generation: {len(docs_to_generate_queries)}")
 
-    generate_and_add_queries(service, config, data_store)
+    generate_and_add_queries(config, data_store, service, docs_to_generate_queries)
 
-    retrieve_and_add_documents(config, data_store)
+    retrieve_and_add_documents(config, data_store, search_engine)
 
-    add_cartesian_product_scores(service, config, data_store)
+    add_cartesian_product_scores(config, data_store, service)
 
     writer.write(config.output_destination)
 

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

@@ -0,0 +1,157 @@
+import json
+from json import JSONDecodeError
+
+import requests
+from urllib.parse import urljoin
+from pydantic import HttpUrl
+from requests.exceptions import HTTPError, ConnectionError, Timeout, RequestException
+from typing import List, Dict, Any, Union
+
+from src.search_engine.search_engine_base import BaseSearchEngine
+from src.model.document import Document
+from src.utils import clean_text
+
+import logging
+log = logging.getLogger(__name__)
+
+
+class ElasticsearchSearchEngine(BaseSearchEngine):
+    """
+    Elasticsearch implementation to search into a given collection
+    """
+    def __init__(self, endpoint: HttpUrl | str):
+        super().__init__(endpoint)
+        self.HEADERS = {'Content-Type': 'application/json'}
+        log.debug(f"Working on endpoint: {self.endpoint}")
+        self.UNIQUE_KEY = "_id"
+
+    def fetch_for_query_generation(self,
+                                   documents_filter: Union[None, List[Dict[str, List[str]]]],
+                                   doc_number: int,
+                                   doc_fields: List[str]) -> List[Document]:
+        """
+        Fetches a set of documents from Elasticsearch for query generation purposes.
+
+        Args:
+            documents_filter (Union[None, List[Dict[str, List[str]]]]): Optional list of field filters to apply.
+                Each filter is a dictionary mapping field names to allowed values.
+            doc_number (int): Number of documents to retrieve.
+            doc_fields (List[str]): List of field names to include in the output.
+
+        Returns:
+            List[Document]: A list of documents formatted as `Document` instances.
+        """
+        # Build base query
+        query = {"match_all": {}}
+
+        # Add filters, if provided
+        filter_clauses = []
+        if documents_filter is not None:
+            for dict_field in documents_filter:
+                for field, values in dict_field.items():
+                    if not values:
+                        continue
+                    filter_clauses.append({"terms": {field: values}})
+
+        # Wrap in a bool query if there are any filters
+        if filter_clauses:
+            query = {
+                "bool": {
+                    "must": {"match_all": {}},
+                    "filter": filter_clauses
+                }
+            }
+
+        # Construct the payload (Elasticsearch query body)
+        payload = {
+            "size": doc_number,
+            "query": query,
+            "_source": doc_fields
+        }
+
+        return self._search(payload)
+
+    def fetch_for_evaluation(self, query_template: str, doc_fields: List[str], keyword: str=None) -> List[Document]:
+        """
+        Executes a search for evaluation using a query template with an optional keyword substitution.
+
+        Args:
+            query_template (str): A JSON-formatted string representing the Elasticsearch query,
+                possibly containing a placeholder for a keyword.
+            doc_fields (List[str]): List of field names to include in the response.
+            keyword (str, optional): A keyword to replace the placeholder in the query.
+                If not provided, a default match_all query is used.
+
+        Returns:
+            List[Document]: A list of documents matching the query.
+        """
+        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]:
+        """
+        Executes the search request to the Elasticsearch `_search` endpoint and parses the response.
+
+        Args:
+            payload (Dict[str, Any]): JSON payload representing the Elasticsearch query.
+
+        Returns:
+            List[Document]: A list of retrieved documents as `Document` instances.
+        """
+        search_url = urljoin(self.endpoint.encoded_string(), '_search')
+
+        try:
+            response = requests.post(search_url, headers=self.HEADERS, json=payload)
+            response.raise_for_status()
+        except (ConnectionError, Timeout, RequestException, HTTPError) as e:
+            log.error(f"ElasticSearch query failed: {e}\nPayload: {payload}")
+            raise
+
+        hits = response.json().get('hits', {}).get('hits', [])
+        result = []
+        for hit in hits:
+            source = hit.get("_source", {})
+            doc_id = source.get("id", hit.get(self.UNIQUE_KEY))
+
+            fields = {
+                key: self._normalize(value)
+                for key, value in source.items()
+                if key != "id"
+            }
+
+            result.append(Document(id=doc_id, fields=fields))
+        return result
+
+    @staticmethod
+    def _normalize(value: Any) -> List[str]:
+        """Normalize a field value into a list of cleaned strings or throws an exception."""
+        try:
+            if value is None:
+                return []
+
+            if isinstance(value, str):
+                return [clean_text(value)]
+
+            if isinstance(value, list):
+                return [clean_text(v) if isinstance(v, str) else str(v) for v in value]
+
+            if isinstance(value, dict):
+                cleaned_dict = {
+                    k: clean_text(v) if isinstance(v, str) else v
+                    for k, v in value.items()
+                }
+                return [json.dumps(cleaned_dict)]
+
+            return [str(value)]
+        except Exception as e:
+            raise ValueError(f"Failed to normalize value: {value}") from e

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

@@ -3,6 +3,8 @@
 from .opensearch_engine import OpenSearchEngine
 from .search_engine_base import BaseSearchEngine
 from .solr_search_engine import SolrSearchEngine
+from .elasticsearch_search_engine import ElasticsearchSearchEngine
+
 import logging
 
 log = logging.getLogger(__name__)
@@ -12,6 +14,7 @@ class SearchEngineFactory:
     SEARCH_ENGINE_REGISTRY = {
         "solr": SolrSearchEngine,
         "opensearch": OpenSearchEngine,
+        "elasticsearch": ElasticsearchSearchEngine
     }
 
     @classmethod