app.py

import os
import json
import gradio as gr
import zipfile
import tempfile
import requests
import urllib.parse
import io

from huggingface_hub import HfApi, login
from PyPDF2 import PdfReader
from langchain_huggingface import HuggingFaceEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_groq import ChatGroq
from dotenv import load_dotenv
from langchain.docstore.document import Document
from langchain.schema import Document
from chunk_python_code import chunk_python_code_with_metadata
from vectorstore import get_chroma_vectorstore

# Load environment variables from .env file
load_dotenv()

# Load configuration from JSON file
with open('config.json') as config_file:
    config = json.load(config_file)

with open("config2.json", "r") as file:
    config2 = json.load(file)

PERSIST_DOC_DIRECTORY = config["persist_doc_directory"]
PERSIST_CODE_DIRECTORY =config["persist_code_directory"]
CHUNK_SIZE = config["chunk_size"]
CHUNK_OVERLAP = config["chunk_overlap"]
EMBEDDING_MODEL_NAME = config["embedding_model"]
LLM_MODEL_NAME = config["llm_model"]
LLM_TEMPERATURE = config["llm_temperature"]
GITLAB_API_URL = config["gitlab_api_url"]
HF_SPACE_NAME = config["hf_space_name"]
DATA_DIR = config["data_dir"]

GROQ_API_KEY = os.environ["GROQ_API_KEY"]
HF_TOKEN = os.environ["HF_Token"]



login(HF_TOKEN)
api = HfApi()

def load_project_id(json_file):
    with open(json_file, 'r') as f:
        data = json.load(f)
    return data['project_id']


def download_gitlab_project_by_version():
    try:
        # Load the configuration from config.json

        # Extract GitLab project information from the config
        api_url = config2['gitlab']['api_url']
        project_id = urllib.parse.quote(config2['gitlab']['project']['id'], safe="")
        version = config2['gitlab']['project']['version']
        
        # Construct the URL for the release's zip file
        url = f"{api_url}/projects/{project_id}/repository/archive.zip?sha={version}"

        # Send GET request to download the zip file
        response = requests.get(url, stream=True)
        archive_bytes = io.BytesIO(response.content)

        
        if response.status_code == 200:
            # Extract filename from content-disposition header
            content_disposition = response.headers.get("content-disposition")
            if content_disposition and "filename=" in content_disposition:
                filename = content_disposition.split("filename=")[-1].strip('"')


        # test
        # target_path = f"{DATA_DIR}/{filename}"
        
        # Check if the request was successful
        if response.status_code == 200:
            api.upload_file(
                path_or_fileobj= archive_bytes,
                path_in_repo= f"{DATA_DIR}/{filename}",
                repo_id=HF_SPACE_NAME,
                repo_type='space'
        )
            print(f"Release {version} downloaded successfully as {file_path}.")
        else:
            print(f"Failed to download the release: {response.status_code} - {response.reason}")
            print(response.text)

    except FileNotFoundError:
        print("The config.json file was not found. Please ensure it exists in the project directory.")
    except json.JSONDecodeError:
        print("Failed to parse the config.json file. Please ensure it contains valid JSON.")
    except Exception as e:
        print(f"An error occurred: {e}")



def download_gitlab_repo():
    print("Start the upload_gitRepository function")
    project_id = load_project_id('repository_ids.json')
    encoded_project_id = urllib.parse.quote_plus(project_id)
    
    # Define the URL to download the repository archive
    archive_url = f"{GITLAB_API_URL}/projects/{encoded_project_id}/repository/archive.zip"
    
    # Download the repository archive
    response = requests.get(archive_url)
    archive_bytes = io.BytesIO(response.content)
    
    # Retrieve the original file name from the response headers
    content_disposition = response.headers.get('content-disposition')
    if content_disposition:
        filename = content_disposition.split('filename=')[-1].strip('\"')
    else:
        filename = 'archive.zip'  # Fallback to a default name if not found

    # Check if the file already exists in the repository
    existing_files = api.list_repo_files(repo_id=HF_SPACE_NAME, repo_type='space')
    target_path = f"{DATA_DIR}/{filename}"

    print(f"Target Path: '{target_path}'")
    print(f"Existing Files: {[repr(file) for file in existing_files]}")
    
    if target_path in existing_files:
        print(f"File '{target_path}' already exists in the repository. Skipping upload...")
    else:
        # Upload the ZIP file to the new folder in the Hugging Face space repository
        print("Uploading File to directory:")
        print(f"Archive Bytes: {repr(archive_bytes.getvalue())[:100]}")  # Show a preview of bytes
        print(f"Target Path in Repo: '{target_path}'")

        api.upload_file(
            path_or_fileobj=archive_bytes,
            path_in_repo=target_path,
            repo_id=HF_SPACE_NAME,
            repo_type='space'
        )
        print("Upload complete")


