MemoryOS 2.0.3__py3-none-any.whl

memos/memories/textual/prefer_text_memory/config.py

@@ -0,0 +1,106 @@
+from typing import Any, ClassVar
+
+from pydantic import Field, field_validator, model_validator
+
+from memos.configs.base import BaseConfig
+
+
+class BaseAdderConfig(BaseConfig):
+    """Base configuration class for Adder."""
+
+
+class NaiveAdderConfig(BaseAdderConfig):
+    """Configuration for Naive Adder."""
+
+    # No additional config needed since components are passed from parent
+
+
+class AdderConfigFactory(BaseConfig):
+    """Factory class for creating Adder configurations."""
+
+    backend: str = Field(..., description="Backend for Adder")
+    config: dict[str, Any] = Field(..., description="Configuration for the Adder backend")
+
+    backend_to_class: ClassVar[dict[str, Any]] = {
+        "naive": NaiveAdderConfig,
+    }
+
+    @field_validator("backend")
+    @classmethod
+    def validate_backend(cls, backend: str) -> str:
+        """Validate the backend field."""
+        if backend not in cls.backend_to_class:
+            raise ValueError(f"Invalid backend: {backend}")
+        return backend
+
+    @model_validator(mode="after")
+    def create_config(self) -> "AdderConfigFactory":
+        config_class = self.backend_to_class[self.backend]
+        self.config = config_class(**self.config)
+        return self
+
+
+class BaseExtractorConfig(BaseConfig):
+    """Base configuration class for Extractor."""
+
+
+class NaiveExtractorConfig(BaseExtractorConfig):
+    """Configuration for Naive Extractor."""
+
+
+class ExtractorConfigFactory(BaseConfig):
+    """Factory class for creating Extractor configurations."""
+
+    backend: str = Field(..., description="Backend for Extractor")
+    config: dict[str, Any] = Field(..., description="Configuration for the Extractor backend")
+
+    backend_to_class: ClassVar[dict[str, Any]] = {
+        "naive": NaiveExtractorConfig,
+    }
+
+    @field_validator("backend")
+    @classmethod
+    def validate_backend(cls, backend: str) -> str:
+        """Validate the backend field."""
+        if backend not in cls.backend_to_class:
+            raise ValueError(f"Invalid backend: {backend}")
+        return backend
+
+    @model_validator(mode="after")
+    def create_config(self) -> "ExtractorConfigFactory":
+        config_class = self.backend_to_class[self.backend]
+        self.config = config_class(**self.config)
+        return self
+
+
+class BaseRetrieverConfig(BaseConfig):
+    """Base configuration class for Retrievers."""
+
+
+class NaiveRetrieverConfig(BaseRetrieverConfig):
+    """Configuration for Naive Retriever."""
+
+
+class RetrieverConfigFactory(BaseConfig):
+    """Factory class for creating Retriever configurations."""
+
+    backend: str = Field(..., description="Backend for Retriever")
+    config: dict[str, Any] = Field(..., description="Configuration for the Retriever backend")
+
+    backend_to_class: ClassVar[dict[str, Any]] = {
+        "naive": NaiveRetrieverConfig,
+    }
+
+    @field_validator("backend")
+    @classmethod
+    def validate_backend(cls, backend: str) -> str:
+        """Validate the backend field."""
+        if backend not in cls.backend_to_class:
+            raise ValueError(f"Invalid backend: {backend}")
+        return backend
+
+    @model_validator(mode="after")
+    def create_config(self) -> "RetrieverConfigFactory":
+        config_class = self.backend_to_class[self.backend]
+        self.config = config_class(**self.config)
+        return self

memos/memories/textual/prefer_text_memory/extractor.py

