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

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

@@ -0,0 +1,263 @@
+"""Internet retrieval module for tree text memory."""
+
+import uuid
+
+from datetime import datetime
+
+import requests
+
+from memos.embedders.factory import OllamaEmbedder
+from memos.memories.textual.item import TextualMemoryItem, TreeNodeTextualMemoryMetadata
+
+
+class GoogleCustomSearchAPI:
+    """Google Custom Search API Client"""
+
+    def __init__(
+        self, api_key: str, search_engine_id: str, max_results: int = 20, num_per_request: int = 10
+    ):
+        """
+        Initialize Google Custom Search API client
+
+        Args:
+            api_key: Google API key
+            search_engine_id: Search engine ID (cx parameter)
+            max_results: Maximum number of results to retrieve
+            num_per_request: Number of results per API request
+        """
+        self.api_key = api_key
+        self.search_engine_id = search_engine_id
+        self.max_results = max_results
+        self.num_per_request = min(num_per_request, 10)  # Google API limits to 10
+        self.base_url = "https://www.googleapis.com/customsearch/v1"
+
+    def search(self, query: str, num_results: int | None = None, start_index: int = 1) -> dict:
+        """
+        Execute search request
+
+        Args:
+            query: Search query
+            num_results: Number of results to return (uses config default if None)
+            start_index: Starting index (default 1)
+
+        Returns:
+            Dictionary containing search results
+        """
+        if num_results is None:
+            num_results = self.num_per_request
+
+        params = {
+            "key": self.api_key,
+            "cx": self.search_engine_id,
+            "q": query,
+            "num": min(num_results, self.num_per_request),
+            "start": start_index,
+        }
+
+        try:
+            response = requests.get(self.base_url, params=params)
+            response.raise_for_status()
+            return response.json()
+        except requests.exceptions.RequestException as e:
+            print(f"Google search request failed: {e}")
+            return {}
+
+    def get_all_results(self, query: str, max_results: int | None = None) -> list[dict]:
+        """
+        Get all search results (with pagination)
+
+        Args:
+            query: Search query
+            max_results: Maximum number of results (uses config default if None)
+
+        Returns:
+            List of all search results
+        """
+        if max_results is None:
+            max_results = self.max_results
+
+        all_results = []
+        start_index = 1
+
+        while len(all_results) < max_results:
+            search_data = self.search(query, start_index=start_index)
+
+            if not search_data or "items" not in search_data:
+                break
+
+            all_results.extend(search_data["items"])
+
+            # Check if there are more results
+            if len(search_data["items"]) < self.num_per_request:
+                break
+
+            start_index += self.num_per_request
+
+            # Avoid infinite loop
+            if start_index > 100:
+                break
+
+        return all_results[:max_results]
+
+
+class InternetGoogleRetriever:
+    """Internet retriever that converts search results to TextualMemoryItem format"""
+
+    def __init__(
+        self,
+        api_key: str,
+        search_engine_id: str,
+        embedder: OllamaEmbedder,
+        max_results: int = 20,
+        num_per_request: int = 10,
+    ):
+        """
+        Initialize internet retriever
+
+        Args:
+            api_key: Google API key
+            search_engine_id: Search engine ID
+            embedder: Embedder instance for generating embeddings
+            max_results: Maximum number of results to retrieve
+            num_per_request: Number of results per API request
+        """
+        self.google_api = GoogleCustomSearchAPI(
+            api_key, search_engine_id, max_results=max_results, num_per_request=num_per_request
+        )
+        self.embedder = embedder
+
+    def retrieve_from_internet(
+        self, query: str, top_k: int = 10, parsed_goal=None
+    ) -> list[TextualMemoryItem]:
+        """
+        Retrieve information from the internet and convert to TextualMemoryItem format
+
+        Args:
+            query: Search query
+            top_k: Number of results to return
+            parsed_goal: Parsed task goal (optional)
+
+        Returns:
+            List of TextualMemoryItem
+        """
+        # Get search results
+        search_results = self.google_api.get_all_results(query, max_results=top_k)
+
+        # Convert to TextualMemoryItem format
+        memory_items = []
+
+        for _, result in enumerate(search_results):
+            # Extract basic information
+            title = result.get("title", "")
+            snippet = result.get("snippet", "")
+            link = result.get("link", "")
+            display_link = result.get("displayLink", "")
+
+            # Combine memory content
+            memory_content = f"Title: {title}\nSummary: {snippet}\nSource: {link}"
+            # Create metadata
+            metadata = TreeNodeTextualMemoryMetadata(
+                user_id=None,
+                session_id=None,
+                status="activated",
+                type="fact",  # Internet search results are usually factual information
+                memory_time=datetime.now().strftime("%Y-%m-%d"),
+                source="web",
+                confidence=85.0,  # Confidence level for internet information
+                entities=self._extract_entities(title, snippet),
+                tags=self._extract_tags(title, snippet, parsed_goal),
+                visibility="public",
+                memory_type="LongTermMemory",  # Internet search results as working memory
+                key=title,
+                sources=[link] if link else [],
+                embedding=self.embedder.embed([memory_content])[0],  # Can add embedding later
+                created_at=datetime.now().isoformat(),
+                usage=[],
+                background=f"Internet search result from {display_link}",
+            )
+
+            # Create TextualMemoryItem
+            memory_item = TextualMemoryItem(
+                id=str(uuid.uuid4()), memory=memory_content, metadata=metadata
+            )
+
+            memory_items.append(memory_item)
+
+        return memory_items
+
+    def _extract_entities(self, title: str, snippet: str) -> list[str]:
+        """
+        Extract entities from title and snippet
+
+        Args:
+            title: Title
+            snippet: Snippet
+
+        Returns:
+            List of entities
+        """
+        # Simple entity extraction logic, can be improved as needed
+        text = f"{title} {snippet}"
+        entities = []
+
+        # Extract possible organization names (with common suffixes)
+        org_suffixes = ["Inc", "Corp", "LLC", "Ltd", "Company", "University", "Institute"]
+        words = text.split()
+        for i, word in enumerate(words):
+            if word in org_suffixes and i > 0:
+                entities.append(f"{words[i - 1]} {word}")
+
+        # Extract possible dates
+        import re
+
+        date_pattern = r"\d{4}-\d{2}-\d{2}|\d{1,2}/\d{1,2}/\d{4}|\w+ \d{1,2}, \d{4}"
+        dates = re.findall(date_pattern, text)
+        entities.extend(dates)
+
+        return entities[:5]  # Limit number of entities
+
+    def _extract_tags(self, title: str, snippet: str, parsed_goal=None) -> list[str]:
+        """
+        Extract tags from title and snippet
+
+        Args:
+            title: Title
+            snippet: Snippet
+            parsed_goal: Parsed task goal
+
+        Returns:
+            List of tags
+        """
+        tags = []
+
+        # Extract tags from parsed goal
+        if parsed_goal:
+            if hasattr(parsed_goal, "topic") and parsed_goal.topic:
+                tags.append(parsed_goal.topic)
+            if hasattr(parsed_goal, "concept") and parsed_goal.concept:
+                tags.append(parsed_goal.concept)
+
+        # Extract keywords from text
+        text = f"{title} {snippet}".lower()
+
+        # Simple keyword extraction
+        keywords = [
+            "news",
+            "report",
+            "article",
+            "study",
+            "research",
+            "analysis",
+            "update",
+            "announcement",
+            "policy",
+            "memo",
+            "document",
+        ]
+
+        for keyword in keywords:
+            if keyword in text:
+                tags.append(keyword)
+
+        # Remove duplicates and limit count
+        return list(set(tags))[:10]

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

@@ -0,0 +1,89 @@
+"""Factory for creating internet retrievers."""
+
+from typing import Any, ClassVar
+
+from memos.configs.internet_retriever import InternetRetrieverConfigFactory
+from memos.embedders.base import BaseEmbedder
+from memos.memories.textual.tree_text_memory.retrieve.internet_retriever import (
+    InternetGoogleRetriever,
+)
+from memos.memories.textual.tree_text_memory.retrieve.xinyusearch import XinyuSearchRetriever
+
+
+class InternetRetrieverFactory:
+    """Factory class for creating internet retriever instances."""
+
+    backend_to_class: ClassVar[dict[str, Any]] = {
+        "google": InternetGoogleRetriever,
+        "bing": InternetGoogleRetriever,  # TODO: Implement BingRetriever
+        "xinyu": XinyuSearchRetriever,
+    }
+
+    @classmethod
+    def from_config(
+        cls, config_factory: InternetRetrieverConfigFactory, embedder: BaseEmbedder
+    ) -> InternetGoogleRetriever | None:
+        """
+        Create internet retriever from configuration.
+
+        Args:
+            config_factory: Internet retriever configuration
+            embedder: Embedder instance for generating embeddings
+
+        Returns:
+            InternetRetriever instance or None if no configuration provided
+        """
+        if config_factory.backend is None:
+            return None
+
+        backend = config_factory.backend
+        if backend not in cls.backend_to_class:
+            raise ValueError(f"Invalid internet retriever backend: {backend}")
+
+        retriever_class = cls.backend_to_class[backend]
+        config = config_factory.config
+
+        # Create retriever with appropriate parameters
+        if backend == "google":
+            return retriever_class(
+                api_key=config.api_key,
+                search_engine_id=config.search_engine_id,
+                embedder=embedder,
+                max_results=config.max_results,
+                num_per_request=config.num_per_request,
+            )
+        elif backend == "bing":
+            # TODO: Implement Bing retriever
+            return retriever_class(
+                api_key=config.api_key,
+                search_engine_id=None,  # Bing doesn't use search_engine_id
+                embedder=embedder,
+                max_results=config.max_results,
+                num_per_request=config.num_per_request,
+            )
+        elif backend == "xinyu":
+            return retriever_class(
+                access_key=config.api_key,  # Use api_key as access_key for xinyu
+                search_engine_id=config.search_engine_id,
+                embedder=embedder,
+                max_results=config.max_results,
+            )
+        else:
+            raise ValueError(f"Unsupported backend: {backend}")
+
+    @classmethod
+    def create_google_retriever(
+        cls, api_key: str, search_engine_id: str, embedder: BaseEmbedder
+    ) -> InternetGoogleRetriever:
+        """
+        Create Google Custom Search retriever.
+
+        Args:
+            api_key: Google API key
+            search_engine_id: Google Custom Search Engine ID
+            embedder: Embedder instance
+
+        Returns:
+            InternetRetriever instance
+        """
+        return InternetGoogleRetriever(api_key, search_engine_id, embedder)

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

