MemoryOS 1.0.0__py3-none-any.whl → 1.1.1__py3-none-any.whl

memos/memories/textual/item.py

@@ -1,13 +1,48 @@
 """Defines memory item types for textual memory."""
 
+import json
 import uuid
 
 from datetime import datetime
-from typing import Literal
+from typing import Any, Literal
 
 from pydantic import BaseModel, ConfigDict, Field, field_validator
 
 
+ALLOWED_ROLES = {"user", "assistant", "system"}
+
+
+class SourceMessage(BaseModel):
+    """
+    Purpose: **memory provenance / traceability**.
+
+    Capture the minimal, reproducible origin context of a memory item so it can be
+    audited, traced, rolled back, or de-duplicated later.
+
+    Fields & conventions:
+        - type: Source kind (e.g., "chat", "doc", "web", "file", "system", ...).
+            If not provided, upstream logic may infer it:
+            presence of `role` ⇒ "chat"; otherwise ⇒ "doc".
+        - role: Conversation role ("user" | "assistant" | "system") when the
+            source is a chat turn.
+        - content: Minimal reproducible snippet from the source. If omitted,
+            upstream may fall back to `doc_path` / `url` / `message_id`.
+        - chat_time / message_id / doc_path: Locators for precisely pointing back
+            to the original record (timestamp, message id, document path).
+        - Extra fields: Allowed (`model_config.extra="allow"`) to carry arbitrary
+            provenance attributes (e.g., url, page, offset, span, local_confidence).
+    """
+
+    type: str | None = "chat"
+    role: Literal["user", "assistant", "system"] | None = None
+    chat_time: str | None = None
+    message_id: str | None = None
+    content: str | None = None
+    doc_path: str | None = None
+
+    model_config = ConfigDict(extra="allow")
+
+
 class TextualMemoryMetadata(BaseModel):
     """Metadata for a memory item.
 
@@ -33,7 +68,7 @@ class TextualMemoryMetadata(BaseModel):
         default=None,
         description="A numeric score (float between 0 and 100) indicating how certain you are about the accuracy or reliability of the memory.",
     )
-    source: Literal["conversation", "retrieved", "web", "file"] | None = Field(
+    source: Literal["conversation", "retrieved", "web", "file", "system"] | None = Field(
         default=None, description="The origin of the memory"
     )
     tags: list[str] | None = Field(
@@ -62,7 +97,7 @@ class TreeNodeTextualMemoryMetadata(TextualMemoryMetadata):
     memory_type: Literal["WorkingMemory", "LongTermMemory", "UserMemory", "OuterMemory"] = Field(
         default="WorkingMemory", description="Memory lifecycle type."
     )
-    sources: list[str] | None = Field(
+    sources: list[SourceMessage] | None = Field(
         default=None, description="Multiple origins of the memory (e.g., URLs, notes)."
     )
     embedding: list[float] | None = Field(
@@ -74,8 +109,8 @@ class TreeNodeTextualMemoryMetadata(TextualMemoryMetadata):
         description="The timestamp of the first creation to the memory. Useful "
         "for tracking memory initialization. Format: ISO 8601.",
     )
-    usage: list[str] | None = Field(
-        default=[],
+    usage: list[str] = Field(
+        default_factory=list,
         description="Usage history of this node",
     )
     background: str | None = Field(
@@ -83,12 +118,40 @@ class TreeNodeTextualMemoryMetadata(TextualMemoryMetadata):
         description="background of this node",
     )
 
-    @field_validator("sources")
+    @field_validator("sources", mode="before")
     @classmethod
-    def validate_sources(cls, v):
-        if v is not None and not isinstance(v, list):
-            raise ValueError("Sources must be a list of strings.")
-        return v
+    def coerce_sources(cls, v):
+        if v is None:
+            return v
+        if not isinstance(v, list):
+            raise TypeError("sources must be a list")
+        out = []
+        for item in v:
+            if isinstance(item, SourceMessage):
+                out.append(item)
+
+            elif isinstance(item, dict):
+                d = dict(item)
+                if d.get("type") is None:
+                    d["type"] = "chat" if d.get("role") in ALLOWED_ROLES else "doc"
+                out.append(SourceMessage(**d))
+
+            elif isinstance(item, str):
+                try:
+                    parsed = json.loads(item)
+                except Exception:
+                    parsed = None
+
+                if isinstance(parsed, dict):
+                    if parsed.get("type") is None:
+                        parsed["type"] = "chat" if parsed.get("role") in ALLOWED_ROLES else "doc"
+                    out.append(SourceMessage(**parsed))
+                else:
+                    out.append(SourceMessage(type="doc", content=item))
+
+            else:
+                out.append(SourceMessage(type="doc", content=str(item)))
+        return out
 
     def __str__(self) -> str:
         """Pretty string representation of the metadata."""
@@ -114,19 +177,17 @@ class TextualMemoryItem(BaseModel):
     id: str = Field(default_factory=lambda: str(uuid.uuid4()))
     memory: str
     metadata: (
-        TextualMemoryMetadata
+        SearchedTreeNodeTextualMemoryMetadata
         | TreeNodeTextualMemoryMetadata
-        | SearchedTreeNodeTextualMemoryMetadata
+        | TextualMemoryMetadata
     ) = Field(default_factory=TextualMemoryMetadata)
 
     model_config = ConfigDict(extra="forbid")
 
+    @field_validator("id")
     @classmethod
-    def validate_id(cls, v):
-        try:
-            uuid.UUID(v)
-        except ValueError as e:
-            raise ValueError("Invalid UUID format") from e
+    def _validate_id(cls, v: str) -> str:
+        uuid.UUID(v)
         return v
 
     @classmethod
@@ -136,6 +197,24 @@ class TextualMemoryItem(BaseModel):
     def to_dict(self) -> dict:
         return self.model_dump(exclude_none=True)
 
+    @field_validator("metadata", mode="before")
+    @classmethod
+    def _coerce_metadata(cls, v: Any):
+        if isinstance(
+            v,
+            SearchedTreeNodeTextualMemoryMetadata
+            | TreeNodeTextualMemoryMetadata
+            | TextualMemoryMetadata,
+        ):
+            return v
+        if isinstance(v, dict):
+            if v.get("relativity") is not None:
+                return SearchedTreeNodeTextualMemoryMetadata(**v)
+            if any(k in v for k in ("sources", "memory_type", "embedding", "background", "usage")):
+                return TreeNodeTextualMemoryMetadata(**v)
+            return TextualMemoryMetadata(**v)
+        return v
+
     def __str__(self) -> str:
         """Pretty string representation of the memory item."""
         return f"<ID: {self.id} | Memory: {self.memory} | Metadata: {self.metadata!s}>"

memos/memories/textual/naive.py

@@ -115,7 +115,7 @@ class NaiveTextMemory(BaseTextMemory):
                 self.memories[i] = memory_dict
                 break
 
-    def search(self, query: str, top_k: int) -> list[TextualMemoryItem]:
+    def search(self, query: str, top_k: int, **kwargs) -> list[TextualMemoryItem]:
         """Search for memories based on a query."""
         sims = [
             (memory, len(set(query.split()) & set(memory["memory"].split())))

memos/memories/textual/tree.py

@@ -2,12 +2,14 @@ import json
 import os
 import shutil
 import tempfile
+import time
 
 from datetime import datetime
 from pathlib import Path
 from typing import Any
 
 from memos.configs.memory import TreeTextMemoryConfig
+from memos.configs.reranker import RerankerConfigFactory
 from memos.embedders.factory import EmbedderFactory, OllamaEmbedder
 from memos.graph_dbs.factory import GraphStoreFactory, Neo4jGraphDB
 from memos.llms.factory import AzureLLM, LLMFactory, OllamaLLM, OpenAILLM
@@ -19,6 +21,7 @@ from memos.memories.textual.tree_text_memory.retrieve.internet_retriever_factory
     InternetRetrieverFactory,
 )
 from memos.memories.textual.tree_text_memory.retrieve.searcher import Searcher
+from memos.reranker.factory import RerankerFactory
 from memos.types import MessageList
 
 
@@ -30,21 +33,59 @@ class TreeTextMemory(BaseTextMemory):
 
     def __init__(self, config: TreeTextMemoryConfig):
         """Initialize memory with the given configuration."""
+        time_start = time.time()
         self.config: TreeTextMemoryConfig = config
         self.extractor_llm: OpenAILLM | OllamaLLM | AzureLLM = LLMFactory.from_config(
             config.extractor_llm
         )
+        logger.info(f"time init: extractor_llm time is: {time.time() - time_start}")
+
+        time_start_ex = time.time()
         self.dispatcher_llm: OpenAILLM | OllamaLLM | AzureLLM = LLMFactory.from_config(
             config.dispatcher_llm
         )
+        logger.info(f"time init: dispatcher_llm time is: {time.time() - time_start_ex}")
+
+        time_start_em = time.time()
         self.embedder: OllamaEmbedder = EmbedderFactory.from_config(config.embedder)
+        logger.info(f"time init: embedder time is: {time.time() - time_start_em}")
+
+        time_start_gs = time.time()
         self.graph_store: Neo4jGraphDB = GraphStoreFactory.from_config(config.graph_db)
+        logger.info(f"time init: graph_store time is: {time.time() - time_start_gs}")
+
+        time_start_rr = time.time()
+        if config.reranker is None:
+            default_cfg = RerankerConfigFactory.model_validate(
+                {
+                    "backend": "cosine_local",
+                    "config": {
+                        "level_weights": {"topic": 1.0, "concept": 1.0, "fact": 1.0},
+                        "level_field": "background",
+                    },
+                }
+            )
+            self.reranker = RerankerFactory.from_config(default_cfg)
+        else:
+            self.reranker = RerankerFactory.from_config(config.reranker)
+        logger.info(f"time init: reranker time is: {time.time() - time_start_rr}")
         self.is_reorganize = config.reorganize
 
+        time_start_mm = time.time()
         self.memory_manager: MemoryManager = MemoryManager(
-            self.graph_store, self.embedder, self.extractor_llm, is_reorganize=self.is_reorganize
+            self.graph_store,
+            self.embedder,
+            self.extractor_llm,
+            memory_size=config.memory_size
+            or {
+                "WorkingMemory": 20,
+                "LongTermMemory": 1500,
+                "UserMemory": 480,
+            },
+            is_reorganize=self.is_reorganize,
         )
-
+        logger.info(f"time init: memory_manager time is: {time.time() - time_start_mm}")
+        time_start_ir = time.time()
         # Create internet retriever if configured
         self.internet_retriever = None
         if config.internet_retriever is not None:
@@ -56,6 +97,7 @@ class TreeTextMemory(BaseTextMemory):
             )
         else:
             logger.info("No internet retriever configured")
+        logger.info(f"time init: internet_retriever time is: {time.time() - time_start_ir}")
 
     def add(self, memories: list[TextualMemoryItem | dict[str, Any]]) -> list[str]:
         """Add memories.
