Skip to content

augmented

Augmented package.

Modules:

  • base –

    Base classes for the augmented step.

  • cohere –

    Classes for augmentation with Cohere embeddings.

  • db –

    Rago DB package.

  • fireworks –

    Classes for augmentation with Fireworks embeddings.

  • openai –

    Classes for augmentation with OpenAI embeddings.

  • sentence_transformer –

    Classes for augmentation with hugging face.

  • spacy –

    Classes for augmentation with SpaCy embeddings.

  • together –

    Classes for augmentation with Together embeddings.

Classes:

  • AugmentedBase –

    Define the base structure for Augmented classes.

  • CohereAug –

    Class for augmentation with Cohere embeddings.

  • FireworksAug –

    Class for augmentation with Fireworks embeddings.

  • OpenAIAug –

    Class for augmentation with OpenAI embeddings.

  • SentenceTransformerAug –

    Class for augmentation with Hugging Face.

  • SpaCyAug –

    Class for augmentation with SpaCy embeddings.

  • TogetherAug –

    Class for augmentation with Together embeddings.

AugmentedBase 

AugmentedBase(
    model_name: Optional[str] = None,
    db: DBBase = FaissDB(),
    top_k: Optional[int] = None,
    api_key: str = '',
    cache: Optional[Cache] = None,
    logs: dict[str, Any] = DEFAULT_LOGS,
)

Bases: RagoBase

Define the base structure for Augmented classes.

Methods:

  • get_embedding –

    Retrieve the embedding for a given text using OpenAI API.

  • search –

    Search an encoded query into vector database.

Source code in src/rago/augmented/base.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def __init__(
    self,
    model_name: Optional[str] = None,
    db: DBBase = FaissDB(),
    top_k: Optional[int] = None,
    api_key: str = '',
    cache: Optional[Cache] = None,
    logs: dict[str, Any] = DEFAULT_LOGS,
) -> None:
    """Initialize AugmentedBase."""
    if logs is DEFAULT_LOGS:
        logs = {}
    super().__init__(api_key=api_key, cache=cache, logs=logs)

    self.db = db

    self.top_k = top_k if top_k is not None else self.default_top_k
    self.model_name = (
        model_name if model_name is not None else self.default_model_name
    )
    self.model = None

    self._validate()
    self._setup()

get_embedding 

get_embedding(content: list[str]) -> EmbeddingType

Retrieve the embedding for a given text using OpenAI API.

Source code in src/rago/augmented/base.py 
75
76
77
def get_embedding(self, content: list[str]) -> EmbeddingType:
    """Retrieve the embedding for a given text using OpenAI API."""
    raise Exception('Method not implemented.')

search abstractmethod 

search(
    query: str, documents: Any, top_k: int = 0
) -> list[str]

Search an encoded query into vector database.

Source code in src/rago/augmented/base.py 
79
80
81
82
83
84
85
86
87
@abstractmethod
def search(
    self,
    query: str,
    documents: Any,
    top_k: int = 0,
) -> list[str]:
    """Search an encoded query into vector database."""
    ...

CohereAug 

CohereAug(
    model_name: Optional[str] = None,
    db: DBBase = FaissDB(),
    top_k: Optional[int] = None,
    api_key: str = '',
    cache: Optional[Cache] = None,
    logs: dict[str, Any] = DEFAULT_LOGS,
)

Bases: AugmentedBase

Class for augmentation with Cohere embeddings.

Methods:

  • get_embedding –

    Retrieve the embedding for given texts using Cohere API.

  • search –

    Search an encoded query into vector database.

Source code in src/rago/augmented/base.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def __init__(
    self,
    model_name: Optional[str] = None,
    db: DBBase = FaissDB(),
    top_k: Optional[int] = None,
    api_key: str = '',
    cache: Optional[Cache] = None,
    logs: dict[str, Any] = DEFAULT_LOGS,
) -> None:
    """Initialize AugmentedBase."""
    if logs is DEFAULT_LOGS:
        logs = {}
    super().__init__(api_key=api_key, cache=cache, logs=logs)

    self.db = db

    self.top_k = top_k if top_k is not None else self.default_top_k
    self.model_name = (
        model_name if model_name is not None else self.default_model_name
    )
    self.model = None

    self._validate()
    self._setup()

