diff --git a/lab-agent-vector-store-completed.ipynb b/lab-agent-vector-store-completed.ipynb new file mode 100644 index 0000000..0f143e4 --- /dev/null +++ b/lab-agent-vector-store-completed.ipynb @@ -0,0 +1,624 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "efad6479-7fb1-40c4-8610-8963f8de307f", + "metadata": {}, + "source": [ + "# Lab | Agent & Vector store\n" + ] + }, + { + "cell_type": "markdown", + "id": "25fb0c6d", + "metadata": {}, + "source": [ + "
\n", + "\n", + "## Intro\n", + "\n", + "In this lab you'll build an AI agent that knows when to consult *different* knowledge bases to answer a question \u2014 instead of relying on a single source of truth.\n", + "\n", + "Here's what to expect:\n", + "\n", + "1. **Follow a full worked demo** \u2014 We'll walk through every step together: ingesting the *state of the union* speech and the *Ruff* docs into two vector stores, wrapping each in a `RetrievalQA` tool, and building an agent that picks the right tool (or both!) depending on the question.\n", + "\n", + "2. **Replicate it yourself with a new dataset** \u2014 Then, you'll swap in a dataset of your choice and rebuild the same pipeline, adapting the prompts and tools along the way.\n", + "\n", + "By the end of this lab, you'll understand how to build multi-source AI agents and be able to apply the pattern to your own datasets." + ] + }, + { + "cell_type": "markdown", + "id": "68b24990", + "metadata": {}, + "source": [ + "
\n", + "\n", + "## Combine agents and vector stores\n", + "\n", + "Let's get into the demo. We'll wrap each vector store in a `RetrievalQA` chain and hand it to an agent as a `Tool`. The agent then decides, at each step, which tool to call based purely on its description \u2014 this is what lets it route between multiple knowledge sources.\n", + "\n", + "There are two flavors of this pattern, both of which we'll try below:\n", + "\n", + "- **Agent as reasoner** \u2014 the agent calls a tool and can keep reasoning afterward (e.g. to combine results from multiple sources).\n", + "- **Agent as router** (`return_direct=True`) \u2014 the agent just picks the right tool and returns its answer immediately, no extra reasoning." + ] + }, + { + "cell_type": "markdown", + "id": "9b22020a", + "metadata": {}, + "source": [ + "
\n", + "\n", + "## Install dependencies\n", + "\n", + "Uncomment and run the cells below to install the dependencies required for this notebook.\n", + "\n", + "Tip: Use a virtual environment to keep this project's dependencies isolated from your system Python and other projects." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12df91e0-7322-43c2-96cd-0159d017a1e1", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "# !pip install \"langchain<0.3\" \"langchain-core<0.3\" \"langchain-community<0.3\" \"langchain-openai<0.2\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "936c402a", + "metadata": {}, + "outputs": [], + "source": [ + "# !pip install python-dotenv==1.2.2 chromadb==1.5.9 beautifulsoup4==4.15.0" + ] + }, + { + "cell_type": "markdown", + "id": "565a6825", + "metadata": {}, + "source": [ + "
\n", + "\n", + "## Initial Setup\n", + "\n", + "Before building anything, we need to load our API credentials, instantiate the LLM we'll use throughout the notebook, and locate the sample document we'll be querying." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2868cff1-a3ef-426f-b8d7-7fe89047a5b7", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "from langchain.chains import RetrievalQA\n", + "from langchain_community.vectorstores import Chroma\n", + "from langchain_openai import OpenAI, OpenAIEmbeddings\n", + "from langchain_text_splitters import CharacterTextSplitter\n", + "from langchain_community.document_loaders import TextLoader" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "aeff3a2d-b0cc-429f-a34c-4381d71f1a5f", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "import os\n", + "from dotenv import load_dotenv, find_dotenv\n", + "_ = load_dotenv(find_dotenv())\n", + "\n", + "OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "14331eec-fd46-42e0-b6e7-adaf21824ef7", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "llm = OpenAI(temperature=0, api_key=OPENAI_API_KEY)" + ] + }, + { + "cell_type": "markdown", + "id": "b3f2857f", + "metadata": {}, + "source": [ + "
\n", + "\n", + "Now let's ingest the state of the union speech: load the raw text, split it into manageable chunks, embed those chunks, and store them in a Chroma vector store." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bbb327d2", + "metadata": {}, + "outputs": [], + "source": [ + "doc_path = \"./datasets/state_of_the_union.txt\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f2675861", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "loader = TextLoader(doc_path)\n", + "documents = loader.load()\n", + "text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)\n", + "texts = text_splitter.split_documents(documents)\n", + "\n", + "embeddings = OpenAIEmbeddings(api_key=OPENAI_API_KEY)\n", + "\n", + "docsearch = Chroma.from_documents(texts, embeddings, collection_name=\"state-of-union\")" + ] + }, + { + "cell_type": "markdown", + "id": "fe780cdc", + "metadata": {}, + "source": [ + "
\n", + "\n", + "## Adding a second knowledge source\n", + "\n", + "To show how an agent can route between multiple tools, let's add a second vector store \u2014 this time built from the Ruff FAQ web page instead of a local file." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bc5403d4", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "state_of_union = RetrievalQA.from_chain_type(\n", + " llm=llm, chain_type=\"stuff\", retriever=docsearch.as_retriever()\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1431cded", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "from langchain_community.document_loaders import WebBaseLoader" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "915d3ff3", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "loader = WebBaseLoader(\"https://beta.ruff.rs/docs/faq/\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "96a2edf8", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "docs = loader.load()\n", + "ruff_texts = text_splitter.split_documents(docs)\n", + "ruff_db = Chroma.from_documents(ruff_texts, embeddings, collection_name=\"ruff\")\n", + "ruff = RetrievalQA.from_chain_type(\n", + " llm=llm, chain_type=\"stuff\", retriever=ruff_db.as_retriever()\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "c0a6c031", + "metadata": {}, + "source": [ + "
\n", + "\n", + "## Create the Agent\n", + "\n", + "With both `RetrievalQA` chains ready, we wrap each one in a `Tool` (giving it a name and a description the agent will use to decide when to call it), then hand both tools to an agent." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "eb142786", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "# Import things that are needed generically\n", + "from langchain.agents import AgentType, Tool, initialize_agent\n", + "from langchain_openai import OpenAI" + ] + }, + { + "cell_type": "markdown", + "id": "6fa28c55", + "metadata": {}, + "source": [ + "
\n", + "\n", + "Let's try it out \u2014 first with a question only the state of the union tool can answer, then one only Ruff can answer. Watch the verbose output to see which tool the agent picks each time." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "850bc4e9", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "tools = [\n", + " Tool(\n", + " name=\"State of Union QA System\",\n", + " func=state_of_union.run,\n", + " description=\"useful for when you need to answer questions about the most recent state of the union address. Input should be a fully formed question.\",\n", + " ),\n", + " Tool(\n", + " name=\"Ruff QA System\",\n", + " func=ruff.run,\n", + " description=\"useful for when you need to answer questions about ruff (a python linter). Input should be a fully formed question.\",\n", + " ),\n", + "]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fc47f230", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "# Construct the agent. We will use the default agent type here.\n", + "# See documentation for a full list of options.\n", + "agent = initialize_agent(\n", + " tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10ca2db8", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "agent.invoke(\n", + " \"What did biden say about ketanji brown jackson in the state of the union address?\"\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4e91b811", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "agent.invoke(\"Why use ruff over flake8?\")" + ] + }, + { + "cell_type": "markdown", + "id": "787a9b5e", + "metadata": {}, + "source": [ + "## Use the Agent solely as a router" + ] + }, + { + "cell_type": "markdown", + "id": "9161ba91", + "metadata": {}, + "source": [ + "
\n", + "\n", + "You can also set `return_direct=True` if you intend to use the agent as a router and just want to directly return the result of the RetrievalQAChain.\n", + "\n", + "Notice that in the above examples the agent did some extra work after querying the RetrievalQAChain. You can avoid that and just return the result directly." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f59b377e", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "tools = [\n", + " Tool(\n", + " name=\"State of Union QA System\",\n", + " func=state_of_union.run,\n", + " description=\"useful for when you need to answer questions about the most recent state of the union address. Input should be a fully formed question.\",\n", + " return_direct=True,\n", + " ),\n", + " Tool(\n", + " name=\"Ruff QA System\",\n", + " func=ruff.run,\n", + " description=\"useful for when you need to answer questions about ruff (a python linter). Input should be a fully formed question.\",\n", + " return_direct=True,\n", + " ),\n", + "]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8615707a", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "agent = initialize_agent(\n", + " tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "36e718a9", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "agent.invoke(\n", + " \"What did biden say about ketanji brown jackson in the state of the union address?\"\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "edfd0a1a", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "agent.invoke(\"Why use ruff over flake8?\")" + ] + }, + { + "cell_type": "markdown", + "id": "49a0cbbe", + "metadata": {}, + "source": [ + "
\n", + "\n", + "## Multi-Hop vector store reasoning\n", + "\n", + "Because vector stores are easily usable as tools in agents, it is easy to use answer multi-hop questions that depend on vector stores using the existing agent framework." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d397a233", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "tools = [\n", + " Tool(\n", + " name=\"State of Union QA System\",\n", + " func=state_of_union.run,\n", + " description=\"useful for when you need to answer questions about the most recent state of the union address. Input should be a fully formed question, not referencing any obscure pronouns from the conversation before.\",\n", + " ),\n", + " Tool(\n", + " name=\"Ruff QA System\",\n", + " func=ruff.run,\n", + " description=\"useful for when you need to answer questions about ruff (a python linter). Input should be a fully formed question, not referencing any obscure pronouns from the conversation before.\",\n", + " ),\n", + "]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "06157240", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "# Construct the agent. We will use the default agent type here.\n", + "# See documentation for a full list of options.\n", + "agent = initialize_agent(\n", + " tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b492b520", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "agent.invoke(\n", + " \"What tool does ruff use to run over Jupyter Notebooks? Did the president mention that tool in the state of the union?\"\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "b26b4963", + "metadata": {}, + "source": [ + "

\n", + "\n", + "---\n", + "\n", + "
\n", + "\n", + "## \ud83d\ude80 Your turn\n", + "\n", + "Time to make this lab your own! \n", + "\n", + "Replace the `state_of_the_union.txt` dataset with something you'd actually enjoy chatting with. A great place to start is the [sonnets.txt dataset](https://github.com/martin-gorner/tensorflow-rnn-shakespeare/blob/master/shakespeare/sonnets.txt) \u2014or any other .txt file from that repository. Of course, you're not limited to those options. Pick any text that interests you and see how your chatbot responds.\n", + "\n", + "Here's what to do:\n", + "\n", + "1. **Get your data** \u2014 Download your chosen `.txt` file into this project folder (or point `TextLoader` at it directly).\n", + "2. **Rebuild the vector store** \u2014 Load, split, and embed your new document, then create a fresh `RetrievalQA` chain for it (give it a descriptive `collection_name`!).\n", + "3. **Rewrite the tool description** \u2014 Update the `Tool`'s `name` and `description` so the agent knows *when* it should reach for this new tool instead of the Ruff or state-of-the-union ones.\n", + "4. **Rebuild the agent** \u2014 Combine your new tool with the existing Ruff tool (or drop it if you'd rather keep just your new dataset + one other source).\n", + "5. **Put it to the test** \u2014 Ask your agent:\n", + " - A direct question that only your new dataset can answer.\n", + " - A question that only the Ruff tool can answer.\n", + " - A multi-hop question that requires combining *both* tools' knowledge, like the Jupyter/Ruff example above.\n", + "6. **Reflect** \u2014 In a markdown cell, briefly note whether the agent picked the right tool(s) each time, and what happened when you set `return_direct=True` vs. not.\n", + "\n", + "\u2b50\ufe0f **Bonus points:**\n", + "- Instead of modifying this same file, create a new file `solution.ipynb` and replicate the process from scratch.\n", + "\n", + "\ud83d\udca1 **Tip:** \n", + "- Watch the `verbose=True` agent logs closely \u2014 they show you the agent's reasoning step by step, which is the best way to understand *why* it picked a particular tool.\n", + "\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Solution \u2014 New Dataset: AI Engineering Notes\n\nFor the exercise I use a small local dataset about AI engineering. Keeping it local makes the lab reproducible and lets the agent route between two clearly different sources:\n\n- **AI Engineering Notes** \u2014 RAG, embeddings, Jupyter, FastAPI, evaluation, and clean architecture.\n- **Ruff FAQ** \u2014 Ruff-specific Python linting knowledge.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "ai_notes_path = \"./datasets/ai_engineering_notes.txt\"\n\nai_loader = TextLoader(ai_notes_path)\nai_documents = ai_loader.load()\n\nai_texts = CharacterTextSplitter(\n chunk_size=500,\n chunk_overlap=50,\n).split_documents(ai_documents)\n\nai_db = Chroma.from_documents(\n ai_texts,\n embeddings,\n collection_name=\"ai-engineering-notes\",\n)\n\nai_notes_qa = RetrievalQA.from_chain_type(\n llm=llm,\n chain_type=\"stuff\",\n retriever=ai_db.as_retriever(search_kwargs={\"k\": 3}),\n)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "exercise_tools = [\n Tool(\n name=\"AI Engineering Notes QA\",\n func=ai_notes_qa.run,\n description=(\n \"Use this tool for questions about AI engineering concepts in the local notes, \"\n \"including RAG, embeddings, vector databases, Jupyter notebooks, FastAPI, \"\n \"evaluation, and clean architecture. Input must be a fully formed question.\"\n ),\n ),\n Tool(\n name=\"Ruff QA System\",\n func=ruff.run,\n description=(\n \"Use this tool for questions specifically about Ruff, the Python linter and formatter. \"\n \"Input must be a fully formed question.\"\n ),\n ),\n]\n\nexercise_agent = initialize_agent(\n exercise_tools,\n llm,\n agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,\n verbose=True,\n)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# 1) A question only the new AI Engineering dataset can answer.\nexercise_agent.invoke(\n \"According to the AI Engineering Notes, what are the main steps in a RAG pipeline?\"\n)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# 2) A question only the Ruff knowledge source can answer.\nexercise_agent.invoke(\n \"Why might a Python developer choose Ruff instead of Flake8?\"\n)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# 3) Multi-hop: requires both tools.\nexercise_agent.invoke(\n \"According to the AI Engineering Notes, what environment is commonly used for AI experiments, \"\n \"and does Ruff support linting code in that environment?\"\n)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Compare with router-only behaviour.\nrouter_tools = [\n Tool(\n name=\"AI Engineering Notes QA\",\n func=ai_notes_qa.run,\n description=\"Use for RAG, embeddings, Jupyter, FastAPI, evaluation, or clean architecture questions.\",\n return_direct=True,\n ),\n Tool(\n name=\"Ruff QA System\",\n func=ruff.run,\n description=\"Use for Ruff Python linter and formatter questions.\",\n return_direct=True,\n ),\n]\n\nrouter_agent = initialize_agent(\n router_tools,\n llm,\n agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,\n verbose=True,\n)\n\nrouter_agent.invoke(\n \"According to the AI Engineering Notes, what are the main steps in a RAG pipeline?\"\n)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Reflection\n\nFor direct questions, the tool descriptions give the agent a strong routing signal. A RAG or embeddings question should go to **AI Engineering Notes QA**, while a Ruff-specific question should go to **Ruff QA System**.\n\nThe multi-hop question is more interesting because it asks first which environment is common for AI experiments and then whether Ruff can lint code in that environment. The first part requires the AI Engineering Notes tool, which identifies **Jupyter notebooks**. The agent can then query the Ruff tool to determine Ruff's support for notebooks. With `return_direct=False`, the agent can call more than one tool and combine their outputs.\n\nWith `return_direct=True`, the selected RetrievalQA chain returns immediately. This is useful when the agent is acting only as a router because it avoids an extra LLM reasoning step. The downside is that it is not suitable for multi-hop questions requiring multiple sources, because the first direct tool result ends the run.\n\nThe main lesson is that tool **names and descriptions are part of the routing interface**. They should be specific, non-overlapping, and describe exactly when the agent should use each source. In production I would also use newer tool-calling agents, isolate vector-store adapters from application logic, and add routing/evaluation tests rather than relying only on manual verbose logs.\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python (LangChain 0.2.x)", + "language": "python", + "name": "langchain-v0.2.x" + }, + "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.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file