fix: nebula multi-embedding & add BochaAI Search Retriever (#195)

* fix: fine-search bug

* feat: multi embedding dimension in nebula

* feat: modify multi-dimension

* fix: multi-embedding error

* feat: update nebula example

* fix: print bug

* feat: add bocha search

* feat: modify bocha search

* feat: finish bocha search

Author: CaralHsi

examples/basic_modules/nebular_example.py

@@ -23,7 +23,7 @@ def show(nebular_data):
 
     tree_config = Neo4jGraphDBConfig.from_json_file("../../examples/data/config/neo4j_config.json")
     tree_config.use_multi_db = True
-    tree_config.db_name = "nebular-show"
+    tree_config.db_name = "nebular-show2"
 
     neo4j_db = Neo4jGraphDB(tree_config)
     neo4j_db.clear()
@@ -42,6 +42,7 @@ def show(nebular_data):
     }
 )
 embedder = EmbedderFactory.from_config(embedder_config)
+embedder_dimension = 3072
 
 
 def embed_memory_item(memory: str) -> list[float]:
@@ -62,7 +63,7 @@ def example_multi_db(db_name: str = "paper"):
             "space": db_name,
             "use_multi_db": True,
             "auto_create": True,
-            "embedding_dimension": 3072,
+            "embedding_dimension": embedder_dimension,
         },
     )
 
@@ -121,7 +122,7 @@ def example_shared_db(db_name: str = "shared-traval-group"):
                 "user_name": user_name,
                 "use_multi_db": False,
                 "auto_create": True,
-                "embedding_dimension": 3072,
+                "embedding_dimension": embedder_dimension,
             },
         )
 
@@ -208,7 +209,7 @@ def example_shared_db(db_name: str = "shared-traval-group"):
             "space": db_name,
             "user_name": user_list[0],
             "auto_create": True,
-            "embedding_dimension": 3072,
+            "embedding_dimension": embedder_dimension,
             "use_multi_db": False,
         },
     )
@@ -238,7 +239,7 @@ def run_user_session(
             "user_name": user_name,
             "use_multi_db": False,
             "auto_create": True,
-            "embedding_dimension": 3072,
+            "embedding_dimension": embedder_dimension,
         },
     )
     graph = GraphStoreFactory.from_config(config)
@@ -404,10 +405,10 @@ def example_complex_shared_db(db_name: str = "shared-traval-group-complex"):
 
 if __name__ == "__main__":
     print("\n=== Example: Multi-DB ===")
-    example_multi_db(db_name="paper")
+    example_multi_db(db_name="paper-new")
 
     print("\n=== Example: Single-DB ===")
-    example_shared_db(db_name="shared_traval_group")
+    example_shared_db(db_name="shared_traval_group-new")
 
     print("\n=== Example: Single-DB-Complex ===")
-    example_complex_shared_db(db_name="shared-traval-group-complex-new11")
+    example_complex_shared_db(db_name="shared-traval-group-complex-new2")

examples/basic_modules/textual_memory_internet_search_example.py

@@ -0,0 +1,52 @@
+"""
+Example: Using InternetRetrieverFactory with BochaAISearchRetriever
+"""
+
+from memos.configs.embedder import EmbedderConfigFactory
+from memos.configs.internet_retriever import InternetRetrieverConfigFactory
+from memos.embedders.factory import EmbedderFactory
+from memos.memories.textual.tree_text_memory.retrieve.internet_retriever_factory import (
+    InternetRetrieverFactory,
+)
+
+
+# ========= 1. Create an embedder =========
+embedder_config = EmbedderConfigFactory.model_validate(
+    {
+        "backend": "ollama",  # Or "sentence_transformer", etc.
+        "config": {
+            "model_name_or_path": "nomic-embed-text:latest",
+        },
+    }
+)
+embedder = EmbedderFactory.from_config(embedder_config)
+
+# ========= 2. Create retriever config for BochaAI =========
+retriever_config = InternetRetrieverConfigFactory.model_validate(
+    {
+        "backend": "bocha",
+        "config": {
+            "api_key": "sk-xxx",  # Your BochaAI API Key
+            "max_results": 5,
+            "reader": {  # Reader config for chunking web content
+                "backend": "simple_struct",
+                "config": {  # your simple struct reader config
+                },
+            },
+        },
+    }
+)
+
+# ========= 3. Build retriever instance via factory =========
+retriever = InternetRetrieverFactory.from_config(retriever_config, embedder)
+
+# ========= 4. Run BochaAI Web Search =========
+print("=== Scenario 1: Web Search (BochaAI) ===")
+query_web = "Alibaba 2024 ESG report"
+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[:500]}...")  # preview first 100 chars
+
+print("==" * 20)

examples/core_memories/tree_textual_memory.py

@@ -1,32 +1,15 @@
 import time
 
 from memos import log
-from memos.configs.embedder import EmbedderConfigFactory
 from memos.configs.mem_reader import SimpleStructMemReaderConfig
 from memos.configs.memory import TreeTextMemoryConfig
-from memos.embedders.factory import EmbedderFactory
 from memos.mem_reader.simple_struct import SimpleStructMemReader
 from memos.memories.textual.tree import TreeTextMemory
 
 
 logger = log.get_logger(__name__)
 
 
-embedder_config = EmbedderConfigFactory.model_validate(
-    {
-        "backend": "ollama",
-        "config": {
-            "model_name_or_path": "nomic-embed-text:latest",
-        },
-    }
-)
-embedder = EmbedderFactory.from_config(embedder_config)
-
-
-def embed_memory_item(memory: str) -> list[float]:
-    return embedder.embed([memory])[0]
-
-
 tree_config = TreeTextMemoryConfig.from_json_file(
     "examples/data/config/tree_config_shared_database.json"
 )

src/memos/configs/internet_retriever.py

@@ -55,6 +55,18 @@ class XinyuSearchConfig(BaseInternetRetrieverConfig):
     )
 
 
+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."""
 
@@ -69,6 +81,7 @@ class InternetRetrieverConfigFactory(BaseConfig):
         "google": GoogleCustomSearchConfig,
         "bing": BingSearchConfig,
         "xinyu": XinyuSearchConfig,
+        "bocha": BochaSearchConfig,
     }
 
     @field_validator("backend")