get_embedding 

get_embedding(content: list[str]) -> EmbeddingType

Retrieve the embedding for given texts using Cohere API.

Source code in src/rago/augmented/cohere.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def get_embedding(self, content: list[str]) -> EmbeddingType:
    """Retrieve the embedding for given texts using Cohere API."""
    cache_key = sha256(''.join(content).encode('utf-8')).hexdigest()
    cached = self._get_cache(cache_key)
    if cached is not None:
        return cast(EmbeddingType, cached)

    model = cast(cohere.Client, self.model)
    response = model.embed(
        texts=content,
        model=self.model_name,
        input_type='search_document',
        embedding_types=['float'],
    )
    result = np.array(response.embeddings.float_, dtype=np.float32)  # type: ignore[union-attr]

    self._save_cache(cache_key, result)

    return result

search 

search(
    query: str, documents: list[str], top_k: int = 0
) -> list[str]

Search an encoded query into vector database.

Source code in src/rago/augmented/cohere.py 
49
50
51
52
53
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 search(
    self, query: str, documents: list[str], top_k: int = 0
) -> list[str]:
    """Search an encoded query into vector database."""
    if not getattr(self, 'db', None):
        raise Exception('Vector database (db) is not initialized.')
    document_encoded = self.get_embedding(documents)
    model = cast(cohere.Client, self.model)
    response = model.embed(
        texts=[query],
        model=self.model_name,
        input_type='search_query',
        embedding_types=['float'],
    )
    query_encoded = np.array(response.embeddings.float_, dtype=np.float32)  # type: ignore[union-attr]

    top_k = top_k or self.top_k or self.default_top_k or 1

    self.db.embed(document_encoded)
    scores, indices = self.db.search(query_encoded, top_k=top_k)

    self.logs['indices'] = indices
    self.logs['scores'] = scores
    self.logs['search_params'] = {
        'query_encoded': query_encoded,
        'top_k': top_k,
    }

    retrieved_docs = [documents[i] for i in indices if i >= 0]

    return retrieved_docs

FireworksAug 

FireworksAug(
    model_name: Optional[str] = None,
    db: DBBase = FaissDB(),
    top_k: Optional[int] = None,
    api_key: str = '',
    cache: Optional[Cache] = None,
    logs: dict[str, Any] = DEFAULT_LOGS,
)

Bases: AugmentedBase

Class for augmentation with Fireworks embeddings.

Methods:

  • get_embedding –

    Retrieve the embedding for given texts using the OpenAI client.

  • search –

    Search an encoded query into vector database.

Source code in src/rago/augmented/base.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def __init__(
    self,
    model_name: Optional[str] = None,
    db: DBBase = FaissDB(),
    top_k: Optional[int] = None,
    api_key: str = '',
    cache: Optional[Cache] = None,
    logs: dict[str, Any] = DEFAULT_LOGS,
) -> None:
    """Initialize AugmentedBase."""
    if logs is DEFAULT_LOGS:
        logs = {}
    super().__init__(api_key=api_key, cache=cache, logs=logs)

    self.db = db

    self.top_k = top_k if top_k is not None else self.default_top_k
    self.model_name = (
        model_name if model_name is not None else self.default_model_name
    )
    self.model = None

    self._validate()
    self._setup()

get_embedding 

get_embedding(content: list[str]) -> EmbeddingType

Retrieve the embedding for given texts using the OpenAI client.

Source code in src/rago/augmented/fireworks.py 
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
def get_embedding(self, content: list[str]) -> EmbeddingType:
    """Retrieve the embedding for given texts using the OpenAI client."""
    cache_key = sha256(''.join(content).encode('utf-8')).hexdigest()
    cached = self._get_cache(cache_key)
    if cached is not None:
        return cast(EmbeddingType, cached)

    # Using the OpenAI embeddings API call for fireworks
    response = self.openai_client.embeddings.create(
        model=self.model_name,
        input=content,
    )
    result = np.array(
        [data.embedding for data in response.data], dtype=np.float32
    )
    self._save_cache(cache_key, result)
    return result

search 

search(
    query: str, documents: list[str], top_k: int = 0
) -> list[str]

