MemoryOS 0.2.2__py3-none-any.whl → 1.0.1__py3-none-any.whl

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

@@ -0,0 +1,231 @@
+"""BochaAI Search API retriever for tree text memory."""
+
+import json
+
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from datetime import datetime
+
+import requests
+
+from memos.embedders.factory import OllamaEmbedder
+from memos.log import get_logger
+from memos.mem_reader.base import BaseMemReader
+from memos.memories.textual.item import TextualMemoryItem
+
+
+logger = get_logger(__name__)
+
+
+class BochaAISearchAPI:
+    """BochaAI Search API Client"""
+
+    def __init__(self, api_key: str, max_results: int = 20):
+        """
+        Initialize BochaAI Search API client.
+
+        Args:
+            api_key: BochaAI API key
+            max_results: Maximum number of search results to retrieve
+        """
+        self.api_key = api_key
+        self.max_results = max_results
+
+        self.web_url = "https://api.bochaai.com/v1/web-search"
+        self.ai_url = "https://api.bochaai.com/v1/ai-search"
+
+        self.headers = {
+            "Authorization": f"Bearer {api_key}",
+            "Content-Type": "application/json",
+        }
+
+    def search_web(self, query: str, summary: bool = True, freshness="noLimit") -> list[dict]:
+        """
+        Perform a Web Search (equivalent to the first curl).
+
+        Args:
+            query: Search query string
+            summary: Whether to include summary in the results
+            freshness: Freshness filter (e.g. 'noLimit', 'day', 'week')
+
+        Returns:
+            A list of search result dicts
+        """
+        body = {
+            "query": query,
+            "summary": summary,
+            "freshness": freshness,
+            "count": self.max_results,
+        }
+        return self._post(self.web_url, body)
+
+    def search_ai(
+        self, query: str, answer: bool = False, stream: bool = False, freshness="noLimit"
+    ) -> list[dict]:
+        """
+        Perform an AI Search (equivalent to the second curl).
+
+        Args:
+            query: Search query string
+            answer: Whether BochaAI should generate an answer
+            stream: Whether to use streaming response
+            freshness: Freshness filter (e.g. 'noLimit', 'day', 'week')
+
+        Returns:
+            A list of search result dicts
+        """
+        body = {
+            "query": query,
+            "freshness": freshness,
+            "count": self.max_results,
+            "answer": answer,
+            "stream": stream,
+        }
+        return self._post(self.ai_url, body)
+
+    def _post(self, url: str, body: dict) -> list[dict]:
+        """Send POST request and parse BochaAI search results."""
+        try:
+            resp = requests.post(url, headers=self.headers, json=body)
+            resp.raise_for_status()
+            raw_data = resp.json()
+
+            # parse the nested structure correctly
+            # ✅ AI Search
+            if "messages" in raw_data:
+                results = []
+                for msg in raw_data["messages"]:
+                    if msg.get("type") == "source" and msg.get("content_type") == "webpage":
+                        try:
+                            content_json = json.loads(msg["content"])
+                            results.extend(content_json.get("value", []))
+                        except Exception as e:
+                            logger.error(f"Failed to parse message content: {e}")
+                return results
+
+            # ✅ Web Search
+            return raw_data.get("data", {}).get("webPages", {}).get("value", [])
+
+        except Exception:
+            import traceback
+
+            logger.error(f"BochaAI search error: {traceback.format_exc()}")
+            return []
+
+
+class BochaAISearchRetriever:
+    """BochaAI retriever that converts search results into TextualMemoryItem objects"""
+
+    def __init__(
+        self,
+        access_key: str,
+        embedder: OllamaEmbedder,
+        reader: BaseMemReader,
+        max_results: int = 20,
+    ):
+        """
+        Initialize BochaAI Search retriever.
+
+        Args:
+            access_key: BochaAI API key
+            embedder: Embedder instance for generating embeddings
+            reader: MemReader instance for processing internet content
+            max_results: Maximum number of search results to retrieve
+        """
+        self.bocha_api = BochaAISearchAPI(access_key, max_results=max_results)
+        self.embedder = embedder
+        self.reader = reader
+
+    def retrieve_from_internet(
+        self, query: str, top_k: int = 10, parsed_goal=None, info=None
+    ) -> list[TextualMemoryItem]:
+        """
+        Default internet retrieval (Web Search).
+        This keeps consistent API with Xinyu and Google retrievers.
+
+        Args:
+            query: Search query
+            top_k: Number of results to retrieve
+            parsed_goal: Parsed task goal (optional)
+            info (dict): Metadata for memory consumption tracking
+
+        Returns:
+            List of TextualMemoryItem
+        """
+        search_results = self.bocha_api.search_ai(query)  # ✅ default to
+        # web-search
+        return self._convert_to_mem_items(search_results, query, parsed_goal, info)
+
+    def retrieve_from_web(
+        self, query: str, top_k: int = 10, parsed_goal=None, info=None
+    ) -> list[TextualMemoryItem]:
+        """Explicitly retrieve using Bocha Web Search."""
+        search_results = self.bocha_api.search_web(query)
+        return self._convert_to_mem_items(search_results, query, parsed_goal, info)
+
+    def retrieve_from_ai(
+        self, query: str, top_k: int = 10, parsed_goal=None, info=None
+    ) -> list[TextualMemoryItem]:
+        """Explicitly retrieve using Bocha AI Search."""
+        search_results = self.bocha_api.search_ai(query)
+        return self._convert_to_mem_items(search_results, query, parsed_goal, info)
+
+    def _convert_to_mem_items(
+        self, search_results: list[dict], query: str, parsed_goal=None, info=None
+    ):
+        """Convert API search results into TextualMemoryItem objects."""
+        memory_items = []
+        if not info:
+            info = {"user_id": "", "session_id": ""}
+
+        with ThreadPoolExecutor(max_workers=8) as executor:
+            futures = [
+                executor.submit(self._process_result, r, query, parsed_goal, info)
+                for r in search_results
+            ]
+            for future in as_completed(futures):
+                try:
+                    memory_items.extend(future.result())
+                except Exception as e:
+                    logger.error(f"Error processing BochaAI search result: {e}")
+
+        # Deduplicate items by memory text
+        unique_memory_items = {item.memory: item for item in memory_items}
+        return list(unique_memory_items.values())
+
+    def _process_result(
+        self, result: dict, query: str, parsed_goal: str, info: None
+    ) -> list[TextualMemoryItem]:
+        """Process one Bocha search result into TextualMemoryItem."""
+        title = result.get("name", "")
+        content = result.get("summary", "") or result.get("snippet", "")
+        summary = result.get("snippet", "")
+        url = result.get("url", "")
+        publish_time = result.get("datePublished", "")
+
+        if publish_time:
+            try:
+                publish_time = datetime.fromisoformat(publish_time.replace("Z", "+00:00")).strftime(
+                    "%Y-%m-%d"
+                )
+            except Exception:
+                publish_time = datetime.now().strftime("%Y-%m-%d")
+        else:
+            publish_time = datetime.now().strftime("%Y-%m-%d")
+
+        # Use reader to split and process the content into chunks
+        read_items = self.reader.get_memory([content], type="doc", info=info)
+
+        memory_items = []
+        for read_item_i in read_items[0]:
+            read_item_i.memory = (
+                f"[Outer internet view] Title: {title}\nNewsTime:"
+                f" {publish_time}\nSummary:"
+                f" {summary}\n"
+                f"Content: {read_item_i.memory}"
+            )
+            read_item_i.metadata.source = "web"
+            read_item_i.metadata.memory_type = "OuterMemory"
+            read_item_i.metadata.sources = [url] if url else []
+            read_item_i.metadata.visibility = "public"
+            memory_items.append(read_item_i)
+        return memory_items

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

