mdb-engine 0.2.1__py3-none-any.whl → 0.2.4__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,32 +445,39 @@ 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(self, text: str | list[str], model: str | None = None) -> list[list[float]]:
         """
-        Generate embeddings for text chunks.
+        Generate embeddings for text or a list of texts.
 
-        Uses the user-provided embedding provider/function.
+        Natural API that works with both single strings and lists.
 
         Args:
-            chunks: List of text chunks to embed
+            text: A single string or list of strings to embed
             model: Optional model identifier (passed to embedding provider)
 
         Returns:
-            List of embedding vectors (each is a list of floats)
+            List of embedding vectors (each is a list of floats).
+            If input was a single string, returns a list containing one vector.
 
         Example:
-            chunks = ["chunk 1", "chunk 2"]
-            vectors = await service.embed_chunks(chunks, model="text-embedding-3-small")
+            # Single string
+            vectors = await service.embed("Hello world", model="text-embedding-3-small")
+            # vectors is [[0.1, 0.2, ...]]
+
+            # List of strings (batch - more efficient)
+            vectors = await service.embed(["chunk 1", "chunk 2"], model="text-embedding-3-small")
+            # vectors is [[0.1, ...], [0.2, ...]]
         """
+        # Normalize to list
+        chunks = [text] if isinstance(text, str) else text
+
         if not chunks:
             return []
 
         try:
             # Use EmbeddingProvider's embed method (handles retries, logging, etc.)
             vectors = await self.embedding_provider.embed(chunks, model=model)
-            logger.info(f"Generated {len(vectors)} embeddings")
+            logger.info(f"Generated {len(vectors)} embedding(s)")
             return vectors
         except (
             AttributeError,
@@ -493,16 +490,36 @@ class EmbeddingService:
             logger.error(f"Error generating embeddings: {e}", exc_info=True)
             raise EmbeddingServiceError(f"Embedding generation failed: {str(e)}") from e
 
+    async def embed_chunks(self, chunks: list[str], model: str | None = None) -> list[list[float]]:
+        """
+        Generate embeddings for text chunks (list only).
+
+        DEPRECATED: Use embed() instead, which accepts both strings and lists.
+        This method is kept for backward compatibility.
+
+        Args:
+            chunks: List of text chunks to embed
+            model: Optional model identifier (passed to embedding provider)
+
+        Returns:
+            List of embedding vectors (each is a list of floats)
+
+        Example:
+            chunks = ["chunk 1", "chunk 2"]
+            vectors = await service.embed_chunks(chunks, model="text-embedding-3-small")
+        """
+        return await self.embed(chunks, model=model)
+
     async def process_and_store(
         self,
         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 +590,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 +643,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 +690,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 +709,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/README.md

@@ -10,6 +10,8 @@ Mem0.ai integration for intelligent memory management in MDB_ENGINE applications
 - **Semantic Search**: Vector-based semantic memory search
 - **Memory Inference**: Optional LLM-based memory inference and summarization
 - **Graph Memory**: Optional graph-based memory relationships (requires graph store config)
+- **Bucket Organization**: Built-in support for organizing memories into buckets (general, file, conversation, etc.)
+- **Dual Storage**: Store both extracted facts AND raw content for richer context retrieval
 
 ## Installation
 
@@ -203,6 +205,92 @@ await memory_service.delete(memory_id="memory_123", user_id="user123")
 await memory_service.delete_all(user_id="user123")
 ```
 
+### Bucket Organization
+
+Organize memories into buckets for better management:
+
+```python
+# Add memory to a bucket
+memory = await memory_service.add(
+    messages=[{"role": "user", "content": "I love Python programming"}],
+    user_id="user123",
+    bucket_id="coding:user123",
+    bucket_type="general",
+    metadata={"category": "coding"}
+)
+
+# Get all buckets for a user
+buckets = await memory_service.get_buckets(user_id="user123")
+
+# Get only file buckets
+file_buckets = await memory_service.get_buckets(
+    user_id="user123",
+    bucket_type="file"
+)
+
+# Get all memories in a specific bucket
+bucket_memories = await memory_service.get_bucket_memories(
+    bucket_id="file:document.pdf:user123",
+    user_id="user123"
+)
+```
+
+### Store Both Facts and Raw Content
+
+Store extracted facts alongside raw content for richer context:
+
+```python
+# Store both extracted facts and raw content
+facts, raw_memory_id = await memory_service.add_with_raw_content(
+    messages=[{"role": "user", "content": "Extract key facts from this document..."}],
+    raw_content="Full document text here...",
+    user_id="user123",
+    bucket_id="file:document.pdf:user123",
+    bucket_type="file",
+    infer=True  # Extract facts
+)
+
+# Later, retrieve raw content when needed
+raw_content = await memory_service.get_raw_content(
+    bucket_id="file:document.pdf:user123",
+    user_id="user123"
+)
+
+# Or include raw content when getting bucket memories
+all_memories = await memory_service.get_bucket_memories(
+    bucket_id="file:document.pdf:user123",
+    user_id="user123",
+    include_raw_content=True
+)
+```
+
+### Bucket Types
+
+Common bucket types:
+- **`general`**: General purpose buckets (e.g., category-based)
+- **`file`**: File-specific buckets (one per uploaded file)
+- **`conversation`**: Conversation-specific buckets
+- **`user`**: User-level buckets
+
+```python
+# General bucket (category-based)
+await memory_service.add(
+    messages=[{"role": "user", "content": "I prefer dark mode"}],
+    user_id="user123",
+    bucket_id="preferences:user123",
+    bucket_type="general"
+)
+
+# File bucket
+await memory_service.add(
+    messages=[{"role": "user", "content": "Document content..."}],
+    user_id="user123",
+    bucket_id="file:report.pdf:user123",
+    bucket_type="file",
+    metadata={"filename": "report.pdf"}
+)
+```
+
 ### Memory Inference
 
 With `infer=True`, the service can generate insights and summaries:
@@ -241,8 +329,11 @@ Mem0MemoryService(
 
 #### Methods
 
-- `add(messages, user_id, metadata=None)` - Add single memory
-- `add_all(memories)` - Add multiple memories
+- `add(messages, user_id, metadata=None, bucket_id=None, bucket_type=None, store_raw_content=False, raw_content=None)` - Add single memory with optional bucket and raw content storage
+- `add_with_raw_content(messages, raw_content, user_id, bucket_id=None, bucket_type=None)` - Store both extracted facts and raw content
+- `get_buckets(user_id, bucket_type=None, limit=None)` - Get all buckets for a user
+- `get_bucket_memories(bucket_id, user_id, include_raw_content=False, limit=None)` - Get all memories in a bucket
+- `get_raw_content(bucket_id, user_id)` - Get raw content for a bucket
 - `search(query, user_id, limit=10, filters=None)` - Search memories
 - `get(memory_id, user_id)` - Get specific memory
 - `get_all(user_id, filters=None)` - Get all memories for user