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_community.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_groq import ChatGroq
from dotenv import load_dotenv
from langchain.docstore.document import Document

# 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)


PERSIST_DIRECTORY = config["persist_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"]
REPOSITORY_DIRECTORY = config["repository_directory"]

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_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"{REPOSITORY_DIRECTORY}/{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, partial_path):
    
    all_files = [] 
    print("inner method of get all files in folder")
    target_dir = os.path.join(temp_dir, partial_path)
    print(target_dir)

    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, partial_paths=None, file_paths=None):
    all_texts = []
    file_references = []

    zip_files = [file for file in os.listdir(directory) if file.endswith('.zip')]

    if not zip_files:
        print("No zip file found in the directory.")
        return all_texts, file_references

    if len(zip_files) > 1:
        print("More than one zip file found.")
        return all_texts, file_references
    else:
        zip_file_path = os.path.join(directory, zip_files[0])

        # Create a temporary directory for the zip file
        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)
            print(f"Extracted {zip_file_path} to {tmpdirname}")
            
            files = []
            
            unzipped_root = os.listdir(tmpdirname)
            if len(unzipped_root) == 1 and os.path.isdir(os.path.join(tmpdirname, unzipped_root[0])):
                tmpsubdirpath= os.path.join(tmpdirname, unzipped_root[0])
            else:
                tmpsubdirpath = tmpdirname
                    
            if not partial_paths and not file_paths:
                for root, _, files_list in os.walk(tmpdirname):
                    for file in files_list:
                        files.append(os.path.join(root, file))
            else: 
                if partial_paths:
                    for partial_path in partial_paths:
                        files += get_all_files_in_folder(tmpsubdirpath, partial_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(f"Paths of files: {iles}")
                file_ext = os.path.splitext(file_path)[1]

                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', '.md', '.txt', '.html', '.json', '.yaml', '.py']:
                        text = f.read().decode('utf-8')
                    elif file_ext in ['.svg']:
                        text = f"SVG file content from {file_path}"
                    elif file_ext in ['.png', '.ico']:
                        text = f"Image metadata from {file_path}"
                    else:
                        continue

                    all_texts.append(text)
                    file_references.append(file_path)

    return all_texts, file_references

import ast

def get_source_segment(source_lines, node):
    start_line, start_col = node.lineno - 1, node.col_offset
    end_line = node.end_lineno - 1 if hasattr(node, 'end_lineno') else node.lineno - 1
    end_col = node.end_col_offset if hasattr(node, 'end_col_offset') else len(source_lines[end_line])

    lines = source_lines[start_line:end_line + 1]
    lines[0] = lines[0][start_col:]
    lines[-1] = lines[-1][:end_col]

    return ''.join(lines)

from langchain.schema import Document

def chunk_python_file_content(content, char_limit=1572):
    source_lines = content.splitlines(keepends=True)
    
    # Parse the content into an abstract syntax tree (AST)
    tree = ast.parse(content)


    chunks = []
    current_chunk = ""
    current_chunk_size = 0
    
    # Find all class definitions and top-level functions in the AST
    class_nodes = [node for node in ast.walk(tree) if isinstance(node, ast.ClassDef)]

    for class_node in class_nodes:
        method_nodes = [node for node in class_node.body if isinstance(node, ast.FunctionDef)]

        if method_nodes:
            first_method_start_line = method_nodes[0].lineno - 1
            class_def_lines = source_lines[class_node.lineno - 1:first_method_start_line]
        else:
            class_def_lines = source_lines[class_node.lineno - 1:class_node.end_lineno]

        class_def = ''.join(class_def_lines)
        class_def_size = len(class_def)
        
        # Add class definition to the current chunk if it fits
        if current_chunk_size + class_def_size <= char_limit:
            current_chunk += f"{class_def.strip()}\n"
            current_chunk_size += class_def_size
        else:
            # Start a new chunk if the class definition exceeds the limit
            if current_chunk:
                chunks.append(current_chunk.strip())
                current_chunk = ""
                current_chunk_size = 0
            current_chunk += f"{class_def.strip()}\n"
            current_chunk_size = class_def_size
        
        for method_node in method_nodes:
            method_def = get_source_segment(source_lines, method_node)
            method_def_size = len(method_def)
            
            # Add method definition to the current chunk if it fits
            if current_chunk_size + method_def_size <= char_limit:
                current_chunk += f"# This is a class method of class: {class_node.name}\n{method_def.strip()}\n"
                current_chunk_size += method_def_size
            else:
                # Start a new chunk if the method definition exceeds the limit
                if current_chunk:
                    chunks.append(current_chunk.strip())
                    current_chunk = ""
                    current_chunk_size = 0
                current_chunk += f"# This is a class method of class: {class_node.name}\n{method_def.strip()}\n"
                current_chunk_size = method_def_size

    if current_chunk:
        chunks.append(current_chunk.strip())

    return chunks



# Split python code into chunks
def split_pythoncode_into_chunks(texts, references, chunk_size, chunk_overlap):
    chunks = []
    
    for text, reference in zip(texts, references):
        file_chunks = chunk_python_file_content(text, char_limit=chunk_size)
        
        for chunk in file_chunks:
            document = Document(page_content=chunk, metadata={"source": reference})
            chunks.append(document)
    
    print(f"Total number of chunks: {len(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}) for chunk in text_splitter.split_text(text)])
    print(f"Total number of chunks: {len(chunks)}")
    return chunks

# Setup Vectorstore
def setup_vectorstore(chunks, model_name, persist_directory):
    embedding_model = HuggingFaceEmbeddings(model_name=model_name)
    vectorstore = Chroma.from_documents(chunks, embedding=embedding_model, persist_directory=persist_directory)
    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):
    results = vectorstore.similarity_search(query, k=k)
    chunks_with_references = [(result.page_content, result.metadata["source"]) for result in results]
    # Print the chosen chunks and their sources to the console
    print("\nChosen chunks and their sources for the query:")
    for chunk, source in chunks_with_references:
        print(f"Source: {source}\nChunk: {chunk}\n")
        print("-" * 50)
    return chunks_with_references

def rag_workflow(query):
    retrieved_doc_chunks = retrieve_from_vectorstore(docstore, query, k=5)
    retrieved_code_chunks = retrieve_from_vectorstore(codestore, query, k=5)
    
    doc_context = "\n\n".join([doc_chunk for doc_chunk, _ in retrieved_doc_chunks])
    code_context = "\n\n".join([code_chunk for code_chunk, _ in retrieved_code_chunks]) 
    
    doc_references = "\n".join([f"[{i+1}] {ref}" for i, (_, ref) in enumerate(retrieved_doc_chunks)])
    code_references = "\n".join([f"[{i+1}] {ref}" for i, (_, ref) in enumerate(retrieved_code_chunks)])
    
    print(f"Context for the query:\n{doc_context}\n")
    
    print(f"References for the query:\n{references}\n")
    
    prompt = f"""You are an expert python developer. Provide a clear and consice answer based only on the information in the retrieved context.
                    The retrieved context contains source code and documenation of an api library. 
                    If no related Information is found from the context to answer the query, reply that you do not know.

            Context:
            {doc_context}

            Query: 
            {query}
    """

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


def initialize():
    global docstore, codestore, chunks, llm

    code_partial_paths = ['kadi_apy/lib/']
    code_file_path = []
    doc_partial_paths = ['docs/source/setup/']
    doc_file_paths = ['docs/source/usage/lib.rst']

    
    code_files, code_file_references = process_directory(REPOSITORY_DIRECTORY, code_partial_paths, code_file_path)
    
    doc_files, doc_file_references = process_directory(REPOSITORY_DIRECTORY, doc_partial_paths, doc_file_paths)
    
    code_chunks = split_pythoncode_into_chunks(code_files, code_file_references, 1500, 0)
    doc_chunks = split_into_chunks(doc_files, doc_file_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 = setup_vectorstore(doc_chunks, EMBEDDING_MODEL_NAME, PERSIST_DIRECTORY)
    codestore = setup_vectorstore(code_chunks, EMBEDDING_MODEL_NAME, PERSIST_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, references = rag_workflow(user_query)
    history[-1] = (user_query, response)

    # Format references for display with text passages
    formatted_references = ""
    docs = query_chroma(vectorstore, user_query, k=5)
    for i, (doc, ref) in enumerate(docs):
        formatted_references += f"""
        <div style="border: 1px solid #ddd; padding: 10px; margin-bottom: 10px; border-radius: 5px;">
            <h3 style="margin-top: 0;">Reference {i+1}</h3>
            <p><strong>Source:</strong> {ref}</p>
            <button onclick="var elem = document.getElementById('text-{i}'); var button = this; if (elem.style.display === 'block') {{ elem.style.display = 'none'; button.innerHTML = '&#9654; show source text'; }} else {{ elem.style.display = 'block'; button.innerHTML = '&#9660; hide source text'; }}">{{'&#9654; show source text'}}</button>
            <div id="text-{i}" style="display: none;">
                <p><strong>Text:</strong> {doc}</p>
            </div>
        </div>
        """

    yield history, formatted_references    

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)
                    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,
                    )

                with gr.Column(scale=3):
                    with gr.Tab("References"):
                        doc_citation = gr.HTML("<p>References used in answering the question will be displayed below.</p>")

            #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()