@@ -5,6 +5,7 @@ from typing import Any, ClassVar
 from memos.configs.internet_retriever import InternetRetrieverConfigFactory
 from memos.embedders.base import BaseEmbedder
 from memos.mem_reader.factory import MemReaderFactory
+from memos.memories.textual.tree_text_memory.retrieve.bochasearch import BochaAISearchRetriever
 from memos.memories.textual.tree_text_memory.retrieve.internet_retriever import (
     InternetGoogleRetriever,
 )
@@ -18,6 +19,7 @@ class InternetRetrieverFactory:
         "google": InternetGoogleRetriever,
         "bing": InternetGoogleRetriever,  # TODO: Implement BingRetriever
         "xinyu": XinyuSearchRetriever,
+        "bocha": BochaAISearchRetriever,
     }
 
     @classmethod
@@ -70,6 +72,13 @@ class InternetRetrieverFactory:
                 reader=MemReaderFactory.from_config(config.reader),
                 max_results=config.max_results,
             )
+        elif backend == "bocha":
+            return retriever_class(
+                access_key=config.api_key,  # Use api_key as access_key for xinyu
+                embedder=embedder,
+                reader=MemReaderFactory.from_config(config.reader),
+                max_results=config.max_results,
+            )
         else:
             raise ValueError(f"Unsupported backend: {backend}")
 

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

