Merge pull request #1131 from seehi/pickup-rag-to-v0.8

Author: garylin2099

metagpt/environment/minecraft_env/minecraft_env.py

@@ -8,14 +8,14 @@
 import time
 from typing import Any, Iterable
 
+from llama_index.vector_stores.chroma import ChromaVectorStore
 from pydantic import ConfigDict, Field
 
 from metagpt.config2 import config as CONFIG
 from metagpt.environment.base_env import Environment
 from metagpt.environment.minecraft_env.const import MC_CKPT_DIR
 from metagpt.environment.minecraft_env.minecraft_ext_env import MinecraftExtEnv
 from metagpt.logs import logger
-from metagpt.rag.vector_stores.chroma import ChromaVectorStore
 from metagpt.utils.common import load_mc_skills_code, read_json_file, write_json_file
 
 

metagpt/rag/factories/index.py

@@ -5,6 +5,7 @@
 from llama_index.core.embeddings import BaseEmbedding
 from llama_index.core.indices.base import BaseIndex
 from llama_index.core.vector_stores.types import BasePydanticVectorStore
+from llama_index.vector_stores.chroma import ChromaVectorStore
 from llama_index.vector_stores.elasticsearch import ElasticsearchStore
 from llama_index.vector_stores.faiss import FaissVectorStore
 
@@ -17,7 +18,6 @@
     ElasticsearchKeywordIndexConfig,
     FAISSIndexConfig,
 )
-from metagpt.rag.vector_stores.chroma import ChromaVectorStore
 
 
 class RAGIndexFactory(ConfigBasedFactory):

metagpt/rag/factories/ranker.py

@@ -3,7 +3,6 @@
 from llama_index.core.llms import LLM
 from llama_index.core.postprocessor import LLMRerank
 from llama_index.core.postprocessor.types import BaseNodePostprocessor
-from llama_index.postprocessor.colbert_rerank import ColbertRerank
 
 from metagpt.rag.factories.base import ConfigBasedFactory
 from metagpt.rag.rankers.object_ranker import ObjectSortPostprocessor
@@ -38,6 +37,12 @@ def _create_llm_ranker(self, config: LLMRankerConfig, **kwargs) -> LLMRerank:
         return LLMRerank(**config.model_dump())
 
     def _create_colbert_ranker(self, config: ColbertRerankConfig, **kwargs) -> LLMRerank:
+        try:
+            from llama_index.postprocessor.colbert_rerank import ColbertRerank
+        except ImportError:
+            raise ImportError(
+                "`llama-index-postprocessor-colbert-rerank` package not found, please run `pip install llama-index-postprocessor-colbert-rerank`"
+            )
         return ColbertRerank(**config.model_dump())
 
     def _create_object_ranker(self, config: ObjectRankerConfig, **kwargs) -> LLMRerank:

metagpt/rag/factories/retriever.py

@@ -6,6 +6,7 @@
 import faiss
 from llama_index.core import StorageContext, VectorStoreIndex
 from llama_index.core.vector_stores.types import BasePydanticVectorStore
+from llama_index.vector_stores.chroma import ChromaVectorStore
 from llama_index.vector_stores.elasticsearch import ElasticsearchStore
 from llama_index.vector_stores.faiss import FaissVectorStore
 
@@ -25,7 +26,6 @@
     FAISSRetrieverConfig,
     IndexRetrieverConfig,
 )
-from metagpt.rag.vector_stores.chroma import ChromaVectorStore
 
 
 class RetrieverFactory(ConfigBasedFactory):

metagpt/rag/vector_stores/__init__.py

@@ -37,8 +37,7 @@ def run(self):
         "llama-index-retrievers-bm25==0.1.3",
         "llama-index-vector-stores-faiss==0.1.1",
         "llama-index-vector-stores-elasticsearch==0.1.6",
-        "llama-index-postprocessor-colbert-rerank==0.1.1",
-        "chromadb==0.4.23",
+        "llama-index-vector-stores-chroma==0.1.6",
     ],
 }