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
743 lines
22 KiB
Plaintext
743 lines
22 KiB
Plaintext
{
|
||
"cells": [
|
||
{
|
||
"attachments": {},
|
||
"cell_type": "markdown",
|
||
"id": "7e9367df",
|
||
"metadata": {},
|
||
"source": [
|
||
"<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/vector_stores/CassandraIndexDemo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "0b692c73",
|
||
"metadata": {},
|
||
"source": [
|
||
"# Cassandra Vector Store"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "1e7787c2",
|
||
"metadata": {},
|
||
"source": [
|
||
"[Apache Cassandra®](https://cassandra.apache.org) is a NoSQL, row-oriented, highly scalable and highly available database. Starting with version 5.0, the database ships with [vector search](https://cassandra.apache.org/doc/trunk/cassandra/vector-search/overview.html) capabilities.\n",
|
||
"\n",
|
||
"DataStax [Astra DB through CQL](https://docs.datastax.com/en/astra-serverless/docs/vector-search/quickstart.html) is a managed serverless database built on Cassandra, offering the same interface and strengths.\n",
|
||
"\n",
|
||
"**This notebook shows the basic usage of the Cassandra Vector Store in LlamaIndex.**\n",
|
||
"\n",
|
||
"To run the full code you need either a running Cassandra cluster equipped with Vector \n",
|
||
"Search capabilities or a DataStax Astra DB instance."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "daff81fe",
|
||
"metadata": {},
|
||
"source": [
|
||
"## Setup"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "47234905",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"%pip install llama-index-vector-stores-cassandra"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "e9f8dbcb",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"!pip install --quiet \"astrapy>=0.5.8\""
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "47264e32",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import os\n",
|
||
"from getpass import getpass\n",
|
||
"\n",
|
||
"from llama_index.core import (\n",
|
||
" VectorStoreIndex,\n",
|
||
" SimpleDirectoryReader,\n",
|
||
" Document,\n",
|
||
" StorageContext,\n",
|
||
")\n",
|
||
"from llama_index.vector_stores.cassandra import CassandraVectorStore"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "7eea233d-9993-4e2f-bbf4-fc335599ed4a",
|
||
"metadata": {},
|
||
"source": [
|
||
"The next step is to initialize CassIO with a global DB connection: this is the only step that is done slightly differently for a Cassandra cluster and Astra DB:"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "3c692310",
|
||
"metadata": {},
|
||
"source": [
|
||
"### Initialization (Cassandra cluster)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "2ea33832-06ad-46eb-a192-290f246ccdb1",
|
||
"metadata": {},
|
||
"source": [
|
||
"In this case, you first need to create a `cassandra.cluster.Session` object,\n",
|
||
"as described in the [Cassandra driver documentation](https://docs.datastax.com/en/developer/python-driver/latest/api/cassandra/cluster/#module-cassandra.cluster).\n",
|
||
"The details vary (e.g. with network settings and authentication), but this might be something like:"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "79475469-a765-4881-a97d-ac66c272c572",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from cassandra.cluster import Cluster\n",
|
||
"\n",
|
||
"cluster = Cluster([\"127.0.0.1\"])\n",
|
||
"session = cluster.connect()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "9bee3cd0-faa1-42c1-9826-1fdea46ea238",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import cassio\n",
|
||
"\n",
|
||
"CASSANDRA_KEYSPACE = input(\"CASSANDRA_KEYSPACE = \")\n",
|
||
"\n",
|
||
"cassio.init(session=session, keyspace=CASSANDRA_KEYSPACE)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "2e0900b8-e0d8-488d-bee5-98c9e76811b8",
|
||
"metadata": {},
|
||
"source": [
|
||
"### Initialization (Astra DB through CQL)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "4cc22cd0-3995-4029-ad9c-edd8d0c7e730",
|
||
"metadata": {},
|
||
"source": [
|
||
"In this case you initialize CassIO with the following connection parameters:\n",
|
||
"\n",
|
||
"- the Database ID, e.g. 01234567-89ab-cdef-0123-456789abcdef\n",
|
||
"- the Token, e.g. AstraCS:6gBhNmsk135.... (it must be a \"Database Administrator\" token)\n",
|
||
"- Optionally a Keyspace name (if omitted, the default one for the database will be used)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "ba118688",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdin",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"ASTRA_DB_ID = 01234567-89ab-cdef-0123-456789abcdef\n",
|
||
"ASTRA_DB_TOKEN = ········\n",
|
||
"ASTRA_DB_KEYSPACE (optional, can be left empty) = \n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"ASTRA_DB_ID = input(\"ASTRA_DB_ID = \")\n",
|
||
"ASTRA_DB_TOKEN = getpass(\"ASTRA_DB_TOKEN = \")\n",
|
||
"\n",
|
||
"desired_keyspace = input(\"ASTRA_DB_KEYSPACE (optional, can be left empty) = \")\n",
|
||
"if desired_keyspace:\n",
|
||
" ASTRA_DB_KEYSPACE = desired_keyspace\n",
|
||
"else:\n",
|
||
" ASTRA_DB_KEYSPACE = None"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "20933a07",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import cassio\n",
|
||
"\n",
|
||
"cassio.init(\n",
|
||
" database_id=ASTRA_DB_ID,\n",
|
||
" token=ASTRA_DB_TOKEN,\n",
|
||
" keyspace=ASTRA_DB_KEYSPACE,\n",
|
||
")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "f9b97a89",
|
||
"metadata": {},
|
||
"source": [
|
||
"### OpenAI key\n",
|
||
"\n",
|
||
"In order to use embeddings by OpenAI you need to supply an OpenAI API Key:"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "0c9f4d21-145a-401e-95ff-ccb259e8ef84",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdin",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"OpenAI API Key: ········\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"os.environ[\"OPENAI_API_KEY\"] = getpass(\"OpenAI API Key:\")"
|
||
]
|
||
},
|
||
{
|
||
"attachments": {},
|
||
"cell_type": "markdown",
|
||
"id": "174ce56b",
|
||
"metadata": {},
|
||
"source": [
|
||
"### Download data"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "37cb4f67",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"--2023-11-10 01:44:05-- https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt\n",
|
||
"Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.109.133, 185.199.110.133, 185.199.111.133, ...\n",
|
||
"Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.109.133|:443... connected.\n",
|
||
"HTTP request sent, awaiting response... 200 OK\n",
|
||
"Length: 75042 (73K) [text/plain]\n",
|
||
"Saving to: ‘data/paul_graham/paul_graham_essay.txt’\n",
|
||
"\n",
|
||
"data/paul_graham/pa 100%[===================>] 73.28K --.-KB/s in 0.01s \n",
|
||
"\n",
|
||
"2023-11-10 01:44:06 (4.80 MB/s) - ‘data/paul_graham/paul_graham_essay.txt’ saved [75042/75042]\n",
|
||
"\n"
|
||
]
|
||
}
|
||
],
|
||
"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": "markdown",
|
||
"id": "59ff935d",
|
||
"metadata": {},
|
||
"source": [
|
||
"## Creating and populating the Vector Store\n",
|
||
"\n",
|
||
"You will now load some essays by Paul Graham from a local file and store them into the Cassandra Vector Store."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "68cbd239-880e-41a3-98d8-dbb3fab55431",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"Total documents: 1\n",
|
||
"First document, id: 12bc6987-366a-49eb-8de0-7b52340e4958\n",
|
||
"First document, hash: abe31930a1775c78df5a5b1ece7108f78fedbf5fe4a9cf58d7a21808fccaef34\n",
|
||
"First document, text (75014 characters):\n",
|
||
"====================\n",
|
||
"\n",
|
||
"\n",
|
||
"What I Worked On\n",
|
||
"\n",
|
||
"February 2021\n",
|
||
"\n",
|
||
"Before college the two main things I worked on, outside of school, were writing and programming. I didn't write essays. I wrote what beginning writers were supposed to write then, and probably still are: short stories. My stories were awful. They had hardly any plot, just characters with strong feelings, which I imagined ma ...\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"# load documents\n",
|
||
"documents = SimpleDirectoryReader(\"./data/paul_graham/\").load_data()\n",
|
||
"print(f\"Total documents: {len(documents)}\")\n",
|
||
"print(f\"First document, id: {documents[0].doc_id}\")\n",
|
||
"print(f\"First document, hash: {documents[0].hash}\")\n",
|
||
"print(\n",
|
||
" \"First document, text\"\n",
|
||
" f\" ({len(documents[0].text)} characters):\\n{'='*20}\\n{documents[0].text[:360]} ...\"\n",
|
||
")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "dd270925",
|
||
"metadata": {},
|
||
"source": [
|
||
"### Initialize the Cassandra Vector Store\n",
|
||
"\n",
|
||
"Creation of the vector store entails creation of the underlying database table if it does not exist yet:"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "afc5c44f",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"cassandra_store = CassandraVectorStore(\n",
|
||
" table=\"cass_v_table\", embedding_dimension=1536\n",
|
||
")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "cbabd1a7",
|
||
"metadata": {},
|
||
"source": [
|
||
"Now wrap this store into an `index` LlamaIndex abstraction for later querying:"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "ca205b88",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"storage_context = StorageContext.from_defaults(vector_store=cassandra_store)\n",
|
||
"\n",
|
||
"index = VectorStoreIndex.from_documents(\n",
|
||
" documents, storage_context=storage_context\n",
|
||
")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "cb11e2e2",
|
||
"metadata": {},
|
||
"source": [
|
||
"Note that the above `from_documents` call does several things at once: it splits the input documents into chunks of manageable size (\"nodes\"), computes embedding vectors for each node, and stores them all in the Cassandra Vector Store."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "04304299-fc3e-40a0-8600-f50c3292767e",
|
||
"metadata": {},
|
||
"source": [
|
||
"## Querying the store"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "b241797e",
|
||
"metadata": {},
|
||
"source": [
|
||
"### Basic querying"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "35369eda",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"The author chose to work on AI because they were inspired by a novel called The Moon is a Harsh Mistress, which featured an intelligent computer, and a PBS documentary that showed Terry Winograd using SHRDLU. These experiences sparked the author's interest in AI and motivated them to pursue it as a field of study and work.\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"query_engine = index.as_query_engine()\n",
|
||
"response = query_engine.query(\"Why did the author choose to work on AI?\")\n",
|
||
"print(response.response)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "48761020",
|
||
"metadata": {},
|
||
"source": [
|
||
"### MMR-based queries\n",
|
||
"\n",
|
||
"The MMR (maximal marginal relevance) method is designed to fetch text chunks from the store that are at the same time relevant to the query but as different as possible from each other, with the goal of providing a broader context to the building of the final answer:"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "eb2054c0",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"The author chose to work on AI because they believed that teaching SHRDLU more words would eventually lead to the development of intelligent programs. They were fascinated by the potential of AI and saw it as an opportunity to expand their understanding of programming and push the limits of what could be achieved.\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"query_engine = index.as_query_engine(vector_store_query_mode=\"mmr\")\n",
|
||
"response = query_engine.query(\"Why did the author choose to work on AI?\")\n",
|
||
"print(response.response)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "4d7bc976",
|
||
"metadata": {},
|
||
"source": [
|
||
"## Connecting to an existing store\n",
|
||
"\n",
|
||
"Since this store is backed by Cassandra, it is persistent by definition. So, if you want to connect to a store that was created and populated previously, here is how:"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "f0aae26e",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"new_store_instance = CassandraVectorStore(\n",
|
||
" table=\"cass_v_table\", embedding_dimension=1536\n",
|
||
")\n",
|
||
"\n",
|
||
"# Create index (from preexisting stored vectors)\n",
|
||
"new_index_instance = VectorStoreIndex.from_vector_store(\n",
|
||
" vector_store=new_store_instance\n",
|
||
")\n",
|
||
"\n",
|
||
"# now you can do querying, etc:\n",
|
||
"query_engine = new_index_instance.as_query_engine(similarity_top_k=5)\n",
|
||
"response = query_engine.query(\n",
|
||
" \"What did the author study prior to working on AI?\"\n",
|
||
")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "1ceec3bf",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"The author studied philosophy prior to working on AI.\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"print(response.response)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "52b975a7",
|
||
"metadata": {},
|
||
"source": [
|
||
"## Removing documents from the index\n",
|
||
"\n",
|
||
"First get an explicit list of pieces of a document, or \"nodes\", from a `Retriever` spawned from the index:"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "75ed7807",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"retriever = new_index_instance.as_retriever(\n",
|
||
" vector_store_query_mode=\"mmr\",\n",
|
||
" similarity_top_k=3,\n",
|
||
" vector_store_kwargs={\"mmr_prefetch_factor\": 4},\n",
|
||
")\n",
|
||
"nodes_with_scores = retriever.retrieve(\n",
|
||
" \"What did the author study prior to working on AI?\"\n",
|
||
")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "6ae9c6b1",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"Found 3 nodes.\n",
|
||
" [0] score = 0.4251742327832831\n",
|
||
" id = 7e628668-58fa-4548-9c92-8c31d315dce0\n",
|
||
" text = What I Worked On\n",
|
||
"\n",
|
||
"February 2021\n",
|
||
"\n",
|
||
"Before college the two main things I worked on, outside o ...\n",
|
||
" [1] score = -0.020323897262800816\n",
|
||
" id = aa279d09-717f-4d68-9151-594c5bfef7ce\n",
|
||
" text = This was now only weeks away. My nice landlady let me leave my stuff in her attic. I had s ...\n",
|
||
" [2] score = 0.011198131320563909\n",
|
||
" id = 50b9170d-6618-4e8b-aaf8-36632e2801a6\n",
|
||
" text = It seemed only a matter of time before we'd have Mike, and when I saw Winograd using SHRDL ...\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"print(f\"Found {len(nodes_with_scores)} nodes.\")\n",
|
||
"for idx, node_with_score in enumerate(nodes_with_scores):\n",
|
||
" print(f\" [{idx}] score = {node_with_score.score}\")\n",
|
||
" print(f\" id = {node_with_score.node.node_id}\")\n",
|
||
" print(f\" text = {node_with_score.node.text[:90]} ...\")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "4bdc78f4-36ee-4358-9050-98f6e6652092",
|
||
"metadata": {},
|
||
"source": [
|
||
"But wait! When using the vector store, you should consider the **document** as the sensible unit to delete, and not any individual node belonging to it. Well, in this case, you just inserted a single text file, so all nodes will have the same `ref_doc_id`:"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "c52bb601-0838-46bb-8b9f-f2012a1c3f84",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"Nodes' ref_doc_id:\n",
|
||
"12bc6987-366a-49eb-8de0-7b52340e4958\n",
|
||
"12bc6987-366a-49eb-8de0-7b52340e4958\n",
|
||
"12bc6987-366a-49eb-8de0-7b52340e4958\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"print(\"Nodes' ref_doc_id:\")\n",
|
||
"print(\"\\n\".join([nws.node.ref_doc_id for nws in nodes_with_scores]))"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "7659d4c3",
|
||
"metadata": {},
|
||
"source": [
|
||
"Now let's say you need to remove the text file you uploaded:"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "2c7aafa0",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"new_store_instance.delete(nodes_with_scores[0].node.ref_doc_id)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "357a624b",
|
||
"metadata": {},
|
||
"source": [
|
||
"Repeat the very same query and check the results now. You should see _no results_ being found:"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "813276ff",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"Found 0 nodes.\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"nodes_with_scores = retriever.retrieve(\n",
|
||
" \"What did the author study prior to working on AI?\"\n",
|
||
")\n",
|
||
"\n",
|
||
"print(f\"Found {len(nodes_with_scores)} nodes.\")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "9fa59402",
|
||
"metadata": {},
|
||
"source": [
|
||
"## Metadata filtering\n",
|
||
"\n",
|
||
"The Cassandra vector store support metadata filtering in the form of exact-match `key=value` pairs at query time. The following cells, which work on a brand new Cassandra table, demonstrate this feature.\n",
|
||
"\n",
|
||
"In this demo, for the sake of brevity, a single source document is loaded (the `../data/paul_graham/paul_graham_essay.txt` text file). Nevertheless, you will attach some custom metadata to the document to illustrate how you can can restrict queries with conditions on the metadata attached to the documents."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "23c6ff1a",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"md_storage_context = StorageContext.from_defaults(\n",
|
||
" vector_store=CassandraVectorStore(\n",
|
||
" table=\"cass_v_table_md\", embedding_dimension=1536\n",
|
||
" )\n",
|
||
")\n",
|
||
"\n",
|
||
"\n",
|
||
"def my_file_metadata(file_name: str):\n",
|
||
" \"\"\"Depending on the input file name, associate a different metadata.\"\"\"\n",
|
||
" if \"essay\" in file_name:\n",
|
||
" source_type = \"essay\"\n",
|
||
" elif \"dinosaur\" in file_name:\n",
|
||
" # this (unfortunately) will not happen in this demo\n",
|
||
" source_type = \"dinos\"\n",
|
||
" else:\n",
|
||
" source_type = \"other\"\n",
|
||
" return {\"source_type\": source_type}\n",
|
||
"\n",
|
||
"\n",
|
||
"# Load documents and build index\n",
|
||
"md_documents = SimpleDirectoryReader(\n",
|
||
" \"./data/paul_graham\", file_metadata=my_file_metadata\n",
|
||
").load_data()\n",
|
||
"md_index = VectorStoreIndex.from_documents(\n",
|
||
" md_documents, storage_context=md_storage_context\n",
|
||
")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "a0a2a030-5621-4844-8971-2b4928ec2ec0",
|
||
"metadata": {},
|
||
"source": [
|
||
"\n",
|
||
"\n",
|
||
"That's it: you can now add filtering to your query engine:"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "d4bfd6f6",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from llama_index.core.vector_stores import ExactMatchFilter, MetadataFilters"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "733467f3",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"Yes, the author appreciated Lisp and painting. They mentioned spending a significant amount of time working on Lisp and even building a new dialect of Lisp called Arc. Additionally, the author mentioned spending most of 2014 painting and experimenting with different techniques.\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"md_query_engine = md_index.as_query_engine(\n",
|
||
" filters=MetadataFilters(\n",
|
||
" filters=[ExactMatchFilter(key=\"source_type\", value=\"essay\")]\n",
|
||
" )\n",
|
||
")\n",
|
||
"md_response = md_query_engine.query(\n",
|
||
" \"did the author appreciate Lisp and painting?\"\n",
|
||
")\n",
|
||
"print(md_response.response)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "4847dc97",
|
||
"metadata": {},
|
||
"source": [
|
||
"To test that the filtering is at play, try to change it to use only `\"dinos\"` documents... there will be no answer this time :)"
|
||
]
|
||
}
|
||
],
|
||
"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
|
||
}
|