[Mem0] Update dependencies and make the package lighter (#1708)

Co-authored-by: Dev-Khant <devkhant24@gmail.com>
This commit is contained in:
Deshraj Yadav
2024-08-14 23:28:07 -07:00
committed by GitHub
parent e35786e567
commit a8ba7abb7d
35 changed files with 634 additions and 1594 deletions

View File

@@ -7,17 +7,26 @@ from mem0.vector_stores.configs import VectorStoreConfig
from mem0.llms.configs import LlmConfig
from mem0.embeddings.configs import EmbedderConfig
class MemoryItem(BaseModel):
id: str = Field(..., description="The unique identifier for the text data")
memory: str = Field(..., description="The memory deduced from the text data") # TODO After prompt changes from platform, update this
memory: str = Field(
..., description="The memory deduced from the text data"
) # TODO After prompt changes from platform, update this
hash: Optional[str] = Field(None, description="The hash of the memory")
# The metadata value can be anything and not just string. Fix it
metadata: Optional[Dict[str, Any]] = Field(None, description="Additional metadata for the text data")
metadata: Optional[Dict[str, Any]] = Field(
None, description="Additional metadata for the text data"
)
score: Optional[float] = Field(
None, description="The score associated with the text data"
)
created_at: Optional[str] = Field(None, description="The timestamp when the memory was created")
updated_at: Optional[str] = Field(None, description="The timestamp when the memory was updated")
created_at: Optional[str] = Field(
None, description="The timestamp when the memory was created"
)
updated_at: Optional[str] = Field(
None, description="The timestamp when the memory was updated"
)
class MemoryConfig(BaseModel):
@@ -36,4 +45,4 @@ class MemoryConfig(BaseModel):
history_db_path: str = Field(
description="Path to the history database",
default=os.path.join(mem0_dir, "history.db"),
)
)

View File

@@ -1,6 +1,7 @@
from abc import ABC
from typing import Optional
class BaseEmbedderConfig(ABC):
"""
Config for Embeddings.
@@ -11,12 +12,10 @@ class BaseEmbedderConfig(ABC):
model: Optional[str] = None,
api_key: Optional[str] = None,
embedding_dims: Optional[int] = None,
# Ollama specific
ollama_base_url: Optional[str] = None,
# Huggingface specific
model_kwargs: Optional[dict] = None
model_kwargs: Optional[dict] = None,
):
"""
Initializes a configuration class instance for the Embeddings.
@@ -33,7 +32,7 @@ class BaseEmbedderConfig(ABC):
:type model_kwargs: Optional[Dict[str, Any]], defaults a dict inside init
"""
self.model = model
self.api_key = api_key
self.embedding_dims = embedding_dims

View File

@@ -1,6 +1,7 @@
from abc import ABC
from typing import Optional
class BaseLlmConfig(ABC):
"""
Config for LLMs.
@@ -14,16 +15,14 @@ class BaseLlmConfig(ABC):
max_tokens: int = 3000,
top_p: float = 0,
top_k: int = 1,
# Openrouter specific
models: Optional[list[str]] = None,
route: Optional[str] = "fallback",
openrouter_base_url: Optional[str] = "https://openrouter.ai/api/v1",
site_url: Optional[str] = None,
app_name: Optional[str] = None,
# Ollama specific
ollama_base_url: Optional[str] = None
ollama_base_url: Optional[str] = None,
):
"""
Initializes a configuration class instance for the LLM.
@@ -55,7 +54,7 @@ class BaseLlmConfig(ABC):
:param ollama_base_url: The base URL of the LLM, defaults to None
:type ollama_base_url: Optional[str], optional
"""
self.model = model
self.temperature = temperature
self.api_key = api_key

View File

