Update app.py

app.py

@@ -280,8 +280,8 @@ def split_into_chunks(texts, references, chunk_size, chunk_overlap):
     print(f"Total number of chunks: {len(chunks)}")
     return chunks
 
-# Setup Chroma
-def setup_chroma(chunks, model_name, persist_directory):
+# 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
@@ -291,7 +291,7 @@ def setup_llm(model_name, temperature, api_key):
     llm = ChatGroq(model=model_name, temperature=temperature, api_key=api_key)
     return llm
 
-def query_chroma(vectorstore, query, k):
+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
@@ -302,14 +302,14 @@ def query_chroma(vectorstore, query, k):
     return chunks_with_references
 
 def rag_workflow(query):
-    retrieved_doc_chunks = query_chroma(vectorstore, query, k=5)
-    #retrieved_code_chunks = query_chroma(codestore, query, k=5)
+    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]) 
+    code_context = "\n\n".join([code_chunk for code_chunk, _ in retrieved_code_chunks]) 
     
-    references = "\n".join([f"[{i+1}] {ref}" for i, (_, ref) in enumerate(docs)])
-
+    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")
     
@@ -332,7 +332,7 @@ def rag_workflow(query):
 
 
 def initialize():
-    global docstore, vectorstore, chunks, llm
+    global docstore, codestore, chunks, llm
 
     code_partial_paths = ['kadi_apy/lib/']
     code_file_path = []
@@ -350,8 +350,9 @@ def initialize():
     print(f"Total number of code_chunks: {len(code_chunks)}")
     print(f"Total number of doc_chunks: {len(doc_chunks)}")
 
-    vectorstore = setup_chroma(doc_chunks, EMBEDDING_MODEL_NAME, PERSIST_DIRECTORY)
-    # codestore = setup_chroma(code_chunks, EMBEDDING_MODEL_NAME, PERSIST_DIRECTORY)
+    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)