Files
wehub-resource-sync a0c8464e58
Build Package / build (ubuntu-latest) (push) Failing after 1s
CodeQL / Analyze (python) (push) Failing after 1s
Core Typecheck / core-typecheck (push) Failing after 1s
Linting / lint (push) Failing after 1s
llama-dev tests / test-llama-dev (push) Failing after 1s
Publish Sub-Package to PyPI if Needed / publish_subpackage_if_needed (push) Has been skipped
Sync Docs to Developer Hub / sync-docs (push) Failing after 0s
Build Package / build (windows-latest) (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 12:26:52 +08:00

400 lines
9.6 KiB
Plaintext

{
"cells": [
{
"attachments": {},
"cell_type": "markdown",
"id": "d6509c3a",
"metadata": {},
"source": [
"<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/observability/WandbCallbackHandler.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
]
},
{
"cell_type": "markdown",
"id": "c0d8b66c",
"metadata": {},
"source": [
"# Wandb Callback Handler\n",
"\n",
"[Weights & Biases Prompts](https://docs.wandb.ai/guides/prompts) is a suite of LLMOps tools built for the development of LLM-powered applications.\n",
"\n",
"The `WandbCallbackHandler` is integrated with W&B Prompts to visualize and inspect the execution flow of your index construction, or querying over your index and more. You can use this handler to persist your created indices as W&B Artifacts allowing you to version control your indices.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "49c3527e",
"metadata": {},
"outputs": [],
"source": [
"%pip install llama-index-callbacks-wandb\n",
"%pip install llama-index-llms-openai"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "612f35ad",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"OpenAI API key configured\n"
]
}
],
"source": [
"import os\n",
"from getpass import getpass\n",
"\n",
"if os.getenv(\"OPENAI_API_KEY\") is None:\n",
" os.environ[\"OPENAI_API_KEY\"] = getpass(\n",
" \"Paste your OpenAI key from:\"\n",
" \" https://platform.openai.com/account/api-keys\\n\"\n",
" )\n",
"assert os.getenv(\"OPENAI_API_KEY\", \"\").startswith(\n",
" \"sk-\"\n",
"), \"This doesn't look like a valid OpenAI API key\"\n",
"print(\"OpenAI API key configured\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "78a29d9a",
"metadata": {},
"outputs": [],
"source": [
"from llama_index.core.callbacks import CallbackManager\n",
"from llama_index.core.callbacks import LlamaDebugHandler\n",
"from llama_index.callbacks.wandb import WandbCallbackHandler\n",
"from llama_index.core import (\n",
" VectorStoreIndex,\n",
" SimpleDirectoryReader,\n",
" SimpleKeywordTableIndex,\n",
" StorageContext,\n",
")\n",
"from llama_index.llms.openai import OpenAI"
]
},
{
"cell_type": "markdown",
"id": "e6feb252",
"metadata": {},
"source": [
"## Setup LLM"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d22fee33",
"metadata": {},
"outputs": [],
"source": [
"from llama_index.core import Settings\n",
"\n",
"Settings.llm = OpenAI(model=\"gpt-4\", temperature=0)"
]
},
{
"cell_type": "markdown",
"id": "8790f4c7",
"metadata": {},
"source": [
"## W&B Callback Manager Setup"
]
},
{
"cell_type": "markdown",
"id": "8a32b984-772e-4832-945e-cb6fc7be9e0b",
"metadata": {},
"source": [
"**Option 1**: Set Global Evaluation Handler"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2a3b9d22-cd67-4fb5-9785-254e58179a02",
"metadata": {},
"outputs": [],
"source": [
"import llama_index.core\n",
"from llama_index.core import set_global_handler\n",
"\n",
"set_global_handler(\"wandb\", run_args={\"project\": \"llamaindex\"})\n",
"wandb_callback = llama_index.core.global_handler"
]
},
{
"cell_type": "markdown",
"id": "d1755516-f8ad-458e-b52f-f7665c023e43",
"metadata": {},
"source": [
"**Option 2**: Manually Configure Callback Handler\n",
"\n",
"Also configure a debugger handler for extra notebook visibility."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "defa9155-daca-4a8f-8ca6-87d1ee98f084",
"metadata": {},
"outputs": [],
"source": [
"llama_debug = LlamaDebugHandler(print_trace_on_end=True)\n",
"\n",
"# wandb.init args\n",
"run_args = dict(\n",
" project=\"llamaindex\",\n",
")\n",
"\n",
"wandb_callback = WandbCallbackHandler(run_args=run_args)\n",
"\n",
"Settings.callback_manager = CallbackManager([llama_debug, wandb_callback])"
]
},
{
"cell_type": "markdown",
"id": "c4cf969a",
"metadata": {},
"source": [
"> After running the above cell, you will get the W&B run page URL. Here you will find a trace table with all the events tracked using [Weights and Biases' Prompts](https://docs.wandb.ai/guides/prompts) feature."
]
},
{
"cell_type": "markdown",
"id": "a4a7c101",
"metadata": {},
"source": [
"## 1. Indexing"
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "e5d31f80",
"metadata": {},
"source": [
"Download Data"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1e7ad71e",
"metadata": {},
"outputs": [],
"source": [
"!mkdir -p 'data/paul_graham/'\n",
"!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d1011596",
"metadata": {},
"outputs": [],
"source": [
"docs = SimpleDirectoryReader(\"./data/paul_graham/\").load_data()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d3d6975c",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"**********\n",
"Trace: index_construction\n",
" |_node_parsing -> 0.295179 seconds\n",
" |_chunking -> 0.293976 seconds\n",
" |_embedding -> 0.494492 seconds\n",
" |_embedding -> 0.346162 seconds\n",
"**********\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[34m\u001b[1mwandb\u001b[0m: Logged trace tree to W&B.\n"
]
}
],
"source": [
"index = VectorStoreIndex.from_documents(docs)"
]
},
{
"cell_type": "markdown",
"id": "0a948efc",
"metadata": {},
"source": [
"### 1.1 Persist Index as W&B Artifacts"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "8ad58e67",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[34m\u001b[1mwandb\u001b[0m: Adding directory to artifact (/Users/loganmarkewich/llama_index/docs/examples/callbacks/wandb/run-20230801_152955-ds93prxa/files/storage)... Done. 0.0s\n"
]
}
],
"source": [
"wandb_callback.persist_index(index, index_name=\"simple_vector_store\")"
]
},
{
"cell_type": "markdown",
"id": "7ed156a6",
"metadata": {},
"source": [
"### 1.2 Download Index from W&B Artifacts"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "dc35f448",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[34m\u001b[1mwandb\u001b[0m: 3 of 3 files downloaded. \n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"**********\n",
"Trace: index_construction\n",
"**********\n"
]
}
],
"source": [
"from llama_index.core import load_index_from_storage\n",
"\n",
"storage_context = wandb_callback.load_storage_context(\n",
" artifact_url=\"ayut/llamaindex/simple_vector_store:v0\"\n",
")\n",
"\n",
"# Load the index and initialize a query engine\n",
"index = load_index_from_storage(\n",
" storage_context,\n",
")"
]
},
{
"cell_type": "markdown",
"id": "ae4de4a9",
"metadata": {},
"source": [
"## 2. Query Over Index"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "42221465",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"**********\n",
"Trace: query\n",
" |_query -> 2.695958 seconds\n",
" |_retrieve -> 0.806379 seconds\n",
" |_embedding -> 0.802871 seconds\n",
" |_synthesize -> 1.8893 seconds\n",
" |_llm -> 1.842434 seconds\n",
"**********\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[34m\u001b[1mwandb\u001b[0m: Logged trace tree to W&B.\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"The text does not provide information on what the author did growing up.\n"
]
}
],
"source": [
"query_engine = index.as_query_engine()\n",
"response = query_engine.query(\"What did the author do growing up?\")\n",
"print(response, sep=\"\\n\")"
]
},
{
"cell_type": "markdown",
"id": "c49ff101",
"metadata": {},
"source": [
"## Close W&B Callback Handler\n",
"\n",
"When we are done tracking our events we can close the wandb run."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "28ef6a7b",
"metadata": {},
"outputs": [],
"source": [
"wandb_callback.finish()"
]
}
],
"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"
}
},
"nbformat": 4,
"nbformat_minor": 5
}