@@ -0,0 +1,61 @@
+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.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

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

@@ -0,0 +1,158 @@
+import concurrent.futures
+
+from memos.embedders.factory import OllamaEmbedder
+from memos.graph_dbs.neo4j import Neo4jGraphDB
+from memos.memories.textual.item import TextualMemoryItem
+from memos.memories.textual.tree_text_memory.retrieve.retrieval_mid_structs import ParsedTaskGoal
+
+
+class GraphMemoryRetriever:
+    """
+    Unified memory retriever that combines both graph-based and vector-based retrieval logic.
+    """
+
+    def __init__(self, graph_store: Neo4jGraphDB, embedder: OllamaEmbedder):
+        self.graph_store = graph_store
+        self.embedder = embedder
+
+    def retrieve(
+        self,
+        query: str,
+        parsed_goal: ParsedTaskGoal,
+        top_k: int,
+        memory_scope: str,
+        query_embedding: list[list[float]] | None = None,
+    ) -> list[TextualMemoryItem]:
+        """
+        Perform hybrid memory retrieval:
+        - Run graph-based lookup from dispatch plan.
+        - Run vector similarity search from embedded query.
+        - Merge and return combined result set.
+
+        Args:
+            query (str): Original task query.
+            parsed_goal (dict): parsed_goal.
+            top_k (int): Number of candidates to return.
+            memory_scope (str): One of ['working', 'long_term', 'user'].
+            query_embedding(list of embedding): list of embedding of query
+
+        Returns:
+            list: Combined memory items.
+        """
+        if memory_scope not in ["WorkingMemory", "LongTermMemory", "UserMemory"]:
+            raise ValueError(f"Unsupported memory scope: {memory_scope}")
+
+        if memory_scope == "WorkingMemory":
+            # For working memory, retrieve all entries (no filtering)
+            working_memories = self.graph_store.get_all_memory_items(scope="WorkingMemory")
+            return [TextualMemoryItem.from_dict(record) for record in working_memories]
+
+        # Step 1: Structured graph-based retrieval
+        graph_results = self._graph_recall(parsed_goal, memory_scope)
+
+        # Step 2: Vector similarity search
+        vector_results = self._vector_recall(query_embedding, memory_scope, top_k)
+
+        # Step 3: Merge and deduplicate results
+        combined = {item.id: item for item in graph_results + vector_results}
+
+        # Debug: 打印在 graph_results 中但不在 combined 中的 id
+        graph_ids = {item.id for item in graph_results}
+        combined_ids = set(combined.keys())
+        lost_ids = graph_ids - combined_ids
+
+        if lost_ids:
+            print(
+                f"[DEBUG] The following nodes were in graph_results but missing in combined: {lost_ids}"
+            )
+
+        return list(combined.values())
+
+    def _graph_recall(
+        self, parsed_goal: ParsedTaskGoal, memory_scope: str
+    ) -> list[TextualMemoryItem]:
+        """
+        Perform structured node-based retrieval from Neo4j.
+        - keys must match exactly (n.key IN keys)
+        - tags must overlap with at least 2 input tags
+        - scope filters by memory_type if provided
+        """
+        candidate_ids = set()
+
+        # 1) key-based OR branch
+        if parsed_goal.keys:
+            key_filters = [
+                {"field": "key", "op": "in", "value": parsed_goal.keys},
+                {"field": "memory_type", "op": "=", "value": memory_scope},
+            ]
+            key_ids = self.graph_store.get_by_metadata(key_filters)
+            candidate_ids.update(key_ids)
+
+        # 2) tag-based OR branch
+        if parsed_goal.tags:
+            tag_filters = [
+                {"field": "tags", "op": "contains", "value": parsed_goal.tags},
+                {"field": "memory_type", "op": "=", "value": memory_scope},
+            ]
+            tag_ids = self.graph_store.get_by_metadata(tag_filters)
+            candidate_ids.update(tag_ids)
+
+        # No matches → return empty
+        if not candidate_ids:
+            return []
+
+        # Load nodes and post-filter
+        node_dicts = self.graph_store.get_nodes(list(candidate_ids))
+
+        final_nodes = []
+        for node in node_dicts:
+            meta = node.get("metadata", {})
+            node_key = meta.get("key")
+            node_tags = meta.get("tags", []) or []
+
+            keep = False
+            # key equals to node_key
+            if parsed_goal.keys and node_key in parsed_goal.keys:
+                keep = True
+            # overlap tags more than 2
+            elif parsed_goal.tags:
+                overlap = len(set(node_tags) & set(parsed_goal.tags))
+                if overlap >= 2:
+                    keep = True
+            if keep:
+                final_nodes.append(TextualMemoryItem.from_dict(node))
+        return final_nodes
+
+    def _vector_recall(
+        self,
+        query_embedding: list[list[float]],
+        memory_scope: str,
+        top_k: int = 20,
+        max_num: int = 5,
+    ) -> list[TextualMemoryItem]:
+        """
+        # TODO: tackle with post-filter and pre-filter(5.18+) better.
+        Perform vector-based similarity retrieval using query embedding.
+        """
+        all_matches = []
+
+        def search_single(vec):
+            return (
+                self.graph_store.search_by_embedding(vector=vec, top_k=top_k, scope=memory_scope)
+                or []
+            )
+
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            futures = [executor.submit(search_single, vec) for vec in query_embedding[:max_num]]
+            for future in concurrent.futures.as_completed(futures):
+                result = future.result()
+                all_matches.extend(result)
+
+        if not all_matches:
+            return []
+
+        # Step 3: Extract matched IDs and retrieve full nodes
+        unique_ids = set({r["id"] for r in all_matches})
+        node_dicts = self.graph_store.get_nodes(list(unique_ids))
+
+        return [TextualMemoryItem.from_dict(record) for record in node_dicts]

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

