from functools import reduce import weaviate as wvt import datetime, os from threading import Lock from time import perf_counter, sleep from weaviate.classes.config import Configure, Property, DataType from weaviate.collections.classes.grpc import MetadataQuery from weaviate.collections.collection import Collection from weaviate.classes.init import AdditionalConfig, Timeout from weaviate.classes.query import Filter from weaviate.config import AdditionalConfig from ..utils.logging import get_logger from ..config import config logger = get_logger("weaviate_service") _get_collection_name = lambda lang: f'{config.weaviate.WEAVIATE_COLLECTION_BASENAME}_{lang}' _collection_names = [_get_collection_name(lang) for lang in config.get('AVAILABLE_LANGUAGES')] def _default_properties() -> list[Property]: return [ Property(name='body', data_type=DataType.TEXT), Property(name='chunk_id', data_type=DataType.TEXT), Property(name='document_id', data_type=DataType.TEXT), Property(name='programs', data_type=DataType.TEXT_ARRAY), Property(name='source', data_type=DataType.TEXT), Property(name='date', data_type=DataType.DATE), ] class WeaviateService: """ Provides an interface for interacting with the Weaviate vector database. Handles initialization, data import, and hybrid queries. """ _instance = None _init_lock = Lock() def __new__(cls): if cls._instance is None: with cls._init_lock: if cls._instance is None: cls._instance = super().__new__(cls) return cls._instance def __init__(self) -> None: """ Initialize the Weaviate service. """ if hasattr(self, '_initialized'): return self._connection_type = 'local' if config.weaviate.LOCAL_DATABASE else 'cloud' self._client = None self._client_lock = Lock() # Some parameters to ensure that the connection will not be closed # during long pauses in conversations self._last_query_time = perf_counter() self._idle_timeout = 25 * 60 self._initialized = True # Initialize the client for the first time logger.info("Initializing Weaviate service...") try: self._init_client() logger.info("Weaviate service initialized successfully") except Exception as e: logger.error(f"Failed to initialize Weaviate service: {e}") raise e def _init_client(self) -> wvt.WeaviateClient: """ Initializes the weaviate client with additional configuration. Performs a warm-up querying to speed-up the subsequent calls. Returns: configured Weaviate client instance on successfull connection. Raises: WeaviateConnectionError of the last failed connection if connection fails after 3 retires. """ # Returns the client if it hasn't been idling for too long if self._client is not None: time_since_query = perf_counter() - self._last_query_time if time_since_query < self._idle_timeout: return self._client # The connection might be closed, clients has to be reconnected logger.warning(f"Client has been idling for too long. Reconnecting to prevent server-side closure...") try: self._client.close() except Exception as _: pass self._client = None # Client initialization with self._client_lock: if self._client: return self._client retries = 0 last_exception: Exception = None while retries < 3: try: if config.weaviate.LOCAL_DATABASE: self._client = wvt.connect_to_local() break self._client = wvt.connect_to_weaviate_cloud( cluster_url=config.weaviate.CLUSTER_URL, auth_credentials=config.weaviate.WEAVIATE_API_KEY, additional_config=AdditionalConfig( timeout=Timeout( init=config.weaviate.INIT_TIMEOUT, query=config.weaviate.QUERY_TIMEOUT, insert=config.weaviate.INSERT_TIMEOUT, ), skip_init_checks=False, ), headers={ "X-HuggingFace-Api-Key": config.weaviate.HUGGING_FACE_API_KEY, }, ) # Warm-up query logger.info("Running warm-up query to initialize server...") try: collection = _get_collection_name(config.get('AVAILABLE_LANGUAGES')[0]) self._client.collections.exists(collection) logger.info("Warm-up finished - server is ready!") except Exception as warmup_err: logger.warning(f"Warm-up query failed (non-critical): {warmup_err}") break except Exception as e: last_exception = e logger.warning(f"Failed to establish connection on try {retries}: {e}") retries += 1 sleep(1) if retries == 3: logger.error(f"Failed to establish connection after 3 retries!") raise last_exception logger.info(f"Successully connected to the {self._connection_type} weaviate database") self._last_query_time = perf_counter() return self._client def _select_collection(self, lang: str) -> tuple[Collection, str]: """ Select a language-specific collection as the active working collection. Args: lang (str): Acceptable language code. Raises: weaviate.exceptions.WeaviateConnectionError: If the specified language collection does not exist. """ if lang not in config.get('AVAILABLE_LANGUAGES'): logger.error(f"No collection for language '{lang}' was found in the database") return None, '' collection_name = _get_collection_name(lang) logger.debug(f"Using collection {collection_name}") client = self._init_client() return client.collections.use(collection_name), collection_name def batch_import(self, data_rows: list, lang: str) -> list: """ Perform a batch import of multiple objects into the current collection. Args: data_rows (list): List of dictionaries representing the data rows to import. lang (str, optional): Language collection to use. If not provided, uses the current one. Returns: list[dict]: List of failed imports with error details, if any. Raises: If no active collection is available or a connection error was catched. """ collection, collection_name = self._select_collection(lang) if collection is None: logger.error("No working collection selected!") return [] import_errors = [] logger.info(f"Batch importing {len(data_rows)} rows into {collection_name}") try: with self._client_lock: with collection.batch.fixed_size(batch_size=100, concurrent_requests=2) as batch: for idx, data_row in enumerate(data_rows): try: batch.add_object(properties=data_row) except Exception as e: import_errors.append({'index': idx, 'chunk_id': data_row['chunk_id'], 'error': str(e)}) if idx % 20 == 0 and idx > 0: if batch.number_errors > 0: logger.info(f"Failed imports at index {idx}: {batch.number_errors}") self._last_query_time = perf_counter() logger.info(f"Batch import finished. Total errors: {len(import_errors)}") except Exception as e: if 'connection' in str(e).lower(): logger.error(f"Connection error during batch import: {e}") self._client = None raise e return import_errors @staticmethod def _create_property_filter(prop, values) -> Filter: match prop: case 'programs': return Filter.by_property('programs').contains_any(values) case 'source': return Filter.by_property('source').contains_any(values) \ if isinstance(values, list) else Filter.by_property('source').equal(values) case _: return None def delete_chunks(self, lang: str, property_filters: dict[str, any] = None) -> int: """ Delete all chunks from the specified collection that match given property filters. Args: lang (str): Language collection to use. property_filters (dict[str, any]): Key-value pairs for filtering. Returns: int: Number of deleted objects (if available, else -1). """ retry_count = 0 max_retries = 2 filters = [self._create_property_filter(prop, values) for prop, values in property_filters.items()] if property_filters else None if filters: filters = [f for f in filters if f is not None] filters = reduce(lambda f1, f2: f1 & f2, filters) if filters else None while retry_count < max_retries: try: collection, collection_name = self._select_collection(lang) if collection is None: logger.error("No working collection selected!") return 0 logger.info(f"Deleting chunks from {collection_name} with filters={property_filters}") with self._client_lock: result = collection.data.delete_many( where=filters ) self._last_query_time = perf_counter() deleted = getattr(result, "objects_deleted", None) if deleted is None: logger.info("Deletion executed (count not returned by client)") return -1 logger.info(f"Deleted {deleted} objects") return deleted except Exception as e: if any(err_type in str(e).lower() for err_type in ['reset', 'closed', 'grpc', 'unavailable']): retry_count += 1 logger.warning(f"Connection error during deletion: {e}. Retrying...") if retry_count == max_retries: raise e else: raise e def ping(self, lang: str) -> dict: try: collection, _ = self._select_collection(lang) with self._client_lock: collection.query.hybrid("health check query") return { 'status': 'OK' } except Exception as e: return { 'status': 'ERROR', 'error': e } def query(self, query: str, lang: str, property_filters: dict[str] = None, limit: int = 5) -> dict: """ Execute a hybrid semantic and keyword query against the active collection with automatic reconnection on idle timeout. Args: query (str): The query string. lang (str, optional): Language collection to use. If not provided, uses the current one. property_filters (dict[str, any]): Key-value pairs for metadata filtering. Keys correspond to document properties (e.g., 'program', 'topic'), and values are the required matches. Multiple filters are combined using logical AND. limit (int, optional): Maximum number of results to return. Defaults to 5. Returns: tuple: A tuple containing the query response and elapsed time. Raises: weaviate.exceptions.WeaviateConnectionError: If no active collection is available. """ retry_count = 0 max_retries = 2 filters = [self._create_property_filter(prop, values) for prop, values in property_filters.items()] if property_filters else None if filters: filters = [f for f in filters if f is not None] filters = reduce(lambda f1, f2: f1 & f2, filters) if filters else None while retry_count < max_retries: try: collection, collection_name = self._select_collection(lang) if collection is None: logger.error("No working collection selected upon starting of the querying!") return [], 0 logger.info(f"Querying collection {collection_name}") query_start_time = perf_counter() with self._client_lock: try: resp = collection.query.hybrid( query=query, filters=filters, limit=limit, return_metadata=MetadataQuery.full() ) except Exception as hybrid_err: if not self._should_fallback_to_bm25(hybrid_err): raise hybrid_err logger.warning( "Hybrid query failed during remote vectorization. " "Falling back to BM25 keyword retrieval: %s", hybrid_err, ) resp = collection.query.bm25( query=query, filters=filters, limit=limit, return_metadata=MetadataQuery.full() ) elapsed = perf_counter() - query_start_time self._last_query_time = perf_counter() logger.info(f"Querying retrieved {len(resp.objects)} objects in {elapsed:3.2f} seconds") return (resp, elapsed) except Exception as e: if any(err_type in str(e).lower() for err_type in ['reset', 'closed', 'grpc', 'unavailable']): retry_count += 1 logger.warning(f"Connection error detected: {e}. Retrying...") if retry_count == max_retries: raise e else: # Probably not a server issue raise e @staticmethod def _should_fallback_to_bm25(error: Exception) -> bool: error_text = str(error).lower() return ( "remote client vectorize" in error_text or "vectorize" in error_text and "401" in error_text or "invalid username or password" in error_text ) def _load_properties(self) -> list[Property]: properties = {} properties_file = os.path.join(config.weaviate.PROPERTIES_PATH, 'properties.yaml') if not os.path.exists(properties_file): logger.warning( f"Optional file 'properties.yaml' is missing on path: {properties_file}. " "Falling back to built-in default properties." ) return _default_properties() try: import yaml with open(properties_file, 'r') as stream: properties = yaml.safe_load(stream) except ModuleNotFoundError: logger.warning( "PyYAML is not installed. Falling back to built-in default properties " "for Weaviate collection creation." ) return _default_properties() except Exception as e: logger.error(f"Failed to load properties from path {properties_file}: {e}") raise e if not properties: logger.warning("properties.yaml is empty. Falling back to built-in default properties.") return _default_properties() final_properties = [] for name, params in properties.items(): try: data_type = params.get('data_type', '') dtype = DataType(data_type) except Exception as e: logger.error(f"Nonexistent datatype {data_type}") raise e final_properties.append(Property( name=name, data_type=dtype, index_filterable=params.get('filterable', True), index_searchable=params.get('searchable', True), skip_vectorization=params.get('skip_vectorization', False), )) return final_properties def _create_collections(self): """ Create and initialize language-specific collections. Creates collections for all available languages with vector configuration. """ properties = self._load_properties() try: client = self._init_client() logger.info('Attempting collections creation...') vector_config = ( Configure.Vectors.text2vec_transformers() if config.weaviate.LOCAL_DATABASE else Configure.Vectors.text2vec_huggingface( name='hsg_rag_embeddings', source_properties=['body'], model=config.processing.EMBEDDING_MODEL, ) ) successful_creations = 0 with self._client_lock: for collection_name in _collection_names: try: client.collections.create( name=collection_name, properties=properties, vector_config=vector_config ) logger.info(f"Created collection {collection_name}") successful_creations += 1 except Exception as e: logger.error(f"Failed to create collection '{collection_name}': {e}") self._last_query_time = perf_counter() if successful_creations == len(_collection_names): logger.info('All collections successfully instantiated') else: logger.warning(f"Only {successful_creations}/{len(_collection_names)} collections created") except Exception as e: logger.error(f"Collections creation failed: {e}") self._client = None raise e def _delete_collections(self): """ Delete all existing collections from the database. Also removes the hash file if it exists. """ try: client = self._init_client() logger.info("Initiating deletion of stored collections...") deleted_count = 0 with self._client_lock: for collection_name in _collection_names: try: if client.collections.exists(collection_name): client.collections.delete(collection_name) logger.info(f"Deleted collection {collection_name}") deleted_count += 1 else: logger.warning(f"Collection {collection_name} does not exist") except Exception as e: logger.error(f"Failed to delete collection {collection_name}: {e}") self._last_query_time = perf_counter() logger.info(f"Deleted {deleted_count}/{len(_collection_names)} collections") except Exception as e: logger.error(f"Collections deletion failed: {e}") self._client = None raise e def _reset_collections(self): self._delete_collections() self._create_collections() def _collect_chunk_ids(self) -> dict: client = self._init_client() try: ids = [] with self._client_lock: for c in client.collections.list_all(simple=False): coll = client.collections.get(c) for obj in coll.iterator(): ids.append(obj.properties['chunk_id']) return ids except Exception as e: logger.error(f"Failed to collect chunk ids: {e}") raise e def _extract_data(self) -> dict: client = self._init_client() try: schema = [] objects = {} with self._client_lock: for c in client.collections.list_all(simple=False): coll = client.collections.get(c) cfg = coll.config.get().to_dict() schema.append(cfg) objects[c] = [] for obj in coll.iterator(include_vector=True): objects[c].append({ "uuid": obj.uuid, "properties": obj.properties, "vector": obj.vector, }) return { 'schema': schema, 'objects': objects, } except Exception as e: logger.error(f"Failed to extract data from database: {e}") raise e def _create_backup(self) -> str: """ Create a backup of the current database state and stores it under selected backup provider. Returns: backup id of the created backup. """ try: if not config.weaviate.BACKUP_METHOD: raise ValueError('Backup method is not selected!') if config.weaviate.BACKUP_METHOD not in config.weaviate.BACKUP_METHODS: raise ValueError(f"Selected backup method 'config.weaviate.BACKUP_METHODS' is not supported!") if not config.weaviate.BACKUP_PATH: raise ValueError("Backup directory is not set!") os.makedirs(config.weaviate.BACKUP_PATH, exist_ok=True) backup_id = f"backup_{datetime.datetime.now().strftime('%Y%m%d%H%M%S%f')}" logger.info(f"Initiating backup creation for {self._connection_type} database...") match config.weaviate.BACKUP_METHOD: case 'manual': import json backup_path = os.path.join(config.weaviate.BACKUP_PATH, backup_id) os.makedirs(backup_path) db_data = self._extract_data() data_backup = { 'creation_date': datetime.datetime.now().isoformat(), } schema_backup_path = os.path.join(backup_path, 'schema.json') with open(schema_backup_path, 'w', encoding='utf-8') as f: json.dump(db_data['schema'], f, indent=2, default=str) objects_backup_path = os.path.join(backup_path, 'objects.json') with open(objects_backup_path, 'w', encoding='utf-8') as f: json.dump(db_data['objects'], f, indent=2, default=str) data_backup_path = os.path.join(backup_path, 'data.json') with open(data_backup_path, 'w', encoding='utf-8') as f: json.dump(data_backup, f, indent=2, default=str) case 's3': client = self._init_client() with self._client_lock: client.backup.create( backup_id=backup_id, backend="s3", include_collections=_collection_names, wait_for_completion=True, ) case _: raise NotImplementedError() self._last_query_time = perf_counter() logger.info(f"Backup '{backup_id}' created successfully") return backup_id except Exception as e: logger.error(f"Backup creation failed: {e}") raise e def _restore_backup(self, backup_id: str): """ Restore the database state from a backup. Restores specified collections from backup. Args: backup_id: ID of the backup to restore from Raises: Exception if backup restoration fails """ self._delete_collections() try: if not config.weaviate.BACKUP_METHOD: raise ValueError('Backup method is not selected!') if config.weaviate.BACKUP_METHOD not in config.weaviate.BACKUP_METHODS: raise ValueError(f"Selected backup method 'config.weaviate.BACKUP_METHODS' is not supported!") if not config.weaviate.BACKUP_PATH: raise ValueError("Backup directory is not set!") os.makedirs(config.weaviate.BACKUP_PATH, exist_ok=True) backup_path = os.path.join(config.weaviate.BACKUP_PATH, backup_id) if not os.path.exists(backup_path): raise RuntimeError(f"Directory for backup 'backup_id' does not exist in the backup directory!") schema_backup_path = os.path.join(backup_path, 'schema.json') if not os.path.exists(schema_backup_path): raise RuntimeError(f"Schema backup is missing in the backup directory!") objects_backup_path = os.path.join(backup_path, 'objects.json') if not os.path.exists(objects_backup_path): raise RuntimeError(f"Objects backup is missing in the backup directory!") client = self._init_client() logger.info(f"Initiating restoration from backup '{backup_id}' for {self._connection_type} database...") with self._client_lock: match config.weaviate.BACKUP_METHOD: case 'manual': import json with open(schema_backup_path) as f: schemas = json.load(f) for cfg in schemas: client.collections.create_from_dict(cfg) with open(objects_backup_path) as f: data = json.load(f) for name, objs in data.items(): logger.info(f"Restoring collection '{name}' with {len(objs)} objects...") coll = client.collections.get(name) with coll.batch.dynamic() as batch: for o in objs: o['properties']['date'] = o['properties']['date'] \ .replace(" ", "T").replace("+00:00", "Z") batch.add_object( uuid=o["uuid"], properties=o["properties"], vector=o["vector"] ) logger.info(f"Collection '{name}' restored successfully") case 's3': client.backup.restore( backup_id=backup_id, backend="s3", wait_for_completion=True, roles_restore="all", users_restore="all", ) case _: raise NotImplementedError() self._last_query_time = perf_counter() logger.info(f"Backup '{backup_id}' restored successfully") except Exception as e: error_msg = str(e).lower() if 'connection' in error_msg: logger.error(f"Connection error during backup restore: {e}. Will reconnect on next operation.") self._client = None logger.error(f"Backup restoration failed: {e}") raise e def _checkhealth(self) -> bool: """ Check the connectivity and health status of the Weaviate database. Verifies: - Connection to the database - Database metadata and version - Existence of all expected collections - Module availability Returns: True if all health checks pass, False otherwise """ try: client = self._init_client() # Check basic connectivity is_connected = False with self._client_lock: is_connected = client.is_connected() connection_status = "✓ OK" if is_connected else "✗ ERROR" logger.info(f"Connection to {self._connection_type} database: {connection_status}") if not is_connected: logger.error("Database connection check failed") return False # Get and log metadata try: with self._client_lock: metainfo = client.get_meta() # Format module information modules = metainfo.get('modules', {}) modules_list = list(modules.keys()) if isinstance(modules, dict) else modules modules_str = ', '.join(str(m) for m in modules_list) if modules_list else 'None' # Truncate long module strings for logging if len(modules_str) > 50: modules_str = modules_str[:47] + '...' # Log connection details if config.weaviate.LOCAL_DATABASE: logger.info( f"Database metadata: " f"HOSTNAME={metainfo.get('hostname', 'unknown')}, " f"VERSION={metainfo.get('version', 'unknown')}, " f"MODULES={modules_str}" ) else: logger.info( f"Database metadata: " f"VERSION={metainfo.get('version', 'unknown')}, " f"MODULES={modules_str}" ) except Exception as e: logger.warning(f"Could not retrieve database metadata: {e}") # Check collection existence all_collections_exist = True with self._client_lock: for collection_name in _collection_names: try: exists = client.collections.exists(collection_name) status = "✓ OK" if exists else "✗ MISSING" logger.info(f"Collection '{collection_name}': {status}") if not exists: all_collections_exist = False except Exception as e: logger.error(f"Error checking collection '{collection_name}': {e}") all_collections_exist = False # Update last health check time self._last_query_time = perf_counter() # Log overall health status if is_connected and all_collections_exist: logger.info("✓ Database health check PASSED - All systems operational") return True else: logger.warning("✗ Database health check FAILED - Some issues detected") return False except Exception as e: error_msg = str(e).lower() if 'connection' in error_msg: logger.error(f"Connection error during health check: {e}. Will reconnect on next operation.") self._client = None logger.error(f"Health check failed: {e}") return False def parse_arguments(): """ Parse command-line arguments for managing Weaviate collections. Returns: argparse.Namespace: Parsed command-line arguments. """ import argparse parser = argparse.ArgumentParser( description='Weaviate database management utility' ) group = parser.add_mutually_exclusive_group() group.add_argument( '-dc', "--delete_collections", action='store_true', help='Delete all collections from the database' ) group.add_argument( '-cc', "--create_collections", action='store_true', help='Initialize collections for different language contents' ) group.add_argument( '-rc', "--redo_collections", action='store_true', help='Delete and recreate all collections' ) group.add_argument( '-ch', "--checkhealth", action='store_true', help='Check database connection and collection existence' ) group.add_argument( '-cb', "--create_backup", action='store_true', help='Create a backup of the current database state' ) group.add_argument( '-rb', "--restore_backup", type=str, metavar='BACKUP_ID', help='Restore database from a backup (provide backup_id)' ) return parser.parse_args() if __name__ == "__main__": args = parse_arguments() service = WeaviateService() if args.create_backup: service._create_backup() if args.restore_backup: service._restore_backup(args.restore_backup) if any([args.delete_collections, args.redo_collections]): service._delete_collections() if any([args.create_collections, args.redo_collections]): service._create_collections() if any([args.checkhealth, args.create_collections, args.redo_collections]): service._checkhealth()