app.py

from flask import Flask, request, jsonify
from langchain_groq import ChatGroq
from langchain_community.llms import Ollama
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from langchain_community.embeddings.fastembed import FastEmbedEmbeddings
from langchain_community.document_loaders import PDFPlumberLoader
from langchain_core.prompts import ChatPromptTemplate
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains import create_retrieval_chain
from langchain.prompts import PromptTemplate
import os
from pathlib import Path

app = Flask(__name__)

folder_path = "db"

cached_llm = Ollama(model="llama3")

embedding = FastEmbedEmbeddings()

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1024, chunk_overlap=80, length_function=len, is_separator_regex=False
)

raw_prompt = PromptTemplate.from_template(
    """ 
    <s>[INST] You are a technical assistant good at searching docuemnts. If you do not have an answer from the provided information say so. [/INST] </s>
    [INST] {input}
           Context: {context}
           Answer:
    [/INST]
"""
)


@app.route("/ai", methods=["POST"])
def aiPost():
    print("Post /ai called")
    json_content = request.json
    query = json_content.get("query")

    print(f"query: {query}")

    response = cached_llm.invoke(query)

    print(response)

    response_answer = {"answer": response}
    return response_answer


@app.route("/ask_pdf_all", methods=["POST"])
def askPDFPost():
    # index = 0
    # print(len(folder_path))
    # lst = os.listdir(folder_path) # your directory path
    # number_files = len(os.listdir(folder_path))
    # print(len(next(os.walk(folder_path))[1]))
    # folder_counter = sum([len(folder) for r, d, folder in os.walk(folder_path)])
    # print(folder_counter)

    # looping subfolder ada berapa
    folder_count = 0
    for folders in os.listdir(folder_path):
        if os.path.isdir(os.path.join(folder_path, folders)):  # if it's a directory
            folder_count += 1  # increment counter

    print("There are {} folders".format(folder_count))

    index = 0
    for file_name in os.listdir(folder_path):
        if os.path.isdir(os.path.join(folder_path, file_name)):  # if it's a directory
    # for index < len(folder_path):
    # for index in range(len(folder_path)):
            print()
            print()
            # file_name = folder_path[index] # index 0 -> file_name = alice
            file_path = os.path.join("db", file_name) # file_path = db/alice
            print(f"Pada file {file_name}")

            # Kodingan LLM nya
            json_content = request.json
            query = json_content.get("query")
            print(f"Jawaban dari pertanyaan : {query}")
            print()

            userdata = {}
            groq_api_key = userdata.get('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')

            print("Loading vector store")
            vector_store = Chroma(
                persist_directory=file_path, 
                embedding_function=embedding,
                )

            retriever = vector_store.as_retriever()

            llm = ChatGroq(
                    groq_api_key=groq_api_key,
                    model_name='mixtral-8x7b-32768'
            )
            
            # Define the RAG template and chain
            rag_template = """Answer the question based only on the following context:
            {context}
            Question: {question}
            """
            rag_prompt = ChatPromptTemplate.from_template(rag_template)

            rag_chain = (
                {"context": retriever, "question": RunnablePassthrough()}
                | rag_prompt
                | llm
                | StrOutputParser()
            )

            # Test the architecture with a simple hard coded question
            result = rag_chain.invoke({"input": query})
            print(result)

            index += 1  # increment counter

    return jsonify(result)
    # return jsonify (folder_count)

@app.route("/ask_pdf_selected", methods=["POST"])
def search():
    json_content = request.json
    filename = json_content.get("fn")
    file_path = os.path.join("db", filename)

    if os.path.exists(file_path):
        #open the file for iteration
        print(f"Content of folder'{filename}'")
        print(os.listdir(f"{folder_path}/{filename}"))
        print()

        print("Post /ask_pdf called")
        query = json_content.get("query")
        print(f"question: {query}")
        print()

        userdata = {}
        groq_api_key = userdata.get('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')

        print("Loading vector store")
        vector_store = Chroma(
            persist_directory=file_path, 
            embedding_function=embedding,
            )

        retriever = vector_store.as_retriever()

        llm = ChatGroq(
                groq_api_key=groq_api_key,
                model_name='mixtral-8x7b-32768'
        )
            # Define the RAG template and chain
        rag_template = """Answer the question based only on the following context:
        {context}
        Question: {question}
        """
        rag_prompt = ChatPromptTemplate.from_template(rag_template)

        rag_chain = (
            {"context": retriever, "question": RunnablePassthrough()}
            | rag_prompt
            | llm
            | StrOutputParser()
        )

        # Test the architecture with a simple hard coded question
        result = rag_chain.invoke({"input": query})
        print(result)

    return jsonify(result)

@app.route("/pdf", methods=["POST"])
def pdfPost():
    file = request.files["file"]
    file_name = file.filename
    save_file = "pdf/" + file_name
    file.save(save_file)
    print(f"filename: {file_name}")

    loader = PDFPlumberLoader(save_file)
    docs = loader.load_and_split()
    print(f"docs len={len(docs)}")

    chunks = text_splitter.split_documents(docs)
    print(f"chunks len={len(chunks)}")

    vector_store = Chroma.from_documents(
        documents=chunks, embedding=embedding, persist_directory=folder_path + f'/{file_name}'
    )

    vector_store.persist()

    response = {
        "status": "Successfully Uploaded", 
        "filename": file_name,
        "doc_len": len(docs),
        "chunks": len(chunks),
    }
    return response


def start_app():
    app.run(host="0.0.0.0", port=8080, debug=True)


if __name__ == "__main__":
    start_app()