qdrantdb

langroid/vector_store/qdrantdb.py

QdrantDB(config=None)

Bases: VectorStore

Source code in langroid/vector_store/qdrantdb.py

def __init__(self, config: QdrantDBConfig | None = None):
    config = config if config is not None else QdrantDBConfig()
    super().__init__(config)
    self.config: QdrantDBConfig = config
    from qdrant_client import QdrantClient

    self.sparse_fastembed_model: Optional["SparseTextEmbedding"] = None
    if self.config.use_sparse_embeddings:
        if self.config.sparse_embedding_backend == "fastembed":
            try:
                from fastembed import SparseTextEmbedding
            except ImportError as e:
                raise ImportError(
                    "To use FastEmbed sparse embeddings, install "
                    "langroid with the [fastembed] extra: "
                    'pip install "langroid[fastembed]"'
                ) from e
            self.sparse_fastembed_model = SparseTextEmbedding(
                model_name=self.config.fastembed_sparse_model
            )
        else:
            try:
                from transformers import AutoModelForMaskedLM, AutoTokenizer
            except ImportError:
                raise ImportError(
                    """
                    To use sparse embeddings,
                    you must install langroid with the [transformers] extra, e.g.:
                    pip install "langroid[transformers]"
                    """
                )

            self.sparse_tokenizer = AutoTokenizer.from_pretrained(  # type: ignore
                self.config.sparse_embedding_model
            )
            self.sparse_model = AutoModelForMaskedLM.from_pretrained(
                self.config.sparse_embedding_model
            )
    self.host = config.host
    self.port = config.port
    load_dotenv()
    key = os.getenv("QDRANT_API_KEY")
    url = os.getenv("QDRANT_API_URL")
    if config.docker:
        if url is None:
            logger.warning(
                f"""The QDRANT_API_URL env variable must be set to use
                QdrantDB in local docker mode. Please set this
                value in your .env file.
                Switching to local storage at {config.storage_path}
                """
            )
            config.cloud = False
        else:
            config.cloud = True
    elif config.cloud and None in [key, url]:
        logger.warning(
            f"""QDRANT_API_KEY, QDRANT_API_URL env variable must be set to use
            QdrantDB in cloud mode. Please set these values
            in your .env file.
            Switching to local storage at {config.storage_path}
            """
        )
        config.cloud = False

    if config.cloud:
        self.client = QdrantClient(
            url=url,
            api_key=key,
            timeout=config.timeout,
        )
    elif config.storage_path == ":memory:":
        self.client = QdrantClient(location=":memory:")
    else:
        try:
            self.client = QdrantClient(
                path=config.storage_path,
            )
        except Exception as e:
            new_storage_path = config.storage_path + ".new"
            logger.warning(
                f"""
                Error connecting to local QdrantDB at {config.storage_path}:
                {e}
                Switching to {new_storage_path}
                """
            )
            self.client = QdrantClient(
                path=new_storage_path,
            )

    # Note: Only create collection if a non-null collection name is provided.
    # This is useful to delay creation of vecdb until we have a suitable
    # collection name (e.g. we could get it from the url or folder path).
    if config.collection_name is not None:
        self.create_collection(
            config.collection_name, replace=config.replace_collection
        )

clone()

Create an independent Qdrant client when running against Qdrant Cloud.

Source code in langroid/vector_store/qdrantdb.py

def clone(self) -> "QdrantDB":
    """Create an independent Qdrant client when running against Qdrant Cloud."""
    if not self.config.cloud:
        return self
    cloned = super().clone()
    assert isinstance(cloned, QdrantDB)
    return cloned

close()

Close the QdrantDB client and release any resources (e.g., file locks). This is especially important for local storage to release the .lock file.

Source code in langroid/vector_store/qdrantdb.py

def close(self) -> None:
    """
    Close the QdrantDB client and release any resources (e.g., file locks).
    This is especially important for local storage to release the .lock file.
    """
    if hasattr(self.client, "close"):
        # QdrantLocal has a close method that releases the lock
        self.client.close()
        logger.info(f"Closed QdrantDB connection for {self.config.storage_path}")

clear_all_collections(really=False, prefix='')

Clear all collections with the given prefix.

Source code in langroid/vector_store/qdrantdb.py

def clear_all_collections(self, really: bool = False, prefix: str = "") -> int:
    """Clear all collections with the given prefix."""

    if not really:
        logger.warning("Not deleting all collections, set really=True to confirm")
        return 0
    coll_names = [
        c for c in self.list_collections(empty=True) if c.startswith(prefix)
    ]
    if len(coll_names) == 0:
        logger.warning(f"No collections found with prefix {prefix}")
        return 0
    n_empty_deletes = 0
    n_non_empty_deletes = 0
    for name in coll_names:
        info = self.client.get_collection(collection_name=name)
        points_count = from_optional(info.points_count, 0)

        n_empty_deletes += points_count == 0
        n_non_empty_deletes += points_count > 0
        self.client.delete_collection(collection_name=name)
    logger.warning(
        f"""
        Deleted {n_empty_deletes} empty collections and
        {n_non_empty_deletes} non-empty collections.
        """
    )
    return n_empty_deletes + n_non_empty_deletes

list_collections(empty=False)

Returns:

Type	Description
List[str]	List of collection names that have at least one vector.

Parameters:

Name	Type	Description	Default
empty	bool	Whether to include empty collections.	False