@@ -96,6 +138,8 @@ class TreeTextMemory(BaseTextMemory):
         mode: str = "fast",
         memory_type: str = "All",
         manual_close_internet: bool = False,
+        moscube: bool = False,
+        search_filter: dict | None = None,
     ) -> list[TextualMemoryItem]:
         """Search for memories based on a query.
         User query -> TaskGoalParser -> MemoryPathResolver ->
@@ -110,6 +154,12 @@ class TreeTextMemory(BaseTextMemory):
             memory_type (str): Type restriction for search.
             ['All', 'WorkingMemory', 'LongTermMemory', 'UserMemory']
             manual_close_internet (bool): If True, the internet retriever will be closed by this search, it high priority than config.
+            moscube (bool): whether you use moscube to answer questions
+            search_filter (dict, optional): Optional metadata filters for search results.
+                - Keys correspond to memory metadata fields (e.g., "user_id", "session_id").
+                - Values are exact-match conditions.
+                Example: {"user_id": "123", "session_id": "abc"}
+                If None, no additional filtering is applied.
         Returns:
             list[TextualMemoryItem]: List of matching memories.
         """
@@ -121,16 +171,20 @@ class TreeTextMemory(BaseTextMemory):
                 self.dispatcher_llm,
                 self.graph_store,
                 self.embedder,
