mdb-engine 0.2.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

mdb_engine/embeddings/service.py

@@ -23,7 +23,7 @@ import os
 import time
 from abc import ABC, abstractmethod
 from datetime import datetime
-from typing import Any, Dict, List, Optional, Union
+from typing import Any
 
 # Optional OpenAI SDK import
 try:
@@ -59,9 +59,7 @@ class BaseEmbeddingProvider(ABC):
     """
 
     @abstractmethod
-    async def embed(
-        self, text: Union[str, List[str]], model: Optional[str] = None
-    ) -> List[List[float]]:
+    async def embed(self, text: str | list[str], model: str | None = None) -> list[list[float]]:
         """
         Generate embeddings for text.
 
@@ -84,7 +82,7 @@ class OpenAIEmbeddingProvider(BaseEmbeddingProvider):
 
     def __init__(
         self,
-        api_key: Optional[str] = None,
+        api_key: str | None = None,
         default_model: str = "text-embedding-3-small",
     ):
         """
@@ -108,9 +106,7 @@ class OpenAIEmbeddingProvider(BaseEmbeddingProvider):
         self.client = AsyncOpenAI(api_key=api_key)
         self.default_model = default_model
 
-    async def embed(
-        self, text: Union[str, List[str]], model: Optional[str] = None
-    ) -> List[List[float]]:
+    async def embed(self, text: str | list[str], model: str | None = None) -> list[list[float]]:
         """Generate embeddings using OpenAI."""
         model = model or self.default_model
 
@@ -149,9 +145,9 @@ class AzureOpenAIEmbeddingProvider(BaseEmbeddingProvider):
 
     def __init__(
         self,
-        api_key: Optional[str] = None,
-        endpoint: Optional[str] = None,
-        api_version: Optional[str] = None,
+        api_key: str | None = None,
+        endpoint: str | None = None,
+        api_version: str | None = None,
         default_model: str = "text-embedding-3-small",
     ):
         """
@@ -191,9 +187,7 @@ class AzureOpenAIEmbeddingProvider(BaseEmbeddingProvider):
         )
         self.default_model = default_model
 
-    async def embed(
-        self, text: Union[str, List[str]], model: Optional[str] = None
-    ) -> List[List[float]]:
+    async def embed(self, text: str | list[str], model: str | None = None) -> list[list[float]]:
         """Generate embeddings using Azure OpenAI."""
         model = model or self.default_model
 
@@ -255,8 +249,8 @@ class EmbeddingProvider:
 
     def __init__(
         self,
-        embedding_provider: Optional[BaseEmbeddingProvider] = None,
-        config: Optional[Dict[str, Any]] = None,
+        embedding_provider: BaseEmbeddingProvider | None = None,
+        config: dict[str, Any] | None = None,
     ):
         """
         Initialize Embedding Provider.
@@ -293,9 +287,7 @@ class EmbeddingProvider:
         # Store config for potential future use
         self.config = config or {}
 
-    async def embed(
-        self, text: Union[str, List[str]], model: Optional[str] = None
-    ) -> List[List[float]]:
+    async def embed(self, text: str | list[str], model: str | None = None) -> list[list[float]]:
         """
         Generates vector embeddings for a string or list of strings.
 
@@ -361,10 +353,10 @@ class EmbeddingService:
 
     def __init__(
         self,
-        embedding_provider: Optional[EmbeddingProvider] = None,
+        embedding_provider: EmbeddingProvider | None = None,
         default_max_tokens: int = 1000,
         default_tokenizer_model: str = "gpt-3.5-turbo",
-        config: Optional[Dict[str, Any]] = None,
+        config: dict[str, Any] | None = None,
     ):
         """
         Initialize Embedding Service.
@@ -397,9 +389,7 @@ class EmbeddingService:
         self.default_max_tokens = default_max_tokens
         self.default_tokenizer_model = default_tokenizer_model
 
-    def _create_splitter(
-        self, max_tokens: int, tokenizer_model: Optional[str] = None
-    ) -> TextSplitter:
+    def _create_splitter(self, max_tokens: int, tokenizer_model: str | None = None) -> TextSplitter:
         """
         Create a TextSplitter instance.
 
