add BM25IndexConfig

metagpt/rag/engines/simple.py

@@ -1,7 +1,8 @@
 """Simple Engine."""
 
 import json
-from typing import Optional
+import os
+from typing import Optional, Union
 
 from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
 from llama_index.core.callbacks.base import CallbackManager
@@ -128,8 +129,8 @@ class SimpleEngine(RetrieverQueryEngine):
             retriever_configs: Configuration for retrievers. If more than one config, will use SimpleHybridRetriever.
             ranker_configs: Configuration for rankers.
         """
-        if not retriever_configs or any(isinstance(config, BM25RetrieverConfig) for config in retriever_configs):
-            raise ValueError("Must provide retriever_configs, and BM25RetrieverConfig is not supported.")
+        if not objs and any(isinstance(config, BM25RetrieverConfig) for config in retriever_configs):
+            raise ValueError("In BM25RetrieverConfig, Objs must not be empty.")
 
         objs = objs or []
         nodes = [ObjectNode(text=obj.rag_key(), metadata=ObjectNode.get_obj_metadata(obj)) for obj in objs]
@@ -182,11 +183,11 @@ class SimpleEngine(RetrieverQueryEngine):
         nodes = [ObjectNode(text=obj.rag_key(), metadata=ObjectNode.get_obj_metadata(obj)) for obj in objs]
         self._save_nodes(nodes)
 
-    def persist(self, persist_dir: str, **kwargs):
+    def persist(self, persist_dir: Union[str, os.PathLike], **kwargs):
         """Persist."""
         self._ensure_retriever_persistable()
 
-        self._persist(persist_dir, **kwargs)
+        self._persist(str(persist_dir), **kwargs)
 
     @classmethod
     def _from_index(

metagpt/rag/factories/index.py

@@ -7,7 +7,12 @@ from llama_index.core.indices.base import BaseIndex
 from llama_index.vector_stores.faiss import FaissVectorStore
 
 from metagpt.rag.factories.base import ConfigFactory
-from metagpt.rag.schema import BaseIndexConfig, ChromaIndexConfig, FAISSIndexConfig
+from metagpt.rag.schema import (
+    BaseIndexConfig,
+    BM25IndexConfig,
+    ChromaIndexConfig,
+    FAISSIndexConfig,
+)
 from metagpt.rag.vector_stores.chroma import ChromaVectorStore
 
 
@@ -16,6 +21,7 @@ class RAGIndexFactory(ConfigFactory):
         creators = {
             FAISSIndexConfig: self._create_faiss,
             ChromaIndexConfig: self._create_chroma,
+            BM25IndexConfig: self._create_bm25,
         }
         super().__init__(creators)
 
@@ -46,5 +52,12 @@ class RAGIndexFactory(ConfigFactory):
         )
         return index
 
+    def _create_bm25(self, config: BM25IndexConfig, **kwargs) -> VectorStoreIndex:
+        embed_model = self._extract_embed_model(config, **kwargs)
+
+        storage_context = StorageContext.from_defaults(persist_dir=config.persist_path)
+        index = load_index_from_storage(storage_context=storage_context, embed_model=embed_model)
+        return index
+
 
 get_index = RAGIndexFactory().get_index

metagpt/rag/factories/retriever.py

@@ -1,5 +1,7 @@
 """RAG Retriever Factory."""
 
+import copy
+
 import chromadb
 import faiss
 from llama_index.core import StorageContext, VectorStoreIndex
@@ -69,8 +71,9 @@ class RetrieverFactory(ConfigFactory):
         return FAISSRetriever(**config.model_dump())
 
     def _create_bm25_retriever(self, config: BM25RetrieverConfig, **kwargs) -> DynamicBM25Retriever:
-        config.index = self._extract_index(config, **kwargs)
-        return DynamicBM25Retriever.from_defaults(**config.model_dump())
+        config.index = copy.deepcopy(self._extract_index(config, **kwargs))
+        nodes = list(config.index.docstore.docs.values())
+        return DynamicBM25Retriever(nodes=nodes, **config.model_dump())
 
     def _create_chroma_retriever(self, config: ChromaRetrieverConfig, **kwargs) -> ChromaRetriever:
         db = chromadb.PersistentClient(path=str(config.persist_path))

metagpt/rag/retrievers/bm25_retriever.py

@@ -1,6 +1,10 @@
 """BM25 retriever."""
+from typing import Callable, Optional
 
-from llama_index.core.schema import BaseNode
+from llama_index.core import VectorStoreIndex
+from llama_index.core.callbacks.base import CallbackManager
+from llama_index.core.constants import DEFAULT_SIMILARITY_TOP_K
+from llama_index.core.schema import BaseNode, IndexNode
 from llama_index.retrievers.bm25 import BM25Retriever
 from rank_bm25 import BM25Okapi
 
@@ -8,8 +12,36 @@ from rank_bm25 import BM25Okapi
 class DynamicBM25Retriever(BM25Retriever):
     """BM25 retriever."""
 
+    def __init__(
+        self,
+        nodes: list[BaseNode],
+        tokenizer: Optional[Callable[[str], list[str]]] = None,
+        similarity_top_k: int = DEFAULT_SIMILARITY_TOP_K,
+        callback_manager: Optional[CallbackManager] = None,
+        objects: Optional[list[IndexNode]] = None,
+        object_map: Optional[dict] = None,
+        verbose: bool = False,
+        index: VectorStoreIndex = None,
+    ) -> None:
+        super().__init__(
+            nodes=nodes,
+            tokenizer=tokenizer,
+            similarity_top_k=similarity_top_k,
+            callback_manager=callback_manager,
+            object_map=object_map,
+            objects=objects,
+            verbose=verbose,
+        )
+        self._index = index
+
     def add_nodes(self, nodes: list[BaseNode], **kwargs) -> None:
-        """Support add nodes"""
+        """Support add nodes."""
         self._nodes.extend(nodes)
         self._corpus = [self._tokenizer(node.get_content()) for node in self._nodes]
         self.bm25 = BM25Okapi(self._corpus)
+
+        self._index.insert_nodes(nodes, **kwargs)
+
+    def persist(self, persist_dir: str, **kwargs) -> None:
+        """Support persist."""
+        self._index.storage_context.persist(persist_dir)

metagpt/rag/schema.py

@@ -89,6 +89,10 @@ class ChromaIndexConfig(VectorIndexConfig):
     collection_name: str = Field(default="metagpt", description="The name of the collection.")
 
 
+class BM25IndexConfig(BaseIndexConfig):
+    """Config for bm25-based index."""
+
+
 class ObjectNodeMetadata(BaseModel):
     """Metadata of ObjectNode."""