+                self.reranker,
                 internet_retriever=None,
+                moscube=moscube,
             )
         else:
             searcher = Searcher(
                 self.dispatcher_llm,
                 self.graph_store,
                 self.embedder,
+                self.reranker,
                 internet_retriever=self.internet_retriever,
+                moscube=moscube,
             )
-        return searcher.search(query, top_k, info, mode, memory_type)
+        return searcher.search(query, top_k, info, mode, memory_type, search_filter)
 
     def get_relevant_subgraph(
         self, query: str, top_k: int = 5, depth: int = 2, center_status: str = "activated"

memos/memories/textual/tree_text_memory/organize/handler.py

@@ -1,5 +1,6 @@
 import json
 import re
+
 from datetime import datetime
 
 from dateutil import parser
@@ -14,6 +15,7 @@ from memos.templates.tree_reorganize_prompts import (
     MEMORY_RELATION_RESOLVER_PROMPT,
 )
 
+
 logger = get_logger(__name__)
 
 
@@ -50,12 +52,12 @@ class NodeHandler:
             ]
             result = self.llm.generate(prompt).strip()
             if result == "contradictory":
-                logger.warning(
+                logger.info(
                     f'detected "{memory.memory}" <==CONFLICT==> "{embedding_candidate.memory}"'
                 )
                 detected_relationships.append([memory, embedding_candidate, "contradictory"])
             elif result == "redundant":
-                logger.warning(
+                logger.info(
                     f'detected "{memory.memory}" <==REDUNDANT==> "{embedding_candidate.memory}"'
                 )
                 detected_relationships.append([memory, embedding_candidate, "redundant"])

memos/memories/textual/tree_text_memory/organize/manager.py

@@ -1,8 +1,10 @@
+import traceback
 import uuid
 
-from concurrent.futures import ThreadPoolExecutor, as_completed
+from concurrent.futures import as_completed
 from datetime import datetime
 
+from memos.context.context import ContextThreadPoolExecutor
 from memos.embedders.factory import OllamaEmbedder
 from memos.graph_dbs.neo4j import Neo4jGraphDB
 from memos.llms.factory import AzureLLM, OllamaLLM, OpenAILLM
@@ -55,24 +57,35 @@ class MemoryManager:
         """
         added_ids: list[str] = []
 
-        with ThreadPoolExecutor(max_workers=8) as executor:
+        with ContextThreadPoolExecutor(max_workers=8) as executor:
             futures = {executor.submit(self._process_memory, m): m for m in memories}
-            for future in as_completed(futures):
+            for future in as_completed(futures, timeout=60):
                 try:
                     ids = future.result()
                     added_ids.extend(ids)
                 except Exception as e:
                     logger.exception("Memory processing error: ", exc_info=e)
 
-        self.graph_store.remove_oldest_memory(
-            memory_type="WorkingMemory", keep_latest=self.memory_size["WorkingMemory"]
-        )
-        self.graph_store.remove_oldest_memory(
-            memory_type="LongTermMemory", keep_latest=self.memory_size["LongTermMemory"]
-        )
-        self.graph_store.remove_oldest_memory(
-            memory_type="UserMemory", keep_latest=self.memory_size["UserMemory"]
-        )
+        try:
+            self.graph_store.remove_oldest_memory(
+                memory_type="WorkingMemory", keep_latest=self.memory_size["WorkingMemory"]
+            )
+        except Exception:
+            logger.warning(f"Remove WorkingMemory error: {traceback.format_exc()}")
+
+        try:
+            self.graph_store.remove_oldest_memory(
+                memory_type="LongTermMemory", keep_latest=self.memory_size["LongTermMemory"]
+            )
+        except Exception:
+            logger.warning(f"Remove LongTermMemory error: {traceback.format_exc()}")
+
+        try:
+            self.graph_store.remove_oldest_memory(
+                memory_type="UserMemory", keep_latest=self.memory_size["UserMemory"]
+            )
+        except Exception:
+            logger.warning(f"Remove UserMemory error: {traceback.format_exc()}")
 
         self._refresh_memory_size()
         return added_ids
@@ -82,12 +95,12 @@ class MemoryManager:
         Replace WorkingMemory
         """
         working_memory_top_k = memories[: self.memory_size["WorkingMemory"]]
-        with ThreadPoolExecutor(max_workers=8) as executor:
+        with ContextThreadPoolExecutor(max_workers=8) as executor:
             futures = [
                 executor.submit(self._add_memory_to_db, memory, "WorkingMemory")
                 for memory in working_memory_top_k
             ]
-            for future in as_completed(futures):
+            for future in as_completed(futures, timeout=60):
                 try:
                     future.result()
                 except Exception as e:
@@ -102,6 +115,7 @@ class MemoryManager:
         """
         Return the cached memory type counts.
         """
+        self._refresh_memory_size()
         return self.current_memory_size
 
     def _refresh_memory_size(self) -> None:

memos/memories/textual/tree_text_memory/organize/relation_reason_detector.py

@@ -46,7 +46,7 @@ class RelationAndReasoningDetector:
                     "sequence_links": [],
                     "aggregate_nodes": [],
                 }
-
+            """
             nearest = self.graph_store.get_neighbors_by_tag(
                 tags=node.metadata.tags,
                 exclude_ids=exclude_ids,
@@ -55,7 +55,6 @@ class RelationAndReasoningDetector:
             )
             nearest = [GraphDBNode(**cand_data) for cand_data in nearest]
 
-            """
             # 1) Pairwise relations (including CAUSE/CONDITION/CONFLICT)
             pairwise = self._detect_pairwise_causal_condition_relations(node, nearest)
             results["relations"].extend(pairwise["relations"])