db

Rago DB package.

Modules:

• base –

  Base classes for database.

• chroma –

  ChromaDB implementation for vector database.

• faiss –

  Module for faiss database.

Classes:

• ChromaDB –

  ChromaDB implementation for vector database.

• DBBase –

  Base class for vector database.

• FaissDB –

  Faiss Database.

ChromaDB

ChromaDB(client: ClientAPI, collection_name: str = 'rago')

Bases: DBBase

ChromaDB implementation for vector database.

Methods:

• embed –

  Embed the documents into the database.

• search –

  Search a query from documents.

Source code in src/rago/augmented/db/chroma.py

def __init__(
    self,
    client: 'ClientAPI',
    collection_name: str = 'rago',
) -> None:
    """Initialize ChromaDB."""
    self.client = client
    self.collection_name = collection_name
    self._setup()

embed

embed(documents: Any) -> None

Embed the documents into the database.

Source code in src/rago/augmented/db/chroma.py

def embed(self, documents: Any) -> None:
    """Embed the documents into the database."""
    if not isinstance(documents, tuple) or len(documents) != 2:
        raise ValueError(
            'documents format must be: (List[str], List[List[float]])'
        )

    documents_list: List[str] = documents[0]
    embeddings_list: List[List[float]] = documents[1]

    # Convert embeddings to numpy array
    embeddings = np.array(embeddings_list, dtype=np.float32)

    self.collection.add(
        documents=documents_list,
        embeddings=embeddings,
        ids=[str(i) for i in range(len(documents_list))],
    )

search

search(
    query_encoded: Any, top_k: int = 2
) -> Tuple[List[float], List[str]]

Search a query from documents.

Source code in src/rago/augmented/db/chroma.py

def search(
    self, query_encoded: Any, top_k: int = 2
) -> Tuple[List[float], List[str]]:
    """Search a query from documents."""
    # Convert query_encoded to numpy array
    query_encoded_np = np.array([query_encoded], dtype=np.float32)

    results = self.collection.query(
        query_embeddings=query_encoded_np.tolist(),
        n_results=top_k,
    )

    # Check if keys exist before accessing them
    distances = results.get('distances', [[]])
    ids = results.get('ids', [[]])

    # Ensure distances and ids are not None before indexing
    distances_list: List[float] = distances[0] if distances else []
    ids_list: List[str] = ids[0] if ids else []

    return distances_list, ids_list

DBBase

Base class for vector database.

Methods:

• embed –

  Embed the documents into the database.

• search –

  Search a query from documents.

embed abstractmethod

embed(documents: Any) -> None

Embed the documents into the database.

Source code in src/rago/augmented/db/base.py

@abstractmethod
def embed(self, documents: Any) -> None:
    """Embed the documents into the database."""
    ...

search abstractmethod

search(
    query_encoded: Any, top_k: int = 2
) -> tuple[
    Iterable[float], Union[Iterable[str], Iterable[int]]
]

Search a query from documents.

Source code in src/rago/augmented/db/base.py

@abstractmethod
def search(
    self, query_encoded: Any, top_k: int = 2
) -> tuple[Iterable[float], Union[Iterable[str], Iterable[int]]]:
    """Search a query from documents."""
    ...

FaissDB

Bases: DBBase

Faiss Database.

Methods:

• embed –

  Embed the documents into the database.

• search –

  Search an encoded query into vector database.

embed

embed(documents: Any) -> None

Embed the documents into the database.

Source code in src/rago/augmented/db/faiss.py

def embed(self, documents: Any) -> None:
    """Embed the documents into the database."""
    self.index = faiss.IndexFlatL2(documents.shape[1])
    self.index.add(documents)

search

search(
    query_encoded: Any, top_k: int = 2
) -> tuple[Iterable[float], Iterable[int]]

Search an encoded query into vector database.

Source code in src/rago/augmented/db/faiss.py

def search(
    self, query_encoded: Any, top_k: int = 2
) -> tuple[Iterable[float], Iterable[int]]:
    """Search an encoded query into vector database."""
    distances, indices = self.index.search(query_encoded, top_k)
    return distances, indices[0]