Update app.py

app.py

@@ -286,85 +286,32 @@ def format_kadi_api_doc_context(docs):
 
                    
 def rag_workflow(query):
+    """
+    RAGChain class to perform the complete RAG workflow.
+    """
+    # Assume 'llm' and 'vector_store' are already initialized instances
+    rag_chain = RAGChain(llm, vector_store)
 
+    # Step 1: Predict which library usage is relevant
+    library_usage_prediction = rag_chain.predict_library_usage(query)
+    print(f"Predicted library usage: {library_usage_prediction}")
 
-    prompt = (
-        f"""The query is: '{query}'.
-            Based on the user's query, assist them by determining which technical document they should read to interact with the software named 'Kadi4Mat'. 
-            There are three different technical documents to choose from:
-                - Document 1: Provides information on how to use a Python library to interact with the HTTP API of 'Kadi4Mat'.
-                - Document 2: Provides information on how to use a Python library to implement custom CLI commands to interact with 'Kadi4Mat'.
-    
-            Your task is to select the single most likely option. 
-                If Document 1 is the best choice, respond with 'kadi-apy python library'. 
-                If Document 2 is the best choice, respond with 'kadi-apy python cli library'. 
-            Respond with only the exact corresponding option and do not include any additional comments, explanations, or text."
-        """
-    )
-    library_usage_prediction = llm.predict(prompt)
-
-    print("METADATA PREDICTION -------------------------:", library_usage_prediction)
-    print(library_usage_prediction)
-
-    rewritten_query_response = llm.invoke(rewrite_prompt)
-    rewritten_query = rewritten_query_response.content.strip()
-    
-    print("Predicted library to search in", library_usage_prediction)
-    print("Rewritten query:", rewritten_query)
-
-    kadi_apy_docs = vector_store.similarity_search(query, k=5, filter={"usage": "doc"})
-    kadi_apy_sourcecode = vector_store.similarity_search(query, k=5, filter={"usage": library_usage_prediction})
-
-    doc_context = format_kadi_api_doc_context(kadi_apy_docs)
-    code_context = format_kadi_apy_library_context(kadi_apy_sourcecode)
-
-
-
-    
-    print("H")
-    print("::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::")
-    for doc in kadi_apy_sourcecode:
-        print(doc.metadata.get("source", "Unknown Type"))
-        print("\n")
-    print("::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::")
-    print("CODE_CONTEST\n:", code_context)
-
-    
-    print("::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::")
-    for doc in kadi_apy_docs:
-        print(doc.metadata.get("source", "Unknown Type"))
-        print("\n")
-    print("::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::")
-
-    print(doc_context)
-    print("DOC_CONTEXT\n:", code_context)
-
-    
-    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.  
-                    
+    # Step 2: Retrieve contexts (documents and code snippets)
+    doc_contexts, code_contexts = rag_chain.retrieve_contexts(query, library_usage_prediction)
+    print("Retrieved Document Contexts:", doc_contexts)
+    print("Retrieved Code Contexts:", code_contexts)
 
-                    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.
+    # Step 3: Format the contexts
+    formatted_doc_context, formatted_code_context = rag_chain.format_context(doc_contexts, code_contexts)
+    print("Formatted Document Contexts:", formatted_doc_context)
+    print("Formatted Code Contexts:", formatted_code_context)
 
-            Doc-context:
-            {doc_context}
+    # Step 4: Generate the final response
+    response = rag_chain.generate_response(query, formatted_doc_context, formatted_code_context)
+    print("Generated Response:", response)
 
-            Code-context:
-            {code_context}
-            
-            Query: 
-            {query}
-    """
+    return response
 
-    
-    response = llm.invoke(prompt)
-    return response.content
 
 
 def get_chroma_vectorstore2(embedding_model):