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_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"]
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)
            
            files = []
            print("tmpdirname: " , tmpdirname)
            unzipped_root = os.listdir(tmpdirname)
            print("unzipped_root ", unzipped_root)
            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):
  #  print("Start setup_vectorstore_function")
  #  embedding_model = HuggingFaceEmbeddings(model_name=model_name)
    
  #  vectorstore = Chroma.from_documents(chunks, embedding=embedding_model, persist_directory=persist_directory)
  #  vectorstore.persist()
   # print("test1", vectorstore._persist_directory)
  #  print("test2",vectorstore.__dir__)
  #  return vectorstore


def setup_vectorstore(chunks, model_name):
    print("Start setup_vectorstore_function")
    
    # Create a temporary directory to use as the persist_directory
    with tempfile.TemporaryDirectory() as temp_dir:
        print(f"Using temporary directory: {temp_dir}")
        
        # Initialize the embedding model
        embedding_model = HuggingFaceEmbeddings(model_name=model_name)
        
        # Set up the vectorstore with the temporary directory
        vectorstore = Chroma.from_documents(chunks, embedding=embedding_model, persist_directory=temp_dir)
        vectorstore.persist()
        
        # Optionally, display the persist directory for debugging
        print("Persist directory:", vectorstore._persist_directory)
        print("Available methods in vectorstore:", dir(vectorstore))
        
        # At this point, you can use your API upload method to upload the persisted vectorstore files
        for root, _, files in os.walk(temp_dir):
            for file_name in files:
                file_path = os.path.join(root, file_name)
                target_path_in_repo = os.path.relpath(file_path, temp_dir)
                print(f"Uploading file: {file_path} -> {target_path_in_repo}")
                api.upload_file(
                    path_or_fileobj=file_path,
                    path_in_repo=target_path_in_repo,
                    repo_id=HF_SPACE_NAME,
                    repo_type="space"
                )
                print(f"Uploaded {file_path} to {target_path_in_repo}")
        
        print("All files uploaded successfully!")


# 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("Document Chunks:\n")
    print("\n\n".join(["="*80 + "\n" + doc_chunk for doc_chunk, _ in retrieved_doc_chunks]))
    print("\nDocument References:\n")
    print(doc_references)

    print("\n" + "="*80 + "\n")  # Separator between doc and code

    print("Code Chunks:\n")
    print("\n\n".join(["="*80 + "\n" + code_chunk for code_chunk, _ in retrieved_code_chunks]))
    print("\nCode References:\n")
    print(code_references)


   # 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. 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 docstore, codestore, chunks, llm

    #code_partial_paths = ['kadi_apy/lib/']
    #code_file_path = []
    doc_partial_paths = []
    #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)
    #codestore = setup_vectorstore(code_chunks, EMBEDDING_MODEL_NAME, PERSIST_CODE_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()