crewAI icon indicating copy to clipboard operation
crewAI copied to clipboard

RAG integration example

Open ryanmac opened this issue 1 year ago • 1 comments

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:

  1. 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.

  2. 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.

  3. 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.

  4. 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.

  5. 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.

  6. Error Handling and Logging: Enhanced error handling and logging mechanisms to ensure smooth operation and ease of debugging when using the RAG model.

  7. 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.

  8. 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:

  1. Import RAGModelHandler: Update your imports to include the RAGModelHandler from the crewai 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
    
  2. Instantiate RAGModelHandler: Create an instance of RAGModelHandler in your script. This instance will manage interactions with the RAG model.

    rag_handler = RAGModelHandler()
    
  3. 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"]
    
  4. 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 ...
    )
    
  5. 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 the FileHandler 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.

ryanmac avatar Feb 06 '24 00:02 ryanmac

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)

deesatzed avatar Feb 07 '24 13:02 deesatzed

We ended up adding RAG as tool on the tools package

joaomdmoura avatar Mar 31 '24 23:03 joaomdmoura