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
326 lines
10 KiB
Plaintext
326 lines
10 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "4e4f9a0f",
|
|
"metadata": {},
|
|
"source": [
|
|
"<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/vector_stores/BagelAutoRetriever.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "307804a3-c02b-4a57-ac0d-172c30ddc851",
|
|
"metadata": {},
|
|
"source": [
|
|
"# Lantern Vector Store (auto-retriever)\n",
|
|
"\n",
|
|
"This guide shows how to perform **auto-retrieval** in LlamaIndex. \n",
|
|
"\n",
|
|
"Many popular vector DBs support a set of metadata filters in addition to a query string for semantic search. Given a natural language query, we first use the LLM to infer a set of metadata filters as well as the right query string to pass to the vector DB (either can also be blank). This overall query bundle is then executed against the vector DB.\n",
|
|
"\n",
|
|
"This allows for more dynamic, expressive forms of retrieval beyond top-k semantic search. The relevant context for a given query may only require filtering on a metadata tag, or require a joint combination of filtering + semantic search within the filtered set, or just raw semantic search.\n",
|
|
"\n",
|
|
"We demonstrate an example with Lantern, but auto-retrieval is also implemented with many other vector DBs (e.g. Pinecone, Chroma, Weaviate, and more)."
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "15119a5b",
|
|
"metadata": {},
|
|
"source": [
|
|
"If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "841c6bb2",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"%pip install llama-index-vector-stores-lantern"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "feb74914",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"!pip install llama-index psycopg2-binary asyncpg"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "d48af8e1",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import logging\n",
|
|
"import sys\n",
|
|
"\n",
|
|
"logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
|
|
"logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "bf49ac18",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# set up OpenAI\n",
|
|
"import os\n",
|
|
"\n",
|
|
"os.environ[\"OPENAI_API_KEY\"] = \"<your-api-key>\"\n",
|
|
"\n",
|
|
"import openai\n",
|
|
"\n",
|
|
"openai.api_key = os.environ[\"OPENAI_API_KEY\"]"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "0ce3143d-198c-4dd2-8e5a-c5cdf94f017a",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import psycopg2\n",
|
|
"from sqlalchemy import make_url\n",
|
|
"\n",
|
|
"connection_string = \"postgresql://postgres:postgres@localhost:5432\"\n",
|
|
"\n",
|
|
"url = make_url(connection_string)\n",
|
|
"\n",
|
|
"db_name = \"postgres\"\n",
|
|
"conn = psycopg2.connect(connection_string)\n",
|
|
"conn.autocommit = True"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "0a2bcc07",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.core import VectorStoreIndex, StorageContext\n",
|
|
"from llama_index.vector_stores.lantern import LanternVectorStore"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "68cbd239-880e-41a3-98d8-dbb3fab55431",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.core.schema import TextNode\n",
|
|
"\n",
|
|
"nodes = [\n",
|
|
" TextNode(\n",
|
|
" text=(\n",
|
|
" \"Michael Jordan is a retired professional basketball player,\"\n",
|
|
" \" widely regarded as one of the greatest basketball players of all\"\n",
|
|
" \" time.\"\n",
|
|
" ),\n",
|
|
" metadata={\n",
|
|
" \"category\": \"Sports\",\n",
|
|
" \"country\": \"United States\",\n",
|
|
" },\n",
|
|
" ),\n",
|
|
" TextNode(\n",
|
|
" text=(\n",
|
|
" \"Angelina Jolie is an American actress, filmmaker, and\"\n",
|
|
" \" humanitarian. She has received numerous awards for her acting\"\n",
|
|
" \" and is known for her philanthropic work.\"\n",
|
|
" ),\n",
|
|
" metadata={\n",
|
|
" \"category\": \"Entertainment\",\n",
|
|
" \"country\": \"United States\",\n",
|
|
" },\n",
|
|
" ),\n",
|
|
" TextNode(\n",
|
|
" text=(\n",
|
|
" \"Elon Musk is a business magnate, industrial designer, and\"\n",
|
|
" \" engineer. He is the founder, CEO, and lead designer of SpaceX,\"\n",
|
|
" \" Tesla, Inc., Neuralink, and The Boring Company.\"\n",
|
|
" ),\n",
|
|
" metadata={\n",
|
|
" \"category\": \"Business\",\n",
|
|
" \"country\": \"United States\",\n",
|
|
" },\n",
|
|
" ),\n",
|
|
" TextNode(\n",
|
|
" text=(\n",
|
|
" \"Rihanna is a Barbadian singer, actress, and businesswoman. She\"\n",
|
|
" \" has achieved significant success in the music industry and is\"\n",
|
|
" \" known for her versatile musical style.\"\n",
|
|
" ),\n",
|
|
" metadata={\n",
|
|
" \"category\": \"Music\",\n",
|
|
" \"country\": \"Barbados\",\n",
|
|
" },\n",
|
|
" ),\n",
|
|
" TextNode(\n",
|
|
" text=(\n",
|
|
" \"Cristiano Ronaldo is a Portuguese professional footballer who is\"\n",
|
|
" \" considered one of the greatest football players of all time. He\"\n",
|
|
" \" has won numerous awards and set multiple records during his\"\n",
|
|
" \" career.\"\n",
|
|
" ),\n",
|
|
" metadata={\n",
|
|
" \"category\": \"Sports\",\n",
|
|
" \"country\": \"Portugal\",\n",
|
|
" },\n",
|
|
" ),\n",
|
|
"]"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "f6021786",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Build Vector Index with Lantern Vector Store\n",
|
|
"\n",
|
|
"Here we load the data into the vector store. As mentioned above, both the text and metadata for each node will get converted into corresponding representations in Lantern. We can now run semantic queries and also metadata filtering on this data from Lantern."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "ba1558b3",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"vector_store = LanternVectorStore.from_params(\n",
|
|
" database=db_name,\n",
|
|
" host=url.host,\n",
|
|
" password=url.password,\n",
|
|
" port=url.port,\n",
|
|
" user=url.username,\n",
|
|
" table_name=\"famous_people\",\n",
|
|
" embed_dim=1536, # openai embedding dimension\n",
|
|
" m=16, # HNSW M parameter\n",
|
|
" ef_construction=128, # HNSW ef construction parameter\n",
|
|
" ef=64, # HNSW ef search parameter\n",
|
|
")\n",
|
|
"\n",
|
|
"storage_context = StorageContext.from_defaults(vector_store=vector_store)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "35369eda",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"index = VectorStoreIndex(nodes, storage_context=storage_context)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "7eb20e49",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Define `VectorIndexAutoRetriever`\n",
|
|
"\n",
|
|
"We define our core `VectorIndexAutoRetriever` module. The module takes in `VectorStoreInfo`,\n",
|
|
"which contains a structured description of the vector store collection and the metadata filters it supports.\n",
|
|
"This information will then be used in the auto-retrieval prompt where the LLM infers metadata filters."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "bedbb693-725f-478f-be26-fa7180ea38b2",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.core.retrievers import VectorIndexAutoRetriever\n",
|
|
"from llama_index.core.vector_stores import MetadataInfo, VectorStoreInfo\n",
|
|
"\n",
|
|
"\n",
|
|
"vector_store_info = VectorStoreInfo(\n",
|
|
" content_info=\"brief biography of celebrities\",\n",
|
|
" metadata_info=[\n",
|
|
" MetadataInfo(\n",
|
|
" name=\"category\",\n",
|
|
" type=\"str\",\n",
|
|
" description=(\n",
|
|
" \"Category of the celebrity, one of [Sports, Entertainment,\"\n",
|
|
" \" Business, Music]\"\n",
|
|
" ),\n",
|
|
" ),\n",
|
|
" MetadataInfo(\n",
|
|
" name=\"country\",\n",
|
|
" type=\"str\",\n",
|
|
" description=(\n",
|
|
" \"Country of the celebrity, one of [United States, Barbados,\"\n",
|
|
" \" Portugal]\"\n",
|
|
" ),\n",
|
|
" ),\n",
|
|
" ],\n",
|
|
")\n",
|
|
"retriever = VectorIndexAutoRetriever(\n",
|
|
" index, vector_store_info=vector_store_info\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "31ac4a3b",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Running over some sample data\n",
|
|
"\n",
|
|
"We try running over some sample data. Note how metadata filters are inferred - this helps with more precise retrieval! "
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "eeb18e9c",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"retriever.retrieve(\"Tell me about two celebrities from United States\")"
|
|
]
|
|
}
|
|
],
|
|
"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"
|
|
},
|
|
"vscode": {
|
|
"interpreter": {
|
|
"hash": "0ac390d292208ca2380c85f5bce7ded36a7a25670a97c40b8009630eb36cb06e"
|
|
}
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 5
|
|
}
|