@@ -0,0 +1,221 @@
+import json
+import uuid
+
+from abc import ABC, abstractmethod
+from concurrent.futures import as_completed
+from datetime import datetime
+from typing import Any
+
+from memos.context.context import ContextThreadPoolExecutor
+from memos.log import get_logger
+from memos.mem_reader.read_multi_modal import detect_lang
+from memos.memories.textual.item import (
+    PreferenceTextualMemoryMetadata,
+    TextualMemoryItem,
+    list_all_fields,
+)
+from memos.memories.textual.prefer_text_memory.spliter import Splitter
+from memos.memories.textual.prefer_text_memory.utils import convert_messages_to_string
+from memos.templates.prefer_complete_prompt import (
+    NAIVE_EXPLICIT_PREFERENCE_EXTRACT_PROMPT,
+    NAIVE_EXPLICIT_PREFERENCE_EXTRACT_PROMPT_ZH,
+    NAIVE_IMPLICIT_PREFERENCE_EXTRACT_PROMPT,
+    NAIVE_IMPLICIT_PREFERENCE_EXTRACT_PROMPT_ZH,
+)
+from memos.types import MessageList
+
+
+logger = get_logger(__name__)
+
+
+class BaseExtractor(ABC):
+    """Abstract base class for extractors."""
+
+    @abstractmethod
+    def __init__(self, llm_provider=None, embedder=None, vector_db=None):
+        """Initialize the extractor."""
+
+
+class NaiveExtractor(BaseExtractor):
+    """Extractor."""
+
+    def __init__(self, llm_provider=None, embedder=None, vector_db=None):
+        """Initialize the extractor."""
+        super().__init__(llm_provider, embedder, vector_db)
+        self.llm_provider = llm_provider
+        self.embedder = embedder
+        self.vector_db = vector_db
+        self.splitter = Splitter()
+
+    def extract_basic_info(self, qa_pair: MessageList) -> dict[str, Any]:
+        """Extract basic information from a QA pair (no LLM needed)."""
+        basic_info = {
+            "dialog_id": str(uuid.uuid4()),
+            "original_text": convert_messages_to_string(qa_pair),
+            "created_at": datetime.now().isoformat(),
+        }
+
+        return basic_info
+
+    def extract_explicit_preference(self, qa_pair: MessageList | str) -> dict[str, Any] | None:
+        """Extract explicit preference from a QA pair."""
+        qa_pair_str = convert_messages_to_string(qa_pair) if isinstance(qa_pair, list) else qa_pair
+        lang = detect_lang(qa_pair_str)
+        _map = {
+            "zh": NAIVE_EXPLICIT_PREFERENCE_EXTRACT_PROMPT_ZH,
+            "en": NAIVE_EXPLICIT_PREFERENCE_EXTRACT_PROMPT,
+        }
+        prompt = _map[lang].replace("{qa_pair}", qa_pair_str)
+
+        try:
+            response = self.llm_provider.generate([{"role": "user", "content": prompt}])
+            if not response:
+                logger.info(
+                    f"[prefer_extractor]: (Error) LLM response content is {response} when extracting explicit preference"
+                )
+                return None
+            response = response.strip().replace("```json", "").replace("```", "").strip()
+            result = json.loads(response)
+            for d in result:
+                d["preference"] = d.pop("explicit_preference")
+            return result
+        except Exception as e:
+            logger.info(f"Error extracting explicit preference: {e}, return None")
+            return None
+
+    def extract_implicit_preference(self, qa_pair: MessageList | str) -> dict[str, Any] | None:
+        """Extract implicit preferences from cluster qa pairs."""
+        if not qa_pair:
+            return None
+        qa_pair_str = convert_messages_to_string(qa_pair) if isinstance(qa_pair, list) else qa_pair
+        lang = detect_lang(qa_pair_str)
+        _map = {
+            "zh": NAIVE_IMPLICIT_PREFERENCE_EXTRACT_PROMPT_ZH,
+            "en": NAIVE_IMPLICIT_PREFERENCE_EXTRACT_PROMPT,
+        }
+        prompt = _map[lang].replace("{qa_pair}", qa_pair_str)
+
+        try:
+            response = self.llm_provider.generate([{"role": "user", "content": prompt}])
+            if not response:
+                logger.info(
+                    f"[prefer_extractor]: (Error) LLM response content is {response} when extracting implicit preference"
+                )
+                return None
+            response = response.strip().replace("```json", "").replace("```", "").strip()
+            result = json.loads(response)
+            for d in result:
+                d["preference"] = d.pop("implicit_preference")
+            return result
+        except Exception as e:
+            logger.info(f"Error extracting implicit preferences: {e}, return None")
+            return None
+
+    def _process_single_chunk_explicit(
+        self, chunk: MessageList, msg_type: str, info: dict[str, Any]
+    ) -> TextualMemoryItem | None:
+        """Process a single chunk and return a TextualMemoryItem."""
+        basic_info = self.extract_basic_info(chunk)
+        if not basic_info["original_text"]:
+            return None
+
+        explicit_pref = self.extract_explicit_preference(basic_info["original_text"])
+        if not explicit_pref:
+            return None
+
+        memories = []
+        for pref in explicit_pref:
+            vector_info = {
+                "embedding": self.embedder.embed([pref["context_summary"]])[0],
+            }
+            user_info = {k: v for k, v in info.items() if k not in list_all_fields()}
+            extract_info = {**basic_info, **pref, **vector_info, **info, "info": user_info}
+
+            metadata = PreferenceTextualMemoryMetadata(
+                type=msg_type, preference_type="explicit_preference", **extract_info
+            )
+            memory = TextualMemoryItem(
+                id=str(uuid.uuid4()), memory=pref["context_summary"], metadata=metadata
+            )
+
+            memories.append(memory)
+
+        return memories
+
+    def _process_single_chunk_implicit(
+        self, chunk: MessageList, msg_type: str, info: dict[str, Any]
+    ) -> TextualMemoryItem | None:
+        basic_info = self.extract_basic_info(chunk)
+        if not basic_info["original_text"]:
+            return None
+        implicit_pref = self.extract_implicit_preference(basic_info["original_text"])
+        if not implicit_pref:
+            return None
+
+        memories = []
+        for pref in implicit_pref:
+            vector_info = {
+                "embedding": self.embedder.embed([pref["context_summary"]])[0],
+            }
+            user_info = {k: v for k, v in info.items() if k not in list_all_fields()}
+            extract_info = {**basic_info, **pref, **vector_info, **info, "info": user_info}
+
+            metadata = PreferenceTextualMemoryMetadata(
+                type=msg_type, preference_type="implicit_preference", **extract_info
+            )
+            memory = TextualMemoryItem(
+                id=str(uuid.uuid4()), memory=pref["context_summary"], metadata=metadata
+            )
+
+            memories.append(memory)
+
+        return memories
+
+    def extract(
+        self,
+        messages: list[MessageList],
+        msg_type: str,
+        info: dict[str, Any],
+        max_workers: int = 10,
+    ) -> list[TextualMemoryItem]:
+        """Extract preference memories based on the messages using thread pool for acceleration."""
+        chunks: list[MessageList] = []
+        for message in messages:
+            chunk = self.splitter.split_chunks(message, split_type="overlap")
+            chunks.extend(chunk)
+        if not chunks:
+            return []
+
+        memories = []
+        with ContextThreadPoolExecutor(max_workers=min(max_workers, len(chunks))) as executor:
+            futures = {
+                executor.submit(self._process_single_chunk_explicit, chunk, msg_type, info): (
+                    "explicit",
+                    chunk,
+                )
+                for chunk in chunks
+            }
+            futures.update(
+                {
+                    executor.submit(self._process_single_chunk_implicit, chunk, msg_type, info): (
+                        "implicit",
+                        chunk,
+                    )
+                    for chunk in chunks
+                }
+            )
+
+            for future in as_completed(futures):
+                try:
+                    memory = future.result()
+                    if memory:
+                        if isinstance(memory, list):
+                            memories.extend(memory)
+                        else:
+                            memories.append(memory)
+                except Exception as e:
+                    task_type, chunk = futures[future]
+                    logger.error(f"Error processing {task_type} chunk: {chunk}\n{e}")
+                    continue
+
+        return memories