Search an encoded query into vector database.

Source code in src/rago/augmented/fireworks.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
def search(
    self, query: str, documents: list[str], top_k: int = 0
) -> list[str]:
    """Search an encoded query into vector database."""
    if not hasattr(self, 'db') or not self.db:
        raise Exception('Vector database (db) is not initialized.')

    document_encoded = self.get_embedding(documents)
    query_encoded = self.get_embedding([query])
    top_k = top_k or self.top_k or self.default_top_k or 1

    self.db.embed(document_encoded)
    scores, indices = self.db.search(query_encoded, top_k=top_k)

    self.logs['indices'] = indices
    self.logs['scores'] = scores
    self.logs['search_params'] = {
        'query_encoded': query_encoded,
        'top_k': top_k,
    }

    retrieved_docs = [documents[i] for i in indices if i >= 0]

    return retrieved_docs

OpenAIAug 

OpenAIAug(
    model_name: Optional[str] = None,
    db: DBBase = FaissDB(),
    top_k: Optional[int] = None,
    api_key: str = '',
    cache: Optional[Cache] = None,
    logs: dict[str, Any] = DEFAULT_LOGS,
)

Bases: AugmentedBase

Class for augmentation with OpenAI embeddings.

Methods:

  • get_embedding –

    Retrieve the embedding for a given text using OpenAI API.

  • search –

    Search an encoded query into vector database.

Source code in src/rago/augmented/base.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def __init__(
    self,
    model_name: Optional[str] = None,
    db: DBBase = FaissDB(),
    top_k: Optional[int] = None,
    api_key: str = '',
    cache: Optional[Cache] = None,
    logs: dict[str, Any] = DEFAULT_LOGS,
) -> None:
    """Initialize AugmentedBase."""
    if logs is DEFAULT_LOGS:
        logs = {}
    super().__init__(api_key=api_key, cache=cache, logs=logs)

    self.db = db

    self.top_k = top_k if top_k is not None else self.default_top_k
    self.model_name = (
        model_name if model_name is not None else self.default_model_name
    )
    self.model = None

    self._validate()
    self._setup()

get_embedding 

get_embedding(content: list[str]) -> EmbeddingType

Retrieve the embedding for a given text using OpenAI API.

Source code in src/rago/augmented/openai.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def get_embedding(self, content: list[str]) -> EmbeddingType:
    """Retrieve the embedding for a given text using OpenAI API."""
    cache_key = sha256(''.join(content).encode('utf-8')).hexdigest()
    cached = self._get_cache(cache_key)
    if cached is not None:
        return cast(EmbeddingType, cached)

    model = cast(openai.OpenAI, self.model)
    response = model.embeddings.create(
        input=content, model=self.model_name
    )
    result = np.array(
        [data.embedding for data in response.data], dtype=np.float32
    )

    self._save_cache(cache_key, result)

    return result

search 

search(
    query: str, documents: list[str], top_k: int = 0
) -> list[str]

Search an encoded query into vector database.

Source code in src/rago/augmented/openai.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
def search(
    self, query: str, documents: list[str], top_k: int = 0
) -> list[str]:
    """Search an encoded query into vector database."""
    if not hasattr(self, 'db') or not self.db:
        raise Exception('Vector database (db) is not initialized.')

    # Encode the documents and query
    document_encoded = self.get_embedding(documents)
    query_encoded = self.get_embedding([query])
    top_k = top_k or self.top_k or self.default_top_k or 1

    self.db.embed(document_encoded)
    scores, indices = self.db.search(query_encoded, top_k=top_k)

    self.logs['indices'] = indices
    self.logs['scores'] = scores
    self.logs['search_params'] = {
        'query_encoded': query_encoded,
        'top_k': top_k,
    }

    retrieved_docs = [documents[i] for i in indices if i >= 0]

    return retrieved_docs

SentenceTransformerAug 

SentenceTransformerAug(
    model_name: Optional[str] = None,
    db: DBBase = FaissDB(),
    top_k: Optional[int] = None,
    api_key: str = '',
    cache: Optional[Cache] = None,
    logs: dict[str, Any] = DEFAULT_LOGS,
)

Bases: AugmentedBase

Class for augmentation with Hugging Face.

