MemoryOS 2.0.3__py3-none-any.whl

memos/multi_mem_cube/views.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Protocol
+
+
+if TYPE_CHECKING:
+    from memos.api.product_models import APIADDRequest, APIFeedbackRequest, APISearchRequest
+
+
+class MemCubeView(Protocol):
+    """
+    A high-level cube view used by AddHandler.
+    It may wrap a single logical cube or multiple cubes,
+    but exposes a unified add_memories interface.
+    """
+
+    def add_memories(self, add_req: APIADDRequest) -> list[dict[str, Any]]:
+        """
+        Process add_req, extract memories and write them into one or more cubes.
+
+        Returns:
+            A list of memory dicts, each item should at least contain:
+            - memory
+            - memory_id
+            - memory_type
+            - cube_id
+        """
+        ...
+
+    def search_memories(self, search_req: APISearchRequest) -> dict[str, Any]:
+        """
+        Process search_req, read memories from one or more cubes and search them.
+
+        Returns:
+            A list of memory dicts, each item should at least contain:
+            - memory
+            - memory_id
+            - memory_type
+            - cube_id
+        """
+        ...
+
+    def feedback_memories(self, feedback_req: APIFeedbackRequest) -> dict[str, Any]:
+        """
+        Process feedback_req, read memories from one or more cubes and feedback them.
+
+        Returns:
+            A list of memory dicts, each item should at least contain:
+            - memory
+            - memory_id
+            - memory_type
+            - cube_id
+        """
+        ...

memos/parsers/base.py

@@ -0,0 +1,15 @@
+from abc import ABC, abstractmethod
+
+from memos.configs.parser import BaseParserConfig
+
+
+class BaseParser(ABC):
+    """Base class for all parsers."""
+
+    @abstractmethod
+    def __init__(self, config: BaseParserConfig):
+        """Initialize the parser with the given configuration."""
+
+    @abstractmethod
+    def parse(self, file_path: str) -> str:
+        """Parse the file at the given path and return its content as a string."""

memos/parsers/factory.py

@@ -0,0 +1,21 @@
+from typing import Any, ClassVar
+
+from memos.configs.parser import ParserConfigFactory
+from memos.memos_tools.singleton import singleton_factory
+from memos.parsers.base import BaseParser
+from memos.parsers.markitdown import MarkItDownParser
+
+
+class ParserFactory(BaseParser):
+    """Factory class for creating Parser instances."""
+
+    backend_to_class: ClassVar[dict[str, Any]] = {"markitdown": MarkItDownParser}
+
+    @classmethod
+    @singleton_factory()
+    def from_config(cls, config_factory: ParserConfigFactory) -> BaseParser:
+        backend = config_factory.backend
+        if backend not in cls.backend_to_class:
+            raise ValueError(f"Invalid backend: {backend}")
+        parser_class = cls.backend_to_class[backend]
+        return parser_class(config_factory.config)

memos/parsers/markitdown.py

@@ -0,0 +1,28 @@
+from memos.configs.parser import MarkItDownParserConfig
+from memos.dependency import require_python_package
+from memos.log import get_logger
+from memos.parsers.base import BaseParser
+
+
+logger = get_logger(__name__)
+
+
+class MarkItDownParser(BaseParser):
+    """MarkItDown Parser class."""
+
+    def __init__(self, config: MarkItDownParserConfig):
+        self.config = config
+
+    @require_python_package(
+        import_name="markitdown",
+        install_command="pip install markitdown[all]",
+        install_link="https://github.com/microsoft/markitdown",
+    )
+    def parse(self, file_path: str) -> str:
+        from markitdown import MarkItDown
+
+        """Parse the file at the given path and return its content as a MarkDown string."""
+        md = MarkItDown(enable_plugins=False)
+        result = md.convert(file_path)
+
+        return result.text_content

memos/reranker/__init__.py

@@ -0,0 +1,4 @@
+from .factory import RerankerFactory
+
+
+__all__ = ["RerankerFactory"]

memos/reranker/base.py

@@ -0,0 +1,25 @@
+# memos/reranker/base.py
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+
+if TYPE_CHECKING:
+    from memos.memories.textual.item import TextualMemoryItem
+
+
+class BaseReranker(ABC):
+    """Abstract interface for memory rerankers."""
+
+    @abstractmethod
+    def rerank(
+        self,
+        query: str,
+        graph_results: list[TextualMemoryItem],
+        top_k: int,
+        search_filter: dict | None = None,
+        **kwargs,
+    ) -> list[tuple[TextualMemoryItem, float]]:
+        """Return top_k (item, score) sorted by score desc."""
+        raise NotImplementedError

memos/reranker/concat.py

