from fastapi import FastAPI, Form, Request, Response, File, Depends, HTTPException, status
from fastapi.responses import RedirectResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from fastapi.encoders import jsonable_encoder
from langchain.llms import CTransformers
from langchain.chains import QAGenerationChain
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.docstore.document import Document
from langchain.document_loaders import PyPDFLoader
from langchain.prompts import PromptTemplate
from langchain.embeddings import HuggingFaceBgeEmbeddings
from langchain.vectorstores import FAISS
from langchain.chains.summarize import load_summarize_chain
from langchain.chains import RetrievalQA
import os 
import json
import time
import uvicorn
import aiofiles
from PyPDF2 import PdfReader
import csv

app = FastAPI()

app.mount("/static", StaticFiles(directory="static"), name="static")

templates = Jinja2Templates(directory="templates")

def load_llm():
    # Load the locally downloaded model here
    llm = CTransformers(
        model = "mistral-7b-instruct-v0.1.Q4_K_S.gguf",
        model_type="mistral",
        max_new_tokens = 1048,
        temperature = 0.3
    )
    return llm

def file_processing(file_path):

    # Load data from PDF
    loader = PyPDFLoader(file_path)
    data = loader.load()

    question_gen = ''

    for page in data:
        question_gen += page.page_content
        
    splitter_ques_gen = RecursiveCharacterTextSplitter(
        chunk_size = 1000,
        chunk_overlap = 100
    )

    chunks_ques_gen = splitter_ques_gen.split_text(question_gen)

    document_ques_gen = [Document(page_content=t) for t in chunks_ques_gen]

    splitter_ans_gen = RecursiveCharacterTextSplitter(
        chunk_size = 300,
        chunk_overlap = 30
    )


    document_answer_gen = splitter_ans_gen.split_documents(
        document_ques_gen
    )

    return document_ques_gen, document_answer_gen

def llm_pipeline(file_path):

    document_ques_gen, document_answer_gen = file_processing(file_path)



    llm_ques_gen_pipeline = load_llm()

    prompt_template = """
    You are an expert at creating questions based on coding materials and documentation.
    Your goal is to prepare a coder or programmer for their exam and coding tests.
    You do this by asking questions about the text below:

    ------------
    {text}
    ------------

    Create questions that will prepare the coders or programmers for their tests.
    Make sure not to lose any important information.

    QUESTIONS:
    """

    PROMPT_QUESTIONS = PromptTemplate(template=prompt_template, input_variables=["text"])

    refine_template = ("""
    You are an expert at creating practice questions based on coding material and documentation.
    Your goal is to help a coder or programmer prepare for a coding test.
    We have received some practice questions to a certain extent: {existing_answer}.
    We have the option to refine the existing questions or add new ones.
    (only if necessary) with some more context below.
    ------------
    {text}
    ------------

    Given the new context, refine the original questions in English.
    If the context is not helpful, please provide the original questions.
    QUESTIONS:
    """
    )

    REFINE_PROMPT_QUESTIONS = PromptTemplate(
        input_variables=["existing_answer", "text"],
        template=refine_template,
    )

    ques_gen_chain = load_summarize_chain(llm = llm_ques_gen_pipeline, 
                                            chain_type = "refine", 
                                            verbose = True, 
                                            question_prompt=PROMPT_QUESTIONS, 
                                            refine_prompt=REFINE_PROMPT_QUESTIONS)

    ques = ques_gen_chain.run(document_ques_gen)

    embeddings = HuggingFaceBgeEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")

    vector_store = FAISS.from_documents(document_answer_gen, embeddings)

    llm_answer_gen = load_llm()

    ques_list = ques.split("\n")
    filtered_ques_list = [element for element in ques_list if element.endswith('?') or element.endswith('.')]

    answer_generation_chain = RetrievalQA.from_chain_type(llm=llm_answer_gen, 
                                                chain_type="stuff", 
                                                retriever=vector_store.as_retriever())

    return answer_generation_chain, filtered_ques_list


def get_csv (file_path):
    answer_generation_chain, ques_list = llm_pipeline(file_path)
    base_folder = 'static/output/'
    if not os.path.isdir(base_folder):
        os.mkdir(base_folder)
    output_file = base_folder+"QA.csv"
    with open(output_file, "w", newline="", encoding="utf-8") as csvfile:
        csv_writer = csv.writer(csvfile)
        csv_writer.writerow(["Question", "Answer"])  # Writing the header row

        for question in ques_list:
            print("Question: ", question)
            answer = answer_generation_chain.run(question)
            print("Answer: ", answer)
            print("--------------------------------------------------\n\n")

            # Save answer to CSV file
            csv_writer.writerow([question, answer])
    return output_file

