MemoryOS 0.0.1__py3-none-any.whl → 0.1.12__py3-none-any.whl

memos/memories/textual/tree.py

@@ -0,0 +1,289 @@
+import json
+import os
+import shutil
+import tempfile
+
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+from memos.configs.memory import TreeTextMemoryConfig
+from memos.embedders.factory import EmbedderFactory, OllamaEmbedder
+from memos.graph_dbs.factory import GraphStoreFactory, Neo4jGraphDB
+from memos.llms.factory import LLMFactory, OllamaLLM, OpenAILLM
+from memos.log import get_logger
+from memos.memories.textual.base import BaseTextMemory
+from memos.memories.textual.item import TextualMemoryItem, TreeNodeTextualMemoryMetadata
+from memos.memories.textual.tree_text_memory.organize.manager import MemoryManager
+from memos.memories.textual.tree_text_memory.retrieve.searcher import Searcher
+from memos.types import MessageList
+
+
+logger = get_logger(__name__)
+
+
+class TreeTextMemory(BaseTextMemory):
+    """General textual memory implementation for storing and retrieving memories."""
+
+    def __init__(self, config: TreeTextMemoryConfig):
+        """Initialize memory with the given configuration."""
+        self.config: TreeTextMemoryConfig = config
+        self.extractor_llm: OpenAILLM | OllamaLLM = LLMFactory.from_config(config.extractor_llm)
+        self.dispatcher_llm: OpenAILLM | OllamaLLM = LLMFactory.from_config(config.dispatcher_llm)
+        self.embedder: OllamaEmbedder = EmbedderFactory.from_config(config.embedder)
+        self.graph_store: Neo4jGraphDB = GraphStoreFactory.from_config(config.graph_db)
+        self.memory_manager: MemoryManager = MemoryManager(self.graph_store, self.embedder)
+
+    def add(self, memories: list[TextualMemoryItem | dict[str, Any]]) -> None:
+        """Add memories.
+        Args:
+            memories: List of TextualMemoryItem objects or dictionaries to add.
+        Later:
+            memory_items = [TextualMemoryItem(**m) if isinstance(m, dict) else m for m in memories]
+            metadata = extract_metadata(memory_items, self.extractor_llm)
+            plan = plan_memory_operations(memory_items, metadata, self.graph_store)
+            execute_plan(memory_items, metadata, plan, self.graph_store)
+        """
+        self.memory_manager.add(memories)
+
+    def replace_working_memory(self, memories: list[TextualMemoryItem]) -> None:
+        self.memory_manager.replace_working_memory(memories)
+
+    def get_working_memory(self) -> list[TextualMemoryItem]:
+        working_memories = self.graph_store.get_all_memory_items(scope="WorkingMemory")
+        items = [TextualMemoryItem.from_dict(record) for record in (working_memories)]
+        # Sort by updated_at in descending order
+        sorted_items = sorted(
+            items, key=lambda x: x.metadata.updated_at or datetime.min, reverse=True
+        )
+        return sorted_items
+
+    def get_current_memory_size(self) -> dict[str, int]:
+        """
+        Get the current size of each memory type.
+        This delegates to the MemoryManager.
+        """
+        return self.memory_manager.get_current_memory_size()
+
+    def search(
+        self, query: str, top_k: int, info=None, mode: str = "fast", memory_type: str = "All"
+    ) -> list[TextualMemoryItem]:
+        """Search for memories based on a query.
+        User query -> TaskGoalParser -> MemoryPathResolver ->
+        GraphMemoryRetriever -> MemoryReranker -> MemoryReasoner -> Final output
+        Args:
+            query (str): The query to search for.
+            top_k (int): The number of top results to return.
+            info (dict): Leave a record of memory consumption.
+            mode (str, optional): The mode of the search.
+            - 'fast': Uses a faster search process, sacrificing some precision for speed.
+            - 'fine': Uses a more detailed search process, invoking large models for higher precision, but slower performance.
+            memory_type (str): Type restriction for search.
+            ['All', 'WorkingMemory', 'LongTermMemory', 'UserMemory']
+        Returns:
+            list[TextualMemoryItem]: List of matching memories.
+        """
+        searcher = Searcher(self.dispatcher_llm, self.graph_store, self.embedder)
+        return searcher.search(query, top_k, info, mode, memory_type)
+
+    def get_relevant_subgraph(
+        self, query: str, top_k: int = 5, depth: int = 2, center_status: str = "activated"
+    ) -> dict[str, Any]:
+        """
+        Find and merge the local neighborhood sub-graphs of the top-k
+        nodes most relevant to the query.
+         Process:
+             1. Embed the user query into a vector representation.
+             2. Use vector similarity search to find the top-k similar nodes.
+             3. For each similar node:
+                 - Ensure its status matches `center_status` (e.g., 'active').
+                 - Retrieve its local subgraph up to `depth` hops.
+                 - Collect the center node, its neighbors, and connecting edges.
+             4. Merge all retrieved subgraphs into a single unified subgraph.
+             5. Return the merged subgraph structure.
+
+         Args:
+             query (str): The user input or concept to find relevant memories for.
+             top_k (int, optional): How many top similar nodes to retrieve. Default is 5.
+             depth (int, optional): The neighborhood depth (number of hops). Default is 2.
+             center_status (str, optional): Status condition the center node must satisfy (e.g., 'active').
+
+         Returns:
+             dict[str, Any]: A subgraph dict with:
+                 - 'core_id': ID of the top matching core node, or None if none found.
+                 - 'nodes': List of unique nodes (core + neighbors) in the merged subgraph.
+                 - 'edges': List of unique edges (as dicts with 'from', 'to', 'type') in the merged subgraph.
+        """
+        # Step 1: Embed query
+        query_embedding = self.embedder.embed([query])[0]
+
+        # Step 2: Get top-1 similar node
+        similar_nodes = self.graph_store.search_by_embedding(query_embedding, top_k=top_k)
+        if not similar_nodes:
+            logger.info("No similar nodes found for query embedding.")
+            return {"core_id": None, "nodes": [], "edges": []}
+
+        # Step 3: Fetch neighborhood
+        all_nodes = {}
+        all_edges = set()
+        cores = []
+
+        for node in similar_nodes:
+            core_id = node["id"]
+            score = node["score"]
+
+            subgraph = self.graph_store.get_subgraph(
+                center_id=core_id, depth=depth, center_status=center_status
+            )
+
+            if not subgraph["core_node"]:
+                logger.info(f"Skipping node {core_id} (inactive or not found).")
+                continue
+
+            core_node = subgraph["core_node"]
+            neighbors = subgraph["neighbors"]
+            edges = subgraph["edges"]
+
+            # Collect nodes
+            all_nodes[core_node["id"]] = core_node
+            for n in neighbors:
+                all_nodes[n["id"]] = n
+
+            # Collect edges
+            for e in edges:
+                all_edges.add((e["source"], e["target"], e["type"]))
+
+            cores.append(
+                {"id": core_id, "score": score, "core_node": core_node, "neighbors": neighbors}
+            )
+
+        top_core = cores[0]
+        return {
+            "core_id": top_core["id"],
+            "nodes": list(all_nodes.values()),
+            "edges": [{"source": f, "target": t, "type": ty} for (f, t, ty) in all_edges],
+        }
+
+    def extract(self, messages: MessageList) -> list[TextualMemoryItem]:
+        raise NotImplementedError
+
+    def update(self, memory_id: str, new_memory: TextualMemoryItem | dict[str, Any]) -> None:
+        raise NotImplementedError
+
+    def get(self, memory_id: str) -> TextualMemoryItem:
+        """Get a memory by its ID."""
+        result = self.graph_store.get_node(memory_id)
+        if result is None:
+            raise ValueError(f"Memory with ID {memory_id} not found")
+        metadata_dict = result.get("metadata", {})
+        return TextualMemoryItem(
+            id=result["id"],
+            memory=result["memory"],
+            metadata=TreeNodeTextualMemoryMetadata(**metadata_dict),
+        )
+
+    def get_by_ids(self, memory_ids: list[str]) -> list[TextualMemoryItem]:
+        raise NotImplementedError
+
+    def get_all(self) -> dict:
+        """Get all memories.
+        Returns:
+            list[TextualMemoryItem]: List of all memories.
+        """
+        all_items = self.graph_store.export_graph()
+        return all_items
+
+    def delete(self, memory_ids: list[str]) -> None:
+        raise NotImplementedError
+
+    def delete_all(self) -> None:
+        """Delete all memories and their relationships from the graph store."""
+        try:
+            self.graph_store.clear()
+            logger.info("All memories and edges have been deleted from the graph.")
+        except Exception as e:
+            logger.error(f"An error occurred while deleting all memories: {e}")
+            raise
+
+    def load(self, dir: str) -> None:
+        try:
+            memory_file = os.path.join(dir, self.config.memory_filename)
+
+            if not os.path.exists(memory_file):
+                logger.warning(f"Memory file not found: {memory_file}")
+                return
+
+            with open(memory_file, encoding="utf-8") as f:
+                memories = json.load(f)
+
+            self.graph_store.import_graph(memories)
+            logger.info(f"Loaded {len(memories)} memories from {memory_file}")
+
+        except FileNotFoundError:
+            logger.error(f"Memory file not found in directory: {dir}")
+        except json.JSONDecodeError as e:
+            logger.error(f"Error decoding JSON from memory file: {e}")
+        except Exception as e:
+            logger.error(f"An error occurred while loading memories: {e}")
+
+    def dump(self, dir: str) -> None:
+        """Dump memories to os.path.join(dir, self.config.memory_filename)"""
+        try:
+            json_memories = self.graph_store.export_graph()
+
+            os.makedirs(dir, exist_ok=True)
+            memory_file = os.path.join(dir, self.config.memory_filename)
+            with open(memory_file, "w", encoding="utf-8") as f:
+                json.dump(json_memories, f, indent=4, ensure_ascii=False)
+
+            logger.info(f"Dumped {len(json_memories.get('nodes'))} memories to {memory_file}")
+
+        except Exception as e:
+            logger.error(f"An error occurred while dumping memories: {e}")
+            raise
+
+    def drop(self, keep_last_n: int = 30) -> None:
+        """
+        Export all memory data to a versioned backup dir and drop the Neo4j database.
+        Only the latest `keep_last_n` backups will be retained.
+        """
+        try:
+            backup_root = Path(tempfile.gettempdir()) / "memos_backups"
+            backup_root.mkdir(parents=True, exist_ok=True)
+
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            backup_dir = backup_root / f"memos_backup_{timestamp}"
+            backup_dir.mkdir()
+
+            logger.info(f"Exporting memory to backup dir: {backup_dir}")
+            self.dump(str(backup_dir))
+
+            # Clean up old backups
+            self._cleanup_old_backups(backup_root, keep_last_n)
+
+            self.graph_store.drop_database()
+            logger.info(f"Database '{self.graph_store.db_name}' dropped after backup.")
+
+        except Exception as e:
+            logger.error(f"Error in drop(): {e}")
+            raise
+
+    @staticmethod
+    def _cleanup_old_backups(root_dir: Path, keep_last_n: int) -> None:
+        """
+        Keep only the latest `keep_last_n` backup directories under `root_dir`.
+        Older ones will be deleted.
+        """
+        backups = sorted(
+            [d for d in root_dir.iterdir() if d.is_dir() and d.name.startswith("memos_backup_")],
+            key=lambda p: p.name,  # name includes timestamp
+            reverse=True,
+        )
+
+        to_delete = backups[keep_last_n:]
+        for old_dir in to_delete:
+            try:
+                shutil.rmtree(old_dir)
+                logger.info(f"Deleted old backup directory: {old_dir}")
+            except Exception as e:
+                logger.warning(f"Failed to delete backup {old_dir}: {e}")

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

@@ -0,0 +1,305 @@
+import uuid
+
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from datetime import datetime
+
+from memos.embedders.factory import OllamaEmbedder
+from memos.graph_dbs.neo4j import Neo4jGraphDB
+from memos.log import get_logger
+from memos.memories.textual.item import TextualMemoryItem, TreeNodeTextualMemoryMetadata
+
+
+logger = get_logger(__name__)
+
+
+class MemoryManager:
+    def __init__(
+        self,
+        graph_store: Neo4jGraphDB,
+        embedder: OllamaEmbedder,
+        memory_size: dict | None = None,
+        threshold: float | None = 0.80,
+        merged_threshold: float | None = 0.92,
+    ):
+        self.graph_store = graph_store
+        self.embedder = embedder
+        self.memory_size = memory_size
+        self.current_memory_size = {
+            "WorkingMemory": 0,
+            "LongTermMemory": 0,
+            "UserMemory": 0,
+        }
+        if not memory_size:
+            self.memory_size = {
+                "WorkingMemory": 20,
+                "LongTermMemory": 10000,
+                "UserMemory": 10000,
+            }
+        self._threshold = threshold
+        self._merged_threshold = merged_threshold
+
+    def add(self, memories: list[TextualMemoryItem]) -> None:
+        """
+        Add new memories in parallel to different memory types (WorkingMemory, LongTermMemory, UserMemory).
+        """
+        with ThreadPoolExecutor(max_workers=8) as executor:
+            futures = [executor.submit(self._process_memory, memory) for memory in memories]
+            for future in as_completed(futures):
+                try:
+                    future.result()
+                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"]
+        )
+
+        self._refresh_memory_size()
+
+    def replace_working_memory(self, memories: list[TextualMemoryItem]) -> None:
+        """
+        Replace WorkingMemory
+        """
+        working_memory_top_k = memories[: self.memory_size["WorkingMemory"]]
+        with ThreadPoolExecutor(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):
+                try:
+                    future.result()
+                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._refresh_memory_size()
+
+    def get_current_memory_size(self) -> dict[str, int]:
+        """
+        Return the cached memory type counts.
+        """
+        return self.current_memory_size
+
+    def _refresh_memory_size(self) -> None:
+        """
+        Query the latest counts from the graph store and update internal state.
+        """
+        results = self.graph_store.get_grouped_counts(group_fields=["memory_type"])
+        self.current_memory_size = {record["memory_type"]: record["count"] for record in results}
+        logger.info(f"[MemoryManager] Refreshed memory sizes: {self.current_memory_size}")
+
+    def _process_memory(self, memory: TextualMemoryItem):
+        """
+        Process and add memory to different memory types (WorkingMemory, LongTermMemory, UserMemory).
+        This method runs asynchronously to process each memory item.
+        """
+        # Add to WorkingMemory
+        self._add_memory_to_db(memory, "WorkingMemory")
+
+        # Add to LongTermMemory and UserMemory
+        if memory.metadata.memory_type in ["LongTermMemory", "UserMemory"]:
+            self._add_to_graph_memory(
+                memory=memory,
+                memory_type=memory.metadata.memory_type,
+            )
+
+    def _add_memory_to_db(self, memory: TextualMemoryItem, memory_type: str):
+        """
+        Add a single memory item to the graph store, with FIFO logic for WorkingMemory.
+        """
+        metadata = memory.metadata.model_copy(update={"memory_type": memory_type}).model_dump(
+            exclude_none=True
+        )
+        metadata["updated_at"] = datetime.now().isoformat()
+        working_memory = TextualMemoryItem(memory=memory.memory, metadata=metadata)
+
+        # Insert node into graph
+        self.graph_store.add_node(working_memory.id, working_memory.memory, metadata)
+
+    def _add_to_graph_memory(self, memory: TextualMemoryItem, memory_type: str):
+        """
+        Generalized method to add memory to a graph-based memory type (e.g., LongTermMemory, UserMemory).
+
+        Parameters:
+        - memory: memory item to insert
+        - memory_type: "LongTermMemory" | "UserMemory"
+        - similarity_threshold: deduplication threshold
+        - topic_summary_prefix: summary node id prefix if applicable
+        - enable_summary_link: whether to auto-link to a summary node
+        """
+        embedding = memory.metadata.embedding
+
+        # Step 1: Find similar nodes for possible merging
+        similar_nodes = self.graph_store.search_by_embedding(
+            vector=embedding,
+            top_k=3,
+            scope=memory_type,
+            threshold=self._threshold,
+            status="activated",
+        )
+
+        if similar_nodes and similar_nodes[0]["score"] > self._merged_threshold:
+            self._merge(memory, similar_nodes)
+        else:
+            node_id = str(uuid.uuid4())
+            # Step 2: Add new node to graph
+            self.graph_store.add_node(
+                node_id, memory.memory, memory.metadata.model_dump(exclude_none=True)
+            )
+
+            # Step 3: Optionally link to a summary node based on topic
+            if memory.metadata.tags:
+                parent_id = self._ensure_structure_path(
+                    memory_type=memory_type, metadata=memory.metadata
+                )
+                if parent_id:
+                    self.graph_store.add_edge(parent_id, node_id, "PARENT")
+
+    def _merge(self, source_node: TextualMemoryItem, similar_nodes: list[dict]) -> None:
+        """
+        TODO: Add node traceability support by optionally preserving source nodes and linking them with MERGED_FROM edges.
+
+        Merge the source memory into the most similar existing node (only one),
+        and establish a MERGED_FROM edge in the graph.
+
+        Parameters:
+            source_node: The new memory item (not yet in the graph)
+            similar_nodes: A list of dicts returned by search_by_embedding(), ordered by similarity
+        """
+        original_node = similar_nodes[0]
+        original_id = original_node["id"]
+        original_data = self.graph_store.get_node(original_id)
+
+        target_text = original_data.get("memory", "")
+        merged_text = f"{target_text}\n⟵MERGED⟶\n{source_node.memory}"
+
+        original_meta = TreeNodeTextualMemoryMetadata(**original_data["metadata"])
+        source_meta = source_node.metadata
+
+        merged_key = source_meta.key or original_meta.key
+        merged_tags = list(set((original_meta.tags or []) + (source_meta.tags or [])))
+        merged_sources = list(set((original_meta.sources or []) + (source_meta.sources or [])))
+        merged_background = f"{original_meta.background}\n⟵MERGED⟶\n{source_meta.background}"
+        merged_embedding = self.embedder.embed([merged_text])[0]
+
+        merged_confidence = float((original_meta.confidence + source_meta.confidence) / 2)
+        merged_usage = list(set((original_meta.usage or []) + (source_meta.usage or [])))
+
+        # Create new merged node
+        merged_id = str(uuid.uuid4())
+        merged_metadata = source_meta.model_copy(
+            update={
+                "embedding": merged_embedding,
+                "updated_at": datetime.now().isoformat(),
+                "key": merged_key,
+                "tags": merged_tags,
+                "sources": merged_sources,
+                "background": merged_background,
+                "confidence": merged_confidence,
+                "usage": merged_usage,
+            }
+        )
+
+        self.graph_store.add_node(
+            merged_id, merged_text, merged_metadata.model_dump(exclude_none=True)
+        )
+
+        # Add traceability edges: both original and new point to merged node
+        self.graph_store.add_edge(original_id, merged_id, type="MERGED_TO")
+        self.graph_store.update_node(original_id, {"status": "archived"})
+        source_id = str(uuid.uuid4())
+        source_metadata = source_node.metadata.model_copy(update={"status": "archived"})
+        self.graph_store.add_node(source_id, source_node.memory, source_metadata.model_dump())
+        self.graph_store.add_edge(source_id, merged_id, type="MERGED_TO")
+        # After creating merged node and tracing lineage
+        self._inherit_edges(original_id, merged_id)
+
+        # Relate other similar nodes to merged if needed
+        for related_node in similar_nodes[1:]:
+            if not self.graph_store.edge_exists(
+                merged_id, related_node["id"], type="ANY", direction="ANY"
+            ):
+                self.graph_store.add_edge(merged_id, related_node["id"], type="RELATE")
+
+    def _inherit_edges(self, from_id: str, to_id: str) -> None:
+        """
+        Migrate all non-lineage edges from `from_id` to `to_id`,
+        and remove them from `from_id` after copying.
+        """
+        edges = self.graph_store.get_edges(from_id, type="ANY", direction="ANY")
+
+        for edge in edges:
+            if edge["type"] == "MERGED_TO":
+                continue  # Keep lineage edges
+
+            new_from = to_id if edge["from"] == from_id else edge["from"]
+            new_to = to_id if edge["to"] == from_id else edge["to"]
+
+            if new_from == new_to:
+                continue
+
+            # Add edge to merged node if it doesn't already exist
+            if not self.graph_store.edge_exists(new_from, new_to, edge["type"], direction="ANY"):
+                self.graph_store.add_edge(new_from, new_to, edge["type"])
+
+            # Remove original edge if it involved the archived node
+            self.graph_store.delete_edge(edge["from"], edge["to"], edge["type"])
+
+    def _ensure_structure_path(
+        self, memory_type: str, metadata: TreeNodeTextualMemoryMetadata
+    ) -> str:
+        """
+        Ensure structural path exists (ROOT → ... → final node), return last node ID.
+
+        Args:
+            path: like ["hobby", "photography"]
+
+        Returns:
+            Final node ID of the structure path.
+        """
+        # Step 1: Try to find an existing memory node with content == tag
+        existing = self.graph_store.get_by_metadata(
+            [
+                {"field": "memory", "op": "=", "value": metadata.key},
+                {"field": "memory_type", "op": "=", "value": memory_type},
+            ]
+        )
+        if existing:
+            node_id = existing[0]  # Use the first match
+        else:
+            # Step 2: If not found, create a new structure node
+            new_node = TextualMemoryItem(
+                memory=metadata.key,
+                metadata=TreeNodeTextualMemoryMetadata(
+                    user_id=metadata.user_id,
+                    session_id=metadata.session_id,
+                    memory_type=memory_type,
+                    status="activated",
+                    tags=[],
+                    key=metadata.key,
+                    embedding=self.embedder.embed([metadata.key])[0],
+                    usage=[],
+                    sources=[],
+                    confidence=0.99,
+                    background="",
+                ),
+            )
+
+            self.graph_store.add_node(
+                id=new_node.id,
+                memory=new_node.memory,
+                metadata=new_node.metadata.model_dump(exclude_none=True),
+            )
+            node_id = new_node.id
+
+        # Step 3: Return this structure node ID as the parent_id
+        return node_id

