mem0ai-azure-mysql 0.1.115.2__py3-none-any.whl → 0.1.116.2__py3-none-any.whl

mem0/configs/llms/ollama.py

@@ -0,0 +1,56 @@
+from typing import Optional
+
+from mem0.configs.llms.base import BaseLlmConfig
+
+
+class OllamaConfig(BaseLlmConfig):
+    """
+    Configuration class for Ollama-specific parameters.
+    Inherits from BaseLlmConfig and adds Ollama-specific settings.
+    """
+
+    def __init__(
+        self,
+        # Base parameters
+        model: Optional[str] = None,
+        temperature: float = 0.1,
+        api_key: Optional[str] = None,
+        max_tokens: int = 2000,
+        top_p: float = 0.1,
+        top_k: int = 1,
+        enable_vision: bool = False,
+        vision_details: Optional[str] = "auto",
+        http_client_proxies: Optional[dict] = None,
+        # Ollama-specific parameters
+        ollama_base_url: Optional[str] = None,
+    ):
+        """
+        Initialize Ollama configuration.
+
+        Args:
+            model: Ollama model to use, defaults to None
+            temperature: Controls randomness, defaults to 0.1
+            api_key: Ollama API key, defaults to None
+            max_tokens: Maximum tokens to generate, defaults to 2000
+            top_p: Nucleus sampling parameter, defaults to 0.1
+            top_k: Top-k sampling parameter, defaults to 1
+            enable_vision: Enable vision capabilities, defaults to False
+            vision_details: Vision detail level, defaults to "auto"
+            http_client_proxies: HTTP client proxy settings, defaults to None
+            ollama_base_url: Ollama base URL, defaults to None
+        """
+        # Initialize base parameters
+        super().__init__(
+            model=model,
+            temperature=temperature,
+            api_key=api_key,
+            max_tokens=max_tokens,
+            top_p=top_p,
+            top_k=top_k,
+            enable_vision=enable_vision,
+            vision_details=vision_details,
+            http_client_proxies=http_client_proxies,
+        )
+
+        # Ollama-specific parameters
+        self.ollama_base_url = ollama_base_url

mem0/configs/llms/openai.py

@@ -0,0 +1,76 @@
+from typing import Any, Callable, List, Optional
+
+from mem0.configs.llms.base import BaseLlmConfig
+
+
+class OpenAIConfig(BaseLlmConfig):
+    """
+    Configuration class for OpenAI and OpenRouter-specific parameters.
+    Inherits from BaseLlmConfig and adds OpenAI-specific settings.
+    """
+
+    def __init__(
+        self,
+        # Base parameters
+        model: Optional[str] = None,
+        temperature: float = 0.1,
+        api_key: Optional[str] = None,
+        max_tokens: int = 2000,
+        top_p: float = 0.1,
+        top_k: int = 1,
+        enable_vision: bool = False,
+        vision_details: Optional[str] = "auto",
+        http_client_proxies: Optional[dict] = None,
+        # OpenAI-specific parameters
+        openai_base_url: Optional[str] = None,
+        models: Optional[List[str]] = None,
+        route: Optional[str] = "fallback",
+        openrouter_base_url: Optional[str] = None,
+        site_url: Optional[str] = None,
+        app_name: Optional[str] = None,
+        # Response monitoring callback
+        response_callback: Optional[Callable[[Any, dict, dict], None]] = None,
+    ):
+        """
+        Initialize OpenAI configuration.
+
+        Args:
+            model: OpenAI model to use, defaults to None
+            temperature: Controls randomness, defaults to 0.1
+            api_key: OpenAI API key, defaults to None
+            max_tokens: Maximum tokens to generate, defaults to 2000
+            top_p: Nucleus sampling parameter, defaults to 0.1
+            top_k: Top-k sampling parameter, defaults to 1
+            enable_vision: Enable vision capabilities, defaults to False
+            vision_details: Vision detail level, defaults to "auto"
+            http_client_proxies: HTTP client proxy settings, defaults to None
+            openai_base_url: OpenAI API base URL, defaults to None
+            models: List of models for OpenRouter, defaults to None
+            route: OpenRouter route strategy, defaults to "fallback"
+            openrouter_base_url: OpenRouter base URL, defaults to None
+            site_url: Site URL for OpenRouter, defaults to None
+            app_name: Application name for OpenRouter, defaults to None
+            response_callback: Optional callback for monitoring LLM responses.
+        """
+        # Initialize base parameters
+        super().__init__(
+            model=model,
+            temperature=temperature,
+            api_key=api_key,
+            max_tokens=max_tokens,
+            top_p=top_p,
+            top_k=top_k,
+            enable_vision=enable_vision,
+            vision_details=vision_details,
+            http_client_proxies=http_client_proxies,
+        )
+
+        # OpenAI-specific parameters
+        self.openai_base_url = openai_base_url
+        self.models = models
+        self.route = route
+        self.openrouter_base_url = openrouter_base_url
+        self.site_url = site_url
+        self.app_name = app_name
+        # Response monitoring
+        self.response_callback = response_callback