@@ -419,9 +409,9 @@ class EmbeddingService:
     async def chunk_text(
         self,
         text_content: str,
-        max_tokens: Optional[int] = None,
-        tokenizer_model: Optional[str] = None,
-    ) -> List[str]:
+        max_tokens: int | None = None,
+        tokenizer_model: str | None = None,
+    ) -> list[str]:
         """
         Split text into semantic chunks.
 
@@ -455,9 +445,7 @@ class EmbeddingService:
             logger.error(f"Error chunking text: {e}", exc_info=True)
             raise EmbeddingServiceError(f"Chunking failed: {str(e)}") from e
 
-    async def embed_chunks(
-        self, chunks: List[str], model: Optional[str] = None
-    ) -> List[List[float]]:
+    async def embed_chunks(self, chunks: list[str], model: str | None = None) -> list[list[float]]:
         """
         Generate embeddings for text chunks.
 
@@ -498,11 +486,11 @@ class EmbeddingService:
         text_content: str,
         source_id: str,
         collection: Any,  # MongoDB collection (AppDB Collection or Motor collection)
-        max_tokens: Optional[int] = None,
-        tokenizer_model: Optional[str] = None,
-        embedding_model: Optional[str] = None,
-        metadata: Optional[Dict[str, Any]] = None,
-    ) -> Dict[str, Any]:
+        max_tokens: int | None = None,
+        tokenizer_model: str | None = None,
+        embedding_model: str | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
         """
         Process text and store chunks with embeddings in MongoDB.
 
@@ -573,7 +561,7 @@ class EmbeddingService:
 
         # Step 3: Prepare documents for insertion
         documents_to_insert = []
-        for i, (chunk_text, vector) in enumerate(zip(chunks, vectors)):
+        for i, (chunk_text, vector) in enumerate(zip(chunks, vectors, strict=False)):
             doc = {
                 "source_id": source_id,
                 "chunk_index": i,
@@ -626,10 +614,10 @@ class EmbeddingService:
     async def process_text(
         self,
         text_content: str,
-        max_tokens: Optional[int] = None,
-        tokenizer_model: Optional[str] = None,
-        embedding_model: Optional[str] = None,
-    ) -> List[Dict[str, Any]]:
+        max_tokens: int | None = None,
+        tokenizer_model: str | None = None,
+        embedding_model: str | None = None,
+    ) -> list[dict[str, Any]]:
         """
         Process text and return chunks with embeddings (without storing).
 
@@ -673,7 +661,7 @@ class EmbeddingService:
 
         # Prepare results
         results = []
