Dage 18: Persistent Storage JSON

rre-dataset-generator/src/model/document.py

@@ -17,6 +17,7 @@ class Document(BaseModel):
         ...,
         description="Fields of the document."
     )
+
     @field_validator('fields')
     @classmethod
     def check_no_empty_fields(cls, v: Dict[str, Any]) -> Dict[str, Any]:

rre-dataset-generator/src/model/query_rating_context.py

@@ -1,5 +1,6 @@
+from __future__ import annotations
 import uuid
-from typing import Dict, List
+from typing import Dict, List, Any
 
 
 class QueryRatingContext:
@@ -12,19 +13,19 @@ class QueryRatingContext:
 
     DOC_NOT_RATED: int = -1  # doc is not yet rated
 
-    def __init__(self, query: str, doc_id: str | None = None):
-        self._id: str = str(uuid.uuid4())
+    def __init__(self, query: str, doc_id: str | None = None, query_id: str | None = None):
+        self._id: str = str(uuid.uuid4()) if query_id is None else str(query_id)
         self._query: str = query
         self._doc_id_to_rating_score: Dict[str, int] = {}
-        # HANDLING NONEs - throwed error in some tests
+        # HANDLING NONEs - threw error in some tests
         if doc_id is not None:
             self._doc_id_to_rating_score[doc_id] = self.DOC_NOT_RATED
 
     def get_query_id(self) -> str:
         """Return the unique identifier for this query context."""
         return self._id
 
-    def get_query(self) -> str:
+    def get_query_text(self) -> str:
         """Return the original query string."""
         return self._query
 
@@ -40,8 +41,33 @@ def add_rating_score(self, doc_id: str, rating_score: int) -> None:
         self._doc_id_to_rating_score[doc_id] = rating_score
 
     def has_rating_score(self, doc_id: str) -> bool:
-        return doc_id in self._doc_id_to_rating_score and self._doc_id_to_rating_score[doc_id] != self.DOC_NOT_RATED
+        return doc_id in self._doc_id_to_rating_score and self._doc_id_to_rating_score.get(doc_id) != self.DOC_NOT_RATED
 
     def get_rating_score(self, doc_id: str) -> int:
-        return self._doc_id_to_rating_score[doc_id]
+        if self.has_rating_score(doc_id):
+            return self._doc_id_to_rating_score[doc_id]
+        raise KeyError(f"Rating for doc_id {doc_id} not found.")
 
