[DAGE-20] changed tests and _search method in both elastic and solr

Author: nicolo-rinaldi

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

@@ -7,6 +7,7 @@
 
 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__)
@@ -114,11 +115,41 @@ def _search(self, payload: Dict[str, Any]) -> List[Document]:
             raise
 
         data = response.json()
-        raw_docs = (data.get('hits') or {}).get('hits') or []
-        reformat_raw_doc = []
-        for doc in raw_docs:
-            clean_doc = dict()
-            clean_doc['id'] = doc[self.UNIQUE_KEY]
-            clean_doc['fields'] = doc['_source']
-            reformat_raw_doc.append(Document(**clean_doc))
-        return reformat_raw_doc
+        hits = (data.get('hits') or {}).get('hits') or []
+        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/solr_search_engine.py

@@ -127,24 +127,39 @@ def _search(self, payload: Dict[str, Any]) -> List[Document]:
             log.error(f"Solr query failed: {e}\nPayload: {payload}")
             raise
 
-        data = response.json()
-        raw_docs = (data.get('response') or {}).get('docs') or []
-        reformat_raw_doc = []
-        for doc in raw_docs:
-             clean_doc = dict()
-             clean_doc['id'] = doc[self.UNIQUE_KEY]
-             clean_doc['fields'] = dict()
-             for k, v in doc.items():
-                if k != self.UNIQUE_KEY:
-                    if isinstance(v, list):
-                        if v:
-                            if isinstance(v[0], str):
-                                clean_doc['fields'][k] = [clean_text(text) for text in v]
-                            else:
-                                clean_doc['fields'][k] = v
-                        else:
-                            log.warning(f"The field {k} is empty, skipped.")
-                    else:
-                        clean_doc['fields'][k] = v
-             reformat_raw_doc.append(Document(**clean_doc))
-        return reformat_raw_doc
+        hits = (response.json().get('response') or {}).get('docs') or []
+        result = []
+        for hit in hits:
+             doc_id = hit.get(self.UNIQUE_KEY)
+             fields = {
+                 key: self._normalize(value)
+                 for key, value in hit.items()
+                 if key != self.UNIQUE_KEY
+             }
+
+             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/tests/unit/test_elasticsearch_search_engine.py

@@ -21,24 +21,29 @@ def elasticsearch_config():
     """Fixture that loads a valid OpenSearch config for unit tests."""
     return Config.load("tests/unit/resources/elasticsearch_good_config.yaml")
 
-def test_elasticsearch_search_engine(monkeypatch, elasticsearch_config):
-    url = "https://fakeurl"
-    search_engine = ElasticsearchSearchEngine(url)
-
-    payload = json.dumps({"match_all": {}})
-
-    mock_doc = {
+@pytest.fixture
+def mock_doc():
+    return {
         "_id": "1",
         '_source': {
-            "mock_title": "A first mocked title",
-            "mock_description": "A first mocked description"
+            "mock_title": ["A first mocked title"],
+            "mock_description": ["A first mocked description"]
         }
     }
-    mock_dict = {
+
+@pytest.fixture
+def mock_dict(mock_doc):
+    return {
         'id': mock_doc['_id'],
         'fields': mock_doc["_source"]
     }
 
+def test_elasticsearch_search_engine_fetch_for_query_generation(monkeypatch, elasticsearch_config, mock_doc, mock_dict):
+    url = "https://fakeurl"
+    search_engine = ElasticsearchSearchEngine(url)
+
+    payload = json.dumps({"match_all": {}})
+
     # apply the monkeypatch for requests.post to mock_post
     monkeypatch.setattr(requests, "post",
                         lambda *args, **kwargs: MockResponseElasticsearchEngine([mock_doc], status_code=200)
@@ -48,6 +53,17 @@ def test_elasticsearch_search_engine(monkeypatch, elasticsearch_config):
                                                       doc_number=elasticsearch_config.doc_number,
                                                       doc_fields=elasticsearch_config.doc_fields)
     assert result[0] == Document(**mock_dict)
+
+def test_elasticsearch_search_engine_fetch_for_evaluation(monkeypatch, elasticsearch_config, mock_doc, mock_dict):
+    url = "https://fakeurl"
+    search_engine = ElasticsearchSearchEngine(url)
+
+    payload = json.dumps({"match_all": {}})
+
+    # apply the monkeypatch for requests.post to mock_post
+    monkeypatch.setattr(requests, "post",
+                        lambda *args, **kwargs: MockResponseElasticsearchEngine([mock_doc], status_code=200)
+                        )
     # search_engine.extract_documents_to_evaluate_system, which contains requests.post, uses the monkeypatch
     result = search_engine.fetch_for_evaluation(keyword="and",
                                                 query_template=elasticsearch_config.query_template,

rre-dataset-generator/tests/unit/test_solr_search_engine.py

@@ -19,22 +19,27 @@ def solr_config():
     """Fixture that loads a valid OpenSearch config for unit tests."""
     return Config.load("tests/unit/resources/solr_good_config.yaml")
 
-def test_solr_search_engine(monkeypatch, solr_config):
-    monkeypatch.setattr(requests, "get", lambda *args, **kwargs: MockResponseUniqueKey(ident="mock_id"))
-    search_engine = SolrSearchEngine("https://fakeurl")
-
-    assert search_engine.UNIQUE_KEY == "mock_id"
-
-    mock_doc = {
+@pytest.fixture
+def mock_doc():
+    return {
         "mock_id": "1",
-        "mock_title": "A first mocked title",
-        "mock_description": "A first mocked description"
+        "mock_title": ["A first mocked title"],
+        "mock_description": ["A first mocked description"]
     }
-    mock_dict = {
+
+@pytest.fixture
+def mock_dict(mock_doc):
+    return {
         'id': mock_doc['mock_id'],
-        'fields': {k:v for k, v in mock_doc.items() if k !='mock_id'}
+        'fields': {k: v for k, v in mock_doc.items() if k != 'mock_id'}
     }
 
+def test_solr_search_engine_fetch_for_query_generation(monkeypatch, solr_config, mock_doc, mock_dict):
+    monkeypatch.setattr(requests, "get", lambda *args, **kwargs: MockResponseUniqueKey(ident="mock_id"))
+    search_engine = SolrSearchEngine("https://fakeurl")
+
+    assert search_engine.UNIQUE_KEY == "mock_id"
+
     # apply the monkeypatch for requests.post to mock_post
     monkeypatch.setattr(requests, "post", lambda *args, **kwargs: MockResponseSolrEngine([mock_doc], status_code=200))
 
@@ -43,6 +48,16 @@ def test_solr_search_engine(monkeypatch, solr_config):
                                                       doc_number=solr_config.doc_number,
                                                       doc_fields=solr_config.doc_fields)
     assert result[0] == Document(**mock_dict)
+
+def test_solr_search_engine_fetch_for_evaluation(monkeypatch, solr_config, mock_doc, mock_dict):
+    monkeypatch.setattr(requests, "get", lambda *args, **kwargs: MockResponseUniqueKey(ident="mock_id"))
+    search_engine = SolrSearchEngine("https://fakeurl")
+
+    assert search_engine.UNIQUE_KEY == "mock_id"
+
+    # apply the monkeypatch for requests.post to mock_post
+    monkeypatch.setattr(requests, "post", lambda *args, **kwargs: MockResponseSolrEngine([mock_doc], status_code=200))
+
     # search_engine.extract_documents_to_evaluate_system, which contains requests.post, uses the monkeypatch
     result = search_engine.fetch_for_evaluation(keyword="and",
                                                 query_template=solr_config.query_template,