# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/indexing/01_indexing_chains.ipynb.

# %% auto 0
__all__ = ['add_ems_to_docs_chain', 'docs_to_faiss_chain', 'pdf_to_faiss_chain', 'index_query_chain']

# %% ../../nbs/indexing/01_indexing_chains.ipynb 2
from ..imports import *
from ..chains import *
from ..utils import *
from .utils import *
from ..pdf.chains import *


# %% ../../nbs/indexing/01_indexing_chains.ipynb 4
def add_ems_to_docs_chain(
    ems_model,
    input_variables=["docs"],
    output_variables=["ems_docs"],
    verbose=False,
):
    "Chain that takes a list of `Documents` and adds `embeddings` to each `Document`."
    return transform_chain(
        add_ems_to_docs,
        transform_kwargs=dict(ems_model=ems_model),
        input_variables=input_variables,
        output_variables=output_variables,
        vars_kwargs_mapping={input_variables[0]: "docs"},
        verbose=verbose,
    )


def docs_to_faiss_chain(
    ems_model,
    index_folder,
    index_name,
    input_variables=["docs"],
    output_variables=["docs"],
    verbose=False,
):
    "Chain that takes a list of `Documents` and adds them to a `FAISS` index in `index_folder`."
    return transform_chain(
        docs_to_faiss,
        input_variables=input_variables,
        output_variables=output_variables,
        transform_kwargs={
            "ems_model": ems_model,
            "index_folder": index_folder,
            "index_name": index_name,
        },
        vars_kwargs_mapping={input_variables[0]: "docs"},
        verbose=verbose,
    )


def pdf_to_faiss_chain(
    ems_model,  # The model to use for vectorestore embeddings.
    index_folder,  # The folder to store the FAISS index.
    index_name,  # The name of the FAISS index.
    input_variables=["pdf_path"],
    output_variables=["docs"],
    chunk_size=200,  # The number of characters per Document.
    chunk_overlap=20,  # The number of characters to overlap between Documents.
    block_size=1500,  # The number of PDFs to process in a single Ray task.
    num_cpus=12,  # The number of CPUs to use for Ray.
    num_gpus=1,  # The number of GPUs to use for Ray.
    verbose=False,
):
    """
    Chain that adds PDFs to `FAISS` indexes in `index_folder`.
    If there are more than `block_size` PDFs, indexing will be distributed using `Ray`.
    """
    docs_chain = pdf_to_docs_chain(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        input_variables=input_variables,
        output_variables=["dc"],
        verbose=verbose,
    )
    faiss_chain = docs_to_faiss_chain(
        ems_model=ems_model,
        index_folder=index_folder,
        index_name=index_name,
        input_variables=["dc"],
        output_variables=output_variables,
        verbose=verbose,
    )
    faiss_chain = ray_chain(
        faiss_chain,
        block_size=block_size,
        num_cpus=num_cpus,
        num_gpus=num_gpus,
        verbose=verbose,
    )
    return SequentialChain(
        chains=[docs_chain, faiss_chain],
        input_variables=input_variables,
        output_variables=output_variables,
        verbose=verbose,
    )


def index_query_chain(
    ems_model,  # The SentenceTransformer model to use for vectorestore embeddings.
    index_folder,  # The folder with the FAISS indexes.
    index_name,  # The name of the FAISS index.
    input_variables=["query", "k"],
    output_variables=["search_results"],
    block_size=10,  # The number of indexes to process in a single Ray task.
    num_cpus=12,  # The number of CPUs to use for Ray.
    num_gpus=1,  # The number of GPUs to use for Ray.
    verbose=False,
):
    """
    Chain that takes a `query` and returns the top `k` results from the `FAISS` indexes in `index_folder`.
    If there are more than `block_size` indexes, search will be distributed using `Ray`.
    """

    index_names_chain = transform_chain(
        index_names,
        transform_kwargs={"index_folder": index_folder, "index_name": index_name},
        input_variables=["k"],
        output_variables=["index_names"],
    )

    search_faiss_chain = transform_chain(
        search_faiss,
        transform_kwargs={"index_folder": index_folder, "ems_model": ems_model},
        input_variables=["index_names", "query", "k"],
        output_variables=["res"],
    )

    search_faiss_chain = ray_chain(
        search_faiss_chain, block_size=block_size, num_cpus=num_cpus, num_gpus=num_gpus
    )

    def flatten_res(res, k):
        if is_list(k):
            k = k[0]
        return [sorted(flatten_list(res), key=lambda x: x[1])[:k]]

    res_chain = transform_chain(
        flatten_res,
        input_variables=["res", "k"],
        output_variables=output_variables,
    )

    return SequentialChain(
        chains=[index_names_chain, search_faiss_chain, res_chain],
        input_variables=input_variables,
        output_variables=output_variables,
        verbose=verbose,
    )

