memorisdk 2.0.0__py3-none-any.whl → 2.1.0__py3-none-any.whl

memori/agents/retrieval_agent.py

@@ -7,7 +7,7 @@ import json
 import threading
 import time
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 import openai
 from loguru import logger
@@ -58,8 +58,8 @@ Be strategic and comprehensive in your search planning."""
 
     def __init__(
         self,
-        api_key: Optional[str] = None,
-        model: Optional[str] = None,
+        api_key: str | None = None,
+        model: str | None = None,
         provider_config: Optional["ProviderConfig"] = None,
     ):
         """
@@ -94,9 +94,19 @@ Be strategic and comprehensive in your search planning."""
         # Background processing
         self._background_executor = None
 
-    def plan_search(
-        self, query: str, context: Optional[str] = None
-    ) -> MemorySearchQuery:
+        # Database type detection for unified search
+        self._database_type = None
+
+    def _detect_database_type(self, db_manager):
+        """Detect database type from db_manager"""
+        if self._database_type is None:
+            self._database_type = getattr(db_manager, "database_type", "sql")
+            logger.debug(
+                f"MemorySearchEngine: Detected database type: {self._database_type}"
+            )
+        return self._database_type
+
+    def plan_search(self, query: str, context: str | None = None) -> MemorySearchQuery:
         """
         Plan search strategy for a user query using OpenAI Structured Outputs with caching
 
@@ -182,13 +192,13 @@ Be strategic and comprehensive in your search planning."""
 
     def execute_search(
         self, query: str, db_manager, namespace: str = "default", limit: int = 10
-    ) -> List[Dict[str, Any]]:
+    ) -> list[dict[str, Any]]:
         """
         Execute intelligent search using planned strategies
 
         Args:
             query: User's search query
-            db_manager: Database manager instance
+            db_manager: Database manager instance (SQL or MongoDB)
             namespace: Memory namespace
             limit: Maximum results to return
 
@@ -196,27 +206,47 @@ Be strategic and comprehensive in your search planning."""
             List of relevant memory items with search metadata
         """
         try:
+            # Detect database type for optimal search strategy
+            db_type = self._detect_database_type(db_manager)
+
             # Plan the search
             search_plan = self.plan_search(query)
             logger.debug(
-                f"Search plan for '{query}': strategies={search_plan.search_strategy}, entities={search_plan.entity_filters}"
+                f"Search plan for '{query}': strategies={search_plan.search_strategy}, entities={search_plan.entity_filters}, db_type={db_type}"
             )
 
             all_results = []
             seen_memory_ids = set()
 
-            # Execute keyword search (primary strategy)
-            if (
-                search_plan.entity_filters
-                or "keyword_search" in search_plan.search_strategy
-            ):
+            # For MongoDB and SQL, use the unified search_memories method as primary strategy
+            # This ensures we use the database's native search capabilities
+            logger.debug(f"Executing unified database search using {db_type} manager")
+            primary_results = db_manager.search_memories(
+                query=search_plan.query_text or query, namespace=namespace, limit=limit
+            )
+            logger.debug(
+                f"Primary database search returned {len(primary_results)} results"
+            )
+
+            # Process primary results and add search metadata
+            for result in primary_results:
+                if (
+                    isinstance(result, dict)
+                    and result.get("memory_id") not in seen_memory_ids
+                ):
+                    seen_memory_ids.add(result["memory_id"])
+                    result["search_strategy"] = f"{db_type}_unified_search"
+                    result["search_reasoning"] = f"Direct {db_type} database search"
+                    all_results.append(result)
+
+            # If we have room for more results and specific entity filters, try keyword search
+            if len(all_results) < limit and search_plan.entity_filters:
                 logger.debug(
-                    f"Executing keyword search for: {search_plan.entity_filters}"
+                    f"Adding targeted keyword search for: {search_plan.entity_filters}"
                 )
                 keyword_results = self._execute_keyword_search(
-                    search_plan, db_manager, namespace, limit
+                    search_plan, db_manager, namespace, limit - len(all_results)
                 )
