crewAI
crewAI copied to clipboard
RAG integration example
Title: Implement RAG Model Support in CrewAI
Description: This commit introduces the integration of the Retrieval-Augmented Generation (RAG) model into the CrewAI framework, significantly enhancing the capabilities of AI agents in handling complex tasks and queries. Key changes and features include:
-
RAG Model Handler Integration: Added
RAGModelHandler
class to manage interactions with the RAG model. This handler facilitates querying the RAG model and processing its responses, enabling agents to leverage external knowledge sources effectively. -
Agent Class Enhancement: Updated the
Agent
class to include the RAG model handler as an attribute. This allows each agent to access and utilize the RAG model for tasks requiring extensive knowledge retrieval and generation. -
File Handling for RAG Queries: Integrated a
FileHandler
within the agent class to manage file paths and content. This feature enables agents to access and refer to specific files, which is crucial for RAG model queries that might require context from external documents. -
Task Execution with RAG Support: Modified the
Task
class to support execution using the RAG model. Tasks can now be designed to utilize the RAG model's capabilities, allowing for more sophisticated and informed responses. -
Backward Compatibility and Flexibility: Ensured that the new RAG model features are backward compatible. Existing users of CrewAI can continue using the framework without the RAG model. The integration is designed to be optional but offers powerful capabilities when utilized.
-
Error Handling and Logging: Enhanced error handling and logging mechanisms to ensure smooth operation and ease of debugging when using the RAG model.
-
Documentation and Examples: Updated the documentation to reflect the new RAG model features and provided examples demonstrating how to use these enhancements in various scenarios.
-
Comprehensive Testing: Added a suite of new tests to validate the functionality of the RAG model integration. These tests cover various scenarios and edge cases to ensure the reliability and robustness of the RAG model implementation within CrewAI.
This implementation aims to broaden the scope of CrewAI's applications, making it a more versatile and powerful tool for developing sophisticated AI agents capable of handling a wide range of tasks with enhanced knowledge retrieval and generation capabilities.
To update your CrewAI script to utilize the new RAG (Retrieval-Augmented Generation) model capabilities, follow these steps:
-
Import RAGModelHandler: Update your imports to include the
RAGModelHandler
from thecrewai
package. This will allow you to create an instance of the RAG model handler in your script.from crewai import Agent, Task, Crew, Process, RAGModelHandler
-
Instantiate RAGModelHandler: Create an instance of
RAGModelHandler
in your script. This instance will manage interactions with the RAG model.rag_handler = RAGModelHandler()
-
Specify File Paths: Define a list of file paths that your agents might need to access for RAG queries. These files will provide the necessary context for the RAG model to generate informed responses.
file_paths = ["./requirements.md"]
-
Integrate RAGModelHandler with Tasks: For each task that requires the RAG model, add the
rag_model_handler=rag_handler
argument to the task's constructor. This links the task with the RAG model, enabling it to leverage the model's capabilities.task_Example = Task( description="Your task description here", agent=agent_Example, rag_model_handler=rag_handler, # Add this line to integrate RAG # ... other task parameters ... )
-
Update Agent File Handling: If your agents need to handle files for RAG queries, ensure that the
Agent
class instances are updated to manage file paths. This can be done by adding file paths to the agent's attributes or by using theFileHandler
class.agent_Example = Agent( # ... agent parameters ... file_paths=file_paths # Add file paths to the agent )
By following these steps, your CrewAI script will be updated to harness the power of the RAG model, enhancing the capabilities of your AI agents with advanced knowledge retrieval and generation. This will allow your agents to produce more informed and context-aware responses, especially for complex tasks requiring external knowledge sources.
Why not use a simple wrapper like this?
import json
import re
import pytest
import os
from crewai import Agent, Task, Crew, Process
from crewai.agents.cache import CacheHandler
import openai
from langchain.llms import Ollama
from langchain.tools import DuckDuckGoSearchRun
from langchain.agents import load_tools
from langchain import hub
from langchain.agents import Tool
from langchain.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings.openai import OpenAIEmbeddings
from langchain.document_loaders import PyPDFLoader
from langchain.schema import Document
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.text_splitter import CharacterTextSplitter
from langchain_community.document_loaders import TextLoader
from langchain_community.embeddings.sentence_transformer import (
SentenceTransformerEmbeddings,
)
from langchain.chat_models import ChatOpenAI
from langchain.schema import HumanMessage
human_tools = load_tools(["human"])
from langchain_community.vectorstores import Chroma
from langchain.document_loaders import PyPDFLoader, DirectoryLoader, PDFMinerLoader
#!pip install pdfminer.six
DDG_tool = DuckDuckGoSearchRun() wiki_tool = WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper())
from langchain_openai import ChatOpenAI G_api ="AIzaSyABLABLABLA" llmG = ChatGoogleGenerativeAI(model="gemini-pro", verbose = True, temperature = 0.16, google_api_key=G_api) llmOI = ChatOpenAI(model_name="gpt-4-0125-preview", temperature=0.1) embedding_function = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
load the document and split it into chunks
loader = TextLoader("datasources/MyTextFile1.txt") documents = loader.load() text_splitter = RecursiveCharacterTextSplitter(chunk_size=400, chunk_overlap=25) docs = text_splitter.split_documents(documents) db = Chroma.from_documents(docs, embedding_function)
persist_directory = "dbpdf"
Path2File = "datasources/MyPDFFile1.pdf"
loader = PDFMinerLoader(Path2File)
documentsP = loader.load()
split it into chunks
text_splitterP = RecursiveCharacterTextSplitter(chunk_size=400, chunk_overlap=25) docsP = text_splitterP.split_documents(documentsP)
load it into Chroma
dbP = Chroma.from_documents(docsP, embedding_function,persist_directory=persist_directory) dbP.persist()
Prepare your RAG prompt format. Make sure to adjust this according to your needs.
TXT_FILE_rag_prompt = "Summarize information about the question. Context: {context}, Question: {question}"
Create the RAG tool with the Chroma vector store.
TXT_FILE_rag_tool = Tool( name="RAG", func=lambda question: rag_prompt.format(context=db.similarity_search(question), question=question), description="Generates responses based on retrieved context and the given question.", vectorstore=db )
Prepare your RAG prompt format. Make sure to adjust this according to your needs.
PDF_FILE_rag_prompt = "Summarize information about the question. Context: {context}, Question: {question}"
Create the RAG tool with the Chroma vector store.
PDF_FILE_rag_tool = Tool( name="RAGpdf", func=lambda question: rag_prompt.format(context=db.similarity_search(question), question=question), description="Generates responses based on retrieved context and the given question.", vectorstore=dbP )
researcher = Agent ( role="Master researcher", goal='Find the best bla bla bla', backstory="""The seasoned bla bla bla""", verbose=True, allow_delegation=False, llm=llmOI, tools=[TXT_FILE_rag_tool] )
writer = Agent( role='Senior Writer', goal = "Expertly rewrites the bla bla bla.""", verbose=True, allow_delegation=True, llm=llmG, tools=[PDF_FILE_rag_tool] )
Note_bla = """bla bla bla note example"""
Create tasks for your agents
task1 = Task( description=f"""task bla bla about : "{Note_bla}" , that will help best bla bla.""", agent=researcher )
task2 = Task( description="""task bla bla about {} bla bla """.format(Note_bla), agent=writer )
Instantiate your crew with a sequential process
crew = Crew( agents=[researcher, writer], tasks=[task1, task2], verbose=2, # You can set it to 1 or 2 to different logging levels manager_llm=llmOI, process = Process.hierarchical ) result = crew.kickoff()
print(result)
We ended up adding RAG as tool on the tools package