@@ -2,15 +2,20 @@ from typing import Optional, ClassVar, Dict, Any
from pydantic import BaseModel, Field, model_validator
class ChromaDbConfig(BaseModel):
try:
from chromadb.api.client import Client
except ImportError:
raise ImportError("Chromadb requires extra dependencies. Install with `pip install chromadb`") from None
raise ImportError(
"Chromadb requires extra dependencies. Install with `pip install chromadb`"
) from None
Client: ClassVar[type] = Client
collection_name: str = Field("mem0", description="Default name for the collection")
client: Optional[Client] = Field(None, description="Existing ChromaDB client instance")
client: Optional[Client] = Field(
None, description="Existing ChromaDB client instance"
)
path: Optional[str] = Field(None, description="Path to the database directory")
host: Optional[str] = Field(None, description="Database connection remote host")
port: Optional[int] = Field(None, description="Database connection remote port")
@@ -29,9 +34,11 @@ class ChromaDbConfig(BaseModel):
input_fields = set(values.keys())
extra_fields = input_fields - allowed_fields
if extra_fields:
raise ValueError(f"Extra fields not allowed: {', '.join(extra_fields)}. Please input only the following fields: {', '.join(allowed_fields)}")
raise ValueError(
f"Extra fields not allowed: {', '.join(extra_fields)}. Please input only the following fields: {', '.join(allowed_fields)}"
)
return values
model_config = {
"arbitrary_types_allowed": True,
}
}

View File

@@ -2,11 +2,14 @@ from typing import Optional, Dict, Any
from pydantic import BaseModel, Field, model_validator
class PGVectorConfig(BaseModel):
dbname: str = Field("postgres", description="Default name for the database")
collection_name: str = Field("mem0", description="Default name for the collection")
embedding_model_dims: Optional[int] = Field(1536, description="Dimensions of the embedding model")
embedding_model_dims: Optional[int] = Field(
1536, description="Dimensions of the embedding model"
)
user: Optional[str] = Field(None, description="Database user")
password: Optional[str] = Field(None, description="Database password")
host: Optional[str] = Field(None, description="Database host. Default is localhost")
@@ -21,7 +24,7 @@ class PGVectorConfig(BaseModel):
if not host and not port:
raise ValueError("Both 'host' and 'port' must be provided.")
return values
@model_validator(mode="before")
@classmethod
def validate_extra_fields(cls, values: Dict[str, Any]) -> Dict[str, Any]:
@@ -29,6 +32,7 @@ class PGVectorConfig(BaseModel):
input_fields = set(values.keys())
extra_fields = input_fields - allowed_fields
if extra_fields:
raise ValueError(f"Extra fields not allowed: {', '.join(extra_fields)}. Please input only the following fields: {', '.join(allowed_fields)}")
raise ValueError(
f"Extra fields not allowed: {', '.join(extra_fields)}. Please input only the following fields: {', '.join(allowed_fields)}"
)
return values

View File

@@ -1,16 +1,24 @@
from pydantic import BaseModel, Field, model_validator
from typing import Optional, ClassVar, Dict, Any
class QdrantConfig(BaseModel):
from qdrant_client import QdrantClient
from qdrant_client import QdrantClient
QdrantClient: ClassVar[type] = QdrantClient
collection_name: str = Field("mem0", description="Name of the collection")
embedding_model_dims: Optional[int] = Field(1536, description="Dimensions of the embedding model")
client: Optional[QdrantClient] = Field(None, description="Existing Qdrant client instance")
embedding_model_dims: Optional[int] = Field(
1536, description="Dimensions of the embedding model"
)
client: Optional[QdrantClient] = Field(
None, description="Existing Qdrant client instance"
)
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("/tmp/qdrant", description="Path for local Qdrant database")
path: Optional[str] = Field(
"/tmp/qdrant", description="Path for local Qdrant database"
)
url: Optional[str] = Field(None, description="Full URL for Qdrant server")
api_key: Optional[str] = Field(None, description="API key for Qdrant server")
on_disk: Optional[bool] = Field(False, description="Enables persistent storage")
@@ -38,9 +46,11 @@ class QdrantConfig(BaseModel):
input_fields = set(values.keys())
extra_fields = input_fields - allowed_fields
if extra_fields:
raise ValueError(f"Extra fields not allowed: {', '.join(extra_fields)}. Please input only the following fields: {', '.join(allowed_fields)}")
raise ValueError(
f"Extra fields not allowed: {', '.join(extra_fields)}. Please input only the following fields: {', '.join(allowed_fields)}"
)
return values
model_config = {
"arbitrary_types_allowed": True,
}
}