mem0/configs/llms/vllm.py

@@ -0,0 +1,56 @@
+from typing import Optional
+
+from mem0.configs.llms.base import BaseLlmConfig
+
+
+class VllmConfig(BaseLlmConfig):
+    """
+    Configuration class for vLLM-specific parameters.
+    Inherits from BaseLlmConfig and adds vLLM-specific settings.
+    """
+
+    def __init__(
+        self,
+        # Base parameters
+        model: Optional[str] = None,
+        temperature: float = 0.1,
+        api_key: Optional[str] = None,
+        max_tokens: int = 2000,
+        top_p: float = 0.1,
+        top_k: int = 1,
+        enable_vision: bool = False,
+        vision_details: Optional[str] = "auto",
+        http_client_proxies: Optional[dict] = None,
+        # vLLM-specific parameters
+        vllm_base_url: Optional[str] = None,
+    ):
+        """
+        Initialize vLLM configuration.
+
+        Args:
+            model: vLLM model to use, defaults to None
+            temperature: Controls randomness, defaults to 0.1
+            api_key: vLLM API key, defaults to None
+            max_tokens: Maximum tokens to generate, defaults to 2000
+            top_p: Nucleus sampling parameter, defaults to 0.1
+            top_k: Top-k sampling parameter, defaults to 1
+            enable_vision: Enable vision capabilities, defaults to False
+            vision_details: Vision detail level, defaults to "auto"
+            http_client_proxies: HTTP client proxy settings, defaults to None
+            vllm_base_url: vLLM base URL, defaults to None
+        """
+        # Initialize base parameters
+        super().__init__(
+            model=model,
+            temperature=temperature,
+            api_key=api_key,
+            max_tokens=max_tokens,
+            top_p=top_p,
+            top_k=top_k,
+            enable_vision=enable_vision,
+            vision_details=vision_details,
+            http_client_proxies=http_client_proxies,
+        )
+
+        # vLLM-specific parameters
+        self.vllm_base_url = vllm_base_url or "http://localhost:8000/v1"

mem0/configs/vector_stores/databricks.py

@@ -0,0 +1,63 @@
+from typing import Any, Dict, Optional
+
+from pydantic import BaseModel, Field, model_validator
+
+from databricks.sdk.service.vectorsearch import EndpointType, VectorIndexType, PipelineType
+
+
+class DatabricksConfig(BaseModel):
+    """Configuration for Databricks Vector Search vector store."""
+
+    workspace_url: str = Field(..., description="Databricks workspace URL")
+    access_token: Optional[str] = Field(None, description="Personal access token for authentication")
+    client_id: Optional[str] = Field(None, description="Databricks Service principal client ID")
+    client_secret: Optional[str] = Field(None, description="Databricks Service principal client secret")
+    azure_client_id: Optional[str] = Field(None, description="Azure AD application client ID (for Azure Databricks)")
+    azure_client_secret: Optional[str] = Field(
+        None, description="Azure AD application client secret (for Azure Databricks)"
+    )
+    endpoint_name: str = Field(..., description="Vector search endpoint name")
+    catalog: str = Field(..., description="The Unity Catalog catalog name")
+    schema: str = Field(..., description="The Unity Catalog schama name")
+    table_name: str = Field(..., description="Source Delta table name")
+    collection_name: str = Field("mem0", description="Vector search index name")
+    index_type: VectorIndexType = Field("DELTA_SYNC", description="Index type: DELTA_SYNC or DIRECT_ACCESS")
+    embedding_model_endpoint_name: Optional[str] = Field(
+        None, description="Embedding model endpoint for Databricks-computed embeddings"
+    )
+    embedding_dimension: int = Field(1536, description="Vector embedding dimensions")
+    endpoint_type: EndpointType = Field("STANDARD", description="Endpoint type: STANDARD or STORAGE_OPTIMIZED")
+    pipeline_type: PipelineType = Field("TRIGGERED", description="Sync pipeline type: TRIGGERED or CONTINUOUS")
+    warehouse_name: Optional[str] = Field(None, description="Databricks SQL warehouse Name")
+    query_type: str = Field("ANN", description="Query type: `ANN` and `HYBRID`")
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_extra_fields(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+        allowed_fields = set(cls.model_fields.keys())
+        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)}"
+            )
+        return values
+
+    @model_validator(mode="after")
+    def validate_authentication(self):
+        """Validate that either access_token or service principal credentials are provided."""
+        has_token = self.access_token is not None
+        has_service_principal = (self.client_id is not None and self.client_secret is not None) or (
+            self.azure_client_id is not None and self.azure_client_secret is not None
+        )
+
+        if not has_token and not has_service_principal:
+            raise ValueError(
+                "Either access_token or both client_id/client_secret or azure_client_id/azure_client_secret must be provided"
+            )
+
+        return self
+
+    model_config = {
+        "arbitrary_types_allowed": True,
+    }

