feat: finish bocha search

Author: CaralHsi

examples/basic_modules/textual_memory_internet_search_example.py

@@ -26,12 +26,12 @@
     {
         "backend": "bocha",
         "config": {
-            "api_key": "sk-xxxx",  # 🔑 Your BochaAI API Key
-            "search_engine_id": "",  # Not required for BochaAI, but field exists for API consistency
+            "api_key": "sk-xxx",  # Your BochaAI API Key
             "max_results": 5,
             "reader": {  # Reader config for chunking web content
-                "backend": "simple",
-                "config": {},
+                "backend": "simple_struct",
+                "config": {  # your simple struct reader config
+                },
             },
         },
     }
@@ -43,21 +43,10 @@
 # ========= 4. Run BochaAI Web Search =========
 print("=== Scenario 1: Web Search (BochaAI) ===")
 query_web = "Alibaba 2024 ESG report"
-results_web = retriever.retrieve_from_web(query_web)
+results_web = retriever.retrieve_from_internet(query_web)
 
 print(f"Retrieved {len(results_web)} memory items.")
 for idx, item in enumerate(results_web, 1):
-    print(f"[{idx}] {item.memory[:100]}...")  # preview first 100 chars
-
-print("==" * 20)
-
-# ========= 5. Run BochaAI AI Search =========
-print("=== Scenario 2: AI Search (BochaAI) ===")
-query_ai = "Weather in Beijing"
-results_ai = retriever.retrieve_from_ai(query_ai)
-
-print(f"Retrieved {len(results_ai)} memory items.")
-for idx, item in enumerate(results_ai, 1):
-    print(f"[{idx}] {item.memory[:100]}...")
+    print(f"[{idx}] {item.memory[:500]}...")  # preview first 100 chars
 
 print("==" * 20)

src/memos/configs/internet_retriever.py

@@ -55,9 +55,17 @@ class XinyuSearchConfig(BaseInternetRetrieverConfig):
     )
 
 
-class BochaSearchConfig(XinyuSearchConfig):
+class BochaSearchConfig(BaseInternetRetrieverConfig):
     """Configuration class for Bocha Search API."""
 
+    max_results: int = Field(default=20, description="Maximum number of results to retrieve")
+    num_per_request: int = Field(default=10, description="Number of results per API request")
+    reader: MemReaderConfigFactory = Field(
+        ...,
+        default_factory=MemReaderConfigFactory,
+        description="Reader configuration",
+    )
+
 
 class InternetRetrieverConfigFactory(BaseConfig):
     """Factory class for creating internet retriever configurations."""

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

@@ -81,12 +81,15 @@ def search_ai(
         return self._post(self.ai_url, body)
 
     def _post(self, url: str, body: dict) -> list[dict]:
-        """Helper method to send POST request and return JSON results."""
+        """Send POST request and parse BochaAI search results."""
         try:
             resp = requests.post(url, headers=self.headers, json=body)
             resp.raise_for_status()
-            data = resp.json()
-            return data.get("results", [])
+            raw_data = resp.json()
+
+            # ✅ parse the nested structure correctly
+            return raw_data.get("data", {}).get("webPages", {}).get("value", [])
+
         except Exception:
             import traceback
 
@@ -100,7 +103,6 @@ class BochaAISearchRetriever:
     def __init__(
         self,
         access_key: str,
-        search_engine_id: str,
         embedder: OllamaEmbedder,
         reader: BaseMemReader,
         max_results: int = 20,
@@ -110,7 +112,6 @@ def __init__(
 
         Args:
             access_key: BochaAI API key
-            search_engine_id: (Not used for Bocha, but kept for API consistency)
             embedder: Embedder instance for generating embeddings
             reader: MemReader instance for processing internet content
             max_results: Maximum number of search results to retrieve
@@ -178,12 +179,22 @@ def _convert_to_mem_items(
     def _process_result(
         self, result: dict, query: str, parsed_goal: str, info: None
     ) -> list[TextualMemoryItem]:
-        """Process a single result into one or more TextualMemoryItems."""
-        title = result.get("title", "")
-        content = result.get("content", "")
-        summary = result.get("summary", "")
+        """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 = datetime.now().strftime("%Y-%m-%d")
+        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)

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

@@ -64,14 +64,21 @@ def from_config(
                 max_results=config.max_results,
                 num_per_request=config.num_per_request,
             )
-        elif backend == "xinyu" or backend == "bocha":
+        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,
                 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}")