DAGE-19: Add OpenSearch support

Author: nseidan

rre-dataset-generator/README.md

@@ -45,6 +45,9 @@ output_destination: "output/generated_dataset.json"
 
 ### 3. Running the Generator
 
+Before running the command below, you need to have running search engine instance (`solr`/`opensearch`/`elasticsearch`/`vespa`).
+
+
 Execute the main script via the `argparse` CLI, pointing to your configuration file:
 ```bash
 uv run dataset_generator.py --config_file config.yaml
@@ -98,4 +101,28 @@ docker compose up --build
 
 This will start 2 services:
  - `solr`, available at http://localhost:8983/solr
- - `solr-init`, loads documents from solr/data/dataset.json only if Solr doesn't index any documents.
+ - `solr-init`, loads documents from solr-init/data/dataset.json.
+
+
+##### Running OpenSearch (Single Node)
+
+To run a local OpenSearch 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.opensearch.yml up --build
+```
+If you have Docker Compose v2 installed, use:
+```bash
+docker compose -f docker-compose.opensearch.yml up --build
+```
+
+This will start 2 services:
+ - `opensearch`, available at http://localhost:9200/
+ - `opensearch-init`, loads documents (`bulk indexing`) from opensearch-init/data/dataset.jsonl.
+

rre-dataset-generator/config.yaml

@@ -9,6 +9,7 @@ query_template: "q=#$query##&df=description&rows=2&q.op=OR&wt=json"
 search_engine_type: "solr"
 
 # Endpoint URL of the search engine collection or index
+# opensearch: http://localhost:9200/testcore/
 search_engine_collection_endpoint: "http://localhost:8983/solr/testcore/"
 
 # (Optional) Filter query to restrict the set of documents used to generate queries.

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

@@ -0,0 +1,113 @@
+import logging
+from typing import List, Dict, Any, Union
+
+import requests
+from pydantic import HttpUrl
+from requests.exceptions import HTTPError, ConnectionError, Timeout, RequestException
+
+from src.model.document import Document
+from src.search_engine.search_engine_base import BaseSearchEngine
+from src.utils import clean_text
+
+log = logging.getLogger(__name__)
+
+
+class OpenSearchEngine(BaseSearchEngine):
+    """
+    OpenSearch implementation to search in a given index.
+    """
+
+    def __init__(self, endpoint: HttpUrl):
+        super().__init__(endpoint)
+        self.HEADERS = {'Content-Type': 'application/json'}
+        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]:
+        filters = []
+        if documents_filter:
+            for field_values in documents_filter:
+                for field, values in field_values.items():
+                    if not values:
+                        continue
+                    if len(values) == 1:
+                        filters.append({"term": {field: values[0]}})
+                    else:
+                        filters.append({"terms": {field: values}})
+
+        fields = doc_fields if self.UNIQUE_KEY in doc_fields else doc_fields + [self.UNIQUE_KEY]
+
+        if filters:
+            query = {
+                "bool": {
+                    "filter": filters
+                }
+            }
+        else:
+            query = {
+                "match_all": {}
+            }
+
+        payload = {
+            "query": query,
+            "_source": fields,
+            "size": doc_number
+        }
+
+        return self._search(payload)
+
+    def fetch_for_evaluation(self, query_template: str, doc_fields: List[str], keyword: str = "*") -> List[Document]:
+        query = query_template.replace(self.PLACEHOLDER, keyword)
+        fields = doc_fields if self.UNIQUE_KEY in doc_fields else doc_fields + [self.UNIQUE_KEY]
+
+        payload = {
+            "query": {
+                "query_string": {
+                    "query": query
+                }
+            },
+            "_source": fields
+        }
+        return self._search(payload)
+
+    def _search(self, payload: Dict[str, Any]) -> List[Document]:
+        search_url = f"{self.endpoint}/_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"OpenSearch 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("_id"))
+
+            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) -> Any:
+        """
+        Normalize a field value:
+        - string → return [cleaned string]
+        - list of strings → return [cleaned strings]
+        - else → return as it is
+        """
+        if isinstance(value, str):
+            return [clean_text(value)]
+        if isinstance(value, list) and all(isinstance(i, str) for i in value):
+            return [clean_text(i) for i in value]
+        return value

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

@@ -1,15 +1,19 @@
 from pydantic import HttpUrl
 
+from .opensearch_engine import OpenSearchEngine
 from .search_engine_base import BaseSearchEngine
 from .solr_search_engine import SolrSearchEngine
 import logging
 
 log = logging.getLogger(__name__)
 
+
 class SearchEngineFactory:
     SEARCH_ENGINE_REGISTRY = {
         "solr": SolrSearchEngine,
+        "opensearch": OpenSearchEngine,
     }
+
     @classmethod
     def build(cls, search_engine_type: str, endpoint: HttpUrl) -> BaseSearchEngine:
         if search_engine_type not in cls.SEARCH_ENGINE_REGISTRY:

rre-dataset-generator/tests/integration/docker-compose.opensearch.yml

@@ -0,0 +1,33 @@
+services:
+  opensearch:
+    image: opensearchproject/opensearch:3.1.0
+    container_name: opensearch
+    environment:
+      - discovery.type=single-node
+      - node.name=opensearch
+      - bootstrap.memory_lock=true
+      - DISABLE_INSTALL_DEMO_CONFIG=true
+      - DISABLE_SECURITY_PLUGIN=true
+    ports:
+      - 9200:9200
+      - 9600:9600
+
+  opensearch-dashboards:
+    image: opensearchproject/opensearch-dashboards:3.1.0
+    container_name: opensearch-dashboards
+    ports:
+      - 5601:5601
+    expose:
+      - "5601"
+    environment:
+      - OPENSEARCH_HOSTS=http://opensearch:9200
+      - DISABLE_SECURITY_DASHBOARDS_PLUGIN=true
+    depends_on:
+      - opensearch
+
+  opensearch-init:
+    build:
+      context: opensearch-init
+    depends_on:
+      - opensearch
+    command: ["sh", "-c", "/opt/rre-dataset-generator/opensearch-init.sh"]

rre-dataset-generator/tests/integration/opensearch-init/Dockerfile

@@ -0,0 +1,4 @@
+FROM curlimages/curl:8.14.1
+
+COPY data/dataset.jsonl /opt/rre-dataset-generator/data/dataset.jsonl
+COPY --chmod=777 opensearch-init.sh /opt/rre-dataset-generator/opensearch-init.sh