Methods:

  • get_embedding –

    Retrieve the embedding for a given text using OpenAI API.

  • search –

    Search an encoded query into vector database.

Source code in src/rago/augmented/base.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def __init__(
    self,
    model_name: Optional[str] = None,
    db: DBBase = FaissDB(),
    top_k: Optional[int] = None,
    api_key: str = '',
    cache: Optional[Cache] = None,
    logs: dict[str, Any] = DEFAULT_LOGS,
) -> None:
    """Initialize AugmentedBase."""
    if logs is DEFAULT_LOGS:
        logs = {}
    super().__init__(api_key=api_key, cache=cache, logs=logs)

    self.db = db

    self.top_k = top_k if top_k is not None else self.default_top_k
    self.model_name = (
        model_name if model_name is not None else self.default_model_name
    )
    self.model = None

    self._validate()
    self._setup()

get_embedding 

get_embedding(content: list[str]) -> EmbeddingType

Retrieve the embedding for a given text using OpenAI API.

Source code in src/rago/augmented/sentence_transformer.py 
24
25
26
27
def get_embedding(self, content: list[str]) -> EmbeddingType:
    """Retrieve the embedding for a given text using OpenAI API."""
    model = cast(SentenceTransformer, self.model)
    return model.encode(content)

search 

search(
    query: str, documents: Any, top_k: int = 0
) -> list[str]

Search an encoded query into vector database.

Source code in src/rago/augmented/sentence_transformer.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
def search(self, query: str, documents: Any, top_k: int = 0) -> list[str]:
    """Search an encoded query into vector database."""
    if not self.model:
        raise Exception('The model was not created.')

    document_encoded = self.get_embedding(documents)
    query_encoded = self.get_embedding([query])
    top_k = top_k or self.top_k or self.default_top_k or 1

    self.db.embed(document_encoded)

    scores, indices = self.db.search(query_encoded, top_k=top_k)

    retrieved_docs = [documents[i] for i in indices]

    self.logs['indices'] = indices
    self.logs['scores'] = scores
    self.logs['search_params'] = {
        'query_encoded': query_encoded,
        'top_k': top_k,
    }

    return retrieved_docs

SpaCyAug 

SpaCyAug(
    model_name: Optional[str] = None,
    db: DBBase = FaissDB(),
    top_k: Optional[int] = None,
    api_key: str = '',
    cache: Optional[Cache] = None,
    logs: dict[str, Any] = DEFAULT_LOGS,
)

Bases: AugmentedBase

Class for augmentation with SpaCy embeddings.

Methods:

  • get_embedding –

    Retrieve the embedding for a given text using SpaCy.

  • search –

    Search an encoded query into vector database.

Source code in src/rago/augmented/base.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def __init__(
    self,
    model_name: Optional[str] = None,
    db: DBBase = FaissDB(),
    top_k: Optional[int] = None,
    api_key: str = '',
    cache: Optional[Cache] = None,
    logs: dict[str, Any] = DEFAULT_LOGS,
) -> None:
    """Initialize AugmentedBase."""
    if logs is DEFAULT_LOGS:
        logs = {}
    super().__init__(api_key=api_key, cache=cache, logs=logs)

    self.db = db

    self.top_k = top_k if top_k is not None else self.default_top_k
    self.model_name = (
        model_name if model_name is not None else self.default_model_name
    )
    self.model = None

    self._validate()
    self._setup()

get_embedding 

get_embedding(content: List[str]) -> EmbeddingType

Retrieve the embedding for a given text using SpaCy.

Source code in src/rago/augmented/spacy.py 
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
53
def get_embedding(self, content: List[str]) -> EmbeddingType:
    """Retrieve the embedding for a given text using SpaCy."""
    cache_key = sha256(''.join(content).encode('utf-8')).hexdigest()
    cached = self._get_cache(cache_key)
    if cached is not None:
        return cast(EmbeddingType, cached)

    model = cast(spacy.language.Language, self.model)
    embeddings = []

    for text in content:
        doc = model(text)

        # Ensure the model has proper vectors
        if not doc.has_vector:
            raise ValueError(f"Text: '{text}' has no valid word vectors!")

        embeddings.append(doc.vector)

    result = np.array(embeddings, dtype=np.float32)

    # Ensure 2D shape (num_texts, embedding_dim)
    if result.ndim == 1:
        result = result.reshape(1, -1)

    self._save_cache(cache_key, result)
    return result