def get_all_files_in_folder(temp_dir, folder_path):
    
    all_files = [] 
    target_dir = os.path.join(temp_dir, folder_path)

    for root, dirs, files in os.walk(target_dir):
        print(f"Files in current directory ({root}): {files}")
        for file in files:
            print(f"Processing file: {file}")
            all_files.append(os.path.join(root, file))

    return all_files

def get_file(temp_dir, file_path):
    full_path = os.path.join(temp_dir, file_path)
    return full_path


def process_directory(directory, folder_paths, file_paths):
    all_texts = []
    file_references = []

    zip_filename = next((file for file in os.listdir(directory) if file.endswith('.zip')), None)
    zip_file_path = os.path.join(directory, zip_filename)

    with tempfile.TemporaryDirectory() as tmpdirname:
        # Unzip the file into the temporary directory
        with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
            zip_ref.extractall(tmpdirname)
            
            files = []
            print("tmpdirname: " , tmpdirname)   
            unzipped_root = os.listdir(tmpdirname)
            print("unzipped_root ", unzipped_root)

            tmpsubdirpath= os.path.join(tmpdirname, unzipped_root[0])  
            print("tempsubdirpath: ", tmpsubdirpath)

            if folder_paths:
                for folder_path in folder_paths:
                    files += get_all_files_in_folder(tmpsubdirpath, folder_path) 
            if file_paths:
                files += [get_file(tmpsubdirpath, file_path) for file_path in file_paths] 

            
            print(f"Total number of files: {len(files)}")
            
            for file_path in files:
              #  print("111111111:", file_path)    
                file_ext = os.path.splitext(file_path)[1]
             #   print("222222222:", file_ext)
                if os.path.getsize(file_path) == 0:
                    print(f"Skipping an empty file: {file_path}")
                    continue

                with open(file_path, 'rb') as f:
                    if file_ext in ['.rst', '.py']:
                        text = f.read().decode('utf-8')
                    
                    all_texts.append(text)
                    print("Filepaths brother:", file_path)
                    relative_path = os.path.relpath(file_path, tmpsubdirpath)
                    print("Relative Filepaths brother:", relative_path)
                    file_references.append(relative_path)
                    
    return all_texts, file_references
                



def split_python_code_into_chunks(texts, file_paths):
    chunks = [] 
    for text, file_path in zip(texts, file_paths):
        document_chunks = chunk_python_code_with_metadata(text, file_path)
        chunks.extend(document_chunks)   
    return chunks


# Split text into chunks
def split_into_chunks(texts, references, chunk_size, chunk_overlap):
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    chunks = []

    for text, reference in zip(texts, references):
        chunks.extend([
            Document(
                page_content=chunk,
                metadata={
                    "source": reference,
                    "usage": "doc"
                }
            ) 
            for chunk in text_splitter.split_text(text)
        ])
    return chunks


# Setup Vectorstore
def embed_documents_into_vectorstore(chunks, model_name, persist_directory):
    print("Start setup_vectorstore_function")
    embedding_model = HuggingFaceEmbeddings(model_name=model_name)   
    vectorstore = get_chroma_vectorstore(embedding_model, persist_directory)
    vectorstore.add_documents(chunks)
    return vectorstore

# Setup LLM
def setup_llm(model_name, temperature, api_key):
    llm = ChatGroq(model=model_name, temperature=temperature, api_key=api_key)
    return llm

def retrieve_from_vectorstore(vectorstore, query, k):
    retrieved_docs = vectorstore.similarity_search(query, k=k)    
    return retrieved_docs



def retrieve_within_kadiApy_docs(vectorstore, query, k):
    filter_criteria = {"usage": "doc"}
    retrieved_docs = vectorstore.similarity_search(query=query, k=k, filter=filter_criteria)
    return retrieved_docs

def retrieve_within_kadiApy_library(vectorstore, query, k):
    filter_criteria = {"usage": "library"}
    retrieved_docs = vectorstore.similarity_search(query=query, k=k, filter=filter_criteria)
    return retrieved_docs

