{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Install dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install -U --quiet 'crewai[tools]' aisuite"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Set environment variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = getpass.getpass(\"Enter your OpenAI API key: \")\n",
    "\n",
    "# Apply a patch to allow nested asyncio loops in Jupyter\n",
    "import nest_asyncio\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create Crew"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Inserting batches in chromadb: 100%|██████████| 1/1 [00:01<00:00,  1.36s/it]\n",
      "Inserting batches in chromadb: 100%|██████████| 1/1 [00:01<00:00,  1.12s/it]\n",
      "Inserting batches in chromadb: 100%|██████████| 1/1 [00:01<00:00,  1.19s/it]\n"
     ]
    }
   ],
   "source": [
    "# Importing Crew related components\n",
    "# Importing CrewAI Flow related components\n",
    "# Importing CrewAI Tools\n",
    "from crewai import Agent, Task, Crew\n",
    "from crewai.flow.flow import Flow, listen, start\n",
    "from crewai_tools import WebsiteSearchTool\n",
    "\n",
    "# Importing AI Suite for adhoc LLM calls and Pydantic\n",
    "from pydantic import BaseModel\n",
    "import aisuite as ai\n",
    "\n",
    "urls = [\n",
    "    \"https://lilianweng.github.io/posts/2023-06-23-agent/\",\n",
    "    \"https://lilianweng.github.io/posts/2023-03-15-prompt-engineering/\",\n",
    "    \"https://lilianweng.github.io/posts/2023-10-25-adv-attack-llm/\",\n",
    "]\n",
    "\n",
    "research_agent = Agent(\n",
    "    role=\"You are a helpful assistant that can answer questions about the web.\",\n",
    "    goal=\"Answer the user's question.\",\n",
    "    backstory=\"You have access to a vast knowledge base of information from the web.\",\n",
    "    tools=[\n",
    "      WebsiteSearchTool(website=urls[0]),\n",
    "      WebsiteSearchTool(website=urls[1]),\n",
    "      WebsiteSearchTool(website=urls[2]),\n",
    "    ],\n",
    "    llm=\"gpt-4o-mini\",\n",
    ")\n",
    "\n",
    "task = Task(\n",
    "  description=\"Answer the following question: {question}\",\n",
    "  expected_output=\"A detailed and accurate answer to the user's question.\",\n",
    "  agent=research_agent,\n",
    ")\n",
    "\n",
    "crew = Crew(\n",
    "    agents=[research_agent],\n",
    "    tasks=[task],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Creating State"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class QAState(BaseModel):\n",
    "  \"\"\"\n",
    "  State for the documentation flow\n",
    "  \"\"\"\n",
    "  question: str = \"What does Lilian Weng say about the types of agent memory?\"\n",
    "  improved_question: str = \"\"\n",
    "  answer: str = \"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Creating Flow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class QAFlow(Flow[QAState]):\n",
    "  @start()\n",
    "  def rewrite_question(self):\n",
    "    print(f\"# Rewriting question: {self.state.question}\")\n",
    "    client = ai.Client()\n",
    "    messages = [\n",
    "        {\n",
    "          \"role\": \"system\",\n",
    "          \"content\": f\"\"\"Look at the input and try to reason about the underlying semantic intent / meaning.\n",
    "            Here is the initial question:\n",
    "            -------\n",
    "            {self.state.question}\n",
    "            -------\n",
    "            Formulate an improved question:\"\"\"\n",
    "        }\n",
    "    ]\n",
    "\n",
    "    response = client.chat.completions.create(\n",
    "        model=\"openai:gpt-4o-mini\",\n",
    "        messages=messages,\n",
    "        temperature=0.3\n",
    "    )\n",
    "\n",
    "    print(response)\n",
    "\n",
    "    improved_question = response.choices[0].message.content\n",
    "    self.state.improved_question = improved_question\n",
    "\n",
    "  @listen(rewrite_question)\n",
    "  def answer_question(self):\n",
    "    print(f\"# Answering question: {self.state.improved_question}\")\n",
    "    result = crew.kickoff(inputs={'question': self.state.improved_question})\n",
    "    self.state.answer = result.raw\n",
    "    return result\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plotting Flow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Plot saved as crewai_flow.html\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"400\"\n",
       "            src=\"crewai_flow.html\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x2a8496cd0>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "flow = QAFlow()\n",
    "flow.plot()\n",
    "\n",
    "# Display the flow visualization using HTML\n",
    "from IPython.display import IFrame\n",
    "IFrame(src='crewai_flow.html', width='100%', height=400)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Kicking off Flow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# Rewriting question: What does Lilian Weng say about the types of agent memory?\n",
      "ChatCompletion(id='chatcmpl-AemJZZeKKjqXr91KZO8zO8a1xcPZl', choices=[Choice(finish_reason='stop', index=0, logprobs=None, message=ChatCompletionMessage(content='What insights does Lilian Weng provide regarding the different types of agent memory?', refusal=None, role='assistant', audio=None, function_call=None, tool_calls=None))], created=1734282205, model='gpt-4o-mini-2024-07-18', object='chat.completion', service_tier=None, system_fingerprint='fp_6fc10e10eb', usage=CompletionUsage(completion_tokens=16, prompt_tokens=56, total_tokens=72, completion_tokens_details=CompletionTokensDetails(accepted_prediction_tokens=0, audio_tokens=0, reasoning_tokens=0, rejected_prediction_tokens=0), prompt_tokens_details=PromptTokensDetails(audio_tokens=0, cached_tokens=0)))\n",
      "# Answering question: What insights does Lilian Weng provide regarding the different types of agent memory?\n",
      "==========\n",
      "Lilian Weng provides insights on different types of agent memory in her work on LLM-powered autonomous agents. She categorizes memory into three main types:\n",
      "\n",
      "1. **Short-term Memory**: This refers to the agent's ability to utilize in-context learning. It allows the agent to process and retain information temporarily for immediate tasks.\n",
      "\n",
      "2. **Long-term Memory**: This memory type allows the agent to retain and recall information over extended periods. It often employs external vector stores for infinite retention and quick retrieval of data.\n",
      "\n",
      "Weng elaborates on the general concept of memory, drawing parallels with human memory systems. She delineates various categories found in human cognitive science:\n",
      "\n",
      "- **Sensory Memory**: The ability to retain impressions of sensory information after stimuli have ceased, typically lasting a few seconds. It includes subcategories like iconic (visual) and echoic (auditory) memory.\n",
      "  \n",
      "- **Short-Term Memory (STM)** or **Working Memory**: This stores information necessary for immediate cognitive tasks, with a capacity of about seven items and a duration of 20-30 seconds.\n",
      "  \n",
      "- **Long-Term Memory (LTM)**: Capable of holding vast amounts of information for long periods ranging from days to decades, it is divided into:\n",
      "  - **Explicit/Declarative Memory**: Recallable memories of facts and events (episodic and semantic).\n",
      "  - **Implicit/Procedural Memory**: Unconscious memory for skills and routines.\n",
      "\n",
      "In analogy, sensory memory can be likened to the initial processing of raw input data, while other forms of memory facilitate more complex actions and decision-making in agents, such as retaining key operational experiences and making informed choices based on past interactions.\n"
     ]
    }
   ],
   "source": [
    "result = flow.kickoff()\n",
    "print(\"=\" * 10)\n",
    "print(result)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