+
+    @classmethod
+    def from_dict(cls, context_as_dict: Dict[str, Any]) -> "QueryRatingContext":
+        query_id = context_as_dict.get("id", None)
+        query_text = context_as_dict.get("query", "")
+        doc_ratings = context_as_dict.get("doc_ratings", {})
+
+        context = cls(query=query_text, query_id=query_id)
+        for doc_id, rating in doc_ratings.items():
+            context.add_doc_id(doc_id)
+            context.add_rating_score(doc_id, rating)
+
+        return context
+
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Convert the QueryRatingContext to a JSON-serializable dictionary.
+        """
+        return {
+            "id": self._id,
+            "query": self._query,
+            "doc_ratings": self._doc_id_to_rating_score
+        }

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

@@ -1,40 +1,55 @@
 from __future__ import annotations
 
+import json
 import logging
-from typing import Dict, List, Optional
+import os
+from pathlib import Path
+from typing import Dict, List, Optional, Any
 
 from src.model.document import Document
 from src.model.query_rating_context import QueryRatingContext
 
 log = logging.getLogger(__name__)
 
 
+TMP_FILE = "./tmp/datastore.json"
+
 class DataStore:
     """
     Stores/retrieves documents, queries, and rating scores.
     """
-
-    def __init__(self):
-        self._documents: Dict[str, Document] = {}  # doc_id -> Document
-        self._queries_by_id: Dict[str, QueryRatingContext] = {}  # query_id → QueryRatingContext
-        self._query_text_to_query_id: Dict[str, str] = {}  # query_text -> query_id
+    def __init__(self, ignore_saved_data: bool = False):
+        self._documents: Dict[str, Document] = {}
+        self._queries_by_id: Dict[str, QueryRatingContext] = {}
+        self._query_text_to_query_id: Dict[str, str] = {}
+
+        # Load from default file if present
+        if not ignore_saved_data and os.path.exists(TMP_FILE):
+            try:
+                self.load_tmp_file_content()
+                log.debug(f"Loaded DataStore from default file: {TMP_FILE}")
+            except Exception as e:
+                log.error(f"Could not load default datastore file '{TMP_FILE}': {e}")
 
     def _get_query_rating_context_by_id(self, query_id: str) -> QueryRatingContext:
         if query_id not in self._queries_by_id:
-            log.error("Query id %s not found in DataStore", query_id)
-            raise KeyError(f"Query id '{query_id}' not found in DataStore")
+            _error_msg = f"Query id {query_id} not found in DataStore"
+            log.error(_error_msg)
+            raise KeyError(_error_msg)
         return self._queries_by_id[query_id]
 
     def _get_document(self, doc_id: str) -> Optional[Document]:
         if doc_id not in self._documents:
-            log.warning("Document id %s not found in DataStore", doc_id)
+            _warning_msg = f"Detected an error when retrieving a document from the data store. Document {doc_id} not found in DataStore"
+            log.warning(_warning_msg)
             return None
         return self._documents[doc_id]
 
     def add_document(self, doc_id: str, document: Document) -> None:
         if doc_id in self._documents:
-            log.error("Document  id %s found in DataStore", doc_id)
-            raise KeyError(f"Document id '{doc_id}' found in DataStore")
+            _error_msg = f"Detected an error when adding document to the data store. Document {doc_id} already present."
+            log.error(_error_msg)
+            raise KeyError(_error_msg)
         self._documents[doc_id] = document
 
     def has_document(self, doc_id: str) -> bool:
@@ -45,13 +60,13 @@ def has_document(self, doc_id: str) -> bool:
 
     def get_document(self, doc_id: str) -> Optional[Document]:
         """
-        Returns Document or None.
+        Returns the Document with the given ID, or None if not found.
         """
         return self._get_document(doc_id)
 
     def get_documents(self) -> List[Document]:
         """
-        Returns a list of Document.
+        Returns a list of Document objects."
         """
         return list(self._documents.values())
 
@@ -69,7 +84,7 @@ def add_query(self, query: str, doc_id: str | None = None) -> str:
             return query_id
 
         # new query rating context
-        context = QueryRatingContext(query, doc_id)
+        context = QueryRatingContext(query=query, doc_id=doc_id)
         query_id = context.get_query_id()
         self._queries_by_id[query_id] = context
         self._query_text_to_query_id[query] = query_id
@@ -92,23 +107,98 @@ def add_rating_score(self, query_id: str, doc_id: str, rating_score: int) -> Non
         Adds rating score associated with the given doc_id and query_id or raises KeyError
         if the query_id is not found.
         """
-        context = self._get_query_rating_context_by_id(query_id)
-
+        context: QueryRatingContext = self._get_query_rating_context_by_id(query_id)
         context.add_rating_score(doc_id, rating_score)
         self._queries_by_id[query_id] = context
 
     def get_rating_score(self, query_id: str, doc_id: str) -> int:
         """
         Returns the rating score for the given (query_id, doc_id) pair or raises KeyError if the query_id is not found.
         """
-        context = self._get_query_rating_context_by_id(query_id)
+        context: QueryRatingContext = self._get_query_rating_context_by_id(query_id)
         return context.get_rating_score(doc_id)
 
     def has_rating_score(self, query_id: str, doc_id: str) -> bool:
         """
         Returns True if the (query_id, doc_id) pair has a rating score (i.e. != -1) or raises KeyError
         if query_id is not found.
         """