@@ -0,0 +1,103 @@
+import re
+
+from typing import Any
+
+from memos.memories.textual.item import SourceMessage
+
+
+_TAG1 = re.compile(r"^\s*\[[^\]]*\]\s*")
+
+
+def get_encoded_tokens(content: str) -> int:
+    """
+    Get encoded tokens.
+    Args:
+        content: str
+    Returns:
+        int: Encoded tokens.
+    """
+    return len(content)
+
+
+def truncate_data(data: list[str | dict[str, Any] | Any], max_tokens: int) -> list[str]:
+    """
+    Truncate data to max tokens.
+    Args:
+        data: List of strings or dictionaries.
+        max_tokens: Maximum number of tokens.
+    Returns:
+        str: Truncated string.
+    """
+    truncated_string = ""
+    for item in data:
+        if isinstance(item, SourceMessage):
+            content = getattr(item, "content", "")
+            chat_time = getattr(item, "chat_time", "")
+            if not content:
+                continue
+            truncated_string += f"[{chat_time}]: {content}\n"
+            if get_encoded_tokens(truncated_string) > max_tokens:
+                break
+    return truncated_string
+
+
+def process_source(
+    items: list[tuple[Any, str | dict[str, Any] | list[Any]]] | None = None,
+    recent_num: int = 10,
+    max_tokens: int = 2048,
+) -> str:
+    """
+    Args:
+        items: List of tuples where each tuple contains (memory, source).
+               source can be str, Dict, or List.
+        recent_num: Number of recent items to concatenate.
+    Returns:
+        str: Concatenated source.
+    """
+    if items is None:
+        items = []
+    concat_data = []
+    memory = None
+    for item in items:
+        memory, source = item
+        concat_data.extend(source[-recent_num:])
+    truncated_string = truncate_data(concat_data, max_tokens)
+    if memory is not None:
+        truncated_string = f"{memory}\n{truncated_string}"
+    return truncated_string
+
+
+def concat_original_source(
+    graph_results: list,
+    rerank_source: str | None = None,
+) -> list[str]:
+    """
+    Merge memory items with original dialogue.
+    Args:
+        graph_results (list[TextualMemoryItem]): List of memory items with embeddings.
+        merge_field (List[str]): List of fields to merge.
+    Returns:
+        list[str]: List of memory and concat orginal memory.
+    """
+    merge_field = []
+    merge_field = ["sources"] if rerank_source is None else rerank_source.split(",")
+    documents = []
+    for item in graph_results:
+        m = item.get("memory") if isinstance(item, dict) else getattr(item, "memory", None)
+
+        memory = _TAG1.sub("", m) if isinstance(m, str) else m
+
+        sources = []
+        for field in merge_field:
+            if isinstance(item, dict):
+                metadata = item.get("metadata", {})
+                source = metadata.get(field) if isinstance(metadata, dict) else None
+            else:
+                source = getattr(item.metadata, field, None) if hasattr(item, "metadata") else None
+
+            if source is None:
+                continue
+            sources.append((memory, source))
+        concat_string = process_source(sources)
+        documents.append(concat_string)
+    return documents

memos/reranker/cosine_local.py

