Skip to content

core

Rago is Retrieval Augmented Generation lightweight framework.

Classes:

  • Rago –

    RAG class.

Rago 

Rago(
    retrieval: RetrievalBase,
    augmented: AugmentedBase,
    generation: GenerationBase,
)

RAG class.

Parameters:

  • retrieval (RetrievalBase) –

    The retrieval component used to fetch relevant data based on the query.

  • augmented (AugmentedBase) –

    The augmentation module responsible for enriching the retrieved data.

  • generation (GenerationBase) –

    The text generation model used to generate a response based on the query and augmented data.

Methods:

  • prompt –

    Run the pipeline for a specific prompt.

Source code in src/rago/core.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
def __init__(
    self,
    retrieval: RetrievalBase,
    augmented: AugmentedBase,
    generation: GenerationBase,
) -> None:
    """Initialize the RAG structure.

    Parameters
    ----------
    retrieval : RetrievalBase
        The retrieval component used to fetch relevant data based
        on the query.
    augmented : AugmentedBase
        The augmentation module responsible for enriching the
        retrieved data.
    generation : GenerationBase
        The text generation model used to generate a response based
        on the query and augmented data.
    """
    self.retrieval = retrieval
    self.augmented = augmented
    self.generation = generation
    self.logs: dict[str, dict[str, Any]] = {
        'retrieval': retrieval.logs,
        'augmented': augmented.logs,
        'generation': generation.logs,
    }

prompt 

prompt(query: str, device: str = 'auto') -> str | BaseModel

Run the pipeline for a specific prompt.

Parameters:

  • query (str) –

    The query or prompt from the user.

  • device (str (default 'auto'), default: 'auto' ) –

    Device for generation (e.g., 'auto', 'cpu', 'cuda'), by default 'auto'.

Returns:

  • str –

    Generated text based on the query and augmented data.

Source code in src/rago/core.py 
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def prompt(self, query: str, device: str = 'auto') -> str | BaseModel:
    """Run the pipeline for a specific prompt.

    Parameters
    ----------
    query : str
        The query or prompt from the user.
    device : str (default 'auto')
        Device for generation (e.g., 'auto', 'cpu', 'cuda'), by
        default 'auto'.

    Returns
    -------
    str
        Generated text based on the query and augmented data.
    """
    ret_data = self.retrieval.get(query)
    self.logs['retrieval']['result'] = ret_data

    aug_data = self.augmented.search(query, ret_data)
    self.logs['augmented']['result'] = aug_data

    gen_data = self.generation.generate(query, context=aug_data)
    self.logs['generation']['result'] = gen_data

    return gen_data