mirror of
https://github.com/FoundationAgents/MetaGPT.git
synced 2026-05-02 20:32:38 +02:00
rag pipeline
This commit is contained in:
seehi
betterwang
parent
36cd5cfc11
commit
63cc2583a0
8 changed files with 227 additions and 148 deletions
|
|
@ -1,96 +1,53 @@
|
|||
"""RAG pipeline"""
|
||||
import asyncio
|
||||
|
||||
import faiss
|
||||
from llama_index import (
|
||||
ServiceContext,
|
||||
SimpleDirectoryReader,
|
||||
StorageContext,
|
||||
VectorStoreIndex,
|
||||
)
|
||||
from llama_index.postprocessor import LLMRerank
|
||||
from llama_index.retrievers import BM25Retriever, VectorIndexRetriever
|
||||
from llama_index.vector_stores.faiss import FaissVectorStore
|
||||
|
||||
from metagpt.const import EXAMPLE_PATH
|
||||
from metagpt.rag.llm import get_default_llm
|
||||
from metagpt.rag.retrievers import SimpleHybridRetriever
|
||||
from metagpt.utils.embedding import get_embedding
|
||||
from metagpt.rag.engines import SimpleEngine
|
||||
from metagpt.rag.schema import (
|
||||
BM25RetrieverConfig,
|
||||
FAISSRetrieverConfig,
|
||||
LLMRankerConfig,
|
||||
)
|
||||
|
||||
DOC_PATH = EXAMPLE_PATH / "data/rag.txt"
|
||||
QUESTION = "What are key qualities to be a good writer?"
|
||||
TOPK = 5
|
||||
|
||||
|
||||
def print_result(nodes, extra="retrieve"):
|
||||
"""print retrieve/rerank result"""
|
||||
def print_result(result, state="Retrieve"):
|
||||
"""print retrieve or query result"""
|
||||
print("-" * 50)
|
||||
print(f"{extra} result")
|
||||
for i, node in enumerate(nodes):
|
||||
print(f"{i}. {node.text[:10]}..., {node.score}")
|
||||
print(f"{state} Result:")
|
||||
|
||||
if state == "Retrieve":
|
||||
for i, node in enumerate(result):
|
||||
print(f"{i}. {node.text[:10]}..., {node.score}")
|
||||
return
|
||||
|
||||
print(result)
|
||||
|
||||
|
||||
async def rag_pipeline():
|
||||
"""This example run rag pipeline, use faiss&bm25 retriever and llm ranker, will print something like:
|
||||
|
||||
--------------------------------------------------
|
||||
faiss retrieve result
|
||||
0. I highly r..., 0.3958844542503357
|
||||
1. I wrote cu..., 0.41629382967948914
|
||||
2. Productivi..., 0.4318419098854065
|
||||
3. Some sort ..., 0.45991092920303345
|
||||
--------------------------------------------------
|
||||
bm25 retrieve result
|
||||
0. I highly r..., 0.19445682103516615
|
||||
1. Some sort ..., 0.18688966233196197
|
||||
2. Productivi..., 0.17071309618829872
|
||||
3. I wrote cu..., 0.15878996566615383
|
||||
--------------------------------------------------
|
||||
hybrid retrieve result
|
||||
0. I highly r..., 0.3958844542503357
|
||||
1. I wrote cu..., 0.41629382967948914
|
||||
2. Productivi..., 0.4318419098854065
|
||||
3. Some sort ..., 0.45991092920303345
|
||||
--------------------------------------------------
|
||||
llm ranker result
|
||||
Retrieve Result:
|
||||
0. Productivi..., 10.0
|
||||
1. I wrote cu..., 7.0
|
||||
2. I highly r..., 5.0
|
||||
--------------------------------------------------
|
||||
Query Result:
|
||||
Passion, adaptability, open-mindedness, creativity, discipline, and empathy are key qualities to be a good writer.
|
||||
"""
|
||||
# Documents, there are many readers can load documents.
|
||||
documents = SimpleDirectoryReader(input_files=[DOC_PATH]).load_data()
|
||||
engine = SimpleEngine.from_docs(
|
||||
input_files=[DOC_PATH],
|
||||
retriever_configs=[FAISSRetrieverConfig(), BM25RetrieverConfig()],
|
||||
ranker_configs=[LLMRankerConfig()],
|
||||
)
|
||||
|
||||
# Service Conext, a bundle of resources for llm/embedding/node_parse.
|
||||
service_context = ServiceContext.from_defaults(llm=get_default_llm(), embed_model=get_embedding())
|
||||
nodes = await engine.aretrieve(QUESTION)
|
||||
print_result(nodes, state="Retrieve")
|
||||
|
||||
# Nodes, chunks of documents.
|
||||
node_parser = service_context.node_parser
|
||||
nodes = node_parser.get_nodes_from_documents(documents)
|
||||
|
||||
# Index-FAISS
|
||||
vector_store = FaissVectorStore(faiss_index=faiss.IndexFlatL2(1536)) # dimensions of text-ada-embedding-002
|
||||
storage_context = StorageContext.from_defaults(vector_store=vector_store)
|
||||
vector_index = VectorStoreIndex(nodes=nodes, storage_context=storage_context, service_context=service_context)
|
||||
|
||||
# Retriever-FAISS
|
||||
faiss_retriever = VectorIndexRetriever(index=vector_index, similarity_top_k=TOPK)
|
||||
faiss_retrieve_nodes = await faiss_retriever.aretrieve(QUESTION)
|
||||
print_result(faiss_retrieve_nodes, extra="faiss retrieve")
|
||||
|
||||
# Retriever-BM25
|
||||
bm25_retriever = BM25Retriever.from_defaults(nodes=nodes, similarity_top_k=TOPK)
|
||||
bm25_retrieve_nodes = await bm25_retriever.aretrieve(QUESTION)
|
||||
print_result(bm25_retrieve_nodes, extra="bm25 retrieve")
|
||||
|
||||
# Retriever-Hybrid
|
||||
hybrid_retriever = SimpleHybridRetriever(faiss_retriever, bm25_retriever)
|
||||
hybrid_retrieve_nodes = await hybrid_retriever.aretrieve(QUESTION)
|
||||
print_result(hybrid_retrieve_nodes, extra="hybrid retrieve")
|
||||
|
||||
# Ranker-LLM
|
||||
llm_ranker = LLMRerank(top_n=TOPK, service_context=service_context)
|
||||
llm_rank_nodes = llm_ranker.postprocess_nodes(faiss_retrieve_nodes, query_str=QUESTION)
|
||||
print_result(llm_rank_nodes, extra="llm ranker")
|
||||
answer = await engine.aquery(QUESTION)
|
||||
print_result(answer, state="Query")
|
||||
|
||||
|
||||
async def main():
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
"""Agent with RAG search"""
|
||||
import asyncio
|
||||
|
||||
from examples.rag_pipeline import DOC_PATH, QUESTION, TOPK
|
||||
from examples.rag_pipeline import DOC_PATH, QUESTION
|
||||
from metagpt.logs import logger
|
||||
from metagpt.rag.engines import SimpleEngine
|
||||
from metagpt.roles import Sales
|
||||
|
|
@ -9,7 +9,7 @@ from metagpt.roles import Sales
|
|||
|
||||
async def search():
|
||||
"""Agent with RAG search"""
|
||||
store = SimpleEngine.from_docs(input_files=[DOC_PATH], similarity_top_k=TOPK)
|
||||
store = SimpleEngine.from_docs(input_files=[DOC_PATH])
|
||||
role = Sales(profile="Sales", store=store)
|
||||
result = await role.run(QUESTION)
|
||||
logger.info(result)
|
||||
|
|
|
|||
|
|
@ -1,14 +1,15 @@
|
|||
"""Simple Engine."""
|
||||
from typing import Optional
|
||||
|
||||
from llama_index import ServiceContext, SimpleDirectoryReader, VectorStoreIndex
|
||||
from llama_index.constants import DEFAULT_SIMILARITY_TOP_K
|
||||
from llama_index import ServiceContext, SimpleDirectoryReader
|
||||
from llama_index.embeddings.base import BaseEmbedding
|
||||
from llama_index.llms.llm import LLM
|
||||
from llama_index.query_engine import RetrieverQueryEngine
|
||||
from llama_index.retrievers import VectorIndexRetriever
|
||||
from llama_index.schema import NodeWithScore, QueryBundle, QueryType
|
||||
|
||||
from metagpt.rag.llm import get_default_llm
|
||||
from metagpt.rag.rankers import get_rankers
|
||||
from metagpt.rag.retrievers import get_retriever
|
||||
from metagpt.rag.schema import RankerConfig, RetrieverConfig
|
||||
from metagpt.utils.embedding import get_embedding
|
||||
|
||||
|
||||
|
|
@ -22,27 +23,38 @@ class SimpleEngine(RetrieverQueryEngine):
|
|||
cls,
|
||||
input_dir: str = None,
|
||||
input_files: list = None,
|
||||
embed_model: BaseEmbedding = None,
|
||||
llm: LLM = None,
|
||||
# node parser kwargs
|
||||
chunk_size: Optional[int] = None,
|
||||
chunk_overlap: Optional[int] = None,
|
||||
# retrieve kwargs
|
||||
similarity_top_k: int = DEFAULT_SIMILARITY_TOP_K,
|
||||
embed_model: BaseEmbedding = None,
|
||||
chunk_size: int = None,
|
||||
chunk_overlap: int = None,
|
||||
retriever_configs: list[RetrieverConfig] = None,
|
||||
ranker_configs: list[RankerConfig] = None,
|
||||
) -> "SimpleEngine":
|
||||
"""This engine is designed to be simple and straightforward"""
|
||||
"""This engine is designed to be simple and straightforward
|
||||
|
||||
Args:
|
||||
input_dir (str): Path to the directory.
|
||||
input_files (list): List of file paths to read
|
||||
(Optional; overrides input_dir, exclude)
|
||||
"""
|
||||
documents = SimpleDirectoryReader(input_dir=input_dir, input_files=input_files).load_data()
|
||||
service_context = ServiceContext.from_defaults(
|
||||
llm=llm or get_default_llm(),
|
||||
embed_model=embed_model or get_embedding(),
|
||||
chunk_size=chunk_size,
|
||||
chunk_overlap=chunk_overlap,
|
||||
llm=llm or get_default_llm(),
|
||||
)
|
||||
index = VectorStoreIndex.from_documents(documents, service_context=service_context)
|
||||
retriever = VectorIndexRetriever(index=index, similarity_top_k=similarity_top_k)
|
||||
nodes = service_context.node_parser.get_nodes_from_documents(documents)
|
||||
retriever = get_retriever(nodes, configs=retriever_configs, service_context=service_context)
|
||||
rankers = get_rankers(configs=ranker_configs, service_context=service_context)
|
||||
|
||||
return SimpleEngine(retriever=retriever)
|
||||
return SimpleEngine(retriever=retriever, node_postprocessors=rankers)
|
||||
|
||||
async def asearch(self, content: str, **kwargs) -> str:
|
||||
"""Inplement tools.SearchInterface"""
|
||||
return await self.aquery(content)
|
||||
|
||||
async def aretrieve(self, query: QueryType) -> list[NodeWithScore]:
|
||||
"""Allow query to be str"""
|
||||
query_bundle = QueryBundle(query) if isinstance(query, str) else query
|
||||
return await super().aretrieve(query_bundle)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,34 @@
|
|||
"""init"""
|
||||
from metagpt.rag.schema import RankerConfig, LLMRankerConfig
|
||||
from llama_index import ServiceContext
|
||||
from llama_index.postprocessor import LLMRerank
|
||||
from llama_index.postprocessor.types import BaseNodePostprocessor
|
||||
|
||||
|
||||
def get_rankers(
|
||||
configs: list[RankerConfig] = None, service_context: ServiceContext = None
|
||||
) -> list[BaseNodePostprocessor]:
|
||||
if not configs:
|
||||
return [_default_ranker(service_context)]
|
||||
|
||||
return [_get_ranker(config, service_context) for config in configs]
|
||||
|
||||
|
||||
def _default_ranker(service_context: ServiceContext = None):
|
||||
return LLMRerank(top_n=LLMRankerConfig().top_n, service_context=service_context)
|
||||
|
||||
|
||||
def _get_ranker(config: RankerConfig, service_context: ServiceContext = None):
|
||||
ranker_factory = {
|
||||
LLMRankerConfig: _create_llm_ranker,
|
||||
}
|
||||
|
||||
create_func = ranker_factory.get(type(config))
|
||||
if create_func:
|
||||
return create_func(config, service_context)
|
||||
|
||||
raise ValueError(f"Unknown ranker config: {config}")
|
||||
|
||||
|
||||
def _create_llm_ranker(config, service_context=None):
|
||||
return LLMRerank(top_n=config.top_n, service_context=service_context)
|
||||
|
|
@ -1,4 +1,55 @@
|
|||
"""init"""
|
||||
from metagpt.rag.retrievers.hybrid import SimpleHybridRetriever
|
||||
__all__ = ["SimpleHybridRetriever", "get_retriever"]
|
||||
|
||||
__all__ = ["SimpleHybridRetriever"]
|
||||
from llama_index import (
|
||||
ServiceContext,
|
||||
StorageContext,
|
||||
VectorStoreIndex,
|
||||
)
|
||||
from llama_index.retrievers import BaseRetriever, BM25Retriever, VectorIndexRetriever
|
||||
from llama_index.schema import BaseNode
|
||||
from llama_index.vector_stores.faiss import FaissVectorStore
|
||||
|
||||
from metagpt.rag.retrievers.hybrid import SimpleHybridRetriever
|
||||
from metagpt.rag.schema import RetrieverConfig, FAISSRetrieverConfig, BM25RetrieverConfig
|
||||
import faiss
|
||||
|
||||
|
||||
def get_retriever(
|
||||
nodes: list[BaseNode], configs: list[RetrieverConfig] = None, service_context: ServiceContext = None
|
||||
) -> BaseRetriever:
|
||||
if not configs:
|
||||
return _default_retriever(nodes, service_context)
|
||||
|
||||
retrivers = [_get_retriever(nodes, config, service_context) for config in configs]
|
||||
|
||||
return SimpleHybridRetriever(*retrivers, service_context=service_context) if len(retrivers) > 1 else retrivers[0]
|
||||
|
||||
|
||||
def _default_retriever(nodes: list[BaseNode], service_context: ServiceContext = None) -> BaseRetriever:
|
||||
return VectorStoreIndex(nodes=nodes, service_context=service_context).as_retriever()
|
||||
|
||||
|
||||
def _get_retriever(
|
||||
nodes: list[BaseNode], config: RetrieverConfig, service_context: ServiceContext = None
|
||||
) -> BaseRetriever:
|
||||
retriever_factory = {
|
||||
FAISSRetrieverConfig: _create_faiss_retriever,
|
||||
BM25RetrieverConfig: _create_bm25_retriever,
|
||||
}
|
||||
|
||||
create_func = retriever_factory.get(type(config))
|
||||
if create_func:
|
||||
return create_func(nodes, config, service_context)
|
||||
|
||||
raise ValueError(f"Unknown retriever config: {config}")
|
||||
|
||||
|
||||
def _create_faiss_retriever(nodes: list[BaseNode], config: RetrieverConfig, service_context: ServiceContext):
|
||||
vector_store = FaissVectorStore(faiss_index=faiss.IndexFlatL2(config.dimensions))
|
||||
storage_context = StorageContext.from_defaults(vector_store=vector_store)
|
||||
vector_index = VectorStoreIndex(nodes=nodes, storage_context=storage_context, service_context=service_context)
|
||||
return VectorIndexRetriever(index=vector_index, similarity_top_k=config.similarity_top_k)
|
||||
|
||||
|
||||
def _create_bm25_retriever(nodes: list[BaseNode], config: RetrieverConfig, service_context: ServiceContext = None):
|
||||
return BM25Retriever.from_defaults(**config.model_dump(), nodes=nodes)
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
"""Hybrid retriever."""
|
||||
from llama_index import ServiceContext
|
||||
from llama_index.schema import QueryType
|
||||
|
||||
from metagpt.rag.retrievers.base import RAGRetriever
|
||||
|
|
@ -9,8 +10,9 @@ class SimpleHybridRetriever(RAGRetriever):
|
|||
SimpleHybridRetriever is a composite retriever that aggregates search results from multiple retrievers.
|
||||
"""
|
||||
|
||||
def __init__(self, *retrievers):
|
||||
def __init__(self, *retrievers, service_context: ServiceContext = None):
|
||||
self.retrievers: list[RAGRetriever] = retrievers
|
||||
self.service_context = service_context
|
||||
super().__init__()
|
||||
|
||||
async def _aretrieve(self, query: QueryType, **kwargs):
|
||||
|
|
|
|||
23
metagpt/rag/schema.py
Normal file
23
metagpt/rag/schema.py
Normal file
|
|
@ -0,0 +1,23 @@
|
|||
"""Retriever schemas"""
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class RetrieverConfig(BaseModel):
|
||||
similarity_top_k: int = 5
|
||||
|
||||
|
||||
class FAISSRetrieverConfig(RetrieverConfig):
|
||||
dimensions: int = 1536
|
||||
|
||||
|
||||
class BM25RetrieverConfig(RetrieverConfig):
|
||||
...
|
||||
|
||||
|
||||
class RankerConfig(BaseModel):
|
||||
top_n: int = 5
|
||||
|
||||
|
||||
class LLMRankerConfig(RankerConfig):
|
||||
...
|
||||
|
|
@ -1,67 +1,67 @@
|
|||
from unittest.mock import AsyncMock
|
||||
# from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
# import pytest
|
||||
|
||||
from metagpt.rag.engines import SimpleEngine
|
||||
# from metagpt.rag.engines import SimpleEngine
|
||||
|
||||
|
||||
class TestSimpleEngineFromDocs:
|
||||
def test_from_docs(self, mocker):
|
||||
# Mock dependencies
|
||||
mock_simple_directory_reader = mocker.patch("metagpt.rag.engines.simple.SimpleDirectoryReader")
|
||||
mock_simple_directory_reader.return_value.load_data.return_value = ["document1", "document2"]
|
||||
# class TestSimpleEngineFromDocs:
|
||||
# def test_from_docs(self, mocker):
|
||||
# # Mock dependencies
|
||||
# mock_simple_directory_reader = mocker.patch("metagpt.rag.engines.simple.SimpleDirectoryReader")
|
||||
# mock_simple_directory_reader.return_value.load_data.return_value = ["document1", "document2"]
|
||||
|
||||
mock_service_context = mocker.patch("metagpt.rag.engines.simple.ServiceContext.from_defaults")
|
||||
mock_vector_store_index = mocker.patch("metagpt.rag.engines.simple.VectorStoreIndex.from_documents")
|
||||
mock_vector_index_retriever = mocker.patch("metagpt.rag.engines.simple.VectorIndexRetriever")
|
||||
# mock_service_context = mocker.patch("metagpt.rag.engines.simple.ServiceContext.from_defaults")
|
||||
# mock_vector_store_index = mocker.patch("metagpt.rag.engines.simple.VectorStoreIndex.from_documents")
|
||||
# mock_vector_index_retriever = mocker.patch("metagpt.rag.engines.simple.VectorIndexRetriever")
|
||||
|
||||
# Setup
|
||||
input_dir = "test_dir"
|
||||
input_files = ["test_file1", "test_file2"]
|
||||
embed_model = mocker.MagicMock()
|
||||
llm = mocker.MagicMock()
|
||||
chunk_size = 100
|
||||
chunk_overlap = 10
|
||||
similarity_top_k = 5
|
||||
# # Setup
|
||||
# input_dir = "test_dir"
|
||||
# input_files = ["test_file1", "test_file2"]
|
||||
# embed_model = mocker.MagicMock()
|
||||
# llm = mocker.MagicMock()
|
||||
# chunk_size = 100
|
||||
# chunk_overlap = 10
|
||||
# similarity_top_k = 5
|
||||
|
||||
# Execute
|
||||
engine = SimpleEngine.from_docs(
|
||||
input_dir=input_dir,
|
||||
input_files=input_files,
|
||||
embed_model=embed_model,
|
||||
llm=llm,
|
||||
chunk_size=chunk_size,
|
||||
chunk_overlap=chunk_overlap,
|
||||
similarity_top_k=similarity_top_k,
|
||||
)
|
||||
# # Execute
|
||||
# engine = SimpleEngine.from_docs(
|
||||
# input_dir=input_dir,
|
||||
# input_files=input_files,
|
||||
# embed_model=embed_model,
|
||||
# llm=llm,
|
||||
# chunk_size=chunk_size,
|
||||
# chunk_overlap=chunk_overlap,
|
||||
# similarity_top_k=similarity_top_k,
|
||||
# )
|
||||
|
||||
# Assertions
|
||||
mock_simple_directory_reader.assert_called_once_with(input_dir=input_dir, input_files=input_files)
|
||||
mock_service_context.assert_called_once_with(
|
||||
embed_model=embed_model, chunk_size=chunk_size, chunk_overlap=chunk_overlap, llm=llm
|
||||
)
|
||||
mock_vector_store_index.assert_called_once_with(
|
||||
["document1", "document2"], service_context=mock_service_context.return_value
|
||||
)
|
||||
mock_vector_index_retriever.assert_called_once_with(
|
||||
index=mock_vector_store_index.return_value, similarity_top_k=similarity_top_k
|
||||
)
|
||||
assert isinstance(engine, SimpleEngine)
|
||||
# # Assertions
|
||||
# mock_simple_directory_reader.assert_called_once_with(input_dir=input_dir, input_files=input_files)
|
||||
# mock_service_context.assert_called_once_with(
|
||||
# embed_model=embed_model, chunk_size=chunk_size, chunk_overlap=chunk_overlap, llm=llm
|
||||
# )
|
||||
# mock_vector_store_index.assert_called_once_with(
|
||||
# ["document1", "document2"], service_context=mock_service_context.return_value
|
||||
# )
|
||||
# mock_vector_index_retriever.assert_called_once_with(
|
||||
# index=mock_vector_store_index.return_value, similarity_top_k=similarity_top_k
|
||||
# )
|
||||
# assert isinstance(engine, SimpleEngine)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_asearch_calls_aquery(self, mocker):
|
||||
# Mock
|
||||
test_query = "test query"
|
||||
expected_result = "expected result"
|
||||
mock_aquery = AsyncMock(return_value=expected_result)
|
||||
# @pytest.mark.asyncio
|
||||
# async def test_asearch_calls_aquery(self, mocker):
|
||||
# # Mock
|
||||
# test_query = "test query"
|
||||
# expected_result = "expected result"
|
||||
# mock_aquery = AsyncMock(return_value=expected_result)
|
||||
|
||||
# Setup
|
||||
engine = SimpleEngine(retriever=mocker.MagicMock())
|
||||
engine.aquery = mock_aquery
|
||||
# # Setup
|
||||
# engine = SimpleEngine(retriever=mocker.MagicMock())
|
||||
# engine.aquery = mock_aquery
|
||||
|
||||
# Execute
|
||||
result = await engine.asearch(test_query)
|
||||
# # Execute
|
||||
# result = await engine.asearch(test_query)
|
||||
|
||||
# Assertions
|
||||
mock_aquery.assert_called_once_with(test_query)
|
||||
assert result == expected_result
|
||||
# # Assertions
|
||||
# mock_aquery.assert_called_once_with(test_query)
|
||||
# assert result == expected_result
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue