klas/t6_mem0
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

adding param and return types (#1689)

Browse Source
This commit is contained in:
dbcontributions
2024-08-14 01:16:55 +05:30
committed by GitHub
parent 64218db7bd
commit a461091ba5
2 changed files with 82 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files

101
mem0/vector_stores/chroma.py
View File

@@ -1,5 +1,5 @@
import logging import logging
from typing import Optional from typing import Optional, List, Dict
from pydantic import BaseModel from pydantic import BaseModel
@@ -15,26 +15,27 @@ from mem0.vector_stores.base import VectorStoreBase
class OutputData(BaseModel): class OutputData(BaseModel):
id: Optional[str] # memory id id: Optional[str] # memory id
score: Optional[float] # distance score: Optional[float] # distance
payload: Optional[dict] # metadata payload: Optional[Dict] # metadata
class ChromaDB(VectorStoreBase): class ChromaDB(VectorStoreBase):
def __init__( def __init__(
self, self,
collection_name, collection_name: str,
client, client: Optional[chromadb.Client] = None,
host, host: Optional[str] = None,
port, port: Optional[int] = None,
path path: Optional[str] = None
): ):
""" """
Initialize the Chromadb vector store. Initialize the Chromadb vector store.
Args: Args:
client (chromadb.Client, optional): Existing chromadb client instance. collection_name (str): Name of the collection.
host (str, optional): Host address for chromadb server. client (chromadb.Client, optional): Existing chromadb client instance. Defaults to None.
port (int, optional): Port for chromadb server. host (str, optional): Host address for chromadb server. Defaults to None.
path (str, optional): Path for local chromadb database. port (int, optional): Port for chromadb server. Defaults to None.
path (str, optional): Path for local chromadb database. Defaults to None.
""" """
if client: if client:
self.client = client self.client = client
@@ -57,15 +58,15 @@ class ChromaDB(VectorStoreBase):
self.collection_name = collection_name self.collection_name = collection_name
self.collection = self.create_col(collection_name) self.collection = self.create_col(collection_name)
def _parse_output(self, data): def _parse_output(self, data: Dict) -> List[OutputData]:
""" """
Parse the output data. Parse the output data.
Args: Args:
data (dict): Output data. data (Dict): Output data.
Returns: Returns:
list: Parsed output data. List[OutputData]: Parsed output data.
""" """
keys = ['ids', 'distances', 'metadatas'] keys = ['ids', 'distances', 'metadatas']
values = [] values = []
@@ -82,21 +83,24 @@ class ChromaDB(VectorStoreBase):
result = [] result = []
for i in range(max_length): for i in range(max_length):
entry = OutputData( entry = OutputData(
id=ids[i] if isinstance(ids, list) and ids and i < len(ids) else None, id=ids[i] if isinstance(ids, list) and ids and i < len(ids) else None,
score=distances[i] if isinstance(distances, list) and distances and i < len(distances) else None, score=distances[i] if isinstance(distances, list) and distances and i < len(distances) else None,
payload=metadatas[i] if isinstance(metadatas, list) and metadatas and i < len(metadatas) else None, payload=metadatas[i] if isinstance(metadatas, list) and metadatas and i < len(metadatas) else None,
) )
result.append(entry) result.append(entry)
return result return result
def create_col(self, name, embedding_fn=None): def create_col(self, name: str, embedding_fn: Optional[callable] = None):
""" """
Create a new collection. Create a new collection.
Args: Args:
name (str): Name of the collection. name (str): Name of the collection.
embedding_fn (function): Embedding function to use. Defaults to None. embedding_fn (Optional[callable]): Embedding function to use. Defaults to None.
Returns:
chromadb.Collection: The created or retrieved collection.
""" """
# Skip creating collection if already exists # Skip creating collection if already exists
collections = self.list_cols() collections = self.list_cols()
@@ -110,102 +114,101 @@ class ChromaDB(VectorStoreBase):
) )
return collection return collection
def insert(self, vectors, payloads=None, ids=None): def insert(self, vectors: List[list], payloads: Optional[List[Dict]] = None, ids: Optional[List[str]] = None):
""" """
Insert vectors into a collection. Insert vectors into a collection.
Args: Args:
vectors (list): List of vectors to insert. vectors (List[list]): List of vectors to insert.
payloads (list, optional): List of payloads corresponding to vectors. Defaults to None. payloads (Optional[List[Dict]], optional): List of payloads corresponding to vectors. Defaults to None.
ids (list, optional): List of IDs corresponding to vectors. Defaults to None. ids (Optional[List[str]], optional): List of IDs corresponding to vectors. Defaults to None.
""" """
self.collection.add(ids=ids, embeddings=vectors, metadatas=payloads) self.collection.add(ids=ids, embeddings=vectors, metadatas=payloads)
def search(self, query, limit=5, filters=None): def search(self, query: List[list], limit: int = 5, filters: Optional[Dict] = None) -> List[OutputData]:
""" """
Search for similar vectors. Search for similar vectors.
Args: Args:
query (list): Query vector. query (List[list]): Query vector.
limit (int, optional): Number of results to return. Defaults to 5. limit (int, optional): Number of results to return. Defaults to 5.
filters (dict, optional): Filters to apply to the search. Defaults to None. filters (Optional[Dict], optional): Filters to apply to the search. Defaults to None.
Returns: Returns:
list: Search results. List[OutputData]: Search results.
""" """
results = self.collection.query(query_embeddings=query, where=filters, n_results=limit) results = self.collection.query(query_embeddings=query, where=filters, n_results=limit)
final_results = self._parse_output(results) final_results = self._parse_output(results)
return final_results return final_results
def delete(self, vector_id): def delete(self, vector_id: str):
""" """
Delete a vector by ID. Delete a vector by ID.
Args: Args:
vector_id (int): ID of the vector to delete. vector_id (str): ID of the vector to delete.
""" """
self.collection.delete(ids=vector_id) self.collection.delete(ids=vector_id)
def update(self, vector_id, vector=None, payload=None): def update(self, vector_id: str, vector: Optional[List[float]] = None, payload: Optional[Dict] = None):
""" """
Update a vector and its payload. Update a vector and its payload.
Args: Args:
vector_id (int): ID of the vector to update. vector_id (str): ID of the vector to update.
vector (list, optional): Updated vector. Defaults to None. vector (Optional[List[float]], optional): Updated vector. Defaults to None.
payload (dict, optional): Updated payload. Defaults to None. payload (Optional[Dict], optional): Updated payload. Defaults to None.
""" """
self.collection.update(ids=vector_id, embeddings=vector, metadatas=payload) self.collection.update(ids=vector_id, embeddings=vector, metadatas=payload)
def get(self, vector_id): def get(self, vector_id: str) -> OutputData:
""" """
Retrieve a vector by ID. Retrieve a vector by ID.
Args: Args:
vector_id (int): ID of the vector to retrieve. vector_id (str): ID of the vector to retrieve.
Returns: Returns:
dict: Retrieved vector. OutputData: Retrieved vector.
""" """
result = self.collection.get(ids=[vector_id]) result = self.collection.get(ids=[vector_id])
return self._parse_output(result)[0] return self._parse_output(result)[0]
def list_cols(self): def list_cols(self) -> List[chromadb.Collection]:
""" """
List all collections. List all collections.
Returns: Returns:
list: List of collection names. List[chromadb.Collection]: List of collections.
""" """
return self.client.list_collections() return self.client.list_collections()
def delete_col(self): def delete_col(self):
""" Delete a collection. """ """
Delete a collection.
"""
self.client.delete_collection(name=self.collection_name) self.client.delete_collection(name=self.collection_name)
def col_info(self): def col_info(self) -> Dict:
""" """
Get information about a collection. Get information about a collection.
Returns: Returns:
dict: Collection information. Dict: Collection information.
""" """
return self.client.get_collection(name=self.collection_name) return self.client.get_collection(name=self.collection_name)
def list(self, filters=None, limit=100): def list(self, filters: Optional[Dict] = None, limit: int = 100) -> List[OutputData]:
""" """
List all vectors in a collection. List all vectors in a collection.
Args: Args:
filters (dict, optional): Filters to apply to the list. filters (Optional[Dict], optional): Filters to apply to the list. Defaults to None.
limit (int, optional): Number of vectors to return. Defaults to 100. limit (int, optional): Number of vectors to return. Defaults to 100.
Returns: Returns:
list: List of vectors. List[OutputData]: List of vectors.
""" """
array = [[0 for _ in range(1536)] for _ in range(1536)] array = [[0 for _ in range(1536)] for _ in range(1536)]
results = self.collection.query(query_embeddings=array, where=filters, n_results=limit) results = self.collection.query(query_embeddings=array, where=filters, n_results=limit)
return [self._parse_output(results)] return self._parse_output(results)

