my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	import logging
#2	import os
#3	import shutil
#4
#5	from qdrant_client import QdrantClient
#6	from qdrant_client.models import (
#7	Distance,
#8	FieldCondition,
#9	Filter,
#10	MatchValue,
#11	PointIdsList,
#12	PointStruct,
#13	Range,
#14	VectorParams,
#15	)
#16
#17	from mem0.vector_stores.base import VectorStoreBase
#18
#19	logger = logging.getLogger(__name__)
#20
#21
#22	class Qdrant(VectorStoreBase):
#23	def __init__(
#24	self,
#25	collection_name: str,
#26	embedding_model_dims: int,
#27	client: QdrantClient = None,
#28	host: str = None,
#29	port: int = None,
#30	path: str = None,
#31	url: str = None,
#32	api_key: str = None,
#33	on_disk: bool = False,
#34	):
#35	"""
#36	Initialize the Qdrant vector store.
#37
#38	Args:
#39	collection_name (str): Name of the collection.
#40	embedding_model_dims (int): Dimensions of the embedding model.
#41	client (QdrantClient, optional): Existing Qdrant client instance. Defaults to None.
#42	host (str, optional): Host address for Qdrant server. Defaults to None.
#43	port (int, optional): Port for Qdrant server. Defaults to None.
#44	path (str, optional): Path for local Qdrant database. Defaults to None.
#45	url (str, optional): Full URL for Qdrant server. Defaults to None.
#46	api_key (str, optional): API key for Qdrant server. Defaults to None.
#47	on_disk (bool, optional): Enables persistent storage. Defaults to False.
#48	"""
#49	if client:
#50	self.client = client
#51	self.is_local = False
#52	else:
#53	params = {}
#54	if api_key:
#55	params["api_key"] = api_key
#56	if url:
#57	params["url"] = url
#58	if host and port:
#59	params["host"] = host
#60	params["port"] = port
#61
#62	if not params:
#63	params["path"] = path
#64	self.is_local = True
#65	if not on_disk:
#66	if os.path.exists(path) and os.path.isdir(path):
#67	shutil.rmtree(path)
#68	else:
#69	self.is_local = False
#70
#71	self.client = QdrantClient(**params)
#72
#73	self.collection_name = collection_name
#74	self.embedding_model_dims = embedding_model_dims
#75	self.on_disk = on_disk
#76	self.create_col(embedding_model_dims, on_disk)
#77
#78	def create_col(self, vector_size: int, on_disk: bool, distance: Distance = Distance.COSINE):
#79	"""
#80	Create a new collection.
#81
#82	Args:
#83	vector_size (int): Size of the vectors to be stored.
#84	on_disk (bool): Enables persistent storage.
#85	distance (Distance, optional): Distance metric for vector similarity. Defaults to Distance.COSINE.
#86	"""
#87	# Skip creating collection if already exists
#88	response = self.list_cols()
#89	for collection in response.collections:
#90	if collection.name == self.collection_name:
#91	logger.debug(f"Collection {self.collection_name} already exists. Skipping creation.")
#92	self._create_filter_indexes()
#93	return
#94
#95	self.client.create_collection(
#96	collection_name=self.collection_name,
#97	vectors_config=VectorParams(size=vector_size, distance=distance, on_disk=on_disk),
#98	)
#99	self._create_filter_indexes()
#100
#101	def _create_filter_indexes(self):
#102	"""Create indexes for commonly used filter fields to enable filtering."""
#103	# Only create payload indexes for remote Qdrant servers
#104	if self.is_local:
#105	logger.debug("Skipping payload index creation for local Qdrant (not supported)")
#106	return
#107
#108	common_fields = ["user_id", "agent_id", "run_id", "actor_id"]
#109
#110	for field in common_fields:
#111	try:
#112	self.client.create_payload_index(
#113	collection_name=self.collection_name,
#114	field_name=field,
#115	field_schema="keyword"
#116	)
#117	logger.info(f"Created index for {field} in collection {self.collection_name}")
#118	except Exception as e:
#119	logger.debug(f"Index for {field} might already exist: {e}")
#120
#121	def insert(self, vectors: list, payloads: list = None, ids: list = None):
#122	"""
#123	Insert vectors into a collection.
#124
#125	Args:
#126	vectors (list): List of vectors to insert.
#127	payloads (list, optional): List of payloads corresponding to vectors. Defaults to None.
#128	ids (list, optional): List of IDs corresponding to vectors. Defaults to None.
#129	"""
#130	logger.info(f"Inserting {len(vectors)} vectors into collection {self.collection_name}")
#131	points = [
#132	PointStruct(
#133	id=idx if ids is None else ids[idx],
#134	vector=vector,
#135	payload=payloads[idx] if payloads else {},
#136	)
#137	for idx, vector in enumerate(vectors)
#138	]
#139	self.client.upsert(collection_name=self.collection_name, points=points)
#140
#141	def _create_filter(self, filters: dict) -> Filter:
#142	"""
#143	Create a Filter object from the provided filters.
#144
#145	Args:
#146	filters (dict): Filters to apply.
#147
#148	Returns:
#149	Filter: The created Filter object.
#150	"""
#151	if not filters:
#152	return None
#153
#154	conditions = []
#155	for key, value in filters.items():
#156	if isinstance(value, dict) and "gte" in value and "lte" in value:
#157	conditions.append(FieldCondition(key=key, range=Range(gte=value["gte"], lte=value["lte"])))
#158	else:
#159	conditions.append(FieldCondition(key=key, match=MatchValue(value=value)))
#160	return Filter(must=conditions) if conditions else None
#161
#162	def search(self, query: str, vectors: list, limit: int = 5, filters: dict = None) -> list:
#163	"""
#164	Search for similar vectors.
#165
#166	Args:
#167	query (str): Query.
#168	vectors (list): Query vector.
#169	limit (int, optional): Number of results to return. Defaults to 5.
#170	filters (dict, optional): Filters to apply to the search. Defaults to None.
#171
#172	Returns:
#173	list: Search results.
#174	"""
#175	query_filter = self._create_filter(filters) if filters else None
#176	hits = self.client.query_points(
#177	collection_name=self.collection_name,
#178	query=vectors,
#179	query_filter=query_filter,
#180	limit=limit,
#181	)
#182	return hits.points
#183
#184	def delete(self, vector_id: int):
#185	"""
#186	Delete a vector by ID.
#187
#188	Args:
#189	vector_id (int): ID of the vector to delete.
#190	"""
#191	self.client.delete(
#192	collection_name=self.collection_name,
#193	points_selector=PointIdsList(
#194	points=[vector_id],
#195	),
#196	)
#197
#198	def update(self, vector_id: int, vector: list = None, payload: dict = None):
#199	"""
#200	Update a vector and its payload.
#201
#202	Args:
#203	vector_id (int): ID of the vector to update.
#204	vector (list, optional): Updated vector. Defaults to None.
#205	payload (dict, optional): Updated payload. Defaults to None.
#206	"""
#207	point = PointStruct(id=vector_id, vector=vector, payload=payload)
#208	self.client.upsert(collection_name=self.collection_name, points=[point])
#209
#210	def get(self, vector_id: int) -> dict:
#211	"""
#212	Retrieve a vector by ID.
#213
#214	Args:
#215	vector_id (int): ID of the vector to retrieve.
#216
#217	Returns:
#218	dict: Retrieved vector.
#219	"""
#220	result = self.client.retrieve(collection_name=self.collection_name, ids=[vector_id], with_payload=True)
#221	return result[0] if result else None
#222
#223	def list_cols(self) -> list:
#224	"""
#225	List all collections.
#226
#227	Returns:
#228	list: List of collection names.
#229	"""
#230	return self.client.get_collections()
#231
#232	def delete_col(self):
#233	"""Delete a collection."""
#234	self.client.delete_collection(collection_name=self.collection_name)
#235
#236	def col_info(self) -> dict:
#237	"""
#238	Get information about a collection.
#239
#240	Returns:
#241	dict: Collection information.
#242	"""
#243	return self.client.get_collection(collection_name=self.collection_name)
#244
#245	def list(self, filters: dict = None, limit: int = 100) -> list:
#246	"""
#247	List all vectors in a collection.
#248
#249	Args:
#250	filters (dict, optional): Filters to apply to the list. Defaults to None.
#251	limit (int, optional): Number of vectors to return. Defaults to 100.
#252
#253	Returns:
#254	list: List of vectors.
#255	"""
#256	query_filter = self._create_filter(filters) if filters else None
#257	result = self.client.scroll(
#258	collection_name=self.collection_name,
#259	scroll_filter=query_filter,
#260	limit=limit,
#261	with_payload=True,
#262	with_vectors=False,
#263	)
#264	return result
#265
#266	def reset(self):
#267	"""Reset the index by deleting and recreating it."""
#268	logger.warning(f"Resetting index {self.collection_name}...")
#269	self.delete_col()
#270	self.create_col(self.embedding_model_dims, self.on_disk)
#271

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public