-        context = self._get_query_rating_context_by_id(query_id)
+        context: QueryRatingContext = self._get_query_rating_context_by_id(query_id)
         return context.has_rating_score(doc_id)
 
+    @staticmethod
+    def ensure_tmp_file_exists() -> Path:
+        """Checks if a file exists on disk and returns its path or create the parent folder.
+        Resolves a given path and logs a warning if the file does not exist.
+
+        Returns:
+            The resolved path as a Path object.
+        """
+        path = Path(TMP_FILE)
+        if not path.exists():
+            path.parent.mkdir(parents=True, exist_ok=True)
+            log.debug(f'Previous file not found in DataStore: {path}')
+        return path
+
+    def save_tmp_file_content(self) -> None:
+        """Saves the current state to a file on disk serializing queries, ratings, and optionally documents.
+
+        Args:
+            filepath: The path to the file where the data will be saved.
+                      If None, a default path is used.
+        """
+        path = self.ensure_tmp_file_exists()
+
+        def default_serializer(obj):
+            """Default function to handle non-serializable objects"""
+            if isinstance(obj, QueryRatingContext):
+                return obj.to_dict()
+            elif isinstance(obj, Document): # more generall: from pydantic import BaseModel
+                return obj.model_dump()
+            else:
+                # Convert to string as fallback
+                return str(obj)
+
+        data = {
+            "queries": self._queries_by_id,
+            "documents": self._documents,
+        }
+
+        # save the content to a file
+        with path.open("w", encoding="utf-8") as f:
+            json.dump(data, f, indent=2, ensure_ascii=False, default=default_serializer)
+
+
+    def load_tmp_file_content(self) -> None:
+        """Loads state from a file on disk loading queries, ratings, and documents from a unified
+        JSON file on disk.
+
+        Args:
+            filepath: The path to the file to load data from. If None, a
+                      default path is used.
+            clear: If True, clears existing data before loading.
+                   Defaults to None.
+        """
+
+        self._documents.clear()
+        self._queries_by_id.clear()
+        self._query_text_to_query_id.clear()
+
+        filepath = self.ensure_tmp_file_exists()
+
+        if not filepath.exists():
+            log.info(f"No datastore file yet at {filepath}, starting fresh.")
+            return
+
+        with filepath.open("r", encoding="utf-8") as f:
+            file_content = json.load(f)
+
+        queries: Dict[str, Dict[str, Any]] = file_content.get("queries", {})
+        for query_id, context_dict in queries.items():
+            context: QueryRatingContext = QueryRatingContext.from_dict(context_dict)
+            self._queries_by_id[query_id] = context
+            self._query_text_to_query_id[context.get_query_text()] = query_id
+
+        documents = file_content.get("documents", {})
+        for doc_id, doc_data in documents.items():
+            self.add_document(doc_id, Document.model_validate(doc_data))

rre-dataset-generator/src/writers/abstract_writer.py

@@ -1,8 +1,6 @@
 from abc import ABC, abstractmethod
 from typing import List, Tuple
 from pathlib import Path
-
-from src.model.query_rating_context import QueryRatingContext
 from src.search_engine.data_store import DataStore
 
 
@@ -28,9 +26,9 @@ def _get_queries_with_ratings(self) -> List[Tuple[str, str, int]]:
         """
         result = []
         for query_ctx in self.datastore.get_queries():
-            query_text = query_ctx.get_query()
+            query_text = query_ctx.get_query_text()
             for doc_id in query_ctx.get_doc_ids():
-                rating = query_ctx.get_rating_score(doc_id)
-                if rating != QueryRatingContext.DOC_NOT_RATED:
+                if query_ctx.has_rating_score(doc_id):
+                    rating = query_ctx.get_rating_score(doc_id)
                     result.append((query_text, doc_id, rating))
         return result