1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980 |
- import time
- from loguru import logger
- from pymilvus import connections, Collection, SearchResult
- from app.core.config import settings
- class Milvus:
- """
- Open-source vector database for unstructured data.
- """
- def __init__(self) -> None:
- self._host = settings.MILVUS_HOST
- self._port = settings.MILVUS_PORT
- async def query_embedding_by_pk(
- self,
- collection_name: str,
- primary_key_name: str,
- pk: int,
- output_field: str = "embeddings",
- ) -> list:
- start = time.perf_counter()
- connections.connect("default", host=self._host, port=self._port)
- logger.debug(time.perf_counter() - start)
- collection = Collection(name=collection_name)
- start = time.perf_counter()
- collection.load()
- logger.debug(time.perf_counter() - start)
- expr = f"{primary_key_name} in [{pk}]"
- start = time.perf_counter()
- res = collection.query(expr=expr, output_fields=[output_field])
- logger.debug(time.perf_counter() - start)
- try:
- emb = res[0].get(output_field)
- except [KeyError, IndexError] as e:
- emb = []
- logger.error(f"Can't find embedding by {pk}, reason: {e}")
- start = time.perf_counter()
- collection.release()
- # connections.disconnect("default")
- logger.debug(time.perf_counter() - start)
- return emb
- async def search(
- self,
- vec_list: list,
- collection_name: str,
- field_name: str,
- limit: int,
- ) -> list:
- connections.connect("default", host=self._host, port=self._port)
- collection = Collection(name=collection_name)
- collection.load()
- SEARCH_PARAM = {
- "metric_type": "L2",
- }
- start = time.perf_counter()
- res = collection.search(
- vec_list,
- field_name,
- param=SEARCH_PARAM,
- limit=limit,
- )
- logger.debug(time.perf_counter() - start)
- pk_list = []
- if isinstance(res, SearchResult):
- for hits in res:
- for hit in hits.ids:
- pk_list.append(hit)
- collection.release()
- # connections.disconnect("default")
- return pk_list
|