{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "source": [
    "# Agents in LlamaIndex\n",
    "\n",
    "This notebook is part of the [Hugging Face Agents Course](https://www.hf.co/learn/agents-course), a free Course from beginner to expert, where you learn to build Agents.\n",
    "\n",
    "![Agents course share](https://huggingface.co/datasets/agents-course/course-images/resolve/main/en/communication/share.png)\n",
    "\n",
    "## Let's install the dependencies\n",
    "\n",
    "We will install the dependencies for this unit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index datasets llama-index-callbacks-arize-phoenix llama-index-vector-stores-chroma llama-index-llms-huggingface-api -U -q"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And, let's log in to Hugging Face to use serverless Inference APIs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import login\n",
    "\n",
    "login()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "source": [
    "## Initialising agents\n",
    "\n",
    "Let's start by initialising an agent. We will use the basic `AgentWorkflow` class to create an agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI\n",
    "from llama_index.core.agent.workflow import AgentWorkflow, ToolCallResult, AgentStream\n",
    "\n",
    "\n",
    "def add(a: int, b: int) -> int:\n",
    "    \"\"\"Add two numbers\"\"\"\n",
    "    return a + b\n",
    "\n",
    "\n",
    "def subtract(a: int, b: int) -> int:\n",
    "    \"\"\"Subtract two numbers\"\"\"\n",
    "    return a - b\n",
    "\n",
    "\n",
    "def multiply(a: int, b: int) -> int:\n",
    "    \"\"\"Multiply two numbers\"\"\"\n",
    "    return a * b\n",
    "\n",
    "\n",
    "def divide(a: int, b: int) -> int:\n",
    "    \"\"\"Divide two numbers\"\"\"\n",
    "    return a / b\n",
    "\n",
    "\n",
    "llm = HuggingFaceInferenceAPI(model_name=\"Qwen/Qwen2.5-Coder-32B-Instruct\")\n",
    "\n",
    "agent = AgentWorkflow.from_tools_or_functions(\n",
    "    tools_or_functions=[subtract, multiply, divide, add],\n",
    "    llm=llm,\n",
    "    system_prompt=\"You are a math agent that can add, subtract, multiply, and divide numbers using provided tools.\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we can run the agent and get the response and reasoning behind the tool calls."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "handler = agent.run(\"What is (2 + 2) * 2?\")\n",
    "async for ev in handler.stream_events():\n",
    "    if isinstance(ev, ToolCallResult):\n",
    "        print(\"\")\n",
    "        print(\"Called tool: \", ev.tool_name, ev.tool_kwargs, \"=>\", ev.tool_output)\n",
    "    elif isinstance(ev, AgentStream):  # showing the thought process\n",
    "        print(ev.delta, end=\"\", flush=True)\n",
    "\n",
    "resp = await handler\n",
    "resp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In a similar fashion, we can pass state and context to the agent.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AgentOutput(response=ChatMessage(role=<MessageRole.ASSISTANT: 'assistant'>, additional_kwargs={}, blocks=[TextBlock(block_type='text', text='Your name is Bob.')]), tool_calls=[], raw={'id': 'chatcmpl-B5sDHfGpSwsVyzvMVH8EWokYwdIKT', 'choices': [{'delta': {'content': None, 'function_call': None, 'refusal': None, 'role': None, 'tool_calls': None}, 'finish_reason': 'stop', 'index': 0, 'logprobs': None}], 'created': 1740739735, 'model': 'gpt-4o-2024-08-06', 'object': 'chat.completion.chunk', 'service_tier': 'default', 'system_fingerprint': 'fp_eb9dce56a8', 'usage': None}, current_agent_name='Agent')"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from llama_index.core.workflow import Context\n",
    "\n",
    "ctx = Context(agent)\n",
    "\n",
    "response = await agent.run(\"My name is Bob.\", ctx=ctx)\n",
    "response = await agent.run(\"What was my name again?\", ctx=ctx)\n",
    "response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating RAG Agents with QueryEngineTools\n",
    "\n",
    "Let's now re-use the `QueryEngine` we defined in the [previous unit on tools](/tools.ipynb) and convert it into a `QueryEngineTool`. We will pass it to the `AgentWorkflow` class to create a RAG agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chromadb\n",
    "\n",
    "from llama_index.core import VectorStoreIndex\n",
    "from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI\n",
    "from llama_index.embeddings.huggingface_api import HuggingFaceInferenceAPIEmbedding\n",
    "from llama_index.core.tools import QueryEngineTool\n",
    "from llama_index.vector_stores.chroma import ChromaVectorStore\n",
    "\n",
    "# Create a vector store\n",
    "db = chromadb.PersistentClient(path=\"./alfred_chroma_db\")\n",
    "chroma_collection = db.get_or_create_collection(\"alfred\")\n",
    "vector_store = ChromaVectorStore(chroma_collection=chroma_collection)\n",
    "\n",
    "# Create a query engine\n",
    "embed_model = HuggingFaceInferenceAPIEmbedding(model_name=\"BAAI/bge-small-en-v1.5\")\n",
    "llm = HuggingFaceInferenceAPI(model_name=\"Qwen/Qwen2.5-Coder-32B-Instruct\")\n",
    "index = VectorStoreIndex.from_vector_store(\n",
    "    vector_store=vector_store, embed_model=embed_model\n",
    ")\n",
    "query_engine = index.as_query_engine(llm=llm)\n",
    "query_engine_tool = QueryEngineTool.from_defaults(\n",
    "    query_engine=query_engine,\n",
    "    name=\"personas\",\n",
    "    description=\"descriptions for various types of personas\",\n",
    "    return_direct=False,\n",
    ")\n",
    "\n",
    "# Create a RAG agent\n",
    "query_engine_agent = AgentWorkflow.from_tools_or_functions(\n",
    "    tools_or_functions=[query_engine_tool],\n",
    "    llm=llm,\n",
    "    system_prompt=\"You are a helpful assistant that has access to a database containing persona descriptions. \",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And, we can once more get the response and reasoning behind the tool calls."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "handler = query_engine_agent.run(\n",
    "    \"Search the database for 'science fiction' and return some persona descriptions.\"\n",
    ")\n",
    "async for ev in handler.stream_events():\n",
    "    if isinstance(ev, ToolCallResult):\n",
    "        print(\"\")\n",
    "        print(\"Called tool: \", ev.tool_name, ev.tool_kwargs, \"=>\", ev.tool_output)\n",
    "    elif isinstance(ev, AgentStream):  # showing the thought process\n",
    "        print(ev.delta, end=\"\", flush=True)\n",
    "\n",
    "resp = await handler\n",
    "resp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating multi-agent systems\n",
    "\n",
    "We can also create multi-agent systems by passing multiple agents to the `AgentWorkflow` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.agent.workflow import (\n",
    "    AgentWorkflow,\n",
    "    ReActAgent,\n",
    ")\n",
    "\n",
    "\n",
    "# Define some tools\n",
    "def add(a: int, b: int) -> int:\n",
    "    \"\"\"Add two numbers.\"\"\"\n",
    "    return a + b\n",
    "\n",
    "\n",
    "def subtract(a: int, b: int) -> int:\n",
    "    \"\"\"Subtract two numbers.\"\"\"\n",
    "    return a - b\n",
    "\n",
    "\n",
    "# Create agent configs\n",
    "# NOTE: we can use FunctionAgent or ReActAgent here.\n",
    "# FunctionAgent works for LLMs with a function calling API.\n",
    "# ReActAgent works for any LLM.\n",
    "calculator_agent = ReActAgent(\n",
    "    name=\"calculator\",\n",
    "    description=\"Performs basic arithmetic operations\",\n",
    "    system_prompt=\"You are a calculator assistant. Use your tools for any math operation.\",\n",
    "    tools=[add, subtract],\n",
    "    llm=llm,\n",
    ")\n",
    "\n",
    "query_agent = ReActAgent(\n",
    "    name=\"info_lookup\",\n",
    "    description=\"Looks up information about XYZ\",\n",
    "    system_prompt=\"Use your tool to query a RAG system to answer information about XYZ\",\n",
    "    tools=[query_engine_tool],\n",
    "    llm=llm,\n",
    ")\n",
    "\n",
    "# Create and run the workflow\n",
    "agent = AgentWorkflow(agents=[calculator_agent, query_agent], root_agent=\"calculator\")\n",
    "\n",
    "# Run the system\n",
    "handler = agent.run(user_msg=\"Can you add 5 and 3?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "async for ev in handler.stream_events():\n",
    "    if isinstance(ev, ToolCallResult):\n",
    "        print(\"\")\n",
    "        print(\"Called tool: \", ev.tool_name, ev.tool_kwargs, \"=>\", ev.tool_output)\n",
    "    elif isinstance(ev, AgentStream):  # showing the thought process\n",
    "        print(ev.delta, end=\"\", flush=True)\n",
    "\n",
    "resp = await handler\n",
    "resp"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}