memos/memories/textual/tree_text_memory/retrieve/reasoner.py

@@ -0,0 +1,64 @@
+import json
+import re
+
+from string import Template
+
+from memos.memories.textual.item import TextualMemoryItem
+from memos.memories.textual.tree_text_memory.retrieve.retrieval_mid_structs import ParsedTaskGoal
+from memos.memories.textual.tree_text_memory.retrieve.utils import REASON_PROMPT
+
+
+class MemoryReasoner:
+    """
+    Memory reasoner that performs reasoning and knowledge synthesis
+    over retrieved memory items using a language model.
+    """
+
+    def __init__(self, llm):
+        self.llm = llm
+
+    def reason(
+        self, query: str, ranked_memories: list, parsed_goal: ParsedTaskGoal
+    ) -> list[TextualMemoryItem]:
+        """
+        Reason across multiple retrieved memory items and synthesize
+        a response or knowledge structure based on query objective.
+
+        Args:
+            query (str): Original user query description.
+            ranked_memories (list): List of relevant memory items.
+            parsed_goal (dict): Structured topic/concept/fact from TaskGoalParser.
+
+        Returns:
+            List of TextualMemoryItem: Refined memory items.
+        """
+        prompt_template = Template(REASON_PROMPT)
+        memory_detailed_str = "\n".join(
+            [
+                f"[{m.id}] ({m.metadata.hierarchy_level}) {m.metadata.key}: {m.memory}"
+                for m in ranked_memories
+            ]
+        )
+        prompt = prompt_template.substitute(task=query, detailed_memory_list=memory_detailed_str)
+
+        response = self.llm.generate([{"role": "user", "content": prompt}])
+        content = response.content if hasattr(response, "content") else response
+
+        # Step 1: Extract selected IDs
+        selected_ids = self._parse_selected_ids(content)
+        id_set = set(selected_ids)
+
+        return [m for m in ranked_memories if m.id in id_set]
+
+    def _parse_selected_ids(self, response_text: str) -> list[str]:
+        """
+        Extracts memory IDs from model response. Supports both simple text list and JSON.
+        """
+        try:
+            parsed = json.loads(response_text)
+            if isinstance(parsed, dict) and "selected_ids" in parsed:
+                return parsed["selected_ids"]
+        except json.JSONDecodeError:
+            pass
+
+        return re.findall(r"[a-f0-9\-]{36}", response_text)  # UUID pattern fallback