Add ChromaDB support (#1612)
This commit is contained in:
Dev Khant
@@ -1,4 +1,3 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
@@ -21,8 +20,7 @@ from mem0.memory.utils import get_update_memory_messages
|
||||
from mem0.vector_stores.configs import VectorStoreConfig
|
||||
from mem0.llms.configs import LlmConfig
|
||||
from mem0.embeddings.configs import EmbedderConfig
|
||||
from mem0.vector_stores.qdrant import Qdrant
|
||||
from mem0.utils.factory import LlmFactory, EmbedderFactory
|
||||
from mem0.utils.factory import LlmFactory, EmbedderFactory, VectorStoreFactory
|
||||
|
||||
# Setup user config
|
||||
setup_config()
|
||||
@@ -57,37 +55,17 @@ class MemoryConfig(BaseModel):
|
||||
description="Path to the history database",
|
||||
default=os.path.join(mem0_dir, "history.db"),
|
||||
)
|
||||
collection_name: str = Field(default="mem0", description="Name of the collection")
|
||||
embedding_model_dims: int = Field(
|
||||
default=1536, description="Dimensions of the embedding model"
|
||||
)
|
||||
|
||||
|
||||
class Memory(MemoryBase):
|
||||
def __init__(self, config: MemoryConfig = MemoryConfig()):
|
||||
self.config = config
|
||||
self.embedding_model = EmbedderFactory.create(self.config.embedder.provider)
|
||||
# Initialize the appropriate vector store based on the configuration
|
||||
vector_store_config = self.config.vector_store.config
|
||||
if self.config.vector_store.provider == "qdrant":
|
||||
self.vector_store = Qdrant(
|
||||
host=vector_store_config.host,
|
||||
port=vector_store_config.port,
|
||||
path=vector_store_config.path,
|
||||
url=vector_store_config.url,
|
||||
api_key=vector_store_config.api_key,
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported vector store type: {self.config.vector_store_type}"
|
||||
)
|
||||
|
||||
self.vector_store = VectorStoreFactory.create(self.config.vector_store.provider, self.config.vector_store.config)
|
||||
self.llm = LlmFactory.create(self.config.llm.provider, self.config.llm.config)
|
||||
self.db = SQLiteManager(self.config.history_db_path)
|
||||
self.collection_name = self.config.collection_name
|
||||
self.vector_store.create_col(
|
||||
name=self.collection_name, vector_size=self.embedding_model.dims
|
||||
)
|
||||
self.collection_name = self.config.vector_store.config.collection_name if "collection_name" in self.config.vector_store.config else "mem0"
|
||||
|
||||
capture_event("mem0.init", self)
|
||||
|
||||
@classmethod
|
||||
|
||||
@@ -44,4 +44,18 @@ class EmbedderFactory:
|
||||
return embedder_instance
|
||||
else:
|
||||
raise ValueError(f"Unsupported Embedder provider: {provider_name}")
|
||||
|
||||
|
||||
class VectorStoreFactory:
|
||||
provider_to_class = {
|
||||
"qdrant": "mem0.vector_stores.qdrant.Qdrant",
|
||||
"chromadb": "mem0.vector_stores.chroma.ChromaDB",
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def create(cls, provider_name, config):
|
||||
class_type = cls.provider_to_class.get(provider_name)
|
||||
if class_type:
|
||||
vector_store_instance = load_class(class_type)
|
||||
return vector_store_instance(**config)
|
||||
else:
|
||||
raise ValueError(f"Unsupported VectorStore provider: {provider_name}")
|
||||
|
||||
224
mem0/vector_stores/chroma.py
Normal file
224
mem0/vector_stores/chroma.py
Normal file
@@ -0,0 +1,224 @@
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
try:
|
||||
import chromadb
|
||||
from chromadb.config import Settings
|
||||
except ImportError:
|
||||
raise ImportError("Chromadb requires extra dependencies. Install with `pip install chromadb`") from None
|
||||
|
||||
from mem0.vector_stores.base import VectorStoreBase
|
||||
|
||||
|
||||
class OutputData(BaseModel):
|
||||
id: Optional[str] # memory id
|
||||
score: Optional[float] # distance
|
||||
payload: Optional[dict] # metadata
|
||||
|
||||
|
||||
class ChromaDB(VectorStoreBase):
|
||||
def __init__(
|
||||
self,
|
||||
collection_name="mem0",
|
||||
client=None,
|
||||
host=None,
|
||||
port=None,
|
||||
path=None
|
||||
):
|
||||
"""
|
||||
Initialize the Qdrant vector store.
|
||||
|
||||
Args:
|
||||
client (QdrantClient, optional): Existing Qdrant client instance. Defaults to None.
|
||||
host (str, optional): Host address for Qdrant server. Defaults to None.
|
||||
port (int, optional): Port for Qdrant server. Defaults to None.
|
||||
path (str, optional): Path for local Qdrant database. Defaults to None.
|
||||
"""
|
||||
if client:
|
||||
self.client = client
|
||||
else:
|
||||
self.settings = Settings(anonymized_telemetry=False)
|
||||
|
||||
if host and port:
|
||||
self.settings.chroma_server_host = host
|
||||
self.settings.chroma_server_http_port = port
|
||||
self.settings.chroma_api_impl = "chromadb.api.fastapi.FastAPI"
|
||||
else:
|
||||
if path is None:
|
||||
path = "db"
|
||||
|
||||
self.settings.persist_directory = path
|
||||
self.settings.is_persistent = True
|
||||
|
||||
self.client = chromadb.Client(self.settings)
|
||||
|
||||
self.collection = self.create_col(collection_name)
|
||||
|
||||
def _parse_output(self, data):
|
||||
"""
|
||||
Parse the output data.
|
||||
|
||||
Args:
|
||||
data (dict): Output data.
|
||||
|
||||
Returns:
|
||||
list: Parsed output data.
|
||||
"""
|
||||
keys = ['ids', 'distances', 'metadatas']
|
||||
values = []
|
||||
|
||||
for key in keys:
|
||||
value = data.get(key, [])
|
||||
if isinstance(value, list) and value and isinstance(value[0], list):
|
||||
value = value[0]
|
||||
values.append(value)
|
||||
|
||||
ids, distances, metadatas = values
|
||||
max_length = max(len(v) for v in values if isinstance(v, list) and v is not None)
|
||||
|
||||
result = []
|
||||
for i in range(max_length):
|
||||
entry = OutputData(
|
||||
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,
|
||||
payload=metadatas[i] if isinstance(metadatas, list) and metadatas and i < len(metadatas) else None,
|
||||
)
|
||||
result.append(entry)
|
||||
|
||||
return result
|
||||
|
||||
def create_col(self, name, embedding_fn=None):
|
||||
"""
|
||||
Create a new collection.
|
||||
|
||||
Args:
|
||||
name (str): Name of the collection.
|
||||
embedding_fn (function): Embedding function to use.
|
||||
"""
|
||||
# Skip creating collection if already exists
|
||||
collections = self.list_cols()
|
||||
for collection in collections:
|
||||
if collection.name == name:
|
||||
logging.debug(f"Collection {name} already exists. Skipping creation.")
|
||||
|
||||
collection = self.client.get_or_create_collection(
|
||||
name=name,
|
||||
embedding_function=embedding_fn,
|
||||
)
|
||||
return collection
|
||||
|
||||
def insert(self, name, vectors, payloads=None, ids=None):
|
||||
"""
|
||||
Insert vectors into a collection.
|
||||
|
||||
Args:
|
||||
name (str): Name of the collection.
|
||||
vectors (list): List of vectors to insert.
|
||||
payloads (list, optional): List of payloads corresponding to vectors. Defaults to None.
|
||||
ids (list, optional): List of IDs corresponding to vectors. Defaults to None.
|
||||
"""
|
||||
|
||||
self.collection.add(ids=ids, embeddings=vectors, metadatas=payloads)
|
||||
|
||||
def search(self, name, query, limit=5, filters=None):
|
||||
"""
|
||||
Search for similar vectors.
|
||||
|
||||
Args:
|
||||
name (str): Name of the collection.
|
||||
query (list): Query vector.
|
||||
limit (int, optional): Number of results to return. Defaults to 5.
|
||||
filters (dict, optional): Filters to apply to the search. Defaults to None.
|
||||
|
||||
Returns:
|
||||
list: Search results.
|
||||
"""
|
||||
results = self.collection.query(query_embeddings=query, where=filters, n_results=limit)
|
||||
final_results = self._parse_output(results)
|
||||
return final_results
|
||||
|
||||
def delete(self, name, vector_id):
|
||||
"""
|
||||
Delete a vector by ID.
|
||||
|
||||
Args:
|
||||
name (str): Name of the collection.
|
||||
vector_id (int): ID of the vector to delete.
|
||||
"""
|
||||
|
||||
self.collection.delete(ids=vector_id)
|
||||
|
||||
def update(self, name, vector_id, vector=None, payload=None):
|
||||
"""
|
||||
Update a vector and its payload.
|
||||
|
||||
Args:
|
||||
name (str): Name of the collection.
|
||||
vector_id (int): ID of the vector to update.
|
||||
vector (list, optional): Updated vector. Defaults to None.
|
||||
payload (dict, optional): Updated payload. Defaults to None.
|
||||
"""
|
||||
|
||||
self.collection.update(ids=vector_id, embeddings=vector, metadatas=payload)
|
||||
|
||||
def get(self, name, vector_id):
|
||||
"""
|
||||
Retrieve a vector by ID.
|
||||
|
||||
Args:
|
||||
name (str): Name of the collection.
|
||||
vector_id (int): ID of the vector to retrieve.
|
||||
|
||||
Returns:
|
||||
dict: Retrieved vector.
|
||||
"""
|
||||
result = self.collection.get(ids=[vector_id])
|
||||
return self._parse_output(result)[0]
|
||||
|
||||
def list_cols(self):
|
||||
"""
|
||||
List all collections.
|
||||
|
||||
Returns:
|
||||
list: List of collection names.
|
||||
"""
|
||||
return self.client.list_collections()
|
||||
|
||||
def delete_col(self, name):
|
||||
"""
|
||||
Delete a collection.
|
||||
|
||||
Args:
|
||||
name (str): Name of the collection to delete.
|
||||
"""
|
||||
self.client.delete_collection(name=name)
|
||||
|
||||
def col_info(self, name):
|
||||
"""
|
||||
Get information about a collection.
|
||||
|
||||
Args:
|
||||
name (str): Name of the collection.
|
||||
|
||||
Returns:
|
||||
dict: Collection information.
|
||||
"""
|
||||
return self.client.get_collection(name=name)
|
||||
|
||||
def list(self, name, filters=None, limit=100):
|
||||
"""
|
||||
List all vectors in a collection.
|
||||
|
||||
Args:
|
||||
name (str): Name of the collection.
|
||||
filters (dict, optional): Filters to apply to the list. Defaults to None.
|
||||
limit (int, optional): Number of vectors to return. Defaults to 100.
|
||||
|
||||
Returns:
|
||||
list: List of vectors.
|
||||
"""
|
||||
array = [[0 for _ in range(1536)] for _ in range(1536)]
|
||||
results = self.collection.query(query_embeddings=array, where=filters, n_results=limit)
|
||||
return [self._parse_output(results)]
|
||||
@@ -4,6 +4,10 @@ from pydantic import BaseModel, Field, field_validator, model_validator
|
||||
|
||||
|
||||
class QdrantConfig(BaseModel):
|
||||
collection_name: str = Field(default="mem0", description="Name of the collection")
|
||||
embedding_model_dims: Optional[int] = Field(
|
||||
default=1536, description="Dimensions of the embedding model"
|
||||
)
|
||||
host: Optional[str] = Field(None, description="Host address for Qdrant server")
|
||||
port: Optional[int] = Field(None, description="Port for Qdrant server")
|
||||
path: Optional[str] = Field(None, description="Path for local Qdrant database")
|
||||
@@ -26,14 +30,36 @@ class QdrantConfig(BaseModel):
|
||||
return values
|
||||
|
||||
|
||||
class ChromaDbConfig(BaseModel):
|
||||
collection_name: str = Field(
|
||||
default="mem0", description="Default name for the collection"
|
||||
)
|
||||
path: Optional[str] = Field(
|
||||
default=None, description="Path to the database directory"
|
||||
)
|
||||
host: Optional[str] = Field(
|
||||
default=None, description="Database connection remote host"
|
||||
)
|
||||
port: Optional[str] = Field(
|
||||
default=None, description="Database connection remote port"
|
||||
)
|
||||
|
||||
@model_validator(mode="before")
|
||||
def check_host_port_or_path(cls, values):
|
||||
host, port, path = values.get("host"), values.get("port"), values.get("path")
|
||||
if not path and not (host and port):
|
||||
raise ValueError("Either 'host' and 'port' or 'path' must be provided.")
|
||||
return values
|
||||
|
||||
|
||||
class VectorStoreConfig(BaseModel):
|
||||
provider: str = Field(
|
||||
description="Provider of the vector store (e.g., 'qdrant', 'chromadb', 'elasticsearch')",
|
||||
default="qdrant",
|
||||
)
|
||||
config: QdrantConfig = Field(
|
||||
config: Optional[dict] = Field(
|
||||
description="Configuration for the specific vector store",
|
||||
default=QdrantConfig(path="/tmp/qdrant"),
|
||||
default={},
|
||||
)
|
||||
|
||||
@field_validator("config")
|
||||
@@ -41,5 +67,7 @@ class VectorStoreConfig(BaseModel):
|
||||
provider = values.data.get("provider")
|
||||
if provider == "qdrant":
|
||||
return QdrantConfig(**v.model_dump())
|
||||
elif provider == "chromadb":
|
||||
return ChromaDbConfig(**v.model_dump())
|
||||
else:
|
||||
raise ValueError(f"Unsupported vector store provider: {provider}")
|
||||
|
||||
@@ -1,9 +1,7 @@
|
||||
import os
|
||||
import shutil
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from qdrant_client import QdrantClient
|
||||
from qdrant_client.models import (
|
||||
Distance,
|
||||
@@ -19,15 +17,11 @@ from qdrant_client.models import (
|
||||
from mem0.vector_stores.base import VectorStoreBase
|
||||
|
||||
|
||||
class QdrantConfig(BaseModel):
|
||||
host: Optional[str] = Field(None, description="Host address for Qdrant server")
|
||||
port: Optional[int] = Field(None, description="Port for Qdrant server")
|
||||
path: Optional[str] = Field(None, description="Path for local Qdrant database")
|
||||
|
||||
|
||||
class Qdrant(VectorStoreBase):
|
||||
def __init__(
|
||||
self,
|
||||
collection_name="mem0",
|
||||
embedding_model_dims=1536,
|
||||
client=None,
|
||||
host="localhost",
|
||||
port=6333,
|
||||
@@ -62,6 +56,8 @@ class Qdrant(VectorStoreBase):
|
||||
params["host"] = host
|
||||
params["port"] = port
|
||||
self.client = QdrantClient(**params)
|
||||
|
||||
self.create_col(collection_name, embedding_model_dims)
|
||||
|
||||
def create_col(self, name, vector_size, distance=Distance.COSINE):
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user