Source code in langroid/vector_store/qdrantdb.py

def list_collections(self, empty: bool = False) -> List[str]:
    """
    Returns:
        List of collection names that have at least one vector.

    Args:
        empty (bool, optional): Whether to include empty collections.
    """

    colls = list(self.client.get_collections())[0][1]
    if empty:
        return [coll.name for coll in colls]
    counts = []
    for coll in colls:
        try:
            counts.append(
                from_optional(
                    self.client.get_collection(
                        collection_name=coll.name
                    ).points_count,
                    0,
                )
            )
        except Exception:
            logger.warning(f"Error getting collection {coll.name}")
            counts.append(0)
    return [coll.name for coll, count in zip(colls, counts) if (count or 0) > 0]

create_collection(collection_name, replace=False)

Create a collection with the given name, optionally replacing an existing collection if replace is True. Args: collection_name (str): Name of the collection to create. replace (bool): Whether to replace an existing collection with the same name. Defaults to False.

Source code in langroid/vector_store/qdrantdb.py

def create_collection(self, collection_name: str, replace: bool = False) -> None:
    """
    Create a collection with the given name, optionally replacing an existing
        collection if `replace` is True.
    Args:
        collection_name (str): Name of the collection to create.
        replace (bool): Whether to replace an existing collection
            with the same name. Defaults to False.
    """
    from qdrant_client.http.models import (
        CollectionStatus,
        Distance,
        SparseIndexParams,
        SparseVectorParams,
        VectorParams,
    )

    self.config.collection_name = collection_name
    if self.client.collection_exists(collection_name=collection_name):
        coll = self.client.get_collection(collection_name=collection_name)
        if (
            coll.status == CollectionStatus.GREEN
            and from_optional(coll.points_count, 0) > 0
        ):
            logger.warning(f"Non-empty Collection {collection_name} already exists")
            if not replace:
                logger.warning("Not replacing collection")
                return
            else:
                logger.warning("Recreating fresh collection")
        self.client.delete_collection(collection_name=collection_name)

    vectors_config = {
        "": VectorParams(
            size=self.embedding_dim,
            distance=Distance.COSINE,
        )
    }
    sparse_vectors_config = None
    if self.config.use_sparse_embeddings:
        sparse_vectors_config = {
            "text-sparse": SparseVectorParams(index=SparseIndexParams())
        }
    self.client.create_collection(
        collection_name=collection_name,
        vectors_config=vectors_config,
        sparse_vectors_config=sparse_vectors_config,
    )
    collection_info = self.client.get_collection(collection_name=collection_name)
    assert collection_info.status == CollectionStatus.GREEN
    assert collection_info.vectors_count in [0, None]
    if settings.debug:
        level = logger.getEffectiveLevel()
        logger.setLevel(logging.INFO)
        logger.info(collection_info)
        logger.setLevel(level)

get_sparse_embeddings(inputs, *, is_query=False)

Encode documents or queries using the configured sparse backend.

Both backends return Qdrant SparseVector objects. FastEmbed uses its query encoder for queries and does not import torch or transformers through this sparse encoding path.

Source code in langroid/vector_store/qdrantdb.py

def get_sparse_embeddings(
    self, inputs: List[str], *, is_query: bool = False
) -> List["SparseVector"]:
    """Encode documents or queries using the configured sparse backend.

    Both backends return Qdrant SparseVector objects. FastEmbed uses its
    query encoder for queries and does not import torch or transformers
    through this sparse encoding path.
    """
    from qdrant_client.http.models import SparseVector

    if not self.config.use_sparse_embeddings or not inputs:
        return []
    if self.config.sparse_embedding_backend == "fastembed":
        if self.sparse_fastembed_model is None:
            raise ValueError("The FastEmbed sparse model is not initialized")
        if is_query:
            embeddings = self.sparse_fastembed_model.query_embed(
                inputs, batch_size=self.config.batch_size
            )
        else:
            embeddings = self.sparse_fastembed_model.embed(
                inputs, batch_size=self.config.batch_size
            )
        return [
            SparseVector(
                indices=embedding.indices.tolist(),
                values=embedding.values.tolist(),
            )
            for embedding in embeddings
        ]

    import torch

    tokens = self.sparse_tokenizer(
        inputs, return_tensors="pt", truncation=True, padding=True
    )
    output = self.sparse_model(**tokens)
    vectors = torch.max(
        torch.log(torch.relu(output.logits) + torch.tensor(1.0))
        * tokens.attention_mask.unsqueeze(-1),
        dim=1,
    )[0].squeeze(dim=1)
    sparse_embeddings = []
    for vec in vectors:
        cols = vec.nonzero().squeeze().cpu().tolist()
        weights = vec[cols].cpu().tolist()
        sparse_embeddings.append(
            SparseVector(
                indices=cols,
                values=weights,
            )
        )
    return sparse_embeddings

is_valid_uuid(uuid_to_test)

Check if a given string is a valid UUID.

Source code in langroid/vector_store/qdrantdb.py

def is_valid_uuid(uuid_to_test: str) -> bool:
    """
    Check if a given string is a valid UUID.
    """
    try:
        uuid_obj = uuid.UUID(uuid_to_test)
        return str(uuid_obj) == uuid_to_test
    except Exception:
        pass
    # Check for valid unsigned 64-bit integer
    try:
        int_value = int(uuid_to_test)
        return 0 <= int_value <= 18446744073709551615
    except ValueError:
        return False