memos/memories/textual/prefer_text_memory/factory.py

@@ -0,0 +1,85 @@
+from typing import Any, ClassVar
+
+from memos.memories.textual.prefer_text_memory.adder import BaseAdder, NaiveAdder
+from memos.memories.textual.prefer_text_memory.config import (
+    AdderConfigFactory,
+    ExtractorConfigFactory,
+    RetrieverConfigFactory,
+)
+from memos.memories.textual.prefer_text_memory.extractor import BaseExtractor, NaiveExtractor
+from memos.memories.textual.prefer_text_memory.retrievers import BaseRetriever, NaiveRetriever
+
+
+class AdderFactory(BaseAdder):
+    """Factory class for creating Adder instances."""
+
+    backend_to_class: ClassVar[dict[str, Any]] = {
+        "naive": NaiveAdder,
+    }
+
+    @classmethod
+    def from_config(
+        cls,
+        config_factory: AdderConfigFactory,
+        llm_provider=None,
+        embedder=None,
+        vector_db=None,
+        text_mem=None,
+    ) -> BaseAdder:
+        """Create a Adder instance from a configuration factory."""
+        backend = config_factory.backend
+        if backend not in cls.backend_to_class:
+            raise ValueError(f"Invalid backend: {backend}")
+        adder_class = cls.backend_to_class[backend]
+        return adder_class(
+            llm_provider=llm_provider, embedder=embedder, vector_db=vector_db, text_mem=text_mem
+        )
+
+
+class ExtractorFactory(BaseExtractor):
+    """Factory class for creating Extractor instances."""
+
+    backend_to_class: ClassVar[dict[str, Any]] = {
+        "naive": NaiveExtractor,
+    }
+
+    @classmethod
+    def from_config(
+        cls,
+        config_factory: ExtractorConfigFactory,
+        llm_provider=None,
+        embedder=None,
+        vector_db=None,
+    ) -> BaseExtractor:
+        """Create a Extractor instance from a configuration factory."""
+        backend = config_factory.backend
+        if backend not in cls.backend_to_class:
+            raise ValueError(f"Invalid backend: {backend}")
+        extractor_class = cls.backend_to_class[backend]
+        return extractor_class(llm_provider=llm_provider, embedder=embedder, vector_db=vector_db)
+
+
+class RetrieverFactory(BaseRetriever):
+    """Factory class for creating Retriever instances."""
+
+    backend_to_class: ClassVar[dict[str, Any]] = {
+        "naive": NaiveRetriever,
+    }
+
+    @classmethod
+    def from_config(
+        cls,
+        config_factory: RetrieverConfigFactory,
+        llm_provider=None,
+        embedder=None,
+        reranker=None,
+        vector_db=None,
+    ) -> BaseRetriever:
+        """Create a Retriever instance from a configuration factory."""
+        backend = config_factory.backend
+        if backend not in cls.backend_to_class:
+            raise ValueError(f"Invalid backend: {backend}")
+        retriever_class = cls.backend_to_class[backend]
+        return retriever_class(
+            llm_provider=llm_provider, embedder=embedder, reranker=reranker, vector_db=vector_db
+        )

