add param annotations for qdrant

This commit is contained in:
root 2023-08-14 21:59:48 +08:00
parent bb7ea6c398
commit 80506ec3ce

View file

@ -10,6 +10,14 @@ from metagpt.document_store.base_store import BaseStore
@dataclass
class QdrantConnection:
"""
Args:
url: qdrant url
host: qdrant host
port: qdrant port
memory: qdrant service use memory mode
api_key: qdrant cloud api_key
"""
url: str = None
host: str = None
port: int = None
@ -37,6 +45,17 @@ class QdrantStore(BaseStore):
force_recreate=False,
**kwargs,
):
"""
create a collection
Args:
collection_name: collection name
vectors_config: VectorParams object,detail in https://github.com/qdrant/qdrant-client
force_recreate: default is False, if True, will delete exists collection,then create it
**kwargs:
Returns:
"""
try:
self.client.get_collection(collection_name)
if force_recreate:
@ -63,6 +82,15 @@ class QdrantStore(BaseStore):
raise Exception(f"Delete collection {collection_name} failed.")
def add(self, collection_name: str, points: List[PointStruct]):
"""
add some vector data to qdrant
Args:
collection_name: collection name
points: list of PointStruct object, about PointStruct detail in https://github.com/qdrant/qdrant-client
Returns: None
"""
# self.client.upload_records()
self.client.upsert(
collection_name,
@ -77,6 +105,18 @@ class QdrantStore(BaseStore):
k=10,
return_vector=False,
):
"""
vector search
Args:
collection_name: qdrant collection name
query: input vector
query_filter: Filter object, detail in https://github.com/qdrant/qdrant-client
k: return the most similar k pieces of data
return_vector: whether return vector
Returns: list of dict
"""
hits = self.client.search(
collection_name=collection_name,
query_vector=query,