MemoryOS 0.2.0__py3-none-any.whl → 0.2.2__py3-none-any.whl

memos/graph_dbs/neo4j_community.py

@@ -0,0 +1,300 @@
+from typing import Any
+
+from memos.configs.graph_db import Neo4jGraphDBConfig
+from memos.graph_dbs.neo4j import Neo4jGraphDB, _prepare_node_metadata
+from memos.log import get_logger
+from memos.vec_dbs.factory import VecDBFactory
+from memos.vec_dbs.item import VecDBItem
+
+
+logger = get_logger(__name__)
+
+
+class Neo4jCommunityGraphDB(Neo4jGraphDB):
+    """
+    Neo4j Community Edition graph memory store.
+
+    Note:
+        This class avoids Enterprise-only features:
+        - No multi-database support
+        - No vector index
+        - No CREATE DATABASE
+    """
+
+    def __init__(self, config: Neo4jGraphDBConfig):
+        assert config.auto_create is False
+        assert config.use_multi_db is False
+        # Init vector database
+        self.vec_db = VecDBFactory.from_config(config.vec_config)
+        # Call parent init
+        super().__init__(config)
+
+    def create_index(
+        self,
+        label: str = "Memory",
+        vector_property: str = "embedding",
+        dimensions: int = 1536,
+        index_name: str = "memory_vector_index",
+    ) -> None:
+        """
+        Create the vector index for embedding and datetime indexes for created_at and updated_at fields.
+        """
+        # Create indexes
+        self._create_basic_property_indexes()
+
+    def add_node(self, id: str, memory: str, metadata: dict[str, Any]) -> None:
+        if not self.config.use_multi_db and self.config.user_name:
+            metadata["user_name"] = self.config.user_name
+
+        # Safely process metadata
+        metadata = _prepare_node_metadata(metadata)
+
+        # Extract required fields
+        embedding = metadata.pop("embedding", None)
+        if embedding is None:
+            raise ValueError(f"Missing 'embedding' in metadata for node {id}")
+
+        # Merge node and set metadata
+        created_at = metadata.pop("created_at")
+        updated_at = metadata.pop("updated_at")
+        vector_sync_status = "success"
+
+        try:
+            # Write to Vector DB
+            item = VecDBItem(
+                id=id,
+                vector=embedding,
+                payload={
+                    "memory": memory,
+                    "vector_sync": vector_sync_status,
+                    **metadata,  # unpack all metadata keys to top-level
+                },
+            )
+            self.vec_db.add([item])
+        except Exception as e:
+            logger.warning(f"[VecDB] Vector insert failed for node {id}: {e}")
+            vector_sync_status = "failed"
+
+        metadata["vector_sync"] = vector_sync_status
+        query = """
+            MERGE (n:Memory {id: $id})
+            SET n.memory = $memory,
+                n.created_at = datetime($created_at),
+                n.updated_at = datetime($updated_at),
+                n += $metadata
+        """
+        with self.driver.session(database=self.db_name) as session:
+            session.run(
+                query,
+                id=id,
+                memory=memory,
+                created_at=created_at,
+                updated_at=updated_at,
+                metadata=metadata,
+            )
+
+    def get_children_with_embeddings(self, id: str) -> list[dict[str, Any]]:
+        where_user = ""
+        params = {"id": id}
+
+        if not self.config.use_multi_db and self.config.user_name:
+            where_user = "AND p.user_name = $user_name AND c.user_name = $user_name"
+            params["user_name"] = self.config.user_name
+
+        query = f"""
+                MATCH (p:Memory)-[:PARENT]->(c:Memory)
+                WHERE p.id = $id {where_user}
+                RETURN c.id AS id, c.memory AS memory
+            """
+
+        with self.driver.session(database=self.db_name) as session:
+            result = session.run(query, params)
+            child_nodes = [{"id": r["id"], "memory": r["memory"]} for r in result]
+
+        # Get embeddings from vector DB
+        ids = [n["id"] for n in child_nodes]
+        vec_items = {v.id: v.vector for v in self.vec_db.get_by_ids(ids)}
+
+        # Merge results
+        for node in child_nodes:
+            node["embedding"] = vec_items.get(node["id"])
+
+        return child_nodes
+
+    # Search / recall operations
+    def search_by_embedding(
+        self,
+        vector: list[float],
+        top_k: int = 5,
+        scope: str | None = None,
+        status: str | None = None,
+        threshold: float | None = None,
+    ) -> list[dict]:
+        """
+        Retrieve node IDs based on vector similarity using external vector DB.
+
+        Args:
+            vector (list[float]): The embedding vector representing query semantics.
+            top_k (int): Number of top similar nodes to retrieve.
+            scope (str, optional): Memory type filter (e.g., 'WorkingMemory', 'LongTermMemory').
+            status (str, optional): Node status filter (e.g., 'activated', 'archived').
+            threshold (float, optional): Minimum similarity score threshold (0 ~ 1).
+
+        Returns:
+            list[dict]: A list of dicts with 'id' and 'score', ordered by similarity.
+
+        Notes:
+            - This method uses an external vector database (not Neo4j) to perform the search.
+            - If 'scope' is provided, it restricts results to nodes with matching memory_type.
+            - If 'status' is provided, it further filters nodes by status.
+            - If 'threshold' is provided, only results with score >= threshold will be returned.
+            - The returned IDs can be used to fetch full node data from Neo4j if needed.
+        """
+        # Build VecDB filter
+        vec_filter = {}
+        if scope:
+            vec_filter["memory_type"] = scope
+        if status:
+            vec_filter["status"] = status
+        vec_filter["vector_sync"] = "success"
+        vec_filter["user_name"] = self.config.user_name
+
+        # Perform vector search
+        results = self.vec_db.search(query_vector=vector, top_k=top_k, filter=vec_filter)
+
+        # Filter by threshold
+        if threshold is not None:
+            results = [r for r in results if r.score is None or r.score >= threshold]
+
+        # Return consistent format
+        return [{"id": r.id, "score": r.score} for r in results]
+
+    def get_all_memory_items(self, scope: str) -> list[dict]:
+        """
+        Retrieve all memory items of a specific memory_type.
+
+        Args:
+            scope (str): Must be one of 'WorkingMemory', 'LongTermMemory', or 'UserMemory'.
+
+        Returns:
+            list[dict]: Full list of memory items under this scope.
+        """
+        if scope not in {"WorkingMemory", "LongTermMemory", "UserMemory"}:
+            raise ValueError(f"Unsupported memory type scope: {scope}")
+
+        where_clause = "WHERE n.memory_type = $scope"
+        params = {"scope": scope}
+
+        if not self.config.use_multi_db and self.config.user_name:
+            where_clause += " AND n.user_name = $user_name"
+            params["user_name"] = self.config.user_name
+
+        query = f"""
+            MATCH (n:Memory)
+            {where_clause}
+            RETURN n
+            """
+
+        with self.driver.session(database=self.db_name) as session:
+            results = session.run(query, params)
+            return [self._parse_node(dict(record["n"])) for record in results]
+
+    def clear(self) -> None:
+        """
+        Clear the entire graph if the target database exists.
+        """
+        # Step 1: clear Neo4j part via parent logic
+        super().clear()
+
+        # Step2: Clear the vector db
+        try:
+            items = self.vec_db.get_by_filter({"user_name": self.config.user_name})
+            if items:
+                self.vec_db.delete([item.id for item in items])
+                logger.info(f"Cleared {len(items)} vectors for user '{self.config.user_name}'.")
+            else:
+                logger.info(f"No vectors to clear for user '{self.config.user_name}'.")
+        except Exception as e:
+            logger.warning(f"Failed to clear vector DB for user '{self.config.user_name}': {e}")
+
+    def drop_database(self) -> None:
+        """
+        Permanently delete the entire database this instance is using.
+        WARNING: This operation is destructive and cannot be undone.
+        """
+        raise ValueError(
+            f"Refusing to drop protected database: {self.db_name} in "
+            f"Shared Database Multi-Tenant mode"
+        )
+
+    # Avoid enterprise feature
+    def _ensure_database_exists(self):
+        pass
+
+    def _create_basic_property_indexes(self) -> None:
+        """
+        Create standard B-tree indexes on memory_type, created_at,
+        and updated_at fields.
+        Create standard B-tree indexes on user_name when use Shared Database
+        Multi-Tenant Mode
+        """
+        # Step 1: Neo4j indexes
+        try:
+            with self.driver.session(database=self.db_name) as session:
+                session.run("""
+                    CREATE INDEX memory_type_index IF NOT EXISTS
+                    FOR (n:Memory) ON (n.memory_type)
+                """)
+                logger.debug("Index 'memory_type_index' ensured.")
+
+                session.run("""
+                    CREATE INDEX memory_created_at_index IF NOT EXISTS
+                    FOR (n:Memory) ON (n.created_at)
+                """)
+                logger.debug("Index 'memory_created_at_index' ensured.")
+
+                session.run("""
+                    CREATE INDEX memory_updated_at_index IF NOT EXISTS
+                    FOR (n:Memory) ON (n.updated_at)
+                """)
+                logger.debug("Index 'memory_updated_at_index' ensured.")
+
+                if not self.config.use_multi_db and self.config.user_name:
+                    session.run(
+                        """
+                        CREATE INDEX memory_user_name_index IF NOT EXISTS
+                        FOR (n:Memory) ON (n.user_name)
+                        """
+                    )
+                logger.debug("Index 'memory_user_name_index' ensured.")
+        except Exception as e:
+            logger.warning(f"Failed to create basic property indexes: {e}")
+
+        # Step 2: VectorDB indexes
+        try:
+            if hasattr(self.vec_db, "ensure_payload_indexes"):
+                self.vec_db.ensure_payload_indexes(["user_name", "memory_type", "status"])
+            else:
+                logger.debug("VecDB does not support payload index creation; skipping.")
+        except Exception as e:
+            logger.warning(f"Failed to create VecDB payload indexes: {e}")
+
+    def _parse_node(self, node_data: dict[str, Any]) -> dict[str, Any]:
+        """Parse Neo4j node and optionally fetch embedding from vector DB."""
+        node = node_data.copy()
+
+        # Convert Neo4j datetime to string
+        for time_field in ("created_at", "updated_at"):
+            if time_field in node and hasattr(node[time_field], "isoformat"):
+                node[time_field] = node[time_field].isoformat()
+        node.pop("user_name", None)
+
+        new_node = {"id": node.pop("id"), "memory": node.pop("memory", ""), "metadata": node}
+        try:
+            vec_item = self.vec_db.get_by_id(new_node["id"])
+            if vec_item and vec_item.vector:
+                new_node["metadata"]["embedding"] = vec_item.vector
+        except Exception as e:
+            logger.warning(f"Failed to fetch vector for node {new_node['id']}: {e}")
+            new_node["metadata"]["embedding"] = None
+        return new_node

