diff --git a/examples/mem_os/multi_user_memos_example.py b/examples/mem_os/multi_user_memos_example.py new file mode 100644 index 000000000..196cb380d --- /dev/null +++ b/examples/mem_os/multi_user_memos_example.py @@ -0,0 +1,125 @@ +""" +Example demonstrating how to use MOSProduct for multi-user scenarios. +""" + +from memos.configs.mem_cube import GeneralMemCubeConfig +from memos.configs.mem_os import MOSConfig +from memos.mem_cube.general import GeneralMemCube +from memos.mem_os.product import MOSProduct + + +def get_config(user_name): + openapi_config = { + "model_name_or_path": "gpt-4o-mini", + "temperature": 0.8, + "max_tokens": 1024, + "top_p": 0.9, + "top_k": 50, + "remove_think_prefix": True, + "api_key": "your-api-key-here", + "api_base": "https://api.openai.com/v1", + } + # Create a default configuration + default_config = MOSConfig( + user_id="root", + chat_model={"backend": "openai", "config": openapi_config}, + mem_reader={ + "backend": "naive", + "config": { + "llm": { + "backend": "openai", + "config": openapi_config, + }, + "embedder": { + "backend": "ollama", + "config": { + "model_name_or_path": "nomic-embed-text:latest", + }, + }, + }, + }, + enable_textual_memory=True, + enable_activation_memory=False, + top_k=5, + max_turns_window=20, + ) + default_cube_config = GeneralMemCubeConfig.model_validate( + { + "user_id": user_name, + "cube_id": f"{user_name}_default_cube", + "text_mem": { + "backend": "tree_text", + "config": { + "extractor_llm": {"backend": "openai", "config": openapi_config}, + "dispatcher_llm": {"backend": "openai", "config": openapi_config}, + "graph_db": { + "backend": "neo4j", + "config": { + "uri": "bolt://localhost:7687", + "user": "neo4j", + "password": "12345678", + "db_name": user_name, + "auto_create": True, + }, + }, + "embedder": { + "backend": "ollama", + "config": { + "model_name_or_path": "nomic-embed-text:latest", + }, + }, + }, + }, + "act_mem": {}, + "para_mem": {}, + } + ) + default_mem_cube = GeneralMemCube(default_cube_config) + return default_config, default_mem_cube + + +def main(): + default_config, default_mem_cube = get_config(user_name="alice") + # Initialize MOSProduct with default config + mos_product = MOSProduct(default_config=default_config) + + # Register first user with default config + result1 = mos_product.user_register( + user_id="alice", + user_name="alice", + interests="I'm interested in machine learning and AI research.", + default_mem_cube=default_mem_cube, + ) + print(f"User registration result: {result1}") + + # Chat with Alice + print("\n=== Chatting with Alice ===") + for response_chunk in mos_product.chat(query="What are my interests?", user_id="alice"): + print(response_chunk, end="") + + # Add memory for Alice + mos_product.add( + user_id="alice", + memory_content="I attended a machine learning conference last week.", + mem_cube_id=result1["default_cube_id"], + ) + + # Search memories for Alice + search_result = mos_product.search(query="conference", user_id="alice") + print(f"\nSearch result for Alice: {search_result}") + + # Search memories for Alice + search_result = mos_product.get_all(query="conference", user_id="alice", memory_type="text_mem") + print(f"\nSearch result for Alice: {search_result}") + + # List all users + users = mos_product.list_users() + print(f"\nAll registered users: {users}") + + # Get user info + alice_info = mos_product.get_user_info("alice") + print(f"\nAlice's info: {alice_info}") + + +if __name__ == "__main__": + main() diff --git a/examples/mem_os/persistent_memos_example.py b/examples/mem_os/persistent_memos_example.py new file mode 100644 index 000000000..16353be6a --- /dev/null +++ b/examples/mem_os/persistent_memos_example.py @@ -0,0 +1,192 @@ +""" +Example demonstrating persistent user management in MemOS. + +This example shows how to use the PersistentUserManager to maintain +user configurations across service restarts. +""" + +import os +import tempfile + +from memos.configs.mem_os import MOSConfig +from memos.mem_os.product import MOSProduct +from memos.mem_user.persistent_user_manager import PersistentUserManager, UserRole + + +def create_sample_config(user_id: str) -> MOSConfig: + """Create a sample configuration for a user.""" + return MOSConfig( + user_id=user_id, + chat_model={ + "backend": "openai", + "config": { + "model_name_or_path": "gpt-3.5-turbo", + "api_key": "your-api-key-here", + "temperature": 0.7, + }, + }, + mem_reader={ + "backend": "naive", + "config": { + "llm": { + "backend": "openai", + "config": { + "model_name_or_path": "gpt-3.5-turbo", + "api_key": "your-api-key-here", + }, + }, + "embedder": { + "backend": "ollama", + "config": { + "model_name_or_path": "nomic-embed-text:latest", + }, + }, + }, + }, + enable_textual_memory=True, + enable_activation_memory=False, + top_k=5, + max_turns_window=20, + ) + + +def demonstrate_persistence(): + """Demonstrate the persistence functionality.""" + print("=== MemOS Persistent User Management Demo ===\n") + + # Create a temporary database for this demo + temp_dir = tempfile.mkdtemp() + db_path = os.path.join(temp_dir, "demo_memos.db") + + try: + # Step 1: Create a persistent user manager + print("1. Creating PersistentUserManager...") + user_manager = PersistentUserManager(db_path=db_path) + print(f" Database created at: {db_path}") + + # Step 2: Create some sample configurations + print("\n2. Creating sample user configurations...") + user_configs = {} + for i in range(3): + user_id = f"user_{i + 1}" + user_name = f"User {i + 1}" + config = create_sample_config(user_id) + user_configs[user_id] = config + + # Create user with configuration + created_id = user_manager.create_user_with_config( + user_name, config, UserRole.USER, user_id + ) + print(f" Created user: {user_name} (ID: {created_id})") + + # Step 3: Verify configurations are saved + print("\n3. Verifying configurations are saved...") + for user_id in user_configs: + config = user_manager.get_user_config(user_id) + if config: + print(f" ✓ Configuration found for {user_id}") + print(f" - Textual memory enabled: {config.enable_textual_memory}") + print(f" - Top-k: {config.top_k}") + else: + print(f" ✗ Configuration not found for {user_id}") + + # Step 4: Simulate service restart by creating a new manager instance + print("\n4. Simulating service restart...") + print(" Creating new PersistentUserManager instance...") + new_user_manager = PersistentUserManager(db_path=db_path) + + # Step 5: Verify configurations are restored + print("\n5. Verifying configurations are restored after restart...") + for user_id in user_configs: + config = new_user_manager.get_user_config(user_id) + if config: + print(f" ✓ Configuration restored for {user_id}") + else: + print(f" ✗ Configuration not restored for {user_id}") + + # Step 6: Create MOSProduct and demonstrate restoration + print("\n6. Creating MOSProduct with persistent user manager...") + default_config = create_sample_config("default_user") + mos_product = MOSProduct(default_config=default_config) + + # The MOSProduct should automatically restore user instances + print(f" Active user instances: {len(mos_product.user_instances)}") + for user_id in mos_product.user_instances: + print(f" - {user_id}") + + # Step 7: Demonstrate configuration update + print("\n7. Demonstrating configuration update...") + user_id = "user_1" + original_config = user_manager.get_user_config(user_id) + if original_config: + # Update configuration + updated_config = original_config.model_copy(deep=True) + updated_config.top_k = 10 + updated_config.enable_activation_memory = True + + success = user_manager.save_user_config(user_id, updated_config) + if success: + print(f" ✓ Updated configuration for {user_id}") + print(f" - New top-k: {updated_config.top_k}") + print(f" - Activation memory: {updated_config.enable_activation_memory}") + else: + print(f" ✗ Failed to update configuration for {user_id}") + + # Step 8: List all configurations + print("\n8. Listing all user configurations...") + all_configs = user_manager.list_user_configs() + print(f" Total configurations: {len(all_configs)}") + for user_id, config in all_configs.items(): + print( + f" - {user_id}: top_k={config.top_k}, textual_memory={config.enable_textual_memory}" + ) + + print("\n=== Demo completed successfully! ===") + print(f"Database file: {db_path}") + print("You can inspect this file to see the persistent data.") + + except Exception as e: + print(f"Error during demo: {e}") + raise + finally: + # Cleanup + if os.path.exists(db_path): + os.remove(db_path) + if os.path.exists(temp_dir): + os.rmdir(temp_dir) + + +def demonstrate_api_usage(): + """Demonstrate how the API would work with persistence.""" + print("\n=== API Usage Example ===") + print(""" + With the new persistent system, your API calls would work like this: + + 1. Register a user (configuration is automatically saved): + POST /product/users/register + { + "user_id": "john_doe", + "user_name": "John Doe", + "interests": "AI, machine learning, programming" + } + + 2. Get user configuration: + GET /product/users/john_doe/config + + 3. Update user configuration: + PUT /product/users/john_doe/config + { + "user_id": "john_doe", + "enable_activation_memory": true, + "top_k": 10, + ... + } + + 4. After service restart, all user instances are automatically restored + and the user can immediately use the system without re-registration. + """) + + +if __name__ == "__main__": + demonstrate_persistence() + demonstrate_api_usage() diff --git a/src/memos/api/config.py b/src/memos/api/config.py new file mode 100644 index 000000000..840f48613 --- /dev/null +++ b/src/memos/api/config.py @@ -0,0 +1,289 @@ +import os + +from typing import Any + +from dotenv import load_dotenv + +from memos.configs.mem_cube import GeneralMemCubeConfig +from memos.configs.mem_os import MOSConfig +from memos.mem_cube.general import GeneralMemCube + + +# Load environment variables +load_dotenv() + + +class APIConfig: + """Centralized configuration management for MemOS APIs.""" + + @staticmethod + def get_openai_config() -> dict[str, Any]: + """Get OpenAI configuration.""" + return { + "model_name_or_path": os.getenv("MOS_OPENAI_MODEL", "gpt-4o-mini"), + "temperature": float(os.getenv("MOS_CHAT_TEMPERATURE", "0.8")), + "max_tokens": int(os.getenv("MOS_MAX_TOKENS", "1024")), + "top_p": float(os.getenv("MOS_TOP_P", "0.9")), + "top_k": int(os.getenv("MOS_TOP_K", "50")), + "remove_think_prefix": True, + "api_key": os.getenv("OPENAI_API_KEY", "your-api-key-here"), + "api_base": os.getenv("OPENAI_API_BASE", "https://api.openai.com/v1"), + } + + @staticmethod + def qwen_config() -> dict[str, Any]: + """Get Qwen configuration.""" + return { + "model_name_or_path": os.getenv("MOS_CHAT_MODEL", "Qwen/Qwen3-1.7B"), + "temperature": float(os.getenv("MOS_CHAT_TEMPERATURE", "0.8")), + "max_tokens": int(os.getenv("MOS_MAX_TOKENS", "4096")), + "remove_think_prefix": True, + } + + @staticmethod + def get_activation_config() -> dict[str, Any]: + """Get Ollama configuration.""" + return { + "backend": "kv_cache", + "config": { + "memory_filename": "activation_memory.pickle", + "extractor_llm": { + "backend": "huggingface_singleton", + "config": { + "model_name_or_path": os.getenv("MOS_CHAT_MODEL", "Qwen/Qwen3-1.7B"), + "temperature": 0.8, + "max_tokens": 1024, + "top_p": 0.9, + "top_k": 50, + "add_generation_prompt": True, + "remove_think_prefix": False, + }, + }, + }, + } + + @staticmethod + def get_neo4j_config() -> dict[str, Any]: + """Get Neo4j configuration.""" + return { + "uri": os.getenv("NEO4J_URI", "bolt://localhost:7687"), + "user": os.getenv("NEO4J_USER", "neo4j"), + "password": os.getenv("NEO4J_PASSWORD", "12345678"), + "auto_create": True, + } + + @staticmethod + def get_scheduler_config() -> dict[str, Any]: + """Get scheduler configuration.""" + return { + "backend": "general_scheduler", + "config": { + "top_k": int(os.getenv("MOS_SCHEDULER_TOP_K", "10")), + "top_n": int(os.getenv("MOS_SCHEDULER_TOP_N", "5")), + "act_mem_update_interval": int( + os.getenv("MOS_SCHEDULER_ACT_MEM_UPDATE_INTERVAL", "300") + ), + "context_window_size": int(os.getenv("MOS_SCHEDULER_CONTEXT_WINDOW_SIZE", "5")), + "activation_mem_size": int(os.getenv("MOS_SCHEDULER_ACTIVATION_MEM_SIZE", "1000")), + "thread_pool_max_workers": int( + os.getenv("MOS_SCHEDULER_THREAD_POOL_MAX_WORKERS", "10") + ), + "consume_interval_seconds": int( + os.getenv("MOS_SCHEDULER_CONSUME_INTERVAL_SECONDS", "3") + ), + "enable_parallel_dispatch": os.getenv( + "MOS_SCHEDULER_ENABLE_PARALLEL_DISPATCH", "true" + ).lower() + == "true", + }, + } + + @staticmethod + def is_scheduler_enabled() -> bool: + """Check if scheduler is enabled via environment variable.""" + return os.getenv("MOS_ENABLE_SCHEDULER", "false").lower() == "true" + + @staticmethod + def get_product_default_config() -> dict[str, Any]: + """Get default configuration for Product API.""" + openai_config = APIConfig.get_openai_config() + qwen_config = APIConfig.qwen_config() + config = { + "user_id": os.getenv("MOS_USER_ID", "root"), + "chat_model": { + "backend": os.getenv("MOS_CHAT_MODEL_PROVIDER", "openai"), + "config": openai_config + if os.getenv("MOS_CHAT_MODEL_PROVIDER", "openai") == "openai" + else qwen_config, + }, + "mem_reader": { + "backend": "simple_struct", + "config": { + "llm": { + "backend": "openai", + "config": openai_config, + }, + "embedder": { + "backend": "ollama", + "config": { + "model_name_or_path": "nomic-embed-text:latest", + "api_base": os.getenv("OLLAMA_API_BASE", "http://localhost:11434"), + }, + }, + "chunker": { + "backend": "sentence", + "config": { + "tokenizer_or_token_counter": "gpt2", + "chunk_size": 512, + "chunk_overlap": 128, + "min_sentences_per_chunk": 1, + }, + }, + }, + }, + "enable_textual_memory": True, + "enable_activation_memory": os.getenv("ENABLE_ACTIVATION_MEMORY", "false").lower() + == "true", + "top_k": int(os.getenv("MOS_TOP_K", "50")), + "max_turns_window": int(os.getenv("MOS_MAX_TURNS_WINDOW", "20")), + } + + # Add scheduler configuration if enabled + if APIConfig.is_scheduler_enabled(): + config["mem_scheduler"] = APIConfig.get_scheduler_config() + config["enable_mem_scheduler"] = True + else: + config["enable_mem_scheduler"] = False + + return config + + @staticmethod + def get_start_default_config() -> dict[str, Any]: + """Get default configuration for Start API.""" + config = { + "user_id": os.getenv("MOS_USER_ID", "default_user"), + "session_id": os.getenv("MOS_SESSION_ID", "default_session"), + "enable_textual_memory": True, + "enable_activation_memory": os.getenv("ENABLE_ACTIVATION_MEMORY", "false").lower() + == "true", + "top_k": int(os.getenv("MOS_TOP_K", "5")), + "chat_model": { + "backend": os.getenv("MOS_CHAT_MODEL_PROVIDER", "openai"), + "config": { + "model_name_or_path": os.getenv("MOS_CHAT_MODEL", "gpt-4o-mini"), + "api_key": os.getenv("OPENAI_API_KEY", "sk-xxxxxx"), + "temperature": float(os.getenv("MOS_CHAT_TEMPERATURE", 0.7)), + "api_base": os.getenv("OPENAI_API_BASE", "http://xxxxxx:3000/v1"), + "max_tokens": int(os.getenv("MOS_MAX_TOKENS", 1024)), + "top_p": float(os.getenv("MOS_TOP_P", 0.9)), + "top_k": int(os.getenv("MOS_TOP_K", 50)), + "remove_think_prefix": True, + }, + }, + } + + # Add scheduler configuration if enabled + if APIConfig.is_scheduler_enabled(): + config["mem_scheduler"] = APIConfig.get_scheduler_config() + config["enable_mem_scheduler"] = True + else: + config["enable_mem_scheduler"] = False + + return config + + @staticmethod + def create_user_config(user_name: str, user_id: str) -> tuple[MOSConfig, GeneralMemCube]: + """Create configuration for a specific user.""" + openai_config = APIConfig.get_openai_config() + neo4j_config = APIConfig.get_neo4j_config() + qwen_config = APIConfig.qwen_config() + # Create MOSConfig + config_dict = { + "user_id": user_id, + "chat_model": { + "backend": os.getenv("MOS_CHAT_MODEL_PROVIDER", "openai"), + "config": openai_config + if os.getenv("MOS_CHAT_MODEL_PROVIDER", "openai") == "openai" + else qwen_config, + }, + "mem_reader": { + "backend": "simple_struct", + "config": { + "llm": { + "backend": "openai", + "config": openai_config, + }, + "embedder": { + "backend": "ollama", + "config": { + "model_name_or_path": "nomic-embed-text:latest", + "api_base": os.getenv("OLLAMA_API_BASE", "http://localhost:11434"), + }, + }, + "chunker": { + "backend": "sentence", + "config": { + "tokenizer_or_token_counter": "gpt2", + "chunk_size": 512, + "chunk_overlap": 128, + "min_sentences_per_chunk": 1, + }, + }, + }, + }, + "enable_textual_memory": True, + "enable_activation_memory": os.getenv("ENABLE_ACTIVATION_MEMORY", "false").lower() + == "true", + "top_k": 30, + "max_turns_window": 20, + } + + # Add scheduler configuration if enabled + if APIConfig.is_scheduler_enabled(): + config_dict["mem_scheduler"] = APIConfig.get_scheduler_config() + config_dict["enable_mem_scheduler"] = True + else: + config_dict["enable_mem_scheduler"] = False + + default_config = MOSConfig(**config_dict) + + # Create MemCube config + default_cube_config = GeneralMemCubeConfig.model_validate( + { + "user_id": user_id, + "cube_id": f"{user_name}_default_cube", + "text_mem": { + "backend": "tree_text", + "config": { + "extractor_llm": {"backend": "openai", "config": openai_config}, + "dispatcher_llm": {"backend": "openai", "config": openai_config}, + "graph_db": { + "backend": "neo4j", + "config": { + "uri": neo4j_config["uri"], + "user": neo4j_config["user"], + "password": neo4j_config["password"], + "db_name": os.getenv( + "NEO4J_DB_NAME", f"db{user_id.replace('-', '')}" + ), # , replace with + "auto_create": neo4j_config["auto_create"], + }, + }, + "embedder": { + "backend": "ollama", + "config": { + "model_name_or_path": "nomic-embed-text:latest", + "api_base": os.getenv("OLLAMA_API_BASE", "http://localhost:11434"), + }, + }, + }, + }, + "act_mem": {} + if os.getenv("ENABLE_ACTIVATION_MEMORY", "false").lower() == "false" + else APIConfig.get_activation_config(), + "para_mem": {}, + } + ) + + default_mem_cube = GeneralMemCube(default_cube_config) + return default_config, default_mem_cube diff --git a/src/memos/api/exceptions.py b/src/memos/api/exceptions.py new file mode 100644 index 000000000..2fd22ad52 --- /dev/null +++ b/src/memos/api/exceptions.py @@ -0,0 +1,28 @@ +import logging + +from fastapi.requests import Request +from fastapi.responses import JSONResponse + + +logger = logging.getLogger(__name__) + + +class APIExceptionHandler: + """Centralized exception handling for MemOS APIs.""" + + @staticmethod + async def value_error_handler(request: Request, exc: ValueError): + """Handle ValueError exceptions globally.""" + return JSONResponse( + status_code=400, + content={"code": 400, "message": str(exc), "data": None}, + ) + + @staticmethod + async def global_exception_handler(request: Request, exc: Exception): + """Handle all unhandled exceptions globally.""" + logger.exception("Unhandled error:") + return JSONResponse( + status_code=500, + content={"code": 500, "message": str(exc), "data": None}, + ) diff --git a/src/memos/api/product_api.py b/src/memos/api/product_api.py new file mode 100644 index 000000000..8823ffc91 --- /dev/null +++ b/src/memos/api/product_api.py @@ -0,0 +1,30 @@ +import logging + +from fastapi import FastAPI + +from memos.api.exceptions import APIExceptionHandler +from memos.api.routers.product_router import router as product_router + + +# Configure logging +logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") +logger = logging.getLogger(__name__) + +app = FastAPI( + title="MemOS Product REST APIs", + description="A REST API for managing multiple users with MemOS Product.", + version="1.0.0", +) + +# Include routers +app.include_router(product_router) + +# Exception handlers +app.exception_handler(ValueError)(APIExceptionHandler.value_error_handler) +app.exception_handler(Exception)(APIExceptionHandler.global_exception_handler) + + +if __name__ == "__main__": + import uvicorn + + uvicorn.run(app, host="0.0.0.0", port=8001) diff --git a/src/memos/api/product_models.py b/src/memos/api/product_models.py new file mode 100644 index 000000000..f8008c708 --- /dev/null +++ b/src/memos/api/product_models.py @@ -0,0 +1,152 @@ +import uuid + +from typing import Generic, Literal, TypeAlias, TypeVar + +from pydantic import BaseModel, Field +from typing_extensions import TypedDict + + +T = TypeVar("T") + + +# ─── Message Types ────────────────────────────────────────────────────────────── + +# Chat message roles +MessageRole: TypeAlias = Literal["user", "assistant", "system"] + + +# Message structure +class MessageDict(TypedDict): + """Typed dictionary for chat message dictionaries.""" + + role: MessageRole + content: str + + +class BaseRequest(BaseModel): + """Base model for all requests.""" + + +class BaseResponse(BaseModel, Generic[T]): + """Base model for all responses.""" + + code: int = Field(200, description="Response status code") + message: str = Field(..., description="Response message") + data: T | None = Field(None, description="Response data") + + +# Product API Models +class UserRegisterRequest(BaseRequest): + """Request model for user registration.""" + + user_id: str = Field( + default_factory=lambda: str(uuid.uuid4()), description="User ID for registration" + ) + user_name: str | None = Field(None, description="User name for registration") + interests: str | None = Field(None, description="User interests") + + +class GetMemoryRequest(BaseRequest): + """Request model for getting memories.""" + + user_id: str = Field(..., description="User ID") + memory_type: Literal["text_mem", "act_mem", "param_mem", "para_mem"] = Field( + ..., description="Memory type" + ) + mem_cube_ids: list[str] | None = Field(None, description="Cube IDs") + search_query: str | None = Field(None, description="Search query") + + +# Start API Models +class Message(BaseModel): + role: str = Field(..., description="Role of the message (user or assistant).") + content: str = Field(..., description="Message content.") + + +class MemoryCreate(BaseRequest): + user_id: str = Field(..., description="User ID") + messages: list[Message] | None = Field(None, description="List of messages to store.") + memory_content: str | None = Field(None, description="Content to store as memory") + doc_path: str | None = Field(None, description="Path to document to store") + mem_cube_id: str | None = Field(None, description="ID of the memory cube") + + +class MemCubeRegister(BaseRequest): + mem_cube_name_or_path: str = Field(..., description="Name or path of the MemCube to register.") + mem_cube_id: str | None = Field(None, description="ID for the MemCube") + + +class ChatRequest(BaseRequest): + """Request model for chat operations.""" + + user_id: str = Field(..., description="User ID") + query: str = Field(..., description="Chat query message") + mem_cube_id: str | None = Field(None, description="Cube ID to use for chat") + history: list[MessageDict] | None = Field(None, description="Chat history") + + +class UserCreate(BaseRequest): + user_name: str | None = Field(None, description="Name of the user") + role: str = Field("user", description="Role of the user") + user_id: str = Field(..., description="User ID") + + +class CubeShare(BaseRequest): + target_user_id: str = Field(..., description="Target user ID to share with") + + +# Response Models +class SimpleResponse(BaseResponse[None]): + """Simple response model for operations without data return.""" + + +class UserRegisterResponse(BaseResponse[dict]): + """Response model for user registration.""" + + +class MemoryResponse(BaseResponse[list]): + """Response model for memory operations.""" + + +class SuggestionResponse(BaseResponse[list]): + """Response model for suggestion operations.""" + + data: dict[str, list[str]] | None = Field(None, description="Response data") + + +class ConfigResponse(BaseResponse[None]): + """Response model for configuration endpoint.""" + + +class SearchResponse(BaseResponse[dict]): + """Response model for search operations.""" + + +class ChatResponse(BaseResponse[str]): + """Response model for chat operations.""" + + +class UserResponse(BaseResponse[dict]): + """Response model for user operations.""" + + +class UserListResponse(BaseResponse[list]): + """Response model for user list operations.""" + + +class MemoryCreateRequest(BaseRequest): + """Request model for creating memories.""" + + user_id: str = Field(..., description="User ID") + messages: list[MessageDict] | None = Field(None, description="List of messages to store.") + memory_content: str | None = Field(None, description="Memory content to store") + doc_path: str | None = Field(None, description="Path to document to store") + mem_cube_id: str | None = Field(None, description="Cube ID") + + +class SearchRequest(BaseRequest): + """Request model for searching memories.""" + + user_id: str = Field(..., description="User ID") + query: str = Field(..., description="Search query") + mem_cube_id: str | None = Field(None, description="Cube ID to search in") diff --git a/src/memos/api/routers/__init__.py b/src/memos/api/routers/__init__.py new file mode 100644 index 000000000..40ed96f43 --- /dev/null +++ b/src/memos/api/routers/__init__.py @@ -0,0 +1 @@ +# API routers module diff --git a/src/memos/api/routers/product_router.py b/src/memos/api/routers/product_router.py new file mode 100644 index 000000000..841954947 --- /dev/null +++ b/src/memos/api/routers/product_router.py @@ -0,0 +1,321 @@ +import json +import logging +import traceback + +from fastapi import APIRouter, HTTPException +from fastapi.responses import StreamingResponse + +from memos.api.config import APIConfig +from memos.api.product_models import ( + BaseResponse, + ChatRequest, + GetMemoryRequest, + MemoryCreateRequest, + MemoryResponse, + SearchRequest, + SearchResponse, + SimpleResponse, + SuggestionResponse, + UserRegisterRequest, + UserRegisterResponse, +) +from memos.configs.mem_os import MOSConfig +from memos.mem_os.product import MOSProduct + + +logger = logging.getLogger(__name__) + +router = APIRouter(prefix="/product", tags=["Product API"]) + +# Initialize MOSProduct instance with lazy initialization +MOS_PRODUCT_INSTANCE = None + + +def get_mos_product_instance(): + """Get or create MOSProduct instance.""" + global MOS_PRODUCT_INSTANCE + if MOS_PRODUCT_INSTANCE is None: + default_config = APIConfig.get_product_default_config() + from memos.configs.mem_os import MOSConfig + + mos_config = MOSConfig(**default_config) + MOS_PRODUCT_INSTANCE = MOSProduct(default_config=mos_config) + logger.info("MOSProduct instance created successfully with inheritance architecture") + return MOS_PRODUCT_INSTANCE + + +get_mos_product_instance() + + +@router.post("/configure", summary="Configure MOSProduct", response_model=SimpleResponse) +async def set_config(config): + """Set MOSProduct configuration.""" + global MOS_PRODUCT_INSTANCE + MOS_PRODUCT_INSTANCE = MOSProduct(default_config=config) + return SimpleResponse(message="Configuration set successfully") + + +@router.post("/users/register", summary="Register a new user", response_model=UserRegisterResponse) +async def register_user(user_req: UserRegisterRequest): + """Register a new user with configuration and default cube.""" + try: + # Get configuration for the user + user_config, default_mem_cube = APIConfig.create_user_config( + user_name=user_req.user_id, user_id=user_req.user_id + ) + mos_product = get_mos_product_instance() + # Register user with default config and mem cube + result = mos_product.user_register( + user_id=user_req.user_id, + user_name=user_req.user_name, + interests=user_req.interests, + config=user_config, + default_mem_cube=default_mem_cube, + ) + + if result["status"] == "success": + return UserRegisterResponse( + message="User registered successfully", + data={"user_id": result["user_id"], "mem_cube_id": result["default_cube_id"]}, + ) + else: + raise HTTPException(status_code=400, detail=result["message"]) + + except Exception as err: + logger.error(f"Failed to register user: {traceback.format_exc()}") + raise HTTPException(status_code=500, detail=str(traceback.format_exc())) from err + + +@router.get( + "/suggestions/{user_id}", summary="Get suggestion queries", response_model=SuggestionResponse +) +async def get_suggestion_queries(user_id: str): + """Get suggestion queries for a specific user.""" + try: + mos_product = get_mos_product_instance() + suggestions = mos_product.get_suggestion_query(user_id) + return SuggestionResponse( + message="Suggestions retrieved successfully", data={"query": suggestions} + ) + except ValueError as err: + raise HTTPException(status_code=404, detail=str(traceback.format_exc())) from err + except Exception as err: + logger.error(f"Failed to get suggestions: {traceback.format_exc()}") + raise HTTPException(status_code=500, detail=str(traceback.format_exc())) from err + + +@router.post("/get_all", summary="Get all memories for user", response_model=MemoryResponse) +async def get_all_memories(memory_req: GetMemoryRequest): + """Get all memories for a specific user.""" + try: + mos_product = get_mos_product_instance() + if memory_req.search_query: + result = mos_product.get_subgraph( + user_id=memory_req.user_id, + query=memory_req.search_query, + mem_cube_ids=memory_req.mem_cube_ids, + ) + return MemoryResponse(message="Memories retrieved successfully", data=result) + else: + result = mos_product.get_all( + user_id=memory_req.user_id, + memory_type=memory_req.memory_type, + mem_cube_ids=memory_req.mem_cube_ids, + ) + return MemoryResponse(message="Memories retrieved successfully", data=result) + + except ValueError as err: + raise HTTPException(status_code=404, detail=str(traceback.format_exc())) from err + except Exception as err: + logger.error(f"Failed to get memories: {traceback.format_exc()}") + raise HTTPException(status_code=500, detail=str(traceback.format_exc())) from err + + +@router.post("/add", summary="add a new memory", response_model=SimpleResponse) +async def create_memory(memory_req: MemoryCreateRequest): + """Create a new memory for a specific user.""" + try: + mos_product = get_mos_product_instance() + mos_product.add( + user_id=memory_req.user_id, + memory_content=memory_req.memory_content, + messages=memory_req.messages, + doc_path=memory_req.doc_path, + mem_cube_id=memory_req.mem_cube_id, + ) + return SimpleResponse(message="Memory created successfully") + + except ValueError as err: + raise HTTPException(status_code=404, detail=str(traceback.format_exc())) from err + except Exception as err: + logger.error(f"Failed to create memory: {traceback.format_exc()}") + raise HTTPException(status_code=500, detail=str(traceback.format_exc())) from err + + +@router.post("/search", summary="Search memories", response_model=SearchResponse) +async def search_memories(search_req: SearchRequest): + """Search memories for a specific user.""" + try: + mos_product = get_mos_product_instance() + result = mos_product.search( + query=search_req.query, + user_id=search_req.user_id, + install_cube_ids=[search_req.mem_cube_id] if search_req.mem_cube_id else None, + ) + return SearchResponse(message="Search completed successfully", data=result) + + except ValueError as err: + raise HTTPException(status_code=404, detail=str(traceback.format_exc())) from err + except Exception as err: + logger.error(f"Failed to search memories: {traceback.format_exc()}") + raise HTTPException(status_code=500, detail=str(traceback.format_exc())) from err + + +@router.post("/chat", summary="Chat with MemOS") +async def chat(chat_req: ChatRequest): + """Chat with MemOS for a specific user. Returns SSE stream.""" + try: + mos_product = get_mos_product_instance() + + def generate_chat_response(): + """Generate chat response as SSE stream.""" + try: + yield from mos_product.chat_with_references( + query=chat_req.query, + user_id=chat_req.user_id, + cube_id=chat_req.mem_cube_id, + history=chat_req.history, + ) + except Exception as e: + logger.error(f"Error in chat stream: {e}") + error_data = f"data: {json.dumps({'type': 'error', 'content': str(traceback.format_exc())})}\n\n" + yield error_data + + return StreamingResponse( + generate_chat_response(), + media_type="text/plain", + headers={ + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "Content-Type": "text/event-stream", + }, + ) + + except ValueError as err: + raise HTTPException(status_code=404, detail=str(traceback.format_exc())) from err + except Exception as err: + logger.error(f"Failed to start chat: {traceback.format_exc()}") + raise HTTPException(status_code=500, detail=str(traceback.format_exc())) from err + + +@router.get("/users", summary="List all users", response_model=BaseResponse[list]) +async def list_users(): + """List all registered users.""" + try: + mos_product = get_mos_product_instance() + users = mos_product.list_users() + return BaseResponse(message="Users retrieved successfully", data=users) + except Exception as err: + logger.error(f"Failed to list users: {traceback.format_exc()}") + raise HTTPException(status_code=500, detail=str(traceback.format_exc())) from err + + +@router.get("/users/{user_id}", summary="Get user info", response_model=BaseResponse[dict]) +async def get_user_info(user_id: str): + """Get user information including accessible cubes.""" + try: + mos_product = get_mos_product_instance() + user_info = mos_product.get_user_info(user_id) + return BaseResponse(message="User info retrieved successfully", data=user_info) + except ValueError as err: + raise HTTPException(status_code=404, detail=str(traceback.format_exc())) from err + except Exception as err: + logger.error(f"Failed to get user info: {traceback.format_exc()}") + raise HTTPException(status_code=500, detail=str(traceback.format_exc())) from err + + +@router.get( + "/configure/{user_id}", summary="Get MOSProduct configuration", response_model=SimpleResponse +) +async def get_config(user_id: str): + """Get MOSProduct configuration.""" + global MOS_PRODUCT_INSTANCE + config = MOS_PRODUCT_INSTANCE.default_config + return SimpleResponse(message="Configuration retrieved successfully", data=config) + + +@router.get( + "/users/{user_id}/config", summary="Get user configuration", response_model=BaseResponse[dict] +) +async def get_user_config(user_id: str): + """Get user-specific configuration.""" + try: + mos_product = get_mos_product_instance() + config = mos_product.get_user_config(user_id) + if config: + return BaseResponse( + message="User configuration retrieved successfully", + data=config.model_dump(mode="json"), + ) + else: + raise HTTPException( + status_code=404, detail=f"Configuration not found for user {user_id}" + ) + except ValueError as err: + raise HTTPException(status_code=404, detail=str(traceback.format_exc())) from err + except Exception as err: + logger.error(f"Failed to get user config: {traceback.format_exc()}") + raise HTTPException(status_code=500, detail=str(traceback.format_exc())) from err + + +@router.put( + "/users/{user_id}/config", summary="Update user configuration", response_model=SimpleResponse +) +async def update_user_config(user_id: str, config_data: dict): + """Update user-specific configuration.""" + try: + mos_product = get_mos_product_instance() + + # Create MOSConfig from the provided data + config = MOSConfig(**config_data) + + # Update the configuration + success = mos_product.update_user_config(user_id, config) + if success: + return SimpleResponse(message="User configuration updated successfully") + else: + raise HTTPException(status_code=500, detail="Failed to update user configuration") + + except ValueError as err: + raise HTTPException(status_code=400, detail=str(traceback.format_exc())) from err + except Exception as err: + logger.error(f"Failed to update user config: {traceback.format_exc()}") + raise HTTPException(status_code=500, detail=str(traceback.format_exc())) from err + + +@router.get( + "/instances/status", summary="Get user configuration status", response_model=BaseResponse[dict] +) +async def get_instance_status(): + """Get information about active user configurations in memory.""" + try: + mos_product = get_mos_product_instance() + status_info = mos_product.get_user_instance_info() + return BaseResponse( + message="User configuration status retrieved successfully", data=status_info + ) + except Exception as err: + logger.error(f"Failed to get user configuration status: {traceback.format_exc()}") + raise HTTPException(status_code=500, detail=str(traceback.format_exc())) from err + + +@router.get("/instances/count", summary="Get active user count", response_model=BaseResponse[int]) +async def get_active_user_count(): + """Get the number of active user configurations in memory.""" + try: + mos_product = get_mos_product_instance() + count = mos_product.get_active_user_count() + return BaseResponse(message="Active user count retrieved successfully", data=count) + except Exception as err: + logger.error(f"Failed to get active user count: {traceback.format_exc()}") + raise HTTPException(status_code=500, detail=str(traceback.format_exc())) from err diff --git a/src/memos/configs/llm.py b/src/memos/configs/llm.py index 26beee2a0..1f75e9948 100644 --- a/src/memos/configs/llm.py +++ b/src/memos/configs/llm.py @@ -54,6 +54,7 @@ class LLMConfigFactory(BaseConfig): "openai": OpenAILLMConfig, "ollama": OllamaLLMConfig, "huggingface": HFLLMConfig, + "huggingface_singleton": HFLLMConfig, # Add singleton support } @field_validator("backend") diff --git a/src/memos/configs/memory.py b/src/memos/configs/memory.py index c728736ca..a0406d7f9 100644 --- a/src/memos/configs/memory.py +++ b/src/memos/configs/memory.py @@ -51,9 +51,9 @@ class KVCacheMemoryConfig(BaseActMemoryConfig): @classmethod def validate_extractor_llm(cls, extractor_llm: LLMConfigFactory) -> LLMConfigFactory: """Validate the extractor_llm field.""" - if extractor_llm.backend != "huggingface": + if extractor_llm.backend not in ["huggingface", "huggingface_singleton"]: raise ConfigurationError( - f"KVCacheMemoryConfig requires extractor_llm backend to be 'huggingface', got '{extractor_llm.backend}'" + f"KVCacheMemoryConfig requires extractor_llm backend to be 'huggingface' or 'huggingface_singleton', got '{extractor_llm.backend}'" ) return extractor_llm @@ -83,9 +83,9 @@ class LoRAMemoryConfig(BaseParaMemoryConfig): @classmethod def validate_extractor_llm(cls, extractor_llm: LLMConfigFactory) -> LLMConfigFactory: """Validate the extractor_llm field.""" - if extractor_llm.backend not in ["huggingface"]: + if extractor_llm.backend not in ["huggingface", "huggingface_singleton"]: raise ConfigurationError( - f"LoRAMemoryConfig requires extractor_llm backend to be 'huggingface', got '{extractor_llm.backend}'" + f"LoRAMemoryConfig requires extractor_llm backend to be 'huggingface' or 'huggingface_singleton', got '{extractor_llm.backend}'" ) return extractor_llm diff --git a/src/memos/llms/factory.py b/src/memos/llms/factory.py index 4435508ba..e73d94283 100644 --- a/src/memos/llms/factory.py +++ b/src/memos/llms/factory.py @@ -3,6 +3,7 @@ from memos.configs.llm import LLMConfigFactory from memos.llms.base import BaseLLM from memos.llms.hf import HFLLM +from memos.llms.hf_singleton import HFSingletonLLM from memos.llms.ollama import OllamaLLM from memos.llms.openai import OpenAILLM @@ -14,6 +15,7 @@ class LLMFactory(BaseLLM): "openai": OpenAILLM, "ollama": OllamaLLM, "huggingface": HFLLM, + "huggingface_singleton": HFSingletonLLM, # Add singleton version } @classmethod diff --git a/src/memos/llms/hf_singleton.py b/src/memos/llms/hf_singleton.py new file mode 100644 index 000000000..af0b6deab --- /dev/null +++ b/src/memos/llms/hf_singleton.py @@ -0,0 +1,114 @@ +import threading + +from typing import ClassVar + +from memos.configs.llm import HFLLMConfig +from memos.llms.hf import HFLLM +from memos.log import get_logger + + +logger = get_logger(__name__) + + +class HFSingletonLLM(HFLLM): + """ + Singleton version of HFLLM that prevents multiple loading of the same model. + This class inherits from HFLLM and adds singleton behavior. + """ + + _instances: ClassVar[dict[str, "HFSingletonLLM"]] = {} + _lock: ClassVar[threading.Lock] = threading.Lock() + + def __new__(cls, config: HFLLMConfig): + """ + Singleton pattern implementation. + Returns existing instance if config already exists, otherwise creates new one. + """ + config_key = cls._get_config_key(config) + + if config_key in cls._instances: + logger.debug(f"Reusing existing HF model: {config.model_name_or_path}") + return cls._instances[config_key] + + with cls._lock: + # Double-check pattern to prevent race conditions + if config_key in cls._instances: + logger.debug(f"Reusing existing HF model: {config.model_name_or_path}") + return cls._instances[config_key] + + logger.info(f"Creating new HF model: {config.model_name_or_path}") + instance = super().__new__(cls) + cls._instances[config_key] = instance + return instance + + def __init__(self, config: HFLLMConfig): + """ + Initialize the singleton HFLLM instance. + Only initializes if this is a new instance. + """ + # Check if already initialized + if hasattr(self, "_initialized"): + return + + # Call parent constructor + super().__init__(config) + self._initialized = True + + @classmethod + def _get_config_key(cls, config: HFLLMConfig) -> str: + """ + Generate a unique key for the HF model configuration. + + Args: + config: The HFLLM configuration + + Returns: + A unique string key representing the configuration + """ + # Create a unique key based on model path and key parameters + key_parts = [config.model_name_or_path] + return "|".join(key_parts) + + @classmethod + def get_instance_count(cls) -> int: + """ + Get the number of unique HF model instances currently managed. + + Returns: + Number of HF model instances + """ + return len(cls._instances) + + @classmethod + def get_instance_info(cls) -> dict[str, str]: + """ + Get information about all managed HF model instances. + + Returns: + Dictionary mapping config keys to model paths + """ + return {key: instance.config.model_name_or_path for key, instance in cls._instances.items()} + + @classmethod + def clear_all(cls) -> None: + """ + Clear all HF model instances from memory. + This should be used carefully as it will force reloading of models. + """ + with cls._lock: + cls._instances.clear() + logger.info("All HF model instances cleared from singleton manager") + + +# Convenience function to get singleton manager info +def get_hf_singleton_info() -> dict[str, int]: + """ + Get information about the HF singleton manager. + + Returns: + Dictionary with instance count and info + """ + return { + "instance_count": HFSingletonLLM.get_instance_count(), + "instance_info": HFSingletonLLM.get_instance_info(), + } diff --git a/src/memos/mem_os/core.py b/src/memos/mem_os/core.py index 07fa4ba97..684d2d58c 100644 --- a/src/memos/mem_os/core.py +++ b/src/memos/mem_os/core.py @@ -30,7 +30,7 @@ class MOSCore: MOSCore acts as an operating system layer for handling and orchestrating MemCube instances. """ - def __init__(self, config: MOSConfig): + def __init__(self, config: MOSConfig, user_manager: UserManager | None = None): self.config = config self.user_id = config.user_id self.session_id = config.session_id @@ -39,7 +39,12 @@ def __init__(self, config: MOSConfig): self.mem_reader = MemReaderFactory.from_config(config.mem_reader) self.chat_history_manager: dict[str, ChatHistory] = {} self._register_chat_history() - self.user_manager = UserManager(user_id=self.user_id if self.user_id else "root") + + # Use provided user_manager or create a new one + if user_manager is not None: + self.user_manager = user_manager + else: + self.user_manager = UserManager(user_id=self.user_id if self.user_id else "root") # Validate user exists if not self.user_manager.validate_user(self.user_id): @@ -427,7 +432,11 @@ def unregister_mem_cube(self, mem_cube_id: str, user_id: str | None = None) -> N raise ValueError(f"MemCube with ID {mem_cube_id} does not exist.") def search( - self, query: str, user_id: str | None = None, install_cube_ids: list[str] | None = None + self, + query: str, + user_id: str | None = None, + install_cube_ids: list[str] | None = None, + top_k: int | None = None, ) -> MOSSearchResult: """ Search for textual memories across all registered MemCubes. @@ -464,18 +473,10 @@ def search( and (mem_cube.text_mem is not None) and self.config.enable_textual_memory ): - memories = mem_cube.text_mem.search(query, top_k=self.config.top_k) - result["text_mem"].append({"cube_id": mem_cube_id, "memories": memories}) - logger.info( - f"🧠 [Memory] Searched memories from {mem_cube_id}:\n{self._str_memories(memories)}\n" + memories = mem_cube.text_mem.search( + query, top_k=top_k if top_k else self.config.top_k ) - if ( - (mem_cube_id in install_cube_ids) - and (mem_cube.act_mem is not None) - and self.config.enable_activation_memory - ): - memories = mem_cube.act_mem.extract(query) - result["act_mem"].append({"cube_id": mem_cube_id, "memories": [memories]}) + result["text_mem"].append({"cube_id": mem_cube_id, "memories": memories}) logger.info( f"🧠 [Memory] Searched memories from {mem_cube_id}:\n{self._str_memories(memories)}\n" ) diff --git a/src/memos/mem_os/product.py b/src/memos/mem_os/product.py index b210c30dd..4c551640a 100644 --- a/src/memos/mem_os/product.py +++ b/src/memos/mem_os/product.py @@ -1,33 +1,553 @@ import json +import os +import time from collections.abc import Generator -from typing import Literal +from datetime import datetime +from typing import Any, Literal + +from transformers import AutoTokenizer from memos.configs.mem_os import MOSConfig +from memos.log import get_logger +from memos.mem_cube.general import GeneralMemCube from memos.mem_os.core import MOSCore -from memos.memories.activation.item import ActivationMemoryItem -from memos.memories.parametric.item import ParametricMemoryItem -from memos.memories.textual.item import TextualMemoryMetadata, TreeNodeTextualMemoryMetadata +from memos.mem_os.utils.format_utils import ( + convert_activation_memory_to_serializable, + convert_graph_to_tree_forworkmem, + filter_nodes_by_tree_ids, + remove_embedding_recursive, + sort_children_by_memory_type, +) +from memos.mem_scheduler.modules.schemas import ANSWER_LABEL, QUERY_LABEL, ScheduleMessageItem +from memos.mem_user.persistent_user_manager import PersistentUserManager, UserRole +from memos.memories.textual.item import ( + TextualMemoryItem, +) from memos.types import MessageList +logger = get_logger(__name__) + +CUBE_PATH = "/tmp/data" +with open("./tmp/fake_data.json") as f: + MOCK_DATA = json.loads(f.read()) + +# Removed ensure_user_instance decorator as it's redundant with MOSCore's built-in validation + + class MOSProduct(MOSCore): """ - The MOSProduct class inherits from MOSCore mainly for product usage. + The MOSProduct class inherits from MOSCore and manages multiple users. + Each user has their own configuration and cube access, but shares the same model instances. """ - def __init__(self, config: MOSConfig): - super().__init__(config) + def __init__(self, default_config: MOSConfig | None = None, max_user_instances: int = 100): + """ + Initialize MOSProduct with an optional default configuration. + + Args: + default_config (MOSConfig | None): Default configuration for new users + max_user_instances (int): Maximum number of user instances to keep in memory + """ + # Initialize with a root config for shared resources + if default_config is None: + # Create a minimal config for root user + root_config = MOSConfig( + user_id="root", + session_id="root_session", + chat_model=default_config.chat_model if default_config else None, + mem_reader=default_config.mem_reader if default_config else None, + enable_mem_scheduler=default_config.enable_mem_scheduler + if default_config + else False, + mem_scheduler=default_config.mem_scheduler if default_config else None, + ) + else: + root_config = default_config.model_copy(deep=True) + root_config.user_id = "root" + root_config.session_id = "root_session" + + # Initialize parent MOSCore with root config + super().__init__(root_config) + + # Product-specific attributes + self.default_config = default_config + self.max_user_instances = max_user_instances + + # User-specific data structures + self.user_configs: dict[str, MOSConfig] = {} + self.user_cube_access: dict[str, set[str]] = {} # user_id -> set of cube_ids + self.user_chat_histories: dict[str, dict] = {} + + # Use PersistentUserManager for user management + self.global_user_manager = PersistentUserManager(user_id="root") + + # Initialize tiktoken for streaming + try: + # Use gpt2 encoding which is more stable and widely compatible + self.tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B") + logger.info("tokenizer initialized successfully for streaming") + except Exception as e: + logger.warning( + f"Failed to initialize tokenizer, will use character-based chunking: {e}" + ) + self.tokenizer = None + + # Restore user instances from persistent storage + self._restore_user_instances() + logger.info(f"User instances restored successfully, now user is {self.mem_cubes.keys()}") + + def _restore_user_instances(self) -> None: + """Restore user instances from persistent storage after service restart.""" + try: + # Get all user configurations from persistent storage + user_configs = self.global_user_manager.list_user_configs() + + # Get the raw database records for sorting by updated_at + session = self.global_user_manager._get_session() + try: + from memos.mem_user.persistent_user_manager import UserConfig + + db_configs = session.query(UserConfig).all() + # Create a mapping of user_id to updated_at timestamp + updated_at_map = {config.user_id: config.updated_at for config in db_configs} + + # Sort by updated_at timestamp (most recent first) and limit by max_instances + sorted_configs = sorted( + user_configs.items(), key=lambda x: updated_at_map.get(x[0], ""), reverse=True + )[: self.max_user_instances] + finally: + session.close() + + for user_id, config in sorted_configs: + if user_id != "root": # Skip root user + try: + # Store user config and cube access + self.user_configs[user_id] = config + self._load_user_cube_access(user_id) + logger.info(f"Restored user configuration for {user_id}") + + except Exception as e: + logger.error(f"Failed to restore user configuration for {user_id}: {e}") + + except Exception as e: + logger.error(f"Error during user instance restoration: {e}") + + def _ensure_user_instance(self, user_id: str, max_instances: int | None = None) -> None: + """ + Ensure user configuration exists, creating it if necessary. + + Args: + user_id (str): The user ID + max_instances (int): Maximum instances to keep in memory (overrides class default) + """ + if user_id in self.user_configs: + return + + # Try to get config from persistent storage first + stored_config = self.global_user_manager.get_user_config(user_id) + if stored_config: + self.user_configs[user_id] = stored_config + self._load_user_cube_access(user_id) + else: + # Use default config + if not self.default_config: + raise ValueError(f"No configuration available for user {user_id}") + user_config = self.default_config.model_copy(deep=True) + user_config.user_id = user_id + user_config.session_id = f"{user_id}_session" + self.user_configs[user_id] = user_config + self._load_user_cube_access(user_id) + + # Apply LRU eviction if needed + max_instances = max_instances or self.max_user_instances + if len(self.user_configs) > max_instances: + # Remove least recently used instance (excluding root) + user_ids = [uid for uid in self.user_configs if uid != "root"] + if user_ids: + oldest_user_id = user_ids[0] + del self.user_configs[oldest_user_id] + if oldest_user_id in self.user_cube_access: + del self.user_cube_access[oldest_user_id] + logger.info(f"Removed least recently used user configuration: {oldest_user_id}") + + def _load_user_cube_access(self, user_id: str) -> None: + """Load user's cube access permissions.""" + try: + # Get user's accessible cubes from persistent storage + accessible_cubes = self.global_user_manager.get_user_cube_access(user_id) + self.user_cube_access[user_id] = set(accessible_cubes) + except Exception as e: + logger.warning(f"Failed to load cube access for user {user_id}: {e}") + self.user_cube_access[user_id] = set() + + def _get_user_config(self, user_id: str) -> MOSConfig: + """Get user configuration.""" + if user_id not in self.user_configs: + self._ensure_user_instance(user_id) + return self.user_configs[user_id] + + def _validate_user_cube_access(self, user_id: str, cube_id: str) -> None: + """Validate user has access to the cube.""" + if user_id not in self.user_cube_access: + self._load_user_cube_access(user_id) + + if cube_id not in self.user_cube_access.get(user_id, set()): + raise ValueError(f"User '{user_id}' does not have access to cube '{cube_id}'") + + def _validate_user_access(self, user_id: str, cube_id: str | None = None) -> None: + """Validate user access using MOSCore's built-in validation.""" + # Use MOSCore's built-in user validation + if cube_id: + self._validate_cube_access(user_id, cube_id) + else: + self._validate_user_exists(user_id) + + def _create_user_config(self, user_id: str, config: MOSConfig) -> MOSConfig: + """Create a new user configuration.""" + # Create a copy of config with the specific user_id + user_config = config.model_copy(deep=True) + user_config.user_id = user_id + user_config.session_id = f"{user_id}_session" + + # Save configuration to persistent storage + self.global_user_manager.save_user_config(user_id, user_config) + + return user_config + + def _get_or_create_user_config( + self, user_id: str, config: MOSConfig | None = None + ) -> MOSConfig: + """Get existing user config or create a new one.""" + if user_id in self.user_configs: + return self.user_configs[user_id] + + # Try to get config from persistent storage first + stored_config = self.global_user_manager.get_user_config(user_id) + if stored_config: + return self._create_user_config(user_id, stored_config) + + # Use provided config or default config + user_config = config or self.default_config + if not user_config: + raise ValueError(f"No configuration provided for user {user_id}") + + return self._create_user_config(user_id, user_config) + + def _load_user_cubes(self, user_id: str) -> None: + """Load all cubes for a user into memory.""" + # Get user's accessible cubes from persistent storage + accessible_cubes = self.global_user_manager.get_user_cubes(user_id) + + for cube in accessible_cubes[:1]: + if cube.cube_id not in self.mem_cubes: + try: + if cube.cube_path and os.path.exists(cube.cube_path): + # Use MOSCore's register_mem_cube method directly + self.register_mem_cube(cube.cube_path, cube.cube_id, user_id) + else: + logger.warning( + f"Cube path {cube.cube_path} does not exist for cube {cube.cube_id}" + ) + except Exception as e: + logger.error(f"Failed to load cube {cube.cube_id} for user {user_id}: {e}") + + def _build_system_prompt(self, user_id: str, memories_all: list[TextualMemoryItem]) -> str: + """ + Build custom system prompt for the user with memory references. + + Args: + user_id (str): The user ID. + memories (list[TextualMemoryItem]): The memories to build the system prompt. + + Returns: + str: The custom system prompt. + """ + + # Build base prompt + base_prompt = ( + "You are a knowledgeable and helpful AI assistant with access to user memories. " + "When responding to user queries, you should reference relevant memories using the provided memory IDs. " + "Use the reference format: [1-n:memoriesID] " + "where refid is a sequential number starting from 1 and increments for each reference in your response, " + "and memoriesID is the specific memory ID provided in the available memories list. " + "For example: [1:abc123], [2:def456], [3:ghi789], [4:jkl101], [5:mno112] " + "Only reference memories that are directly relevant to the user's question. " + "Make your responses natural and conversational while incorporating memory references when appropriate." + ) + + # Add memory context if available + if memories_all: + memory_context = "\n\n## Available ID Memories:\n" + for i, memory in enumerate(memories_all, 1): + # Format: [memory_id]: memory_content + memory_id = f"{memory.id.split('-')[0]}" if hasattr(memory, "id") else f"mem_{i}" + memory_content = memory.memory if hasattr(memory, "memory") else str(memory) + memory_context += f"{memory_id}: {memory_content}\n" + return base_prompt + memory_context + + return base_prompt + + def _process_streaming_references_complete(self, text_buffer: str) -> tuple[str, str]: + """ + Complete streaming reference processing to ensure reference tags are never split. + + Args: + text_buffer (str): The accumulated text buffer. + + Returns: + tuple[str, str]: (processed_text, remaining_buffer) + """ + import re + + # Pattern to match complete reference tags: [refid:memoriesID] + complete_pattern = r"\[\d+:[^\]]+\]" + + # Find all complete reference tags + complete_matches = list(re.finditer(complete_pattern, text_buffer)) + + if complete_matches: + # Find the last complete tag + last_match = complete_matches[-1] + end_pos = last_match.end() + + # Return text up to the end of the last complete tag + processed_text = text_buffer[:end_pos] + remaining_buffer = text_buffer[end_pos:] + return processed_text, remaining_buffer + + # Check for incomplete reference tags + # Look for opening bracket with number and colon + opening_pattern = r"\[\d+:" + opening_matches = list(re.finditer(opening_pattern, text_buffer)) + + if opening_matches: + # Find the last opening tag + last_opening = opening_matches[-1] + opening_start = last_opening.start() + + # Check if we have a complete opening pattern + if last_opening.end() <= len(text_buffer): + # We have a complete opening pattern, keep everything in buffer + return "", text_buffer + else: + # Incomplete opening pattern, return text before it + return text_buffer[:opening_start], text_buffer[opening_start:] + + # Check for partial opening pattern (starts with [ but not complete) + if "[" in text_buffer: + ref_start = text_buffer.find("[") + return text_buffer[:ref_start], text_buffer[ref_start:] + + # No reference tags found, return all text + return text_buffer, "" + + def _extract_references_from_response(self, response: str) -> list[dict]: + """ + Extract reference information from the response. + + Args: + response (str): The complete response text. + + Returns: + list[dict]: List of reference information. + """ + import re + + references = [] + # Pattern to match [refid:memoriesID] + pattern = r"\[(\d+):([^\]]+)\]" + + matches = re.findall(pattern, response) + for ref_number, memory_id in matches: + references.append({"memory_id": memory_id, "reference_number": int(ref_number)}) + + return references + + def _chunk_response_with_tiktoken( + self, response: str, chunk_size: int = 5 + ) -> Generator[str, None, None]: + """ + Chunk response using tiktoken for proper token-based streaming. + + Args: + response (str): The response text to chunk. + chunk_size (int): Number of tokens per chunk. + + Yields: + str: Chunked text pieces. + """ + if self.tokenizer: + # Use tiktoken for proper token-based chunking + print(response) + tokens = self.tokenizer.encode(response) + + for i in range(0, len(tokens), chunk_size): + token_chunk = tokens[i : i + chunk_size] + chunk_text = self.tokenizer.decode(token_chunk) + yield chunk_text + else: + # Fallback to character-based chunking + char_chunk_size = chunk_size * 4 # Approximate character to token ratio + for i in range(0, len(response), char_chunk_size): + yield response[i : i + char_chunk_size] + + def _send_message_to_scheduler( + self, + user_id: str, + mem_cube_id: str, + query: str, + label: str, + ): + """ + Send message to scheduler. + args: + user_id: str, + mem_cube_id: str, + query: str, + """ + + if self.enable_mem_scheduler and (self.mem_scheduler is not None): + message_item = ScheduleMessageItem( + user_id=user_id, + mem_cube_id=mem_cube_id, + mem_cube=self.mem_cubes[mem_cube_id], + label=label, + content=query, + timestamp=datetime.now(), + ) + self.mem_scheduler.submit_messages(messages=[message_item]) + + def register_mem_cube( + self, mem_cube_name_or_path: str, mem_cube_id: str | None = None, user_id: str | None = None + ) -> None: + """ + Register a MemCube with the MOS. + + Args: + mem_cube_name_or_path (str): The name or path of the MemCube to register. + mem_cube_id (str, optional): The identifier for the MemCube. If not provided, a default ID is used. + """ + + if mem_cube_id in self.mem_cubes: + logger.info(f"MemCube with ID {mem_cube_id} already in MOS, skip install.") + else: + if os.path.exists(mem_cube_name_or_path): + self.mem_cubes[mem_cube_id] = GeneralMemCube.init_from_dir(mem_cube_name_or_path) + else: + logger.warning( + f"MemCube {mem_cube_name_or_path} does not exist, try to init from remote repo." + ) + self.mem_cubes[mem_cube_id] = GeneralMemCube.init_from_remote_repo( + mem_cube_name_or_path + ) + + def user_register( + self, + user_id: str, + user_name: str | None = None, + config: MOSConfig | None = None, + interests: str | None = None, + default_mem_cube: GeneralMemCube | None = None, + ) -> dict[str, str]: + """Register a new user with configuration and default cube. + + Args: + user_id (str): The user ID for registration. + user_name (str): The user name for registration. + config (MOSConfig | None, optional): User-specific configuration. Defaults to None. + interests (str | None, optional): User interests as string. Defaults to None. + + Returns: + dict[str, str]: Registration result with status and message. + """ + try: + # Use provided config or default config + user_config = config or self.default_config + if not user_config: + return { + "status": "error", + "message": "No configuration provided for user registration", + } + if not user_name: + user_name = user_id + + # Create user with configuration using persistent user manager + self.global_user_manager.create_user_with_config( + user_id, user_config, UserRole.USER, user_id + ) + + # Create user configuration + user_config = self._create_user_config(user_id, user_config) + + # Create a default cube for the user using MOSCore's methods + default_cube_name = f"{user_name}_default_cube" + mem_cube_name_or_path = f"{CUBE_PATH}/{default_cube_name}" + default_cube_id = self.create_cube_for_user( + cube_name=default_cube_name, owner_id=user_id, cube_path=mem_cube_name_or_path + ) + + if default_mem_cube: + try: + default_mem_cube.dump(mem_cube_name_or_path) + except Exception as e: + print(e) + + # Register the default cube with MOS TODO overide + self.register_mem_cube(mem_cube_name_or_path, default_cube_id, user_id) + + # Add interests to the default cube if provided + if interests: + self.add(memory_content=interests, mem_cube_id=default_cube_id, user_id=user_id) + + return { + "status": "success", + "message": f"User {user_name} registered successfully with default cube {default_cube_id}", + "user_id": user_id, + "default_cube_id": default_cube_id, + } + + except Exception as e: + return {"status": "error", "message": f"Failed to register user: {e!s}"} def get_suggestion_query(self, user_id: str) -> list[str]: """Get suggestion query from LLM. Args: - user_id (str, optional): Custom user ID. + user_id (str): User ID. Returns: list[str]: The suggestion query list. """ + suggestion_prompt = """ + You are a helpful assistant that can help users to generate suggestion query + I will get some user recently memories, + you should generate some suggestion query , the query should be user what to query, + user recently memories is : + {memories} + please generate 3 suggestion query, + output should be a json format, the key is "query", the value is a list of suggestion query. + + example: + {{ + "query": ["query1", "query2", "query3"] + }} + """ + memories = "\n".join( + [ + m.memory + for m in super().search("my recently memories", user_id=user_id, top_k=10)[ + "text_mem" + ][0]["memories"] + ] + ) + message_list = [{"role": "system", "content": suggestion_prompt.format(memories=memories)}] + response = self.chat_llm.generate(message_list) + response_json = json.loads(response) + + return response_json["query"] + def chat( self, query: str, @@ -38,43 +558,165 @@ def chat( """Chat with LLM SSE Type. Args: query (str): Query string. - user_id (str, optional): Custom user ID. + user_id (str): User ID. cube_id (str, optional): Custom cube ID for user. history (list[dict], optional): Chat history. Returns: Generator[str, None, None]: The response string generator. """ - memories_list = self.search(query)["act_mem"] - content_list = [] - for memory in memories_list: - content_list.append(memory.content) - yield f"data: {json.dumps({'type': 'metadata', 'content': content_list})}\n\n" - llm_response = super().chat(query, user_id) - for chunk in llm_response: - chunk_data: str = f"data: {json.dumps({'type': 'text', 'content': chunk})}\n\n" + # Use MOSCore's built-in validation + if cube_id: + self._validate_cube_access(user_id, cube_id) + else: + self._validate_user_exists(user_id) + + # Load user cubes if not already loaded + self._load_user_cubes(user_id) + time_start = time.time() + memories_list = super().search(query, user_id)["text_mem"] + # Get response from parent MOSCore (returns string, not generator) + response = super().chat(query, user_id) + time_end = time.time() + + # Use tiktoken for proper token-based chunking + for chunk in self._chunk_response_with_tiktoken(response, chunk_size=5): + chunk_data = f"data: {json.dumps({'type': 'text', 'content': chunk})}\n\n" yield chunk_data - reference = [{"id": "1234"}] + + # Prepare reference data + reference = [] + for memories in memories_list: + memories_json = memories.model_dump() + memories_json["metadata"]["ref_id"] = f"[{memories.id.split('-')[0]}]" + memories_json["metadata"]["embedding"] = [] + memories_json["metadata"]["sources"] = [] + reference.append(memories_json) + yield f"data: {json.dumps({'type': 'reference', 'content': reference})}\n\n" + total_time = round(float(time_end - time_start), 1) + + yield f"data: {json.dumps({'type': 'time', 'content': {'total_time': total_time, 'speed_improvement': '23%'}})}\n\n" yield f"data: {json.dumps({'type': 'end'})}\n\n" - def get_all( + def chat_with_references( self, + query: str, user_id: str, - memory_type: Literal["text_mem", "act_mem", "param_mem"], cube_id: str | None = None, - ) -> list[ - dict[ - str, - str - | list[ - TextualMemoryMetadata - | TreeNodeTextualMemoryMetadata - | ActivationMemoryItem - | ParametricMemoryItem - ], + history: MessageList | None = None, + ) -> Generator[str, None, None]: + """ + Chat with LLM with memory references and streaming output. + + Args: + query (str): Query string. + user_id (str): User ID. + cube_id (str, optional): Custom cube ID for user. + history (MessageList, optional): Chat history. + + Returns: + Generator[str, None, None]: The response string generator with reference processing. + """ + + self._load_user_cubes(user_id) + + time_start = time.time() + memories_list = super().search( + query, user_id, install_cube_ids=[cube_id] if cube_id else None + )["text_mem"][0]["memories"] + + # Build custom system prompt with relevant memories + system_prompt = self._build_system_prompt(user_id, memories_list) + + # Get chat history + target_user_id = user_id if user_id is not None else self.user_id + if target_user_id not in self.chat_history_manager: + self._register_chat_history(target_user_id) + + chat_history = self.chat_history_manager[target_user_id] + current_messages = [ + {"role": "system", "content": system_prompt}, + *chat_history.chat_history, + {"role": "user", "content": query}, ] - ]: + + # Generate response with custom prompt + past_key_values = None + if self.config.enable_activation_memory: + # Handle activation memory (copy MOSCore logic) + for mem_cube_id, mem_cube in self.mem_cubes.items(): + if mem_cube.act_mem and mem_cube_id == cube_id: + kv_cache = next(iter(mem_cube.act_mem.get_all()), None) + past_key_values = ( + kv_cache.memory if (kv_cache and hasattr(kv_cache, "memory")) else None + ) + if past_key_values is not None: + logger.info("past_key_values is not None will apply to chat") + else: + logger.info("past_key_values is None will not apply to chat") + break + response = self.chat_llm.generate(current_messages, past_key_values=past_key_values) + else: + response = self.chat_llm.generate(current_messages) + + time_end = time.time() + + # Simulate streaming output with proper reference handling using tiktoken + + # Initialize buffer for streaming + buffer = "" + + # Use tiktoken for proper token-based chunking + for chunk in self._chunk_response_with_tiktoken(response, chunk_size=5): + buffer += chunk + + # Process buffer to ensure complete reference tags + processed_chunk, remaining_buffer = self._process_streaming_references_complete(buffer) + + if processed_chunk: + chunk_data = f"data: {json.dumps({'type': 'text', 'data': processed_chunk}, ensure_ascii=False)}\n\n" + yield chunk_data + buffer = remaining_buffer + + # Process any remaining buffer + if buffer: + processed_chunk, remaining_buffer = self._process_streaming_references_complete(buffer) + if processed_chunk: + chunk_data = f"data: {json.dumps({'type': 'text', 'data': processed_chunk}, ensure_ascii=False)}\n\n" + yield chunk_data + + # Prepare reference data + reference = [] + for memories in memories_list: + memories_json = memories.model_dump() + memories_json["metadata"]["ref_id"] = f"{memories.id.split('-')[0]}" + memories_json["metadata"]["embedding"] = [] + memories_json["metadata"]["sources"] = [] + memories_json["metadata"]["memory"] = memories.memory + reference.append({"metadata": memories_json["metadata"]}) + + yield f"data: {json.dumps({'type': 'reference', 'data': reference})}\n\n" + total_time = round(float(time_end - time_start), 1) + yield f"data: {json.dumps({'type': 'time', 'data': {'total_time': total_time, 'speed_improvement': '23%'}})}\n\n" + chat_history.chat_history.append({"role": "user", "content": query}) + chat_history.chat_history.append({"role": "assistant", "content": response}) + self._send_message_to_scheduler( + user_id=user_id, mem_cube_id=cube_id, query=query, label=QUERY_LABEL + ) + self._send_message_to_scheduler( + user_id=user_id, mem_cube_id=cube_id, query=response, label=ANSWER_LABEL + ) + self.chat_history_manager[user_id] = chat_history + + yield f"data: {json.dumps({'type': 'end'})}\n\n" + + def get_all( + self, + user_id: str, + memory_type: Literal["text_mem", "act_mem", "param_mem", "para_mem"], + mem_cube_ids: list[str] | None = None, + ) -> list[dict[str, Any]]: """Get all memory items for a user. Args: @@ -83,7 +725,233 @@ def get_all( memory_type (Literal["text_mem", "act_mem", "param_mem"]): The type of memory to get. Returns: - list[TextualMemoryMetadata | TreeNodeTextualMemoryMetadata | ActivationMemoryItem | ParametricMemoryItem]: A list of memory items. + list[dict[str, Any]]: A list of memory items with cube_id and memories structure. + """ + + # Load user cubes if not already loaded + self._load_user_cubes(user_id) + memory_list = super().get_all( + mem_cube_id=mem_cube_ids[0] if mem_cube_ids else None, user_id=user_id + )[memory_type] + reformat_memory_list = [] + if memory_type == "text_mem": + for memory in memory_list: + memories = remove_embedding_recursive(memory["memories"]) + custom_type_ratios = { + "WorkingMemory": 0.20, + "LongTermMemory": 0.40, + "UserMemory": 0.40, + } + tree_result = convert_graph_to_tree_forworkmem( + memories, target_node_count=150, type_ratios=custom_type_ratios + ) + memories_filtered = filter_nodes_by_tree_ids(tree_result, memories) + children = tree_result["children"] + children_sort = sort_children_by_memory_type(children) + tree_result["children"] = children_sort + memories_filtered["tree_structure"] = tree_result + reformat_memory_list.append( + {"cube_id": memory["cube_id"], "memories": [memories_filtered]} + ) + elif memory_type == "act_mem": + reformat_memory_list.append( + { + "cube_id": "xxxxxxxxxxxxxxxx" if not mem_cube_ids else mem_cube_ids[0], + "memories": MOCK_DATA, + } + ) + elif memory_type == "para_mem": + cache_item = self.mem_cubes[mem_cube_ids[0]].act_mem.extract("这是一个小问题哈哈哈哈") + self.mem_cubes[mem_cube_ids[0]].act_mem.add([cache_item]) + act_mem_params = self.mem_cubes[mem_cube_ids[0]].act_mem.get_all() + # Convert activation memory to serializable format + serializable_act_mem = convert_activation_memory_to_serializable(act_mem_params) + reformat_memory_list.append( + { + "cube_id": "xxxxxxxxxxxxxxxx" if not mem_cube_ids else mem_cube_ids[0], + "memories": serializable_act_mem, + } + ) + return reformat_memory_list + + def _get_subgraph( + self, query: str, mem_cube_id: str, user_id: str | None = None, top_k: int = 5 + ) -> list[dict[str, Any]]: + result = {"para_mem": [], "act_mem": [], "text_mem": []} + if self.config.enable_textual_memory and self.mem_cubes[mem_cube_id].text_mem: + result["text_mem"].append( + { + "cube_id": mem_cube_id, + "memories": self.mem_cubes[mem_cube_id].text_mem.get_relevant_subgraph( + query, top_k=top_k + ), + } + ) + return result + + def get_subgraph( + self, + user_id: str, + query: str, + mem_cube_ids: list[str] | None = None, + ) -> list[dict[str, Any]]: + """Get all memory items for a user. + + Args: + user_id (str): The ID of the user. + cube_id (str | None, optional): The ID of the cube. Defaults to None. + mem_cube_ids (list[str], optional): The IDs of the cubes. Defaults to None. + + Returns: + list[dict[str, Any]]: A list of memory items with cube_id and memories structure. """ - memory_list = super().get_all(user_id, cube_id)[memory_type] - return memory_list + + # Load user cubes if not already loaded + self._load_user_cubes(user_id) + memory_list = self._get_subgraph( + query=query, mem_cube_id=mem_cube_ids[0], user_id=user_id, top_k=20 + )["text_mem"] + reformat_memory_list = [] + for memory in memory_list: + memories = remove_embedding_recursive(memory["memories"]) + custom_type_ratios = {"WorkingMemory": 0.20, "LongTermMemory": 0.40, "UserMemory": 0.4} + tree_result = convert_graph_to_tree_forworkmem( + memories, target_node_count=150, type_ratios=custom_type_ratios + ) + memories_filtered = filter_nodes_by_tree_ids(tree_result, memories) + children = tree_result["children"] + children_sort = sort_children_by_memory_type(children) + tree_result["children"] = children_sort + memories_filtered["tree_structure"] = tree_result + reformat_memory_list.append( + {"cube_id": memory["cube_id"], "memories": [memories_filtered]} + ) + + return reformat_memory_list + + def search( + self, query: str, user_id: str, install_cube_ids: list[str] | None = None, top_k: int = 20 + ): + """Search memories for a specific user.""" + # Validate user access + self._validate_user_access(user_id) + + # Load user cubes if not already loaded + self._load_user_cubes(user_id) + search_result = super().search(query, user_id, install_cube_ids, top_k) + text_memory_list = search_result["text_mem"] + reformat_memory_list = [] + for memory in text_memory_list: + memories_list = [] + for data in memory["memories"]: + memories = data.model_dump() + memories["ref_id"] = f"[{memories['id'].split('-')[0]}]" + memories["metadata"]["embedding"] = [] + memories["metadata"]["sources"] = [] + memories["metadata"]["ref_id"] = f"[{memories['id'].split('-')[0]}]" + memories["metadata"]["id"] = memories["id"] + memories["metadata"]["memory"] = memories["memory"] + memories_list.append(memories) + reformat_memory_list.append({"cube_id": memory["cube_id"], "memories": memories_list}) + search_result["text_mem"] = reformat_memory_list + + return search_result + + def add( + self, + user_id: str, + messages: MessageList | None = None, + memory_content: str | None = None, + doc_path: str | None = None, + mem_cube_id: str | None = None, + ): + """Add memory for a specific user.""" + # Use MOSCore's built-in user/cube validation + if mem_cube_id: + self._validate_cube_access(user_id, mem_cube_id) + else: + self._validate_user_exists(user_id) + + # Load user cubes if not already loaded + self._load_user_cubes(user_id) + + result = super().add(messages, memory_content, doc_path, mem_cube_id, user_id) + + return result + + def list_users(self) -> list: + """List all registered users.""" + return self.global_user_manager.list_users() + + def get_user_info(self, user_id: str) -> dict: + """Get user information including accessible cubes.""" + # Use MOSCore's built-in user validation + # Validate user access + self._validate_user_access(user_id) + + result = super().get_user_info() + + return result + + def share_cube_with_user(self, cube_id: str, owner_user_id: str, target_user_id: str) -> bool: + """Share a cube with another user.""" + # Use MOSCore's built-in cube access validation + self._validate_cube_access(owner_user_id, cube_id) + + result = super().share_cube_with_user(cube_id, target_user_id) + + return result + + def clear_user_chat_history(self, user_id: str) -> None: + """Clear chat history for a specific user.""" + # Validate user access + self._validate_user_access(user_id) + + super().clear_messages(user_id) + + def update_user_config(self, user_id: str, config: MOSConfig) -> bool: + """Update user configuration. + + Args: + user_id (str): The user ID. + config (MOSConfig): The new configuration. + + Returns: + bool: True if successful, False otherwise. + """ + try: + # Save to persistent storage + success = self.global_user_manager.save_user_config(user_id, config) + if success: + # Update in-memory config + self.user_configs[user_id] = config + logger.info(f"Updated configuration for user {user_id}") + + return success + except Exception as e: + logger.error(f"Failed to update user config for {user_id}: {e}") + return False + + def get_user_config(self, user_id: str) -> MOSConfig | None: + """Get user configuration. + + Args: + user_id (str): The user ID. + + Returns: + MOSConfig | None: The user's configuration or None if not found. + """ + return self.global_user_manager.get_user_config(user_id) + + def get_active_user_count(self) -> int: + """Get the number of active user configurations in memory.""" + return len(self.user_configs) + + def get_user_instance_info(self) -> dict[str, Any]: + """Get information about user configurations in memory.""" + return { + "active_instances": len(self.user_configs), + "max_instances": self.max_user_instances, + "user_ids": list(self.user_configs.keys()), + "lru_order": list(self.user_configs.keys()), # OrderedDict maintains insertion order + } diff --git a/src/memos/mem_os/utils/format_utils.py b/src/memos/mem_os/utils/format_utils.py new file mode 100644 index 000000000..ebf18ba22 --- /dev/null +++ b/src/memos/mem_os/utils/format_utils.py @@ -0,0 +1,1146 @@ +import math +import random + +from typing import Any + +from memos.log import get_logger +from memos.memories.activation.item import KVCacheItem + + +logger = get_logger(__name__) + + +def extract_node_name(memory: str) -> str: + """Extract the first two words from memory as node_name""" + if not memory: + return "" + + words = [word.strip() for word in memory.split() if word.strip()] + + if len(words) >= 2: + return " ".join(words[:2]) + elif len(words) == 1: + return words[0] + else: + return "" + + +def analyze_tree_structure_enhanced(nodes: list[dict], edges: list[dict]) -> dict: + """Enhanced tree structure analysis, focusing on branching degree and leaf distribution""" + # Build adjacency list + adj_list = {} + reverse_adj = {} + for edge in edges: + source, target = edge["source"], edge["target"] + adj_list.setdefault(source, []).append(target) + reverse_adj.setdefault(target, []).append(source) + + # Find all nodes and root nodes + all_nodes = {node["id"] for node in nodes} + target_nodes = {edge["target"] for edge in edges} + root_nodes = all_nodes - target_nodes + + subtree_analysis = {} + + def analyze_subtree_enhanced(root_id: str) -> dict: + """Enhanced subtree analysis, focusing on branching degree and structure quality""" + visited = set() + max_depth = 0 + leaf_count = 0 + total_nodes = 0 + branch_nodes = 0 # Number of branch nodes with multiple children + chain_length = 0 # Longest single chain length + width_per_level = {} # Width per level + + def dfs(node_id: str, depth: int, chain_len: int): + nonlocal max_depth, leaf_count, total_nodes, branch_nodes, chain_length + + if node_id in visited: + return + + visited.add(node_id) + total_nodes += 1 + max_depth = max(max_depth, depth) + chain_length = max(chain_length, chain_len) + + # Record number of nodes per level + width_per_level[depth] = width_per_level.get(depth, 0) + 1 + + children = adj_list.get(node_id, []) + + if not children: # Leaf node + leaf_count += 1 + elif len(children) > 1: # Branch node + branch_nodes += 1 + # Reset chain length because we encountered a branch + for child in children: + dfs(child, depth + 1, 0) + else: # Single child node (chain structure) + for child in children: + dfs(child, depth + 1, chain_len + 1) + + dfs(root_id, 0, 0) + + # Calculate structure quality metrics + avg_width = sum(width_per_level.values()) / len(width_per_level) if width_per_level else 0 + max_width = max(width_per_level.values()) if width_per_level else 0 + + # Calculate branch density: ratio of branch nodes to total nodes + branch_density = branch_nodes / total_nodes if total_nodes > 0 else 0 + + # Calculate depth-width ratio: ideal tree should have moderate depth and good breadth + depth_width_ratio = max_depth / max_width if max_width > 0 else max_depth + + quality_score = calculate_enhanced_quality( + max_depth, + leaf_count, + total_nodes, + branch_nodes, + chain_length, + branch_density, + depth_width_ratio, + max_width, + ) + + return { + "root_id": root_id, + "max_depth": max_depth, + "leaf_count": leaf_count, + "total_nodes": total_nodes, + "branch_nodes": branch_nodes, + "max_chain_length": chain_length, + "branch_density": branch_density, + "max_width": max_width, + "avg_width": avg_width, + "depth_width_ratio": depth_width_ratio, + "nodes_in_subtree": list(visited), + "quality_score": quality_score, + "width_per_level": width_per_level, + } + + for root_id in root_nodes: + subtree_analysis[root_id] = analyze_subtree_enhanced(root_id) + + return subtree_analysis + + +def calculate_enhanced_quality( + max_depth: int, + leaf_count: int, + total_nodes: int, + branch_nodes: int, + max_chain_length: int, + branch_density: float, + depth_width_ratio: float, + max_width: int, +) -> float: + """Enhanced quality calculation, prioritizing branching degree and leaf distribution""" + + if total_nodes <= 1: + return 0.1 + + # 1. Branch quality score (weight: 35%) + # Branch node count score + branch_count_score = min(branch_nodes * 3, 15) # 3 points per branch node, max 15 points + + # Branch density score: ideal density between 20%-60% + if 0.2 <= branch_density <= 0.6: + branch_density_score = 10 + elif branch_density > 0.6: + branch_density_score = max(5, 10 - (branch_density - 0.6) * 20) + else: + branch_density_score = branch_density * 25 # Linear growth for 0-20% + + branch_score = (branch_count_score + branch_density_score) * 0.35 + + # 2. Leaf quality score (weight: 25%) + # Leaf count score + leaf_count_score = min(leaf_count * 2, 20) + + # Leaf distribution score: ideal leaf ratio 30%-70% of total nodes + leaf_ratio = leaf_count / total_nodes + if 0.3 <= leaf_ratio <= 0.7: + leaf_ratio_score = 10 + elif leaf_ratio > 0.7: + leaf_ratio_score = max(3, 10 - (leaf_ratio - 0.7) * 20) + else: + leaf_ratio_score = leaf_ratio * 20 # Linear growth for 0-30% + + leaf_score = (leaf_count_score + leaf_ratio_score) * 0.25 + + # 3. Structure balance score (weight: 25%) + # Depth score: moderate depth is best (3-8 layers) + if 3 <= max_depth <= 8: + depth_score = 15 + elif max_depth < 3: + depth_score = max_depth * 3 # Lower score for 1-2 layers + else: + depth_score = max(5, 15 - (max_depth - 8) * 1.5) # Gradually reduce score beyond 8 layers + + # Width score: larger max width is better, but with upper limit + width_score = min(max_width * 1.5, 15) + + # Depth-width ratio penalty: too large ratio means tree is too "thin" + if depth_width_ratio > 3: + ratio_penalty = (depth_width_ratio - 3) * 2 + structure_score = max(0, (depth_score + width_score - ratio_penalty)) * 0.25 + else: + structure_score = (depth_score + width_score) * 0.25 + + # 4. Chain structure penalty (weight: 15%) + # Longest single chain length penalty: overly long chains severely affect display + if max_chain_length <= 2: + chain_penalty_score = 10 + elif max_chain_length <= 5: + chain_penalty_score = 8 - (max_chain_length - 2) + else: + chain_penalty_score = max(0, 3 - (max_chain_length - 5) * 0.5) + + chain_score = chain_penalty_score * 0.15 + + # 5. Comprehensive calculation + total_score = branch_score + leaf_score + structure_score + chain_score + + # Special case severe penalties + if max_chain_length > total_nodes * 0.8: # If more than 80% are single chains + total_score *= 0.3 + elif branch_density < 0.1 and total_nodes > 5: # Large tree with almost no branches + total_score *= 0.5 + + return total_score + + +def sample_nodes_with_type_balance( + nodes: list[dict], + edges: list[dict], + target_count: int = 150, + type_ratios: dict[str, float] | None = None, +) -> tuple[list[dict], list[dict]]: + """ + Balanced sampling based on type ratios and tree quality + + Args: + nodes: List of nodes + edges: List of edges + target_count: Target number of nodes + type_ratios: Expected ratio for each type, e.g. {'WorkingMemory': 0.15, 'EpisodicMemory': 0.30, ...} + """ + if len(nodes) <= target_count: + return nodes, edges + + # Default type ratio configuration + if type_ratios is None: + type_ratios = { + "WorkingMemory": 0.10, # 10% + "EpisodicMemory": 0.25, # 25% + "SemanticMemory": 0.25, # 25% + "ProceduralMemory": 0.20, # 20% + "EmotionalMemory": 0.15, # 15% + "MetaMemory": 0.05, # 5% + } + + print( + f"Starting type-balanced sampling, original nodes: {len(nodes)}, target nodes: {target_count}" + ) + print(f"Target type ratios: {type_ratios}") + + # Analyze current node type distribution + current_type_counts = {} + nodes_by_type = {} + + for node in nodes: + memory_type = node.get("metadata", {}).get("memory_type", "Unknown") + current_type_counts[memory_type] = current_type_counts.get(memory_type, 0) + 1 + if memory_type not in nodes_by_type: + nodes_by_type[memory_type] = [] + nodes_by_type[memory_type].append(node) + + print(f"Current type distribution: {current_type_counts}") + + # Calculate target node count for each type + type_targets = {} + remaining_target = target_count + + for memory_type, ratio in type_ratios.items(): + if memory_type in nodes_by_type: + target_for_type = int(target_count * ratio) + # Ensure not exceeding the actual node count for this type + target_for_type = min(target_for_type, len(nodes_by_type[memory_type])) + type_targets[memory_type] = target_for_type + remaining_target -= target_for_type + + # Handle types not in ratio configuration + other_types = set(nodes_by_type.keys()) - set(type_ratios.keys()) + if other_types and remaining_target > 0: + per_other_type = max(1, remaining_target // len(other_types)) + for memory_type in other_types: + allocation = min(per_other_type, len(nodes_by_type[memory_type])) + type_targets[memory_type] = allocation + remaining_target -= allocation + + # If there's still remaining, distribute proportionally to main types + if remaining_target > 0: + main_types = [t for t in type_ratios if t in nodes_by_type] + if main_types: + extra_per_type = remaining_target // len(main_types) + for memory_type in main_types: + additional = min( + extra_per_type, + len(nodes_by_type[memory_type]) - type_targets.get(memory_type, 0), + ) + type_targets[memory_type] = type_targets.get(memory_type, 0) + additional + + print(f"Target node count for each type: {type_targets}") + + # Perform subtree quality sampling for each type + selected_nodes = [] + + for memory_type, target_for_type in type_targets.items(): + if target_for_type <= 0 or memory_type not in nodes_by_type: + continue + + type_nodes = nodes_by_type[memory_type] + print(f"\n--- Processing {memory_type} type: {len(type_nodes)} -> {target_for_type} ---") + + if len(type_nodes) <= target_for_type: + selected_nodes.extend(type_nodes) + print(f" Select all: {len(type_nodes)} nodes") + else: + # Use enhanced subtree quality sampling + type_selected = sample_by_enhanced_subtree_quality(type_nodes, edges, target_for_type) + selected_nodes.extend(type_selected) + print(f" Sampled selection: {len(type_selected)} nodes") + + # Filter edges + selected_node_ids = {node["id"] for node in selected_nodes} + filtered_edges = [ + edge + for edge in edges + if edge["source"] in selected_node_ids and edge["target"] in selected_node_ids + ] + + print(f"\nFinal selected nodes: {len(selected_nodes)}") + print(f"Final edges: {len(filtered_edges)}") + + # Verify final type distribution + final_type_counts = {} + for node in selected_nodes: + memory_type = node.get("metadata", {}).get("memory_type", "Unknown") + final_type_counts[memory_type] = final_type_counts.get(memory_type, 0) + 1 + + print(f"Final type distribution: {final_type_counts}") + for memory_type, count in final_type_counts.items(): + percentage = count / len(selected_nodes) * 100 + target_percentage = type_ratios.get(memory_type, 0) * 100 + print( + f" {memory_type}: {count} nodes ({percentage:.1f}%, target: {target_percentage:.1f}%)" + ) + + return selected_nodes, filtered_edges + + +def sample_by_enhanced_subtree_quality( + nodes: list[dict], edges: list[dict], target_count: int +) -> list[dict]: + """Sample using enhanced subtree quality""" + if len(nodes) <= target_count: + return nodes + + # Analyze subtree structure + subtree_analysis = analyze_tree_structure_enhanced(nodes, edges) + + if not subtree_analysis: + # If no subtree structure, sample by node importance + return sample_nodes_by_importance(nodes, edges, target_count) + + # Sort subtrees by quality score + sorted_subtrees = sorted( + subtree_analysis.items(), key=lambda x: x[1]["quality_score"], reverse=True + ) + + print(" Subtree quality ranking:") + for i, (root_id, analysis) in enumerate(sorted_subtrees[:5]): + print( + f" #{i + 1} Root node {root_id}: Quality={analysis['quality_score']:.2f}, " + f"Depth={analysis['max_depth']}, Branches={analysis['branch_nodes']}, " + f"Leaves={analysis['leaf_count']}, Max Width={analysis['max_width']}" + ) + + # Greedy selection of high-quality subtrees + selected_nodes = [] + selected_node_ids = set() + + for root_id, analysis in sorted_subtrees: + subtree_nodes = analysis["nodes_in_subtree"] + new_nodes = [node_id for node_id in subtree_nodes if node_id not in selected_node_ids] + + if not new_nodes: + continue + + remaining_quota = target_count - len(selected_nodes) + + if len(new_nodes) <= remaining_quota: + # Entire subtree can be added + for node_id in new_nodes: + node = next((n for n in nodes if n["id"] == node_id), None) + if node: + selected_nodes.append(node) + selected_node_ids.add(node_id) + print(f" Select entire subtree {root_id}: +{len(new_nodes)} nodes") + else: + # Subtree too large, need partial selection + if analysis["quality_score"] > 5: # Only partial selection for high-quality subtrees + subtree_node_objects = [n for n in nodes if n["id"] in new_nodes] + partial_selection = select_best_nodes_from_subtree( + subtree_node_objects, edges, remaining_quota, root_id + ) + + selected_nodes.extend(partial_selection) + for node in partial_selection: + selected_node_ids.add(node["id"]) + print( + f" Partial selection of subtree {root_id}: +{len(partial_selection)} nodes" + ) + + if len(selected_nodes) >= target_count: + break + + # If target count not reached, supplement with remaining nodes + if len(selected_nodes) < target_count: + remaining_nodes = [n for n in nodes if n["id"] not in selected_node_ids] + remaining_count = target_count - len(selected_nodes) + additional = sample_nodes_by_importance(remaining_nodes, edges, remaining_count) + selected_nodes.extend(additional) + print(f" Supplementary selection: +{len(additional)} nodes") + + return selected_nodes + + +def select_best_nodes_from_subtree( + subtree_nodes: list[dict], edges: list[dict], max_count: int, root_id: str +) -> list[dict]: + """Select the most important nodes from subtree, prioritizing branch structure""" + if len(subtree_nodes) <= max_count: + return subtree_nodes + + # Build internal connection relationships of subtree + subtree_node_ids = {node["id"] for node in subtree_nodes} + subtree_edges = [ + edge + for edge in edges + if edge["source"] in subtree_node_ids and edge["target"] in subtree_node_ids + ] + + # Calculate importance score for each node + node_scores = [] + + for node in subtree_nodes: + node_id = node["id"] + + # Out-degree and in-degree + out_degree = sum(1 for edge in subtree_edges if edge["source"] == node_id) + in_degree = sum(1 for edge in subtree_edges if edge["target"] == node_id) + + # Content length score + content_score = min(len(node.get("memory", "")), 300) / 15 + + # Branch node bonus + branch_bonus = out_degree * 8 if out_degree > 1 else 0 + + # Root node bonus + root_bonus = 15 if node_id == root_id else 0 + + # Connection importance + connection_score = (out_degree + in_degree) * 3 + + # Leaf node moderate bonus (ensure certain number of leaf nodes) + leaf_bonus = 5 if out_degree == 0 and in_degree > 0 else 0 + + total_score = content_score + connection_score + branch_bonus + root_bonus + leaf_bonus + node_scores.append((node, total_score)) + + # Sort by score and select + node_scores.sort(key=lambda x: x[1], reverse=True) + selected = [node for node, _ in node_scores[:max_count]] + + return selected + + +def sample_nodes_by_importance( + nodes: list[dict], edges: list[dict], target_count: int +) -> list[dict]: + """Sample by node importance (for cases without tree structure)""" + if len(nodes) <= target_count: + return nodes + + node_scores = [] + + for node in nodes: + node_id = node["id"] + out_degree = sum(1 for edge in edges if edge["source"] == node_id) + in_degree = sum(1 for edge in edges if edge["target"] == node_id) + content_score = min(len(node.get("memory", "")), 200) / 10 + connection_score = (out_degree + in_degree) * 5 + random_score = random.random() * 10 + + total_score = content_score + connection_score + random_score + node_scores.append((node, total_score)) + + node_scores.sort(key=lambda x: x[1], reverse=True) + return [node for node, _ in node_scores[:target_count]] + + +# Modified main function to use new sampling strategy +def convert_graph_to_tree_forworkmem( + json_data: dict[str, Any], + target_node_count: int = 150, + type_ratios: dict[str, float] | None = None, +) -> dict[str, Any]: + """ + Enhanced graph-to-tree conversion function, prioritizing branching degree and type balance + """ + original_nodes = json_data.get("nodes", []) + original_edges = json_data.get("edges", []) + + print(f"Original node count: {len(original_nodes)}") + print(f"Target node count: {target_node_count}") + filter_original_edges = [] + for original_edge in original_edges: + if original_edge["type"] == "PARENT": + filter_original_edges.append(original_edge) + original_edges = filter_original_edges + # Use enhanced type-balanced sampling + if len(original_nodes) > target_node_count: + nodes, edges = sample_nodes_with_type_balance( + original_nodes, original_edges, target_node_count, type_ratios + ) + else: + nodes, edges = original_nodes, original_edges + + # The rest of tree structure building remains unchanged... + # [Original tree building code here] + + # Create node mapping table + node_map = {} + for node in nodes: + memory = node.get("memory", "") + node_map[node["id"]] = { + "id": node["id"], + "value": memory, + "frequency": random.randint(1, 100), + "node_name": extract_node_name(memory), + "memory_type": node.get("metadata", {}).get("memory_type", "Unknown"), + "children": [], + } + + # Build parent-child relationship mapping + children_map = {} + parent_map = {} + + for edge in edges: + source = edge["source"] + target = edge["target"] + if source not in children_map: + children_map[source] = [] + children_map[source].append(target) + parent_map[target] = source + + # Find root nodes + all_node_ids = set(node_map.keys()) + children_node_ids = set(parent_map.keys()) + root_node_ids = all_node_ids - children_node_ids + + # Separate WorkingMemory and other root nodes + working_memory_roots = [] + other_roots = [] + + for root_id in root_node_ids: + if node_map[root_id]["memory_type"] == "WorkingMemory": + working_memory_roots.append(root_id) + else: + other_roots.append(root_id) + + def build_tree(node_id: str) -> dict[str, Any]: + """Recursively build tree structure""" + if node_id not in node_map: + return None + + children_ids = children_map.get(node_id, []) + children = [] + for child_id in children_ids: + child_tree = build_tree(child_id) + if child_tree: + children.append(child_tree) + + node = { + "id": node_id, + "node_name": node_map[node_id]["node_name"], + "value": node_map[node_id]["value"], + "memory_type": node_map[node_id]["memory_type"], + "frequency": node_map[node_id]["frequency"], + } + + if children: + node["children"] = children + + return node + + # Build root tree list + root_trees = [] + for root_id in other_roots: + tree = build_tree(root_id) + if tree: + root_trees.append(tree) + + # Handle WorkingMemory + if working_memory_roots: + working_memory_children = [] + for wm_root_id in working_memory_roots: + tree = build_tree(wm_root_id) + if tree: + working_memory_children.append(tree) + + working_memory_node = { + "id": "WorkingMemory", + "node_name": "WorkingMemory", + "value": "WorkingMemory", + "memory_type": "WorkingMemory", + "children": working_memory_children, + "frequency": 0, + } + + root_trees.append(working_memory_node) + + # Create total root node + result = { + "id": "root", + "node_name": "root", + "value": "root", + "memory_type": "Root", + "children": root_trees, + "frequency": 0, + } + + return result + + +def print_tree_structure(node: dict[str, Any], level: int = 0, max_level: int = 5): + """Print the first few layers of tree structure for easy viewing""" + if level > max_level: + return + + indent = " " * level + node_id = node.get("id", "unknown") + node_name = node.get("node_name", "") + node_value = node.get("value", "") + memory_type = node.get("memory_type", "Unknown") + + # Determine display method based on whether there are children + children = node.get("children", []) + if children: + # Intermediate node, display name, type and child count + print(f"{indent}- {node_name} [{memory_type}] ({len(children)} children)") + print(f"{indent} ID: {node_id}") + display_value = node_value[:80] + "..." if len(node_value) > 80 else node_value + print(f"{indent} Value: {display_value}") + + if level < max_level: + for child in children: + print_tree_structure(child, level + 1, max_level) + elif level == max_level: + print(f"{indent} ... (expansion limited)") + else: + # Leaf node, display name, type and value + display_value = node_value[:80] + "..." if len(node_value) > 80 else node_value + print(f"{indent}- {node_name} [{memory_type}]: {display_value}") + print(f"{indent} ID: {node_id}") + + +def analyze_final_tree_quality(tree_data: dict[str, Any]) -> dict: + """Analyze final tree quality, including type diversity, branch structure, etc.""" + stats = { + "total_nodes": 0, + "by_type": {}, + "by_depth": {}, + "max_depth": 0, + "total_leaves": 0, + "total_branches": 0, # Number of branch nodes with multiple children + "subtrees": [], + "type_diversity": {}, + "structure_quality": {}, + "chain_analysis": {}, # Single chain structure analysis + } + + def analyze_subtree(node, depth=0, parent_path="", chain_length=0): + stats["total_nodes"] += 1 + stats["max_depth"] = max(stats["max_depth"], depth) + + # Count by type + memory_type = node.get("memory_type", "Unknown") + stats["by_type"][memory_type] = stats["by_type"].get(memory_type, 0) + 1 + + # Count by depth + stats["by_depth"][depth] = stats["by_depth"].get(depth, 0) + 1 + + children = node.get("children", []) + current_path = ( + f"{parent_path}/{node.get('node_name', 'unknown')}" + if parent_path + else node.get("node_name", "root") + ) + + # Analyze node type + if not children: # Leaf node + stats["total_leaves"] += 1 + # Record chain length + if "max_chain_length" not in stats["chain_analysis"]: + stats["chain_analysis"]["max_chain_length"] = 0 + stats["chain_analysis"]["max_chain_length"] = max( + stats["chain_analysis"]["max_chain_length"], chain_length + ) + elif len(children) == 1: # Single child node (chain) + # Continue calculating chain length + for child in children: + analyze_subtree(child, depth + 1, current_path, chain_length + 1) + return # Early return to avoid duplicate processing + else: # Branch node (multiple children) + stats["total_branches"] += 1 + # Reset chain length + chain_length = 0 + + # If it's the root node of a major subtree, analyze its characteristics + if depth <= 2 and children: # Major subtree + subtree_depth = 0 + subtree_leaves = 0 + subtree_nodes = 0 + subtree_branches = 0 + subtree_types = {} + subtree_max_width = 0 + width_per_level = {} + + def count_subtree(subnode, subdepth): + nonlocal \ + subtree_depth, \ + subtree_leaves, \ + subtree_nodes, \ + subtree_branches, \ + subtree_max_width + subtree_nodes += 1 + subtree_depth = max(subtree_depth, subdepth) + + # Count type distribution within subtree + sub_memory_type = subnode.get("memory_type", "Unknown") + subtree_types[sub_memory_type] = subtree_types.get(sub_memory_type, 0) + 1 + + # Count width per level + width_per_level[subdepth] = width_per_level.get(subdepth, 0) + 1 + subtree_max_width = max(subtree_max_width, width_per_level[subdepth]) + + subchildren = subnode.get("children", []) + if not subchildren: + subtree_leaves += 1 + elif len(subchildren) > 1: + subtree_branches += 1 + + for child in subchildren: + count_subtree(child, subdepth + 1) + + count_subtree(node, 0) + + # Calculate subtree quality metrics + branch_density = subtree_branches / subtree_nodes if subtree_nodes > 0 else 0 + leaf_ratio = subtree_leaves / subtree_nodes if subtree_nodes > 0 else 0 + depth_width_ratio = ( + subtree_depth / subtree_max_width if subtree_max_width > 0 else subtree_depth + ) + + stats["subtrees"].append( + { + "root": node.get("node_name", "unknown"), + "type": memory_type, + "depth": subtree_depth, + "leaves": subtree_leaves, + "nodes": subtree_nodes, + "branches": subtree_branches, + "branch_density": branch_density, + "leaf_ratio": leaf_ratio, + "max_width": subtree_max_width, + "depth_width_ratio": depth_width_ratio, + "path": current_path, + "type_distribution": subtree_types, + "quality_score": calculate_enhanced_quality( + subtree_depth, + subtree_leaves, + subtree_nodes, + subtree_branches, + 0, + branch_density, + depth_width_ratio, + subtree_max_width, + ), + } + ) + + # Recursively analyze child nodes + for child in children: + analyze_subtree(child, depth + 1, current_path, 0) # Reset chain length + + analyze_subtree(tree_data) + + # Calculate overall structure quality + if stats["total_nodes"] > 1: + branch_density = stats["total_branches"] / stats["total_nodes"] + leaf_ratio = stats["total_leaves"] / stats["total_nodes"] + + # Calculate average width per level + total_width = sum(stats["by_depth"].values()) + avg_width = total_width / len(stats["by_depth"]) if stats["by_depth"] else 0 + max_width = max(stats["by_depth"].values()) if stats["by_depth"] else 0 + + stats["structure_quality"] = { + "branch_density": branch_density, + "leaf_ratio": leaf_ratio, + "avg_width": avg_width, + "max_width": max_width, + "depth_width_ratio": stats["max_depth"] / max_width + if max_width > 0 + else stats["max_depth"], + "is_well_balanced": 0.2 <= branch_density <= 0.6 and 0.3 <= leaf_ratio <= 0.7, + } + + # Calculate type diversity metrics + total_types = len(stats["by_type"]) + if total_types > 1: + # Calculate uniformity of type distribution (Shannon diversity index) + shannon_diversity = 0 + for count in stats["by_type"].values(): + if count > 0: + p = count / stats["total_nodes"] + shannon_diversity -= p * math.log2(p) + + # Normalize diversity index (0-1 range) + max_diversity = math.log2(total_types) if total_types > 1 else 0 + normalized_diversity = shannon_diversity / max_diversity if max_diversity > 0 else 0 + + stats["type_diversity"] = { + "total_types": total_types, + "shannon_diversity": shannon_diversity, + "normalized_diversity": normalized_diversity, + "distribution_balance": min(stats["by_type"].values()) / max(stats["by_type"].values()) + if max(stats["by_type"].values()) > 0 + else 0, + } + + # Single chain analysis + total_single_child_nodes = sum( + 1 for subtree in stats["subtrees"] if subtree.get("branch_density", 0) < 0.1 + ) + stats["chain_analysis"].update( + { + "single_chain_subtrees": total_single_child_nodes, + "chain_subtree_ratio": total_single_child_nodes / len(stats["subtrees"]) + if stats["subtrees"] + else 0, + } + ) + + return stats + + +def print_tree_analysis(tree_data: dict[str, Any]): + """Print enhanced tree analysis results""" + stats = analyze_final_tree_quality(tree_data) + + print("\n" + "=" * 60) + print("🌳 Enhanced Tree Structure Quality Analysis Report") + print("=" * 60) + + # Basic statistics + print("\n📊 Basic Statistics:") + print(f" Total nodes: {stats['total_nodes']}") + print(f" Max depth: {stats['max_depth']}") + print( + f" Leaf nodes: {stats['total_leaves']} ({stats['total_leaves'] / stats['total_nodes'] * 100:.1f}%)" + ) + print( + f" Branch nodes: {stats['total_branches']} ({stats['total_branches'] / stats['total_nodes'] * 100:.1f}%)" + ) + + # Structure quality assessment + structure = stats.get("structure_quality", {}) + if structure: + print("\n🏗️ Structure Quality Assessment:") + print( + f" Branch density: {structure['branch_density']:.3f} ({'✅ Good' if 0.2 <= structure['branch_density'] <= 0.6 else '⚠️ Needs improvement'})" + ) + print( + f" Leaf ratio: {structure['leaf_ratio']:.3f} ({'✅ Good' if 0.3 <= structure['leaf_ratio'] <= 0.7 else '⚠️ Needs improvement'})" + ) + print(f" Max width: {structure['max_width']}") + print( + f" Depth-width ratio: {structure['depth_width_ratio']:.2f} ({'✅ Good' if structure['depth_width_ratio'] <= 3 else '⚠️ Too thin'})" + ) + print( + f" Overall balance: {'✅ Good' if structure['is_well_balanced'] else '⚠️ Needs improvement'}" + ) + + # Single chain analysis + chain_analysis = stats.get("chain_analysis", {}) + if chain_analysis: + print("\n🔗 Single Chain Structure Analysis:") + print(f" Longest chain: {chain_analysis.get('max_chain_length', 0)} layers") + print(f" Single chain subtrees: {chain_analysis.get('single_chain_subtrees', 0)}") + print( + f" Single chain subtree ratio: {chain_analysis.get('chain_subtree_ratio', 0) * 100:.1f}%" + ) + + if chain_analysis.get("max_chain_length", 0) > 5: + print(" ⚠️ Warning: Overly long single chain structure may affect display") + elif chain_analysis.get("chain_subtree_ratio", 0) > 0.3: + print( + " ⚠️ Warning: Too many single chain subtrees, suggest increasing branch structure" + ) + else: + print(" ✅ Single chain structure well controlled") + + # Type diversity + type_div = stats.get("type_diversity", {}) + if type_div: + print("\n🎨 Type Diversity Analysis:") + print(f" Total types: {type_div['total_types']}") + print(f" Diversity index: {type_div['shannon_diversity']:.3f}") + print(f" Normalized diversity: {type_div['normalized_diversity']:.3f}") + print(f" Distribution balance: {type_div['distribution_balance']:.3f}") + + # Type distribution + print("\n📋 Type Distribution Details:") + for mem_type, count in sorted(stats["by_type"].items(), key=lambda x: x[1], reverse=True): + percentage = count / stats["total_nodes"] * 100 + print(f" {mem_type}: {count} nodes ({percentage:.1f}%)") + + # Depth distribution + print("\n📏 Depth Distribution:") + for depth in sorted(stats["by_depth"].keys()): + count = stats["by_depth"][depth] + print(f" Depth {depth}: {count} nodes") + + # Major subtree analysis + if stats["subtrees"]: + print("\n🌲 Major Subtree Analysis (sorted by quality):") + sorted_subtrees = sorted( + stats["subtrees"], key=lambda x: x.get("quality_score", 0), reverse=True + ) + for i, subtree in enumerate(sorted_subtrees[:8]): # Show first 8 + quality = subtree.get("quality_score", 0) + print(f" #{i + 1} {subtree['root']} [{subtree['type']}]:") + print(f" Quality score: {quality:.2f}") + print( + f" Structure: Depth={subtree['depth']}, Branches={subtree['branches']}, Leaves={subtree['leaves']}" + ) + print( + f" Density: Branch density={subtree.get('branch_density', 0):.3f}, Leaf ratio={subtree.get('leaf_ratio', 0):.3f}" + ) + + if quality > 15: + print(" ✅ High quality subtree") + elif quality > 8: + print(" 🟡 Medium quality subtree") + else: + print(" 🔴 Low quality subtree") + + print("\n" + "=" * 60) + + +def remove_embedding_recursive(memory_info: dict) -> Any: + """remove the embedding from the memory info + Args: + memory_info: product memory info + + Returns: + Any: product memory info without embedding + """ + if isinstance(memory_info, dict): + new_dict = {} + for key, value in memory_info.items(): + if key != "embedding": + new_dict[key] = remove_embedding_recursive(value) + return new_dict + elif isinstance(memory_info, list): + return [remove_embedding_recursive(item) for item in memory_info] + else: + return memory_info + + +def remove_embedding_from_memory_items(memory_items: list[Any]) -> list[dict]: + """Batch remove embedding fields from multiple TextualMemoryItem objects""" + clean_memories = [] + + for item in memory_items: + memory_dict = item.model_dump() + + # Remove embedding from metadata + if "metadata" in memory_dict and "embedding" in memory_dict["metadata"]: + del memory_dict["metadata"]["embedding"] + + clean_memories.append(memory_dict) + + return clean_memories + + +def sort_children_by_memory_type(children: list[dict[str, Any]]) -> list[dict[str, Any]]: + """ + sort the children by the memory_type + Args: + children: the children of the node + Returns: + the sorted children + """ + if not children: + return children + + def get_sort_key(child): + memory_type = child.get("memory_type", "Unknown") + # Sort directly by memory_type string, same types will naturally cluster together + return memory_type + + # Sort by memory_type + sorted_children = sorted(children, key=get_sort_key) + + return sorted_children + + +def extract_all_ids_from_tree(tree_node): + """ + Recursively traverse tree structure to extract all node IDs + + Args: + tree_node: Tree node (dictionary format) + + Returns: + set: Set containing all node IDs + """ + ids = set() + + # Add current node ID (if exists) + if "id" in tree_node: + ids.add(tree_node["id"]) + + # Recursively process child nodes + if tree_node.get("children"): + for child in tree_node["children"]: + ids.update(extract_all_ids_from_tree(child)) + + return ids + + +def filter_nodes_by_tree_ids(tree_data, nodes_data): + """ + Filter nodes list based on IDs used in tree structure + + Args: + tree_data: Tree structure data (dictionary) + nodes_data: Data containing nodes list (dictionary) + + Returns: + dict: Filtered nodes data, maintaining original structure + """ + # Extract all IDs used in the tree + used_ids = extract_all_ids_from_tree(tree_data) + + # Filter nodes list, keeping only nodes with IDs used in the tree + filtered_nodes = [node for node in nodes_data["nodes"] if node["id"] in used_ids] + + # Return result maintaining original structure + return {"nodes": filtered_nodes} + + +def convert_activation_memory_to_serializable( + act_mem_items: list[KVCacheItem], +) -> list[dict[str, Any]]: + """ + Convert activation memory items to a serializable format. + + Args: + act_mem_items: List of KVCacheItem objects + + Returns: + List of dictionaries with serializable data + """ + serializable_items = [] + + for item in act_mem_items: + # Extract basic information that can be serialized + serializable_item = { + "id": item.id, + "metadata": item.metadata, + "memory_info": { + "type": "DynamicCache", + "key_cache_layers": len(item.memory.key_cache) if item.memory else 0, + "value_cache_layers": len(item.memory.value_cache) if item.memory else 0, + "device": str(item.memory.key_cache[0].device) + if item.memory and item.memory.key_cache + else "unknown", + "dtype": str(item.memory.key_cache[0].dtype) + if item.memory and item.memory.key_cache + else "unknown", + }, + } + + # Add tensor shape information if available + if item.memory and item.memory.key_cache: + key_shapes = [] + value_shapes = [] + + for i, key_tensor in enumerate(item.memory.key_cache): + if key_tensor is not None: + key_shapes.append({"layer": i, "shape": list(key_tensor.shape)}) + + if i < len(item.memory.value_cache) and item.memory.value_cache[i] is not None: + value_shapes.append( + {"layer": i, "shape": list(item.memory.value_cache[i].shape)} + ) + + serializable_item["memory_info"]["key_shapes"] = key_shapes + serializable_item["memory_info"]["value_shapes"] = value_shapes + + serializable_items.append(serializable_item) + + return serializable_items + + +def convert_activation_memory_summary(act_mem_items: list[KVCacheItem]) -> dict[str, Any]: + """ + Create a summary of activation memory for API responses. + + Args: + act_mem_items: List of KVCacheItem objects + + Returns: + Dictionary with summary information + """ + if not act_mem_items: + return {"total_items": 0, "summary": "No activation memory items found"} + + total_items = len(act_mem_items) + total_layers = 0 + total_parameters = 0 + + for item in act_mem_items: + if item.memory and item.memory.key_cache: + total_layers += len(item.memory.key_cache) + + # Calculate approximate parameter count + for key_tensor in item.memory.key_cache: + if key_tensor is not None: + total_parameters += key_tensor.numel() + + for value_tensor in item.memory.value_cache: + if value_tensor is not None: + total_parameters += value_tensor.numel() + + return { + "total_items": total_items, + "total_layers": total_layers, + "total_parameters": total_parameters, + "summary": f"Activation memory contains {total_items} items with {total_layers} layers and approximately {total_parameters:,} parameters", + } diff --git a/src/memos/mem_user/persistent_user_manager.py b/src/memos/mem_user/persistent_user_manager.py new file mode 100644 index 000000000..e3c476262 --- /dev/null +++ b/src/memos/mem_user/persistent_user_manager.py @@ -0,0 +1,260 @@ +"""Persistent user management system for MemOS with configuration storage. + +This module extends the base UserManager to provide persistent storage +for user configurations and MOS instances. +""" + +import json + +from datetime import datetime +from typing import Any + +from sqlalchemy import Column, String, Text + +from memos.configs.mem_os import MOSConfig +from memos.log import get_logger +from memos.mem_user.user_manager import Base, UserManager + + +logger = get_logger(__name__) + + +class UserConfig(Base): + """User configuration model for the database.""" + + __tablename__ = "user_configs" + + user_id = Column(String, primary_key=True) + config_data = Column(Text, nullable=False) # JSON string of MOSConfig + created_at = Column(String, nullable=False) # ISO format timestamp + updated_at = Column(String, nullable=False) # ISO format timestamp + + def __repr__(self): + return f"" + + +class PersistentUserManager(UserManager): + """Extended UserManager with configuration persistence.""" + + def __init__(self, db_path: str | None = None, user_id: str = "root"): + """Initialize the persistent user manager. + + Args: + db_path (str, optional): Path to the SQLite database file. + If None, uses default path in MEMOS_DIR. + user_id (str, optional): User ID. If None, uses default user ID. + """ + super().__init__(db_path, user_id) + + # Create user_configs table + Base.metadata.create_all(bind=self.engine) + logger.info("PersistentUserManager initialized with configuration storage") + + def _convert_datetime_strings(self, obj: Any) -> Any: + """Recursively convert datetime strings back to datetime objects in config dict. + + Args: + obj: The object to process (dict, list, or primitive type) + + Returns: + The object with datetime strings converted to datetime objects + """ + if isinstance(obj, dict): + result = {} + for key, value in obj.items(): + if key == "created_at" and isinstance(value, str): + try: + result[key] = datetime.fromisoformat(value) + except ValueError: + # If parsing fails, keep the original string + result[key] = value + else: + result[key] = self._convert_datetime_strings(value) + return result + elif isinstance(obj, list): + return [self._convert_datetime_strings(item) for item in obj] + else: + return obj + + def save_user_config(self, user_id: str, config: MOSConfig) -> bool: + """Save user configuration to database. + + Args: + user_id (str): The user ID. + config (MOSConfig): The user's MOS configuration. + + Returns: + bool: True if successful, False otherwise. + """ + session = self._get_session() + try: + # Convert config to JSON string with proper datetime handling + config_dict = config.model_dump(mode="json") + config_json = json.dumps(config_dict, indent=2) + + from datetime import datetime + + now = datetime.now().isoformat() + + # Check if config already exists + existing_config = ( + session.query(UserConfig).filter(UserConfig.user_id == user_id).first() + ) + + if existing_config: + # Update existing config + existing_config.config_data = config_json + existing_config.updated_at = now + logger.info(f"Updated configuration for user {user_id}") + else: + # Create new config + user_config = UserConfig( + user_id=user_id, config_data=config_json, created_at=now, updated_at=now + ) + session.add(user_config) + logger.info(f"Saved new configuration for user {user_id}") + + session.commit() + return True + + except Exception as e: + session.rollback() + logger.error(f"Error saving user config for {user_id}: {e}") + return False + finally: + session.close() + + def get_user_config(self, user_id: str) -> MOSConfig | None: + """Get user configuration from database. + + Args: + user_id (str): The user ID. + + Returns: + MOSConfig | None: The user's configuration or None if not found. + """ + session = self._get_session() + try: + user_config = session.query(UserConfig).filter(UserConfig.user_id == user_id).first() + + if user_config: + config_dict = json.loads(user_config.config_data) + # Convert datetime strings back to datetime objects + config_dict = self._convert_datetime_strings(config_dict) + return MOSConfig(**config_dict) + return None + + except Exception as e: + logger.error(f"Error loading user config for {user_id}: {e}") + return None + finally: + session.close() + + def delete_user_config(self, user_id: str) -> bool: + """Delete user configuration from database. + + Args: + user_id (str): The user ID. + + Returns: + bool: True if successful, False otherwise. + """ + session = self._get_session() + try: + user_config = session.query(UserConfig).filter(UserConfig.user_id == user_id).first() + + if user_config: + session.delete(user_config) + session.commit() + logger.info(f"Deleted configuration for user {user_id}") + return True + return False + + except Exception as e: + session.rollback() + logger.error(f"Error deleting user config for {user_id}: {e}") + return False + finally: + session.close() + + def list_user_configs(self) -> dict[str, MOSConfig]: + """List all user configurations. + + Returns: + Dict[str, MOSConfig]: Dictionary mapping user_id to MOSConfig. + """ + session = self._get_session() + try: + user_configs = session.query(UserConfig).all() + result = {} + + for user_config in user_configs: + try: + config_dict = json.loads(user_config.config_data) + # Convert datetime strings back to datetime objects + config_dict = self._convert_datetime_strings(config_dict) + result[user_config.user_id] = MOSConfig(**config_dict) + except Exception as e: + logger.error(f"Error parsing config for user {user_config.user_id}: {e}") + continue + + return result + + except Exception as e: + logger.error(f"Error listing user configs: {e}") + return {} + finally: + session.close() + + def create_user_with_config( + self, user_name: str, config: MOSConfig, role=None, user_id: str | None = None + ) -> str: + """Create a new user with configuration. + + Args: + user_name (str): Name of the user. + config (MOSConfig): The user's configuration. + role: User role (optional, uses default from UserManager). + user_id (str, optional): Custom user ID. + + Returns: + str: The created user ID. + + Raises: + ValueError: If user_name already exists. + """ + # Create user using parent method + created_user_id = self.create_user(user_name, role, user_id) + + # Save configuration + if not self.save_user_config(created_user_id, config): + logger.error(f"Failed to save configuration for user {created_user_id}") + + return created_user_id + + def delete_user(self, user_id: str) -> bool: + """Delete a user and their configuration. + + Args: + user_id (str): The user ID. + + Returns: + bool: True if successful, False otherwise. + """ + # Delete configuration first + self.delete_user_config(user_id) + + # Delete user using parent method + return super().delete_user(user_id) + + def get_user_cube_access(self, user_id: str) -> list[str]: + """Get list of cube IDs that a user has access to. + + Args: + user_id (str): The user ID. + + Returns: + list[str]: List of cube IDs the user can access. + """ + cubes = self.get_user_cubes(user_id) + return [cube.cube_id for cube in cubes]