class LongRAGPack(BaseLlamaPack):
"""
Implements Long RAG.
This implementation is based on the following paper: https://arxiv.org/pdf/2406.15319
"""
def __init__(
self,
data_dir: str,
llm: t.Optional[LLM] = None,
chunk_size: t.Optional[int] = DEFAULT_CHUNK_SIZE,
similarity_top_k: int = DEFAULT_TOP_K,
small_chunk_size: int = DEFAULT_SMALL_CHUNK_SIZE,
index: t.Optional[VectorStoreIndex] = None,
index_kwargs: t.Optional[t.Dict[str, t.Any]] = None,
verbose: bool = False,
):
"""
Constructor.
Args:
data_dir (str): Data directory
llm (t.Optional[LLM]): LLM
chunk_size (Optional[int], optional): Splits each doc to chunk_size to demonstrate grouping. Set to None to disable splitting then grouping. Defaults to DEFAULT_CHUNK_SIZE.
similarity_top_k (int, optional): Top k. Defaults to DEFAULT_TOP_K.
small_chunk_size (int, optional): Small chunk size to split large documents into smaller embeddings of small_chunk_size. Defaults to DEFAULT_SMALL_CHUNK_SIZE.
index (Optional[VectorStoreIndex], optional): Vector index to use (from persist dir). If None, creates a new vector index. Defaults to None
index_kwargs (Optional[Dict[str, Any]], optional): Kwargs to use when constructing VectorStoreIndex. Defaults to None.
verbose (bool, Optional): Verbose mode. Defaults to False
"""
# initialize workflow
self._wf = LongRAGWorkflow(verbose=verbose)
# initialize vars
self._data_dir = data_dir
self._llm = llm or Settings.llm
self._chunk_size = chunk_size
self._similarity_top_k = similarity_top_k
self._small_chunk_size = small_chunk_size
# run wf initialization
result = asyncio_run(
self._wf.run(
data_dir=self._data_dir,
llm=self._llm,
chunk_size=self._chunk_size,
similarity_top_k=self._similarity_top_k,
small_chunk_size=self._small_chunk_size,
index=index,
index_kwargs=index_kwargs,
)
)
self._retriever = result["retriever"]
self._query_eng = result["query_engine"]
self._index = result["index"]
def get_modules(self) -> t.Dict[str, t.Any]:
"""Get Modules."""
return {
"query_engine": self._query_eng,
"llm": self._llm,
"retriever": self._retriever,
"index": self._index,
"workflow": self._wf,
}
def run(self, query: str, *args: t.Any, **kwargs: t.Any) -> t.Any:
"""Runs pipeline."""
return asyncio_run(self._wf.run(query_str=query))