memos/llms/base.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from collections.abc import Generator
 
 from memos.configs.llm import BaseLLMConfig
 from memos.types import MessageList
@@ -14,3 +15,11 @@ class BaseLLM(ABC):
     @abstractmethod
     def generate(self, messages: MessageList, **kwargs) -> str:
         """Generate a response from the LLM."""
+
+    @abstractmethod
+    def generate_stream(self, messages: MessageList, **kwargs) -> Generator[str, None, None]:
+        """
+        (Optional) Generate a streaming response from the LLM.
+        Subclasses should override this if they support streaming.
+        By default, this raises NotImplementedError.
+        """

memos/llms/deepseek.py

@@ -0,0 +1,54 @@
+from collections.abc import Generator
+
+from memos.configs.llm import DeepSeekLLMConfig
+from memos.llms.openai import OpenAILLM
+from memos.llms.utils import remove_thinking_tags
+from memos.log import get_logger
+from memos.types import MessageList
+
+
+logger = get_logger(__name__)
+
+
+class DeepSeekLLM(OpenAILLM):
+    """DeepSeek LLM via OpenAI-compatible API."""
+
+    def __init__(self, config: DeepSeekLLMConfig):
+        super().__init__(config)
+
+    def generate(self, messages: MessageList) -> str:
+        """Generate a response from DeepSeek."""
+        response = self.client.chat.completions.create(
+            model=self.config.model_name_or_path,
+            messages=messages,
+            temperature=self.config.temperature,
+            max_tokens=self.config.max_tokens,
+            top_p=self.config.top_p,
+            extra_body=self.config.extra_body,
+        )
+        logger.info(f"Response from DeepSeek: {response.model_dump_json()}")
+        response_content = response.choices[0].message.content
+        if self.config.remove_think_prefix:
+            return remove_thinking_tags(response_content)
+        else:
+            return response_content
+
+    def generate_stream(self, messages: MessageList, **kwargs) -> Generator[str, None, None]:
+        """Stream response from DeepSeek."""
+        response = self.client.chat.completions.create(
+            model=self.config.model_name_or_path,
+            messages=messages,
+            stream=True,
+            temperature=self.config.temperature,
+            max_tokens=self.config.max_tokens,
+            top_p=self.config.top_p,
+            extra_body=self.config.extra_body,
+        )
+        # Streaming chunks of text
+        for chunk in response:
+            delta = chunk.choices[0].delta
+            if hasattr(delta, "reasoning_content") and delta.reasoning_content:
+                yield delta.reasoning_content
+
+            if hasattr(delta, "content") and delta.content:
+                yield delta.content

memos/llms/factory.py

@@ -2,9 +2,13 @@ from typing import Any, ClassVar
 
 from memos.configs.llm import LLMConfigFactory
 from memos.llms.base import BaseLLM
+from memos.llms.deepseek import DeepSeekLLM
 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
+from memos.llms.openai import AzureLLM, OpenAILLM
+from memos.llms.qwen import QwenLLM
+from memos.llms.vllm import VLLMLLM
 
 
 class LLMFactory(BaseLLM):
@@ -12,8 +16,13 @@ class LLMFactory(BaseLLM):
 
     backend_to_class: ClassVar[dict[str, Any]] = {
         "openai": OpenAILLM,
+        "azure": AzureLLM,
         "ollama": OllamaLLM,
         "huggingface": HFLLM,
+        "huggingface_singleton": HFSingletonLLM,  # Add singleton version
+        "vllm": VLLMLLM,
+        "qwen": QwenLLM,
+        "deepseek": DeepSeekLLM,
     }
 
     @classmethod

memos/llms/hf.py

@@ -1,4 +1,5 @@
-import torch
+from collections.abc import Generator
+from typing import Any
 
 from transformers import (
     AutoModelForCausalLM,
@@ -71,6 +72,26 @@ class HFLLM(BaseLLM):
         else:
             return self._generate_with_cache(prompt, past_key_values)
 
+    def generate_stream(
+        self, messages: MessageList, past_key_values: DynamicCache | None = None
+    ) -> Generator[str, None, None]:
+        """
+        Generate a streaming response from the model.
+        Args:
+            messages (MessageList): Chat messages for prompt construction.
+            past_key_values (DynamicCache | None): Optional KV cache for fast generation.
+        Yields:
+            str: Streaming model response chunks.
+        """
+        prompt = self.tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=self.config.add_generation_prompt
+        )
+        logger.info(f"HFLLM streaming prompt: {prompt}")
+        if past_key_values is None:
+            yield from self._generate_full_stream(prompt)
+        else:
+            yield from self._generate_with_cache_stream(prompt, past_key_values)
+
     def _generate_full(self, prompt: str) -> str:
         """
         Generate output from scratch using the full prompt.
@@ -104,6 +125,73 @@ class HFLLM(BaseLLM):
             else response
         )
 
+    def _generate_full_stream(self, prompt: str) -> Generator[str, None, None]:
+        """
+        Generate output from scratch using the full prompt with streaming.
+        Args:
+            prompt (str): The input prompt string.
+        Yields:
+            str: Streaming response chunks.
+        """
+        import torch
+
+        inputs = self.tokenizer([prompt], return_tensors="pt").to(self.model.device)
+
+        # Get generation parameters
+        max_new_tokens = getattr(self.config, "max_tokens", 128)
+        remove_think_prefix = getattr(self.config, "remove_think_prefix", False)
+
+        # Manual streaming generation
+        generated_ids = inputs.input_ids.clone()
+        accumulated_text = ""
+
+        for _ in range(max_new_tokens):
+            # Forward pass
+            with torch.no_grad():
+                outputs = self.model(
+                    input_ids=generated_ids,
+                    use_cache=True,
+                    return_dict=True,
+                )
+
+            # Get next token logits
+            next_token_logits = outputs.logits[:, -1, :]
+
+            # Apply logits processors if sampling
+            if getattr(self.config, "do_sample", True):
+                batch_size, _ = next_token_logits.size()
+                dummy_ids = torch.zeros(
+                    (batch_size, 1), dtype=torch.long, device=next_token_logits.device
+                )
+                filtered_logits = self.logits_processors(dummy_ids, next_token_logits)
+                probs = torch.softmax(filtered_logits, dim=-1)
+                next_token = torch.multinomial(probs, num_samples=1)
+            else:
+                next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
+
+            # Check for EOS token
+            if self._should_stop(next_token):
+                break
+
+            # Append new token
+            generated_ids = torch.cat([generated_ids, next_token], dim=-1)
+
+            # Decode and yield the new token
+            new_token_text = self.tokenizer.decode(next_token[0], skip_special_tokens=True)
+            if new_token_text:  # Only yield non-empty tokens
+                accumulated_text += new_token_text
+
+                # Apply thinking tag removal if enabled
+                if remove_think_prefix:
+                    processed_text = remove_thinking_tags(accumulated_text)
+                    # Only yield the difference (new content)
+                    if len(processed_text) > len(accumulated_text) - len(new_token_text):
+                        yield processed_text[len(accumulated_text) - len(new_token_text) :]
+                    else:
+                        yield new_token_text
+                else:
+                    yield new_token_text
+
     def _generate_with_cache(self, query: str, kv: DynamicCache) -> str:
         """
         Generate output incrementally using an existing KV cache.