-                logger.debug(f"Keyword search returned {len(keyword_results)} results")
 
                 for result in keyword_results:
                     if (
@@ -230,20 +260,17 @@ Be strategic and comprehensive in your search planning."""
                         )
                         all_results.append(result)
 
-            # Execute category-based search
-            if (
+            # If we have room for more results, try category-based search
+            if len(all_results) < limit and (
                 search_plan.category_filters
                 or "category_filter" in search_plan.search_strategy
             ):
                 logger.debug(
-                    f"Executing category search for: {[c.value for c in search_plan.category_filters]}"
+                    f"Adding category search for: {[c.value for c in search_plan.category_filters]}"
                 )
                 category_results = self._execute_category_search(
                     search_plan, db_manager, namespace, limit - len(all_results)
                 )
-                logger.debug(
-                    f"Category search returned {len(category_results)} results"
-                )
 
                 for result in category_results:
                     if (
@@ -257,20 +284,17 @@ Be strategic and comprehensive in your search planning."""
                         )
                         all_results.append(result)
 
-            # Execute importance-based search
-            if (
+            # If we have room for more results, try importance-based search
+            if len(all_results) < limit and (
                 search_plan.min_importance > 0.0
                 or "importance_filter" in search_plan.search_strategy
             ):
                 logger.debug(
-                    f"Executing importance search with min_importance: {search_plan.min_importance}"
+                    f"Adding importance search with min_importance: {search_plan.min_importance}"
                 )
                 importance_results = self._execute_importance_search(
                     search_plan, db_manager, namespace, limit - len(all_results)
                 )
-                logger.debug(
-                    f"Importance search returned {len(importance_results)} results"
-                )
 
                 for result in importance_results:
                     if (
@@ -284,22 +308,6 @@ Be strategic and comprehensive in your search planning."""
                         )
                         all_results.append(result)
 
-            # If no specific strategies worked, do a general search
-            if not all_results:
-                logger.debug(
-                    "No results from specific strategies, executing general search"
-                )
-                general_results = db_manager.search_memories(
-                    query=search_plan.query_text, namespace=namespace, limit=limit
-                )
-                logger.debug(f"General search returned {len(general_results)} results")
-
-                for result in general_results:
-                    if isinstance(result, dict):
-                        result["search_strategy"] = "general_search"
-                        result["search_reasoning"] = "General content search"
-                        all_results.append(result)
-
             # Filter out any non-dictionary results before processing
             valid_results = []
             for result in all_results:
@@ -362,7 +370,7 @@ Be strategic and comprehensive in your search planning."""
 
     def _execute_keyword_search(
         self, search_plan: MemorySearchQuery, db_manager, namespace: str, limit: int
-    ) -> List[Dict[str, Any]]:
+    ) -> list[dict[str, Any]]:
         """Execute keyword-based search"""
         keywords = search_plan.entity_filters
         if not keywords:
@@ -398,7 +406,7 @@ Be strategic and comprehensive in your search planning."""
 
     def _execute_category_search(
         self, search_plan: MemorySearchQuery, db_manager, namespace: str, limit: int
-    ) -> List[Dict[str, Any]]:
+    ) -> list[dict[str, Any]]:
         """Execute category-based search"""
         categories = (
             [cat.value for cat in search_plan.category_filters]
@@ -411,32 +419,96 @@ Be strategic and comprehensive in your search planning."""
 
         # This would need to be implemented in the database manager
         # For now, get all memories and filter by category
+        logger.debug(
+            f"Searching memories by categories: {categories} in namespace: {namespace}"
+        )
         all_results = db_manager.search_memories(
             query="", namespace=namespace, limit=limit * 3
         )
 
+        logger.debug(
+            f"Retrieved {len(all_results)} total results for category filtering"
+        )
+
         filtered_results = []
-        for result in all_results:
+        for i, result in enumerate(all_results):
+            logger.debug(f"Processing result {i+1}/{len(all_results)}: {type(result)}")
+
             # Extract category from processed_data if it's stored as JSON
             try:
-                if "processed_data" in result:
+                memory_category = None
+
+                # Check processed_data field first
+                if "processed_data" in result and result["processed_data"]:
                     processed_data = result["processed_data"]
+                    logger.debug(
+                        f"Found processed_data: {type(processed_data)} - {str(processed_data)[:100]}..."
+                    )
+
                     # Handle both dict and JSON string formats
                     if isinstance(processed_data, str):
-                        processed_data = json.loads(processed_data)
-                    elif not isinstance(processed_data, dict):
-                        continue  # Skip if neither dict nor string
+                        try:
+                            processed_data = json.loads(processed_data)
+                        except json.JSONDecodeError as je:
+                            logger.debug(f"JSON decode error for processed_data: {je}")
+                            continue
+
+                    if isinstance(processed_data, dict):
+                        # Try multiple possible category locations
+                        category_paths = [
+                            ["category", "primary_category"],
+                            ["category"],
+                            ["primary_category"],
+                            ["metadata", "category"],
+                            ["classification", "category"],
+                        ]
+
+                        for path in category_paths:
+                            temp_data = processed_data
+                            try:
+                                for key in path:
+                                    temp_data = temp_data.get(key, {})
+                                if isinstance(temp_data, str) and temp_data:
+                                    memory_category = temp_data
+                                    logger.debug(
+                                        f"Found category via path {path}: {memory_category}"
+                                    )
+                                    break
+                            except (AttributeError, TypeError):
+                                continue
+                    else:
+                        logger.debug(
+                            f"processed_data is not a dict after parsing: {type(processed_data)}"
+                        )
+                        continue
 
-                    memory_category = processed_data.get("category", {}).get(
-                        "primary_category", ""
+                # Fallback: check direct category field
+                if not memory_category and "category" in result and result["category"]:
+                    memory_category = result["category"]
+                    logger.debug(f"Found category via direct field: {memory_category}")
+
+                # Check if the found category matches any of our target categories
+                if memory_category:
+                    logger.debug(
+                        f"Comparing memory category '{memory_category}' against target categories {categories}"
                     )
                     if memory_category in categories:
                         filtered_results.append(result)
-                elif result.get("category") in categories:
-                    filtered_results.append(result)
-            except (json.JSONDecodeError, KeyError, AttributeError):
+                        logger.debug(f"✓ Category match found: {memory_category}")
+                    else:
+                        logger.debug(
+                            f"✗ Category mismatch: {memory_category} not in {categories}"
+                        )
+                else:
+                    logger.debug("No category found in result")
+
+            except Exception as e:
+                logger.debug(f"Error processing result {i+1}: {e}")
                 continue
 
+        logger.debug(
+            f"Category filtering complete: {len(filtered_results)} results match categories {categories}"
+        )
         return filtered_results[:limit]
 
     def _detect_structured_output_support(self) -> bool:
@@ -609,7 +681,7 @@ Be strategic and comprehensive in your search planning."""
 }"""
 
     def _create_search_query_from_dict(
-        self, data: Dict[str, Any], original_query: str
+        self, data: dict[str, Any], original_query: str
     ) -> MemorySearchQuery:
         """
         Create MemorySearchQuery from dictionary with proper validation and defaults
@@ -651,7 +723,7 @@ Be strategic and comprehensive in your search planning."""
 
     def _execute_importance_search(
         self, search_plan: MemorySearchQuery, db_manager, namespace: str, limit: int
-    ) -> List[Dict[str, Any]]:
+    ) -> list[dict[str, Any]]:
         """Execute importance-based search"""
         min_importance = max(
             search_plan.min_importance, 0.7
@@ -692,7 +764,7 @@ Be strategic and comprehensive in your search planning."""
 
     async def execute_search_async(
         self, query: str, db_manager, namespace: str = "default", limit: int = 10
-    ) -> List[Dict[str, Any]]:
+    ) -> list[dict[str, Any]]:
         """
         Async version of execute_search for better performance in background processing
         """
@@ -807,7 +879,7 @@ Be strategic and comprehensive in your search planning."""
 
     def search_memories(
         self, query: str, max_results: int = 5, namespace: str = "default"
-    ) -> List[Dict[str, Any]]:
+    ) -> list[dict[str, Any]]:
         """
         Simple search interface for compatibility with memory tools
 
@@ -821,8 +893,11 @@ Be strategic and comprehensive in your search planning."""
         """
         # This is a compatibility method that uses the database manager directly
         # We'll need the database manager to be injected or passed
-        # For now, return empty list and log the issue
-        logger.warning(f"search_memories called without database manager: {query}")
+        # For now, return empty list and log the issue with parameters
+        logger.warning(
+            f"search_memories called without database manager: query='{query}', "
+            f"max_results={max_results}, namespace='{namespace}'"
+        )
         return []
 
 

memori/config/manager.py

@@ -4,7 +4,7 @@ Configuration manager for Memoriai
 
 import os
 from pathlib import Path
-from typing import Any, Dict, Optional, Union
+from typing import Any, Optional
 
 from loguru import logger
 
@@ -16,7 +16,7 @@ class ConfigManager:
     """Central configuration manager for Memoriai"""
 
     _instance: Optional["ConfigManager"] = None
-    _settings: Optional[MemoriSettings] = None
+    _settings: MemoriSettings | None = None
 
     def __new__(cls) -> "ConfigManager":
         """Singleton pattern for configuration manager"""
@@ -50,7 +50,7 @@ class ConfigManager:
             logger.warning(f"Failed to load configuration from environment: {e}")
             raise ConfigurationError(f"Environment configuration error: {e}")
 
-    def load_from_file(self, config_path: Union[str, Path]) -> None:
+    def load_from_file(self, config_path: str | Path) -> None:
         """Load configuration from file"""
         try:
             config_path = Path(config_path)
@@ -122,8 +122,8 @@ class ConfigManager:
         self._settings = MemoriSettings(**merged_dict)
 
     def _deep_merge_dicts(
-        self, base: Dict[str, Any], override: Dict[str, Any]
-    ) -> Dict[str, Any]:
+        self, base: dict[str, Any], override: dict[str, Any]
+    ) -> dict[str, Any]:
         """Deep merge two dictionaries"""
         result = base.copy()
 
@@ -139,7 +139,7 @@ class ConfigManager:
 
         return result
 
-    def save_to_file(self, config_path: Union[str, Path], format: str = "json") -> None:
+    def save_to_file(self, config_path: str | Path, format: str = "json") -> None:
         """Save current configuration to file"""
         if self._settings is None:
             raise ConfigurationError("No configuration loaded to save")
@@ -236,7 +236,7 @@ class ConfigManager:
             logger.error(f"Configuration validation failed: {e}")
             return False
 
-    def get_config_info(self) -> Dict[str, Any]:
+    def get_config_info(self) -> dict[str, Any]:
         """Get information about current configuration"""
         return {
             "loaded": self._settings is not None,

memori/config/memory_manager.py

@@ -6,7 +6,7 @@ a clean interface for memory management operations.
 """
 
 import uuid
-from typing import Any, Dict, List, Optional
+from typing import Any
 
 from loguru import logger
 
@@ -26,25 +26,25 @@ class MemoryManager:
         self,
         database_connect: str = "sqlite:///memori.db",
         template: str = "basic",
-        mem_prompt: Optional[str] = None,
+        mem_prompt: str | None = None,
         conscious_ingest: bool = False,
         auto_ingest: bool = False,
-        namespace: Optional[str] = None,
+        namespace: str | None = None,
         shared_memory: bool = False,
-        memory_filters: Optional[List[str]] = None,
-        user_id: Optional[str] = None,
+        memory_filters: list[str] | None = None,
+        user_id: str | None = None,
         verbose: bool = False,
-        provider_config: Optional[Any] = None,
+        provider_config: Any | None = None,
         # Additional parameters for compatibility
-        openai_api_key: Optional[str] = None,
-        api_key: Optional[str] = None,
-        api_type: Optional[str] = None,
-        base_url: Optional[str] = None,
-        azure_endpoint: Optional[str] = None,
-        azure_deployment: Optional[str] = None,
-        api_version: Optional[str] = None,
-        azure_ad_token: Optional[str] = None,
-        organization: Optional[str] = None,
+        openai_api_key: str | None = None,
+        api_key: str | None = None,
+        api_type: str | None = None,
+        base_url: str | None = None,
+        azure_endpoint: str | None = None,
+        azure_deployment: str | None = None,
+        api_version: str | None = None,
+        azure_ad_token: str | None = None,
+        organization: str | None = None,
         **kwargs,
     ):
         """
@@ -114,7 +114,7 @@ class MemoryManager:
 
         logger.debug("MemoryManager configured with Memori instance")
 
-    def enable(self, interceptors: Optional[List[str]] = None) -> Dict[str, Any]:
+    def enable(self, interceptors: list[str] | None = None) -> dict[str, Any]:
         """
         Enable memory recording using LiteLLM native callbacks.
 
