@@ -1,7 +1,7 @@
|
||||
import copy
|
||||
import os
|
||||
import uuid
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
from typing import Any, Optional, Union
|
||||
|
||||
try:
|
||||
from qdrant_client import QdrantClient
|
||||
@@ -69,14 +69,14 @@ class QdrantDB(BaseVectorDB):
|
||||
def _get_or_create_collection(self):
|
||||
return f"{self.config.collection_name}-{self.embedder.vector_dimension}".lower().replace("_", "-")
|
||||
|
||||
def get(self, ids: Optional[List[str]] = None, where: Optional[Dict[str, any]] = None, limit: Optional[int] = None):
|
||||
def get(self, ids: Optional[list[str]] = None, where: Optional[dict[str, any]] = None, limit: Optional[int] = None):
|
||||
"""
|
||||
Get existing doc ids present in vector database
|
||||
|
||||
:param ids: _list of doc ids to check for existence
|
||||
:type ids: List[str]
|
||||
:type ids: list[str]
|
||||
:param where: to filter data
|
||||
:type where: Dict[str, any]
|
||||
:type where: dict[str, any]
|
||||
:param limit: The number of entries to be fetched
|
||||
:type limit: Optional int, defaults to None
|
||||
:return: All the existing IDs
|
||||
@@ -122,21 +122,21 @@ class QdrantDB(BaseVectorDB):
|
||||
|
||||
def add(
|
||||
self,
|
||||
embeddings: List[List[float]],
|
||||
documents: List[str],
|
||||
metadatas: List[object],
|
||||
ids: List[str],
|
||||
**kwargs: Optional[Dict[str, any]],
|
||||
embeddings: list[list[float]],
|
||||
documents: list[str],
|
||||
metadatas: list[object],
|
||||
ids: list[str],
|
||||
**kwargs: Optional[dict[str, any]],
|
||||
):
|
||||
"""add data in vector database
|
||||
:param embeddings: list of embeddings for the corresponding documents to be added
|
||||
:type documents: List[List[float]]
|
||||
:type documents: list[list[float]]
|
||||
:param documents: list of texts to add
|
||||
:type documents: List[str]
|
||||
:type documents: list[str]
|
||||
:param metadatas: list of metadata associated with docs
|
||||
:type metadatas: List[object]
|
||||
:type metadatas: list[object]
|
||||
:param ids: ids of docs
|
||||
:type ids: List[str]
|
||||
:type ids: list[str]
|
||||
"""
|
||||
embeddings = self.embedder.embedding_fn(documents)
|
||||
|
||||
@@ -159,25 +159,25 @@ class QdrantDB(BaseVectorDB):
|
||||
|
||||
def query(
|
||||
self,
|
||||
input_query: List[str],
|
||||
input_query: list[str],
|
||||
n_results: int,
|
||||
where: Dict[str, any],
|
||||
where: dict[str, any],
|
||||
citations: bool = False,
|
||||
**kwargs: Optional[Dict[str, Any]],
|
||||
) -> Union[List[Tuple[str, Dict]], List[str]]:
|
||||
**kwargs: Optional[dict[str, Any]],
|
||||
) -> Union[list[tuple[str, dict]], list[str]]:
|
||||
"""
|
||||
query contents from vector database based on vector similarity
|
||||
:param input_query: list of query string
|
||||
:type input_query: List[str]
|
||||
:type input_query: list[str]
|
||||
:param n_results: no of similar documents to fetch from database
|
||||
:type n_results: int
|
||||
:param where: Optional. to filter data
|
||||
:type where: Dict[str, any]
|
||||
:type where: dict[str, any]
|
||||
:param citations: we use citations boolean param to return context along with the answer.
|
||||
:type citations: bool, default is False.
|
||||
:return: The content of the document that matched your query,
|
||||
along with url of the source and doc_id (if citations flag is true)
|
||||
:rtype: List[str], if citations=False, otherwise List[Tuple[str, str, str]]
|
||||
:rtype: list[str], if citations=False, otherwise list[tuple[str, str, str]]
|
||||
"""
|
||||
query_vector = self.embedder.embedding_fn([input_query])[0]
|
||||
keys = set(where.keys() if where is not None else set())
|
||||
|
||||
Reference in New Issue
Block a user