@@ -113,6 +201,8 @@ class HFLLM(BaseLLM):
         Returns:
             str: Model response.
         """
+        import torch
+
         query_ids = self.tokenizer(
             query, return_tensors="pt", add_special_tokens=False
         ).input_ids.to(self.model.device)
@@ -137,10 +227,70 @@ class HFLLM(BaseLLM):
             else response
         )
 
-    @torch.no_grad()
-    def _prefill(
-        self, input_ids: torch.Tensor, kv: DynamicCache
-    ) -> tuple[torch.Tensor, DynamicCache]:
+    def _generate_with_cache_stream(
+        self, query: str, kv: DynamicCache
+    ) -> Generator[str, None, None]:
+        """
+        Generate output incrementally using an existing KV cache with streaming.
+        Args:
+            query (str): The new user query string.
+            kv (DynamicCache): The prefilled KV cache.
+        Yields:
+            str: Streaming response chunks.
+        """
+        query_ids = self.tokenizer(
+            query, return_tensors="pt", add_special_tokens=False
+        ).input_ids.to(self.model.device)
+
+        max_new_tokens = getattr(self.config, "max_tokens", 128)
+        remove_think_prefix = getattr(self.config, "remove_think_prefix", False)
+
+        # Initial forward pass
+        logits, kv = self._prefill(query_ids, kv)
+        next_token = self._select_next_token(logits)
+
+        # Yield first token
+        first_token_text = self.tokenizer.decode(next_token[0], skip_special_tokens=True)
+        accumulated_text = ""
+        if first_token_text:
+            accumulated_text += first_token_text
+            if remove_think_prefix:
+                processed_text = remove_thinking_tags(accumulated_text)
+                if len(processed_text) > len(accumulated_text) - len(first_token_text):
+                    yield processed_text[len(accumulated_text) - len(first_token_text) :]
+                else:
+                    yield first_token_text
+            else:
+                yield first_token_text
+
+        generated = [next_token]
+
+        # Continue generation
+        for _ in range(max_new_tokens - 1):
+            if self._should_stop(next_token):
+                break
+            logits, kv = self._prefill(next_token, kv)
+            next_token = self._select_next_token(logits)
+
+            # Decode and yield the new token
+            new_token_text = self.tokenizer.decode(next_token[0], skip_special_tokens=True)
+            if new_token_text:
+                accumulated_text += new_token_text
+
+                # Apply thinking tag removal if enabled
+                if remove_think_prefix:
+                    processed_text = remove_thinking_tags(accumulated_text)
+                    # Only yield the difference (new content)
+                    if len(processed_text) > len(accumulated_text) - len(new_token_text):
+                        yield processed_text[len(accumulated_text) - len(new_token_text) :]
+                    else:
+                        yield new_token_text
+                else:
+                    yield new_token_text
+
+            generated.append(next_token)
+
+    def _prefill(self, input_ids: Any, kv: DynamicCache) -> tuple[Any, DynamicCache]:
         """
         Forward the model once, returning last-step logits and updated KV cache.
         Args:
@@ -149,15 +299,18 @@ class HFLLM(BaseLLM):
         Returns:
             tuple[torch.Tensor, DynamicCache]: (last-step logits, updated KV cache)
         """
-        out = self.model(
-            input_ids=input_ids,
-            use_cache=True,
-            past_key_values=kv,
-            return_dict=True,
-        )
+        import torch
+
+        with torch.no_grad():
+            out = self.model(
+                input_ids=input_ids,
+                use_cache=True,
+                past_key_values=kv,
+                return_dict=True,
+            )
         return out.logits[:, -1, :], out.past_key_values
 
-    def _select_next_token(self, logits: torch.Tensor) -> torch.Tensor:
+    def _select_next_token(self, logits: Any) -> Any:
         """
         Select the next token from logits using sampling or argmax, depending on config.
         Args:
@@ -165,6 +318,8 @@ class HFLLM(BaseLLM):
         Returns:
             torch.Tensor: Selected token ID(s).
         """
+        import torch
+
         if getattr(self.config, "do_sample", True):
             batch_size, _ = logits.size()
             dummy_ids = torch.zeros((batch_size, 1), dtype=torch.long, device=logits.device)
@@ -173,7 +328,7 @@ class HFLLM(BaseLLM):
             return torch.multinomial(probs, num_samples=1)
         return torch.argmax(logits, dim=-1, keepdim=True)
 
-    def _should_stop(self, token: torch.Tensor) -> bool:
+    def _should_stop(self, token: Any) -> bool:
         """
         Check if the given token is the EOS (end-of-sequence) token.
         Args:
@@ -197,6 +352,8 @@ class HFLLM(BaseLLM):
         Returns:
             DynamicCache: The constructed KV cache object.
         """
+        import torch
+
         # Accept multiple input types and convert to standard chat messages
         if isinstance(messages, str):
             messages = [