Spaces:
Sleeping
Sleeping
| 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) | |