search 

search(
    query: str, documents: list[str], top_k: int = 0
) -> list[str]

Search an encoded query into vector database.

Source code in src/rago/augmented/spacy.py 
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 search(
    self, query: str, documents: list[str], top_k: int = 0
) -> list[str]:
    """Search an encoded query into vector database."""
    if not hasattr(self, 'db') or not self.db:
        raise Exception('Vector database (db) is not initialized.')

    # Encode the documents and query
    document_encoded = self.get_embedding(documents)
    query_encoded = self.get_embedding([query])
    top_k = top_k or self.top_k or self.default_top_k or 1

    self.db.embed(document_encoded)
    scores, indices = self.db.search(query_encoded, top_k=top_k)

    self.logs['indices'] = indices
    self.logs['scores'] = scores
    self.logs['search_params'] = {
        'query_encoded': query_encoded,
        'top_k': top_k,
    }

    retrieved_docs = [documents[i] for i in indices if i >= 0]

    return retrieved_docs

TogetherAug 

TogetherAug(
    model_name: Optional[str] = None,
    db: DBBase = FaissDB(),
    top_k: Optional[int] = None,
    api_key: str = '',
    cache: Optional[Cache] = None,
    logs: dict[str, Any] = DEFAULT_LOGS,
)

Bases: AugmentedBase

Class for augmentation with Together embeddings.

Methods:

  • get_embedding –

    Retrieve the embedding for given texts using Together API.

  • search –

    Search an encoded query into vector database.

Source code in src/rago/augmented/base.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def __init__(
    self,
    model_name: Optional[str] = None,
    db: DBBase = FaissDB(),
    top_k: Optional[int] = None,
    api_key: str = '',
    cache: Optional[Cache] = None,
    logs: dict[str, Any] = DEFAULT_LOGS,
) -> None:
    """Initialize AugmentedBase."""
    if logs is DEFAULT_LOGS:
        logs = {}
    super().__init__(api_key=api_key, cache=cache, logs=logs)

    self.db = db

    self.top_k = top_k if top_k is not None else self.default_top_k
    self.model_name = (
        model_name if model_name is not None else self.default_model_name
    )
    self.model = None

    self._validate()
    self._setup()

get_embedding 

get_embedding(content: list[str]) -> EmbeddingType

Retrieve the embedding for given texts using Together API.

Source code in src/rago/augmented/together.py 
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
def get_embedding(self, content: list[str]) -> EmbeddingType:
    """Retrieve the embedding for given texts using Together API."""
    cache_key = sha256(''.join(content).encode('utf-8')).hexdigest()
    cached = self._get_cache(cache_key)
    if cached is not None:
        return cast(EmbeddingType, cached)

    client = cast(Together, self.model)
    all_embeddings = []
    for text in content:
        response = client.embeddings.create(
            model=self.model_name, input=text
        )
        embedding = response.data[0].embedding
        all_embeddings.append(embedding)
    result = np.array(all_embeddings, dtype=np.float32)

    self._save_cache(cache_key, result)

    return result

search 

search(
    query: str, documents: list[str], top_k: int = 0
) -> list[str]

Search an encoded query into vector database.

Source code in src/rago/augmented/together.py 
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
def search(
    self, query: str, documents: list[str], top_k: int = 0
) -> list[str]:
    """Search an encoded query into vector database."""
    if not hasattr(self, 'db') or not self.db:
        raise Exception('Vector database (db) is not initialized.')
    document_encoded = self.get_embedding(documents)
    query_encoded = self.get_embedding([query])
    top_k = top_k or self.top_k or self.default_top_k or 1

    self.db.embed(document_encoded)
    scores, indices = self.db.search(query_encoded, top_k=top_k)

    self.logs['indices'] = indices
    self.logs['scores'] = scores
    self.logs['search_params'] = {
        'query_encoded': query_encoded,
        'top_k': top_k,
    }

    retrieved_docs = [documents[i] for i in indices if i >= 0]

    return retrieved_docs