@@ -0,0 +1,111 @@
+import numpy as np
+
+from memos.embedders.factory import OllamaEmbedder
+from memos.llms.factory import OllamaLLM, OpenAILLM
+from memos.memories.textual.item import TextualMemoryItem
+from memos.memories.textual.tree_text_memory.retrieve.retrieval_mid_structs import ParsedTaskGoal
+
+
+def batch_cosine_similarity(
+    query_vec: list[float], candidate_vecs: list[list[float]]
+) -> list[float]:
+    """
+    Compute cosine similarity between a single query vector and multiple candidate vectors using NumPy.
+
+    Args:
+        query_vec (list[float]): The query embedding.
+        candidate_vecs (list[list[float]]): A list of memory embeddings.
+
+    Returns:
+        list[float]: Cosine similarity scores for each candidate.
+    """
+    query = np.array(query_vec)
+    candidates = np.array(candidate_vecs)
+
+    # Normalize query and candidates
+    query_norm = np.linalg.norm(query)
+    candidates_norm = np.linalg.norm(candidates, axis=1)
+
+    # Compute dot products
+    dot_products = np.dot(candidates, query)
+
+    # Avoid division by zero
+    eps = 1e-10
+    similarities = dot_products / (candidates_norm * query_norm + eps)
+
+    return similarities.tolist()
+
+
+class MemoryReranker:
+    """
+    Rank retrieved memory cards by structural priority and contextual similarity.
+    """
+
+    def __init__(self, llm: OpenAILLM | OllamaLLM, embedder: OllamaEmbedder):
+        self.llm = llm
+        self.embedder = embedder
+
+        # Structural priority weights
+        self.level_weights = {
+            "topic": 1.0,
+            "concept": 1.0,
+            "fact": 1.0,
+        }
+
+    def rerank(
+        self,
+        query: str,
+        query_embedding: list[float],
+        graph_results: list,
+        top_k: int,
+        parsed_goal: ParsedTaskGoal,
+    ) -> list[tuple[TextualMemoryItem, float]]:
+        """
+        Rerank memory items by relevance to task.
+
+        Args:
+            query (str): Original task.
+            query_embedding(list[float]): embedding of query
+            graph_results (list): Combined retrieval results.
+            top_k (int): Number of top results to return.
+            parsed_goal (dict): Structured task representation.
+
+        Returns:
+            list(tuple): Ranked list of memory items with similarity score.
+        """
+        # Step 1: Filter out items without embeddings
+        items_with_embeddings = [item for item in graph_results if item.metadata.embedding]
+        embeddings = [item.metadata.embedding for item in items_with_embeddings]
+
+        if not embeddings:
+            return graph_results[:top_k]
+
+        # Step 2: Compute cosine similarities
+        similarity_scores = batch_cosine_similarity(query_embedding, embeddings)
+
+        # Step 3: Apply structural weight boost
+        def get_weight(item: TextualMemoryItem) -> float:
+            level = item.metadata.background
+            return self.level_weights.get(level, 1.0)
+
+        weighted_scores = [
+            sim * get_weight(item)
+            for sim, item in zip(similarity_scores, items_with_embeddings, strict=False)
+        ]
+
+        # Step 4: Sort by weighted score
+        sorted_items = sorted(
+            zip(items_with_embeddings, weighted_scores, strict=False),
+            key=lambda pair: pair[1],
+            reverse=True,
+        )
+
+        # Step 5: Return top-k items with fallback
+        top_items = sorted_items[:top_k]
+
+        if len(top_items) < top_k:
+            selected_items = [item for item, _ in top_items]
+            remaining = [(item, -1.0) for item in graph_results if item not in selected_items]
+            top_items.extend(remaining[: top_k - len(top_items)])
+
+        return top_items  # list of (item, score)

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

@@ -0,0 +1,13 @@
+from dataclasses import dataclass, field
+
+
+@dataclass
+class ParsedTaskGoal:
+    """
+    Goal structure for both Fast & LLM.
+    """
+
+    memories: list[str] = field(default_factory=list)
+    keys: list[str] = field(default_factory=list)
+    tags: list[str] = field(default_factory=list)
+    goal_type: str | None = None  # e.g., 'default', 'explanation', etc.