@app.get("/")
async def index(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})

@app.post("/upload")
async def chat(request: Request, pdf_file: bytes = File(), filename: str = Form(...)):
    base_folder = 'static/docs/'
    if not os.path.isdir(base_folder):
        os.mkdir(base_folder)
    pdf_filename = os.path.join(base_folder, filename)

    async with aiofiles.open(pdf_filename, 'wb') as f:
        await f.write(pdf_file)
    response_data = jsonable_encoder(json.dumps({"msg": 'success',"pdf_filename": pdf_filename}))
    res = Response(response_data)
    return res


@app.post("/analyze")
async def chat(request: Request, pdf_filename: str = Form(...)):
    output_file = get_csv(pdf_filename)
    response_data = jsonable_encoder(json.dumps({"output_file": output_file}))
    res = Response(response_data)
    return res

if __name__ == "__main__":
    uvicorn.run("app:app", host='0.0.0.0', port=8000, reload=True)
langchain.chains.summarize.load_summarize_chain langchain.text_splitter.RecursiveCharacterTextSplitter langchain.embeddings.HuggingFaceBgeEmbeddings langchain.docstore.document.Document langchain.llms.CTransformers langchain.vectorstores.FAISS.from_documents langchain.document_loaders.PyPDFLoader langchain.prompts.PromptTemplate
[((911, 920), 'fastapi.FastAPI', 'FastAPI', ([], {}), '()\n', (918, 920), False, 'from fastapi import FastAPI, Form, Request, Response, File, Depends, HTTPException, status\n'), ((1008, 1046), 'fastapi.templating.Jinja2Templates', 'Jinja2Templates', ([], {'directory': '"""templates"""'}), "(directory='templates')\n", (1023, 1046), False, 'from fastapi.templating import Jinja2Templates\n'), ((945, 976), 'fastapi.staticfiles.StaticFiles', 'StaticFiles', ([], {'directory': '"""static"""'}), "(directory='static')\n", (956, 976), False, 'from fastapi.staticfiles import StaticFiles\n'), ((1123, 1247), 'langchain.llms.CTransformers', 'CTransformers', ([], {'model': '"""mistral-7b-instruct-v0.1.Q4_K_S.gguf"""', 'model_type': '"""mistral"""', 'max_new_tokens': '(1048)', 'temperature': '(0.3)'}), "(model='mistral-7b-instruct-v0.1.Q4_K_S.gguf', model_type=\n    'mistral', max_new_tokens=1048, temperature=0.3)\n", (1136, 1247), False, 'from langchain.llms import CTransformers\n'), ((1385, 1407), 'langchain.document_loaders.PyPDFLoader', 'PyPDFLoader', (['file_path'], {}), '(file_path)\n', (1396, 1407), False, 'from langchain.document_loaders import PyPDFLoader\n'), ((1562, 1628), 'langchain.text_splitter.RecursiveCharacterTextSplitter', 'RecursiveCharacterTextSplitter', ([], {'chunk_size': '(1000)', 'chunk_overlap': '(100)'}), '(chunk_size=1000, chunk_overlap=100)\n', (1592, 1628), False, 'from langchain.text_splitter import RecursiveCharacterTextSplitter\n'), ((1831, 1895), 'langchain.text_splitter.RecursiveCharacterTextSplitter', 'RecursiveCharacterTextSplitter', ([], {'chunk_size': '(300)', 'chunk_overlap': '(30)'}), '(chunk_size=300, chunk_overlap=30)\n', (1861, 1895), False, 'from langchain.text_splitter import RecursiveCharacterTextSplitter\n'), ((2734, 2800), 'langchain.prompts.PromptTemplate', 'PromptTemplate', ([], {'template': 'prompt_template', 'input_variables': "['text']"}), "(template=prompt_template, input_variables=['text'])\n", (2748, 2800), False, 'from langchain.prompts import PromptTemplate\n'), ((3479, 3569), 'langchain.prompts.PromptTemplate', 'PromptTemplate', ([], {'input_variables': "['existing_answer', 'text']", 'template': 'refine_template'}), "(input_variables=['existing_answer', 'text'], template=\n    refine_template)\n", (3493, 3569), False, 'from langchain.prompts import PromptTemplate\n'), ((3615, 3779), 'langchain.chains.summarize.load_summarize_chain', 'load_summarize_chain', ([], {'llm': 'llm_ques_gen_pipeline', 'chain_type': '"""refine"""', 'verbose': '(True)', 'question_prompt': 'PROMPT_QUESTIONS', 'refine_prompt': 'REFINE_PROMPT_QUESTIONS'}), "(llm=llm_ques_gen_pipeline, chain_type='refine',\n    verbose=True, question_prompt=PROMPT_QUESTIONS, refine_prompt=\n    REFINE_PROMPT_QUESTIONS)\n", (3635, 3779), False, 'from langchain.chains.summarize import load_summarize_chain\n'), ((4033, 4111), 'langchain.embeddings.HuggingFaceBgeEmbeddings', 'HuggingFaceBgeEmbeddings', ([], {'model_name': '"""sentence-transformers/all-mpnet-base-v2"""'}), "(model_name='sentence-transformers/all-mpnet-base-v2')\n", (4057, 4111), False, 'from langchain.embeddings import HuggingFaceBgeEmbeddings\n'), ((4134, 4187), 'langchain.vectorstores.FAISS.from_documents', 'FAISS.from_documents', (['document_answer_gen', 'embeddings'], {}), '(document_answer_gen, embeddings)\n', (4154, 4187), False, 'from langchain.vectorstores import FAISS\n'), ((5686, 5692), 'fastapi.File', 'File', ([], {}), '()\n', (5690, 5692), False, 'from fastapi import FastAPI, Form, Request, Response, File, Depends, HTTPException, status\n'), ((5710, 5719), 'fastapi.Form', 'Form', (['...'], {}), '(...)\n', (5714, 5719), False, 'from fastapi import FastAPI, Form, Request, Response, File, Depends, HTTPException, status\n'), ((5847, 5882), 'os.path.join', 'os.path.join', (['base_folder', 'filename'], {}), '(base_folder, filename)\n', (5859, 5882), False, 'import os\n'), ((6084, 6107), 'fastapi.Response', 'Response', (['response_data'], {}), '(response_data)\n', (6092, 6107), False, 'from fastapi import FastAPI, Form, Request, Response, File, Depends, HTTPException, status\n'), ((6205, 6214), 'fastapi.Form', 'Form', (['...'], {}), '(...)\n', (6209, 6214), False, 'from fastapi import FastAPI, Form, Request, Response, File, Depends, HTTPException, status\n'), ((6349, 6372), 'fastapi.Response', 'Response', (['response_data'], {}), '(response_data)\n', (6357, 6372), False, 'from fastapi import FastAPI, Form, Request, Response, File, Depends, HTTPException, status\n'), ((6424, 6486), 'uvicorn.run', 'uvicorn.run', (['"""app:app"""'], {'host': '"""0.0.0.0"""', 'port': '(8000)', 'reload': '(True)'}), "('app:app', host='0.0.0.0', port=8000, reload=True)\n", (6435, 6486), False, 'import uvicorn\n'), ((1754, 1778), 'langchain.docstore.document.Document', 'Document', ([], {'page_content': 't'}), '(page_content=t)\n', (1762, 1778), False, 'from langchain.docstore.document import Document\n'), ((4812, 4838), 'os.path.isdir', 'os.path.isdir', (['base_folder'], {}), '(base_folder)\n', (4825, 4838), False, 'import os\n'), ((4849, 4870), 'os.mkdir', 'os.mkdir', (['base_folder'], {}), '(base_folder)\n', (4857, 4870), False, 'import os\n'), ((5008, 5027), 'csv.writer', 'csv.writer', (['csvfile'], {}), '(csvfile)\n', (5018, 5027), False, 'import csv\n'), ((5768, 5794), 'os.path.isdir', 'os.path.isdir', (['base_folder'], {}), '(base_folder)\n', (5781, 5794), False, 'import os\n'), ((5805, 5826), 'os.mkdir', 'os.mkdir', (['base_folder'], {}), '(base_folder)\n', (5813, 5826), False, 'import os\n'), ((5901, 5934), 'aiofiles.open', 'aiofiles.open', (['pdf_filename', '"""wb"""'], {}), "(pdf_filename, 'wb')\n", (5914, 5934), False, 'import aiofiles\n'), ((6012, 6072), 'json.dumps', 'json.dumps', (["{'msg': 'success', 'pdf_filename': pdf_filename}"], {}), "({'msg': 'success', 'pdf_filename': pdf_filename})\n", (6022, 6072), False, 'import json\n'), ((6296, 6336), 'json.dumps', 'json.dumps', (["{'output_file': output_file}"], {}), "({'output_file': output_file})\n", (6306, 6336), False, 'import json\n')]