def retrieve_within_kadiApy_cli_library(vectorstore, query, k):
    filter_criteria = {"usage": "cli_library"}
    retrieved_docs = vectorstore.similarity_search(query=query, k=k, filter=filter_criteria)
    return retrieved_docs

def retrieve_within_kadiApy_cli_library_excluding_cli_commands(vectorstore, query, k):
    filter_criteria = {"usage": "cli_library"}
    filtered_docs = [doc for doc in documents if "command" not in doc.metadata]
    return filtered_docs
    
    
def retrieve_kadiApy_cli_commands(vectorstore, query, k):
    filter_criteria = {"type": "command"}
    results = vectorstore.similarity_search(query=query, k=k, filter=filter_criteria)
    return results



def retrieve_docs_from_vectorstore(vectorstore, query, k): 
    return vectorstore.similarity_search(query, k=k)




def format_kadi_apy_library_context(docs):
    doc_context = []
    
    for doc in docs:
        # Extract metadata information
        class_info = doc.metadata.get("class", "Unknown Class")
        type_info = doc.metadata.get("type", "Unknown Type")
        source_info = doc.metadata.get("source", "Unknown Type")
        # Format metadata and document content
        #print("YYYYYYYEEEEEEEEEEEEEEE222222222222222222222222222222:}\n\n", doc.page_content)
        formatted_doc = f"# source: {source_info}\n# class: {class_info}\n# type: {type_info}\n{doc.page_content}\n\n\n"
        doc_context.append(formatted_doc)
    
    return doc_context


def format_kadi_api_doc_context(docs):
    doc_context = []

    for doc in docs:
        source_info = doc.metadata.get("source", "Unknown Type")
        #print("YYYYYYYEEEEEEEEEEEEEEE:}\n\n", doc.page_content)
        formatted_doc = f"# source: {source_info}\n{doc.page_content}\n\n\n"
        doc_context.append(formatted_doc)
    
    return doc_context


                   
def rag_workflow(query):
    


    
    rewritten_query_response = llm.invoke("Hello")
    rewritten_query = rewritten_query_response.content.strip()


    prompt = (
        f"""The query is: '{query}'.
            Based on the user's query, assist them by determining which technical document they should read to interact with the software named 'Kadi4Mat'. 
            There are three different technical documents to choose from:
                - Document 1: Provides information on how to use a Python library to interact with the HTTP API of 'Kadi4Mat'.
                - Document 2: Provides information on how to use a Python library to implement custom CLI commands to interact with 'Kadi4Mat'.
                - Document 3: Provides prewritten CLI command lines and guides on how to use the command-line interface to interact with 'Kadi4Mat'.
    
            Your task is to select the single most likely option. 
                If Document 1 is the best choice, respond with 'kadi-apy python library'. 
                If Document 2 is the best choice, respond with 'kadi-apy python cli library'. 
                If Document 3 is the best choice, respond with 'kadi-apy command line interface'. 
            Respond with only the exact corresponding option and do not include any additional comments, explanations, or text."
        """
    )
    metadata_prediction = llm.predict(prompt)

    print("METADATA PREDICTION -------------------------:", metadata_prediction)
    print(metadata_prediction)
    
    kadi_apy_docs = retrieve_within_kadiApy_docs (vector_store, query, k = 5)
    kadi_apy_library_docs = vector_store.similarity_search(query, filter={"usage": metadata_prediction})
        
    doc_context = format_kadi_api_doc_context(kadi_apy_docs)
    code_context = format_kadi_apy_library_context(kadi_apy_library_docs)

    
    # Assuming kadi_apy_library_docs is a list of Document objects

    print("HERE WE GHOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOOO")
    print("::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::")
    for doc in kadi_apy_library_docs:
        print(doc.metadata.get("source", "Unknown Type"))
        print("\n")
    print("::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::")



    
    prompt = f"""You are an expert python developer. You are assisting in generating code for users who wants to make use of "kadi-apy", an API library. 
                    "Doc-context:" provides you with information how to use this API library by givnig code examples and code documentation.
                    "Code-context:" provides you information of API methods and classes from the "kadi-apy" library. 
                    Based on the retrieved contexts and the guidelines answer the query.  
                    

                    General Guidelines:
                    - If no related information is found from the contexts to answer the query, reply that you do not know.
                    
                    Guidelines when generating code:
                    - First display the full code and then follow with a well structured explanation of the generated code.

            Doc-context:
            {doc_context}

            Code-context:
            {code_context}
            
            Query: 
            {query}
    """

    
    response = llm.invoke(prompt)
    return response.content


