# Databricks notebook source
from src.retriever import init_vectorDB_from_doc, retriever

from transformers import AutoTokenizer, pipeline
from langchain_core.prompts import ChatPromptTemplate
from typing import List,Optional, Tuple # import the Tuple type
from langchain.docstore.document import Document as LangchainDocument
from langchain_community.vectorstores import FAISS
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains import create_retrieval_chain
def promt_template():
    prompt_in_chat_format = """Using the information contained in the given context, give a comprehensive answer to the question.
            Respond only to the question asked, response should be concise and relevant to the question.
            Provide the number of the source document when relevant.If the answer cannot be deduced from the context, do not give an answer. Please answer in french
            {context}
             """
           
    prompt = ChatPromptTemplate.from_messages(
        [
            ("system",prompt_in_chat_format),
            ("human", "{input}")
        ])
    #RAG_PROMPT_TEMPLATE = tokenizer.apply_chat_template(
    #prompt_in_chat_format, tokenize=False, add_generation_prompt=True)
    return prompt 

def answer_with_rag(
    query: str, retriever,llm

) -> Tuple[str, List[LangchainDocument]]:
    # Build the final prompt
    #relevant_docs= retriever(query,vectorDB,reranker,num_doc_before_rerank,num_final_relevant_docs,rerank) 
    #context = "\nExtracted documents:\n"
    #context += "".join([f"Document {str(i)}:::\n" + doc for i, doc in enumerate(relevant_docs)])
    #print("=> Context:")
    #print(context)
    RAG_PROMPT_TEMPLATE = promt_template()
    document_chain = create_stuff_documents_chain(llm, RAG_PROMPT_TEMPLATE)
    retrieval_chain=create_retrieval_chain(retriever,document_chain)
    print("=> Final prompt:")
    #print(final_prompt)
    # Redact an answer
    print("=> Generating answer...")
    response=retrieval_chain.invoke({'input':query})

    return response['answer'], response["context"]