Create ragchain.py

ragchain.py

@@ -0,0 +1,86 @@
+class RAGChain:
+    
+    def __init__(self, llm, vector_store):
+        """
+        Initialize the RAGChain with an LLM instance and a vector store.
+        """
+        self.llm = llm
+        self.vector_store = vector_store
+
+    def predict_library_usage(self, query):
+        """
+        Use the LLM to predict the relevant library for the user's query.
+        """
+        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."
+            """
+        )
+        return self.llm.predict(prompt)
+
+    def retrieve_contexts(self, query, library_usage_prediction):
+        """
+        Retrieve relevant documents and source code based on the query and library usage prediction.
+        """
+        doc_contexts = self.vector_store.similarity_search(query, k=5, filter={"usage": "doc"})
+        code_contexts = self.vector_store.similarity_search(query, k=5, filter={"usage": library_usage_prediction})
+        
+        return doc_contexts, code_contexts
+
+    def format_context(self, doc_contexts, code_contexts):
+        """
+        Format the retrieved document and code contexts.
+        """
+        doc_context = format_kadi_api_doc_context(doc_contexts)
+        code_context = format_kadi_apy_library_context(code_contexts)
+        
+        return doc_context, code_context
+
+    def generate_response(self, query, doc_context, code_context):
+        """
+        Generate a response using the retrieved contexts and the LLM.
+        """
+        prompt = f"""You are an expert python developer. You are assisting in generating code for users who want to programmatically 
+                    make use of api of a software. There is a specific Python library named "kadiAPY" designed to interact with 
+                    the API of the software. It provides an object-oriented approach for interfacing with the API.
+
+                    You are given "Documentation Snippets" and "Code Snippets"
+                        "Documentation Snippets:" Contains a collection of potentially useful snippets, including code examples and documentation excerpts of "kadiAPY"
+                        "Code Snippets:" Contains potentially useful snippets from the source code of "kadiAPY"
+                        
+                    Based on the retrieved snippets and the guidelines answer the "query".  
+                        
+                        General Guidelines:
+                        - If no related information is found from the snippets 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.
+                        
+                Documentation Snippets:
+                {doc_context}
+                Code Snippets:
+                {code_context}
+                
+                Query: 
+                {query}
+        """
+        return self.llm.invoke(prompt).content
+
+    def rag_workflow(self, query):
+        """
+        Complete the RAG workflow: predict library usage, retrieve contexts, and generate a response.
+        """
+        library_usage_prediction = self.predict_library_usage(query)
+        
+        doc_contexts, code_contexts = self.retrieve_contexts(query, library_usage_prediction)
+        doc_context, code_context = self.format_context(doc_contexts, code_contexts)
+        
+        return self.generate_response(query, doc_context, code_context)