Repository: stanford-oval/storm
File: knowledge_storm/utils.py (Python)

Code before:
from trafilatura import extract
from tqdm import tqdm
import pandas as pd

logging.getLogger("httpx").setLevel(logging.WARNING)  # Disable INFO logging for httpx.


def load_api_key(toml_file_path):
    try:
        with open(toml_file_path, 'r') as file:
            data = toml.load(file)
    except FileNotFoundError:
        print(f"File not found: {toml_file_path}", file=sys.stderr)
        return
    except toml.TomlDecodeError:
        print(f"Error decoding TOML file: {toml_file_path}", file=sys.stderr)
        return
    # Set environment variables
    for key, value in data.items():
        os.environ[key] = str(value)


def makeStringRed(message):
    return f"\033[91m {message}\033[00m"

class VectorStoreManager:
    @staticmethod
    def _check_create_collection(client: QdrantClient, collection_name: str, model: HuggingFaceEmbeddings):
        """
        Check if the Qdrant collection exists and create it if it does not.
        """
        if client is None:
            raise ValueError("Qdrant client is not initialized.")
        if client.collection_exists(collection_name=f"{collection_name}"):
            print(f"Collection {collection_name} exists. Loading the collection...")
            return Qdrant(
                client=client,
                collection_name=collection_name,
                embeddings=model,
            )
        else:
            print(f"Collection {collection_name} does not exist. Creating the collection...")
            # create the collection
            client.create_collection(
                collection_name=f"{collection_name}",
                vectors_config=models.VectorParams(size=1024, distance=models.Distance.COSINE),
            )
            return Qdrant(
                client=client,
                collection_name=collection_name,
                embeddings=model,

Review comment (suggestion):
Rename to `QdrantVectorStoreManager`. Also, please add a few comments to describe the functionality of this class and how it can be connected with `VectorRM`.

Code after:
from qdrant_client import QdrantClient, models
from tqdm import tqdm
from trafilatura import extract

logging.getLogger("httpx").setLevel(logging.WARNING)  # Disable INFO logging for httpx.


def load_api_key(toml_file_path):
    try:
        with open(toml_file_path, 'r') as file:
            data = toml.load(file)
    except FileNotFoundError:
        print(f"File not found: {toml_file_path}", file=sys.stderr)
        return
    except toml.TomlDecodeError:
        print(f"Error decoding TOML file: {toml_file_path}", file=sys.stderr)
        return
    # Set environment variables
    for key, value in data.items():
        os.environ[key] = str(value)


def makeStringRed(message):
    return f"\033[91m {message}\033[00m"


class QdrantVectorStoreManager:
    """
    Helper class for managing the Qdrant vector store, can be used with `VectorRM` in rm.py.
    
    Before you initialize `VectorRM`, call `create_or_update_vector_store` to create or update the vector store.
    Once you have the vector store, you can initialize `VectorRM` with the vector store path or the Qdrant server URL.
    """
    @staticmethod
    def _check_create_collection(client: QdrantClient, collection_name: str, model: HuggingFaceEmbeddings):
        """Check if the Qdrant collection exists and create it if it does not."""
        if client is None:
            raise ValueError("Qdrant client is not initialized.")
        if client.collection_exists(collection_name=f"{collection_name}"):
            print(f"Collection {collection_name} exists. Loading the collection...")
            return Qdrant(
                client=client,
                collection_name=collection_name,
                embeddings=model,
            )
        else:
            print(f"Collection {collection_name} does not exist. Creating the collection...")
            # create the collection
            client.create_collection(
                collection_name=f"{collection_name}",
                vectors_config=models.VectorParams(size=1024, distance=models.Distance.COSINE),