my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	import logging
#2	import uuid
#3	from typing import List, Optional
#4
#5	from pydantic import BaseModel
#6
#7	try:
#8	import vecs
#9	except ImportError:
#10	raise ImportError("The 'vecs' library is required. Please install it using 'pip install vecs'.")
#11
#12	from mem0.configs.vector_stores.supabase import IndexMeasure, IndexMethod
#13	from mem0.vector_stores.base import VectorStoreBase
#14
#15	logger = logging.getLogger(__name__)
#16
#17
#18	class OutputData(BaseModel):
#19	id: Optional[str]
#20	score: Optional[float]
#21	payload: Optional[dict]
#22
#23
#24	class Supabase(VectorStoreBase):
#25	def __init__(
#26	self,
#27	connection_string: str,
#28	collection_name: str,
#29	embedding_model_dims: int,
#30	index_method: IndexMethod = IndexMethod.AUTO,
#31	index_measure: IndexMeasure = IndexMeasure.COSINE,
#32	):
#33	"""
#34	Initialize the Supabase vector store using vecs.
#35
#36	Args:
#37	connection_string (str): PostgreSQL connection string
#38	collection_name (str): Collection name
#39	embedding_model_dims (int): Dimension of the embedding vector
#40	index_method (IndexMethod): Index method to use. Defaults to AUTO.
#41	index_measure (IndexMeasure): Distance measure to use. Defaults to COSINE.
#42	"""
#43	self.db = vecs.create_client(connection_string)
#44	self.collection_name = collection_name
#45	self.embedding_model_dims = embedding_model_dims
#46	self.index_method = index_method
#47	self.index_measure = index_measure
#48
#49	collections = self.list_cols()
#50	if collection_name not in collections:
#51	self.create_col(embedding_model_dims)
#52
#53	def _preprocess_filters(self, filters: Optional[dict] = None) -> Optional[dict]:
#54	"""
#55	Preprocess filters to be compatible with vecs.
#56
#57	Args:
#58	filters (Dict, optional): Filters to preprocess. Multiple filters will be
#59	combined with AND logic.
#60	"""
#61	if filters is None:
#62	return None
#63
#64	if len(filters) == 1:
#65	# For single filter, keep the simple format
#66	key, value = next(iter(filters.items()))
#67	return {key: {"$eq": value}}
#68
#69	# For multiple filters, use $and clause
#70	return {"$and": [{key: {"$eq": value}} for key, value in filters.items()]}
#71
#72	def create_col(self, embedding_model_dims: Optional[int] = None) -> None:
#73	"""
#74	Create a new collection with vector support.
#75	Will also initialize vector search index.
#76
#77	Args:
#78	embedding_model_dims (int, optional): Dimension of the embedding vector.
#79	If not provided, uses the dimension specified in initialization.
#80	"""
#81	dims = embedding_model_dims or self.embedding_model_dims
#82	if not dims:
#83	raise ValueError(
#84	"embedding_model_dims must be provided either during initialization or when creating collection"
#85	)
#86
#87	logger.info(f"Creating new collection: {self.collection_name}")
#88	try:
#89	self.collection = self.db.get_or_create_collection(name=self.collection_name, dimension=dims)
#90	self.collection.create_index(method=self.index_method.value, measure=self.index_measure.value)
#91	logger.info(f"Successfully created collection {self.collection_name} with dimension {dims}")
#92	except Exception as e:
#93	logger.error(f"Failed to create collection: {str(e)}")
#94	raise
#95
#96	def insert(
#97	self, vectors: List[List[float]], payloads: Optional[List[dict]] = None, ids: Optional[List[str]] = None
#98	):
#99	"""
#100	Insert vectors into the collection.
#101
#102	Args:
#103	vectors (List[List[float]]): List of vectors to insert
#104	payloads (List[Dict], optional): List of payloads corresponding to vectors
#105	ids (List[str], optional): List of IDs corresponding to vectors
#106	"""
#107	logger.info(f"Inserting {len(vectors)} vectors into collection {self.collection_name}")
#108
#109	if not ids:
#110	ids = [str(uuid.uuid4()) for _ in vectors]
#111	if not payloads:
#112	payloads = [{} for _ in vectors]
#113
#114	records = [(id, vector, payload) for id, vector, payload in zip(ids, vectors, payloads)]
#115
#116	self.collection.upsert(records)
#117
#118	def search(
#119	self, query: str, vectors: List[float], limit: int = 5, filters: Optional[dict] = None
#120	) -> List[OutputData]:
#121	"""
#122	Search for similar vectors.
#123
#124	Args:
#125	query (str): Query.
#126	vectors (List[float]): Query vector.
#127	limit (int, optional): Number of results to return. Defaults to 5.
#128	filters (Dict, optional): Filters to apply to the search. Defaults to None.
#129
#130	Returns:
#131	List[OutputData]: Search results
#132	"""
#133	filters = self._preprocess_filters(filters)
#134	results = self.collection.query(
#135	data=vectors, limit=limit, filters=filters, include_metadata=True, include_value=True
#136	)
#137
#138	return [OutputData(id=str(result[0]), score=float(result[1]), payload=result[2]) for result in results]
#139
#140	def delete(self, vector_id: str):
#141	"""
#142	Delete a vector by ID.
#143
#144	Args:
#145	vector_id (str): ID of the vector to delete
#146	"""
#147	self.collection.delete([(vector_id,)])
#148
#149	def update(self, vector_id: str, vector: Optional[List[float]] = None, payload: Optional[dict] = None):
#150	"""
#151	Update a vector and/or its payload.
#152
#153	Args:
#154	vector_id (str): ID of the vector to update
#155	vector (List[float], optional): Updated vector
#156	payload (Dict, optional): Updated payload
#157	"""
#158	if vector is None:
#159	# If only updating metadata, we need to get the existing vector
#160	existing = self.get(vector_id)
#161	if existing and existing.payload:
#162	vector = existing.payload.get("vector", [])
#163
#164	if vector:
#165	self.collection.upsert([(vector_id, vector, payload or {})])
#166
#167	def get(self, vector_id: str) -> Optional[OutputData]:
#168	"""
#169	Retrieve a vector by ID.
#170
#171	Args:
#172	vector_id (str): ID of the vector to retrieve
#173
#174	Returns:
#175	Optional[OutputData]: Retrieved vector data or None if not found
#176	"""
#177	result = self.collection.fetch([(vector_id,)])
#178	if not result:
#179	return []
#180
#181	record = result[0]
#182	return OutputData(id=str(record.id), score=None, payload=record.metadata)
#183
#184	def list_cols(self) -> List[str]:
#185	"""
#186	List all collections.
#187
#188	Returns:
#189	List[str]: List of collection names
#190	"""
#191	return self.db.list_collections()
#192
#193	def delete_col(self):
#194	"""Delete the collection."""
#195	self.db.delete_collection(self.collection_name)
#196
#197	def col_info(self) -> dict:
#198	"""
#199	Get information about the collection.
#200
#201	Returns:
#202	Dict: Collection information including name and configuration
#203	"""
#204	info = self.collection.describe()
#205	return {
#206	"name": info.name,
#207	"count": info.vectors,
#208	"dimension": info.dimension,
#209	"index": {"method": info.index_method, "metric": info.distance_metric},
#210	}
#211
#212	def list(self, filters: Optional[dict] = None, limit: int = 100) -> List[OutputData]:
#213	"""
#214	List vectors in the collection.
#215
#216	Args:
#217	filters (Dict, optional): Filters to apply
#218	limit (int, optional): Maximum number of results to return. Defaults to 100.
#219
#220	Returns:
#221	List[OutputData]: List of vectors
#222	"""
#223	filters = self._preprocess_filters(filters)
#224	query = [0] * self.embedding_model_dims
#225	ids = self.collection.query(
#226	data=query, limit=limit, filters=filters, include_metadata=True, include_value=False
#227	)
#228	ids = [id[0] for id in ids]
#229	records = self.collection.fetch(ids=ids)
#230
#231	return [[OutputData(id=str(record[0]), score=None, payload=record[2]) for record in records]]
#232
#233	def reset(self):
#234	"""Reset the index by deleting and recreating it."""
#235	logger.warning(f"Resetting index {self.collection_name}...")
#236	self.delete_col()
#237	self.create_col(self.embedding_model_dims)
#238

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public