-        for i, (chunk_text, vector) in enumerate(zip(chunks, vectors)):
+        for i, (chunk_text, vector) in enumerate(zip(chunks, vectors, strict=False)):
             results.append(
                 {
                     "chunk_index": i,
@@ -692,8 +680,8 @@ class EmbeddingService:
 
 # Dependency injection helper
 def get_embedding_service(
-    embedding_provider: Optional[BaseEmbeddingProvider] = None,
-    config: Optional[Dict[str, Any]] = None,
+    embedding_provider: BaseEmbeddingProvider | None = None,
+    config: dict[str, Any] | None = None,
 ) -> EmbeddingService:
     """
     Create EmbeddingService instance with auto-detected or provided embedding provider.

mdb_engine/exceptions.py

@@ -5,7 +5,7 @@ These exceptions provide more specific error types while maintaining
 backward compatibility with RuntimeError.
 """
 
-from typing import Any, Dict, List, Optional
+from typing import Any
 
 
 class MongoDBEngineError(RuntimeError):
@@ -21,7 +21,7 @@ class MongoDBEngineError(RuntimeError):
                  collection_name, etc.)
     """
 
-    def __init__(self, message: str, context: Optional[Dict[str, Any]] = None) -> None:
+    def __init__(self, message: str, context: dict[str, Any] | None = None) -> None:
         """
         Initialize the exception.
 
@@ -58,9 +58,9 @@ class InitializationError(MongoDBEngineError):
     def __init__(
         self,
         message: str,
-        mongo_uri: Optional[str] = None,
-        db_name: Optional[str] = None,
-        context: Optional[Dict[str, Any]] = None,
+        mongo_uri: str | None = None,
+        db_name: str | None = None,
+        context: dict[str, Any] | None = None,
     ) -> None:
         """
         Initialize the initialization error.
@@ -99,10 +99,10 @@ class ManifestValidationError(MongoDBEngineError):
     def __init__(
         self,
         message: str,
-        error_paths: Optional[List[str]] = None,
-        manifest_slug: Optional[str] = None,
-        schema_version: Optional[str] = None,
-        context: Optional[Dict[str, Any]] = None,
+        error_paths: list[str] | None = None,
+        manifest_slug: str | None = None,
+        schema_version: str | None = None,
+        context: dict[str, Any] | None = None,
     ) -> None:
         """
         Initialize the manifest validation error.
@@ -144,9 +144,9 @@ class ConfigurationError(MongoDBEngineError):
     def __init__(
         self,
         message: str,
-        config_key: Optional[str] = None,
-        config_value: Optional[Any] = None,
-        context: Optional[Dict[str, Any]] = None,
+        config_key: str | None = None,
+        config_value: Any | None = None,
+        context: dict[str, Any] | None = None,
     ) -> None:
         """
         Initialize the configuration error.
@@ -185,10 +185,10 @@ class QueryValidationError(MongoDBEngineError):
     def __init__(
         self,
         message: str,
-        query_type: Optional[str] = None,
-        operator: Optional[str] = None,
-        path: Optional[str] = None,
-        context: Optional[Dict[str, Any]] = None,
+        query_type: str | None = None,
+        operator: str | None = None,
+        path: str | None = None,
+        context: dict[str, Any] | None = None,
     ) -> None:
         """
         Initialize the query validation error.
@@ -231,10 +231,10 @@ class ResourceLimitExceeded(MongoDBEngineError):
     def __init__(
         self,
         message: str,
-        limit_type: Optional[str] = None,
-        limit_value: Optional[Any] = None,
-        actual_value: Optional[Any] = None,
-        context: Optional[Dict[str, Any]] = None,
+        limit_type: str | None = None,
+        limit_value: Any | None = None,
+        actual_value: Any | None = None,
+        context: dict[str, Any] | None = None,
     ) -> None:
         """
         Initialize the resource limit exceeded error.

mdb_engine/indexes/helpers.py

@@ -6,14 +6,14 @@ in index creation and management.
 """
 
 import logging
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
 
 logger = logging.getLogger(__name__)
 
 
 def normalize_keys(
-    keys: Union[Dict[str, Any], List[Tuple[str, Any]]],
-) -> List[Tuple[str, Any]]:
+    keys: dict[str, Any] | list[tuple[str, Any]],
+) -> list[tuple[str, Any]]:
     """
     Normalize index keys to a consistent format.
 
@@ -28,7 +28,7 @@ def normalize_keys(
     return keys
 
 
-def keys_to_dict(keys: Union[Dict[str, Any], List[Tuple[str, Any]]]) -> Dict[str, Any]:
+def keys_to_dict(keys: dict[str, Any] | list[tuple[str, Any]]) -> dict[str, Any]:
     """
     Convert index keys to dictionary format for comparison.
 
