diff --git a/docs/docs/examples/vector_stores/VespaIndexDemo.ipynb b/docs/docs/examples/vector_stores/VespaIndexDemo.ipynb new file mode 100644 index 0000000000000..9d3c3e59fd194 --- /dev/null +++ b/docs/docs/examples/vector_stores/VespaIndexDemo.ipynb @@ -0,0 +1,537 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "d792fb5b", + "metadata": {}, + "source": [ + "\n", + " \n", + " \n", + " \"#Vespa\"\n", + "\n" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "23cf319b", + "metadata": {}, + "source": [ + "\"Open\n" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "307804a3-c02b-4a57-ac0d-172c30ddc851", + "metadata": {}, + "source": [ + "# Vespa Vector Store demo\n" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "5508d8ac", + "metadata": {}, + "source": [ + "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0beb6603", + "metadata": {}, + "outputs": [], + "source": [ + "%pip install llama-index-vector-stores-vespa llama-index pyvespa" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "f7010b1d-d1bb-4f08-9309-a328bb4ea396", + "metadata": {}, + "source": [ + "#### Setting up API key\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "08ad68ce", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "import openai\n", + "\n", + "os.environ[\"OPENAI_API_KEY\"] = \"sk-...\"\n", + "openai.api_key = os.environ[\"OPENAI_API_KEY\"]" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "8ee4473a-094f-4d0a-a825-e1213db07240", + "metadata": {}, + "source": [ + "#### Load documents, build the VectorStoreIndex\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0a2bcc07", + "metadata": {}, + "outputs": [], + "source": [ + "from llama_index.core import VectorStoreIndex\n", + "from llama_index.vector_stores.vespa import VespaVectorStore\n", + "from IPython.display import Markdown, display" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "3a41a70d", + "metadata": {}, + "source": [ + "## Defining some sample data\n", + "\n", + "Let's insert some documents.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "df6b6d46", + "metadata": {}, + "outputs": [], + "source": [ + "from llama_index.core.schema import TextNode\n", + "\n", + "nodes = [\n", + " TextNode(\n", + " text=\"The Shawshank Redemption\",\n", + " metadata={\n", + " \"author\": \"Stephen King\",\n", + " \"theme\": \"Friendship\",\n", + " \"year\": 1994,\n", + " },\n", + " ),\n", + " TextNode(\n", + " text=\"The Godfather\",\n", + " metadata={\n", + " \"director\": \"Francis Ford Coppola\",\n", + " \"theme\": \"Mafia\",\n", + " \"year\": 1972,\n", + " },\n", + " ),\n", + " TextNode(\n", + " text=\"Inception\",\n", + " metadata={\n", + " \"director\": \"Christopher Nolan\",\n", + " \"theme\": \"Fiction\",\n", + " \"year\": 2010,\n", + " },\n", + " ),\n", + " TextNode(\n", + " text=\"To Kill a Mockingbird\",\n", + " metadata={\n", + " \"author\": \"Harper Lee\",\n", + " \"theme\": \"Mafia\",\n", + " \"year\": 1960,\n", + " },\n", + " ),\n", + " TextNode(\n", + " text=\"1984\",\n", + " metadata={\n", + " \"author\": \"George Orwell\",\n", + " \"theme\": \"Totalitarianism\",\n", + " \"year\": 1949,\n", + " },\n", + " ),\n", + " TextNode(\n", + " text=\"The Great Gatsby\",\n", + " metadata={\n", + " \"author\": \"F. Scott Fitzgerald\",\n", + " \"theme\": \"The American Dream\",\n", + " \"year\": 1925,\n", + " },\n", + " ),\n", + " TextNode(\n", + " text=\"Harry Potter and the Sorcerer's Stone\",\n", + " metadata={\n", + " \"author\": \"J.K. Rowling\",\n", + " \"theme\": \"Fiction\",\n", + " \"year\": 1997,\n", + " },\n", + " ),\n", + "]" + ] + }, + { + "cell_type": "markdown", + "id": "31fe4378", + "metadata": {}, + "source": [ + "### Initilizing the VespaVectorStore\n" + ] + }, + { + "cell_type": "markdown", + "id": "a0be7d09", + "metadata": {}, + "source": [ + "To make it really simple to get started, we provide a template Vespa application that will be deployed upon initializing the vector store.\n", + "\n", + "This is a huge abstraction and there are endless opportunities to tailor and customize the Vespa application to your needs. But for now, let's keep it simple and initialize with the default template.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "30b0b2e3", + "metadata": {}, + "outputs": [], + "source": [ + "from llama_index.core import StorageContext\n", + "\n", + "vector_store = VespaVectorStore()\n", + "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n", + "index = VectorStoreIndex(nodes, storage_context=storage_context)" + ] + }, + { + "cell_type": "markdown", + "id": "71a4a3ec", + "metadata": {}, + "source": [ + "### Deleting documents\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f4637a79", + "metadata": {}, + "outputs": [], + "source": [ + "node_to_delete = nodes[0].node_id\n", + "node_to_delete" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "84a97903", + "metadata": {}, + "outputs": [], + "source": [ + "vector_store.delete(ref_doc_id=node_to_delete)" + ] + }, + { + "cell_type": "markdown", + "id": "03315550", + "metadata": {}, + "source": [ + "## Querying\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "74cabf95", + "metadata": {}, + "outputs": [], + "source": [ + "from llama_index.core.vector_stores.types import (\n", + " VectorStoreQuery,\n", + " VectorStoreQueryMode,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d6401e25", + "metadata": {}, + "outputs": [], + "source": [ + "query = VectorStoreQuery(\n", + " query_str=\"Great Gatsby\",\n", + " mode=VectorStoreQueryMode.TEXT_SEARCH,\n", + " similarity_top_k=1,\n", + ")\n", + "result = vector_store.query(query)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "09f1bf81", + "metadata": {}, + "outputs": [], + "source": [ + "result" + ] + }, + { + "cell_type": "markdown", + "id": "77d2528e", + "metadata": {}, + "source": [ + "## As retriever\n" + ] + }, + { + "cell_type": "markdown", + "id": "a8d7aca1", + "metadata": {}, + "source": [ + "### Default query mode (text search)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5a71818e", + "metadata": {}, + "outputs": [], + "source": [ + "retriever = index.as_retriever(vector_store_query_mode=\"default\")\n", + "results = retriever.retrieve(\"Who directed inception?\")\n", + "display(Markdown(f\"**Retrieved nodes:**\\n {results}\"))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bfe83ebf", + "metadata": {}, + "outputs": [], + "source": [ + "retriever = index.as_retriever(vector_store_query_mode=\"semantic_hybrid\")\n", + "results = retriever.retrieve(\"Who wrote Harry Potter?\")\n", + "display(Markdown(f\"**Retrieved nodes:**\\n {results}\"))" + ] + }, + { + "cell_type": "markdown", + "id": "c8aa36e8", + "metadata": {}, + "source": [ + "### As query engine\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c1bd18f8", + "metadata": {}, + "outputs": [], + "source": [ + "query_engine = index.as_query_engine()\n", + "response = query_engine.query(\"Who directed inception?\")\n", + "display(Markdown(f\"**Response:** {response}\"))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "aede9cf6", + "metadata": {}, + "outputs": [], + "source": [ + "query_engine = index.as_query_engine(\n", + " vector_store_query_mode=\"semantic_hybrid\", verbose=True\n", + ")\n", + "response = query_engine.query(\n", + " \"When was the book about the wizard boy published and what was it called?\"\n", + ")\n", + "display(Markdown(f\"**Response:** {response}\"))\n", + "display(Markdown(f\"**Sources:** {response.source_nodes}\"))" + ] + }, + { + "cell_type": "markdown", + "id": "90081efd", + "metadata": {}, + "source": [ + "## Using metadata filters\n", + "\n", + "**NOTE**: This metadata filtering is done by llama-index, outside of vespa. For native and much more performant filtering, you should use Vespa's own filtering capabilities.\n", + "\n", + "See [Vespa's documentation](https://docs.vespa.ai/en/reference/query-language-reference.html) for more information.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0663ab38", + "metadata": {}, + "outputs": [], + "source": [ + "from llama_index.core.vector_stores import (\n", + " FilterOperator,\n", + " FilterCondition,\n", + " MetadataFilter,\n", + " MetadataFilters,\n", + ")\n", + "\n", + "# Let's define a filter that will only allow nodes that has the theme \"Fiction\" OR is published after 1997\n", + "\n", + "filters = MetadataFilters(\n", + " filters=[\n", + " MetadataFilter(key=\"theme\", value=\"Fiction\"),\n", + " MetadataFilter(key=\"year\", value=1997, operator=FilterOperator.GT),\n", + " ],\n", + " condition=FilterCondition.OR,\n", + ")\n", + "\n", + "retriever = index.as_retriever(filters=filters)\n", + "result = retriever.retrieve(\"Harry Potter\")\n", + "display(Markdown(f\"**Result:** {result}\"))" + ] + }, + { + "cell_type": "markdown", + "id": "414e6d78", + "metadata": {}, + "source": [ + "## Abstraction level of this integration\n", + "\n", + "To make it really simple to get started, we provide a template Vespa application that will be deployed upon initializing the vector store. This removes some of the complexity of setting up Vespa for the first time, but for serious use cases, we strongly recommend that you read the [Vespa documentation](docs.vespa.ai) and tailor the application to your needs.\n", + "\n", + "### The template\n", + "\n", + "The provided template Vespa application can be seen below:\n", + "\n", + "```python\n", + "from vespa.package import (\n", + " ApplicationPackage,\n", + " Field,\n", + " Schema,\n", + " Document,\n", + " HNSW,\n", + " RankProfile,\n", + " Component,\n", + " Parameter,\n", + " FieldSet,\n", + " GlobalPhaseRanking,\n", + " Function,\n", + ")\n", + "\n", + "hybrid_template = ApplicationPackage(\n", + " name=\"hybridsearch\",\n", + " schema=[\n", + " Schema(\n", + " name=\"doc\",\n", + " document=Document(\n", + " fields=[\n", + " Field(name=\"id\", type=\"string\", indexing=[\"summary\"]),\n", + " Field(name=\"metadata\", type=\"string\", indexing=[\"summary\"]),\n", + " Field(\n", + " name=\"text\",\n", + " type=\"string\",\n", + " indexing=[\"index\", \"summary\"],\n", + " index=\"enable-bm25\",\n", + " bolding=True,\n", + " ),\n", + " Field(\n", + " name=\"embedding\",\n", + " type=\"tensor(x[384])\",\n", + " indexing=[\n", + " \"input text\",\n", + " \"embed\",\n", + " \"index\",\n", + " \"attribute\",\n", + " ],\n", + " ann=HNSW(distance_metric=\"angular\"),\n", + " is_document_field=False,\n", + " ),\n", + " ]\n", + " ),\n", + " fieldsets=[FieldSet(name=\"default\", fields=[\"text\", \"metadata\"])],\n", + " rank_profiles=[\n", + " RankProfile(\n", + " name=\"bm25\",\n", + " inputs=[(\"query(q)\", \"tensor(x[384])\")],\n", + " functions=[Function(name=\"bm25sum\", expression=\"bm25(text)\")],\n", + " first_phase=\"bm25sum\",\n", + " ),\n", + " RankProfile(\n", + " name=\"semantic\",\n", + " inputs=[(\"query(q)\", \"tensor(x[384])\")],\n", + " first_phase=\"closeness(field, embedding)\",\n", + " ),\n", + " RankProfile(\n", + " name=\"fusion\",\n", + " inherits=\"bm25\",\n", + " inputs=[(\"query(q)\", \"tensor(x[384])\")],\n", + " first_phase=\"closeness(field, embedding)\",\n", + " global_phase=GlobalPhaseRanking(\n", + " expression=\"reciprocal_rank_fusion(bm25sum, closeness(field, embedding))\",\n", + " rerank_count=1000,\n", + " ),\n", + " ),\n", + " ],\n", + " )\n", + " ],\n", + " components=[\n", + " Component(\n", + " id=\"e5\",\n", + " type=\"hugging-face-embedder\",\n", + " parameters=[\n", + " Parameter(\n", + " \"transformer-model\",\n", + " {\n", + " \"url\": \"https://github.com/vespa-engine/sample-apps/raw/master/simple-semantic-search/model/e5-small-v2-int8.onnx\"\n", + " },\n", + " ),\n", + " Parameter(\n", + " \"tokenizer-model\",\n", + " {\n", + " \"url\": \"https://raw.githubusercontent.com/vespa-engine/sample-apps/master/simple-semantic-search/model/tokenizer.json\"\n", + " },\n", + " ),\n", + " ],\n", + " )\n", + " ],\n", + ")\n", + "```\n", + "\n", + "Note that the fields `id`, `metadata`, `text`, and `embedding` are required for the integration to work.\n", + "The schema name must also be `doc`, and the rank profiles must be named `bm25`, `semantic`, and `fusion`.\n", + "\n", + "Other than that you are free to modify as you see fit by switching out embedding models, adding more fields, or changing the ranking expressions.\n", + "\n", + "For more details, check out this Pyvespa example notebook on [hybrid search](https://pyvespa.readthedocs.io/en/latest/getting-started-pyvespa.html).\n" + ] + } + ], + "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 +} diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/.gitignore b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/.gitignore new file mode 100644 index 0000000000000..990c18de22908 --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/.gitignore @@ -0,0 +1,153 @@ +llama_index/_static +.DS_Store +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +bin/ +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +etc/ +include/ +lib/ +lib64/ +parts/ +sdist/ +share/ +var/ +wheels/ +pip-wheel-metadata/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +.ruff_cache + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints +notebooks/ + +# IPython +profile_default/ +ipython_config.py + +# pyenv +.python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ +pyvenv.cfg + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# Jetbrains +.idea +modules/ +*.swp + +# VsCode +.vscode + +# pipenv +Pipfile +Pipfile.lock + +# pyright +pyrightconfig.json diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/BUILD b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/BUILD new file mode 100644 index 0000000000000..72b97f383ad9b --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/BUILD @@ -0,0 +1,4 @@ +poetry_requirements( + name="poetry", + module_mapping={"pyvespa": ["vespa"]}, +) diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/CHANGELOG.md b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/CHANGELOG.md new file mode 100644 index 0000000000000..633b40f215c09 --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/CHANGELOG.md @@ -0,0 +1,9 @@ +# CHANGELOG — llama-index-vector-stores-opensearch + +## [0.1.2] + +- Adds OpensearchVectorClient as top-level import + +## [0.1.1] + +- Fixes strict equality in dependency of llama-index-core diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/Makefile b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/Makefile new file mode 100644 index 0000000000000..b9eab05aa3706 --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/Makefile @@ -0,0 +1,17 @@ +GIT_ROOT ?= $(shell git rev-parse --show-toplevel) + +help: ## Show all Makefile targets. + @grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[33m%-30s\033[0m %s\n", $$1, $$2}' + +format: ## Run code autoformatters (black). + pre-commit install + git ls-files | xargs pre-commit run black --files + +lint: ## Run linters: pre-commit (black, ruff, codespell) and mypy + pre-commit install && git ls-files | xargs pre-commit run --show-diff-on-failure --files + +test: ## Run tests via pytest. + pytest tests + +watch-docs: ## Build and watch documentation. + sphinx-autobuild docs/ docs/_build/html --open-browser --watch $(GIT_ROOT)/llama_index/ diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/README.md b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/README.md new file mode 100644 index 0000000000000..0956439e74517 --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/README.md @@ -0,0 +1,143 @@ + + + + #Vespa + + +# LlamaIndex Vector_Stores Integration: Vespa + +[Vespa.ai](https://vespa.ai/) is an open-source big data serving engine. It is designed for low-latency and high-throughput serving of data and models. Vespa.ai is used by many companies to serve search results, recommendations, and rankings for billions of documents and users, expecting response times in the milliseconds. + +This integration allows you to use Vespa.ai as a vector store for LlamaIndex. Vespa has integrated support for [embedding inference](https://docs.vespa.ai/en/embedding.html), so you don't need to run a separate service for these tasks. + +Huggingface 🤗 embedders are supported, as well as SPLADE and ColBERT. + +## Abstraction level of this integration + +To make it really simple to get started, we provide a template Vespa application that will be deployed upon initializing the vector store. This removes some of the complexity of setting up Vespa for the first time, but for serious use cases, we strongly recommend that you read the [Vespa documentation](docs.vespa.ai) and tailor the application to your needs. + +## The template + +The provided template Vespa application can be seen below: + +```python +from vespa.package import ( + ApplicationPackage, + Field, + Schema, + Document, + HNSW, + RankProfile, + Component, + Parameter, + FieldSet, + GlobalPhaseRanking, + Function, +) + +hybrid_template = ApplicationPackage( + name="hybridsearch", + schema=[ + Schema( + name="doc", + document=Document( + fields=[ + Field(name="id", type="string", indexing=["summary"]), + Field( + name="metadata", type="string", indexing=["summary"] + ), + Field( + name="text", + type="string", + indexing=["index", "summary"], + index="enable-bm25", + bolding=True, + ), + Field( + name="embedding", + type="tensor(x[384])", + indexing=[ + "input text", + "embed", + "index", + "attribute", + ], + ann=HNSW(distance_metric="angular"), + is_document_field=False, + ), + ] + ), + fieldsets=[FieldSet(name="default", fields=["text", "metadata"])], + rank_profiles=[ + RankProfile( + name="bm25", + inputs=[("query(q)", "tensor(x[384])")], + functions=[ + Function(name="bm25sum", expression="bm25(text)") + ], + first_phase="bm25sum", + ), + RankProfile( + name="semantic", + inputs=[("query(q)", "tensor(x[384])")], + first_phase="closeness(field, embedding)", + ), + RankProfile( + name="fusion", + inherits="bm25", + inputs=[("query(q)", "tensor(x[384])")], + first_phase="closeness(field, embedding)", + global_phase=GlobalPhaseRanking( + expression="reciprocal_rank_fusion(bm25sum, closeness(field, embedding))", + rerank_count=1000, + ), + ), + ], + ) + ], + components=[ + Component( + id="e5", + type="hugging-face-embedder", + parameters=[ + Parameter( + "transformer-model", + { + "url": "https://github.com/vespa-engine/sample-apps/raw/master/simple-semantic-search/model/e5-small-v2-int8.onnx" + }, + ), + Parameter( + "tokenizer-model", + { + "url": "https://raw.githubusercontent.com/vespa-engine/sample-apps/master/simple-semantic-search/model/tokenizer.json" + }, + ), + ], + ) + ], +) +``` + +Note that the fields `id`, `metadata`, `text`, and `embedding` are required for the integration to work. +The schema name must also be `doc`, and the rank profiles must be named `bm25`, `semantic`, and `fusion`. + +Other than that you are free to modify as you see fit by switching out embedding models, adding more fields, or changing the ranking expressions. + +For more details, check out this Pyvespa example notebook on [hybrid search](https://pyvespa.readthedocs.io/en/latest/getting-started-pyvespa.html). + +## Going to production + +If you are ready to graduate to a production setup, we highly recommend to check out the [Vespa Cloud](https://cloud.vespa.ai/) service, where we manage all infrastructure and operations for you. Free trials are available. + +## Next steps + +There are many awesome features in Vespa, that are not exposed directly in this integration, check out [Pyvespa examples](https://pyvespa.readthedocs.io/en/latest/examples/pyvespa-examples.html) for some inspiration on what you can do with Vespa. + +Teasers: + +- Binary + Matryoshka embeddings. +- ColBERT. +- ONNX models. +- XGBoost and lightGBM models for ranking. +- Multivector indexing. +- and much more. diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/llama_index/vector_stores/vespa/BUILD b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/llama_index/vector_stores/vespa/BUILD new file mode 100644 index 0000000000000..db46e8d6c978c --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/llama_index/vector_stores/vespa/BUILD @@ -0,0 +1 @@ +python_sources() diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/llama_index/vector_stores/vespa/__init__.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/llama_index/vector_stores/vespa/__init__.py new file mode 100644 index 0000000000000..7472eed62961b --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/llama_index/vector_stores/vespa/__init__.py @@ -0,0 +1,6 @@ +from llama_index.vector_stores.vespa.base import ( + VespaVectorStore, +) +from llama_index.vector_stores.vespa.templates import hybrid_template + +__all__ = ["VespaVectorStore", "hybrid_template"] diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/llama_index/vector_stores/vespa/base.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/llama_index/vector_stores/vespa/base.py new file mode 100644 index 0000000000000..74ec8d657dc06 --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/llama_index/vector_stores/vespa/base.py @@ -0,0 +1,505 @@ +"""Vespa vector store.""" + +from typing import Any, List, Optional, Callable + + +from llama_index.core.schema import BaseNode, MetadataMode +from llama_index.core.vector_stores.types import ( + VectorStore, + VectorStoreQuery, + VectorStoreQueryMode, + VectorStoreQueryResult, +) +from llama_index.core.vector_stores.utils import ( + node_to_metadata_dict, + metadata_dict_to_node, +) + +from llama_index.vector_stores.vespa.templates import hybrid_template + +import asyncio +import logging +import json +import sys + +try: + from vespa.application import Vespa + from vespa.package import ApplicationPackage + from vespa.io import VespaResponse + from vespa.deployment import VespaCloud, VespaDocker +except ImportError: + raise ModuleNotFoundError( + "pyvespa not installed. Please install it via `pip install pyvespa`" + ) + +logger = logging.getLogger(__name__) +handler = logging.StreamHandler(stream=sys.stdout) +logger.addHandler(handler) +logger.setLevel(logging.INFO) + + +def callback(response: VespaResponse, id: str): + if not response.is_successful(): + logger.debug( + f"Failed to feed document {id} with status code {response.status_code}: Reason {response.get_json()}" + ) + + +class VespaVectorStore(VectorStore): + """ + Vespa vector store. + + Can be initialized in several ways: + 1. (Default) Initialize Vespa vector store with default hybrid template and local (docker) deployment. + 2. Initialize by providing an application package created in pyvespa (can be deployed locally or to Vespa cloud). + 3. Initialize from previously deployed Vespa application by providing URL. (Local or cloud deployment). + + The application must be set up with the following fields: + - id: Document id + - text: Text field + - embedding: Field to store embedding vectors. + - metadata: Metadata field (all metadata will be stored here) + + The application must be set up with the following rank profiles: + - bm25: For text search + - semantic: For semantic search + - fusion: For semantic hybrid search + + When creating a VectorStoreIndex from VespaVectorStore, the index will add documents to the Vespa application. + Be ware that the Vespa container will be reused if not deleted between deployments, to avoid data duplication. + During query time, the index queries the Vespa application to get the top k most relevant hits. + + Args: + application_package (ApplicationPackage): Application package + deployment_target (str): Deployment target, either `local` or `cloud` + port (int): Port that Vespa application will run on. Only applicable if deployment_target is `local` + default_schema_name (str): Schema name in Vespa application + namespace (str): Namespace in Vespa application. See https://docs.vespa.ai/en/documents.html#namespace. Defaults to `default`. + embeddings_outside_vespa (bool): Whether embeddings are created outside Vespa, or not. + url (Optional[str]): URL of deployed Vespa application. + groupname (Optional[str]): Group name in Vespa application, only applicable in `streaming` mode, see https://pyvespa.readthedocs.io/en/latest/examples/scaling-personal-ai-assistants-with-streaming-mode-cloud.html#A-summary-of-Vespa-streaming-mode + tenant (Optional[str]): Tenant for Vespa application. Applicable only if deployment_target is `cloud` + key_location (Optional[str]): Location of the control plane key used for signing HTTP requests to the Vespa Cloud. + key_content (Optional[str]): Content of the control plane key used for signing HTTP requests to the Vespa Cloud. Use only when key file is not available. + auth_client_token_id (Optional[str]): Use token based data plane authentication. This is the token name configured in the Vespa Cloud Console. This is used to configure Vespa services.xml. The token is given read and write permissions. + kwargs (Any): Additional kwargs for Vespa application + + Examples: + `pip install llama-index-vector-stores-vespa` + + ```python + from llama_index.core import VectorStoreIndex + from llama_index.vector_stores.vespa import VespaVectorStore + + vector_store = VespaVectorStore() + storage_context = StorageContext.from_defaults(vector_store=vector_store) + index = VectorStoreIndex(nodes, storage_context=storage_context) + retriever = index.as_retriever() + retriever.retrieve("Who directed inception?") + + ``` + """ + + stores_text: bool = True + is_embedding_query: bool = False + flat_metadata: bool = True + + def __init__( + self, + application_package: ApplicationPackage = hybrid_template, + namespace: str = "default", + default_schema_name: str = "doc", + deployment_target: str = "local", # "local" or "cloud" + port: int = 8080, + embeddings_outside_vespa: bool = False, + url: Optional[str] = None, + groupname: Optional[str] = None, + tenant: Optional[str] = None, + application: Optional[str] = "hybridsearch", + key_location: Optional[str] = None, + key_content: Optional[str] = None, + auth_client_token_id: Optional[str] = None, + **kwargs: Any, + ) -> None: + # Verify that application_package is an instance of ApplicationPackage + if not isinstance(application_package, ApplicationPackage): + raise ValueError( + "application_package must be an instance of vespa.package.ApplicationPackage" + ) + if application_package == hybrid_template: + logger.info( + "Using default hybrid template. Please make sure that the Vespa application is set up with the correct schema and rank profile." + ) + # Initialize all parameters + self.application_package = application_package + self.deployment_target = deployment_target + self.default_schema_name = default_schema_name + self.namespace = namespace + self.embeddings_outside_vespa = embeddings_outside_vespa + self.port = port + self.url = url + self.groupname = groupname + self.tenant = tenant + self.application = application + self.key_location = key_location + self.key_content = key_content + self.auth_client_token_id = auth_client_token_id + self.kwargs = kwargs + if self.url is None: + self.app = self._deploy() + else: + self.app = self._try_get_running_app() + + @property + def client(self) -> Vespa: + """Get client.""" + return self.app + + def _try_get_running_app(self) -> Vespa: + app = Vespa(url=f"{self.url}:{self.port}") + status = app.get_application_status() + if status.status_code == 200: + return app + else: + raise ConnectionError( + f"Vespa application not running on url {self.url} and port {self.port}. Please start Vespa application first." + ) + + def _deploy(self) -> Vespa: + if self.deployment_target == "cloud": + app = self._deploy_app_cloud() + elif self.deployment_target == "local": + app = self._deploy_app_local() + else: + raise ValueError( + f"Deployment target {self.deployment_target} not supported. Please choose either `local` or `cloud`." + ) + return app + + def _deploy_app_local(self) -> Vespa: + logger.info(f"Deploying Vespa application {self.application} to Vespa Docker.") + return VespaDocker(port=8080).deploy(self.application_package) + + def _deploy_app_cloud(self) -> Vespa: + logger.info(f"Deploying Vespa application {self.application} to Vespa Cloud.") + return VespaCloud( + tenant=self.tenant, + application=self.application, + application_package=self.application_package, + key_location=self.key_location, + key_content=self.key_content, + auth_client_token_id=self.auth_client_token_id, + **self.kwargs, + ).deploy() + + def add( + self, + nodes: List[BaseNode], + schema: Optional[str] = None, + callback: Optional[Callable[[VespaResponse, str], None]] = callback, + ) -> List[str]: + """ + Add nodes to vector store. + + Args: + nodes (List[BaseNode]): List of nodes to add + schema (Optional[str]): Schema name in Vespa application to add nodes to. Defaults to `default_schema_name`. + """ + # Create vespa iterable from nodes + ids = [] + data_to_insert = [] + for node in nodes: + metadata = node_to_metadata_dict( + node, remove_text=False, flat_metadata=self.flat_metadata + ) + logger.debug(f"Metadata: {metadata}") + entry = { + "id": node.node_id, + "fields": { + "id": node.node_id, + "text": node.get_content(metadata_mode=MetadataMode.NONE) or "", + "metadata": json.dumps(metadata), + }, + } + if self.embeddings_outside_vespa: + entry["fields"]["embedding"] = node.get_embedding() + data_to_insert.append(entry) + ids.append(node.node_id) + + self.app.feed_iterable( + data_to_insert, + schema=schema or self.default_schema_name, + namespace=self.namespace, + operation_type="feed", + callback=callback, + ) + return ids + + async def async_add( + self, + nodes: List[BaseNode], + schema: Optional[str] = None, + callback: Optional[Callable[[VespaResponse, str], None]] = callback, + max_connections: int = 10, + num_concurrent_requests: int = 1000, + total_timeout: int = 60, + **kwargs: Any, + ) -> List[str]: + """ + Add nodes to vector store asynchronously. + + Args: + nodes (List[BaseNode]): List of nodes to add + schema (Optional[str]): Schema name in Vespa application to add nodes to. Defaults to `default_schema_name`. + max_connections (int): Maximum number of connections to Vespa application + num_concurrent_requests (int): Maximum number of concurrent requests + total_timeout (int): Total timeout for all requests + kwargs (Any): Additional kwargs for Vespa application + """ + semaphore = asyncio.Semaphore(max_concurrent_requests) + ids = [] + data_to_insert = [] + for node in nodes: + metadata = node_to_metadata_dict( + node, remove_text=False, flat_metadata=self.flat_metadata + ) + logger.debug(f"Metadata: {metadata}") + entry = { + "id": node.node_id, + "fields": { + "id": node.node_id, + "text": node.get_content(metadata_mode=MetadataMode.NONE) or "", + "metadata": json.dumps(metadata), + }, + } + if self.embeddings_outside_vespa: + entry["fields"]["embedding"] = node.get_embedding() + data_to_insert.append(entry) + ids.append(node.node_id) + + async with self.app.asyncio( + connections=max_connections, total_timeout=total_timeout + ) as async_app: + for doc in data_to_insert: + async with semaphore: + task = asyncio.create_task( + async_app.feed_data_point( + data_id=doc["id"], + fields=doc["fields"], + schema=schema or self.default_schema_name, + namespace=self.namespace, + timeout=10, + ) + ) + tasks.append(task) + + results = await asyncio.wait(tasks, return_when=asyncio.ALL_COMPLETED) + for result in results: + if result.exception(): + raise result.exception + return ids + + def delete( + self, + ref_doc_id: str, + namespace: Optional[str] = None, + **delete_kwargs: Any, + ) -> None: + """ + Delete nodes using with ref_doc_id. + """ + response: VespaResponse = self.app.delete_data( + schema=self.default_schema_name, + namespace=namespace or self.namespace, + data_id=ref_doc_id, + kwargs=delete_kwargs, + ) + if not response.is_successful(): + raise ValueError( + f"Delete request failed: {response.status_code}, response payload: {response.json}" + ) + logger.info(f"Deleted node with id {ref_doc_id}") + + async def adelete( + self, + ref_doc_id: str, + namespace: Optional[str] = None, + **delete_kwargs: Any, + ) -> None: + """ + Delete nodes using with ref_doc_id. + NOTE: this is not implemented for all vector stores. If not implemented, + it will just call delete synchronously. + """ + logger.info("Async delete not implemented. Will call delete synchronously.") + self.delete(ref_doc_id, **delete_kwargs) + + def _create_query_body( + self, + query: VectorStoreQuery, + sources_str: str, + rank_profile: Optional[str] = None, + create_embedding: bool = True, + vector_top_k: int = 10, + ) -> dict: + """ + Create query parameters for Vespa. + + Args: + query (VectorStoreQuery): VectorStoreQuery object + sources_str (str): Sources string + rank_profile (Optional[str]): Rank profile to use. If not provided, default rank profile is used. + create_embedding (bool): Whether to create embedding + vector_top_k (int): Number of top k vectors to return + + Returns: + dict: Query parameters + """ + logger.info(f"Query: {query}") + if query.filters: + logger.warning("Filter support not implemented yet. Will be ignored.") + if query.alpha: + logger.warning( + "Alpha support not implemented. Must be defined in Vespa rank profile. " + "See for example https://pyvespa.readthedocs.io/en/latest/examples/evaluating-with-snowflake-arctic-embed.html" + ) + + if query.query_embedding is None and not create_embedding: + raise ValueError( + "Input embedding must be provided if embeddings are not created outside Vespa" + ) + + base_params = { + "hits": query.similarity_top_k, + "ranking.profile": rank_profile + or self._get_default_rank_profile(query.mode), + "query": query.query_str, + "tracelevel": 9, + } + logger.debug(query.mode) + if query.mode in [ + VectorStoreQueryMode.TEXT_SEARCH, + VectorStoreQueryMode.DEFAULT, + ]: + query_params = {"yql": f"select * from {sources_str} where userQuery()"} + elif query.mode in [ + VectorStoreQueryMode.SEMANTIC_HYBRID, + VectorStoreQueryMode.HYBRID, + ]: + if not query.embedding_field: + embedding_field = "embedding" + logger.warning( + f"Embedding field not provided. Using default embedding field {embedding_field}" + ) + query_params = { + "yql": f"select * from {sources_str} where {self._build_query_filter(query.mode, embedding_field, vector_top_k, query.similarity_top_k)}", + "input.query(q)": ( + f"embed({query.query_str})" + if create_embedding + else query.query_embedding + ), + } + else: + raise NotImplementedError( + f"Query mode {query.mode} not implemented for Vespa yet. Contributions are welcome!" + ) + + return {**base_params, **query_params} + + def _get_default_rank_profile(self, mode): + return { + VectorStoreQueryMode.TEXT_SEARCH: "bm25", + VectorStoreQueryMode.SEMANTIC_HYBRID: "fusion", + VectorStoreQueryMode.HYBRID: "fusion", + VectorStoreQueryMode.DEFAULT: "bm25", + }.get(mode) + + def _build_query_filter( + self, mode, embedding_field, vector_top_k, similarity_top_k + ): + """ + Build query filter for Vespa query. + The part after "select * from {sources_str} where" in the query. + """ + if mode in [ + VectorStoreQueryMode.SEMANTIC_HYBRID, + VectorStoreQueryMode.HYBRID, + ]: + return f"rank({{targetHits:{vector_top_k}}}nearestNeighbor({embedding_field},q), userQuery()) limit {similarity_top_k}" + else: + raise ValueError(f"Query mode {mode} not supported.") + + def query( + self, + query: VectorStoreQuery, + sources: Optional[List[str]] = None, + rank_profile: Optional[str] = None, + vector_top_k: int = 10, + **kwargs: Any, + ) -> VectorStoreQueryResult: + """Query vector store.""" + logger.debug(f"Query: {query}") + sources_str = ",".join(sources) if sources else "sources *" + mode = query.mode + body = self._create_query_body( + query=query, + sources_str=sources_str, + rank_profile=rank_profile, + create_embedding=not self.embeddings_outside_vespa, + vector_top_k=vector_top_k, + ) + logger.info(f"Vespa Query body:\n {body}") + with self.app.syncio() as session: + response = session.query( + body=body, + ) + if not response.is_successful(): + raise ValueError( + f"Query request failed: {response.status_code}, response payload: {response.get_json()}" + ) + logger.debug("Response:") + logger.debug(response.json) + logger.debug("Hits:") + logger.debug(response.hits) + nodes = [] + ids: List[str] = [] + similarities: List[float] = [] + for hit in response.hits: + response_fields: dict = hit.get("fields", {}) + metadata = response_fields.get("metadata", {}) + metadata = json.loads(metadata) + logger.debug(f"Metadata: {metadata}") + node = metadata_dict_to_node(metadata) + text = response_fields.get("body", "") + node.set_content(text) + nodes.append(node) + ids.append(response_fields.get("id")) + similarities.append(hit["relevance"]) + return VectorStoreQueryResult(nodes=nodes, ids=ids, similarities=similarities) + + async def aquery( + self, + query: VectorStoreQuery, + sources: Optional[List[str]] = None, + rank_profile: Optional[str] = None, + vector_top_k: int = 10, + **kwargs: Any, + ) -> VectorStoreQueryResult: + """ + Asynchronously query vector store. + NOTE: this is not implemented for all vector stores. If not implemented, + it will just call query synchronously. + """ + logger.info("Async query not implemented. Will call query synchronously.") + return self.query( + query=query, + sources=sources, + rank_profile=rank_profile, + vector_top_k=vector_top_k, + **kwargs, + ) + + def persist( + self, + ) -> None: + return NotImplemented("Persist is not implemented for VespaVectorStore") diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/llama_index/vector_stores/vespa/templates.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/llama_index/vector_stores/vespa/templates.py new file mode 100644 index 0000000000000..18570bd9e620c --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/llama_index/vector_stores/vespa/templates.py @@ -0,0 +1,96 @@ +try: + from vespa.package import ( + ApplicationPackage, + Field, + Schema, + Document, + HNSW, + RankProfile, + Component, + Parameter, + FieldSet, + GlobalPhaseRanking, + Function, + ) +except ImportError: + raise ModuleNotFoundError( + "pyvespa not installed. Please install it via `pip install pyvespa`" + ) + +hybrid_template = ApplicationPackage( + name="hybridsearch", + schema=[ + Schema( + name="doc", + document=Document( + fields=[ + Field(name="id", type="string", indexing=["summary"]), + Field(name="metadata", type="string", indexing=["summary"]), + Field( + name="text", + type="string", + indexing=["index", "summary"], + index="enable-bm25", + bolding=True, + ), + Field( + name="embedding", + type="tensor(x[384])", + indexing=[ + "input text", + "embed", + "index", + "attribute", + ], + ann=HNSW(distance_metric="angular"), + is_document_field=False, + ), + ] + ), + fieldsets=[FieldSet(name="default", fields=["text", "metadata"])], + rank_profiles=[ + RankProfile( + name="bm25", + inputs=[("query(q)", "tensor(x[384])")], + functions=[Function(name="bm25sum", expression="bm25(text)")], + first_phase="bm25sum", + ), + RankProfile( + name="semantic", + inputs=[("query(q)", "tensor(x[384])")], + first_phase="closeness(field, embedding)", + ), + RankProfile( + name="fusion", + inherits="bm25", + inputs=[("query(q)", "tensor(x[384])")], + first_phase="closeness(field, embedding)", + global_phase=GlobalPhaseRanking( + expression="reciprocal_rank_fusion(bm25sum, closeness(field, embedding))", + rerank_count=1000, + ), + ), + ], + ) + ], + components=[ + Component( + id="e5", + type="hugging-face-embedder", + parameters=[ + Parameter( + "transformer-model", + { + "url": "https://github.com/vespa-engine/sample-apps/raw/master/simple-semantic-search/model/e5-small-v2-int8.onnx" + }, + ), + Parameter( + "tokenizer-model", + { + "url": "https://raw.githubusercontent.com/vespa-engine/sample-apps/master/simple-semantic-search/model/tokenizer.json" + }, + ), + ], + ) + ], +) diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/pyproject.toml b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/pyproject.toml new file mode 100644 index 0000000000000..9830a544a05c4 --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/pyproject.toml @@ -0,0 +1,65 @@ +[build-system] +build-backend = "poetry.core.masonry.api" +requires = ["poetry-core"] + +[tool.codespell] +check-filenames = true +check-hidden = true +skip = "*.csv,*.html,*.json,*.jsonl,*.pdf,*.txt,*.ipynb" + +[tool.llamahub] +contains_example = true +import_path = "llama_index.vector_stores.vespa" + +[tool.llamahub.class_authors] +VespaVectorStore = "thomasht86" + +[tool.mypy] +disallow_untyped_defs = true +exclude = ["_static", "build", "examples", "notebooks", "venv"] +ignore_missing_imports = true +python_version = "3.8" + +[tool.poetry] +authors = ["Thomas Thoresen "] +description = "llama-index vector_stores vespa integration" +exclude = ["**/BUILD"] +license = "MIT" +name = "llama-index-vector-stores-vespa" +readme = "README.md" +version = "0.0.1" + +[tool.poetry.dependencies] +python = ">=3.8.1,<4.0" +llama-index-core = "^0.10.1" +pyvespa = "^0.40.0" + +[tool.poetry.group.dev.dependencies] +ipython = "8.10.0" +jupyter = "^1.0.0" +mypy = "0.991" +pre-commit = "3.2.0" +pylint = "2.15.10" +pytest = "7.2.1" +pytest-asyncio = "0.23.6" +pytest-mock = "3.11.1" +pyvespa = "0.40.0" +ruff = "0.0.292" +tree-sitter-languages = "^1.8.0" +types-Deprecated = ">=0.1.0" +types-PyYAML = "^6.0.12.12" +types-protobuf = "^4.24.0.4" +types-redis = "4.5.5.0" +types-requests = "2.28.11.8" +types-setuptools = "67.1.0.0" + +[tool.poetry.group.dev.dependencies.black] +extras = ["jupyter"] +version = "<=23.9.1,>=23.7.0" + +[tool.poetry.group.dev.dependencies.codespell] +extras = ["toml"] +version = ">=v2.2.6" + +[[tool.poetry.packages]] +include = "llama_index/" diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/tests/BUILD b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/tests/BUILD new file mode 100644 index 0000000000000..dabf212d7e716 --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/tests/BUILD @@ -0,0 +1 @@ +python_tests() diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/tests/__init__.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/tests/__init__.py new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/tests/test_vespavectorstore.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/tests/test_vespavectorstore.py new file mode 100644 index 0000000000000..44bbc6b90d5b4 --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/tests/test_vespavectorstore.py @@ -0,0 +1,154 @@ +import asyncio +import pytest + +from llama_index.core.schema import TextNode +from llama_index.core.vector_stores.types import ( + VectorStoreQuery, + VectorStoreQueryMode, +) +from vespa.application import ApplicationPackage +from llama_index.vector_stores.vespa import VespaVectorStore, hybrid_template + +try: + # Should be installed as pyvespa-dependency + import docker + + client = docker.from_env() + docker_available = client.ping() +except Exception: + docker_available = False + + +# Assuming Vespa services are mocked or local Vespa Docker is used +@pytest.fixture(scope="session") +def vespa_app(): + app_package: ApplicationPackage = hybrid_template + return VespaVectorStore(application_package=app_package, deployment_target="local") + + +@pytest.fixture(scope="session") +def nodes() -> list: + return [ + TextNode( + text="The Shawshank Redemption", + metadata={ + "id": "1", + "author": "Stephen King", + "theme": "Friendship", + "year": 1994, + }, + ), + TextNode( + text="The Godfather", + metadata={ + "id": "2", + "director": "Francis Ford Coppola", + "theme": "Mafia", + "year": 1972, + }, + ), + TextNode( + text="Inception", + metadata={ + "id": "3", + "director": "Christopher Nolan", + "theme": "Fiction", + "year": 2010, + }, + ), + TextNode( + text="To Kill a Mockingbird", + metadata={ + "id": "4", + "author": "Harper Lee", + "theme": "Mafia", + "year": 1960, + }, + ), + TextNode( + text="1984", + metadata={ + "id": "5", + "author": "George Orwell", + "theme": "Totalitarianism", + "year": 1949, + }, + ), + TextNode( + text="The Great Gatsby", + metadata={ + "id": "6", + "author": "F. Scott Fitzgerald", + "theme": "The American Dream", + "year": 1925, + }, + ), + TextNode( + text="Harry Potter and the Sorcerer's Stone", + metadata={ + "id": "7", + "author": "J.K. Rowling", + "theme": "Fiction", + "year": 1997, + }, + ), + ] + + +@pytest.fixture(scope="session") +def added_node_ids(vespa_app, nodes): + return vespa_app.add(nodes) + # Assume returned `inserted_ids` is a list of IDs that match the order of `nodes` + + +@pytest.mark.skipif(not docker_available, reason="Docker not available") +def test_query_text_search(vespa_app, added_node_ids): + query = VectorStoreQuery( + query_str="Inception", # Ensure the query matches the case used in the nodes + mode="text_search", + similarity_top_k=1, + ) + result = vespa_app.query(query) + assert len(result.nodes) == 1 + node_metadata = result.nodes[0].metadata + assert node_metadata["id"] == "3", "Expected Inception node" + + +@pytest.mark.skipif(not docker_available, reason="Docker not available") +def test_query_vector_search(vespa_app, added_node_ids): + query = VectorStoreQuery( + query_str="magic, wizardry", + mode="semantic_hybrid", + similarity_top_k=1, + ) + result = vespa_app.query(query) + assert len(result.nodes) == 1, "Expected 1 result" + node_metadata = result.nodes[0].metadata + print(node_metadata) + assert node_metadata["id"] == "7", "Expected Harry Potter node" + + +@pytest.mark.skipif(not docker_available, reason="Docker not available") +def test_delete_node(vespa_app, added_node_ids): + # Testing the deletion of a node + vespa_app.delete(ref_doc_id=added_node_ids[1]) + query = VectorStoreQuery( + query_str="Godfather", + mode=VectorStoreQueryMode.TEXT_SEARCH, + similarity_top_k=1, + ) + result = vespa_app.query(query) + assert ( + len(result.nodes) == 0 + ), f"Deleted node still present in the vector store: {result.nodes}" + + +@pytest.mark.skipif(not docker_available, reason="Docker not available") +@pytest.mark.asyncio() +async def test_async_add_and_query(vespa_app, nodes): + # Testing async add and query + await asyncio.gather(*[vespa_app.async_add(nodes)]) + query = VectorStoreQuery(query_str="Harry Potter", similarity_top_k=1) + result = await vespa_app.aquery(query) + assert len(result.nodes) == 1 + assert result.nodes[0].node_id == "7"