testing vector db and rag

Author: VergaJU

hissmed_literature_assistant.py

@@ -32,36 +32,27 @@ def hisscat_download_literature(user_message, cat):
     return user_message
 
 
-@tool
-def rag_tool(query, cat) -> str:
-    """
-    run this tools whenever message starts with Please tell me about
-
-    Process a query using Retrieval Augmented Generation (RAG) and cite documents from memory.
-    """
-    # Initialize the Agent and Memory
-    agent = Agent(cat)
-    memory = Memory(cat)
-    settings = cat.mad_hatter.get_plugin().load_settings()
-    top_references = settings['top_references']
-    # Retrieve relevant documents based on the query
-    relevant_documents = memory.retrieve_documents(query, max_results=top_references)
+@hook
+def before_cat_sends_message(msg, cat):
+    # Extract the user query from the message
+    user_query = msg['content']
+    
+    # Retrieve the relevant context documents from memory
+    context_documents = cat.memory.get('relevant_documents', '')  # Adjust this to your memory retrieval method
 
-    # Prepare the prompt for the LLM
-    prompt = f"You are a research assistant that, based on the following documents, answer the query.
-                If the information required is not in the documents answer 'I am sorry, I need more information'.
-                The query: '{query}'.\n\n"
-    for doc in relevant_documents:
-        prompt += f"Document: {doc.title}\nContent: {doc.content}\n\n"
+    # Create the final prompt
+    prompt = f"""You are a research assistant who wants to give accurate answers considering only the provided context. Write a paragraph with the answer and cite the documents you used. Don't invent anything, if you can't answer just say 'I am sorry I don't have enough information'.
 
-    # Generate the response using the LLM
-    response = agent.generate_response(prompt)
+Documents:
+{context_documents}
 
-    # Add citations to the response
-    citations = [f"[{i+1}] {doc.title}" for i, doc in enumerate(relevant_documents)]
-    response += "\n\nCitations:\n" + "\n".join(citations)
+Question:
+{user_query}
+"""
 
-    return response
+    # Update the message content with the new prompt
+    msg['content'] = prompt
+    return msg
 
 
 @tool
@@ -80,6 +71,7 @@ def delete_declarative_memory(cat):
         return "Failed to delete entries in declarative memory."
 
 
+
 def load_files_to_memory(cat, directory='./literature'):
     """
     Load supported files from the specified directory into the Cat's declarative memory.
@@ -90,12 +82,20 @@ def load_files_to_memory(cat, directory='./literature'):
     """
     # Initialize the Rabbit Hole
     rabbit_hole = RabbitHole(cat)
-
+    settings = cat.mad_hatter.get_plugin().load_settings()
+    chunk_size = settings['chunk_size']
+    chunk_overlap = settings['chunk_overlap']
     # Iterate through files in the directory
     for filename in os.listdir(directory):
         if filename.endswith(('.txt', '.md', '.pdf', '.html')):
             file_path = os.path.join(directory, filename)
-            # Load the file into the declarative memory
-            with open(file_path, 'rb') as file:
-                rabbit_hole.insert_memory(file)
-            print(f"{datetime.now()} Loaded {filename} into memory.")
+            try:
+                # Load the file into the declarative memory
+                with open(file_path, 'rb') as file:
+                    rabbit_hole.ingest_file(file, chunk_size=chunk_size, chunk_overlap=chunk_overlap)
+                log.info(f"{datetime.now()} Loaded {filename} into memory.")
+            except Exception as e:
+                log.error(f"Error loading {filename}: {e}")
+
+
+