MifosBot/Mobile-Wallet
0
1from dotenv import load_dotenv2import os3load_dotenv()4 5OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")6 7from langchain_text_splitters import RecursiveCharacterTextSplitter8from langchain_text_splitters import Language9from langchain_openai import OpenAIEmbeddings10from langchain_community.vectorstores import Chroma11from langchain_openai import ChatOpenAI12from langchain.chains import RetrievalQA13import chromadb14import gradio as gr15import tqdm16 17def read_file(file_path):18 with open(file_path, "r", encoding="utf-8") as file:19 return file.read()20 21def infer_module_name(file_path):22 path_parts = file_path.split(os.sep)23 if "src" in path_parts:24 src_index = path_parts.index("src")25 return "/".join(path_parts[src_index+1:-1])26 return "root"27 28def process_files(root_dir, file_extension, language=None):29 if language:30 splitter = RecursiveCharacterTextSplitter.from_language(31 language=language, chunk_size=3000, chunk_overlap=10032 )33 else:34 splitter = RecursiveCharacterTextSplitter(35 chunk_size=3000, chunk_overlap=10036 )37 38 all_docs = []39 40 for root, _, files in os.walk(root_dir):41 for file in files:42 if file.endswith(file_extension):43 file_path = os.path.join(root, file)44 file_name = os.path.basename(file_path)45 folder_path = root46 module_name = infer_module_name(file_path)47 content = read_file(file_path)48 content = f"file name: {file_name}\n path: {folder_path}\n {content}"49 50 docs = splitter.create_documents(51 [content],52 metadatas=[{53 'source': file_name, 54 'type': file_extension[1:],55 'module': module_name, 56 'folder_path': folder_path 57 }]58 )59 all_docs.extend(docs)60 61 return all_docs62 63def process_all_files(root_directory):64 ts_docs = process_files(root_directory, '.ts', Language.TS)65 html_docs = process_files(root_directory, '.html', Language.HTML)66 txt_docs = process_files(root_directory, '.txt')67 md_docs = process_files(root_directory, '.md')68 js_docs = process_files(root_directory, '.js', Language.JS)69 kt_docs = process_files(root_directory, '.kt', Language.KOTLIN)70 71 all_docs = ts_docs + html_docs + txt_docs + md_docs + js_docs + kt_docs72 return all_docs73 74def initialize_or_load_database():75 model_name = 'text-embedding-3-large'76 embeddings = OpenAIEmbeddings(77 model=model_name,78 openai_api_key=os.environ.get('OPENAI_API_KEY')79 )80 81 chroma_client = chromadb.PersistentClient(path="./mobile_wallet_vector_storage")82 collection_name = "all_files"83 84 if os.path.exists("collection_storage.txt"):85 with open("collection_storage.txt", "r") as f:86 collection_storage_name, collection_storage_id = f.read().splitlines()87 print("Loading existing vector database...")88 docsearch = Chroma(89 client=chroma_client,90 collection_name=collection_name,91 embedding_function=embeddings92 )93 else:94 print("Creating new vector database...")95 root_directory = "mobile-wallet" 96 all_documents = process_all_files(root_directory)97 print(f"Total number of chunks across all files: {len(all_documents)}")98 print("Total number of files: ", len(set([doc.metadata['source'] for doc in all_documents])))99 100 docsearch = Chroma.from_documents(101 documents=all_documents,102 embedding=embeddings,103 collection_name=collection_name,104 client=chroma_client105 )106 107 collection_storage_name = chroma_client.list_collections()[0].name108 collection_storage_id = chroma_client.list_collections()[0].id109 # print("name: ", collection_storage_name)110 # print("id: ", collection_storage_id)111 112 with open("collection_storage.txt", "w") as f:113 f.write(f"{collection_storage_name}\n{collection_storage_id}")114 115 return docsearch116 117docsearch = initialize_or_load_database()118 119llm = ChatOpenAI(120 openai_api_key=os.environ.get('OPENAI_API_KEY'),121 model_name='gpt-4o-mini',122 temperature=0.3123)124 125qa = RetrievalQA.from_chain_type(126 llm=llm,127 chain_type="stuff", 128 retriever=docsearch.as_retriever(),129 return_source_documents=True130)131 132def get_top_20_embeddings(query):133 docs_and_scores = docsearch.similarity_search_with_score(query, k=20) 134 return docs_and_scores135 136 137def get_parent_document_embeddings(query, num_docs=5):138 139 docs_and_scores = docsearch.similarity_search_with_score(query, k=num_docs)140 141 parent_docs = {}142 143 for doc, score in docs_and_scores:144 parent_doc_key = doc.metadata['source'] 145 if parent_doc_key not in parent_docs:146 parent_docs[parent_doc_key] = (doc, score)147 148 return list(parent_docs.values())149 150def get_top_5_parent_documents(query):151 return get_parent_document_embeddings(query, num_docs=5)152 153def answer_question_with_parent_docs(question):154 top_5_results = get_top_5_parent_documents(question)155 156 context = "\n".join([doc.page_content for doc, _ in top_5_results])157 print("Context: ", context)158 159 query_data = (160 "You are an expert in project structure and various file types including TypeScript, HTML, Markdown, JS and Kotlin."161 "When answering questions, focus on the file organization, key components of the codebase, and the structure of the project."162 "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."163 "If you're unsure of the answer, suggest referring to the Mifos Slack Channel."164 "\nContext:\n" + context + "\n" + question165 )166 167 response = qa.invoke(query_data)168 169 # top_20_results = get_top_5_parent_documents(question)170 # print("Top 5 matching parent documents:")171 # for i, (doc, score) in enumerate(top_20_results, 1):172 # print(f"{i}. Document: {doc.page_content[:1000]}...")173 # print(f" Metadata: {doc.metadata}")174 # print(f" Similarity Score: {score}")175 # print()176 177 return response['result']178 179 180interface = gr.Interface(181 fn=answer_question_with_parent_docs, 182 inputs=gr.Textbox(label="Ask a question about the files"),183 outputs=gr.Textbox(label="Answer"),184 title="Mifos Mobile-Wallet Chatbot",185 description="Ask questions about Kotlin in Mifos Mobile-Wallet",186)187 188if __name__ == "__main__":189 interface.launch()190