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

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

@@ -73,10 +73,12 @@ class RelationAndReasoningDetector:
             results["sequence_links"].extend(seq)
             """
 
+            """
             # 4) Aggregate
             agg = self._detect_aggregate_node_for_group(node, nearest, min_group_size=5)
             if agg:
                 results["aggregate_nodes"].append(agg)
+            """
 
         except Exception as e:
             logger.error(

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

@@ -3,7 +3,7 @@ import threading
 import time
 import traceback
 
-from collections import Counter, defaultdict
+from collections import defaultdict
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from queue import PriorityQueue
 from typing import Literal
@@ -17,8 +17,7 @@ from memos.graph_dbs.neo4j import Neo4jGraphDB
 from memos.llms.base import BaseLLM
 from memos.log import get_logger
 from memos.memories.textual.item import TreeNodeTextualMemoryMetadata
-from memos.memories.textual.tree_text_memory.organize.conflict import ConflictHandler
-from memos.memories.textual.tree_text_memory.organize.redundancy import RedundancyHandler
+from memos.memories.textual.tree_text_memory.organize.handler import NodeHandler
 from memos.memories.textual.tree_text_memory.organize.relation_reason_detector import (
     RelationAndReasoningDetector,
 )
@@ -63,10 +62,10 @@ class GraphStructureReorganizer:
         self.relation_detector = RelationAndReasoningDetector(
             self.graph_store, self.llm, self.embedder
         )
-        self.conflict = ConflictHandler(graph_store=graph_store, llm=llm, embedder=embedder)
-        self.redundancy = RedundancyHandler(graph_store=graph_store, llm=llm, embedder=embedder)
+        self.resolver = NodeHandler(graph_store=graph_store, llm=llm, embedder=embedder)
 
         self.is_reorganize = is_reorganize
+        self._reorganize_needed = True
         if self.is_reorganize:
             # ____ 1. For queue message driven thread ___________
             self.thread = threading.Thread(target=self._run_message_consumer_loop)
@@ -125,13 +124,17 @@ class GraphStructureReorganizer:
         """
         import schedule
 
-        schedule.every(600).seconds.do(self.optimize_structure, scope="LongTermMemory")
-        schedule.every(600).seconds.do(self.optimize_structure, scope="UserMemory")
+        schedule.every(100).seconds.do(self.optimize_structure, scope="LongTermMemory")
+        schedule.every(100).seconds.do(self.optimize_structure, scope="UserMemory")
 
         logger.info("Structure optimizer schedule started.")
         while not getattr(self, "_stop_scheduler", False):
-            schedule.run_pending()
-            time.sleep(1)
+            if self._reorganize_needed:
+                logger.info("[Reorganizer] Triggering optimize_structure due to new nodes.")
+                self.optimize_structure(scope="LongTermMemory")
+                self.optimize_structure(scope="UserMemory")
+                self._reorganize_needed = False
+            time.sleep(30)
 
     def stop(self):
         """
@@ -148,45 +151,31 @@ class GraphStructureReorganizer:
         logger.info("Structure optimizer stopped.")
 
     def handle_message(self, message: QueueMessage):
-        handle_map = {
-            "add": self.handle_add,
-            "remove": self.handle_remove,
-            "merge": self.handle_merge,
-        }
+        handle_map = {"add": self.handle_add, "remove": self.handle_remove}
         handle_map[message.op](message)
         logger.debug(f"message queue size: {self.queue.qsize()}")
 
     def handle_add(self, message: QueueMessage):
         logger.debug(f"Handling add operation: {str(message)[:500]}")
-        # ———————— 1. check for conflicts ————————
         added_node = message.after_node[0]
-        conflicts = self.conflict.detect(added_node, scope=added_node.metadata.memory_type)
-        if conflicts:
-            for added_node, existing_node in conflicts:
-                self.conflict.resolve(added_node, existing_node)
-                logger.info(f"Resolved conflict between {added_node.id} and {existing_node.id}.")
-
-        # ———————— 2. check for redundancy ————————
-        redundancies = self.redundancy.detect(added_node, scope=added_node.metadata.memory_type)
-        if redundancies:
-            for added_node, existing_node in redundancies:
-                self.redundancy.resolve_two_nodes(added_node, existing_node)
-                logger.info(f"Resolved redundancy between {added_node.id} and {existing_node.id}.")
+        detected_relationships = self.resolver.detect(
+            added_node, scope=added_node.metadata.memory_type
+        )
+        if detected_relationships:
+            for added_node, existing_node, relation in detected_relationships:
+                self.resolver.resolve(added_node, existing_node, relation)
+
+        self._reorganize_needed = True
 
     def handle_remove(self, message: QueueMessage):
         logger.debug(f"Handling remove operation: {str(message)[:50]}")
 
-    def handle_merge(self, message: QueueMessage):
-        after_node = message.after_node[0]
-        logger.debug(f"Handling merge operation: <{after_node.memory}>")
-        self.redundancy.resolve_one_node(after_node)
-
     def optimize_structure(
         self,
         scope: str = "LongTermMemory",
         local_tree_threshold: int = 10,
-        min_cluster_size: int = 3,
-        min_group_size: int = 5,
+        min_cluster_size: int = 4,
+        min_group_size: int = 20,
     ):
         """
         Periodically reorganize the graph:
@@ -253,7 +242,7 @@ class GraphStructureReorganizer:
                     except Exception as e:
                         logger.warning(
                             f"[Reorganize] Cluster processing "
-                            f"failed: {e}, trace: {traceback.format_exc()}"
+                            f"failed: {e}, cluster_nodes: {cluster_nodes}, trace: {traceback.format_exc()}"
                         )
                 logger.info("[GraphStructure Reorganize] Structure optimization finished.")
 
@@ -271,29 +260,23 @@ class GraphStructureReorganizer:
         if len(cluster_nodes) <= min_cluster_size:
             return
 
-        if len(cluster_nodes) <= local_tree_threshold:
-            # Small cluster ➜ single parent
-            parent_node = self._summarize_cluster(cluster_nodes, scope)
-            self._create_parent_node(parent_node)
-            self._link_cluster_nodes(parent_node, cluster_nodes)
-        else:
-            # Large cluster ➜ local sub-clustering
-            sub_clusters = self._local_subcluster(cluster_nodes)
-            sub_parents = []
-
-            for sub_nodes in sub_clusters:
-                if len(sub_nodes) < min_cluster_size:
-                    continue  # Skip tiny noise
-                sub_parent_node = self._summarize_cluster(sub_nodes, scope)
-                self._create_parent_node(sub_parent_node)
-                self._link_cluster_nodes(sub_parent_node, sub_nodes)
-                sub_parents.append(sub_parent_node)
-
-            if sub_parents:
-                cluster_parent_node = self._summarize_cluster(cluster_nodes, scope)
-                self._create_parent_node(cluster_parent_node)
-                for sub_parent in sub_parents:
-                    self.graph_store.add_edge(cluster_parent_node.id, sub_parent.id, "PARENT")
+        # Large cluster ➜ local sub-clustering
+        sub_clusters = self._local_subcluster(cluster_nodes)
+        sub_parents = []
+
+        for sub_nodes in sub_clusters:
+            if len(sub_nodes) < min_cluster_size:
+                continue  # Skip tiny noise
+            sub_parent_node = self._summarize_cluster(sub_nodes, scope)
+            self._create_parent_node(sub_parent_node)
+            self._link_cluster_nodes(sub_parent_node, sub_nodes)
+            sub_parents.append(sub_parent_node)
+
+        if sub_parents and len(sub_parents) >= min_cluster_size:
+            cluster_parent_node = self._summarize_cluster(cluster_nodes, scope)
+            self._create_parent_node(cluster_parent_node)
+            for sub_parent in sub_parents:
+                self.graph_store.add_edge(cluster_parent_node.id, sub_parent.id, "PARENT")
 
         logger.info("Adding relations/reasons")
         nodes_to_check = cluster_nodes
@@ -350,7 +333,9 @@ class GraphStructureReorganizer:
 
             logger.info("[Reorganizer] Cluster relation/reasoning done.")
 
-    def _local_subcluster(self, cluster_nodes: list[GraphDBNode]) -> list[list[GraphDBNode]]:
+    def _local_subcluster(
+        self, cluster_nodes: list[GraphDBNode], max_length: int = 8000
+    ) -> (list)[list[GraphDBNode]]:
         """
         Use LLM to split a large cluster into semantically coherent sub-clusters.
         """