@@ -43,7 +43,7 @@ def keys_to_dict(keys: Union[Dict[str, Any], List[Tuple[str, Any]]]) -> Dict[str
     return {k: v for k, v in keys}
 
 
-def is_id_index(keys: Union[Dict[str, Any], List[Tuple[str, Any]]]) -> bool:
+def is_id_index(keys: dict[str, Any] | list[tuple[str, Any]]) -> bool:
     """
     Check if index keys target the _id field (which MongoDB creates automatically).
 
@@ -63,10 +63,10 @@ def is_id_index(keys: Union[Dict[str, Any], List[Tuple[str, Any]]]) -> bool:
 async def check_and_update_index(
     index_manager: Any,
     index_name: str,
-    expected_keys: Union[Dict[str, Any], List[Tuple[str, Any]]],
-    expected_options: Optional[Dict[str, Any]] = None,
+    expected_keys: dict[str, Any] | list[tuple[str, Any]],
+    expected_options: dict[str, Any] | None = None,
     log_prefix: str = "",
-) -> Tuple[bool, Optional[Dict[str, Any]]]:
+) -> tuple[bool, dict[str, Any] | None]:
     """
     Check if an index exists and matches the expected definition.
 
@@ -118,11 +118,11 @@ async def check_and_update_index(
 
 
 def validate_index_definition_basic(
-    index_def: Dict[str, Any],
+    index_def: dict[str, Any],
     index_name: str,
-    required_fields: List[str],
+    required_fields: list[str],
     log_prefix: str = "",
-) -> Tuple[bool, Optional[str]]:
+) -> tuple[bool, str | None]:
     """
     Basic validation for index definitions.
 

mdb_engine/indexes/manager.py

@@ -8,7 +8,7 @@ This module is part of MDB_ENGINE - MongoDB Engine.
 
 import json
 import logging
-from typing import Any, Dict, List
+from typing import Any
 
 from motor.motor_asyncio import AsyncIOMotorDatabase
 from pymongo.errors import (
@@ -44,7 +44,7 @@ logger = logging.getLogger(__name__)
 
 async def _handle_regular_index(
     index_manager: AsyncAtlasIndexManager,
-    index_def: Dict[str, Any],
+    index_def: dict[str, Any],
     index_name: str,
     log_prefix: str,
 ) -> None:
@@ -156,7 +156,7 @@ async def _handle_regular_index(
 
 async def _handle_ttl_index(
     index_manager: AsyncAtlasIndexManager,
-    index_def: Dict[str, Any],
+    index_def: dict[str, Any],
     index_name: str,
     log_prefix: str,
 ) -> None:
@@ -203,7 +203,7 @@ async def _handle_ttl_index(
 
 async def _handle_partial_index(
     index_manager: AsyncAtlasIndexManager,
-    index_def: Dict[str, Any],
+    index_def: dict[str, Any],
     index_name: str,
     log_prefix: str,
 ) -> None:
@@ -269,7 +269,7 @@ async def _handle_partial_index(
 
 async def _handle_text_index(
     index_manager: AsyncAtlasIndexManager,
-    index_def: Dict[str, Any],
+    index_def: dict[str, Any],
     index_name: str,
     log_prefix: str,
 ) -> None:
@@ -335,7 +335,7 @@ async def _handle_text_index(
 
 async def _handle_geospatial_index(
     index_manager: AsyncAtlasIndexManager,
-    index_def: Dict[str, Any],
+    index_def: dict[str, Any],
     index_name: str,
     log_prefix: str,
 ) -> None:
@@ -400,7 +400,7 @@ async def _handle_geospatial_index(
 
 async def _handle_search_index(
     index_manager: AsyncAtlasIndexManager,
-    index_def: Dict[str, Any],
+    index_def: dict[str, Any],
     index_name: str,
     index_type: str,
     slug: str,
@@ -502,7 +502,7 @@ async def _handle_search_index(
 
 async def _handle_hybrid_index(
     index_manager: AsyncAtlasIndexManager,
-    index_def: Dict[str, Any],
+    index_def: dict[str, Any],
     index_name: str,
     slug: str,
     log_prefix: str,
@@ -692,7 +692,7 @@ async def run_index_creation_for_collection(
     db: AsyncIOMotorDatabase,
     slug: str,
     collection_name: str,
-    index_definitions: List[Dict[str, Any]],
+    index_definitions: list[dict[str, Any]],
 ):
     """Create or update indexes for a collection based on index definitions."""
     log_prefix = f"[{slug} -> {collection_name}]"

mdb_engine/memory/service.py

@@ -9,7 +9,7 @@ mem0 handles embeddings and LLM via environment variables (.env).
 import logging
 import os
 import tempfile
-from typing import Any, Dict, List, Optional, Union
+from typing import Any
 
 # Set MEM0_DIR environment variable early to avoid permission issues
 # mem0 tries to create .mem0 directory at import time, so we set this before any import
@@ -74,7 +74,7 @@ def _detect_provider_from_env() -> str:
         return "openai"
 
 
-def _detect_embedding_dimensions(model_name: str) -> Optional[int]:
+def _detect_embedding_dimensions(model_name: str) -> int | None:
     """
     Auto-detect embedding dimensions from model name.
 
@@ -123,7 +123,7 @@ class Mem0MemoryServiceError(Exception):
 
 def _build_vector_store_config(
     db_name: str, collection_name: str, mongo_uri: str, embedding_model_dims: int
-) -> Dict[str, Any]:
+) -> dict[str, Any]:
     """Build vector store configuration for mem0."""
     return {
         "vector_store": {
@@ -138,7 +138,7 @@ def _build_vector_store_config(
     }
 
 
-def _build_embedder_config(provider: str, embedding_model: str, app_slug: str) -> Dict[str, Any]:
+def _build_embedder_config(provider: str, embedding_model: str, app_slug: str) -> dict[str, Any]:
     """Build embedder configuration for mem0."""
     clean_embedding_model = embedding_model.replace("azure/", "").replace("openai/", "")
     if provider == "azure":
@@ -190,7 +190,7 @@ def _build_embedder_config(provider: str, embedding_model: str, app_slug: str) -
 
 def _build_llm_config(
     provider: str, chat_model: str, temperature: float, app_slug: str
-) -> Dict[str, Any]:
+) -> dict[str, Any]:
     """Build LLM configuration for mem0."""
     clean_chat_model = chat_model.replace("azure/", "").replace("openai/", "")
     if provider == "azure":
@@ -245,7 +245,7 @@ def _build_llm_config(
     return config
 
 
-def _initialize_memory_instance(mem0_config: Dict[str, Any], app_slug: str) -> tuple:
+def _initialize_memory_instance(mem0_config: dict[str, Any], app_slug: str) -> tuple:
     """Initialize Mem0 Memory instance and return (instance, init_method)."""
     logger.debug(
         "Initializing Mem0 Memory with config structure",
@@ -330,7 +330,7 @@ class Mem0MemoryService:
         mongo_uri: str,
         db_name: str,
         app_slug: str,
-        config: Optional[Dict[str, Any]] = None,
+        config: dict[str, Any] | None = None,
     ):
         """
         Initialize Mem0 Memory Service.
@@ -498,11 +498,11 @@ class Mem0MemoryService:
 
     def add(
         self,
-        messages: Union[str, List[Dict[str, str]]],
-        user_id: Optional[str] = None,
-        metadata: Optional[Dict[str, Any]] = None,
+        messages: str | list[dict[str, str]],
+        user_id: str | None = None,
+        metadata: dict[str, Any] | None = None,
         **kwargs,
-    ) -> List[Dict[str, Any]]:
+    ) -> list[dict[str, Any]]:
         """
         Add memories from messages or text.
 
@@ -677,13 +677,13 @@ class Mem0MemoryService:
 
     def get_all(
         self,
-        user_id: Optional[str] = None,
-        limit: Optional[int] = None,
+        user_id: str | None = None,
+        limit: int | None = None,
         retry_on_empty: bool = True,
         max_retries: int = 2,
         retry_delay: float = 0.5,
         **kwargs,
-    ) -> List[Dict[str, Any]]:
+    ) -> list[dict[str, Any]]:
         """
         Get all memories for a user.
 
@@ -755,7 +755,7 @@ class Mem0MemoryService:
                     result = self.memory.get_all(
                         user_id=str(user_id), limit=limit, **kwargs
                     )  # Ensure string
-                    result_length = len(result) if isinstance(result, (list, dict)) else "N/A"
+                    result_length = len(result) if isinstance(result, list | dict) else "N/A"
                     logger.debug(
                         f"🟢 RESULT RECEIVED: type={type(result).__name__}, "
                         f"length={result_length}",
@@ -764,7 +764,7 @@ class Mem0MemoryService:
                             "user_id": user_id,
                             "result_type": type(result).__name__,
                             "result_length": (
-                                len(result) if isinstance(result, (list, dict)) else 0
+                                len(result) if isinstance(result, list | dict) else 0
                             ),
                             "attempt": attempt + 1,
                         },
@@ -791,7 +791,7 @@ class Mem0MemoryService:
                         "result_type": str(type(result)),
                         "is_dict": isinstance(result, dict),
                         "is_list": isinstance(result, list),
-                        "result_length": (len(result) if isinstance(result, (list, dict)) else 0),
+                        "result_length": (len(result) if isinstance(result, list | dict) else 0),
                     },
                 )
 
@@ -857,12 +857,12 @@ class Mem0MemoryService:
     def search(
         self,
         query: str,
-        user_id: Optional[str] = None,
-        limit: Optional[int] = None,
-        metadata: Optional[Dict[str, Any]] = None,
-        filters: Optional[Dict[str, Any]] = None,
+        user_id: str | None = None,
+        limit: int | None = None,
+        metadata: dict[str, Any] | None = None,
+        filters: dict[str, Any] | None = None,
         **kwargs,
-    ) -> List[Dict[str, Any]]:
+    ) -> list[dict[str, Any]]:
         """
         Search for relevant memories using semantic search.
 
@@ -970,7 +970,7 @@ class Mem0MemoryService:
             )
             raise Mem0MemoryServiceError(f"Failed to search memories: {e}") from e
 