memos/memories/textual/prefer_text_memory/retrievers.py

@@ -0,0 +1,177 @@
+import os
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+from memos.context.context import ContextThreadPoolExecutor
+from memos.memories.textual.item import PreferenceTextualMemoryMetadata, TextualMemoryItem
+from memos.vec_dbs.item import MilvusVecDBItem
+
+
+class BaseRetriever(ABC):
+    """Abstract base class for retrievers."""
+
+    @abstractmethod
+    def __init__(self, llm_provider=None, embedder=None, reranker=None, vector_db=None):
+        """Initialize the retriever."""
+
+    @abstractmethod
+    def retrieve(
+        self,
+        query: str,
+        top_k: int,
+        info: dict[str, Any] | None = None,
+        search_filter: dict[str, Any] | None = None,
+    ) -> list[TextualMemoryItem]:
+        """Retrieve memories from the retriever."""
+
+
+class NaiveRetriever(BaseRetriever):
+    """Naive retriever."""
+
+    def __init__(self, llm_provider=None, embedder=None, reranker=None, vector_db=None):
+        """Initialize the naive retriever."""
+        super().__init__(llm_provider, embedder, reranker, vector_db)
+        self.reranker = reranker
+        self.vector_db = vector_db
+        self.embedder = embedder
+
+    def _naive_reranker(
+        self, query: str, prefs_mem: list[TextualMemoryItem], top_k: int, **kwargs: Any
+    ) -> list[TextualMemoryItem]:
+        if self.reranker:
+            prefs_mem_reranked = []
+            prefs_mem_tuple = self.reranker.rerank(query, prefs_mem, top_k)
+            for item, score in prefs_mem_tuple:
+                item.metadata.score = score
+                prefs_mem_reranked.append(item)
+        return prefs_mem_reranked
+
+    def _original_text_reranker(
+        self,
+        query: str,
+        prefs_mem: list[TextualMemoryItem],
+        prefs: list[MilvusVecDBItem],
+        top_k: int,
+        **kwargs: Any,
+    ) -> list[TextualMemoryItem]:
+        if self.reranker:
+            from copy import deepcopy
+
+            prefs_mem_for_reranker = deepcopy(prefs_mem)
+            for pref_mem, pref in zip(prefs_mem_for_reranker, prefs, strict=False):
+                pref_mem.memory = pref_mem.memory + "\n" + pref.original_text
+            reranked_results = self.reranker.rerank(query, prefs_mem_for_reranker, top_k)
+            prefs_mem_for_reranker = [item for item, _ in reranked_results]
+            prefs_ids = [item.id for item in prefs_mem_for_reranker]
+            prefs_dict = {item.id: item for item in prefs_mem}
+
+            # Create mapping from id to score from reranked results
+            reranked_scores = {item.id: score for item, score in reranked_results}
+
+            # Assign scores to the original items
+            result_items = []
+            for item_id in prefs_ids:
+                if item_id in prefs_dict:
+                    original_item = prefs_dict[item_id]
+                    original_item.metadata.score = reranked_scores.get(item_id)
+                    result_items.append(original_item)
+            return result_items
+        return prefs_mem
+
+    def retrieve(
+        self,
+        query: str,
+        top_k: int,
+        info: dict[str, Any] | None = None,
+        search_filter: dict[str, Any] | None = None,
+    ) -> list[TextualMemoryItem]:
+        """Retrieve memories from the naive retriever."""
+        # TODO: un-support rewrite query and session filter now
+        if info:
+            info = info.copy()  # Create a copy to avoid modifying the original
+            info.pop("chat_history", None)
+            info.pop("session_id", None)
+        search_filter = {"and": [info, search_filter]}
+        query_embeddings = self.embedder.embed([query])  # Pass as list to get list of embeddings
+        query_embedding = query_embeddings[0]  # Get the first (and only) embedding
+
+        # Use thread pool to parallelize the searches
+        with ContextThreadPoolExecutor(max_workers=2) as executor:
+            # Submit all search tasks
+            future_explicit = executor.submit(
+                self.vector_db.search,
+                query_embedding,
+                query,
+                "explicit_preference",
+                top_k * 2,
+                search_filter,
+            )
+            future_implicit = executor.submit(
+                self.vector_db.search,
+                query_embedding,
+                query,
+                "implicit_preference",
+                top_k * 2,
+                search_filter,
+            )
+
+            # Wait for all results
+            explicit_prefs = future_explicit.result()
+            implicit_prefs = future_implicit.result()
+
+        # sort by score
+        explicit_prefs.sort(key=lambda x: x.score, reverse=True)
+        implicit_prefs.sort(key=lambda x: x.score, reverse=True)
+
+        explicit_prefs_mem = [
+            TextualMemoryItem(
+                id=pref.id,
+                memory=pref.memory,
+                metadata=PreferenceTextualMemoryMetadata(**pref.payload),
+            )
+            for pref in explicit_prefs
+            if pref.payload.get("preference", None)
+        ]
+
+        implicit_prefs_mem = [
+            TextualMemoryItem(
+                id=pref.id,
+                memory=pref.memory,
+                metadata=PreferenceTextualMemoryMetadata(**pref.payload),
+            )
+            for pref in implicit_prefs
+            if pref.payload.get("preference", None)
+        ]
+
+        reranker_map = {
+            "naive": self._naive_reranker,
+            "original_text": self._original_text_reranker,
+        }
+        reranker_func = reranker_map["naive"]
+        prefs_mem_explicit = reranker_func(
+            query=query,
+            prefs_mem=explicit_prefs_mem,
+            prefs=explicit_prefs,
+            top_k=top_k,
+        )
+        prefs_mem_implicit = reranker_func(
+            query=query,
+            prefs_mem=implicit_prefs_mem,
+            prefs=implicit_prefs,
+            top_k=top_k,
+        )
+
+        # filter explicit mem by score bigger than threshold
+        prefs_mem_explicit = [
+            item
+            for item in prefs_mem_explicit
+            if item.metadata.score >= float(os.getenv("PREFERENCE_SEARCH_THRESHOLD", 0.0))
+        ]
+        prefs_mem_implicit = [
+            item
+            for item in prefs_mem_implicit
+            if item.metadata.score >= float(os.getenv("PREFERENCE_SEARCH_THRESHOLD", 0.0))
+        ]
+
+        return prefs_mem_explicit + prefs_mem_implicit