diff --git a/metagpt/rag/factories/ranker.py b/metagpt/rag/factories/ranker.py index 07cb1b929..476fe8c1a 100644 --- a/metagpt/rag/factories/ranker.py +++ b/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 @@ class RankerFactory(ConfigBasedFactory): 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: diff --git a/setup.py b/setup.py index 4fa5499da..3eab2b6a0 100644 --- a/setup.py +++ b/setup.py @@ -38,7 +38,6 @@ extras_require = { "llama-index-vector-stores-faiss==0.1.1", "llama-index-vector-stores-elasticsearch==0.1.6", "llama-index-vector-stores-chroma==0.1.6", - "llama-index-postprocessor-colbert-rerank==0.1.1", ], }