{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from crewai import Agent, Task, Crew\n",
    "\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "llm = ChatOpenAI(\n",
    "    openai_api_base=\"https://api.groq.com/openai/v1\",\n",
    "    openai_api_key=os.environ['GROQ_API_KEY'],\n",
    "    model_name=\"llama3-8b-8192\",\n",
    "    temperature=0,\n",
    "    max_tokens=1000,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from crewai_tools import PDFSearchTool\n",
    "\n",
    "rag_tool = PDFSearchTool(pdf='/content/17.pdf',\n",
    "    config=dict(\n",
    "        llm=dict(\n",
    "            provider=\"groq\", # or google, openai, anthropic, llama2, ...\n",
    "            config=dict(\n",
    "                model=\"llama3-8b-8192\",\n",
    "                # temperature=0.5,\n",
    "                # top_p=1,\n",
    "                # stream=true,\n",
    "            ),\n",
    "        ),\n",
    "        embedder=dict(\n",
    "            provider=\"huggingface\", # or openai, ollama, ...\n",
    "            config=dict(\n",
    "                model=\"BAAI/bge-small-en-v1.5\",\n",
    "                #task_type=\"retrieval_document\",\n",
    "                # title=\"Embeddings\",\n",
    "            ),\n",
    "        ),\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rag_tool.run(\"How does exercise price determine for ESOP?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from langchain_community.tools.tavily_search import TavilySearchResults\n",
    "os.environ['TAVILY_API_KEY'] = userdata.get('TAVILY_API_KEY')\n",
    "web_search_tool = TavilySearchResults(k=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "web_search_tool.run(\"How does exercise price determine for ESOP?\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from crewai_tools  import tool\n",
    "@tool\n",
    "def router_tool(question):\n",
    "  \"\"\"Router Function\"\"\"\n",
    "  if 'ESOP' in question:\n",
    "    return 'vectorstore'\n",
    "  else:\n",
    "    return 'web_search'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Router Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Router_Agent = Agent(\n",
    "  role='Router',\n",
    "  goal='Route user question to a vectorstore or web search',\n",
    "  backstory=(\n",
    "    \"You are an expert at routing a user question to a vectorstore or web search.\"\n",
    "    \"Use the vectorstore for questions on concepta related to Retrieval-Augmented Generation.\"\n",
    "    \"You do not need to be stringent with the keywords in the question related to these topics. Otherwise, use web-search.\"\n",
    "  ),\n",
    "  verbose=True,\n",
    "  allow_delegation=False,\n",
    "  llm=llm,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Retriever_Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Retriever_Agent = Agent(\n",
    "role=\"Retriever\",\n",
    "goal=\"Use the information retrieved from the vectorstore to answer the question\",\n",
    "backstory=(\n",
    "    \"You are an assistant for question-answering tasks.\"\n",
    "    \"Use the information present in the retrieved context to answer the question.\"\n",
    "    \"You have to provide a clear concise answer.\"\n",
    "),\n",
    "verbose=True,\n",
    "allow_delegation=False,\n",
    "llm=llm,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Grader Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Grader_agent =  Agent(\n",
    "  role='Answer Grader',\n",
    "  goal='Filter out erroneous retrievals',\n",
    "  backstory=(\n",
    "    \"You are a grader assessing relevance of a retrieved document to a user question.\"\n",
    "    \"If the document contains keywords related to the user question, grade it as relevant.\"\n",
    "    \"It does not need to be a stringent test.You have to make sure that the answer is relevant to the question.\"\n",
    "  ),\n",
    "  verbose=True,\n",
    "  allow_delegation=False,\n",
    "  llm=llm,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Hallucination Grader Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hallucination_grader = Agent(\n",
    "    role=\"Hallucination Grader\",\n",
    "    goal=\"Filter out hallucination\",\n",
    "    backstory=(\n",
    "        \"You are a hallucination grader assessing whether an answer is grounded in / supported by a set of facts.\"\n",
    "        \"Make sure you meticulously review the answer and check if the response provided is in alignmnet with the question asked\"\n",
    "    ),\n",
    "    verbose=True,\n",
    "    allow_delegation=False,\n",
    "    llm=llm,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Answer Grader Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "answer_grader = Agent(\n",
    "    role=\"Answer Grader\",\n",
    "    goal=\"Filter out hallucination from the answer.\",\n",
    "    backstory=(\n",
    "        \"You are a grader assessing whether an answer is useful to resolve a question.\"\n",
    "        \"Make sure you meticulously review the answer and check if it makes sense for the question asked\"\n",
    "        \"If the answer is relevant generate a clear and concise response.\"\n",
    "        \"If the answer gnerated is not relevant then perform a websearch using 'web_search_tool'\"\n",
    "    ),\n",
    "    verbose=True,\n",
    "    allow_delegation=False,\n",
    "    llm=llm,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Router Task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "router_task = Task(\n",
    "    description=(\"Analyse the keywords in the question {question}\"\n",
    "    \"Based on the keywords decide whether it is eligible for a vectorstore search or a web search.\"\n",
    "    \"Return a single word 'vectorstore' if it is eligible for vectorstore search.\"\n",
    "    \"Return a single word 'websearch' if it is eligible for web search.\" \n",
    "    \"Do not provide any other premable or explaination.\"\n",
    "    ),\n",
    "    expected_output=(\"Give a binary choice 'websearch' or 'vectorstore' based on the question\"\n",
    "    \"Do not provide any other premable or explaination.\"),\n",
    "    agent=Router_Agent,\n",
    "    tools=[router_tool],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Retriever Task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever_task = Task(\n",
    "    description=(\"Based on the response from the router task extract information for the question {question} with the help of the respective tool.\"\n",
    "    \"Use the web_serach_tool to retrieve information from the web in case the router task output is 'websearch'.\"\n",
    "    \"Use the rag_tool to retrieve information from the vectorstore in case the router task output is 'vectorstore'.\"\n",
    "    ),\n",
    "    expected_output=(\"You should analyse the output of the 'router_task'\"\n",
    "    \"If the response is 'websearch' then use the web_search_tool to retrieve information from the web.\"\n",
    "    \"If the response is 'vectorstore' then use the rag_tool to retrieve information from the vectorstore.\"\n",
    "    \"Return a claer and consise text as response.\"),\n",
    "    agent=Retriever_Agent,\n",
    "    context=[router_task],\n",
    "   #tools=[retriever_tool],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Grader Task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grader_task = Task(\n",
    "    description=(\"Based on the response from the retriever task for the quetion {question} evaluate whether the retrieved content is relevant to the question.\"\n",
    "    ),\n",
    "    expected_output=(\"Binary score 'yes' or 'no' score to indicate whether the document is relevant to the question\"\n",
    "    \"You must answer 'yes' if the response from the 'retriever_task' is in alignment with the question asked.\"\n",
    "    \"You must answer 'no' if the response from the 'retriever_task' is not in alignment with the question asked.\"\n",
    "    \"Do not provide any preamble or explanations except for 'yes' or 'no'.\"),\n",
    "    agent=Grader_agent,\n",
    "    context=[retriever_task],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Hallucination Grader Task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hallucination_task = Task(\n",
    "    description=(\"Based on the response from the grader task for the quetion {question} evaluate whether the answer is grounded in / supported by a set of facts.\"),\n",
    "    expected_output=(\"Binary score 'yes' or 'no' score to indicate whether the answer is sync with the question asked\"\n",
    "    \"Respond 'yes' if the answer is in useful and contains fact about the question asked.\"\n",
    "    \"Respond 'no' if the answer is not useful and does not contains fact about the question asked.\"\n",
    "    \"Do not provide any preamble or explanations except for 'yes' or 'no'.\"),\n",
    "    agent=hallucination_grader,\n",
    "    context=[grader_task],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Answer grader Task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "answer_task = Task( \n",
    "    description=(\"Based on the response from the hallucination task for the quetion {question} evaluate whether the answer is useful to resolve the question.\"\n",
    "    \"If the answer is 'yes' return a clear and concise answer.\"\n",
    "    \"If the answer is 'no' then perform a 'websearch' and return the response\"),\n",
    "    expected_output=(\"Return a clear and concise response if the response from 'hallucination_task' is 'yes'.\"\n",
    "    \"Perform a web search using 'web_search_tool' and return ta clear and concise response only if the response from 'hallucination_task' is 'no'.\"\n",
    "    \"Otherwise respond as 'Sorry! unable to find a valid response'.\"),  \n",
    "    context=[hallucination_task],\n",
    "    agent=answer_grader,\n",
    "    #tools=[answer_grader_tool],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Setup the Crew"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rag_crew = Crew(\n",
    "    agents=[Router_Agent, Retriever_Agent, Grader_agent, hallucination_grader, answer_grader],\n",
    "    tasks=[router_task, retriever_task, grader_task, hallucination_task, answer_task],\n",
    "    verbose=True,\n",
    " \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs ={\"question\":\"Does the ESOP supplement the salary of an employee?\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = rag_crew.kickoff(inputs=inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env_crewai",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}