57
mem0/vector_stores/qdrant.py
View File

@@ -20,27 +20,29 @@ from mem0.vector_stores.base import VectorStoreBase
class Qdrant(VectorStoreBase): class Qdrant(VectorStoreBase):
def __init__( def __init__(
self, self,
collection_name, collection_name: str,
embedding_model_dims, embedding_model_dims: int,
client, client: QdrantClient = None,
host, host: str = None,
port, port: int = None,
path, path: str = None,
url, url: str = None,
api_key, api_key: str = None,
on_disk on_disk: bool = False
): ):
""" """
Initialize the Qdrant vector store. Initialize the Qdrant vector store.
Args: Args:
client (QdrantClient, optional): Existing Qdrant client instance. collection_name (str): Name of the collection.
host (str, optional): Host address for Qdrant server. embedding_model_dims (int): Dimensions of the embedding model.
port (int, optional): Port for Qdrant server. client (QdrantClient, optional): Existing Qdrant client instance. Defaults to None.
path (str, optional): Path for local Qdrant database. host (str, optional): Host address for Qdrant server. Defaults to None.
url (str, optional): Full URL for Qdrant server. port (int, optional): Port for Qdrant server. Defaults to None.
api_key (str, optional): API key for Qdrant server. path (str, optional): Path for local Qdrant database. Defaults to None.
on_disk (bool, optional): Enables persistant storage. url (str, optional): Full URL for Qdrant server. Defaults to None.
api_key (str, optional): API key for Qdrant server. Defaults to None.
on_disk (bool, optional): Enables persistent storage. Defaults to False.
""" """
if client: if client:
self.client = client self.client = client
@@ -64,13 +66,13 @@ class Qdrant(VectorStoreBase):
self.collection_name = collection_name self.collection_name = collection_name
self.create_col(embedding_model_dims, on_disk) self.create_col(embedding_model_dims, on_disk)
def create_col(self, vector_size, on_disk, distance=Distance.COSINE): def create_col(self, vector_size: int, on_disk: bool, distance: Distance = Distance.COSINE):
""" """
Create a new collection. Create a new collection.
Args: Args:
name (str): Name of the collection.
vector_size (int): Size of the vectors to be stored. vector_size (int): Size of the vectors to be stored.
on_disk (bool): Enables persistent storage.
distance (Distance, optional): Distance metric for vector similarity. Defaults to Distance.COSINE. distance (Distance, optional): Distance metric for vector similarity. Defaults to Distance.COSINE.
""" """
# Skip creating collection if already exists # Skip creating collection if already exists
@@ -85,7 +87,7 @@ class Qdrant(VectorStoreBase):
vectors_config=VectorParams(size=vector_size, distance=distance, on_disk=on_disk), vectors_config=VectorParams(size=vector_size, distance=distance, on_disk=on_disk),
) )
def insert(self, vectors, payloads=None, ids=None): def insert(self, vectors: list, payloads: list = None, ids: list = None):
""" """
Insert vectors into a collection. Insert vectors into a collection.
@@ -104,7 +106,7 @@ class Qdrant(VectorStoreBase):
] ]
self.client.upsert(collection_name=self.collection_name, points=points) self.client.upsert(collection_name=self.collection_name, points=points)
def _create_filter(self, filters): def _create_filter(self, filters: dict) -> Filter:
""" """
Create a Filter object from the provided filters. Create a Filter object from the provided filters.
@@ -128,7 +130,7 @@ class Qdrant(VectorStoreBase):
) )
return Filter(must=conditions) if conditions else None return Filter(must=conditions) if conditions else None
def search(self, query, limit=5, filters=None): def search(self, query: list, limit: int = 5, filters: dict = None) -> list:
""" """
Search for similar vectors. Search for similar vectors.
@@ -149,7 +151,7 @@ class Qdrant(VectorStoreBase):
) )
return hits return hits
def delete(self, vector_id): def delete(self, vector_id: int):
""" """
Delete a vector by ID. Delete a vector by ID.
@@ -163,7 +165,7 @@ class Qdrant(VectorStoreBase):
), ),
) )
def update(self, vector_id, vector=None, payload=None): def update(self, vector_id: int, vector: list = None, payload: dict = None):
""" """
Update a vector and its payload. Update a vector and its payload.
@@ -175,7 +177,7 @@ class Qdrant(VectorStoreBase):
point = PointStruct(id=vector_id, vector=vector, payload=payload) point = PointStruct(id=vector_id, vector=vector, payload=payload)
self.client.upsert(collection_name=self.collection_name, points=[point]) self.client.upsert(collection_name=self.collection_name, points=[point])
def get(self, vector_id): def get(self, vector_id: int) -> dict:
""" """
Retrieve a vector by ID. Retrieve a vector by ID.
@@ -190,7 +192,7 @@ class Qdrant(VectorStoreBase):
) )
return result[0] if result else None return result[0] if result else None
def list_cols(self): def list_cols(self) -> list:
""" """
List all collections. List all collections.
@@ -203,7 +205,7 @@ class Qdrant(VectorStoreBase):
""" Delete a collection. """ """ Delete a collection. """
self.client.delete_collection(collection_name=self.collection_name) self.client.delete_collection(collection_name=self.collection_name)
def col_info(self): def col_info(self) -> dict:
""" """
Get information about a collection. Get information about a collection.
@@ -212,11 +214,12 @@ class Qdrant(VectorStoreBase):
""" """
return self.client.get_collection(collection_name=self.collection_name) return self.client.get_collection(collection_name=self.collection_name)
def list(self, filters=None, limit=100): def list(self, filters: dict = None, limit: int = 100) -> list:
""" """
List all vectors in a collection. List all vectors in a collection.
Args: Args:
filters (dict, optional): Filters to apply to the list. Defaults to None.
limit (int, optional): Number of vectors to return. Defaults to 100. limit (int, optional): Number of vectors to return. Defaults to 100.
Returns: Returns: