diff --git a/examples/rag_pipeline.py b/examples/rag_pipeline.py index 7dbca35a6..75888e480 100644 --- a/examples/rag_pipeline.py +++ b/examples/rag_pipeline.py @@ -18,13 +18,13 @@ from metagpt.rag.schema import ( ) from metagpt.utils.exceptions import handle_exception +LLM_TIP = "If you not sure, just answer I don't know." + DOC_PATH = EXAMPLE_DATA_PATH / "rag/writer.txt" -QUESTION = "What are key qualities to be a good writer?" +QUESTION = f"What are key qualities to be a good writer? {LLM_TIP}" TRAVEL_DOC_PATH = EXAMPLE_DATA_PATH / "rag/travel.txt" -TRAVEL_QUESTION = "What does Bob like?" - -LLM_TIP = "If you not sure, just answer I don't know." +TRAVEL_QUESTION = f"What does Bob like? {LLM_TIP}" class Player(BaseModel): @@ -40,21 +40,21 @@ class Player(BaseModel): class RAGExample: - """Show how to use RAG. + """Show how to use RAG.""" - Default engine use LLM Reranker, if the answer from the LLM is incorrect, may encounter `IndexError: list index out of range`. - """ - - def __init__(self, engine: SimpleEngine = None): + def __init__(self, engine: SimpleEngine = None, use_llm_ranker: bool = False): self._engine = engine + self._use_llm_ranker = use_llm_ranker @property def engine(self): if not self._engine: + ranker_configs = [LLMRankerConfig()] if self._use_llm_ranker else None + self._engine = SimpleEngine.from_docs( input_files=[DOC_PATH], retriever_configs=[FAISSRetrieverConfig()], - ranker_configs=[LLMRankerConfig()], + ranker_configs=ranker_configs, ) return self._engine @@ -105,7 +105,7 @@ class RAGExample: """ self._print_title("Add Docs") - travel_question = f"{TRAVEL_QUESTION}{LLM_TIP}" + travel_question = f"{TRAVEL_QUESTION}" travel_filepath = TRAVEL_DOC_PATH logger.info("[Before add docs]") @@ -240,8 +240,14 @@ class RAGExample: async def main(): - """RAG pipeline.""" - e = RAGExample() + """RAG pipeline. + + Note: + 1. If `use_llm_ranker` is True, then will use LLM Reranker to get better result, but it is not always guaranteed that the output will be parseable for reranking, + prefer `gpt-4-turbo`, otherwise might encounter `IndexError: list index out of range` or `ValueError: invalid literal for int() with base 10`. + """ + e = RAGExample(use_llm_ranker=False) + await e.run_pipeline() await e.add_docs() await e.add_objects()