refactor: remove Optional from client attribute and streamline class logic

Author: mehmetcanay

datastew/repository/weaviate.py

@@ -74,15 +74,15 @@ def __init__(
         self.terminology_schema = terminology_schema
         self.concept_schema = concept_schema
         self.mapping_schema = mapping_schema
-        self.client: Optional[WeaviateClient] = None
+        self.client: WeaviateClient
         self.headers = None
         if self.use_weaviate_vectorizer:
             if huggingface_key:
                 self.headers = {"X-HuggingFace-Api-Key": huggingface_key}
         if self.mode == "memory":
-            self._connect_to_memory(path, http_port, grpc_port)
+            self.client = self._connect_to_memory(path, http_port, grpc_port)
         elif self.mode == "remote":
-            self._connect_to_remote(path, port)
+            self.client = self._connect_to_remote(path, port)
         else:
             raise ValueError(f"Repository mode {mode} is not defined. Use either disk or remote.")
 
@@ -113,13 +113,10 @@ def get_iterator(self, collection: Literal["Concept", "Mapping", "Terminology"])
         """Retrieves an iterator for the specified collection from the client.
 
         :param collection: The collection type to retrieve the iterator for.
-        :raises ValueError: If the client is not initialized or is invalid.
         :raises ValueError: If the specified collection is not supported.
         :return: An iterator for the specified collection from the client's collection,
             with vectors and references as needed.
         """
-        if not self.client:
-            raise ValueError("Client is not initialized or is invalid.")
         if collection == "Concept":
             return_references = QueryReference(link_on="hasTerminology")
         elif collection == "Mapping":
@@ -138,12 +135,9 @@ def get_all_sentence_embedders(self) -> List[str]:
         `self.use_weaviate_vectorizer` is True, it retrieves the names from the vector keys in the embeddings returned
         by an object in the collection.
 
-        :raises ValueError: If the client is not initialized.
         :raises RuntimeError: If there is an issue fetching sentence embedders or vector configurations.
         :return: A list of sentence embedder names.
         """
-        if not self.client:
-            raise ValueError("Client is not initialized or is invalid.")
         sentence_embedders = set()
         mapping_collection = self.client.collections.get("Mapping")
         try:
@@ -164,8 +158,6 @@ def get_all_sentence_embedders(self) -> List[str]:
         return list(sentence_embedders)
 
     def get_concept(self, concept_id: str) -> Concept:
-        if not self.client:
-            raise ValueError("Client is not initialized or is invalid.")
         try:
             if not self._concept_exists(concept_id):
                 raise RuntimeError(f"Concept {concept_id} does not exists")
@@ -190,8 +182,6 @@ def get_concept(self, concept_id: str) -> Concept:
         return concept
 
     def get_concepts(self, limit: int = 100, offset: int = 0, terminology_name: Optional[str] = None) -> Page[Concept]:
-        if not self.client:
-            raise ValueError("Client is not initialized or is invalid.")
         try:
             concept_collection = self.client.collections.get("Concept")
 
@@ -234,8 +224,6 @@ def get_concepts(self, limit: int = 100, offset: int = 0, terminology_name: Opti
         return Page[Concept](items=concepts, limit=limit, offset=offset, total_count=total_count)
 
     def get_terminology(self, terminology_name: str) -> Terminology:
-        if not self.client:
-            raise ValueError("Client is not initialized or is invalid.")
         try:
             if not self._terminology_exists(terminology_name):
                 raise RuntimeError(f"Terminology {terminology_name} does not exists")
@@ -252,8 +240,6 @@ def get_terminology(self, terminology_name: str) -> Terminology:
         return terminology
 
     def get_all_terminologies(self) -> List[Terminology]:
-        if not self.client:
-            raise ValueError("Client is not initialized or is invalid.")
         terminologies = []
         try:
             terminology_collection = self.client.collections.get("Terminology")
@@ -281,16 +267,12 @@ def get_mappings(
             `use_weaviate_vectorizer` is `True`, defaults to None
         :param limit: The maximum number of mappings to return, defaults to 1000
         :param offset: The number of mappings to skip before returning results, defaults to 0
-        :raises ValueError: If the client is not initialized or is invalid.
         :raises ValueError: If the terminology is not found.
         :raises ValueError: If the sentence embedder is not found.
         :raises ValueError: If `sentence_embedder` is `None` and `use_weaviate_vectorizer` is `True`.
         :raises RuntimeError: If the fetch operation fails.
         :return: A page object containing a list of Mapping objects, along with pagination details.
         """
-        if not self.client:
-            raise ValueError("Client is not initialized or is invalid.")
-
         mappings = []  # List to store fetched mappings
         filters = None  # List to store filters for query
         target_vector = True  # Whether to include vectors in the response
@@ -410,9 +392,6 @@ def get_closest_mappings(
         :raises RuntimeError: If the fetch operation fails
         :return: A list of Mapping or MappingResult objects based on whether similarity scores are included.
         """
-        if not self.client:
-            raise ValueError("Client is not initialized or is invalid.")
-
         mappings = []  # List to store fetched mappings
         filters = None
         target_vector = None
@@ -511,8 +490,6 @@ def get_closest_mappings(
         return mappings
 
     def store(self, model_object_instance: Union[Terminology, Concept, Mapping]):
-        if not self.client:
-            raise ValueError("Client is not initialized or is invalid.")
         try:
             if isinstance(model_object_instance, Terminology):
                 properties = {"name": model_object_instance.name}
@@ -590,14 +567,10 @@ def import_from_jsonl(
         :param jsonl_path: Path to the JSONL file.
         :param object_type: Literal specifying the object type, must be "terminology", "concept", or "mapping".
         :param chunk_size: The number of items to process in each batch, defaults to 100.
-        :raises ValueError: If the client is not initialized or is invalid.
         :raises ValueError: If 'id' or 'properties' is missing in a JSON object.
         :raises ValueError: If there is an invalid JSON object.
         :raises RuntimeError: If an unexpected error occurs during import.
         """
-        if not self.client:
-            raise ValueError("Client is not initialized or is invalid.")
-
         try:
             collection = self.client.collections.get(object_type.capitalize())
             chunk = []
@@ -635,18 +608,10 @@ def close(self):
         self.shut_down()
 
     def shut_down(self):
-        if not self.client:
-            raise ValueError("Client is not initialized or is invalid.")
         self.client.close()
 
     def clear_all(self):
-        """Deletes all data and schema classes (Mapping, Concept, Terminology) and re-creates them.
-
-        :raises ValueError: If the Weaviate client is not initialized.
-        """
-        if not self.client:
-            raise ValueError("Client is not initialized or is invalid.")
-
+        """Deletes all data and schema classes (Mapping, Concept, Terminology) and re-creates them."""
         for schema in [self.mapping_schema, self.concept_schema, self.terminology_schema]:
             class_name = schema.schema["class"]
             if self.client.collections.exists(class_name):
@@ -665,8 +630,6 @@ def clear_all(self):
         self._create_schema_if_not_exists(self.mapping_schema.schema)
 
     def _sentence_embedder_exists(self, name: str) -> bool:
-        if not self.client:
-            raise ValueError("Client is not initialized or is invalid.")
         try:
             mapping = self.client.collections.get("Mapping")
             if not self.use_weaviate_vectorizer:
@@ -682,8 +645,6 @@ def _sentence_embedder_exists(self, name: str) -> bool:
             raise RuntimeError(f"Failed to check if sentence embedder exists: {e}")
 
     def _terminology_exists(self, name: str) -> bool:
-        if not self.client:
-            raise ValueError("Client is not initialized or is invalid.")
         try:
             terminology = self.client.collections.get("Terminology")
             response = terminology.query.fetch_objects(filters=Filter.by_property("name").equal(name))
@@ -695,8 +656,6 @@ def _terminology_exists(self, name: str) -> bool:
             raise RuntimeError(f"Failed to check if terminology exists: {e}")
 
     def _concept_exists(self, concept_id: str) -> bool:
-        if not self.client:
-            raise ValueError("Client is not initialized or is invalid.")
         try:
             concept = self.client.collections.get("Concept")
             response = concept.query.fetch_objects(filters=Filter.by_property("conceptID").equal(concept_id))
@@ -717,8 +676,6 @@ def _mapping_exists(self, mapping: Mapping) -> bool:
             exists, otherwise `False`.
         """
         try:
-            if not self.client:
-                raise ValueError("Client is not initialized or is invalid.")
             # Check if the concept exists first
             concept_properties = {
                 "conceptID": mapping.concept.concept_identifier,
@@ -761,8 +718,6 @@ def _uuid_exists(self, collection_name: str, uuid: str) -> bool:
         :return: True if the object with the given UUID exists, False otherwise.
         """
         try:
-            if not self.client:
-                raise ValueError("Client is not initialized or is invalid.")
             collection = self.client.collections.get(collection_name)
             obj = collection.query.fetch_object_by_id(uuid)
             return obj is not None
@@ -788,8 +743,6 @@ def _create_schema_if_not_exists(self, schema):
         :raises RuntimeError: If there is an issue checking for or creating the schema in Weaviate, such as connection
             error.
         """
-        if not self.client:
-            raise ValueError("Client is not initialized or is invalid.")
         references = None
         vectorizer_config = None
         class_name = schema["class"]
@@ -818,28 +771,26 @@ def _is_port_in_use(port) -> bool:
         with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
             return s.connect_ex(("localhost", port)) == 0
 
-    def _connect_to_memory(self, path: str, http_port: int, grpc_port: int):
+    def _connect_to_memory(self, path: str, http_port: int, grpc_port: int) -> WeaviateClient:
         try:
             if path is None:
                 raise ValueError("Path must be provided for disk mode.")
             if self._is_port_in_use(http_port) and self._is_port_in_use(grpc_port):
-                if self.client:
-                    self.client.close()
-                self.client = weaviate.connect_to_local(port=http_port, grpc_port=grpc_port, headers=self.headers)
+                return weaviate.connect_to_local(port=http_port, grpc_port=grpc_port, headers=self.headers)
             else:
-                self.client = weaviate.connect_to_embedded(
+                return weaviate.connect_to_embedded(
                     persistence_data_path=path,
                     headers=self.headers,
                     environment_variables={"ENABLE_MODULES": "text2vec-ollama"},
                 )
         except Exception as e:
             raise ConnectionError(f"Failed to initialize Weaviate client: {e}")
 
-    def _connect_to_remote(self, path: str, port: int):
+    def _connect_to_remote(self, path: str, port: int) -> WeaviateClient:
         try:
             if path is None:
                 raise ValueError("Remote URL must be provided for remote mode.")
-            self.client = weaviate.connect_to_local(host=path, port=port, headers=self.headers)
+            return weaviate.connect_to_local(host=path, port=port, headers=self.headers)
         except Exception as e:
             raise ConnectionError(f"Failed to initialize Weaviate client: {e}")