@@ -44,16 +44,23 @@ class GraphMemoryRetriever:
 
         if memory_scope == "WorkingMemory":
             # For working memory, retrieve all entries (no filtering)
-            working_memories = self.graph_store.get_all_memory_items(scope="WorkingMemory")
+            working_memories = self.graph_store.get_all_memory_items(
+                scope="WorkingMemory", include_embedding=True
+            )
             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)
+        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
+            # Structured graph-based retrieval
+            future_graph = executor.submit(self._graph_recall, parsed_goal, memory_scope)
+            # Vector similarity search
+            future_vector = executor.submit(
+                self._vector_recall, query_embedding, memory_scope, top_k
+            )
 
-        # Step 2: Vector similarity search
-        vector_results = self._vector_recall(query_embedding, memory_scope, top_k)
+            graph_results = future_graph.result()
+            vector_results = future_vector.result()
 
-        # Step 3: Merge and deduplicate results
+        # Merge and deduplicate by ID
         combined = {item.id: item for item in graph_results + vector_results}
 
         graph_ids = {item.id for item in graph_results}
@@ -67,6 +74,51 @@ class GraphMemoryRetriever:
 
         return list(combined.values())
 
+    def retrieve_from_cube(
+        self,
+        top_k: int,
+        memory_scope: str,
+        query_embedding: list[list[float]] | None = None,
+        cube_name: str = "memos_cube01",
+    ) -> 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:
+            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
+            cube_name: specify cube_name
+
+        Returns:
+            list: Combined memory items.
+        """
+        if memory_scope not in ["WorkingMemory", "LongTermMemory", "UserMemory"]:
+            raise ValueError(f"Unsupported memory scope: {memory_scope}")
+
+        graph_results = self._vector_recall(
+            query_embedding, memory_scope, top_k, cube_name=cube_name
+        )
+
+        for result_i in graph_results:
+            result_i.metadata.memory_type = "OuterMemory"
+        # Merge and deduplicate by ID
+        combined = {item.id: item for item in graph_results}
+
+        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]:
@@ -101,7 +153,7 @@ class GraphMemoryRetriever:
             return []
 
         # Load nodes and post-filter
-        node_dicts = self.graph_store.get_nodes(list(candidate_ids))
+        node_dicts = self.graph_store.get_nodes(list(candidate_ids), include_embedding=True)
 
         final_nodes = []
         for node in node_dicts:
@@ -127,7 +179,8 @@ class GraphMemoryRetriever:
         query_embedding: list[list[float]],
         memory_scope: str,
         top_k: int = 20,
-        max_num: int = 5,
+        max_num: int = 3,
+        cube_name: str | None = None,
     ) -> list[TextualMemoryItem]:
         """
         # TODO: tackle with post-filter and pre-filter(5.18+) better.
@@ -137,7 +190,9 @@ class GraphMemoryRetriever:
 
         def search_single(vec):
             return (
-                self.graph_store.search_by_embedding(vector=vec, top_k=top_k, scope=memory_scope)
+                self.graph_store.search_by_embedding(
+                    vector=vec, top_k=top_k, scope=memory_scope, cube_name=cube_name
+                )
                 or []
             )
 
@@ -152,6 +207,8 @@ class GraphMemoryRetriever:
 
         # 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))
+        node_dicts = self.graph_store.get_nodes(
+            list(unique_ids), include_embedding=True, cube_name=cube_name
+        )
 
         return [TextualMemoryItem.from_dict(record) for record in node_dicts]

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

@@ -78,7 +78,7 @@ class MemoryReranker:
         embeddings = [item.metadata.embedding for item in items_with_embeddings]
 
         if not embeddings:
-            return graph_results[:top_k]
+            return [(item, 0.5) for item in graph_results[:top_k]]
 
         # Step 2: Compute cosine similarities
         similarity_scores = batch_cosine_similarity(query_embedding, embeddings)