mem0/configs/vector_stores/elasticsearch.py

@@ -19,6 +19,7 @@ class ElasticsearchConfig(BaseModel):
     custom_search_query: Optional[Callable[[List[float], int, Optional[Dict]], Dict]] = Field(
         None, description="Custom search query function. Parameters: (query, limit, filters) -> Dict"
     )
+    headers: Optional[Dict[str, str]] = Field(None, description="Custom headers to include in requests")
 
     @model_validator(mode="before")
     @classmethod
@@ -33,6 +34,23 @@ class ElasticsearchConfig(BaseModel):
 
         return values
 
+    @model_validator(mode="before")
+    @classmethod
+    def validate_headers(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+        """Validate headers format and content"""
+        headers = values.get("headers")
+        if headers is not None:
+            # Check if headers is a dictionary
+            if not isinstance(headers, dict):
+                raise ValueError("headers must be a dictionary")
+            
+            # Check if all keys and values are strings
+            for key, value in headers.items():
+                if not isinstance(key, str) or not isinstance(value, str):
+                    raise ValueError("All header keys and values must be strings")
+        
+        return values
+
     @model_validator(mode="before")
     @classmethod
     def validate_extra_fields(cls, values: Dict[str, Any]) -> Dict[str, Any]:

mem0/configs/vector_stores/milvus.py

@@ -25,6 +25,7 @@ class MilvusDBConfig(BaseModel):
     collection_name: str = Field("mem0", description="Name of the collection")
     embedding_model_dims: int = Field(1536, description="Dimensions of the embedding model")
     metric_type: str = Field("L2", description="Metric type for similarity search")
+    db_name: str = Field("", description="Name of the database")
 
     @model_validator(mode="before")
     @classmethod

mem0/configs/vector_stores/pgvector.py

@@ -13,15 +13,28 @@ class PGVectorConfig(BaseModel):
     port: Optional[int] = Field(None, description="Database port. Default is 1536")
     diskann: Optional[bool] = Field(True, description="Use diskann for approximate nearest neighbors search")
     hnsw: Optional[bool] = Field(False, description="Use hnsw for faster search")
+    # New SSL and connection options
+    sslmode: Optional[str] = Field(None, description="SSL mode for PostgreSQL connection (e.g., 'require', 'prefer', 'disable')")
+    connection_string: Optional[str] = Field(None, description="PostgreSQL connection string (overrides individual connection parameters)")
+    connection_pool: Optional[Any] = Field(None, description="psycopg2 connection pool object (overrides connection string and individual parameters)")
 
     @model_validator(mode="before")
     def check_auth_and_connection(cls, values):
+        # If connection_pool is provided, skip validation of individual connection parameters
+        if values.get("connection_pool") is not None:
+            return values
+        
+        # If connection_string is provided, skip validation of individual connection parameters
+        if values.get("connection_string") is not None:
+            return values
+        
+        # Otherwise, validate individual connection parameters
         user, password = values.get("user"), values.get("password")
         host, port = values.get("host"), values.get("port")
         if not user and not password:
-            raise ValueError("Both 'user' and 'password' must be provided.")
+            raise ValueError("Both 'user' and 'password' must be provided when not using connection_string or connection_pool.")
         if not host and not port:
-            raise ValueError("Both 'host' and 'port' must be provided.")
+            raise ValueError("Both 'host' and 'port' must be provided when not using connection_string or connection_pool.")
         return values
 
     @model_validator(mode="before")

mem0/configs/vector_stores/pinecone.py

@@ -18,6 +18,7 @@ class PineconeConfig(BaseModel):
     metric: str = Field("cosine", description="Distance metric for vector similarity")
     batch_size: int = Field(100, description="Batch size for operations")
     extra_params: Optional[Dict[str, Any]] = Field(None, description="Additional parameters for Pinecone client")
+    namespace: Optional[str] = Field(None, description="Namespace for the collection")
 
     @model_validator(mode="before")
     @classmethod

mem0/embeddings/azure_openai.py

@@ -24,13 +24,10 @@ class AzureOpenAIEmbedding(EmbeddingBase):
 
         auth_kwargs = {}
         if api_key:
-            print("Using API key for Azure OpenAI Embedding authentication.")
             auth_kwargs["api_key"] = api_key
         elif azure_ad_token:
-            print("Using Azure AD token for Azure OpenAI Embedding authentication.")
             auth_kwargs["azure_ad_token"] = azure_ad_token
         else:
-            print("Using DefaultAzureCredential for Azure OpenAI Embedding authentication.")
             credential = DefaultAzureCredential()
             token_provider = get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
             auth_kwargs["azure_ad_token_provider"] = token_provider

mem0/embeddings/ollama.py

@@ -36,7 +36,7 @@ class OllamaEmbedding(EmbeddingBase):
         Ensure the specified model exists locally. If not, pull it from Ollama.
         """
         local_models = self.client.list()["models"]
-        if not any(model.get("name") == self.config.model for model in local_models):
+        if not any(model.get("name") == self.config.model or model.get("model") == self.config.model for model in local_models):
             self.client.pull(self.config.model)
 
     def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):

mem0/graphs/configs.py

@@ -5,12 +5,24 @@ from pydantic import BaseModel, Field, field_validator, model_validator
 from mem0.llms.configs import LlmConfig
 
 
+class RerankConfig(BaseModel):
+    provider: str = Field(
+        description="Provider of the rerank model (e.g., 'openai', 'azure', 'cohere')",
+        default="cohere", 
+    )
+    config: Optional[dict] = Field(
+        description="Configuration for the specific rerank model", default={}
+    )
+
 class Neo4jConfig(BaseModel):
     url: Optional[str] = Field(None, description="Host address for the graph database")
     username: Optional[str] = Field(None, description="Username for the graph database")
     password: Optional[str] = Field(None, description="Password for the graph database")
     database: Optional[str] = Field(None, description="Database for the graph database")
     base_label: Optional[bool] = Field(None, description="Whether to use base node label __Entity__ for all entities")
+    similarity_threshold: float = Field(0.7, description="Threshold for the similarity of nodes")
+    top_k: int = Field(5, description="Number of top scored results to return")
+    rerank: Optional[RerankConfig] = Field(None, description="Rerank configuration")
 
     @model_validator(mode="before")
     def check_host_port_or_path(cls, values):
@@ -21,6 +33,12 @@ class Neo4jConfig(BaseModel):
         )
         if not url or not username or not password:
             raise ValueError("Please provide 'url', 'username' and 'password'.")
+        
+        if values.get("rerank") is not None:
+            values["rerank"] = RerankConfig(**values.get("rerank"))
+            if values["rerank"].provider not in ("cohere"):
+                raise ValueError("Invalid rerank provider. Supported providers are: cohere")
+
         return values
 
 
@@ -70,12 +88,16 @@ class NeptuneConfig(BaseModel):
             )
 
 
+class KuzuConfig(BaseModel):
+    db: Optional[str] = Field(":memory:", description="Path to a Kuzu database file")
+
+
 class GraphStoreConfig(BaseModel):
     provider: str = Field(
-        description="Provider of the data store (e.g., 'neo4j', 'memgraph', 'neptune')",
+        description="Provider of the data store (e.g., 'neo4j', 'memgraph', 'neptune', 'kuzu')",
         default="neo4j",
     )
-    config: Union[Neo4jConfig, MemgraphConfig, NeptuneConfig] = Field(
+    config: Union[Neo4jConfig, MemgraphConfig, NeptuneConfig, KuzuConfig] = Field(
         description="Configuration for the specific data store", default=None
     )
     llm: Optional[LlmConfig] = Field(description="LLM configuration for querying the graph store", default=None)
@@ -92,5 +114,7 @@ class GraphStoreConfig(BaseModel):
             return MemgraphConfig(**v.model_dump())
         elif provider == "neptune":
             return NeptuneConfig(**v.model_dump())
+        elif provider == "kuzu":
+            return KuzuConfig(**v.model_dump())
         else:
             raise ValueError(f"Unsupported graph store provider: {provider}")

mem0/graphs/neptune/main.py

@@ -1,4 +1,5 @@
 import logging
+
 from .base import NeptuneBase
 
 try:

mem0/graphs/tools.py

@@ -249,23 +249,23 @@ RELATIONS_STRUCT_TOOL = {
                     "items": {
                         "type": "object",
                         "properties": {
-                            "source_entity": {
+                            "source": {
                                 "type": "string",
                                 "description": "The source entity of the relationship.",
                             },
-                            "relatationship": {
+                            "relationship": {
                                 "type": "string",
                                 "description": "The relationship between the source and destination entities.",
                             },
-                            "destination_entity": {
+                            "destination": {
                                 "type": "string",
                                 "description": "The destination entity of the relationship.",
                             },
                         },
                         "required": [
-                            "source_entity",
-                            "relatationship",
-                            "destination_entity",
+                            "source",
+                            "relationship",
+                            "destination",
                         ],
                         "additionalProperties": False,
                     },

mem0/llms/anthropic.py

@@ -1,17 +1,37 @@
 import os
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 
 try:
     import anthropic
 except ImportError:
     raise ImportError("The 'anthropic' library is required. Please install it using 'pip install anthropic'.")
 
+from mem0.configs.llms.anthropic import AnthropicConfig
 from mem0.configs.llms.base import BaseLlmConfig
 from mem0.llms.base import LLMBase
 
 
 class AnthropicLLM(LLMBase):
-    def __init__(self, config: Optional[BaseLlmConfig] = None):
+    def __init__(self, config: Optional[Union[BaseLlmConfig, AnthropicConfig, Dict]] = None):
+        # Convert to AnthropicConfig if needed
+        if config is None:
+            config = AnthropicConfig()
+        elif isinstance(config, dict):
+            config = AnthropicConfig(**config)
+        elif isinstance(config, BaseLlmConfig) and not isinstance(config, AnthropicConfig):
+            # Convert BaseLlmConfig to AnthropicConfig
+            config = AnthropicConfig(
+                model=config.model,
+                temperature=config.temperature,
+                api_key=config.api_key,
+                max_tokens=config.max_tokens,
+                top_p=config.top_p,
+                top_k=config.top_k,
+                enable_vision=config.enable_vision,
+                vision_details=config.vision_details,
+                http_client_proxies=config.http_client,
+            )
+
         super().__init__(config)
 
         if not self.config.model:
@@ -26,6 +46,7 @@ class AnthropicLLM(LLMBase):
         response_format=None,
         tools: Optional[List[Dict]] = None,
         tool_choice: str = "auto",
+        **kwargs,
     ):
         """
         Generate a response based on the given messages using Anthropic.
@@ -35,6 +56,7 @@ class AnthropicLLM(LLMBase):
             response_format (str or object, optional): Format of the response. Defaults to "text".
             tools (list, optional): List of tools that the model can call. Defaults to None.
             tool_choice (str, optional): Tool choice method. Defaults to "auto".
+            **kwargs: Additional Anthropic-specific parameters.
 
         Returns:
             str: The generated response.
@@ -48,14 +70,15 @@ class AnthropicLLM(LLMBase):
             else:
                 filtered_messages.append(message)
 
-        params = {
-            "model": self.config.model,
-            "messages": filtered_messages,
-            "system": system_message,
-            "temperature": self.config.temperature,
-            "max_tokens": self.config.max_tokens,
-            "top_p": self.config.top_p,
-        }
+        params = self._get_supported_params(messages=messages, **kwargs)
+        params.update(
+            {
+                "model": self.config.model,
+                "messages": filtered_messages,
+                "system": system_message,
+            }
+        )
+
         if tools:  # TODO: Remove tools if no issues found with new memory addition logic
             params["tools"] = tools
             params["tool_choice"] = tool_choice