-    def get(self, memory_id: str, user_id: Optional[str] = None, **kwargs) -> Dict[str, Any]:
+    def get(self, memory_id: str, user_id: str | None = None, **kwargs) -> dict[str, Any]:
         """
         Get a single memory by ID.
 
@@ -1037,11 +1037,11 @@ class Mem0MemoryService:
     def update(
         self,
         memory_id: str,
-        data: Union[str, List[Dict[str, str]]],
-        user_id: Optional[str] = None,
-        metadata: Optional[Dict[str, Any]] = None,
+        data: str | list[dict[str, str]],
+        user_id: str | None = None,
+        metadata: dict[str, Any] | None = None,
         **kwargs,
-    ) -> Dict[str, Any]:
+    ) -> dict[str, Any]:
         """
         Update a memory by ID with new data.
 
@@ -1120,7 +1120,7 @@ class Mem0MemoryService:
             )
             raise Mem0MemoryServiceError(f"Failed to update memory: {e}") from e
 
-    def delete(self, memory_id: str, user_id: Optional[str] = None, **kwargs) -> bool:
+    def delete(self, memory_id: str, user_id: str | None = None, **kwargs) -> bool:
         """
         Delete a memory by ID.
 
@@ -1169,7 +1169,7 @@ class Mem0MemoryService:
             )
             raise Mem0MemoryServiceError(f"Failed to delete memory: {e}") from e
 
-    def delete_all(self, user_id: Optional[str] = None, **kwargs) -> bool:
+    def delete_all(self, user_id: str | None = None, **kwargs) -> bool:
         """
         Delete all memories for a user.
 
@@ -1205,7 +1205,7 @@ class Mem0MemoryService:
 
 
 def get_memory_service(
-    mongo_uri: str, db_name: str, app_slug: str, config: Optional[Dict[str, Any]] = None
+    mongo_uri: str, db_name: str, app_slug: str, config: dict[str, Any] | None = None
 ) -> Mem0MemoryService:
     """
     Get or create a Mem0MemoryService instance (cached).