def initialize():
    global vector_store, chunks, llm

    download_gitlab_project_by_version()
    
    code_file_paths = ['kadi_apy']
    doc_file_path = ['docs/source/']

    
    code_texts, code_references = process_directory(DATA_DIR, code_partial_paths, code_file_paths)
    print("LEEEEEEEEEEEENGTH of code_texts: ", len(code_texts))

    
    doc_texts, kadiAPY_doc_references = process_directory(DATA_DIR, doc_partial_paths, doc_file_paths)
    print("LEEEEEEEEEEEENGTH of doc_files: ", len(doc_texts))
    
    code_chunks = split_python_code_into_chunks(code_texts, code_references)
    doc_chunks = split_into_chunks(doc_texts, kadiAPY_doc_references, CHUNK_SIZE, CHUNK_OVERLAP)

    print(f"Total number of code_chunks: {len(code_chunks)}")
    print(f"Total number of doc_chunks: {len(doc_chunks)}")

    #docstore = embed_documents_into_vectorstore(kadiAPY_code_chunks, EMBEDDING_MODEL_NAME, PERSIST_DOC_DIRECTORY)
    #codestore = embed_documents_into_vectorstore(kadiAPY_doc_chunks, EMBEDDING_MODEL_NAME, PERSIST_CODE_DIRECTORY)

    vector_store = embed_documents_into_vectorstore(
                    chunks= doc_chunks + code_chunks, 
                    model_name= EMBEDDING_MODEL_NAME, 
                    persist_directory= PERSIST_DOC_DIRECTORY 
                    )
    
    llm = setup_llm(LLM_MODEL_NAME, LLM_TEMPERATURE, GROQ_API_KEY)


initialize()

# Gradio utils
def check_input_text(text):
    if not text:
        gr.Warning("Please input a question.")
        raise TypeError
    return True

def add_text(history, text):
    history = history + [(text, None)]
    yield history, ""




import gradio as gr


def bot_kadi(history):
    user_query = history[-1][0]
    response = rag_workflow(user_query)
    history[-1] = (user_query, response)

    yield history  

def main():
    with gr.Blocks() as demo:
        gr.Markdown("## Kadi4Mat - AI Chat-Bot")
        gr.Markdown("AI assistant for Kadi4Mat based on RAG architecture powered by LLM")

        with gr.Tab("Kadi4Mat - AI Assistant"):
            with gr.Row():
                with gr.Column(scale=10):
                    chatbot = gr.Chatbot([], elem_id="chatbot", label="Kadi Bot", bubble_full_width=False, show_copy_button=True, height=600)
                    user_txt = gr.Textbox(label="Question", placeholder="Type in your question and press Enter or click Submit")

                    with gr.Row():
                        with gr.Column(scale=1):
                            submit_btn = gr.Button("Submit", variant="primary")
                        with gr.Column(scale=1):
                            clear_btn = gr.Button("Clear", variant="stop")

                    gr.Examples(
                        examples=[
                            "Who is working on Kadi4Mat?",
                            "How do i install the Kadi-Apy library?",
                            "How do i install the Kadi-Apy library for development?",
                            "I need a method to upload a file to a record",
                        ],
                        inputs=user_txt,
                        outputs=chatbot,
                        fn=add_text,
                        label="Try asking...",
                        cache_examples=False,
                        examples_per_page=3,
                    )

            user_txt.submit(check_input_text, user_txt, None).success(add_text, [chatbot, user_txt], [chatbot, user_txt]).then(bot_kadi, [chatbot], [chatbot])
            submit_btn.click(check_input_text, user_txt, None).success(add_text, [chatbot, user_txt], [chatbot, user_txt]).then(bot_kadi, [chatbot], [chatbot])
            #user_txt.submit(check_input_text, user_txt, None).success(add_text, [chatbot, user_txt], [chatbot, user_txt]).then(bot_kadi, [chatbot], [chatbot, doc_citation])
            #submit_btn.click(check_input_text, user_txt, None).success(add_text, [chatbot, user_txt], [chatbot, user_txt]).then(bot_kadi, [chatbot], [chatbot, doc_citation])
            clear_btn.click(lambda: None, None, chatbot, queue=False)

    demo.launch() 

    
if __name__ == "__main__":
    main()