@@ -364,7 +349,9 @@ class GraphStructureReorganizer:
             scene_lines.append(line)
 
         joined_scene = "\n".join(scene_lines)
-        prompt = LOCAL_SUBCLUSTER_PROMPT.replace("{joined_scene}", joined_scene)
+        if len(joined_scene) > max_length:
+            logger.warning(f"Sub-cluster too long: {joined_scene}")
+        prompt = LOCAL_SUBCLUSTER_PROMPT.replace("{joined_scene}", joined_scene[:max_length])
 
         messages = [{"role": "user", "content": prompt}]
         response_text = self.llm.generate(messages)
@@ -389,12 +376,12 @@ class GraphStructureReorganizer:
         install_command="pip install scikit-learn",
         install_link="https://scikit-learn.org/stable/install.html",
     )
-    def _partition(self, nodes, min_cluster_size: int = 3, max_cluster_size: int = 20):
+    def _partition(self, nodes, min_cluster_size: int = 10, max_cluster_size: int = 20):
         """
         Partition nodes by:
-        1) Frequent tags (top N & above threshold)
-        2) Remaining nodes by embedding clustering (MiniBatchKMeans)
-        3) Small clusters merged or assigned to 'Other'
+        - If total nodes <= max_cluster_size -> return all nodes in one cluster.
+        - If total nodes > max_cluster_size -> cluster by embeddings, recursively split.
+        - Only keep clusters with size > min_cluster_size.
 
         Args:
             nodes: List of GraphDBNode
@@ -405,105 +392,80 @@ class GraphStructureReorganizer:
         """
         from sklearn.cluster import MiniBatchKMeans
 
-        # 1) Count all tags
-        tag_counter = Counter()
-        for node in nodes:
-            for tag in node.metadata.tags:
-                tag_counter[tag] += 1
-
-        # Select frequent tags
-        top_n_tags = {tag for tag, count in tag_counter.most_common(50)}
-        threshold_tags = {tag for tag, count in tag_counter.items() if count >= 50}
-        frequent_tags = top_n_tags | threshold_tags
-
-        # Group nodes by tags
-        tag_groups = defaultdict(list)
-
-        for node in nodes:
-            for tag in node.metadata.tags:
-                if tag in frequent_tags:
-                    tag_groups[tag].append(node)
-                    break
-
-        filtered_tag_clusters = []
-        assigned_ids = set()
-        for tag, group in tag_groups.items():
-            if len(group) >= min_cluster_size:
-                # Split large groups into chunks of at most max_cluster_size
-                for i in range(0, len(group), max_cluster_size):
-                    sub_group = group[i : i + max_cluster_size]
-                    filtered_tag_clusters.append(sub_group)
-                    assigned_ids.update(n.id for n in sub_group)
-            else:
-                logger.info(f"... dropped tag {tag} due to low size ...")
-
-        logger.info(
-            f"[MixedPartition] Created {len(filtered_tag_clusters)} clusters from tags. "
-            f"Nodes grouped by tags: {len(assigned_ids)} / {len(nodes)}"
-        )
-
-        # Remaining nodes -> embedding clustering
-        remaining_nodes = [n for n in nodes if n.id not in assigned_ids]
-        logger.info(
-            f"[MixedPartition] Remaining nodes for embedding clustering: {len(remaining_nodes)}"
-        )
-
-        embedding_clusters = []
+        if len(nodes) <= max_cluster_size:
+            logger.info(
+                f"[KMeansPartition] Node count {len(nodes)} <= {max_cluster_size}, skipping KMeans."
+            )
+            return [nodes]
 
-        def recursive_clustering(nodes_list):
+        def recursive_clustering(nodes_list, depth=0):
             """Recursively split clusters until each is <= max_cluster_size."""
+            indent = "  " * depth
+            logger.info(
+                f"{indent}[Recursive] Start clustering {len(nodes_list)} nodes at depth {depth}"
+            )
+
             if len(nodes_list) <= max_cluster_size:
+                logger.info(
+                    f"{indent}[Recursive] Node count <= {max_cluster_size}, stop splitting."
+                )
                 return [nodes_list]
-
             # Try kmeans with k = ceil(len(nodes) / max_cluster_size)
-            x = np.array([n.metadata.embedding for n in nodes_list if n.metadata.embedding])
-            if len(x) < 2:
+            x_nodes = [n for n in nodes_list if n.metadata.embedding]
+            x = np.array([n.metadata.embedding for n in x_nodes])
+
+            if len(x) < min_cluster_size:
+                logger.info(
+                    f"{indent}[Recursive] Too few embeddings ({len(x)}), skipping clustering."
+                )
                 return [nodes_list]
 
             k = min(len(x), (len(nodes_list) + max_cluster_size - 1) // max_cluster_size)
-            k = max(1, min(k, len(x)))
+            k = max(1, k)
 
             try:
+                logger.info(f"{indent}[Recursive] Clustering with k={k} on {len(x)} points.")
                 kmeans = MiniBatchKMeans(n_clusters=k, batch_size=256, random_state=42)
                 labels = kmeans.fit_predict(x)
 
                 label_groups = defaultdict(list)
-                for node, label in zip(nodes_list, labels, strict=False):
+                for node, label in zip(x_nodes, labels, strict=False):
                     label_groups[label].append(node)
 
+                # Map: label -> nodes with no embedding (fallback group)
+                no_embedding_nodes = [n for n in nodes_list if not n.metadata.embedding]
+                if no_embedding_nodes:
+                    logger.warning(
+                        f"{indent}[Recursive] {len(no_embedding_nodes)} nodes have no embedding. Added to largest cluster."
+                    )
+                    # Assign to largest cluster
+                    largest_label = max(label_groups.items(), key=lambda kv: len(kv[1]))[0]
+                    label_groups[largest_label].extend(no_embedding_nodes)
+
                 result = []
-                for sub_group in label_groups.values():
-                    result.extend(recursive_clustering(sub_group))
+                for label, sub_group in label_groups.items():
+                    logger.info(f"{indent}  Cluster-{label}: {len(sub_group)} nodes")
+                    result.extend(recursive_clustering(sub_group, depth=depth + 1))
                 return result
+
             except Exception as e:
-                logger.warning(f"Clustering failed: {e}, falling back to single cluster.")
+                logger.warning(
+                    f"{indent}[Recursive] Clustering failed: {e}, fallback to one cluster."
+                )
                 return [nodes_list]
 
-        if remaining_nodes:
-            clusters = recursive_clustering(remaining_nodes)
-            embedding_clusters.extend(clusters)
-            logger.info(
-                f"[MixedPartition] Created {len(embedding_clusters)} clusters from embeddings."
-            )
-
-        # Merge all clusters
-        all_clusters = filtered_tag_clusters + embedding_clusters
+        raw_clusters = recursive_clustering(nodes)
+        filtered_clusters = [c for c in raw_clusters if len(c) > min_cluster_size]
 
-        # Handle small clusters (< min_cluster_size)
-        final_clusters = []
-        small_nodes = []
-        for group in all_clusters:
-            if len(group) < min_cluster_size:
-                small_nodes.extend(group)
-            else:
-                final_clusters.append(group)
+        logger.info(f"[KMeansPartition] Total clusters before filtering: {len(raw_clusters)}")
+        for i, cluster in enumerate(raw_clusters):
+            logger.info(f"[KMeansPartition]   Cluster-{i}: {len(cluster)} nodes")
 
-        if small_nodes:
-            final_clusters.append(small_nodes)
-            logger.info(f"[MixedPartition] {len(small_nodes)} nodes assigned to 'Other' cluster.")
+        logger.info(
+            f"[KMeansPartition] Clusters after filtering (>{min_cluster_size}): {len(filtered_clusters)}"
+        )
 
-        logger.info(f"[MixedPartition] Total final clusters: {len(final_clusters)}")
-        return final_clusters
+        return filtered_clusters
 
     def _summarize_cluster(self, cluster_nodes: list[GraphDBNode], scope: str) -> GraphDBNode:
         """
@@ -600,7 +562,7 @@ class GraphStructureReorganizer:
         for i, node in enumerate(message.after_node or []):
             if not isinstance(node, str):
                 continue
-            raw_node = self.graph_store.get_node(node)
+            raw_node = self.graph_store.get_node(node, include_embedding=True)
             if raw_node is None:
                 logger.debug(f"Node with ID {node} not found in the graph store.")
                 message.after_node[i] = None

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

@@ -0,0 +1,229 @@
+"""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"Title: {title}\nNewsTime: {publish_time}\nSummary: {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}")