@@ -0,0 +1,102 @@
+# memos/reranker/cosine_local.py
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from memos.log import get_logger
+from memos.utils import timed
+
+from .base import BaseReranker
+
+
+if TYPE_CHECKING:
+    from memos.memories.textual.item import TextualMemoryItem
+
+try:
+    import numpy as _np
+
+    _HAS_NUMPY = True
+except Exception:
+    _HAS_NUMPY = False
+
+logger = get_logger(__name__)
+
+
+def _cosine_one_to_many(q: list[float], m: list[list[float]]) -> list[float]:
+    """
+    Compute cosine similarities between a single vector q and a matrix m (rows are candidates).
+    """
+    if not _HAS_NUMPY:
+
+        def dot(a, b):  # lowercase per N806
+            return sum(x * y for x, y in zip(a, b, strict=False))
+
+        def norm(a):  # lowercase per N806
+            return sum(x * x for x in a) ** 0.5
+
+        qn = norm(q) or 1e-10
+        sims = []
+        for v in m:
+            vn = norm(v) or 1e-10
+            sims.append(dot(q, v) / (qn * vn))
+        return sims
+
+    qv = _np.asarray(q, dtype=float)  # lowercase
+    mv = _np.asarray(m, dtype=float)  # lowercase
+    qn = _np.linalg.norm(qv) or 1e-10
+    mn = _np.linalg.norm(mv, axis=1)  # lowercase
+    dots = mv @ qv
+    return (dots / (mn * qn + 1e-10)).tolist()
+
+
+class CosineLocalReranker(BaseReranker):
+    def __init__(
+        self,
+        level_weights: dict[str, float] | None = None,
+        level_field: str = "background",
+        **kwargs,
+    ):
+        self.level_weights = level_weights or {"topic": 1.0, "concept": 1.0, "fact": 1.0}
+        self.level_field = level_field
+
+    @timed
+    def rerank(
+        self,
+        query: str,
+        graph_results: list,
+        top_k: int,
+        **kwargs,
+    ) -> list[tuple[TextualMemoryItem, float]]:
+        if not graph_results:
+            return []
+
+        query_embedding: list[float] | None = kwargs.get("query_embedding")
+        if not query_embedding:
+            return [(item, 0.0) for item in graph_results[:top_k]]
+
+        items_with_emb = [
+            it
+            for it in graph_results
+            if getattr(it, "metadata", None) and getattr(it.metadata, "embedding", None)
+        ]
+        if not items_with_emb:
+            return [(item, 0.5) for item in graph_results[:top_k]]
+
+        cand_vecs = [it.metadata.embedding for it in items_with_emb]
+        sims = _cosine_one_to_many(query_embedding, cand_vecs)
+
+        def get_weight(it: TextualMemoryItem) -> float:
+            level = getattr(it.metadata, self.level_field, None)
+            return self.level_weights.get(level, 1.0)
+
+        weighted = [sim * get_weight(it) for sim, it in zip(sims, items_with_emb, strict=False)]
+        scored_pairs = list(zip(items_with_emb, weighted, strict=False))
+        scored_pairs.sort(key=lambda x: x[1], reverse=True)
+
+        top_items = scored_pairs[:top_k]
+        if len(top_items) < top_k:
+            chosen = {it.id for it, _ in top_items}
+            remain = [(it, -1.0) for it in graph_results if it.id not in chosen]
+            top_items.extend(remain[: top_k - len(top_items)])
+        logger.info(f"CosineLocalReranker rerank result: {top_items[:1]}")
+        return top_items

memos/reranker/factory.py

@@ -0,0 +1,72 @@
+# memos/reranker/factory.py
+from __future__ import annotations
+
+import json
+
+from typing import TYPE_CHECKING, Any
+
+# Import singleton decorator
+from memos.memos_tools.singleton import singleton_factory
+
+from .cosine_local import CosineLocalReranker
+from .http_bge import HTTPBGEReranker
+from .http_bge_strategy import HTTPBGERerankerStrategy
+from .noop import NoopReranker
+
+
+if TYPE_CHECKING:
+    from memos.configs.reranker import RerankerConfigFactory
+
+    from .base import BaseReranker
+
+
+class RerankerFactory:
+    @staticmethod
+    @singleton_factory("RerankerFactory")
+    def from_config(cfg: RerankerConfigFactory | None) -> BaseReranker | None:
+        if not cfg:
+            return None
+
+        backend = (cfg.backend or "").lower()
+        c: dict[str, Any] = cfg.config or {}
+
+        headers_extra = c.get("headers_extra")
+        if isinstance(headers_extra, str):
+            try:
+                headers_extra = json.loads(headers_extra)
+            except Exception:
+                headers_extra = None
+
+        if backend in {"http_bge", "bge"}:
+            return HTTPBGEReranker(
+                reranker_url=c.get("url") or c.get("endpoint") or c.get("reranker_url"),
+                model=c.get("model", "bge-reranker-v2-m3"),
+                timeout=int(c.get("timeout", 10)),
+                max_query_tokens=min(max(c.get("max_query_tokens", 8000), 100), 8000),
+                concate_len=min(max(c.get("concate_len", 1000), 4), 8000),
+                headers_extra=headers_extra,
+                rerank_source=c.get("rerank_source"),
+            )
+
+        if backend in {"cosine_local", "cosine"}:
+            return CosineLocalReranker(
+                level_weights=c.get("level_weights"),
+                level_field=c.get("level_field", "background"),
+            )
+
+        if backend in {"noop", "none", "disabled"}:
+            return NoopReranker()
+
+        if backend in {"http_bge_strategy", "bge_strategy"}:
+            return HTTPBGERerankerStrategy(
+                reranker_url=c.get("url") or c.get("endpoint") or c.get("reranker_url"),
+                model=c.get("model", "bge-reranker-v2-m3"),
+                timeout=int(c.get("timeout", 10)),
+                max_query_tokens=min(max(c.get("max_query_tokens", 8000), 100), 8000),
+                concate_len=min(max(c.get("concate_len", 1000), 4), 8000),
+                headers_extra=headers_extra,
+                rerank_source=c.get("rerank_source"),
+                reranker_strategy=c.get("reranker_strategy"),
+            )
+
+        raise ValueError(f"Unknown reranker backend: {cfg.backend}")