app.py

from dotenv import load_dotenv
import os
load_dotenv()

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_text_splitters import Language
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_openai import ChatOpenAI
from langchain.chains import RetrievalQA
import chromadb
import gradio as gr
import tqdm

def read_file(file_path):
    with open(file_path, "r", encoding="utf-8") as file:
        return file.read()

def infer_module_name(file_path):
    path_parts = file_path.split(os.sep)
    if "src" in path_parts:
        src_index = path_parts.index("src")
        return "/".join(path_parts[src_index+1:-1])
    return "root"

def process_files(root_dir, file_extension, language=None):
    if language:
        splitter = RecursiveCharacterTextSplitter.from_language(
            language=language, chunk_size=3000, chunk_overlap=100
        )
    else:
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=3000, chunk_overlap=100
        )
    
    all_docs = []

    for root, _, files in os.walk(root_dir):
        for file in files:
            if file.endswith(file_extension):
                file_path = os.path.join(root, file)
                file_name = os.path.basename(file_path)
                folder_path = root
                module_name = infer_module_name(file_path)
                content = read_file(file_path)
                content = f"file name: {file_name}\n path: {folder_path}\n {content}"

                docs = splitter.create_documents(
                    [content],
                    metadatas=[{
                        'source': file_name, 
                        'type': file_extension[1:],
                        'module': module_name,  
                        'folder_path': folder_path  
                    }]
                )
                all_docs.extend(docs)

    return all_docs

def process_all_files(root_directory):
    ts_docs = process_files(root_directory, '.ts', Language.TS)
    html_docs = process_files(root_directory, '.html', Language.HTML)
    txt_docs = process_files(root_directory, '.txt')
    md_docs = process_files(root_directory, '.md')
    js_docs = process_files(root_directory, '.js', Language.JS)
    kt_docs = process_files(root_directory, '.kt', Language.KOTLIN)

    all_docs = ts_docs + html_docs + txt_docs + md_docs + js_docs + kt_docs
    return all_docs

def initialize_or_load_database():
    model_name = 'text-embedding-3-large'
    embeddings = OpenAIEmbeddings(
        model=model_name,
        openai_api_key=os.environ.get('OPENAI_API_KEY')
    )

    chroma_client = chromadb.PersistentClient(path="./mifos-mobile_vector_storage")
    collection_name = "all_files"

    if os.path.exists("collection_storage.txt"):
        with open("collection_storage.txt", "r") as f:
            collection_storage_name, collection_storage_id = f.read().splitlines()
        print("Loading existing vector database...")
        docsearch = Chroma(
            client=chroma_client,
            collection_name=collection_name,
            embedding_function=embeddings
        )
    else:
        print("Creating new vector database...")
        root_directory = "mifos-mobile"    
        all_documents = process_all_files(root_directory)
        print(f"Total number of chunks across all files: {len(all_documents)}")
        print("Total number of files: ", len(set([doc.metadata['source'] for doc in all_documents])))

        docsearch = Chroma.from_documents(
            documents=all_documents,
            embedding=embeddings,
            collection_name=collection_name,
            client=chroma_client
        )

        collection_storage_name = chroma_client.list_collections()[0].name
        collection_storage_id = chroma_client.list_collections()[0].id
        print("name: ", collection_storage_name)
        print("id: ", collection_storage_id)
        
        with open("collection_storage.txt", "w") as f:
            f.write(f"{collection_storage_name}\n{collection_storage_id}")

    return docsearch

docsearch = initialize_or_load_database()

llm = ChatOpenAI(
    openai_api_key=os.environ.get('OPENAI_API_KEY'),
    model_name='gpt-4o-mini',
    temperature=0.3
)

qa = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",  
    retriever=docsearch.as_retriever(),
    return_source_documents=True
)

def get_top_20_embeddings(query):
    docs_and_scores = docsearch.similarity_search_with_score(query, k=20) 
    return docs_and_scores


def get_parent_document_embeddings(query, num_docs=5):
    docs_and_scores = docsearch.similarity_search_with_score(query, k=num_docs)
    parent_docs = {}
    
    for doc, score in docs_and_scores:
        parent_doc_key = doc.metadata['source'] 
        if parent_doc_key not in parent_docs:
            parent_docs[parent_doc_key] = (doc, score)
    
    return list(parent_docs.values())

def get_top_5_parent_documents(query):
    return get_parent_document_embeddings(query, num_docs=5)

def answer_question_with_parent_docs(question):
    top_5_results = get_top_5_parent_documents(question)
    
    context = "\n".join([doc.page_content for doc, _ in top_5_results])
    # print("Context: ", context)
    
    query_data = (
        "You are an expert in project structure and various file types including TypeScript, HTML, Markdown, JS and Kotlin."
        "When answering questions, focus on the file organization, key components of the codebase, and the structure of the project."
        "For general queries,like hi,hello etc, provide a brief answer, but for questions about project structure, include module names, file paths, and folder organization."
        "If you're unsure of the answer, suggest referring to the Mifos Slack Channel."
        "\nContext:\n" + context + "\n" + question
    )

    response = qa.invoke(query_data)
    
    top_20_results = get_top_5_parent_documents(question)
    # print("Top 5 matching parent documents:")
    # for i, (doc, score) in enumerate(top_20_results, 1):
    #     print(f"{i}. Document: {doc.page_content[:1000]}...")
    #     print(f"   Metadata: {doc.metadata}")
    #     print(f"   Similarity Score: {score}")
    #     print()
    
    return response['result']


interface = gr.Interface(
    fn=answer_question_with_parent_docs, 
    inputs=gr.Textbox(label="Ask a question about the files"),
    outputs=gr.Textbox(label="Answer"),
    title="Mifos-Mobile Chatbot",
    description="Ask questions about Kotlin in Mifos-Mobile",
)

if __name__ == "__main__":
    interface.launch()