Fixed implementation - now shoud be simplier

Author: MrCabss69

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

@@ -12,6 +12,9 @@
 configure_logging(level=logging.INFO)
 log = logging.getLogger(__name__)
 
+
+TMP_FILE = "./test_file.json"
+
 class DataStore:
     """
     Stores/retrieves documents, queries, and rating scores.
@@ -117,72 +120,75 @@ def has_rating_score(self, query_id: str, doc_id: str) -> bool:
         return context.has_rating_score(doc_id)
 
 
-    def save_queries_and_docs(self, filepath: str | Path) -> None:
-        """
-        Saves all query contexts (query_id, text, doc_ids) to a JSON file.
-        """
-        data = [
-            {
-                "query_id": ctx.get_query_id(),
-                "query_text": ctx.get_query(),
-                "doc_ids": ctx.get_doc_ids()
-            }
-            for ctx in self._queries_by_id.values()
-        ]
+    @staticmethod
+    def query_context_docs_to_dict(query: QueryRatingContext, documents: List[Document]) -> dict:
+        return {
+            "query_id": query.get_query_id(),
+            "query_text": query.get_query(),
+            "doc_ids": query.get_doc_ids(),
+            "doc_ratings": query._doc_id_to_rating_score.copy(),  # copy to avoid mutation
+            "documents": [doc.model_dump_json() for doc in documents if doc is not None]
+        }
+
+    @staticmethod
+    def dict_to_query_context_docs(dict_: dict) -> tuple[str, str, list[str], dict[str, int], list[dict]]:
+        query_id = dict_["query_id"]
+        query_text = dict_["query_text"]
+        doc_ids = dict_["doc_ids"]
+        doc_ratings = dict_["doc_ratings"]
+        documents = dict_.get("documents", [])
+        return query_id, query_text, doc_ids, doc_ratings, documents
+
+
+    def save_tmp_file_content(self, filepath: str | Path = None) -> None:
+        """
+        Save _queries_by_id, _query_text_to_query_id, and self._documents from a unified JSON file in disk
+        """
+        global TMP_FILE
+        if filepath is None:
+            filepath = TMP_FILE
+
+        all_content = []
+        for query_ctx in self._queries_by_id.values():
+            docs_ = [self.get_document(doc_id) for doc_id in query_ctx.get_doc_ids()]
+            data = self.query_context_docs_to_dict(query_ctx, docs_)
+            all_content.append(data)
+
         with open(filepath, "w", encoding="utf-8") as f:
-            json.dump(data, f, indent=2)
+            json.dump(all_content, f, indent=2)
 
 
-    def load_queries_and_docs(self, filepath: str | Path) -> None:
+
+    def load_tmp_file_content(self, filepath: str | Path = None) -> None:
         """
