Skip to content

Longrag

LongRAGPack #

Bases: BaseLlamaPack

Implements Long RAG.

This implementation is based on the following paper: https://arxiv.org/pdf/2406.15319

Source code in llama-index-packs/llama-index-packs-longrag/llama_index/packs/longrag/base.py
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
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))

get_modules #

get_modules() -> Dict[str, Any]

Get Modules.

Source code in llama-index-packs/llama-index-packs-longrag/llama_index/packs/longrag/base.py
381
382
383
384
385
386
387
388
389
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,
    }

run #

run(query: str, *args: Any, **kwargs: Any) -> Any

Runs pipeline.

Source code in llama-index-packs/llama-index-packs-longrag/llama_index/packs/longrag/base.py
391
392
393
def run(self, query: str, *args: t.Any, **kwargs: t.Any) -> t.Any:
    """Runs pipeline."""
    return asyncio_run(self._wf.run(query_str=query))