@@ -162,7 +162,7 @@ class MemoryManager:
             logger.error(f"Failed to enable MemoryManager: {e}")
             return {"success": False, "message": str(e)}
 
-    def disable(self) -> Dict[str, Any]:
+    def disable(self) -> dict[str, Any]:
         """
         Disable memory recording using LiteLLM native callbacks.
 
@@ -194,7 +194,7 @@ class MemoryManager:
             logger.error(f"Failed to disable MemoryManager: {e}")
             return {"success": False, "message": str(e)}
 
-    def get_status(self) -> Dict[str, Dict[str, Any]]:
+    def get_status(self) -> dict[str, dict[str, Any]]:
         """
         Get status of memory recording system.
 
@@ -220,7 +220,7 @@ class MemoryManager:
             }
         }
 
-    def get_health(self) -> Dict[str, Any]:
+    def get_health(self) -> dict[str, Any]:
         """
         Get health check of the memory management system.
 
@@ -263,8 +263,8 @@ class MemoryManager:
         self,
         user_input: str,
         ai_output: str,
-        model: Optional[str] = None,
-        metadata: Optional[Dict[str, Any]] = None,
+        model: str | None = None,
+        metadata: dict[str, Any] | None = None,
     ) -> str:
         """
         Record a conversation (placeholder for future implementation).
@@ -279,10 +279,10 @@ class MemoryManager:
         self,
         query: str,
         limit: int = 5,
-        memory_types: Optional[List[str]] = None,
-        categories: Optional[List[str]] = None,
-        min_importance: Optional[float] = None,
-    ) -> List[Dict[str, Any]]:
+        memory_types: list[str] | None = None,
+        categories: list[str] | None = None,
+        min_importance: float | None = None,
+    ) -> list[dict[str, Any]]:
         """
         Search memories (placeholder for future implementation).
 

memori/config/settings.py

@@ -4,7 +4,7 @@ Pydantic-based configuration settings for Memoriai
 
 from enum import Enum
 from pathlib import Path
-from typing import Any, Dict, Optional, Union
+from typing import Any
 
 from pydantic import BaseModel, Field, validator
 
@@ -25,6 +25,7 @@ class DatabaseType(str, Enum):
     SQLITE = "sqlite"
     POSTGRESQL = "postgresql"
     MYSQL = "mysql"
+    MONGODB = "mongodb"
 
 
 class RetentionPolicy(str, Enum):
@@ -63,7 +64,13 @@ class DatabaseSettings(BaseModel):
             raise ValueError("Connection string cannot be empty")
 
         # Basic validation for supported protocols
-        valid_prefixes = ["sqlite://", "sqlite:///", "postgresql://", "mysql://"]
+        valid_prefixes = [
+            "sqlite://",
+            "sqlite:///",
+            "postgresql://",
+            "mysql://",
+            "mongodb://",
+        ]
         if not any(v.startswith(prefix) for prefix in valid_prefixes):
             raise ValueError(f"Unsupported database type in connection string: {v}")
 
@@ -73,7 +80,7 @@ class DatabaseSettings(BaseModel):
 class AgentSettings(BaseModel):
     """AI agent configuration settings"""
 
-    openai_api_key: Optional[str] = Field(
+    openai_api_key: str | None = Field(
         default=None, description="OpenAI API key for memory processing"
     )
     default_model: str = Field(
@@ -208,7 +215,7 @@ class MemoriSettings(BaseModel):
     integrations: IntegrationSettings = Field(default_factory=IntegrationSettings)
 
     # Custom settings
-    custom_settings: Dict[str, Any] = Field(
+    custom_settings: dict[str, Any] = Field(
         default_factory=dict, description="Custom user-defined settings"
     )
 
@@ -225,7 +232,7 @@ class MemoriSettings(BaseModel):
         return cls()
 
     @classmethod
-    def from_file(cls, config_path: Union[str, Path]) -> "MemoriSettings":
+    def from_file(cls, config_path: str | Path) -> "MemoriSettings":
         """Load settings from JSON/YAML file"""
         import json
         from pathlib import Path
@@ -251,7 +258,7 @@ class MemoriSettings(BaseModel):
 
         return cls(**data)
 
-    def to_file(self, config_path: Union[str, Path], format: str = "json") -> None:
+    def to_file(self, config_path: str | Path, format: str = "json") -> None:
         """Save settings to file"""
         import json
         from pathlib import Path