-        Loads query contexts from a JSON file. Reconstructs query_id, query text, and associated doc_ids.
+        Reconstruct _queries_by_id, _query_text_to_query_id, and self._documents from disk file
         """
-        with open(filepath, "r", encoding="utf-8") as f:
-            data = json.load(f)
-
-        for item in data:
-            context = QueryRatingContext(item["query_text"], doc_id=None)
-            context._id = item["query_id"] 
-            for doc_id in item["doc_ids"]:
-                context.add_doc_id(doc_id)
-            self._queries_by_id[context.get_query_id()] = context
-            self._query_text_to_query_id[context.get_query()] = context.get_query_id()
-
-
-    def save_rating_triples(self, filepath: str | Path) -> None:
-        """
-        Saves all (query_id, doc_id, score) triples to a JSON file.
-        """
-        triples = []
-        for _context in self._queries_by_id.values():
-            query_id = _context.get_query_id()
-            for doc_id in _context.get_doc_ids():
-                try:
-                    score = _context.get_rating_score(doc_id)
-                    triples.append({
-                        "query_id": query_id,
-                        "doc_id": doc_id,
-                        "score": score
-                    })
-                except KeyError:
-                    continue
-        with open(filepath, "w", encoding="utf-8") as f:
-            json.dump(triples, f, indent=2)
+        global TMP_FILE
 
+        if filepath is None:
+            filepath = TMP_FILE
 
-    def load_rating_triples(self, filepath: str | Path) -> None:
-        """
-        Loads rating triples from a JSON file and updates existing QueryRatingContext entries.
-        """
         with open(filepath, "r", encoding="utf-8") as f:
-            triples = json.load(f)
-
-        for triple in triples:
-            query_id = triple["query_id"]
-            doc_id = triple["doc_id"]
-            score = triple["score"]
-            if query_id not in self._queries_by_id:
-                raise ValueError(f"Query ID {query_id} not found when loading triples.")
-            self._queries_by_id[query_id].add_rating_score(doc_id, score)
+            all_content = json.load(f)
+        
+        # Loop to reconstruct each register
+        for entry in all_content:
+            query_id, query_text, doc_ids, doc_ratings, documents = self.dict_to_query_context_docs(entry)
+
+            # Register docs
+            for doc_data in documents:
+                document = Document.model_validate(doc_data)
+                self.add_document(document.id, document)
+
+            # Reconstruct the QueryRating
+            context = QueryRatingContext(query=query_text)
+            context._id = query_id  # generated UUID
+
+            for doc_id in doc_ids:
+                context.add_doc_id(doc_id)
 
+            for doc_id, rating in doc_ratings.items():
+                context.add_rating_score(doc_id, rating)
 
+            self._queries_by_id[query_id] = context
+            self._query_text_to_query_id[query_text] = query_id

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

@@ -61,48 +61,69 @@ def test_add_and_get_query_expect_query_stored_in_data_store_and_check_same_quer
 
 
 # tmp_path: pytest fixture with a temporal dir
-def test_save_and_load_queries_and_ratings_roundtrip(tmp_path):
+def test_save_tmp_file_content_expect_json_created(tmp_path):
+    ds = DataStore()
 
-    # Step 1: build the intial instance
-    ds1 = DataStore()
-
-    # Step 2: add Documents
+    # Add Documents
     doc1 = Document(id="d1", fields={"title": "AI", "text": "Deep learning"})
     doc2 = Document(id="d2", fields={"title": "LLMs", "text": "Transformers"})
-    ds1.add_document(doc1.id, doc1)
-    ds1.add_document(doc2.id, doc2)
-
-    # Step 3: add queries and ratings
-    qid1 = ds1.add_query("artificial intelligence", doc1.id)
-    ds1.add_rating_score(qid1, doc1.id, 1)
-    ds1.add_rating_score(qid1, doc2.id, 0)
-
-    qid2 = ds1.add_query("transformer models", doc2.id)
-    ds1.add_rating_score(qid2, doc2.id, 1)
-
-    # Step 4: save into disk
-    queries_path = tmp_path / "queries.json"
-    triples_path = tmp_path / "triples.json"
-
-    ds1.save_queries_and_docs(queries_path)
-    ds1.save_rating_triples(triples_path)
-
-    # Step 5: Verify saving
-    assert queries_path.exists()
-    assert triples_path.exists()
-    assert json.loads(queries_path.read_text(encoding="utf-8"))  # assert not empty 
-    assert json.loads(triples_path.read_text(encoding="utf-8"))  # assert not empty 
-
-    # Step 6: load a new datastore - avoid caching artifacts
-    ds2 = DataStore()
-    ds2.load_queries_and_docs(queries_path)
-    ds2.load_rating_triples(triples_path)
-
-    # Step 7: verify integrity after loading
-    assert ds2.get_query(qid1).get_query() == "artificial intelligence"
-    assert set(ds2.get_query(qid1).get_doc_ids()) == {"d1", "d2"}
-    assert ds2.get_rating_score(qid1, "d1") == 1
-    assert ds2.get_rating_score(qid1, "d2") == 0
-
-    assert ds2.get_query(qid2).get_query() == "transformer models"
-    assert ds2.get_rating_score(qid2, "d2") == 1
+    ds.add_document(doc1.id, doc1)
+    ds.add_document(doc2.id, doc2)
+
+    # Add Queries and Ratings
+    qid1 = ds.add_query("artificial intelligence", doc1.id)
+    ds.add_rating_score(qid1, "d1", 1)
+    ds.add_rating_score(qid1, "d2", 0)
+
+    qid2 = ds.add_query("transformer models", doc2.id)
+    ds.add_rating_score(qid2, "d2", 1)
+
+    # Save to disk
+    path = tmp_path / "datastore.json"
+    ds.save_tmp_file_content(path)
+
+    # Assertions
+    assert path.exists()
+    data = json.loads(path.read_text(encoding="utf-8"))
+    assert isinstance(data, list)
+    assert len(data) == 2
+    for entry in data:
+        assert "query_id" in entry
+        assert "query_text" in entry
+        assert "doc_ids" in entry
+        assert "doc_ratings" in entry
+        assert "documents" in entry
+
+
+def test_load_tmp_file_content_expect_data_restored_correctly(tmp_path):
+    # Create known content
+    content = [
+        {
+            "query_id": "q1",
+            "query_text": "ai",
+            "doc_ids": ["d1", "d2"],
+            "doc_ratings": {"d1": 1, "d2": 0},
+            "documents": [
+                {"id": "d1", "fields": {"title": "AI", "text": "Deep learning"}},
+                {"id": "d2", "fields": {"title": "LLMs", "text": "Transformers"}},
+            ]
+        }
+    ]
+
+    # Write file manually
+    path = tmp_path / "datastore.json"
+    path.write_text(json.dumps(content, indent=2), encoding="utf-8")
+
+    # Load from disk
+    ds = DataStore()
+    ds.load_tmp_file_content(path)
+
+    # Assertions
+    assert ds.get_query("q1").get_query() == "ai"
+    assert ds.get_rating_score("q1", "d1") == 1
+    assert ds.get_rating_score("q1", "d2") == 0
+
+    doc1 = ds.get_document("d1")
+    doc2 = ds.get_document("d2")
+    assert doc1.fields["title"] == "AI"
+    assert doc2.fields["text"] == "Transformers"