{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "29317d99-1605-410d-b838-fbc3469e4856",
   "metadata": {},
   "source": [
    "# Local Tool Calling Agent\n",
    "\n",
    "Here, we'll build a [tool calling agent](https://python.langchain.com/v0.1/docs/modules/agents/agent_types/tool_calling/) using local models.\n",
    "\n",
    "We'll use the [new fine-tune from Groq](https://wow.groq.com/introducing-llama-3-groq-tool-use-models/) with tool calling via Ollama:\n",
    "\n",
    "Access the model:\n",
    "\n",
    "```\n",
    "ollama pull llama3-groq-tool-use\n",
    "ollama pull llama3.1\n",
    "```\n",
    "\n",
    "And also, we'll use the Ollama partner package.\n",
    "\n",
    "This notebook accompanies the video here:\n",
    "\n",
    "https://www.youtube.com/watch?v=Nfk99Fz8H9k"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "666175b4-f056-489d-9d94-a0a666270794",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-xxx\"\n",
    "os.environ[\"TAVILY_API_KEY\"] = \"tvly-xxx\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "120c1da8-e45e-4ffa-9ac1-a536026c7e1c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m24.0\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.1.2\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install -qU langchain-ollama"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "32c0504b-007a-4af6-9976-c7294ed26b73",
   "metadata": {},
   "outputs": [],
   "source": [
    "# /// LLM ///\n",
    "\n",
    "from langchain_ollama import ChatOllama\n",
    "\n",
    "llm = ChatOllama(\n",
    "    # model=\"llama3-groq-tool-use\",\n",
    "    model=\"llama3.1\",\n",
    "    temperature=0,\n",
    ")\n",
    "\n",
    "# /// Retriever tool ///\n",
    "\n",
    "from langchain_community.document_loaders import WebBaseLoader\n",
    "from langchain_community.vectorstores import SKLearnVectorStore\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "\n",
    "# List of URLs to load documents from\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",
    "# Load documents from the URLs\n",
    "docs = [WebBaseLoader(url).load() for url in urls]\n",
    "docs_list = [item for sublist in docs for item in sublist]\n",
    "\n",
    "# Initialize a text splitter with specified chunk size and overlap\n",
    "text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(\n",
    "    chunk_size=250, chunk_overlap=0\n",
    ")\n",
    "\n",
    "# Split the documents into chunks\n",
    "doc_splits = text_splitter.split_documents(docs_list)\n",
    "\n",
    "# Add the document chunks to the \"vector store\" using OpenAIEmbeddings\n",
    "vectorstore = SKLearnVectorStore.from_documents(\n",
    "    documents=doc_splits,\n",
    "    embedding=OpenAIEmbeddings(),\n",
    ")\n",
    "retriever = vectorstore.as_retriever(k=4)\n",
    "\n",
    "\n",
    "# Define a tool, which we will connect to our agent\n",
    "def retrieve_documents(query: str) -> list:\n",
    "    \"\"\"Retrieve documents from the vector store based on the query.\"\"\"\n",
    "    return retriever.invoke(query)\n",
    "\n",
    "\n",
    "# /// Search Tool\n",
    "\n",
    "from langchain_community.tools.tavily_search import TavilySearchResults\n",
    "\n",
    "from langchain.schema import Document\n",
    "\n",
    "web_search_tool = TavilySearchResults()\n",
    "\n",
    "\n",
    "def web_search(query: str) -> str:\n",
    "    \"\"\"Run web search on the question.\"\"\"\n",
    "    web_results = web_search_tool.invoke({\"query\": query})\n",
    "    return [\n",
    "        Document(page_content=d[\"content\"], metadata={\"url\": d[\"url\"]})\n",
    "        for d in web_results\n",
    "    ]\n",
    "\n",
    "# Tool list\n",
    "tools = [retrieve_documents, web_search]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "30052f47-2b5d-46f5-9873-eb716145cda1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Annotated, List\n",
    "\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.runnables import Runnable, RunnableConfig\n",
    "from langgraph.graph.message import AnyMessage, add_messages\n",
    "from typing_extensions import TypedDict\n",
    "\n",
    "class State(TypedDict):\n",
    "    messages: Annotated[list[AnyMessage], add_messages]\n",
    "\n",
    "class Assistant:\n",
    "    def __init__(self, runnable: Runnable):\n",
    "        \"\"\"\n",
    "        Initialize the Assistant with a runnable object.\n",
    "\n",
    "        Args:\n",
    "            runnable (Runnable): The runnable instance to invoke.\n",
    "        \"\"\"\n",
    "        self.runnable = runnable\n",
    "\n",
    "    def __call__(self, state: State, config: RunnableConfig):\n",
    "        \"\"\"\n",
    "        Call method to invoke the LLM and handle its responses.\n",
    "        Re-prompt the assistant if the response is not a tool call or meaningful text.\n",
    "\n",
    "        Args:\n",
    "            state (State): The current state containing messages.\n",
    "            config (RunnableConfig): The configuration for the runnable.\n",
    "\n",
    "        Returns:\n",
    "            dict: The final state containing the updated messages.\n",
    "        \"\"\"\n",
    "        while True:\n",
    "            result = self.runnable.invoke(state)  # Invoke the LLM\n",
    "            if not result.tool_calls and (\n",
    "                not result.content\n",
    "                or isinstance(result.content, list)\n",
    "                and not result.content[0].get(\"text\")\n",
    "            ):\n",
    "                messages = state[\"messages\"] + [(\"user\", \"Respond with a real output.\")]\n",
    "                state = {**state, \"messages\": messages}\n",
    "            else:\n",
    "                break\n",
    "        return {\"messages\": result}\n",
    "\n",
    "\n",
    "# Create the primary assistant prompt template\n",
    "primary_assistant_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"You are a helpful assistant tasked with answering user questions. \"\n",
    "            \"You have access to two tools: retrieve_documents and web_search. \"\n",
    "            \"For any user questions about LLM agents, use the retrieve_documents tool to get information for a vectorstore. \"\n",
    "            \"For any other questions, such as questions about current events, use the web_search tool to get information from the web. \",\n",
    "        ),\n",
    "        (\"placeholder\", \"{messages}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Prompt our LLM and bind tools\n",
    "assistant_runnable = primary_assistant_prompt | llm.bind_tools(tools)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "40504a0b-8a99-4420-a6bf-561c62e893d1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 10812 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from langchain_core.messages import ToolMessage\n",
    "from langchain_core.runnables import RunnableLambda\n",
    "from langgraph.prebuilt import ToolNode\n",
    "\n",
    "\n",
    "def create_tool_node_with_fallback(tools: list) -> dict:\n",
    "    return ToolNode(tools).with_fallbacks(\n",
    "        [RunnableLambda(handle_tool_error)], exception_key=\"error\"\n",
    "    )\n",
    "\n",
    "\n",
    "def handle_tool_error(state: State) -> dict:\n",
    "    error = state.get(\"error\")\n",
    "    tool_calls = state[\"messages\"][-1].tool_calls\n",
    "    return {\n",
    "        \"messages\": [\n",
    "            ToolMessage(\n",
    "                content=f\"Error: {repr(error)}\\n please fix your mistakes.\",\n",
    "                tool_call_id=tc[\"id\"],\n",
    "            )\n",
    "            for tc in tool_calls\n",
    "        ]\n",
    "    }\n",
    "\n",
    "\n",
    "from IPython.display import Image, display\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from langgraph.graph import END, START, StateGraph\n",
    "from langgraph.prebuilt import tools_condition\n",
    "\n",
    "# Graph\n",
    "builder = StateGraph(State)\n",
    "\n",
    "# Define nodes: these do the work\n",
    "builder.add_node(\"assistant\", Assistant(assistant_runnable))\n",
    "builder.add_node(\"tools\", create_tool_node_with_fallback(tools))\n",
    "\n",
    "# Define edges: these determine how the control flow moves\n",
    "builder.add_edge(START, \"assistant\")\n",
    "builder.add_conditional_edges(\n",
    "    \"assistant\",\n",
    "    # If the latest message (result) from assistant is a tool call -> tools_condition routes to tools\n",
    "    # If the latest message (result) from assistant is a not a tool call -> tools_condition routes to END\n",
    "    tools_condition,\n",
    ")\n",
    "builder.add_edge(\"tools\", \"assistant\")\n",
    "\n",
    "# The checkpointer lets the graph persist its state\n",
    "memory = MemorySaver()\n",
    "react_graph = builder.compile(checkpointer=memory)\n",
    "\n",
    "# Show\n",
    "display(Image(react_graph.get_graph(xray=True).draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "43c633d5-e7a7-4b7c-8dc7-760a3b032e95",
   "metadata": {},
   "outputs": [],
   "source": [
    "import uuid\n",
    "\n",
    "def predict_react_agent_answer(example: dict):\n",
    "    \"\"\"Use this for answer evaluation\"\"\"\n",
    "\n",
    "    config = {\"configurable\": {\"thread_id\": str(uuid.uuid4())}}\n",
    "    messages = react_graph.invoke({\"messages\": (\"user\", example[\"input\"])}, config)\n",
    "    return {\"response\": messages[\"messages\"][-1].content, \"messages\": messages}\n",
    "\n",
    "\n",
    "example = {\"input\": \"Get me information about the the types of LLM agent memory?\"}\n",
    "response = predict_react_agent_answer(example)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf82fa52-9e6c-4f37-94ae-91450dac602e",
   "metadata": {},
   "source": [
    "See trace with llama3.1 here:\n",
    "\n",
    "https://smith.langchain.com/public/44d0c7dd-a756-47ad-8025-ee7ae6469ecb/r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "cd74a0b3-be40-46cd-97bf-ef9676878289",
   "metadata": {},
   "outputs": [],
   "source": [
    "example = {\"input\": \"Get me information about the current weather in SF.\"}\n",
    "response = predict_react_agent_answer(example)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cac91bf-c975-44a2-a9fd-99706fee5735",
   "metadata": {},
   "source": [
    "See trace with llama3.1 here:\n",
    "\n",
    "https://smith.langchain.com/public/7a4938e3-f94f-4e04-a162-bf592fba4643/r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74b813cb-18ed-42d8-b313-6ee56ded4bcc",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
