Forward backend options during streaming speech generation (#1389) (#1390)

Switch SQLite ANN connections when load() is given a new path

.github/workflows/build.yml

@@ -10,7 +10,7 @@ jobs:
       matrix:
         os: [ubuntu-latest, macos-latest, windows-latest]
 
-    timeout-minutes: 60
+    timeout-minutes: 75
     steps:
       - name: Checkout code
         uses: actions/checkout@v6
@@ -18,7 +18,7 @@ jobs:
       - name: Install Python
         uses: actions/setup-python@v6
         with:
-          python-version: "3.10"
+          python-version: "3.11"
 
       - name: Install Java
         uses: actions/setup-java@v5
@@ -52,7 +52,7 @@ jobs:
       - name: Build
         run: |
           pip install -U wheel
-          pip install .[all,dev] "litellm<1.97.0"
+          pip install .[all,dev]
           pip cache purge
           python -c "import nltk; nltk.download(['punkt', 'punkt_tab', 'averaged_perceptron_tagger_eng'])"
           python --version
@@ -68,7 +68,9 @@ jobs:
         if: matrix.os == 'ubuntu-latest'
 
       - name: Test Coverage
-        run: coveralls --service=github
+        run: |
+          coverage report --fail-under=100
+          coveralls --service=github
         if: matrix.os == 'ubuntu-latest'
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/docs.yml

@@ -13,7 +13,7 @@ jobs:
       - name: Install Python
         uses: actions/setup-python@v6
         with:
-          python-version: "3.10"
+          python-version: "3.11"
 
       - name: Install txtai
         run: |

.github/workflows/minimal.yml

@@ -13,7 +13,7 @@ jobs:
       - name: Install Python
         uses: actions/setup-python@v6
         with:
-          python-version: "3.10"
+          python-version: "3.11"
 
       - name: Minimal install
         run: |

.pylintrc

@@ -12,6 +12,8 @@ max-locals=40
 max-returns=10
 max-attributes=20
 min-public-methods=0
+max-public-methods=80
 
 [FORMAT]
 max-line-length=150
+max-module-lines=1500

README.md

@@ -49,7 +49,7 @@ Summary of txtai features:
 - 🔋 Batteries included with defaults to get up and running fast
 - ☁️ Run local or scale out with container orchestration
 
-txtai is built with Python 3.10+, [Hugging Face Transformers](https://github.com/huggingface/transformers), [Sentence Transformers](https://github.com/UKPLab/sentence-transformers) and [FastAPI](https://github.com/tiangolo/fastapi). txtai is open-source under an Apache 2.0 license.
+txtai is built with Python 3.11+, [Hugging Face Transformers](https://github.com/huggingface/transformers), [Sentence Transformers](https://github.com/UKPLab/sentence-transformers) and [FastAPI](https://github.com/tiangolo/fastapi). txtai is open-source under an Apache 2.0 license.
 
 > [!NOTE]
 >
@@ -190,7 +190,7 @@ The easiest way to install is via pip and PyPI
 pip install txtai
 ```
 
-Python 3.10+ is supported. Using a Python [virtual environment](https://docs.python.org/3/library/venv.html) is recommended.
+Python 3.11+ is supported. Using a Python [virtual environment](https://docs.python.org/3/library/venv.html) is recommended.
 
 See the detailed [install instructions](https://neuml.github.io/txtai/install) for more information covering [optional dependencies](https://neuml.github.io/txtai/install/#optional-dependencies), [environment specific prerequisites](https://neuml.github.io/txtai/install/#environment-specific-prerequisites), [installing from source](https://neuml.github.io/txtai/install/#install-from-source), [conda support](https://neuml.github.io/txtai/install/#conda), [lightweight minimal installation](https://neuml.github.io/txtai/install/#minimal-install) and how to [run with containers](https://neuml.github.io/txtai/cloud).
 

docker/base/Dockerfile

@@ -1,5 +1,5 @@
 # Set base image
-ARG BASE_IMAGE=python:3.10-slim
+ARG BASE_IMAGE=python:3.11-slim
 FROM $BASE_IMAGE
 
 # Install GPU-enabled version of PyTorch if set
@@ -8,7 +8,7 @@ ARG GPU
 # Target CPU architecture
 ARG TARGETARCH
 
-# Set Python version (i.e. 3, 3.10)
+# Set Python version (i.e. 3, 3.11)
 ARG PYTHON_VERSION=3
 
 # List of txtai components to install
@@ -27,7 +27,7 @@ RUN \
     # Install txtai project and dependencies
     ln -s /usr/bin/python${PYTHON_VERSION} /usr/bin/python && \
     python -m pip install --no-cache-dir -U pip wheel setuptools && \
-    if [ -z ${GPU} ]; then pip install --no-cache-dir torch==2.12.1+cpu torchvision==0.27.1+cpu -f https://download.pytorch.org/whl/torch -f https://download.pytorch.org/whl/torchvision; fi && \
+    if [ -z ${GPU} ]; then pip install --no-cache-dir torch==2.14.1+cpu torchvision==0.29.1+cpu -f https://download.pytorch.org/whl/torch -f https://download.pytorch.org/whl/torchvision; fi && \
     python -m pip install --no-cache-dir txtai${COMPONENTS} && \
     python -c "import sys, importlib.util as util; 1 if util.find_spec('nltk') else sys.exit(); import nltk; nltk.download(['punkt', 'punkt_tab', 'averaged_perceptron_tagger_eng'])" && \
     \

docker/minimal/Dockerfile

@@ -1,11 +1,11 @@
 # Set base image
-ARG BASE_IMAGE=python:3.10-slim
+ARG BASE_IMAGE=python:3.11-slim
 FROM $BASE_IMAGE
 
 # Target CPU architecture
 ARG TARGETARCH
 
-# Set Python version (i.e. 3, 3.10)
+# Set Python version (i.e. 3, 3.11)
 ARG PYTHON_VERSION=3
 
 # Locale environment variables

docs/cloud.md

@@ -30,8 +30,8 @@ Examples build commands below.
 # Get Dockerfile
 wget https://raw.githubusercontent.com/neuml/txtai/master/docker/base/Dockerfile
 
-# Build Ubuntu 22.04 image running Python 3.10
-docker build -t txtai --build-arg BASE_IMAGE=ubuntu:22.04 --build-arg PYTHON_VERSION=3.10 .
+# Build Ubuntu 22.04 image running Python 3.11
+docker build -t txtai --build-arg BASE_IMAGE=ubuntu:22.04 --build-arg PYTHON_VERSION=3.11 .
 
 # Build image with GPU support
 docker build -t txtai --build-arg GPU=1 .

docs/embeddings/configuration/ann.md

@@ -4,7 +4,7 @@ Approximate Nearest Neighbor (ANN) index configuration for storing vector embedd
 
 ## backend
 ```yaml
-backend: faiss|hnsw|annoy|ggml|milvus|numpy|torch|turbovec|zvec|pgvector|sqlite|custom
+backend: faiss|hnsw|annoy|ggml|milvus|numpy|rabitq|torch|turbovec|zvec|pgvector|sqlite|custom
 ```
 
 Sets the ANN backend. Defaults to `faiss`. Additional backends are available via the [ann](../../../install/#ann) extras package. Set custom backends via setting this parameter to the fully resolvable class string.
@@ -91,6 +91,26 @@ numpy:
 
 The NumPy backend is a k-nearest neighbors backend. It's designed for simplicity and works well with smaller datasets that fit into memory.
 
+### rabitq
+
+```yaml
+rabitq:
+    mode: index mode (ivf or hnsw) - defaults to "ivf"
+    nbits: number of quantization bits per dimension (int) - defaults to 1,
+           supports 1 to 9 in both modes, ivf mode also supports 32 which
+           keeps raw vectors for reranking
+    clusters: number of IVF clusters (int) - defaults to
+              max(1, min(round(4 * sqrt(embeddings count)), embeddings count))
+    nprobe: search probe setting for ivf mode (int) - defaults to
+            max(1, round(num_clusters/16))
+    m: M param for hnsw mode (int) - defaults to 16
+    efconstruction: ef_construction param for hnsw mode (int) - defaults to 200
+    efsearch: ef search param for hnsw mode (int) - defaults to None and not set
+    randomseed: random-seed param for hnsw mode (int) - defaults to 100
+```
+
+The [rabitq](https://github.com/VectorDB-NTU/RaBitQ-Library) backend is a quantized index powered by the [RaBitQ algorithm](https://doi.org/10.1145/3725413). It supports ivf and hnsw search modes. Vectors are stored with 1-bit quantization by default. More bits increase accuracy at the cost of memory. The index is stored as a single file and no vectors are kept outside of it. With nbits set to 32, the vectors are stored in the index for reranking. The ivf mode supports upserts and deletes. Deleted rows keep their storage. Rows added after the first index are assigned to the clusters trained when the index was built, so recall can drop if the new data differs from the original data. Rebuild the index to train new clusters. Note that hnsw indexes can not be modified after creation, upserts/deletes and other modifications are not supported.
+
 ### torch
 
 ```yaml
@@ -119,7 +139,9 @@ The [turbovec](https://github.com/RyanCodrai/turbovec) backend is a k-nearest ne
 
 ```yaml
 zvec:
+    efconstruction: ef_construction param for HnswIndexParam (int) - defaults to 200
     m: number of HNSW links per element (int) - defaults to 50
+    efsearch: ef search param for HnswQueryParam (int) - defaults to 300
 ```
 
 The [zvec](https://github.com/alibaba/zvec) backend is an embedded, path-based vector index.

docs/embeddings/configuration/scoring.md

@@ -36,7 +36,7 @@ batch: Sets the transform batch size
 encodebatch: Sets the encode batch size
 vectors: additional model init args
 encodeargs: additional encode() args
-backend: ivfsparse|pgsparse
+backend: ivfsparse|pgsparse|zvecsparse
 ```
 
 Sparse vector scoring options. The sparse scoring instance combines a sparse vector model with a sparse approximate nearest neighbor index (ANN). This method supports both vector normalization and score normalization.
@@ -54,7 +54,7 @@ Score normalization scales the output between 0 and 1. This setting supports:
 ```yaml
 ivfsparse:
   sample: percent of data to use for model training (0.0 - 1.0)
-  nfeatures: top n features to use for model training (int)
+  nfeatures: top n features to use for model training, defaults to all features (int)
   nlist: desired number of clusters (int)
   nprobe: search probe setting (int)
   minpoints: minimum number of points for a cluster (int)
@@ -66,6 +66,16 @@ Inverted file (IVF) index with flat vector file storage and sparse array support
 
 Sparse ANN backed by Postgres. Supports same options as the [pgvector](../ann/#pgvector) ANN.
 
+#### zvecsparse
+```yaml
+zvecsparse:
+  efconstruction: ef_construction param for HnswIndexParam (int) - defaults to 200
+  m: number of HNSW links per element (int) - defaults to 50
+  efsearch: ef search param for HnswQueryParam (int) - defaults to 300
+```
+
+Sparse ANN backed by [zvec](https://github.com/alibaba/zvec). Stores sparse vectors in an embedded, path-based vector index. Requires the [ann](../../../install/#ann) extras package.
+
 ## terms
 ```yaml
 terms: boolean|dict
@@ -81,6 +91,15 @@ Supports a `dict` with the parameters `cachelimit` and `cutoff`.
 
 When `terms` is set to `True`, default parameters are used for the `cachelimit` and `cutoff`. Normally, these defaults are sufficient.
 
+### Wildcard queries
+
+Keyword queries containing `*` expand matching terms using SQL `LIKE`. Within these
+expressions, `%` matches any sequence and `_` matches one character. To match a literal
+percent sign, underscore or backslash, use `\%`, `\_` or `\\`, respectively.
+For example, `report\_*` matches `report_2026`, while `report_*` also matches `reportX2026`.
+Use a tokenizer that preserves these characters, such as `tokenizer: {whitespace: true}`,
+or pass pre-tokenized query terms. Queries without `*` retain exact-term matching.
+
 ## normalize
 ```yaml
 normalize: boolean|str|dict

docs/embeddings/configuration/vectors.md

@@ -176,7 +176,8 @@ muvera:
 ```
 
 Settings to control the size of MUVERA fixed dimensional outputs. Default is 20 * 2^5 * 16 = 10,240 dimensions.
-Set `muvera` to `false` to disable it or `true` to use the default settings.
+Set `muvera` to `false` to disable it or `true` to use the default settings. `false` only makes sense when the multi-vector output
+is consumed directly, for example through `PoolingFactory`. Embeddings vectors require `muvera` or `lemur`.
 
 ### lemur
 ```yaml

docs/embeddings/query.md

@@ -54,6 +54,8 @@ The number of candidates should be larger than the desired number of results whe
 - For a single query filter clause, the default is the query limit
 - With multiple filtering clauses, the default is 10x the query limit
 
+When a query orders results by score, its `offset` adds to the query limit.
+
 The index name is only applicable when [subindexes](../configuration/general/#indexes) are enabled. This specifies the index to use for the query.
 
 Weights sets the hybrid score weights when an index has both a sparse and dense index.

docs/index.md

@@ -51,7 +51,7 @@ Summary of txtai features:
 - 🔋 Batteries included with defaults to get up and running fast
 - ☁️ Run local or scale out with container orchestration
 
-txtai is built with Python 3.10+, [Hugging Face Transformers](https://github.com/huggingface/transformers), [Sentence Transformers](https://github.com/UKPLab/sentence-transformers) and [FastAPI](https://github.com/tiangolo/fastapi). txtai is open-source under an Apache 2.0 license.
+txtai is built with Python 3.11+, [Hugging Face Transformers](https://github.com/huggingface/transformers), [Sentence Transformers](https://github.com/UKPLab/sentence-transformers) and [FastAPI](https://github.com/tiangolo/fastapi). txtai is open-source under an Apache 2.0 license.
 
 !!! note
 

docs/install.md

@@ -9,7 +9,7 @@ The easiest way to install is via pip and PyPI
 pip install txtai
 ```
 
-Python 3.10+ is supported. Using a Python [virtual environment](https://docs.python.org/3/library/venv.html) is recommended.
+Python 3.11+ is supported. Using a Python [virtual environment](https://docs.python.org/3/library/venv.html) is recommended.
 
 ## Optional dependencies
 

docs/workflow/task/index.md

@@ -42,6 +42,19 @@ workflow:
 
 ::: txtai.workflow.Task.__init__
 
+## Multi-action task inputs
+
+A multi-action `Task` requires reusable inputs. Passing a single-pass iterator, such as a generator or the output of a `StreamTask`, raises a `ValueError` before consuming inputs or running any actions. Automatically buffering an iterator could exhaust memory, so tasks do not convert these inputs to a list.
+
+Iterators passed to a `Workflow` remain supported: the workflow collects bounded batches before passing them to its tasks. For example:
+
+```python
+workflow = Workflow([Task([list, list])], batch=2)
+list(workflow(iter([1, 2, 3])))  # [(1, 1), (2, 2), (3, 3)]
+```
+
+This batching occurs at the workflow input; it does not batch iterator outputs produced by tasks inside that workflow. Single-action tasks continue to accept iterator inputs.
+
 ## Multi-action task concurrency
 
 The default processing mode is to run actions sequentially. Multiprocessing support is already built in at a number of levels. Any of the GPU models will maximize GPU utilization for example and even in CPU mode, concurrency is utilized. But there are still use cases for task action concurrency. For example, if the system has multiple GPUs, the task runs external sequential code, or the task has a large number of I/O tasks.
@@ -73,7 +86,7 @@ workflow = Workflow([Task(lambda x: [y * 3 for y in x], unpack=False, column=0)]
 list(workflow([(2, 8)]))
 ```
 
-For the example input tuple of (2, 2), the workflow will only select the first element (2) and run the task against that element. 
+For the example input tuple of (2, 8), the workflow will only select the first element (2) and run the task against that element.
 
 ```python
 workflow = Workflow([Task([lambda x: [y * 3 for y in x], 

examples/01_Introducing_txtai.ipynb

@@ -27,7 +27,7 @@
         "- 🔋 Batteries included with defaults to get up and running fast\n",
         "- ☁️ Run local or scale out with container orchestration\n",
         "\n",
-        "txtai is built with Python 3.10+, [Hugging Face Transformers](https://github.com/huggingface/transformers), [Sentence Transformers](https://github.com/UKPLab/sentence-transformers) and [FastAPI](https://github.com/tiangolo/fastapi). txtai is open-source under an Apache 2.0 license.\n",
+        "txtai is built with Python 3.11+, [Hugging Face Transformers](https://github.com/huggingface/transformers), [Sentence Transformers](https://github.com/UKPLab/sentence-transformers) and [FastAPI](https://github.com/tiangolo/fastapi). txtai is open-source under an Apache 2.0 license.\n",
         "\n",
         "> [NeuML](https://neuml.com) is the company behind txtai and we provide AI consulting services around our stack. [Schedule a meeting](https://cal.com/neuml/intro) or [send a message](mailto:[email protected]) to learn more.\n",
         ">\n",

examples/52_Build_RAG_pipelines_with_txtai.ipynb

@@ -125,7 +125,7 @@
             "## Install\n",
             "The easiest way to install is via pip and PyPI\n",
             "pip install txtai\n",
-            "Python 3.10+ is supported. Using a Python virtual environment is **recommended** .\n",
+            "Python 3.11+ is supported. Using a Python virtual environment is **recommended** .\n",
             "See the detailed install instructions for more information covering optional dependencies, environment specific prerequisites, installing from source, conda support and how to run with containers.\n",
             "\n",
             "\n",
@@ -525,8 +525,8 @@
           "name": "stdout",
           "output_type": "stream",
           "text": [
-            "ANSWER: Python 3.10 and later versions (Python 3.10+) are supported.\n",
-            "CITATION: [{'id': '24', 'text': 'Python 3.10+ is supported. Using a Python virtual environment is recommended.'}]\n"
+            "ANSWER: Python 3.11 and later versions (Python 3.11+) are supported.\n",
+            "CITATION: [{'id': '24', 'text': 'Python 3.11+ is supported. Using a Python virtual environment is recommended.'}]\n"
           ]
         }
       ],

examples/73_Chunking_your_data_for_RAG.ipynb

@@ -100,12 +100,12 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "{'indexid': 0, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': '**GitHub - neuml/txtai: 💡 All-in-one open-source embeddings database for semantic search, LLM orchestration and language model workflows**\\n\\n*💡 All-in-one open-source embeddings database for semantic search, LLM orchestration and language model workflows - neuml/txtai*\\n\\n\\n\\n**All-in-one embeddings database** \\ntxtai is an all-in-one embeddings database for semantic search, LLM orchestration and language model workflows.\\n\\nEmbeddings databases are a union of vector indexes (sparse and dense), graph networks and relational databases.\\n\\nThis foundation enables vector search and/or serves as a powerful knowledge source for large language model (LLM) applications.\\n\\nBuild autonomous agents, retrieval augmented generation (RAG) processes, multi-model workflows and more.\\n\\nSummary of txtai features:\\n\\n- 🔎 Vector search with SQL, object storage, topic modeling, graph analysis and multimodal indexing\\n- 📄 Create embeddings for text, documents, audio, images and video\\n- 💡 Pipelines powered by language models that run LLM prompts, question-answering, labeling, transcription, translation, summarization and more\\n- ↪️️ Workflows to join pipelines together and aggregate business logic. txtai processes can be simple microservices or multi-model workflows.\\n- 🤖 Agents that intelligently connect embeddings, pipelines, workflows and other agents together to autonomously solve complex problems\\n- ⚙️ Build with Python or YAML. API bindings available for [JavaScript](https://github.com/neuml/txtai.js) , [Java](https://github.com/neuml/txtai.java) , [Rust](https://github.com/neuml/txtai.rs) and [Go](https://github.com/neuml/txtai.go) .\\n- 🔋 Batteries included with defaults to get up and running fast\\n- ☁️ Run local or scale out with container orchestration\\ntxtai is built with Python 3.10+, [Hugging Face Transformers](https://github.com/huggingface/transformers) , [Sentence Transformers](https://github.'}\n",
+      "{'indexid': 0, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': '**GitHub - neuml/txtai: 💡 All-in-one open-source embeddings database for semantic search, LLM orchestration and language model workflows**\\n\\n*💡 All-in-one open-source embeddings database for semantic search, LLM orchestration and language model workflows - neuml/txtai*\\n\\n\\n\\n**All-in-one embeddings database** \\ntxtai is an all-in-one embeddings database for semantic search, LLM orchestration and language model workflows.\\n\\nEmbeddings databases are a union of vector indexes (sparse and dense), graph networks and relational databases.\\n\\nThis foundation enables vector search and/or serves as a powerful knowledge source for large language model (LLM) applications.\\n\\nBuild autonomous agents, retrieval augmented generation (RAG) processes, multi-model workflows and more.\\n\\nSummary of txtai features:\\n\\n- 🔎 Vector search with SQL, object storage, topic modeling, graph analysis and multimodal indexing\\n- 📄 Create embeddings for text, documents, audio, images and video\\n- 💡 Pipelines powered by language models that run LLM prompts, question-answering, labeling, transcription, translation, summarization and more\\n- ↪️️ Workflows to join pipelines together and aggregate business logic. txtai processes can be simple microservices or multi-model workflows.\\n- 🤖 Agents that intelligently connect embeddings, pipelines, workflows and other agents together to autonomously solve complex problems\\n- ⚙️ Build with Python or YAML. API bindings available for [JavaScript](https://github.com/neuml/txtai.js) , [Java](https://github.com/neuml/txtai.java) , [Rust](https://github.com/neuml/txtai.rs) and [Go](https://github.com/neuml/txtai.go) .\\n- 🔋 Batteries included with defaults to get up and running fast\\n- ☁️ Run local or scale out with container orchestration\\ntxtai is built with Python 3.11+, [Hugging Face Transformers](https://github.com/huggingface/transformers) , [Sentence Transformers](https://github.'}\n",
       "{'indexid': 1, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': 'com/UKPLab/sentence-transformers) and [FastAPI](https://github.com/tiangolo/fastapi) . txtai is open-source under an Apache 2.0 license.\\n\\n*Interested in an easy and secure way to run hosted txtai applications? Then join the* [txtai.cloud](https://txtai.cloud) *preview to learn more.* \\n\\n## Why txtai?\\nNew vector databases, LLM frameworks and everything in between are sprouting up daily. Why build with txtai?\\n\\n- Up and running in minutes with [pip](https://neuml.github.io/txtai/install/) or [Docker](https://neuml.github.io/txtai/cloud/) \\n```\\n# Get started in a couple lines\\nimport txtai\\n\\nembeddings = txtai.Embeddings()\\nembeddings.index([\"Correct\", \"Not what we hoped\"])\\nembeddings.search(\"positive\", 1)\\n#[(0, 0.29862046241760254)]\\n```\\n\\n- Built-in API makes it easy to develop applications using your programming language of choice\\n```\\n# app.yml\\nembeddings:\\n path: sentence-transformers/all-MiniLM-L6-v2\\n```\\n\\n```\\nCONFIG=app.yml uvicorn \"txtai.api:app\"\\ncurl -X GET \"http://localhost:8000/search?query=positive\"\\n```\\n\\n- Run local - no need to ship data off to disparate remote services\\n- Work with micromodels all the way up to large language models (LLMs)\\n- Low footprint - install additional dependencies and scale up when needed\\n- [Learn by example](https://neuml.github.io/txtai/examples) - notebooks cover all available functionality\\n\\n## Use Cases\\nThe following sections introduce common txtai use cases. A comprehensive set of over 60 [example notebooks and applications](https://neuml.github.io/txtai/examples) are also available.\\n\\n\\n### Semantic Search\\nBuild semantic/similarity/vector/neural search applications.'}\n",
       "{'indexid': 2, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': 'Traditional search systems use keywords to find data. Semantic search has an understanding of natural language and identifies results that have the same meaning, not necessarily the same keywords.\\n\\nGet started with the following examples.\\n\\n|Notebook|Description||\\n|---|---|---|\\n|[Introducing txtai](https://github.com/neuml/txtai/blob/master/examples/01_Introducing_txtai.ipynb) |Overview of the functionality provided by txtai||\\n|[Similarity search with images](https://github.com/neuml/txtai/blob/master/examples/13_Similarity_search_with_images.ipynb) |Embed images and text into the same space for search||\\n|[Build a QA database](https://github.com/neuml/txtai/blob/master/examples/34_Build_a_QA_database.ipynb) |Question matching with semantic search||\\n|[Semantic Graphs](https://github.com/neuml/txtai/blob/master/examples/38_Introducing_the_Semantic_Graph.ipynb) |Explore topics, data connectivity and run network analysis||\\n\\n### LLM Orchestration\\nAutonomous agents, retrieval augmented generation (RAG), chat with your data, pipelines and workflows that interface with large language models (LLMs).\\n\\nSee below to learn more.\\n\\n|Notebook|Description||\\n|---|---|---|\\n|[Prompt templates and task chains](https://github.com/neuml/txtai/blob/master/examples/44_Prompt_templates_and_task_chains.ipynb) |Build model prompts and connect tasks together with workflows||\\n|[Integrate LLM frameworks](https://github.com/neuml/txtai/blob/master/examples/53_Integrate_LLM_Frameworks.ipynb) |Integrate llama.cpp, LiteLLM and custom generation frameworks||\\n|[Build knowledge graphs with LLMs](https://github.'}\n",
       "{'indexid': 3, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': 'com/neuml/txtai/blob/master/examples/57_Build_knowledge_graphs_with_LLM_driven_entity_extraction.ipynb) |Build knowledge graphs with LLM-driven entity extraction||\\n\\n#### Agents\\nAgents connect embeddings, pipelines, workflows and other agents together to autonomously solve complex problems.\\n\\ntxtai agents are built on top of the Transformers Agent framework. This supports all LLMs txtai supports (Hugging Face, llama.cpp, OpenAI / Claude / AWS Bedrock via LiteLLM).\\n\\nSee the link below to learn more.\\n\\n|Notebook|Description||\\n|---|---|---|\\n|[Analyzing Hugging Face Posts with Graphs and Agents](https://github.com/neuml/txtai/blob/master/examples/68_Analyzing_Hugging_Face_Posts_with_Graphs_and_Agents.ipynb) |Explore a rich dataset with Graph Analysis and Agents||\\n|[Granting autonomy to agents](https://github.com/neuml/txtai/blob/master/examples/69_Granting_autonomy_to_agents.ipynb) |Agents that iteratively solve problems as they see fit||\\n|[Analyzing LinkedIn Company Posts with Graphs and Agents](https://github.com/neuml/txtai/blob/master/examples/71_Analyzing_LinkedIn_Company_Posts_with_Graphs_and_Agents.ipynb) |Exploring how to improve social media engagement with AI||\\n\\n#### Retrieval augmented generation\\nRetrieval augmented generation (RAG) reduces the risk of LLM hallucinations by constraining the output with a knowledge base as context. RAG is commonly used to \"chat with your data\".\\n\\nA novel feature of txtai is that it can provide both an answer and source citation.\\n\\n|Notebook|Description||\\n|---|---|---|\\n|[Build RAG pipelines with txtai](https://github.'}\n",
       "{'indexid': 4, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': 'com/neuml/txtai/blob/master/examples/52_Build_RAG_pipelines_with_txtai.ipynb) |Guide on retrieval augmented generation including how to create citations||\\n|[How RAG with txtai works](https://github.com/neuml/txtai/blob/master/examples/63_How_RAG_with_txtai_works.ipynb) |Create RAG processes, API services and Docker instances||\\n|[Advanced RAG with graph path traversal](https://github.com/neuml/txtai/blob/master/examples/58_Advanced_RAG_with_graph_path_traversal.ipynb) |Graph path traversal to collect complex sets of data for advanced RAG||\\n|[Speech to Speech RAG](https://github.com/neuml/txtai/blob/master/examples/65_Speech_to_Speech_RAG.ipynb) |Full cycle speech to speech workflow with RAG||\\n\\n### Language Model Workflows\\nLanguage model workflows, also known as semantic workflows, connect language models together to build intelligent applications.\\n\\nWhile LLMs are powerful, there are plenty of smaller, more specialized models that work better and faster for specific tasks. This includes models for extractive question-answering, automatic summarization, text-to-speech, transcription and translation.\\n\\n|Notebook|Description||\\n|---|---|---|\\n|[Run pipeline workflows](https://github.com/neuml/txtai/blob/master/examples/14_Run_pipeline_workflows.ipynb) |Simple yet powerful constructs to efficiently process data||\\n|[Building abstractive text summaries](https://github.com/neuml/txtai/blob/master/examples/09_Building_abstractive_text_summaries.ipynb) |Run abstractive text summarization||\\n|[Transcribe audio to text](https://github.com/neuml/txtai/blob/master/examples/11_Transcribe_audio_to_text.'}\n",
-      "{'indexid': 5, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': 'ipynb) |Convert audio files to text||\\n|[Translate text between languages](https://github.com/neuml/txtai/blob/master/examples/12_Translate_text_between_languages.ipynb) |Streamline machine translation and language detection||\\n\\n## Installation\\nThe easiest way to install is via pip and PyPI\\n\\n```\\npip install txtai\\n```\\n\\nPython 3.10+ is supported. Using a Python [virtual environment](https://docs.python.org/3/library/venv.html) is recommended.\\n\\nSee the detailed [install instructions](https://neuml.github.io/txtai/install) for more information covering [optional dependencies](https://neuml.github.io/txtai/install/#optional-dependencies) , [environment specific prerequisites](https://neuml.github.io/txtai/install/#environment-specific-prerequisites) , [installing from source](https://neuml.github.io/txtai/install/#install-from-source) , [conda support](https://neuml.github.io/txtai/install/#conda) and how to [run with containers](https://neuml.github.io/txtai/cloud) .\\n\\n\\n## Model guide\\nSee the table below for the current recommended models. These models all allow commercial use and offer a blend of speed and performance.\\n\\n|Component|Model(s)|\\n|---|---|\\n|[Embeddings](https://neuml.github.io/txtai/embeddings) |[all-MiniLM-L6-v2](https://hf.co/sentence-transformers/all-MiniLM-L6-v2) |\\n|[Image Captions](https://neuml.github.io/txtai/pipeline/image/caption) |[BLIP](https://hf.co/Salesforce/blip-image-captioning-base) |'}\n",
+      "{'indexid': 5, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': 'ipynb) |Convert audio files to text||\\n|[Translate text between languages](https://github.com/neuml/txtai/blob/master/examples/12_Translate_text_between_languages.ipynb) |Streamline machine translation and language detection||\\n\\n## Installation\\nThe easiest way to install is via pip and PyPI\\n\\n```\\npip install txtai\\n```\\n\\nPython 3.11+ is supported. Using a Python [virtual environment](https://docs.python.org/3/library/venv.html) is recommended.\\n\\nSee the detailed [install instructions](https://neuml.github.io/txtai/install) for more information covering [optional dependencies](https://neuml.github.io/txtai/install/#optional-dependencies) , [environment specific prerequisites](https://neuml.github.io/txtai/install/#environment-specific-prerequisites) , [installing from source](https://neuml.github.io/txtai/install/#install-from-source) , [conda support](https://neuml.github.io/txtai/install/#conda) and how to [run with containers](https://neuml.github.io/txtai/cloud) .\\n\\n\\n## Model guide\\nSee the table below for the current recommended models. These models all allow commercial use and offer a blend of speed and performance.\\n\\n|Component|Model(s)|\\n|---|---|\\n|[Embeddings](https://neuml.github.io/txtai/embeddings) |[all-MiniLM-L6-v2](https://hf.co/sentence-transformers/all-MiniLM-L6-v2) |\\n|[Image Captions](https://neuml.github.io/txtai/pipeline/image/caption) |[BLIP](https://hf.co/Salesforce/blip-image-captioning-base) |'}\n",
       "{'indexid': 6, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': '|[Labels - Zero Shot](https://neuml.github.io/txtai/pipeline/text/labels) |[BART-Large-MNLI](https://hf.co/facebook/bart-large) |\\n|[Labels - Fixed](https://neuml.github.io/txtai/pipeline/text/labels) |Fine-tune with [training pipeline](https://neuml.github.io/txtai/pipeline/train/trainer) |\\n|[Large Language Model (LLM)](https://neuml.github.io/txtai/pipeline/text/llm) |[Llama 3.1 Instruct](https://hf.co/meta-llama/Llama-3.1-8B-Instruct) |\\n|[Summarization](https://neuml.github.io/txtai/pipeline/text/summary) |[DistilBART](https://hf.co/sshleifer/distilbart-cnn-12-6) |\\n|[Text-to-Speech](https://neuml.github.io/txtai/pipeline/audio/texttospeech) |[ESPnet JETS](https://hf.co/NeuML/ljspeech-jets-onnx) |\\n|[Transcription](https://neuml.github.io/txtai/pipeline/audio/transcription) |[Whisper](https://hf.co/openai/whisper-base) |\\n|[Translation](https://neuml.github.io/txtai/pipeline/text/translation) |[OPUS Model Series](https://hf.co/Helsinki-NLP) |\\nModels can be loaded as either a path from the Hugging Face Hub or a local directory. Model paths are optional, defaults are loaded when not specified. For tasks with no recommended model, txtai uses the default models as shown in the Hugging Face Tasks guide.\\n\\nSee the following links to learn more.\\n\\n\\n## Powered by txtai\\nThe following applications are powered by txtai.'}\n",
       "{'indexid': 7, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': \"|Application|Description|\\n|---|---|\\n|[rag](https://github.com/neuml/rag) |Retrieval Augmented Generation (RAG) application|\\n|[ragdata](https://github.com/neuml/ragdata) |Build knowledge bases for RAG|\\n|[paperai](https://github.com/neuml/paperai) |Semantic search and workflows for medical/scientific papers|\\n|[annotateai](https://github.com/neuml/annotateai) |Automatically annotate papers with LLMs|\\nIn addition to this list, there are also many other [open-source projects](https://github.com/neuml/txtai/network/dependents) , [published research](https://scholar.google.com/scholar?q=txtai&hl=en&as_ylo=2022) and closed proprietary/commercial projects that have built on txtai in production.\\n\\n\\n## Further Reading\\n- [Tutorial series on Hashnode](https://neuml.hashnode.dev/series/txtai-tutorial) | [dev.to](https://dev.to/neuml/tutorial-series-on-txtai-ibg) \\n- [What's new in txtai 8.0](https://medium.com/neuml/whats-new-in-txtai-8-0-2d7d0ab4506b) | [7.0](https://medium.com/neuml/whats-new-in-txtai-7-0-855ad6a55440) | [6.0](https://medium.com/neuml/whats-new-in-txtai-6-0-7d93eeedf804) | [5.0](https://medium.com/neuml/whats-new-in-txtai-5-0-e5c75a13b101) | [4.0](https://medium.\"}\n",
       "{'indexid': 8, 'id': '0', 'url': 'https://github.com/neuml/txtai', 'text': 'com/neuml/whats-new-in-txtai-4-0-bbc3a65c3d1c) \\n- [Getting started with semantic search](https://medium.com/neuml/getting-started-with-semantic-search-a9fd9d8a48cf) | [workflows](https://medium.com/neuml/getting-started-with-semantic-workflows-2fefda6165d9) | [rag](https://medium.com/neuml/getting-started-with-rag-9a0cca75f748) \\n\\n## Documentation\\n[Full documentation on txtai](https://neuml.github.io/txtai) including configuration settings for embeddings, pipelines, workflows, API and a FAQ with common questions/issues is available.\\n\\n\\n## Contributing\\nFor those who would like to contribute to txtai, please see [this guide](https://github.com/neuml/.github/blob/master/CONTRIBUTING.md) .'}\n",

examples/84_Agent_Tools.ipynb

@@ -714,7 +714,7 @@
        "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Observations: ## Search Results\n",
        "\n",
        "|txtai · PyPI](https://pypi.org/project/txtai/)\n",
-       "☁️ Run local or scale out with container orchestration txtai is built with Python 3.10+, Hugging Face Transformers,\n",
+       "☁️ Run local or scale out with container orchestration txtai is built with Python 3.11+, Hugging Face Transformers,\n",
        "Sentence Transformers and FastAPI. txtai is open-source under an Apache 2.0 license. |!NOTE] NeuML is the company \n",
        "behind txtai and we provide AI consulting services around our stack. Schedule a meeting or send a message to ...\n",
        "\n",
@@ -767,7 +767,7 @@
        "Observations: ## Search Results\n",
        "\n",
        "|txtai · PyPI](https://pypi.org/project/txtai/)\n",
-       "☁️ Run local or scale out with container orchestration txtai is built with Python 3.10+, Hugging Face Transformers,\n",
+       "☁️ Run local or scale out with container orchestration txtai is built with Python 3.11+, Hugging Face Transformers,\n",
        "Sentence Transformers and FastAPI. txtai is open-source under an Apache 2.0 license. |!NOTE] NeuML is the company \n",
        "behind txtai and we provide AI consulting services around our stack. Schedule a meeting or send a message to ...\n",
        "\n",
@@ -896,7 +896,7 @@
        "|Rust](https://github.com/neuml/txtai.rs) and |Go](https://github.com/neuml/txtai.go) .\n",
        "- 🔋 Batteries included with defaults to get up and running fast\n",
        "- ☁️ Run local or scale out with container orchestration\n",
-       "txtai is built with Python 3.10+, |Hugging Face Transformers](https://github.com/huggingface/transformers) , \n",
+       "txtai is built with Python 3.11+, |Hugging Face Transformers](https://github.com/huggingface/transformers) , \n",
        "|Sentence Transformers](https://github.com/UKPLab/sentence-transformers) and \n",
        "|FastAPI](https://github.com/tiangolo/fastapi) . txtai is open-source under an Apache 2.0 license.\n",
        "\n",
@@ -1060,7 +1060,7 @@
        "pip install txtai\n",
        "```\n",
        "\n",
-       "Python 3.10+ is supported. Using a Python |virtual environment](https://docs.python.org/3/library/venv.html) is \n",
+       "Python 3.11+ is supported. Using a Python |virtual environment](https://docs.python.org/3/library/venv.html) is \n",
        "recommended.\n",
        "\n",
        "See the detailed |install instructions](https://neuml.github.io/txtai/install) for more information covering \n",
@@ -1169,7 +1169,7 @@
        "|Rust](https://github.com/neuml/txtai.rs) and |Go](https://github.com/neuml/txtai.go) .\n",
        "- 🔋 Batteries included with defaults to get up and running fast\n",
        "- ☁️ Run local or scale out with container orchestration\n",
-       "txtai is built with Python 3.10+, |Hugging Face Transformers](https://github.com/huggingface/transformers) , \n",
+       "txtai is built with Python 3.11+, |Hugging Face Transformers](https://github.com/huggingface/transformers) , \n",
        "|Sentence Transformers](https://github.com/UKPLab/sentence-transformers) and \n",
        "|FastAPI](https://github.com/tiangolo/fastapi) . txtai is open-source under an Apache 2.0 license.\n",
        "\n",
@@ -1333,7 +1333,7 @@
        "pip install txtai\n",
        "```\n",
        "\n",
-       "Python 3.10+ is supported. Using a Python |virtual environment](https://docs.python.org/3/library/venv.html) is \n",
+       "Python 3.11+ is supported. Using a Python |virtual environment](https://docs.python.org/3/library/venv.html) is \n",
        "recommended.\n",
        "\n",
        "See the detailed |install instructions](https://neuml.github.io/txtai/install) for more information covering \n",
@@ -1554,7 +1554,7 @@
        "|Rust](https://github.com/neuml/txtai.rs) and |Go](https://github.com/neuml/txtai.go) .\n",
        "- 🔋 Batteries included with defaults to get up and running fast\n",
        "- ☁️ Run local or scale out with container orchestration\n",
-       "txtai is built with Python 3.10+, |Hugging Face Transformers](https://github.com/huggingface/transformers) , \n",
+       "txtai is built with Python 3.11+, |Hugging Face Transformers](https://github.com/huggingface/transformers) , \n",
        "|Sentence Transformers](https://github.com/UKPLab/sentence-transformers) and \n",
        "|FastAPI](https://github.com/tiangolo/fastapi) . txtai is open-source under an Apache 2.0 license.\n",
        "\n",
@@ -1599,7 +1599,7 @@
        "|Rust](https://github.com/neuml/txtai.rs) and |Go](https://github.com/neuml/txtai.go) .\n",
        "- 🔋 Batteries included with defaults to get up and running fast\n",
        "- ☁️ Run local or scale out with container orchestration\n",
-       "txtai is built with Python 3.10+, |Hugging Face Transformers](https://github.com/huggingface/transformers) , \n",
+       "txtai is built with Python 3.11+, |Hugging Face Transformers](https://github.com/huggingface/transformers) , \n",
        "|Sentence Transformers](https://github.com/UKPLab/sentence-transformers) and \n",
        "|FastAPI](https://github.com/tiangolo/fastapi) . txtai is open-source under an Apache 2.0 license.\n",
        "\n",
@@ -1673,7 +1673,7 @@
        "pip install txtai\n",
        "```\n",
        "\n",
-       "Python 3.10+ is supported. Using a Python |virtual environment](https://docs.python.org/3/library/venv.html) is \n",
+       "Python 3.11+ is supported. Using a Python |virtual environment](https://docs.python.org/3/library/venv.html) is \n",
        "recommended.\n",
        "\n",
        "## Optional dependencies\n",
@@ -1847,7 +1847,7 @@
        "pip install txtai\n",
        "```\n",
        "\n",
-       "Python 3.10+ is supported. Using a Python |virtual environment](https://docs.python.org/3/library/venv.html) is \n",
+       "Python 3.11+ is supported. Using a Python |virtual environment](https://docs.python.org/3/library/venv.html) is \n",
        "recommended.\n",
        "\n",
        "## Optional dependencies\n",
@@ -2043,7 +2043,7 @@
        "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">╭─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮\n",
        "│ Calling tool: 'write' with arguments: {'path': 'txtai_facts.md', 'content': '# txtai - All-in-One AI            │\n",
        "│ Framework\\n\\n## Overview\\n\\ntxtai is an all-in-one open-source AI framework for semantic search, LLM            │\n",
-       "│ orchestration, and language model workflows. It\\'s built with Python 3.10+, Hugging Face Transformers, Sentence │\n",
+       "│ orchestration, and language model workflows. It\\'s built with Python 3.11+, Hugging Face Transformers, Sentence │\n",
        "│ Transformers, and FastAPI, and is licensed under Apache 2.0.\\n\\n## Key Features\\n\\n- **Vector Search**:         │\n",
        "│ Semantic search with SQL, object storage, topic modeling, graph analysis, and multimodal indexing\\n-            │\n",
        "│ **Embeddings**: Create embeddings for text, documents, audio, images, and video\\n- **Language Model             │\n",
@@ -2063,7 +2063,7 @@
        "│ Installation\\n\\nThe easiest way to install is via pip:\\n\\n```bash\\npip install txtai\\n```\\n\\nOptional           │\n",
        "│ dependencies can be installed as extras:\\n\\n```bash\\npip install txtai[all]  # Install all dependencies\\npip    │\n",
        "│ install txtai[api]  # Serve txtai via a web API\\npip install txtai[workflow]  # All workflow tasks\\n```\\n\\n##   │\n",
-       "│ Technology Stack\\n\\n- Built with Python 3.10+\\n- Hugging Face Transformers\\n- Sentence Transformers\\n-          │\n",
+       "│ Technology Stack\\n\\n- Built with Python 3.11+\\n- Hugging Face Transformers\\n- Sentence Transformers\\n-          │\n",
        "│ FastAPI\\n\\n## Supported Models\\n\\n- Embeddings: all-MiniLM-L6-v2\\n- Image Captions: BLIP\\n- Labels (Zero Shot): │\n",
        "│ BART-Large-MNLI\\n- Large Language Model (LLM): gpt-oss-20b\\n- Summarization: DistilBART\\n- Text-to-Speech:      │\n",
        "│ ESPnet JETS\\n- Transcription: Whisper\\n- Translation: OPUS Model Series\\n\\n## Company Behind txtai\\n\\nNeuML is  │\n",
@@ -2081,7 +2081,7 @@
        "╭─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮\n",
        "│ Calling tool: 'write' with arguments: {'path': 'txtai_facts.md', 'content': '# txtai - All-in-One AI            │\n",
        "│ Framework\\n\\n## Overview\\n\\ntxtai is an all-in-one open-source AI framework for semantic search, LLM            │\n",
-       "│ orchestration, and language model workflows. It\\'s built with Python 3.10+, Hugging Face Transformers, Sentence │\n",
+       "│ orchestration, and language model workflows. It\\'s built with Python 3.11+, Hugging Face Transformers, Sentence │\n",
        "│ Transformers, and FastAPI, and is licensed under Apache 2.0.\\n\\n## Key Features\\n\\n- **Vector Search**:         │\n",
        "│ Semantic search with SQL, object storage, topic modeling, graph analysis, and multimodal indexing\\n-            │\n",
        "│ **Embeddings**: Create embeddings for text, documents, audio, images, and video\\n- **Language Model             │\n",
@@ -2101,7 +2101,7 @@
        "│ Installation\\n\\nThe easiest way to install is via pip:\\n\\n```bash\\npip install txtai\\n```\\n\\nOptional           │\n",
        "│ dependencies can be installed as extras:\\n\\n```bash\\npip install txtai[all]  # Install all dependencies\\npip    │\n",
        "│ install txtai[api]  # Serve txtai via a web API\\npip install txtai[workflow]  # All workflow tasks\\n```\\n\\n##   │\n",
-       "│ Technology Stack\\n\\n- Built with Python 3.10+\\n- Hugging Face Transformers\\n- Sentence Transformers\\n-          │\n",
+       "│ Technology Stack\\n\\n- Built with Python 3.11+\\n- Hugging Face Transformers\\n- Sentence Transformers\\n-          │\n",
        "│ FastAPI\\n\\n## Supported Models\\n\\n- Embeddings: all-MiniLM-L6-v2\\n- Image Captions: BLIP\\n- Labels (Zero Shot): │\n",
        "│ BART-Large-MNLI\\n- Large Language Model (LLM): gpt-oss-20b\\n- Summarization: DistilBART\\n- Text-to-Speech:      │\n",
        "│ ESPnet JETS\\n- Transcription: Whisper\\n- Translation: OPUS Model Series\\n\\n## Company Behind txtai\\n\\nNeuML is  │\n",
@@ -2252,7 +2252,7 @@
     "\n",
     "## Overview\n",
     "\n",
-    "txtai is an all-in-one open-source AI framework for semantic search, LLM orchestration, and language model workflows. It's built with Python 3.10+, Hugging Face Transformers, Sentence Transformers, and FastAPI, and is licensed under Apache 2.0.\n",
+    "txtai is an all-in-one open-source AI framework for semantic search, LLM orchestration, and language model workflows. It's built with Python 3.11+, Hugging Face Transformers, Sentence Transformers, and FastAPI, and is licensed under Apache 2.0.\n",
     "\n",
     "## Key Features\n",
     "\n",

setup.py

@@ -62,6 +62,7 @@ extras["ann"] = [
     "hnswlib>=0.5.0",
     "milvus-lite>=3.1.1",
     "pgvector>=0.4.1",
+    "rabitqlib>=0.5.1",
     "scikit-learn>=0.23.1",
     "scipy>=1.4.1",
     "sqlalchemy>=2.0.20",
@@ -142,6 +143,7 @@ extras["vectors"] = [
     "litellm>=1.37.16",
     "llama-cpp-python>=0.2.75",
     "model2vec>=0.3.0",
+    "onnxruntime>=1.11.0",
     "scikit-learn>=0.23.1",
     "scipy>=1.4.1",
     "sentence-transformers>=5.0.0",
@@ -179,7 +181,7 @@ extras["all"] = (
 
 setup(
     name=package,
-    version="9.13.0",
+    version="9.15.0",
     author="NeuML",
     description="All-in-one open-source AI framework for semantic search, LLM orchestration and language model workflows",
     long_description=DESCRIPTION,
@@ -194,7 +196,7 @@ setup(
     packages=find_packages(where="src/python"),
     package_dir={"": "src/python"},
     keywords="search embedding machine-learning nlp",
-    python_requires=">=3.10",
+    python_requires=">=3.11",
     install_requires=install,
     extras_require=extras,
     classifiers=[

src/python/txtai/agent/tool/bash.py

@@ -14,12 +14,13 @@ class BashTool(Tool):
     """
 
     # pylint: disable=W0231
-    def __init__(self, allowed=None):
+    def __init__(self, allowed=None, timeout=30):
         """
         Creates a BashTool.
 
         Args:
             allowed: list of allowed commands to run, has limited set of defaults which are a best effort not a sandbox
+            timeout: maximum number of seconds to wait for a command to complete, defaults to 30, None disables the timeout
         """
 
         # Tool parameters
@@ -31,7 +32,10 @@ class BashTool(Tool):
         self.output_type = "any"
 
         # Default list of allowed commands
-        self.allowed = allowed if allowed else ["cat", "cut", "diff", "grep", "head", "ls", "tail"]
+        self.allowed = allowed if allowed is not None else ["cat", "cut", "diff", "grep", "head", "ls", "tail"]
+
+        # Maximum time to wait for a command, None waits indefinitely
+        self.timeout = timeout
 
         # Validate parameters and initialize tool
         super().__init__()
@@ -41,6 +45,9 @@ class BashTool(Tool):
         """
         Runs a shell command as a subprocess.
 
+        stdin is closed and a timeout is applied. Otherwise an allowed command called with no file arguments
+        (i.e. `cat`) reads from stdin and blocks the agent indefinitely.
+
         Args:
             command: command arguments as a list
 
@@ -50,6 +57,11 @@ class BashTool(Tool):
 
         output = None
         if command and command[0] in self.allowed:
-            output = subprocess.run(command, capture_output=True, text=True, check=False).stdout
+            try:
+                output = subprocess.run(
+                    command, capture_output=True, text=True, encoding="utf-8", check=False, stdin=subprocess.DEVNULL, timeout=self.timeout
+                ).stdout
+            except subprocess.TimeoutExpired:
+                output = f"Command timed out after {self.timeout} seconds"
 
         return output

src/python/txtai/ann/dense/__init__.py

@@ -9,6 +9,7 @@ from .hnsw import HNSW
 from .milvus import Milvus
 from .numpy import NumPy
 from .pgvector import PGVector
+from .rabitq import RabitQ
 from .torch import Torch
 from .turbovec import TurboVec
 from .zvec import Zvec

src/python/txtai/ann/dense/factory.py

@@ -12,6 +12,7 @@ from .hnsw import HNSW
 from .milvus import Milvus
 from .numpy import NumPy
 from .pgvector import PGVector
+from .rabitq import RabitQ
 from .sqlite import SQLite
 from .torch import Torch
 from .turbovec import TurboVec
@@ -35,6 +36,7 @@ class ANNFactory:
             ANN
         """
 
+        # pylint: disable=R0912
         # ANN instance
         ann = None
         backend = config.get("backend", "faiss" if FAISS else "numpy")
@@ -44,16 +46,18 @@ class ANNFactory:
             ann = Annoy(config)
         elif backend == "faiss":
             ann = Faiss(config)
+        elif backend == "ggml":
+            ann = GGML(config)
         elif backend == "hnsw":
             ann = HNSW(config)
         elif backend == "milvus":
             ann = Milvus(config)
-        elif backend == "ggml":
-            ann = GGML(config)
         elif backend == "numpy":
             ann = NumPy(config)
         elif backend == "pgvector":
             ann = PGVector(config)
+        elif backend == "rabitq":
+            ann = RabitQ(config)
         elif backend == "sqlite":
             ann = SQLite(config)
         elif backend == "torch":

src/python/txtai/ann/dense/faiss.py

@@ -142,7 +142,7 @@ class Faiss(ANN):
 
         # Derive quantization. Prefer backend-specific setting. Fallback to root-level parameter.
         quantize = self.setting("quantize", self.config.get("quantize"))
-        quantize = 8 if isinstance(quantize, bool) else quantize
+        quantize = 8 if isinstance(quantize, bool) and quantize else quantize
 
         # Get storage setting
         storage = f"SQ{quantize}" if quantize else "Flat"

src/python/txtai/ann/dense/ggml.py

@@ -155,7 +155,7 @@ class GGMLTensors:
         gguf = ggml.gguf_init_from_file(path.encode("utf-8"), params)
 
         # Load tensors from GGUF file
-        self.loadtensors(context)
+        self.loadtensors(gguf, context)
 
         # Create graph operation
         self.graph, self.output = self.creategraph()
@@ -218,11 +218,12 @@ class GGMLTensors:
             ids: ids to delete
         """
 
-        shape = utils.get_shape(self.data)
+        rows = self.rows()
 
-        # Filter any index greater than size of array
-        ids = [x for x in ids if x < shape[1]]
-        self.deletes.extend(ids)
+        # Set index ids as deleted, ignoring ids outside the array and ids already marked deleted
+        for x in ids:
+            if x < rows and x not in self.deletes:
+                self.deletes.append(x)
 
     def search(self, queries):
         """
@@ -250,8 +251,7 @@ class GGMLTensors:
             ggml.ggml_backend_graph_compute(self.backend, self.graph)
 
             # Get size of embeddings data
-            size = utils.get_shape(self.data)
-            size = size[1] if len(size) > 1 else 1
+            size = self.rows()
 
             # Get and return results
             results = np.zeros((batch.shape[0], size), dtype=np.float32)
@@ -272,7 +272,7 @@ class GGMLTensors:
             count
         """
 
-        return utils.get_shape(self.data)[1] - len(self.deletes) if self.data else 0
+        return self.rows() - len(self.deletes) if self.data else 0
 
     def save(self, path):
         """
@@ -288,9 +288,12 @@ class GGMLTensors:
         # Init and save data tensor
         gguf = ggml.gguf_init_empty()
 
-        # Add the data tensor
-        ggml.ggml_set_name(self.data, b"data")
-        ggml.gguf_add_tensor(gguf, self.data)
+        # Add the data tensor. GGUF can't store a tensor with zero rows, only save the dimensions for an empty index.
+        if self.rows():
+            ggml.ggml_set_name(self.data, b"data")
+            ggml.gguf_add_tensor(gguf, self.data)
+        else:
+            ggml.gguf_set_val_u32(gguf, b"dimensions", utils.get_shape(self.data)[0])
 
         # Optionally create and add the deletes tensor
         if self.deletes:
@@ -413,11 +416,12 @@ class GGMLTensors:
         # Copy embeddings data
         self.copy(data, self.data, tensortype, 0)
 
-    def loadtensors(self, context):
+    def loadtensors(self, gguf, context):
         """
         Loads existing tensors from context.
 
         Args:
+            gguf: gguf context
             context: ggml context
         """
 
@@ -436,6 +440,10 @@ class GGMLTensors:
 
             # Copy tensor data to backend
             ggml.ggml_backend_tensor_set(self.data, ggml.ggml_get_data(data), 0, ggml.ggml_nbytes(data))
+        else:
+            # Empty index, create tensors using the saved dimensions
+            dimensions = ggml.gguf_get_val_u32(gguf, ggml.gguf_find_key(gguf, b"dimensions"))
+            self.createtensors(np.zeros((0, dimensions), dtype=np.float32))
 
         # Load deletes tensor
         data = ggml.ggml_get_tensor(context, b"deletes")
@@ -465,8 +473,7 @@ class GGMLTensors:
         queries = ggml.ggml_new_tensor_2d(context, ggml.GGML_TYPE_F32, data.shape[1], self.querysize)
 
         # Embeddings data with space for both existing and new data
-        shape = utils.get_shape(self.data)
-        merge = ggml.ggml_new_tensor_2d(context, tensortype, data.shape[1], data.shape[0] + shape[1])
+        merge = ggml.ggml_new_tensor_2d(context, tensortype, data.shape[1], data.shape[0] + self.rows())
 
         # Create new buffer
         buffer = ggml.ggml_backend_alloc_ctx_tensors(context, backend)
@@ -512,7 +519,11 @@ class GGMLTensors:
 
         # Read tensor type
         tensortype = self.quantize
-        tensortype = "Q8_0" if isinstance(tensortype, bool) else f"Q{int(tensortype)}_0" if isinstance(tensortype, int) else tensortype
+        tensortype = (
+            "Q8_0"
+            if isinstance(tensortype, bool) and tensortype
+            else f"Q{int(tensortype)}_0" if isinstance(tensortype, int) and not isinstance(tensortype, bool) else tensortype
+        )
         tensortype = tensortype.upper() if tensortype else "F32"
 
         # Validate tensor type
@@ -572,3 +583,13 @@ class GGMLTensors:
         """
 
         return [queries[x : x + self.querysize] for x in range(0, len(queries), self.querysize)]
+
+    def rows(self):
+        """
+        Number of rows in the data tensor. Tensors with a single row or no rows report a 1-D shape, so the shape can't be used.
+
+        Returns:
+            row count
+        """
+
+        return ggml.ggml_nrows(self.data)

src/python/txtai/ann/dense/hnsw.py

@@ -35,6 +35,9 @@ class HNSW(ANN):
         self.backend = Index(dim=self.config["dimensions"], space=self.config["metric"])
         self.backend.load_index(path)
 
+        # Default the delete counter for configurations that don't have it
+        self.config["deletes"] = self.config.get("deletes", 0)
+
     def index(self, embeddings):
         # Inner product is equal to cosine similarity on normalized vectors
         self.config["metric"] = "ip"

src/python/txtai/ann/dense/milvus.py

@@ -83,8 +83,9 @@ class Milvus(ANN):
         self.metadata()
 
     def delete(self, ids):
-        if ids:
-            self.collection.delete(ids)
+        # Check length, ids can be a NumPy array
+        if len(ids):
+            self.collection.delete([int(x) for x in ids])
 
     def search(self, queries, limit):
         matches = self.collection.search(queries.tolist(), top_k=limit, metric_type="IP", anns_field="embedding")
@@ -138,3 +139,6 @@ class Milvus(ANN):
 
         self.backend = milvus_lite.MilvusLite(self.path)
         self.collection = self.backend.get_collection("txtai")
+
+        # Reopened collections with an index start released, load to enable search
+        self.collection.load()

src/python/txtai/ann/dense/numpy.py

@@ -60,11 +60,11 @@ class NumPy(ANN):
         self.metadata()
 
     def delete(self, ids):
-        # Filter any index greater than size of array
-        ids = [x for x in ids if x < self.backend.shape[0]]
+        # Ignore IDs outside the array; negative IDs must not wrap to live rows.
+        ids = [x for x in ids if 0 <= x < self.backend.shape[0]]
 
-        # Clear specified ids
-        self.backend[ids] = self.tensor(self.zeros((len(ids), self.backend.shape[1])))
+        # Clear specified ids, zeros must match the array data type (e.g. uint8 for quantized data)
+        self.backend[ids] = self.tensor(self.zeros((len(ids), self.backend.shape[1]), dtype=self.backend.dtype))
 
     def search(self, queries, limit):
         if self.qbits:

src/python/txtai/ann/dense/pgvector.py

@@ -24,7 +24,6 @@ from ...util import Library
 np = Library().numpy()
 
 
-# pylint: disable=R0904
 class PGVector(ANN):
     """
     Builds an ANN index backed by a Postgres database.
@@ -72,7 +71,8 @@ class PGVector(ANN):
         self.metadata()
 
     def delete(self, ids):
-        self.database.execute(delete(self.table).where(self.table.c["indexid"].in_(ids)))
+        # Bind ids as int, NumPy integers don't match the indexid column
+        self.database.execute(delete(self.table).where(self.table.c["indexid"].in_([int(x) for x in ids])))
 
     def search(self, queries, limit):
         results = []

src/python/txtai/ann/dense/rabitq.py

@@ -0,0 +1,176 @@
+"""
+RabitQ module
+"""
+
+import math
+
+# Conditional import
+try:
+    from rabitqlib import HnswIndex, IvfIndex
+    from sklearn.cluster import KMeans
+
+    RABITQ = True
+except ImportError:
+    RABITQ = False
+
+from ..base import ANN
+
+# Core library imports
+from ...util import Library
+
+library = Library()
+np = library.numpy()
+
+# Sentinel for unfilled top-k slots in native search results
+SENTINEL = 0xFFFFFFFF
+
+
+class RabitQ(ANN):
+    """
+    Builds an ANN index using the rabitqlib library (RaBitQ quantization).
+
+    Only the native index is stored, as a single file. No vectors are kept outside of it. ivf indexes support append
+    and delete natively. The native id of a row is its index id. hnsw indexes can't be extended or have rows removed,
+    so append and delete are not supported in that mode.
+    """
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        if not RABITQ:
+            raise ImportError('rabitqlib is not available - install "ann" extra to enable')
+
+        # Number of rows removed from the native index (ivf mode only)
+        self.config["deletes"] = self.config.get("deletes", 0)
+
+    def load(self, path):
+        self.close()
+
+        # Invalid modes raise ValueError before native code runs
+        self.backend = IvfIndex.load(path) if self.mode() == "ivf" else HnswIndex.load(path)
+
+    def index(self, embeddings):
+        self.close()
+
+        # Mode is validated before any state is stored
+        mode = self.mode()
+
+        vectors = np.ascontiguousarray(embeddings, dtype=np.float32)
+        rows = vectors.shape[0]
+
+        # Quantization bits per dimension, validated by rabitqlib
+        nbits = self.setting("nbits", 1)
+
+        # Default clusters to 4 * sqrt(N), bounded to [1, N] so KMeans always fits
+        clusters = self.setting("clusters", None)
+        clusters = max(1, min(clusters if clusters else round(4 * math.sqrt(rows)), rows))
+
+        cluster = KMeans(n_clusters=clusters, random_state=0, n_init=1).fit(vectors)
+        centroids = np.ascontiguousarray(cluster.cluster_centers_, dtype=np.float32)
+        clusterids = np.ascontiguousarray(cluster.labels_, dtype=np.uint32)
+
+        dim = self.config["dimensions"]
+        if mode == "ivf":
+            backend = IvfIndex(dim=dim, max_elements=rows, num_clusters=clusters, nbits=nbits, metric="ip")
+            settings = {"mode": mode, "nbits": nbits, "clusters": clusters, "nprobe": self.setting("nprobe", self.nprobe(clusters))}
+        else:
+            backend = HnswIndex(
+                dim=dim,
+                max_elements=rows,
+                M=self.setting("m", 16),
+                ef_construction=self.setting("efconstruction", 200),
+                nbits=nbits,
+                metric="ip",
+                random_seed=self.setting("randomseed", 100),
+            )
+            settings = {
+                "mode": mode,
+                "nbits": nbits,
+                "m": self.setting("m", 16),
+                "efconstruction": self.setting("efconstruction", 200),
+                "efsearch": self.setting("efsearch", None),
+                "randomseed": self.setting("randomseed", 100),
+            }
+
+        backend.build(vectors, centroids, clusterids)
+        self.backend = backend
+
+        # Add id offset and index build metadata
+        self.config["offset"] = rows
+        self.config["deletes"] = 0
+        self.metadata(settings)
+
+    def append(self, embeddings):
+        if self.mode() != "ivf":
+            super().append(embeddings)
+
+        # Native ids continue from the last row ever added
+        embeddings = np.ascontiguousarray(embeddings, dtype=np.float32)
+        self.backend.add(embeddings)
+
+        # Update id offset and index metadata
+        self.config["offset"] = self.backend.max_elements
+        self.metadata()
+
+    def delete(self, ids):
+        if self.mode() != "ivf":
+            super().delete(ids)
+
+        # Native remove hides rows from search. Unknown ids are silently ignored.
+        ids = [uid for uid in ids if 0 <= uid < self.backend.max_elements]
+        if ids:
+            self.config["deletes"] += self.backend.remove(np.array(ids, dtype=np.int64))
+
+    def search(self, queries, limit):
+        if not self.count():
+            return [[] for _ in queries]
+
+        queries = np.ascontiguousarray(queries, dtype=np.float32)
+
+        # The native index errors when k exceeds its row count
+        k = min(limit, self.count())
+
+        if self.mode() == "ivf":
+            nprobe = self.setting("nprobe", self.nprobe(self.backend.num_clusters))
+            ids, distances = self.backend.search(queries, k, nprobe)
+        else:
+            ids, distances = self.backend.search(queries, k, self.setting("efsearch", None) or 0)
+
+        # Map results to [(id, score)], dropping unfilled sentinel slots
+        return [
+            [(int(uid), float(1.0 - dist)) for uid, dist in zip(uids, dists) if uid != SENTINEL]
+            for uids, dists in zip(ids.tolist(), distances.tolist())
+        ]
+
+    def count(self):
+        return self.backend.max_elements - self.config["deletes"] if self.backend is not None else 0
+
+    def save(self, path):
+        self.backend.save(path)
+
+    def mode(self):
+        """
+        Returns the configured index mode.
+
+        Returns:
+            index mode, only "ivf" or "hnsw" are accepted
+        """
+
+        mode = self.setting("mode", "ivf")
+        if mode not in ("ivf", "hnsw"):
+            raise ValueError(f'Invalid RabitQ mode: {mode}. Supported modes are "ivf" and "hnsw".')
+
+        return mode
+
+    def nprobe(self, clusters):
+        """
+        Returns the default number of clusters to probe at search time.
+
+        Args:
+            clusters: number of IVF clusters
+
+        Returns:
+            default nprobe
+        """
+
+        return max(1, round(clusters / 16))

src/python/txtai/ann/dense/sqlite.py

@@ -32,9 +32,11 @@ class SQLite(ANN):
 
         # Quantization setting
         self.quantize = self.setting("quantize")
-        self.quantize = 8 if isinstance(self.quantize, bool) else int(self.quantize) if self.quantize else None
+        self.quantize = 8 if isinstance(self.quantize, bool) and self.quantize else int(self.quantize) if self.quantize else None
 
     def load(self, path):
+        if self.path != path:
+            self.close()
         self.path = path
 
     def index(self, embeddings):
@@ -55,7 +57,8 @@ class SQLite(ANN):
         self.metadata()
 
     def delete(self, ids):
-        self.database().executemany(self.deletesql(), [(x,) for x in ids])
+        # Bind ids as int, sqlite3 binds NumPy integers as blobs that match the wrong rows
+        self.database().executemany(self.deletesql(), [(int(x),) for x in ids])
 
     def search(self, queries, limit):
         results = []
@@ -73,6 +76,9 @@ class SQLite(ANN):
         return self.cursor.fetchone()[0]
 
     def save(self, path):
+        # Loading defers opening the connection until the first database operation.
+        self.database()
+
         # Temporary database
         if not self.path:
             # Save temporary database
@@ -193,10 +199,13 @@ class SQLite(ANN):
             connection.execute(self.tablesql())
 
             # The backup call will hang if there are uncommitted changes, need to copy over
-            # with iterdump (which is much slower)
-            for sql in self.connection.iterdump():
-                if self.tosql('insert into "{table}"') in sql.lower():
-                    connection.execute(sql)
+            # the rows (which is much slower). Quantized vectors must be inserted with their type.
+            embedding = "vec_bit(?)" if self.quantize == 1 else "vec_int8(?)" if self.quantize == 8 else "?"
+            rows = self.connection.execute(self.tosql("SELECT indexid, embedding FROM {table}"))
+            connection.executemany(self.tosql(f"INSERT INTO {{table}}(indexid, embedding) VALUES (?, {embedding})"), rows)
+
+            # Commit copied rows, otherwise they are rolled back when the connection is closed
+            connection.commit()
         else:
             # Database is up to date, can do a more efficient copy with SQLite C API
             self.connection.backup(connection)
@@ -254,7 +263,12 @@ class SQLite(ANN):
             SELECT
         """
 
-        return self.tosql(("SELECT indexid, 1 - distance FROM {table} " f"WHERE embedding MATCH {self.embeddingsql()} AND k = ? ORDER BY distance"))
+        # BIT distance counts differing bits; normalize it before converting to similarity.
+        # SQLite binary dimensions are the number of bits, not the packed byte count.
+        distance = f"distance / {float(self.config['dimensions'])}" if self.quantize == 1 else "distance"
+        return self.tosql(
+            f"SELECT indexid, 1 - ({distance}) FROM {{table}} " f"WHERE embedding MATCH {self.embeddingsql()} AND k = ? ORDER BY distance"
+        )
 
     def countsql(self):
         """

src/python/txtai/ann/dense/zvec.py

@@ -49,6 +49,7 @@ class Zvec(ANN):
         self.close()
 
         # Lookup index settings
+        efconstruction = self.setting("efconstruction", 200)
         m = self.setting("m", 50)
 
         # Create collection
@@ -58,9 +59,9 @@ class Zvec(ANN):
             name="txtai",
             vectors=zvec.VectorSchema(
                 "embedding",
-                zvec.DataType.VECTOR_FP32,
+                self.datatype(),
                 self.config["dimensions"],
-                index_param=zvec.HnswIndexParam(metric_type=zvec.MetricType.IP, m=m),
+                index_param=zvec.HnswIndexParam(metric_type=zvec.MetricType.IP, m=m, ef_construction=efconstruction),
             ),
         )
         self.backend = zvec.create_and_open(path=self.path, schema=schema)
@@ -68,26 +69,36 @@ class Zvec(ANN):
         # Add items - position in embeddings is used as the id
         self.insert(embeddings, 0)
 
+        # Build HNSW index
+        self.backend.optimize()
+
         # Add id offset and index build metadata
         self.config["offset"] = embeddings.shape[0]
-        self.metadata({"m": m, "zvec": zvec.__version__})
+        self.metadata({"efconstruction": efconstruction, "m": m, "zvec": zvec.__version__})
 
     def append(self, embeddings):
         self.insert(embeddings, self.config["offset"])
 
+        # Add new items to HNSW index
+        self.backend.optimize()
+
         # Update id offset and index metadata
         self.config["offset"] += embeddings.shape[0]
         self.metadata()
 
     def delete(self, ids):
-        if ids:
+        # Check length, ids can be a NumPy array
+        if len(ids):
             self.backend.delete([str(uid) for uid in ids])
 
     def search(self, queries, limit):
+        # Lookup search settings
+        param = zvec.HnswQueryParam(ef=self.setting("efsearch", 300))
+
         results = []
         for query in queries:
             matches = self.backend.query(
-                zvec.Query(field_name="embedding", vector=query.tolist()),
+                zvec.Query(field_name="embedding", vector=self.prepare(query), param=param),
                 topk=limit,
             )
             results.append([(int(match.id), float(match.score)) for match in matches])
@@ -128,5 +139,28 @@ class Zvec(ANN):
             for start in range(0, embeddings.shape[0], 1024):
                 batch = embeddings[start : start + 1024]
                 self.backend.insert(
-                    [zvec.Doc(id=str(offset + start + uid), vectors={"embedding": embedding.tolist()}) for uid, embedding in enumerate(batch)]
+                    [zvec.Doc(id=str(offset + start + uid), vectors={"embedding": self.prepare(embedding)}) for uid, embedding in enumerate(batch)]
                 )
+
+    def datatype(self):
+        """
+        Gets the vector data type.
+
+        Returns:
+            zvec vector data type
+        """
+
+        return zvec.DataType.VECTOR_FP32
+
+    def prepare(self, data):
+        """
+        Prepares data for the vector field.
+
+        Args:
+            data: input data
+
+        Returns:
+            data formatted for zvec
+        """
+
+        return data.tolist()

src/python/txtai/ann/sparse/__init__.py

@@ -5,3 +5,4 @@ Sparse ANN imports
 from .factory import SparseANNFactory
 from .ivfsparse import IVFSparse
 from .pgsparse import PGSparse
+from .zvecsparse import ZvecSparse

src/python/txtai/ann/sparse/factory.py

@@ -7,6 +7,7 @@ from ...util import Resolver
 from ..base import ANN
 from .ivfsparse import IVFSparse
 from .pgsparse import PGSparse
+from .zvecsparse import ZvecSparse
 
 
 class SparseANNFactory:
@@ -35,6 +36,8 @@ class SparseANNFactory:
             ann = IVFSparse(config)
         elif backend == "pgsparse":
             ann = PGSparse(config)
+        elif backend == "zvecsparse":
+            ann = ZvecSparse(config)
         else:
             ann = SparseANNFactory.resolve(backend, config)
 

src/python/txtai/ann/sparse/ivfsparse.py

@@ -115,8 +115,12 @@ class IVFSparse(ANN):
         self.metadata()
 
     def delete(self, ids):
-        # Set index ids as deleted
-        self.deletes.extend(ids)
+        # Set index ids as deleted, ignoring ids that were never indexed and ids already marked deleted
+        # Ids are stored as int, the deletes list is serialized with msgpack
+        size = self.size()
+        for x in ids:
+            if x < size and x not in self.deletes:
+                self.deletes.append(int(x))
 
     def search(self, queries, limit):
         results = []
@@ -204,10 +208,10 @@ class IVFSparse(ANN):
             cluster centroids
         """
 
-        # Select top n most important features that contribute to L2 vector norm
-        indices = np.argsort(-norm(train, axis=0))[: self.setting("nfeatures", 25)]
-        data = train[:, indices]
-        data = train
+        # Select top n most important features that contribute to L2 vector norm.
+        # nfeatures is opt-in: the full feature space is used when it isn't set.
+        nfeatures = self.setting("nfeatures")
+        data = train[:, np.argsort(-norm(train, axis=0))[:nfeatures]] if nfeatures else train
 
         # Cluster data using k-means
         kmeans = MiniBatchKMeans(n_clusters=clusters, random_state=0, n_init=5).fit(data)

src/python/txtai/ann/sparse/zvecsparse.py

@@ -0,0 +1,35 @@
+"""
+ZvecSparse module
+"""
+
+# Conditional import
+try:
+    import zvec
+except ImportError:
+    pass
+
+from ..dense import Zvec
+
+
+class ZvecSparse(Zvec):
+    """
+    Builds a Sparse ANN index using the zvec library.
+    """
+
+    def search(self, queries, limit):
+        # Lookup search settings
+        param = zvec.HnswQueryParam(ef=self.setting("efsearch", 300))
+
+        results = []
+        for query in queries:
+            # zvec requires at least one non-zero value per query
+            matches = self.backend.query(zvec.Query(field_name="embedding", vector=self.prepare(query), param=param), topk=limit) if query.nnz else []
+            results.append([(int(match.id), float(match.score)) for match in matches])
+
+        return results
+
+    def datatype(self):
+        return zvec.DataType.SPARSE_VECTOR_FP32
+
+    def prepare(self, data):
+        return {int(index): float(value) for index, value in zip(data.indices, data.data)}

src/python/txtai/api/base.py

@@ -32,7 +32,7 @@ class API(Application):
         weights = self.weights(request.query_params.get("weights") if request and hasattr(request, "query_params") else weights)
         index = request.query_params.get("index") if request and hasattr(request, "query_params") else index
         parameters = request.query_params.get("parameters") if request and hasattr(request, "query_params") else parameters
-        graph = request.query_params.get("graph") if request and hasattr(request, "query_params") else graph
+        graph = self.graph(request.query_params.get("graph") if request and hasattr(request, "query_params") else graph)
 
         # Decode parameters
         parameters = json.loads(parameters) if parameters and isinstance(parameters, str) else parameters
@@ -157,3 +157,17 @@ class API(Application):
         """
 
         return float(weights) if weights else weights
+
+    def graph(self, graph):
+        """
+        Parses the graph parameter from the request.
+
+        Args:
+            graph: graph parameter
+
+        Returns:
+            graph as a bool
+        """
+
+        # Query parameters are strings - accept the same values the batchsearch endpoint coerces
+        return graph.strip().lower() in ("1", "true", "yes", "on") if isinstance(graph, str) else bool(graph)

src/python/txtai/api/cluster.py

@@ -56,17 +56,19 @@ class Cluster:
         """
 
         # Build URL
-        action = f"search?query={urllib.parse.quote_plus(query)}"
+        params = {"query": query}
         if limit:
-            action += f"&limit={limit}"
-        if weights:
-            action += f"&weights={weights}"
+            params["limit"] = limit
+        if weights is not None:
+            params["weights"] = weights
         if index:
-            action += f"&index={index}"
+            params["index"] = index
         if parameters:
-            action += f"&parameters={json.dumps(parameters) if isinstance(parameters, dict) else parameters}"
+            params["parameters"] = json.dumps(parameters) if isinstance(parameters, dict) else parameters
         if graph is not None:
-            action += f"&graph={graph}"
+            params["graph"] = graph
+
+        action = f"search?{urllib.parse.urlencode(params)}"
 
         # Run query and flatten results into single results list
         results = []
@@ -100,7 +102,7 @@ class Cluster:
         params = {"queries": queries}
         if limit:
             params["limit"] = limit
-        if weights:
+        if weights is not None:
             params["weights"] = weights
         if index:
             params["index"] = index

src/python/txtai/api/routers/openai.py

@@ -165,9 +165,9 @@ class ChatResponse:
         return {
             "id": str(uuid.uuid4()),
             "object": "chat.completion",
-            "created": int(time.time() * 1000),
+            "created": int(time.time()),
             "model": model,
-            "choices": [{"id": 0, "message": {"role": "assistant", "content": result}, "finish_reason": "stop"}],
+            "choices": [{"index": 0, "message": {"role": "assistant", "content": result}, "finish_reason": "stop"}],
         }
 
 
@@ -177,15 +177,38 @@ class StreamingChatResponse:
     """
 
     def __call__(self, model, result):
+        # All chunks of a response share the same id and creation time
+        uid, created = str(uuid.uuid4()), int(time.time())
+
         for chunk in result:
-            yield "data: " + json.dumps(
-                {
-                    "id": str(uuid.uuid4()),
-                    "object": "chat.completion.chunk",
-                    "created": int(time.time() * 1000),
-                    "model": model,
-                    "choices": [{"id": 0, "delta": {"content": chunk}}],
-                }
-            ) + "\n\n"
+            yield self.chunk(uid, created, model, {"content": chunk}, None)
+
+        # Final chunk signals the end of the response
+        yield self.chunk(uid, created, model, {}, "stop")
 
         yield "data: [DONE]\n\n"
+
+    def chunk(self, uid, created, model, delta, reason):
+        """
+        Builds a server-sent event for a single chat completion chunk.
+
+        Args:
+            uid: response id
+            created: response creation time in seconds
+            model: model name
+            delta: chunk content
+            reason: finish reason, None until the final chunk
+
+        Returns:
+            server-sent event
+        """
+
+        data = {
+            "id": uid,
+            "object": "chat.completion.chunk",
+            "created": created,
+            "model": model,
+            "choices": [{"index": 0, "delta": delta, "finish_reason": reason}],
+        }
+
+        return "data: " + json.dumps(data) + "\n\n"

src/python/txtai/app/base.py

@@ -18,7 +18,6 @@ from ..util import Library
 yaml = Library().yaml()
 
 
-# pylint: disable=R0904
 class Application:
     """
     Builds YAML-configured txtai applications.
@@ -259,7 +258,7 @@ class Application:
         """
 
         # Check for task shorthand syntax
-        task = {"action": task} if isinstance(task, (str, list)) else task
+        task = {"action": task} if isinstance(task, (str, list)) else task.copy()
 
         if "action" in task:
             action = task["action"]

src/python/txtai/data/base.py

@@ -38,7 +38,10 @@ class Data:
             (train, validation)
         """
 
-        return (self.prepare(train, self.process, workers, batch), self.prepare(validation, self.process, workers, batch) if validation else None)
+        return (
+            self.prepare(train, self.process, workers, batch),
+            self.prepare(validation, self.process, workers, batch) if validation is not None else None,
+        )
 
     def prepare(self, data, fn, workers, batch):
         """
@@ -54,7 +57,8 @@ class Data:
             tokens
         """
 
-        if hasattr(data, "map"):
+        # Check for column_names, pandas DataFrames also have a map method
+        if hasattr(data, "column_names"):
             # Hugging Face dataset
             tokens = data.map(fn, batched=True, batch_size=batch, num_proc=workers, remove_columns=data.column_names)
         else:
@@ -93,11 +97,11 @@ class Data:
         column = self.columns[-1]
 
         # Return length of labels if it's an array
-        length = self.length(data[column][0] if hasattr(data, "columns") else data[0][column])
+        length = self.length(next(iter(data[column])) if hasattr(data, "columns") else data[0][column])
         if length:
             return length
 
-        if hasattr(data, "map"):
+        if hasattr(data, "column_names"):
             # Hugging Face dataset
             labels = sorted(data.unique(self.columns[-1]))
         elif hasattr(data, "columns"):

src/python/txtai/data/questions.py

@@ -72,14 +72,20 @@ class Questions(Data):
                 while sequences[end] != (1 if self.rpad else 0):
                     end -= 1
 
-                # Map start character and end character to matching token index
-                while start < len(offset) and offset[start][0] <= startchar:
-                    start += 1
-                tokenized["start_positions"].append(start - 1)
+                # Label with CLS token if the answer isn't fully in this span. Long contexts are split into
+                # overlapping spans with stride and most spans don't contain the answer.
+                if not (offset[start][0] <= startchar and offset[end][1] >= endchar):
+                    tokenized["start_positions"].append(clstoken)
+                    tokenized["end_positions"].append(clstoken)
+                else:
+                    # Map start character and end character to matching token index
+                    while start < len(offset) and offset[start][0] <= startchar:
+                        start += 1
+                    tokenized["start_positions"].append(start - 1)
 
-                while offset[end][1] >= endchar:
-                    end -= 1
-                tokenized["end_positions"].append(end + 1)
+                    while offset[end][1] >= endchar:
+                        end -= 1
+                    tokenized["end_positions"].append(end + 1)
 
         return tokenized
 

src/python/txtai/data/texts.py

@@ -104,7 +104,7 @@ class Texts(Data):
             nextlength = len(rows[x + 1]) if x < len(rows) - 1 else 0
 
             # New row
-            if (length + nextlength) >= self.maxlength:
+            if (length + nextlength) > self.maxlength:
                 for column in columns:
                     results[column].append(list(chain(*inputs[column][index : x + 1])))
 

src/python/txtai/database/duckdb.py

@@ -18,6 +18,7 @@ except ImportError:
 
 from .embedded import Embedded
 from .schema import Statement
+from .sql import SQLError
 
 
 class DuckDB(Embedded):
@@ -29,6 +30,11 @@ class DuckDB(Embedded):
     DELETE_DOCUMENT = "DELETE FROM documents WHERE id = ?"
     DELETE_OBJECT = "DELETE FROM objects WHERE id = ?"
 
+    # Added to errors that abort the transaction
+    ABORTED = (
+        "This error aborted the DuckDB transaction. Unsaved changes are lost and later operations fail until the index is loaded or built again."
+    )
+
     def __init__(self, config):
         super().__init__(config)
 
@@ -36,8 +42,15 @@ class DuckDB(Embedded):
             raise ImportError('DuckDB is not available - install "database" extra to enable')
 
     def execute(self, function, *args):
-        # Call parent method with DuckDB compatible arguments
-        return super().execute(function, *self.formatargs(args))
+        try:
+            # Call parent method with DuckDB compatible arguments
+            return super().execute(function, *self.formatargs(args))
+        except SQLError as e:
+            # Runtime errors abort the transaction. DuckDB has no savepoints, so it can't be recovered without a rollback.
+            if self.aborted():
+                raise SQLError(f"{e}\n\n{DuckDB.ABORTED}") from None
+
+            raise
 
     def insertdocument(self, uid, data, tags, entry):
         # Delete existing document
@@ -123,6 +136,20 @@ class DuckDB(Embedded):
 
         return connection
 
+    def aborted(self):
+        """
+        Checks if the current transaction is aborted.
+
+        Returns:
+            True if the current transaction is aborted
+        """
+
+        try:
+            self.connection.execute("SELECT 1")
+            return False
+        except duckdb.TransactionException:
+            return True
+
     def formatargs(self, args):
         """
         DuckDB doesn't support named parameters. This method replaces named parameters with question marks
@@ -139,19 +166,16 @@ class DuckDB(Embedded):
             # Unpack query args
             query, parameters = args
 
-            # Iterate over parameters
-            #   - Replace named parameters with ?'s
-            #   - Build list of value with position indexes
+            # Match all named parameters in a single pass, so values are listed in the order they appear in the query
+            #   - Match on a word boundary, bind parameters can be followed by any non-word character, i.e. `in (:x)`
             params = []
-            for key, value in parameters.items():
-                pattern = rf"\:{key}(?=\s|$)"
-                match = re.search(pattern, query)
-                if match:
-                    query = re.sub(pattern, "?", query, count=1)
-                    params.append((match.start(), value))
+            if parameters:
+                pattern = r"\:(" + "|".join(re.escape(key) for key in parameters) + r")\b"
+                params = [parameters[match.group(1)] for match in re.finditer(pattern, query)]
+                query = re.sub(pattern, "?", query)
 
             # Repack query and parameter list
-            args = (query, [value for _, value in sorted(params, key=lambda x: x[0])])
+            args = (query, params)
 
         return args
 

src/python/txtai/database/embedded.py

@@ -31,6 +31,9 @@ class Embedded(RDBMS):
         self.path = path
 
     def save(self, path):
+        # Initialize connection if not open
+        self.initialize()
+
         # Temporary database
         if not self.path:
             # Save temporary database

src/python/txtai/database/rdbms.py

@@ -9,7 +9,6 @@ from .base import Database
 from .schema import Statement
 
 
-# pylint: disable=R0904
 class RDBMS(Database):
     """
     Base relational database class. A relational database uses SQL to insert, update, delete and select from a
@@ -186,8 +185,9 @@ class RDBMS(Database):
         if indexids:
             select = f"{self.resolve('indexid')}, {self.resolve('score')}"
 
-        # Use JOIN when documents table is used to utilize indexes, default to LEFT JOIN
-        join = "JOIN" if any(x and self.jsonprefix() in x for x in [where, groupby, orderby]) else "LEFT JOIN"
+        # Use JOIN when a JSON field is filtered, since that filter already excludes rows without a
+        # documents row. GROUP BY/ORDER BY don't filter, so they keep the LEFT JOIN.
+        join = "JOIN" if where and self.jsonprefix() in where else "LEFT JOIN"
 
         # Build query text
         query = Statement.TABLE_CLAUSE % (select, join)

src/python/txtai/database/sql/aggregate.py

@@ -6,6 +6,7 @@ import itertools
 import operator
 
 from .base import SQL
+from .error import SQLError
 
 
 class Aggregate(SQL):
@@ -62,18 +63,41 @@ class Aggregate(SQL):
 
         aggregates = {}
         for column in columns:
-            column = column.lower()
-            if column.startswith(("count(", "sum(", "total(")):
+            # Match function names case-insensitively but keep the column name as returned
+            name = column.lower()
+            if name.startswith(("count(", "sum(", "total(")):
                 aggregates[column] = sum
-            elif column.startswith("max("):
+            elif name.startswith("max("):
                 aggregates[column] = max
-            elif column.startswith("min("):
+            elif name.startswith("min("):
                 aggregates[column] = min
-            elif column.startswith("avg("):
-                aggregates[column] = lambda x: sum(x) / len(x)
+            elif name.startswith("avg("):
+                aggregates[column] = self.avg
 
         return aggregates
 
+    def avg(self, values, counts):
+        """
+        Combines per-shard/per-group average values into a single average, weighted by a
+        matching count(*) column when one was selected. Without it, there is no way to
+        recombine the averages correctly, so this raises instead of returning a wrong number.
+
+        Args:
+            values: list of per-shard/per-group average values
+            counts: list of matching count(*) values, or None
+
+        Returns:
+            combined average
+        """
+
+        if len(values) == 1:
+            return values[0]
+
+        if not counts or sum(counts) == 0:
+            raise SQLError("avg() requires a count(*) column to combine results from multiple shards")
+
+        return sum(value * count for value, count in zip(values, counts)) / sum(counts)
+
     def aggregate(self, query, results, columns, aggcolumns):
         """
         Merges aggregate columns in results.
@@ -94,6 +118,9 @@ class Aggregate(SQL):
         else:
             results = [results]
 
+        # Column providing the row count each group's values were computed over, used to weight avg() columns
+        countcolumn = next((column for column in columns if column.lower() == "count(*)"), None)
+
         # Compute column values
         rows = []
         for result in results:
@@ -101,9 +128,10 @@ class Aggregate(SQL):
             row = {}
             for column in columns:
                 if column in aggcolumns:
-                    # Calculate aggregate value
-                    function = aggcolumns[column]
-                    row[column] = function([r[column] for r in result])
+                    # Calculate aggregate value, skip NULLs from shards with no matching rows
+                    matches = [r for r in result if r[column] is not None]
+                    values, counts = [r[column] for r in matches], [r[countcolumn] for r in matches] if countcolumn else None
+                    row[column] = (self.avg(values, counts) if column.lower().startswith("avg(") else aggcolumns[column](values)) if values else None
                 else:
                     # Non aggregate column value repeat, use first value
                     row[column] = result[0][column]
@@ -155,9 +183,9 @@ class Aggregate(SQL):
                 clause = clause.rsplit(" ")[0]
                 reverse = True
 
-            # Order by columns must be in select clause
+            # Order by columns must be in select clause. NULLs sort first like in SQLite (last when descending).
             if clause in query["select"]:
-                results = sorted(results, key=operator.itemgetter(clause), reverse=reverse)
+                results = sorted(results, key=lambda row, column=clause: (row[column] is not None, row[column]), reverse=reverse)
 
         return results
 

src/python/txtai/database/sql/base.py

@@ -5,6 +5,7 @@ SQL module
 from io import StringIO
 from shlex import shlex
 
+from .error import SQLError
 from .expression import Expression
 
 
@@ -42,8 +43,11 @@ class SQL:
 
         clauses = None
         if self.issql(query):
-            # Ignore multiple statements
-            query = query.split(";")[0]
+            # Multiple statements are not supported
+            components = [x for x in query.split(";") if x.strip()]
+            if len(components) > 1:
+                raise SQLError("Multiple SQL statements are not supported. Pass values that contain a semicolon as a bind parameter.")
+            query = components[0]
 
             # Tokenize query
             tokens, positions = self.tokenize(query)
@@ -66,6 +70,10 @@ class SQL:
             if similar:
                 clauses["similar"] = similar
 
+            # Add select column aliases, if any. Only a query with an offset uses them.
+            if aliases and clauses["offset"]:
+                clauses["aliases"] = aliases
+
         # Return clauses, default to full query if this is not a SQL query
         return clauses if clauses else {"similar": [[query]]}
 

src/python/txtai/database/sql/token.py

@@ -26,6 +26,9 @@ class Token:
     # Default list of sort order operators
     SORT_ORDER = ["asc", "desc"]
 
+    # Default list of boolean literals
+    BOOLEANS = ["true", "false"]
+
     @staticmethod
     def get(tokens, x):
         """
@@ -197,8 +200,8 @@ class Token:
             True if this token is a literal, False otherwise
         """
 
-        # Literals are wrapped in quotes, parens, wildcards or numeric.
-        return token and (token.startswith(("'", '"', ",", "(", ")", "*")) or token.replace(".", "", 1).isdigit())
+        # Literals are wrapped in quotes, parens, wildcards, numeric or boolean.
+        return token and (token.startswith(("'", '"', ",", "(", ")", "*")) or token.replace(".", "", 1).isdigit() or token.lower() in Token.BOOLEANS)
 
     @staticmethod
     def islogicseparator(token):

src/python/txtai/embeddings/base.py

@@ -18,7 +18,6 @@ from .index import Action, Configuration, Functions, Indexes, IndexIds, Reducer,
 from .search import Explain, Ids, Query, Search, Terms
 
 
-# pylint: disable=C0302,R0904
 class Embeddings:
     """
     Embeddings databases are the engine that delivers semantic search. Data is transformed into embeddings vectors where similar concepts
@@ -136,7 +135,7 @@ class Embeddings:
                 self.ann.index(embeddings)
 
             # Save indexids-ids mapping for indexes with no database, except when this is a reindex
-            if ids and not reindex and not self.database:
+            if ids is not None and not reindex and not self.database:
                 self.ids = self.createids(ids)
 
         # Index scoring, if necessary
@@ -549,6 +548,9 @@ class Embeddings:
         if cloud:
             path = cloud.load(path)
 
+        # Reset archive since this replaces the current index
+        self.archive = None
+
         # Check if this is an archive file and extract
         path, apath = self.checkarchive(path)
         if apath:
@@ -931,7 +933,8 @@ class Embeddings:
         """
 
         # Create archive instance, if necessary
-        self.archive = ArchiveFactory.create()
+        if not self.archive:
+            self.archive = ArchiveFactory.create()
 
         # Check if path is an archive file
         if self.archive.isarchive(path):

src/python/txtai/embeddings/index/transform.py

@@ -147,10 +147,10 @@ class Transform:
                 if not self.indexing and not document[1].get(self.text):
                     document[1][self.text] = str(document[0])
 
-                if self.text in document[1]:
+                if document[1].get(self.text) is not None:
                     yield (document[0], document[1][self.text], document[2])
                     offset += 1
-                elif self.object in document[1]:
+                elif document[1].get(self.object) is not None:
                     yield (document[0], document[1][self.object], document[2])
                     offset += 1
             else:

src/python/txtai/embeddings/search/base.py

@@ -92,7 +92,7 @@ class Search:
         Args:
             queries: list of queries
             limit: maximum results
-            weights: hybrid score weights
+            weights: hybrid score weights, either one set of weights for all queries or a list of [dense, sparse] weights per query
             index: index name
 
         Returns:
@@ -114,9 +114,13 @@ class Search:
             if isinstance(weights, (int, float)):
                 weights = [weights, 1 - weights]
 
+            # Use the same weights for each query, unless a list of weights per query is passed
+            if not isinstance(weights[0], (list, tuple)):
+                weights = [weights] * len(queries)
+
             # Create weighted scores via hybrid fusion strategy
             fusion = Hybrid(self.scoring)
-            return [fusion(vectors, weights, limit) for vectors in zip(dense, sparse)]
+            return [fusion(vectors, weights[x], limit) for x, vectors in enumerate(zip(dense, sparse))]
 
         # Raise an error if when no indexes are available
         if not sparse and not dense:
@@ -228,8 +232,8 @@ class Search:
         # Override limit with query limit, if applicable
         limit = max(limit, self.limit(queries))
 
-        # Bulk index scan
-        scan = Scan(self.search, limit, weights, index)(queries, parameters)
+        # Bulk index scan. The database query skips OFFSET rows, so add them to the number of candidates.
+        scan = Scan(self.search, limit + self.offset(queries), weights, index)(queries, parameters)
 
         # Combine index search results with database search results
         results = []
@@ -296,6 +300,32 @@ class Search:
 
         return qlimit
 
+    def offset(self, queries):
+        """
+        Parses the largest OFFSET clause from queries ordered by score.
+
+        Args:
+            queries: list of queries
+
+        Returns:
+            largest offset number or 0 if not found
+        """
+
+        qoffset = 0
+        for query in queries:
+            # Skip non-numeric offsets (e.g. a ":n" bind parameter)
+            o = query.get("offset")
+            o = int(o) if o and o.isdigit() else 0
+
+            # Only a query ordered first by score desc counts its offset. Any other order, or a select alias named score that replaces
+            # the similarity score in the order, must see the same candidates on every page.
+            if o > qoffset:
+                orderby = query.get("orderby")
+                if (not orderby or orderby.split(",")[0].strip().lower() == "score desc") and "score" not in query.get("aliases", {}):
+                    qoffset = o
+
+        return qoffset
+
     def graphsearch(self, queries, limit, weights, index):
         """
         Executes an index + graph search.
@@ -316,8 +346,8 @@ class Search:
         # Override limit with query limit, if applicable
         limit = max(limit, self.limit(queries))
 
-        # Bulk index scan
-        scan = Scan(self.search, limit, weights, index)(queries, None)
+        # Bulk index scan. The graph query drops SKIP rows, so add them to the number of candidates.
+        scan = Scan(self.search, limit + self.offset(queries), weights, index)(queries, None)
 
         # Combine index search results with database search results
         for x, query in enumerate(queries):

src/python/txtai/embeddings/search/explain.py

@@ -86,8 +86,9 @@ class Explain:
             # Parse query
             query = self.database.parse(query)
 
-            # Extract query from similar clause
-            query = " ".join([" ".join(clause) for clause in query["similar"]]) if "similar" in query else None
+            # Extract query text from similar clauses. Only the first argument is query text, the rest are
+            # candidates, index and weights.
+            query = " ".join([clause[0] for clause in query["similar"]]) if "similar" in query else None
 
         # Return original texts if query, text or score not present
         if not query or not texts or "score" not in texts[0] or "text" not in texts[0]:

src/python/txtai/embeddings/search/hybrid.py

@@ -194,7 +194,8 @@ class LogOdds:
         array = [s for s in raw if s > 0]
         if array:
             median = sorted(array)[len(array) // 2]
-            std = (sum((x - sum(array) / len(array)) ** 2 for x in array) / len(array)) ** 0.5
+            mean = sum(array) / len(array)
+            std = (sum((x - mean) ** 2 for x in array) / len(array)) ** 0.5
             alpha = 1.0 / std if std > 0 else 1.0
 
         return median, alpha

src/python/txtai/embeddings/search/scan.py

@@ -61,9 +61,9 @@ class Scan:
 
             candidates = max(candidates) if candidates else default
 
-            # Query weights to pass to batch search
-            weights = [query.weights for query in iqueries if query.weights is not None]
-            weights = max(weights) if weights else self.weights
+            # Query weights to pass to batch search. Each query keeps its own weights, queries without weights use the default.
+            weights = [query.weights if query.weights is not None else self.weights for query in iqueries]
+            weights = [[weight, 1 - weight] if isinstance(weight, (int, float)) else weight for weight in weights]
 
             # Index to run query against
             index = index if index else self.index

src/python/txtai/embeddings/search/terms.py

@@ -37,9 +37,10 @@ class Terms:
                 # Parse query
                 parse = self.database.parse(query)
 
-                # Join terms from similar clauses (absent when the query has no similar() clause)
+                # Join query text from similar clauses (absent when the query has no similar() clause).
+                # Only the first argument is query text, the rest are candidates, index and weights.
                 if "similar" in parse:
-                    terms.append(" ".join(" ".join(s) for s in parse["similar"]))
+                    terms.append(" ".join(s[0] for s in parse["similar"]))
 
             return terms
 

src/python/txtai/graph/base.py

@@ -7,7 +7,6 @@ from collections import Counter
 from .topics import Topics
 
 
-# pylint: disable=R0904
 class Graph:
     """
     Base class for Graph instances. This class builds graph networks. Supports topic modeling
@@ -460,6 +459,10 @@ class Graph:
 
                     # Also delete topic, if it's empty
                     if not self.topics[topic]:
+                        # Categories are stored by topic position
+                        if self.categories:
+                            self.categories.pop(list(self.topics).index(topic))
+
                         self.topics.pop(topic)
 
                 # Delete node
@@ -757,7 +760,7 @@ class Graph:
             ids = ids.keys() if ids else None
 
             # Infer topic
-            topic = Counter(self.attribute(x, "topic") for x in ids).most_common(1)[0][0] if ids else None
+            topic = self.commonattribute(ids, "topic")
             if topic:
                 # Add id to topic list and set topic attribute
                 self.topics[topic].append(node)
@@ -767,5 +770,21 @@ class Graph:
                 self.addattribute(node, "topicrank", len(self.topics[topic]) - 1)
 
                 # Infer category
-                category = Counter(self.attribute(x, "category") for x in ids).most_common(1)[0][0]
-                self.addattribute(node, "category", category)
+                category = self.commonattribute([x for x in ids if self.attribute(x, "topic") == topic], "category")
+                if category:
+                    self.addattribute(node, "category", category)
+
+    def commonattribute(self, ids, attribute):
+        """
+        Gets the most common non-empty attribute value across a list of nodes.
+
+        Args:
+            ids: list of node ids
+            attribute: attribute name
+
+        Returns:
+            most common value or None if no node has the attribute
+        """
+
+        values = Counter(value for value in (self.attribute(x, attribute) for x in ids or []) if value is not None)
+        return values.most_common(1)[0][0] if values else None

src/python/txtai/graph/networkx.py

@@ -24,7 +24,6 @@ from .base import Graph
 from .query import Query
 
 
-# pylint: disable=R0904
 class NetworkX(Graph):
     """
     Graph instance backed by NetworkX.
@@ -131,7 +130,7 @@ class NetworkX(Graph):
                 for value in column:
                     if isinstance(value, list):
                         # Path group
-                        nodes.update([node for node in value if node and not isinstance(node, dict)])
+                        nodes.update([node for node in value if node is not None and not isinstance(node, dict)])
                     elif isinstance(value, dict):
                         # Nodes by id attribute
                         nodes.update(uid for uid, attr in self.scan(data=True) if attr["id"] == value["id"])

src/python/txtai/graph/query.py

@@ -88,7 +88,7 @@ class Query:
         """
 
         # Parameters
-        where, limit, nodes, similar = None, None, [], []
+        where, limit, offset, nodes, similar = None, None, None, [], []
 
         # Parse where clause
         match = re.search(r"where(.+?)return", query, flags=re.DOTALL | re.IGNORECASE)
@@ -100,6 +100,13 @@ class Query:
         if match:
             limit = match.group(1)
 
+        # Parse skip clause when results are ordered first by score desc. Graph results with no ORDER BY do not follow score order.
+        # Most queries have no SKIP, so check for the keyword before the regex.
+        if "skip" in query.lower():
+            match = re.search(r"skip\s+(\d+)", query, flags=re.DOTALL | re.IGNORECASE)
+            if match and re.search(r"order\s+by\s+\w+\.score\s+desc", query, flags=re.DOTALL | re.IGNORECASE):
+                offset = match.group(1)
+
         # Parse similar clauses
         for x, match in enumerate(re.finditer(r"similar\((.+?)\)", query, flags=re.DOTALL | re.IGNORECASE)):
             # Replace similar clause with placeholder
@@ -115,6 +122,7 @@ class Query:
             "query": query,
             "where": where,
             "limit": limit,
+            "offset": offset,
             "nodes": nodes,
             "similar": similar,
         }

src/python/txtai/graph/rdbms.py

@@ -59,6 +59,10 @@ class RDBMS(NetworkX):
         else:
             yield from super().scan(attribute, data)
 
+    def addattribute(self, node, field, value):
+        if self.hasnode(node):
+            self.graph.backend.add_node(node, {field: value})
+
     def load(self, path):
         # Create graph instance
         self.graph, self.database = self.connect()

src/python/txtai/graph/topics.py

@@ -82,8 +82,10 @@ class Topics:
             # Get topn terms
             terms = self.topn(idf, topn)
 
-            # Sort community by score descending
-            community = [uid for uid, _ in scoring.search(terms, len(community))]
+            # Sort community by score descending. Ensure both scored and unscored nodes are kept in community.
+            ranked = [uid for uid, _ in scoring.search(terms, len(community))]
+            unranked = sorted(set(community) - set(ranked), key=lambda x: centrality[x], reverse=True)
+            community = ranked + unranked
         else:
             # No text found for topic, generate topic name
             terms = ["topic", str(index)]

src/python/txtai/models/pooling/base.py

@@ -36,7 +36,9 @@ class Pooling(Module):
         super().__init__()
 
         self.model = Models.load(path, modelargs=modelargs)
-        self.tokenizer = Models.tokenizer(tokenizer if tokenizer else path)
+        self.tokenizer = Models.tokenizer(
+            tokenizer if tokenizer else path, trust_remote_code=modelargs.get("trust_remote_code", False) if modelargs else False
+        )
         self.device = Models.device(device)
 
         # Detect unbounded tokenizer typically found in older models

src/python/txtai/models/pooling/factory.py

@@ -85,16 +85,20 @@ class PoolingFactory:
         # Load 1_Pooling/config.json file
         config = PoolingFactory.load(path, "1_Pooling/config.json")
 
+        # Models trained with Sentence Transformers v5.4+ use `pooling_mode` parameter
+        # Support both the new and legacy parameters
+        mode = config.get("pooling_mode") if config else None
+
         # Set to CLS pooling if it's enabled and mean pooling is disabled
-        if config and config.get("pooling_mode_cls_token") and not config.get("pooling_mode_mean_tokens"):
+        if config and config.get("pooling_mode_cls_token", mode == "cls") and not config.get("pooling_mode_mean_tokens", mode == "mean"):
             method = "clspooling"
 
         # Set to last token pooling if it's enabled and mean pooling is disabled
-        if config and config.get("pooling_mode_lasttoken") and not config.get("pooling_mode_mean_tokens"):
+        if config and config.get("pooling_mode_lasttoken", mode == "lasttoken") and not config.get("pooling_mode_mean_tokens", mode == "mean"):
             method = "lastpooling"
 
         # Set to max pooling if it's enabled and mean pooling is disabled
-        if config and config.get("pooling_mode_max_tokens") and not config.get("pooling_mode_mean_tokens"):
+        if config and config.get("pooling_mode_max_tokens", mode == "max") and not config.get("pooling_mode_mean_tokens", mode == "mean"):
             method = "maxpooling"
 
         # Check for late interaction pooling

src/python/txtai/models/pooling/late.py

@@ -291,7 +291,8 @@ class LatePooling(Pooling):
             config = self.load(path, "artifact.metadata")
             params = ["query_token_id", "query_maxlen", "doc_token_id", "doc_maxlen"]
 
-        return [config.get(p) for p in params]
+        # Use defaults when the settings file is missing
+        return [config.get(p) for p in params] if config else [None] * len(params)
 
     def loadlinear(self, path, models):
         """

src/python/txtai/models/pooling/lemur.py

@@ -209,6 +209,10 @@ class Lemur:
         with open(config, encoding="utf-8") as source:
             self.config = json.load(source)
 
+        # Encoder repos also have config.json and model.safetensors, check this is a LEMUR artifact
+        if "modeltype" not in self.config:
+            raise ValueError(f"{path} is not a LEMUR artifact, create one with LemurTrainer")
+
         self.model = LemurModel(**self.config).to(self.device)
         training = self.config.get("training", {})
         self.selectedepoch = training.get("selectedepoch")

src/python/txtai/pipeline/audio/signal.py

@@ -169,11 +169,15 @@ class Signal:
             audio with silence removed
         """
 
-        # Process in 20ms chunks
+        # Empty audio has no chunks to concatenate.
+        if len(audio) == 0:
+            return audio
+
+        # Process in chunks, including the final partial chunk
         n, offset = int(rate * (20 / 1000.0) * 2), 0
 
         chunks = []
-        while offset + n <= len(audio):
+        while offset < len(audio):
             # Calculate energy for chunk and detection result
             chunk = audio[offset : offset + n]
             energyfreq = Signal.energy(chunk, rate)

src/python/txtai/pipeline/audio/texttospeech.py

@@ -92,7 +92,7 @@ class TextToSpeech(Pipeline):
 
         # Streaming response
         if stream:
-            return self.stream(texts, speaker, encoding)
+            return self.stream(texts, speaker, encoding, **kwargs)
 
         # Transform text to speech
         results = [self.execute(x, speaker, encoding, **kwargs) for x in texts]
@@ -137,7 +137,7 @@ class TextToSpeech(Pipeline):
 
         return exists
 
-    def stream(self, texts, speaker, encoding):
+    def stream(self, texts, speaker, encoding, **kwargs):
         """
         Iterates over texts, splits into segments and yields snippets of audio.
         This method is designed to integrate with streaming LLM generation.
@@ -146,6 +146,7 @@ class TextToSpeech(Pipeline):
             texts: list of input texts
             speaker: speaker id
             encoding: audio encoding format
+            kwargs: additional keyword args
 
         Returns:
             snippets of audio as NumPy arrays or audio bytes depending on encoding parameter
@@ -157,11 +158,11 @@ class TextToSpeech(Pipeline):
 
             if x == "\n" or (x.strip().endswith(".") and len([y for y in buffer if y]) > 2):
                 data, buffer = "".join(buffer), []
-                yield self.execute(data, speaker, encoding)
+                yield self.execute(data, speaker, encoding, **kwargs)
 
         if buffer:
             data = "".join(buffer)
-            yield self.execute(data, speaker, encoding)
+            yield self.execute(data, speaker, encoding, **kwargs)
 
     def execute(self, text, speaker, encoding, **kwargs):
         """

src/python/txtai/pipeline/data/htmltomd.py

@@ -4,6 +4,8 @@ HTMLToMarkdown module
 
 import re
 
+from textwrap import indent
+
 # Conditional import
 try:
     from bs4 import BeautifulSoup, NavigableString
@@ -58,6 +60,14 @@ class HTMLToMarkdown(Pipeline):
         for script in soup.find_all(["script", "style"]):
             script.decompose()
 
+        # Replace line breaks with newlines. Table cells and headings must stay on a single line, use a space there.
+        for br in soup.find_all("br"):
+            parent = br.parent
+            br.replace_with(" " if br.find_parent(["th", "td", "h1", "h2", "h3", "h4", "h5", "h6"]) else "\n")
+
+            # Merge the line break with the text around it
+            parent.smooth()
+
         # Check for article sections
         article = next((x for x in ["article", "main"] if soup.find(x)), None)
 
@@ -184,12 +194,17 @@ class HTMLToMarkdown(Pipeline):
             target, text = x if x.name else node, x.text
 
             if text.strip():
+                # Leading whitespace separates this text from the previous word, keep it outside of the formatting
+                space = " " if text[0].isspace() and texts and not texts[-1][-1:].isspace() else ""
+
                 if target.name in ("b", "strong"):
-                    text = f"**{text.strip()}** "
+                    text = f"{space}**{text.strip()}** "
                 elif target.name in ("i", "em"):
-                    text = f"*{text.strip()}* "
+                    text = f"{space}*{text.strip()}* "
                 elif target.name == "a":
-                    text = f"[{text.strip()}]({target.get('href')}) "
+                    # Links without a target are added as plain text
+                    href = target.get("href")
+                    text = f"{space}[{text.strip()}]({href}) " if href else f"{space}{text.strip()} "
 
             texts.append(text)
 
@@ -253,21 +268,54 @@ class HTMLToMarkdown(Pipeline):
             list as markdown
         """
 
-        elements = []
-        for x, element in enumerate(node.find_all("li")):
-            # Unordered lists use dashes. Ordered lists use numbers.
-            prefix = "-" if node.name == "ul" else f"{x + 1}."
+        elements, count = [], 0
+        if node.name == "ol":
+            try:
+                count = int(node.get("start", 1)) - 1
+            except ValueError:
+                pass
 
-            # List item text
-            text = self.process(element, article)
+        for element in node.find_all(("li", "ul", "ol"), recursive=False):
+            # Nested list that is a direct child of this list
+            if element.name != "li":
+                self.nested(element, article, elements)
+            else:
+                # Unordered lists use dashes. Ordered lists use numbers.
+                count += 1
+                prefix = "-" if node.name == "ul" else f"{count}."
 
-            # Add list element
-            if text:
-                elements.append(f"{prefix} {text}")
+                # Detach nested lists, these are formatted separately
+                nested = [child.extract() for child in element.find_all(("ul", "ol"), recursive=False)]
+
+                # List item text
+                text = self.process(element, article)
+
+                # Add list element
+                if text:
+                    elements.append(f"{prefix} {text}")
+
+                # Add nested lists
+                for child in nested:
+                    self.nested(child, article, elements)
 
         # Join elements together as string
         return "\n".join(elements)
 
+    def nested(self, node, article, elements):
+        """
+        Nested list handler. Formats a nested list and adds it to elements indented with a tab,
+        which is preserved by text cleaning.
+
+        Args:
+            node: nested list node
+            article: True if the main section node is an article
+            elements: list of formatted list elements
+        """
+
+        text = self.items(node, article)
+        if text:
+            elements.append(indent(text, "\t"))
+
     def code(self, node):
         """
         Code block handler. This method transforms a HTML pre or code block into a Markdown formatted

src/python/txtai/pipeline/data/tabular.py

@@ -99,12 +99,15 @@ class Tabular(Pipeline):
         columns = self.textcolumns
         if not columns:
             columns = list(df.columns)
-            if self.idcolumn:
+            if self.idcolumn is not None:
                 columns.remove(self.idcolumn)
 
         # Transform into (id, text, tag) tuples
-        for index, row in df.iterrows():
-            uid = row[self.idcolumn] if self.idcolumn else index
+        # Tuple iteration preserves column types instead of coercing each row to one dtype.
+        for index, values in zip(df.index, df.itertuples(index=False, name=None)):
+            row = dict(zip(df.columns, values))
+            # Coerce NaN ids to None so the row-index fallback applies.
+            uid = self.column(row[self.idcolumn]) if self.idcolumn is not None else index
             uid = uid if uid is not None else index
             text = self.concat(row, columns)
 
@@ -112,10 +115,10 @@ class Tabular(Pipeline):
 
             # Also add row for content
             if isinstance(self.content, list):
-                row = {column: self.column(value) for column, value in row.to_dict().items() if column in self.content}
+                row = {column: self.column(value) for column, value in row.items() if column in self.content}
                 rows.append((uid, row, None))
             elif self.content:
-                row = {column: self.column(value) for column, value in row.to_dict().items()}
+                row = {column: self.column(value) for column, value in row.items()}
                 rows.append((uid, row, None))
 
         return rows

src/python/txtai/pipeline/data/tokenizer.py

@@ -82,7 +82,7 @@ class Tokenizer(Pipeline):
             #  - Strings to be at least 2 characters long AND
             #  - At least 1 non-trailing alpha character in string
             # Note: The standard Python re module is much faster than regex for this expression
-            self.alphanum = re.compile(r"^\d*[a-z][\-.0-9:_a-z]{1,}$")
+            self.alphanum = re.compile(r"^\d*[a-zA-Z][\-.0-9:_a-zA-Z]{1,}$")
         elif regexp:
             # Regular expression for tokenization
             self.regexp = regex.compile(regexp)

src/python/txtai/pipeline/hfmodel.py

@@ -64,7 +64,11 @@ class HFModel(Tensors):
         model = self.prepare(model)
 
         # Load tokenizer
-        tokenizer = transformers.AutoTokenizer.from_pretrained(tokenizer) if isinstance(tokenizer, str) else tokenizer
+        tokenizer = (
+            transformers.AutoTokenizer.from_pretrained(tokenizer, trust_remote_code=kwargs.get("trust_remote_code", False))
+            if isinstance(tokenizer, str)
+            else tokenizer
+        )
 
         return model, tokenizer
 

src/python/txtai/pipeline/image/caption.py

@@ -38,8 +38,8 @@ class Caption(HFModel):
             self.model, self.tokenizer, self.processor = path
         else:
             self.model = transformers.AutoModelForImageTextToText.from_pretrained(path, **kwargs)
-            self.tokenizer = transformers.AutoTokenizer.from_pretrained(path)
-            self.processor = transformers.AutoImageProcessor.from_pretrained(path)
+            self.tokenizer = transformers.AutoTokenizer.from_pretrained(path, trust_remote_code=kwargs.get("trust_remote_code", False))
+            self.processor = transformers.AutoImageProcessor.from_pretrained(path, trust_remote_code=kwargs.get("trust_remote_code", False))
 
         # Move model to device
         self.model = self.model.to(self.device)

src/python/txtai/pipeline/llm/generation.py

@@ -173,12 +173,15 @@ class Generation:
         """
 
         # Consume "thinking" tokens
-        text, buffer = None, ""
+        text, buffer = "", ""
         for chunk in results:
             buffer += chunk
             text = self.cleanthink(buffer)
             if text != buffer:
                 break
+            prefix = buffer.lstrip()
+            if not any(marker.startswith(prefix) or prefix.startswith(marker) for marker in ("<think>", "<|")):
+                break
 
         # Yield remaining tokens
         yield from text.lstrip()

src/python/txtai/pipeline/text/labels.py

@@ -120,19 +120,17 @@ class Labels(HFPipeline):
         result = [(config.label2id.get(x["label"], 0), x["score"]) for x in result]
 
         if labels:
+            # Map lowercase label names to label ids
+            names = {x.lower(): uid for x, uid in config.label2id.items()}
+
             matches = []
             for label in labels:
-                # Lookup label keys from model config
-                if label.isdigit():
-                    label = int(label)
-                    keys = list(config.id2label.keys())
-                else:
-                    label = label.lower()
-                    keys = [x.lower() for x in config.label2id.keys()]
+                # Resolve label to a label id, by id or by name
+                uid = int(label) if label.isdigit() else names.get(label.lower())
 
-                # Find and add label match
-                if label in keys:
-                    matches.append(keys.index(label))
+                # Add label match, if found
+                if uid in config.id2label:
+                    matches.append(uid)
 
             return [(label, score) for label, score in result if label in matches]
 

src/python/txtai/pipeline/text/lateencoder.py

@@ -95,12 +95,21 @@ class LateEncoder(Pipeline):
             list of (id, score)
         """
 
+        # Padding rows are all zeros, true token rows are unit normalized
+        qmask, dmask = queries.abs().sum(axis=-1) > 0, data.abs().sum(axis=-1) > 0
+
         # Compute bulk dot product using einstein notation
-        scores = torch.einsum("ash,bth->abst", queries, data).max(axis=-1).values.mean(axis=-1)
+        scores = torch.einsum("ash,bth->abst", queries, data)
+
+        # Max score per query token, ignoring padded data tokens
+        scores = scores.masked_fill(~dmask[None, :, None, :], torch.finfo(scores.dtype).min).max(axis=-1).values
+
+        # Mean score over true query tokens, ignoring padded query tokens
+        scores = scores.masked_fill(~qmask[:, None, :], 0).sum(axis=-1) / qmask.sum(axis=-1, keepdim=True).clamp(min=1)
         scores = scores.cpu().numpy()
 
         # Get top n matching indices and scores
-        indices = np.argpartition(-scores, limit if limit and limit < scores.shape[0] else scores.shape[0] - 1)[:, :limit]
+        indices = np.argpartition(-scores, limit if limit and limit < scores.shape[1] else scores.shape[1] - 1)[:, :limit]
         scores = np.take_along_axis(scores, indices, axis=1)
 
         # argpartition doesn't order within the top n - sort each row by score descending

src/python/txtai/pipeline/text/reranker.py

@@ -2,7 +2,18 @@
 Reranker module
 """
 
+import hashlib
+import json
+import logging
+from collections import OrderedDict
+from threading import RLock
+
 from ..base import Pipeline
+from ...util import Library
+
+library = Library()
+np, torch, safetensors = library.numpy(), library.torch(), library.safetensors()
+logger = logging.getLogger(__name__)
 
 
 class Reranker(Pipeline):
@@ -11,16 +22,23 @@ class Reranker(Pipeline):
     embeddings instance for this to work properly.
     """
 
-    def __init__(self, embeddings, similarity):
+    # Cache lock shared by all instances, since rerankers can share one similarity pipeline
+    lock = RLock()
+
+    def __init__(self, embeddings, similarity, cache=None):
         """
         Creates a Reranker pipeline.
 
         Args:
             embeddings: embeddings instance (content must be enabled)
             similarity: similarity instance
+            cache: optional document token cache; True bounds at 1000, False/None disable; other values use int(cache), non-positive disables
         """
 
         self.embeddings, self.similarity = embeddings, similarity
+        self.cachesize = 1000 if cache is True else 0 if cache is None or cache is False else int(cache)
+        self.cache = OrderedDict() if self.cachesize > 0 else None
+        self.namespace = None
 
     # pylint: disable=W0222
     def __call__(self, query, limit=3, factor=10, **kwargs):
@@ -45,13 +63,110 @@ class Reranker(Pipeline):
         # Re-rank using similarity pipeline
         ranked = []
         for x, result in enumerate(results):
-            texts = [row["text"] for row in result]
+            # Only score when a search has results
+            if result:
+                texts = [row["text"] for row in result]
 
-            # Score results and merge
-            for uid, score in self.similarity(queries[x], texts):
-                result[uid]["score"] = score
+                if self.cache is not None and (namespace := self.modelkey()) and all("id" in row for row in result):
+                    with self.lock:
+                        texts = self.vectors(result, namespace)
+                        for uid, score in self.similarity(queries[x], texts):
+                            result[uid]["score"] = score
+                else:
+                    # Score results and merge
+                    for uid, score in self.similarity(queries[x], texts):
+                        result[uid]["score"] = score
 
             # Sort and take top n sorted results
             ranked.append(sorted(result, key=lambda row: row["score"], reverse=True)[:limit])
 
         return ranked[0] if isinstance(query, str) else ranked
+
+    def modelkey(self):
+        """Compute a stable namespace for batch-independent document encoding."""
+
+        encoder = getattr(self.similarity, "lateencoder", None)
+        if not encoder:
+            return None
+
+        pooling = encoder.model
+        center = getattr(pooling, "center", None)
+        if center and center["scope"] == "batch":
+            return None
+
+        mean = center.get("mean") if center else None
+        identity = {
+            "pooling": type(pooling).__name__,
+            "model": getattr(pooling.model.config, "_name_or_path"),
+            "revision": getattr(pooling.model.config, "_commit_hash", None),
+            "tokenizer": pooling.tokenizer.name_or_path,
+            "scope": center["scope"] if center else None,
+            "mean": hashlib.sha256(np.asarray(mean).tobytes()).hexdigest() if mean is not None else None,
+        }
+        return hashlib.sha256(json.dumps(identity, sort_keys=True).encode("utf-8")).hexdigest()
+
+    def reset(self, namespace):
+        """Clear entries when the encoder namespace changes; caller holds the cache lock."""
+
+        if self.namespace != namespace:
+            self.cache.clear()
+            self.namespace = namespace
+
+    def vectors(self, rows, namespace):
+        """Gather cached vectors and encode misses once, preserving candidate positions."""
+
+        with self.lock:
+            self.reset(namespace)
+            keys, vectors, misses = [], [], []
+            for row in rows:
+                key = (namespace, row["id"], hashlib.sha256(row["text"].encode("utf-8")).hexdigest())
+                vector = self.cache.get(key)
+                keys.append(key)
+                vectors.append(vector)
+                if vector is None:
+                    misses.append(len(vectors) - 1)
+                else:
+                    self.cache.move_to_end(key)
+
+            if misses:
+                encoded = self.similarity.encode([rows[x]["text"] for x in misses], "data")
+                for position, vector in zip(misses, encoded):
+                    vectors[position] = vector[vector.abs().sum(dim=-1) > 0].detach().to(device="cpu", dtype=torch.float32)
+
+            # Keep local references through padding before LRU insertion can evict any current candidate.
+            batch = torch.zeros((len(vectors), max(vector.shape[0] for vector in vectors), vectors[0].shape[1]), dtype=torch.float32)
+            for position, vector in enumerate(vectors):
+                batch[position, : vector.shape[0]] = vector
+
+            for position in misses:
+                self.cache.update({keys[position]: vectors[position]})
+                while len(self.cache) > self.cachesize:
+                    self.cache.popitem(last=False)
+
+            return batch.to(self.similarity.lateencoder.device)
+
+    def save(self, path):
+        """Save an independent safetensors snapshot; disabled or unsupported caches do nothing."""
+
+        if self.cache is not None and (namespace := self.modelkey()):
+            with self.lock:
+                self.reset(namespace)
+                metadata = {"version": "1", "namespace": namespace, "keys": json.dumps([list(key[1:]) for key in self.cache.keys()])}
+                safetensors.numpy.save_file({str(x): vector.numpy() for x, vector in enumerate(self.cache.values())}, path, metadata)
+
+    def load(self, path):
+        """Load a bounded snapshot, warning on model mismatch; disabled or unsupported caches do nothing."""
+
+        if self.cache is not None and (namespace := self.modelkey()):
+            with self.lock:
+                self.cache.clear()
+                self.namespace = namespace
+                with safetensors.safe_open(path, framework="np") as source:
+                    metadata = source.metadata()
+                    if metadata["namespace"] != namespace:
+                        logger.warning("Ignoring reranker cache snapshot with a different encoder namespace")
+                        return
+                    for x, key in enumerate(json.loads(metadata["keys"])):
+                        self.cache.update({(namespace, *key): torch.from_numpy(source.get_tensor(str(x)))})
+                        while len(self.cache) > self.cachesize:
+                            self.cache.popitem(last=False)

src/python/txtai/pipeline/text/similarity.py

@@ -57,7 +57,7 @@ class Similarity(Labels):
             return self.crossencoder(query, texts, multilabel, labels=labels)
 
         if self.lateencoder:
-            return self.lateencoder(query, texts)
+            return self.lateencoder(query, texts, **kwargs)
 
         # Call Labels pipeline for texts using input query as the candidate label
         scores = super().__call__(texts, [query] if isinstance(query, str) else query, multilabel, **kwargs)

src/python/txtai/pipeline/text/summary.py

@@ -91,14 +91,20 @@ class Summary(HFModel):
         if prefix:
             inputs = [prefix + x for x in inputs]
 
-        # Tokenize inputs
-        tokens = self.tokenizer(inputs, truncation=kwargs.pop("truncation", False), return_tensors="pt").to(self.device)
+        truncation = kwargs.pop("truncation", False)
 
-        # Generate outputs
-        outputs = self.model.generate(**tokens, **kwargs)
+        results = []
+        for batch in self.batch(inputs, self.batchsize):
+            # Tokenize inputs, padding texts of different lengths within a batch
+            tokens = self.tokenizer(batch, padding=True, truncation=truncation, return_tensors="pt").to(self.device)
 
-        # Decode and return
-        return [self.tokenizer.decode(x, skip_special_tokens=True) for x in outputs]
+            # Generate outputs
+            outputs = self.model.generate(**tokens, **kwargs)
+
+            # Decode outputs
+            results.extend(self.tokenizer.decode(x, skip_special_tokens=True) for x in outputs)
+
+        return results
 
     def clean(self, text):
         """

src/python/txtai/scoring/terms.py

@@ -13,6 +13,7 @@ from collections import Counter
 from threading import RLock
 
 # Core library imports
+from ..serialize import SerializeFactory
 from ..util import Library
 
 np = Library().numpy()
@@ -37,7 +38,7 @@ class Terms:
 
     INSERT_TERM = "INSERT OR REPLACE INTO terms VALUES (?, ?, ?)"
     SELECT_TERMS = "SELECT ids, freqs FROM terms WHERE term = ?"
-    WILDCARD_TERMS = "SELECT term FROM terms WHERE term LIKE ?"
+    WILDCARD_TERMS = "SELECT term FROM terms WHERE term LIKE ? ESCAPE '\\'"
 
     # Documents table
     CREATE_DOCUMENTS = """
@@ -129,14 +130,21 @@ class Terms:
             ids: ids to delete
         """
 
-        # Set index ids as deleted, ignoring ids that were never indexed and ids
-        # already marked deleted. Skipping duplicates keeps count() accurate when
-        # the same id is deleted more than once (e.g. repeated upsert/delete flows).
-        for i in ids:
-            if i in self.ids:
-                indexid = self.ids.index(i)
-                if indexid not in self.deletes:
+        # Walk each requested id from its last match to mark every index id, including re-added ids
+        deletes = set(self.deletes)
+        for uid in ids:
+            start = 0
+            while True:
+                try:
+                    indexid = self.ids.index(uid, start)
+                except ValueError:
+                    break
+
+                if indexid not in deletes:
                     self.deletes.append(indexid)
+                    deletes.add(indexid)
+
+                start = indexid + 1
 
     def index(self):
         """
@@ -271,9 +279,13 @@ class Terms:
             # Index id - length
             self.lengths.append(length)
 
-        # Cast ids to int if every id is an integer
-        if all(uid.isdigit() for uid in self.ids):
+        # Retain integer inference for legacy indexes with plain TEXT ids.
+        if all(isinstance(uid, str) and uid.isdigit() for uid in self.ids):
             self.ids = [int(uid) for uid in self.ids]
+        else:
+            # Serialized ids preserve types, including numeric strings and mixed string/integer ids.
+            serializer = SerializeFactory.create()
+            self.ids = [serializer.loadbytes(uid) if isinstance(uid, bytes) else uid for uid in self.ids]
 
         # Clear cache
         self.weights.cache_clear()
@@ -286,12 +298,16 @@ class Terms:
             path: path to write terms database
         """
 
+        # Initialize database, if necessary
+        self.initialize()
+
         # Clear documents table
         self.cursor.execute(Terms.DELETE_DOCUMENTS)
 
-        # Save document attributes
+        # Serialize ids as BLOBs so SQLite TEXT affinity does not discard their types.
+        serializer = SerializeFactory.create()
         for i, uid in enumerate(self.ids):
-            self.cursor.execute(Terms.INSERT_DOCUMENT, [i, uid, 1 if i in self.deletes else 0, self.lengths[i]])
+            self.cursor.execute(Terms.INSERT_DOCUMENT, [i, serializer.savebytes(uid), 1 if i in self.deletes else 0, self.lengths[i]])
 
         # Temporary database
         if not self.path:
@@ -438,6 +454,7 @@ class Terms:
         Expands a list of terms using the following term expansion rules.
 
           - Asterisks (*) for wildcard expressions
+          - In wildcard expressions, SQL LIKE % and _ remain wildcards; backslash escapes literal characters
 
         Args:
             terms: list of terms
@@ -548,6 +565,9 @@ class Terms:
             # Update scores
             scores[uids] += freq * weights
 
+        # Common term scores can restore deleted candidates; clear them after merging.
+        scores[self.deletes] = 0
+
     def candidates(self, scores, topn):
         """
         Gets the topn scored candidates. This method ignores deleted documents.

src/python/txtai/vectors/base.py

@@ -96,16 +96,23 @@ class Vectors:
             model
         """
 
+        # Build a cache key from settings that change the loaded model when a configuration is present.
+        if self.config:
+            config = {"path": path, "method": self.config.get("method"), "vectors": self.config.get("vectors", {})}
+            key = json.dumps(config, sort_keys=True, default=str)
+        else:
+            key = path
+
         # Check if model is cached
-        if self.models and path in self.models:
-            return self.models[path]
+        if self.models and key in self.models:
+            return self.models[key]
 
         # Create new model
         model = self.loadmodel(path)
 
         # Store model in cache
         if self.models is not None and path:
-            self.models[path] = model
+            self.models[key] = model
 
         return model
 
@@ -378,6 +385,9 @@ class Vectors:
         embeddings = self.encode(data, category)
 
         if embeddings is not None:
+            if embeddings.ndim != 2:
+                raise ValueError(f"Expected 2 dimensional vectors, produced {embeddings.ndim} dimensions")
+
             # Truncate embeddings, if necessary
             if self.dimensionality and self.dimensionality < embeddings.shape[1]:
                 embeddings = self.truncate(embeddings)

src/python/txtai/vectors/dense/factory.py

@@ -11,6 +11,7 @@ from .litellm import LiteLLM
 from .litert import LiteRT
 from .llama import LlamaCpp
 from .m2v import Model2Vec
+from .onnx import ONNX
 from .sbert import STVectors
 from .words import WordVectors
 
@@ -57,6 +58,10 @@ class VectorsFactory:
         if method == "model2vec":
             return Model2Vec(config, scoring, models)
 
+        # ONNX vectors
+        if method == "onnx":
+            return ONNX(config, scoring, models)
+
         # Sentence Transformers vectors
         if method == "sentence-transformers":
             return STVectors(config, scoring, models) if config and config.get("path") else None
@@ -99,6 +104,8 @@ class VectorsFactory:
                     method = "llama.cpp"
                 elif Model2Vec.ismodel(path):
                     method = "model2vec"
+                elif ONNX.ismodel(path):
+                    method = "onnx"
                 elif WordVectors.ismodel(path):
                     method = "words"
                 else:

src/python/txtai/vectors/dense/litert.py

@@ -13,7 +13,7 @@ try:
 except ImportError:
     LITERT = False
 
-from ...util import Download, Library
+from ...util import Download, DownloadError, Library
 
 from ..base import Vectors
 
@@ -51,12 +51,23 @@ class LiteRT(Vectors):
         # Check if this is a local path, otherwise download from the HF Hub
         model = path if os.path.exists(path) else Download()(path)
 
-        # Also local tokenizer file
-        tokenizer = os.path.dirname(path) + "/" + "tokenizer.json"
-        tokenizer = tokenizer if os.path.exists(tokenizer) else Download()(tokenizer)
+        # Derive tokenizer file path
+        tokenizer = self.config.get("tokenizer")
+        if not tokenizer:
+            tokenizer = os.path.dirname(path) + "/" + "tokenizer.json"
+            try:
+                tokenizer = tokenizer if os.path.exists(tokenizer) else Download()(tokenizer)
+
+            # Model is in a subdirectory of a HF Hub repo, read the tokenizer from the repo root
+            except DownloadError:
+                # Local models don't have a HF Hub repo root
+                if os.path.exists(path):
+                    raise
+
+                tokenizer = Download()("/".join(path.split("/")[:2]), "tokenizer.json")
 
         # Load tokenizer and model
-        tokenizer = Tokenizer.from_file(tokenizer)
+        tokenizer = Tokenizer.from_file(tokenizer) if os.path.exists(tokenizer) else Tokenizer.from_pretrained(tokenizer)
         model = CompiledModel.from_file(
             model,
             HardwareAccelerator.GPU | HardwareAccelerator.NPU | HardwareAccelerator.CPU if self.config.get("gpu", True) else HardwareAccelerator.CPU,

src/python/txtai/vectors/dense/m2v.py

@@ -41,8 +41,8 @@ class Model2Vec(Vectors):
                     config = json.load(f)
                     return config.get("model_type") == "model2vec"
 
-        # Ignore this error - invalid repo or directory
-        except DownloadError:
+        # Ignore invalid repo/directory and when HF Hub is not available
+        except (DownloadError, ImportError):
             pass
 
         return False

src/python/txtai/vectors/dense/onnx.py

@@ -0,0 +1,170 @@
+"""
+ONNX module
+"""
+
+import os
+
+# Conditional import
+try:
+    import onnxruntime as ort
+
+    from tokenizers import Tokenizer
+
+    ONNX_RUNTIME = True
+except ImportError:
+    ONNX_RUNTIME = False
+
+from ...util import Download, DownloadError, Library
+
+from ..base import Vectors
+
+# Core library imports
+np = Library().numpy()
+
+
+class ONNX(Vectors):
+    """
+    Builds vectors using ONNX Runtime.
+    """
+
+    @staticmethod
+    def ismodel(path):
+        """
+        Checks if path is an ONNX model.
+
+        Args:
+            path: input path
+
+        Returns:
+            True if this is an ONNX model, False otherwise
+        """
+
+        return isinstance(path, str) and path.lower().endswith(".onnx")
+
+    def __init__(self, config, scoring, models):
+        # Check before parent constructor since it calls loadmodel
+        if not ONNX_RUNTIME:
+            raise ImportError('onnxruntime is not available - install "vectors" extra to enable')
+
+        super().__init__(config, scoring, models)
+
+    def loadmodel(self, path):
+        # Check if this is a local path, otherwise download from the HF Hub
+        model = path if os.path.exists(path) else Download()(path)
+
+        # Create ONNX session
+        session = ort.InferenceSession(model, ort.SessionOptions(), self.providers())
+
+        # Load tokenizer. ONNX supports dynamic shapes, so pad to the longest sequence in each batch
+        tokenizer = self.loadtokenizer(path)
+        tokenizer.enable_padding()
+
+        maxlength = self.config.get("maxlength")
+        if maxlength:
+            tokenizer.enable_truncation(max_length=maxlength)
+
+        return (session, tokenizer)
+
+    def loadtokenizer(self, path):
+        """
+        Loads the tokenizer for a model. Reads the `tokenizer` configuration option when set,
+        otherwise falls back to a tokenizer.json file stored alongside the model or at the root of its HF Hub repo.
+
+        Args:
+            path: model path
+
+        Returns:
+            Tokenizer
+        """
+
+        # Derive tokenizer file path
+        tokenizer = self.config.get("tokenizer")
+        if not tokenizer:
+            tokenizer = os.path.dirname(path) + "/" + "tokenizer.json"
+            try:
+                tokenizer = tokenizer if os.path.exists(tokenizer) else Download()(tokenizer)
+
+            # Model is in a subdirectory of a HF Hub repo, read the tokenizer from the repo root
+            except DownloadError:
+                # Local models don't have a HF Hub repo root
+                if os.path.exists(path):
+                    raise
+
+                tokenizer = Download()("/".join(path.split("/")[:2]), "tokenizer.json")
+
+        # Local tokenizer file vs a model id resolved through the HF Hub
+        return Tokenizer.from_file(tokenizer) if os.path.exists(tokenizer) else Tokenizer.from_pretrained(tokenizer)
+
+    def providers(self):
+        """
+        Returns a list of available and usable providers.
+
+        Returns:
+            list of available and usable providers
+        """
+
+        # Prefer the CUDA provider when it's available, it requires onnxruntime-gpu
+        if self.config.get("gpu", True) and "CUDAExecutionProvider" in ort.get_available_providers():
+            return ["CUDAExecutionProvider", "CPUExecutionProvider"]
+
+        return ["CPUExecutionProvider"]
+
+    def encode(self, data, category=None):
+        # Unpack model
+        session, tokenizer = self.model
+
+        results = []
+        for start in range(0, len(data), self.encodebatch):
+            batch = data[start : start + self.encodebatch]
+
+            inputs, masks = self.inputs(session, tokenizer.encode_batch(batch))
+
+            # Run model and pool the outputs into a single vector per input
+            results.append(self.pool(session.run(None, inputs)[0], masks))
+
+        return np.vstack(results)
+
+    def inputs(self, session, encoding):
+        """
+        Builds model inputs from a tokenized batch. Only inputs declared by the model are set,
+        since not every exported model takes token_type_ids.
+
+        Args:
+            session: inference session
+            encoding: tokenized batch
+
+        Returns:
+            (model inputs, attention masks)
+        """
+
+        masks = np.array([x.attention_mask for x in encoding], dtype=np.int64)
+        available = {x.name for x in session.get_inputs()}
+
+        features = {
+            "input_ids": np.array([x.ids for x in encoding], dtype=np.int64),
+            "attention_mask": masks,
+            "token_type_ids": np.array([x.type_ids for x in encoding], dtype=np.int64),
+        }
+
+        return {name: value for name, value in features.items() if name in available}, masks
+
+    def pool(self, outputs, masks):
+        """
+        Builds a single vector per input. Token-level outputs are mean pooled with the attention
+        mask, outputs that are already pooled are returned as is.
+
+        Args:
+            outputs: model outputs
+            masks: attention masks
+
+        Returns:
+            embeddings
+        """
+
+        # Model already returns a vector per input
+        if outputs.ndim == 2:
+            return outputs
+
+        # Mean pooling, ignoring padding tokens
+        masks = np.expand_dims(masks, -1).astype(outputs.dtype)
+        return np.sum(outputs * masks, axis=1) / np.clip(masks.sum(axis=1), a_min=1e-9, a_max=None)

src/python/txtai/vectors/dense/words.py

@@ -97,8 +97,8 @@ class WordVectors(Vectors):
                     config = json.load(f)
                     return config.get("model_type") == "staticvectors"
 
-        # Ignore this error - invalid repo or directory
-        except DownloadError:
+        # Ignore invalid repo/directory and when HF Hub is not available
+        except (DownloadError, ImportError):
             pass
 
         return False

src/python/txtai/vectors/sparse/base.py

@@ -2,6 +2,8 @@
 SparseVectors module
 """
 
+import os
+
 # Conditional import
 try:
     from scipy.sparse import csr_matrix, vstack
@@ -56,6 +58,10 @@ class SparseVectors(Vectors):
                 data = self.loadembeddings(queue)
                 embeddings = vstack((embeddings, data)) if embeddings is not None else data
 
+        # Remove temporary file (if checkpointing is disabled)
+        if not checkpoint:
+            os.remove(stream)
+
         # Return sparse array
         return (ids, dimensions, embeddings)
 

src/python/txtai/version.py

@@ -3,4 +3,4 @@ Version strings
 """
 
 # Current version tag
-__version__ = "9.13.0"
+__version__ = "9.15.0"

src/python/txtai/workflow/factory.py

@@ -24,16 +24,20 @@ class WorkflowFactory:
             Workflow
         """
 
+        # Copy configuration before resolving tasks
+        config = config.copy()
+
         # Resolve workflow tasks
         tasks = []
         for tconfig in config["tasks"]:
+            tconfig = tconfig.copy()
             task = tconfig.pop("task") if "task" in tconfig else ""
             tasks.append(TaskFactory.create(tconfig, task))
 
         config["tasks"] = tasks
 
         if "stream" in config:
-            sconfig = config["stream"]
+            sconfig = config["stream"].copy()
             task = sconfig.pop("task") if "task" in sconfig else "stream"
 
             config["stream"] = TaskFactory.create(sconfig, task)

src/python/txtai/workflow/task/base.py

@@ -6,6 +6,8 @@ import logging
 import re
 import types
 
+from collections.abc import Iterator
+
 # Core library imports
 from ...util import Library
 
@@ -265,14 +267,21 @@ class Task:
         Executes action(s) on elements.
 
         Args:
-            elements: list of data elements
+            elements: iterable data elements
             executor: execute instance, enables concurrent task actions
 
         Returns:
             transformed data elements
+
+        Raises:
+            ValueError: if iterator inputs are passed to a multi-action task
         """
 
         if self.action:
+            # Reject single-pass inputs without consuming or buffering them.
+            if len(self.action) > 1 and isinstance(elements, Iterator):
+                raise ValueError("Iterator inputs are not supported for multi-action tasks")
+
             # Run actions
             outputs = []
             for x, action in enumerate(self.action):
@@ -359,7 +368,14 @@ class Task:
             return outputs
 
         if self.merge == "vstack":
-            return self.vstack(outputs)
+            results = self.vstack(outputs)
+
+            if isinstance(results, np.ndarray) or torch.is_tensor(results):
+                size = len(self.action)
+                results = [OneToMany(list(results[x : x + size])) for x in range(0, len(results), size)]
+
+            return results
+
         if self.merge == "concat":
             return self.concat(outputs)
 
@@ -474,7 +490,8 @@ class Task:
             list of concat outputs
         """
 
-        return [". ".join([str(y) for y in x if y]) for x in self.hstack(outputs)]
+        # Only skip empty values: 0 and False are data, as the tabular pipeline's concat already treats them
+        return [". ".join([str(y) for y in x if y is not None and y != ""]) for x in self.hstack(outputs)]
 
 
 class OneToMany:

src/python/txtai/workflow/task/factory.py

@@ -47,6 +47,9 @@ class TaskFactory:
             Task
         """
 
+        # Copy configuration before binding arguments
+        config = config.copy()
+
         # Create lambda function if additional arguments present
         if "args" in config:
             args = config.pop("args")

src/python/txtai/workflow/task/retrieve.py

@@ -44,11 +44,11 @@ class RetrieveTask(Task):
         self.safeoutput = SafeOpen(safeopen=directory if safeopen else safeopen, allowurl=False)
 
     def prepare(self, element):
-        # Extract file path from URL
-        path = urlparse(element).path
-
         # Validate input element
-        url, _ = self.safeinput.valid(element)
+        url, exists = self.safeinput.valid(element)
+
+        # Preserve local filenames; only parse remote inputs as URLs
+        path = url if exists else urlparse(element).path
         if url:
             if self.flatten:
                 # Flatten directory structure (default)

src/python/txtai/workflow/task/service.py

@@ -19,7 +19,7 @@ class ServiceTask(Task):
     Task to runs requests against remote service urls.
     """
 
-    def register(self, url=None, method=None, params=None, batch=True, extract=None):
+    def register(self, url=None, method=None, params=None, batch=True, extract=None, timeout=30):
         """
         Adds service parameters to task. Checks if required dependencies are installed.
 
@@ -29,6 +29,7 @@ class ServiceTask(Task):
             params: default query parameters
             batch: if True, all elements are passed in a single batch request, otherwise a service call is executed per element
             extract: list of sections to extract from response
+            timeout: request timeout in seconds, defaults to 30
         """
 
         if not XML_TO_DICT:
@@ -43,6 +44,9 @@ class ServiceTask(Task):
         # If True, all elements are passed in a single batch request, otherwise a service call is executed per element
         self.batch = batch
 
+        # Request timeout in seconds. Prevents a slow or unresponsive endpoint from blocking the workflow indefinitely.
+        self.timeout = timeout
+
         # Save sections to extract. Supports both a single string and a hierarchical list of sections.
         self.extract = extract
         if self.extract:
@@ -50,7 +54,10 @@ class ServiceTask(Task):
 
     def execute(self, elements, executor=None):
         if self.batch:
-            elements = self.request(elements)
+            response = self.request(elements)
+
+            # Empty response body returns None for each element, same as batch=False
+            elements = [None] * len(elements) if response is None else response
         else:
             elements = [self.request(element) for element in elements]
 
@@ -83,12 +90,16 @@ class ServiceTask(Task):
 
         # Run request
         if self.method and self.method.lower() == "get":
-            response = requests.get(self.url, params=params)
+            response = requests.get(self.url, params=params, timeout=self.timeout)
         else:
-            response = requests.post(self.url, json=params)
+            response = requests.post(self.url, json=params, timeout=self.timeout)
+
+        # Empty response body (e.g. 204 No Content), nothing to parse
+        if not response.content:
+            return None
 
         # Parse data based on content-type
-        mimetype = response.headers["Content-Type"].split(";")[0]
+        mimetype = response.headers.get("Content-Type", "").split(";")[0]
         if mimetype.lower().endswith("xml"):
             data = xmltodict.parse(response.text)
         else:

src/python/txtai/workflow/task/storage.py

@@ -54,8 +54,8 @@ class StorageTask(Task):
         outputs = []
         for element in elements:
             if self.matches(element):
-                # Get directory listing and run actions
-                outputs.extend(super().__call__(self.list(element), executor))
+                # Unpack tagged input, list container and re-tag each listing to match input format
+                outputs.extend(self.pack(element, result) for result in super().__call__(self.list(self.upack(element, True)), executor))
             else:
                 outputs.append(element)
 

src/python/txtai/workflow/task/template.py

@@ -37,7 +37,7 @@ class TemplateTask(Task):
     def prepare(self, element):
         # Check if element matches any processing rules
         match = self.match(element)
-        if match:
+        if match is not None:
             return match
 
         # Apply template processing, if applicable
@@ -109,8 +109,7 @@ class RagTask(TemplateTask):
             params.pop("query", None)
             params["text"] = params.pop("question")
 
-            element["question"] = super().prepare(params)
-            return element
+            return {**element, "question": super().prepare(params)}
 
         # Default mode is to use element text for both query and question
         return {"query": element, "question": super().prepare(element)}

test/python/testagent.py

@@ -6,6 +6,7 @@ import os
 import tempfile
 import unittest
 
+from threading import Thread
 from unittest.mock import patch
 
 from datetime import datetime
@@ -13,7 +14,7 @@ from datetime import datetime
 from smolagents import CodeAgent, PythonInterpreterTool
 
 from txtai.agent import Agent
-from txtai.agent.tool import ToolFactory
+from txtai.agent.tool import BashTool, ToolFactory
 from txtai.embeddings import Embeddings
 
 # agents.md content
@@ -189,6 +190,77 @@ class TestAgent(unittest.TestCase):
         self.assertEqual(tool.name, "bash")
         self.assertIsNot(tool, ToolFactory.default("bash"))
 
+    def testToolsBashStdin(self):
+        """
+        Test default bash tool doesn't block reading stdin
+        """
+
+        tool = ToolFactory.default("bash")
+
+        # Replace stdin with a pipe that never reaches EOF, simulating a long-lived
+        # parent process (i.e. an API service) that keeps stdin open
+        read, write = os.pipe()
+        stdin = os.dup(0)
+
+        result = {}
+
+        try:
+            os.dup2(read, 0)
+
+            # `cat` with no file arguments reads stdin - run in a thread to avoid blocking the test suite
+            thread = Thread(target=lambda: result.update(output=tool(["cat"])), daemon=True)
+            thread.start()
+            thread.join(timeout=30)
+        finally:
+            os.dup2(stdin, 0)
+            for descriptor in (stdin, read, write):
+                os.close(descriptor)
+
+        # Command must return and not block forever reading stdin
+        self.assertFalse(thread.is_alive())
+        self.assertEqual(result["output"], "")
+
+    def testToolsBashTimeout(self):
+        """
+        Test bash tool commands are stopped once they exceed the timeout
+        """
+
+        tool = BashTool(allowed=["sleep"], timeout=1)
+
+        self.assertIn("timed out", tool(["sleep", "30"]))
+
+        # Timeout is disabled when set to None
+        self.assertIsNone(BashTool(timeout=None).timeout)
+
+    def testToolsBashAllowedEmpty(self):
+        """
+        Test an empty allowed commands list permits no commands
+        """
+
+        # An empty allow list is an explicit choice and must be honored, not replaced by the defaults
+        self.assertEqual(BashTool(allowed=[]).allowed, [])
+
+        # No command runs when nothing is allowed
+        self.assertIsNone(BashTool(allowed=[])(["cat", "file.txt"]))
+
+        # The defaults still apply when allowed isn't set, including an explicit None
+        self.assertEqual(BashTool().allowed, ["cat", "cut", "diff", "grep", "head", "ls", "tail"])
+        self.assertEqual(BashTool(allowed=None).allowed, ["cat", "cut", "diff", "grep", "head", "ls", "tail"])
+
+    @patch("subprocess.run")
+    def testToolsBashEncoding(self, run):
+        """
+        Test bash tool decodes subprocess output as UTF-8 instead of the locale encoding
+        """
+
+        run.return_value.stdout = "output"
+
+        tool = BashTool(allowed=["cat"])
+        tool(["cat", "file.txt"])
+
+        # Subprocess output must be decoded as UTF-8, not left to the locale's default encoding
+        self.assertEqual(run.call_args.kwargs.get("encoding"), "utf-8")
+
     def testToolsEmbeddings(self):
         """
         Test adding Embeddings as a tool

test/python/testann/testdense.py

@@ -1,707 +0,0 @@
-"""
-Dense ANN module tests
-"""
-
-import os
-import platform
-import sys
-import tempfile
-import unittest
-
-from unittest.mock import patch
-
-import numpy as np
-
-from sqlalchemy.dialects.postgresql import BIT
-from sqlalchemy.ext.compiler import compiles
-
-from txtai.ann import ANNFactory, ANN
-from txtai.serialize import SerializeFactory
-
-
-# pylint: disable=R0904
-class TestDense(unittest.TestCase):
-    """
-    Dense ANN tests.
-    """
-
-    def testAnnoy(self):
-        """
-        Test Annoy backend
-        """
-
-        self.runTests("annoy", None, False)
-
-    def testAnnoyCustom(self):
-        """
-        Test Annoy backend with custom settings
-        """
-
-        # Test with custom settings
-        self.runTests("annoy", {"annoy": {"ntrees": 2, "searchk": 1}}, False)
-
-    def testCustomBackend(self):
-        """
-        Test resolving a custom backend
-        """
-
-        self.runTests("txtai.ann.Faiss")
-
-    def testCustomBackendInvalid(self):
-        """
-        Test resolving an invalid backend
-        """
-
-        with self.assertRaises(ImportError):
-            ANNFactory.create({"backend": "pprint.pprint"})
-
-    def testCustomBackendNotFound(self):
-        """
-        Test resolving an unresolvable backend
-        """
-
-        with self.assertRaises(ImportError):
-            ANNFactory.create({"backend": "notfound.ann"})
-
-    def testFaiss(self):
-        """
-        Test Faiss backend
-        """
-
-        self.runTests("faiss")
-
-    def testFaissBinary(self):
-        """
-        Test Faiss backend with a binary hash index
-        """
-
-        ann = ANNFactory.create({"backend": "faiss", "quantize": 1, "dimensions": 240 * 8, "faiss": {"components": "BHash32"}})
-
-        # Generate and index dummy data
-        data = np.random.rand(100, 240).astype(np.uint8)
-        ann.index(data)
-
-        # Generate query vector and test search
-        query = np.random.rand(240).astype(np.uint8)
-        self.assertGreater(ann.search(np.array([query]), 1)[0][0][1], 0)
-
-    def testFaissCustom(self):
-        """
-        Test Faiss backend with custom settings
-        """
-
-        # Test with custom settings
-        self.runTests("faiss", {"faiss": {"nprobe": 2, "components": "PCA16,IDMap,SQ8", "sample": 1.0}}, False)
-        self.runTests("faiss", {"faiss": {"components": "IVF,SQ8"}}, False)
-
-    @patch("platform.system")
-    def testFaissMacOS(self, system):
-        """
-        Test Faiss backend with macOS
-        """
-
-        # Run test
-        system.return_value = "Darwin"
-
-        # pylint: disable=C0415, W0611
-        # Force reload of class
-        name = "txtai.ann.dense.faiss"
-        module = sys.modules[name]
-        del sys.modules[name]
-        import txtai.ann.dense.faiss
-
-        # Run tests
-        self.runTests("faiss")
-
-        # Restore original module
-        sys.modules[name] = module
-
-    @unittest.skipIf(os.name == "nt", "mmap not supported on Windows")
-    def testFaissMmap(self):
-        """
-        Test Faiss backend with mmap enabled
-        """
-
-        # Test to with mmap enabled
-        self.runTests("faiss", {"faiss": {"mmap": True}}, False)
-
-    def testFaissNprobeClamp(self):
-        """
-        Test Faiss default nprobe clamp
-        """
-
-        ann, _ = self.faissmodel(5001, {"sample": 0.01})
-        self.assertGreater(round(ann.cells(ann.count()) / 16), ann.backend.nlist)
-        self.assertEqual(ann.nprobe(), ann.backend.nlist)
-
-    def testFaissNprobeConfigured(self):
-        """
-        Test configured Faiss nprobe values
-        """
-
-        for nprobe in [1, 3]:
-            with self.subTest(nprobe=nprobe):
-                ann, _ = self.faissmodel(100, {"components": "IVF2,Flat", "nprobe": nprobe})
-                self.assertEqual(ann.nprobe(), nprobe)
-
-    def testFaissNprobeResults(self):
-        """
-        Test Faiss default nprobe clamp results
-        """
-
-        ann, rng = self.faissmodel(5001, {"sample": 0.01})
-        queries = rng.random((10, 8), dtype=np.float32)
-        self.normalize(queries)
-
-        default = ann.search(queries, 5)
-        ann.config["faiss"]["nprobe"] = ann.backend.nlist
-        full = ann.search(queries, 5)
-
-        self.assertEqual(default, full)
-
-    def testFaissNprobeSmall(self):
-        """
-        Test Faiss default nprobe clamp on a small index
-        """
-
-        ann, _ = self.faissmodel(100, {"components": "IVF2,Flat"})
-        self.assertEqual(ann.backend.nlist, 2)
-        self.assertEqual(ann.nprobe(), ann.backend.nlist)
-
-    def testFaissNprobeUnchanged(self):
-        """
-        Test Faiss default nprobe without a clamp
-        """
-
-        ann, _ = self.faissmodel(5001, {"components": "IVF16,Flat"})
-        expected = round(ann.cells(ann.count()) / 16)
-        self.assertLessEqual(expected, ann.backend.nlist)
-        self.assertEqual(ann.nprobe(), expected)
-
-    def testGGML(self):
-        """
-        Test GGML backend
-        """
-
-        self.runTests("ggml")
-
-    def testGGMLQuantization(self):
-        """
-        Test GGML backend with quantization enabled
-        """
-
-        ann = ANNFactory.create({"backend": "ggml", "ggml": {"quantize": "Q4_0"}})
-
-        # Generate and index dummy data
-        data = np.random.rand(100, 256).astype(np.float32)
-        ann.index(data)
-
-        # Test save and load
-        index = os.path.join(tempfile.gettempdir(), "ggml.q4_0.v1")
-        ann.save(index)
-        ann.load(index)
-
-        # Generate query vector and test search
-        query = np.random.rand(256).astype(np.float32)
-        self.normalize(query)
-        self.assertGreater(ann.search(np.array([query]), 1)[0][0][1], 0)
-
-        # Validate count
-        self.assertEqual(ann.count(), 100)
-
-        # Test delete
-        ann.delete([0])
-        self.assertEqual(ann.count(), 99)
-
-        # Save updated index with deletes and reload
-        index = os.path.join(tempfile.gettempdir(), "ggml.q4_0.v2")
-        ann.save(index)
-        ann.load(index)
-        ann.index(data)
-
-    def testGGMLInvalid(self):
-        """
-        Test invalid GGML configurations
-        """
-
-        data = np.random.rand(100, 240).astype(np.float32)
-
-        with self.assertRaises(ValueError):
-            ann = ANNFactory.create({"backend": "ggml", "ggml": {"quantize": "NOEXIST", "gpu": False}})
-            ann.index(data)
-
-        with self.assertRaises(ValueError):
-            ann = ANNFactory.create({"backend": "ggml", "ggml": {"quantize": "Q4_K"}})
-            ann.index(data)
-
-    def testHnsw(self):
-        """
-        Test Hnswlib backend
-        """
-
-        self.runTests("hnsw")
-
-    def testHnswCustom(self):
-        """
-        Test Hnswlib backend with custom settings
-        """
-
-        # Test with custom settings
-        self.runTests("hnsw", {"hnsw": {"efconstruction": 100, "m": 8, "randomseed": 0, "efsearch": 15}})
-
-    def testMilvus(self):
-        """
-        Test milvus-lite backend
-        """
-
-        self.runTests("milvus")
-
-    def testMilvusCustom(self):
-        """
-        Test milvus-lite backend with custom settings
-        """
-
-        self.runTests("milvus", {"milvus": {"m": 16}})
-
-        # Test invalid file path handled
-        with self.assertRaises(FileNotFoundError):
-            ann = ANNFactory.create({"backend": "milvus"})
-            ann.load("non-exist-path")
-
-    def testNotImplemented(self):
-        """
-        Test exceptions for non-implemented methods
-        """
-
-        ann = ANN({})
-
-        self.assertRaises(NotImplementedError, ann.load, None)
-        self.assertRaises(NotImplementedError, ann.index, None)
-        self.assertRaises(NotImplementedError, ann.append, None)
-        self.assertRaises(NotImplementedError, ann.delete, None)
-        self.assertRaises(NotImplementedError, ann.search, None, None)
-        self.assertRaises(NotImplementedError, ann.count)
-        self.assertRaises(NotImplementedError, ann.save, None)
-
-    def testNumPy(self):
-        """
-        Test NumPy backend
-        """
-
-        self.runTests("numpy")
-
-    @patch.dict(os.environ, {"ALLOW_PICKLE": "True"})
-    def testNumPyLegacy(self):
-        """
-        Test NumPy backend with legacy pickled data
-        """
-
-        serializer = SerializeFactory.create("pickle", allowpickle=True)
-
-        # Create output directory
-        output = os.path.join(tempfile.gettempdir(), "ann.npy")
-        path = os.path.join(output, "embeddings")
-        os.makedirs(output, exist_ok=True)
-
-        # Generate data and save as pickle
-        data = np.random.rand(100, 240).astype(np.float32)
-        serializer.save(data, path)
-
-        ann = ANNFactory.create({"backend": "numpy"})
-        ann.load(path)
-
-        # Validate count
-        self.assertEqual(ann.count(), 100)
-
-    def testNumPySafetensors(self):
-        """
-        Test NumPy backend with safetensors storage
-        """
-
-        ann = ANNFactory.create({"backend": "numpy", "numpy": {"safetensors": True}})
-
-        # Generate and index dummy data
-        data = np.random.rand(100, 240).astype(np.float32)
-        ann.index(data)
-
-        # Test save and load
-        index = os.path.join(tempfile.gettempdir(), "numpy.safetensors")
-        ann.save(index)
-        ann.load(index)
-
-        # Generate query vector and test search
-        query = np.random.rand(240).astype(np.float32)
-        self.normalize(query)
-        self.assertGreater(ann.search(np.array([query]), 1)[0][0][1], 0)
-
-        # Validate count
-        self.assertEqual(ann.count(), 100)
-
-    def testNumPyQuantizeDelete(self):
-        """
-        Test that deleted rows don't resurface in search results for quantized (hamming) NumPy indexes
-        """
-
-        ann = ANNFactory.create({"backend": "numpy", "quantize": 1, "dimensions": 4})
-
-        # Index uint8 vectors
-        data = np.array([[1, 0, 0, 0], [0, 255, 0, 0], [0, 0, 255, 0], [0, 0, 0, 255]], dtype=np.uint8)
-        ann.index(data)
-
-        # Delete the first row and search with its (former) vector
-        ann.delete([0])
-        results = ann.search(np.array([data[0]]), 4)[0]
-
-        # Deleted row must score 0 so the caller's score > 0 filter drops it, same as the dot-product path
-        self.assertEqual(dict(results).get(0), 0)
-        self.assertEqual(ann.count(), 3)
-
-    @patch("sqlalchemy.orm.Query.limit")
-    def testPGVector(self, query):
-        """
-        Test PGVector backend
-        """
-
-        # pylint: disable=W0613
-        @compiles(BIT, "sqlite")
-        def compile_bit_sqlite(type_, compiler, **kw):
-            return "BLOB"
-
-        # Generate test record
-        data = np.random.rand(1, 240).astype(np.float32)
-
-        # Mock database query
-        query.return_value = [(x, -1.0) for x in range(data.shape[0])]
-
-        configs = [
-            ("full", {"dimensions": 240}, {}, data),
-            ("half", {"dimensions": 240}, {"precision": "half"}, data),
-            ("binary", {"quantize": 1, "dimensions": 240 * 8}, {}, data.astype(np.uint8)),
-        ]
-
-        # Create ANN
-        for name, config, pgvector, data in configs:
-            path = os.path.join(tempfile.gettempdir(), f"pgvector.{name}.sqlite")
-            ann = ANNFactory.create(
-                {**{"backend": "pgvector", "pgvector": {**{"url": f"sqlite:///{path}", "schema": "txtai"}, **pgvector}}, **config}
-            )
-
-            # Test indexing
-            ann.index(data)
-            ann.append(data)
-
-            # Validate search results
-            self.assertEqual(ann.search(data, 1), [[(0, 1.0)]])
-
-            # Validate save/load/delete
-            ann.save(None)
-            ann.load(None)
-
-            # Validate count
-            self.assertEqual(ann.count(), 2)
-
-            # Test delete
-            ann.delete([0])
-            self.assertEqual(ann.count(), 1)
-
-            # Close ANN
-            ann.close()
-
-    @unittest.skipIf(platform.system() == "Darwin", "SQLite extensions not supported on macOS")
-    def testSQLite(self):
-        """
-        Test SQLite backend
-        """
-
-        self.runTests("sqlite")
-
-    @unittest.skipIf(platform.system() == "Darwin", "SQLite extensions not supported on macOS")
-    def testSQLiteCustom(self):
-        """
-        Test SQLite backend with custom settings
-        """
-
-        # Test with custom settings
-        self.runTests("sqlite", {"sqlite": {"quantize": 1}})
-        self.runTests("sqlite", {"sqlite": {"quantize": 8}})
-
-        # Test saving to a new path
-        model = self.backend("sqlite")
-        expected = model.count() - 1
-
-        # Test save variations
-        index = os.path.join(tempfile.gettempdir(), "ann.sqlite")
-        new = os.path.join(tempfile.gettempdir(), "ann.sqlite.new")
-
-        # Save new
-        model.save(index)
-
-        # Save to same path
-        model.save(index)
-
-        # Delete id
-        model.delete([0])
-
-        # Save to another path
-        model.load(index)
-        model.save(new)
-
-        self.assertEqual(model.count(), expected)
-
-    def testTorch(self):
-        """
-        Test Torch backend
-        """
-
-        self.runTests("torch")
-
-    @unittest.skipIf(platform.system() == "Darwin", "Torch quantization not supported on macOS")
-    def testTorchQuantization(self):
-        """
-        Test Torch backend with quantization enabled
-        """
-
-        for qtype in ["fp4", "nf4", "int8"]:
-            ann = ANNFactory.create({"backend": "torch", "torch": {"quantize": {"type": qtype}}})
-
-            # Generate and index dummy data
-            data = np.random.rand(100, 240).astype(np.float32)
-            ann.index(data)
-
-            # Test save and load
-            index = os.path.join(tempfile.gettempdir(), f"{qtype}.safetensors")
-            ann.save(index)
-            ann.load(index)
-
-            # Generate query vector and test search
-            query = np.random.rand(240).astype(np.float32)
-            self.normalize(query)
-            self.assertGreater(ann.search(np.array([query]), 1)[0][0][1], 0)
-
-            # Validate count
-            self.assertEqual(ann.count(), 100)
-
-            # Test delete
-            ann.delete([0])
-            self.assertEqual(ann.count(), 99)
-
-    def testTurboVec(self):
-        """
-        Test turbovec backend
-        """
-
-        self.runTests("turbovec")
-
-    def testZvec(self):
-        """
-        Test zvec backend
-        """
-
-        self.runTests("zvec")
-
-    def testZvecCustom(self):
-        """
-        Test zvec backend with custom settings
-        """
-
-        self.runTests("zvec", {"zvec": {"m": 16}})
-
-        # Test invalid file path handled
-        with self.assertRaises(FileNotFoundError):
-            ann = ANNFactory.create({"backend": "zvec"})
-            ann.load("non-exist-path")
-
-    def runTests(self, name, params=None, update=True):
-        """
-        Runs a series of standard backend tests.
-
-        Args:
-            name: backend name
-            params: additional config parameters
-            update: If append/delete options should be tested
-        """
-
-        self.assertEqual(self.backend(name, params).config["backend"], name)
-        self.assertEqual(self.save(name, params).count(), 10000)
-
-        if update:
-            self.assertEqual(self.append(name, params, 500).count(), 10500)
-            self.assertEqual(self.delete(name, params, [0, 1]).count(), 9998)
-            self.assertEqual(self.delete(name, params, [100000]).count(), 10000)
-
-        self.assertGreater(self.search(name, params), 0)
-        self.assertGreater(self.limit(name, params), 0)
-
-    def backend(self, name, params=None, length=10000):
-        """
-        Test a backend.
-
-        Args:
-            name: backend name
-            params: additional config parameters
-            length: number of rows to generate
-
-        Returns:
-            ANN model
-        """
-
-        # Generate test data
-        data = np.random.rand(length, 240).astype(np.float32)
-        self.normalize(data)
-
-        config = {"backend": name, "dimensions": data.shape[1]}
-        if params:
-            config.update(params)
-
-        model = ANNFactory.create(config)
-        model.index(data)
-
-        return model
-
-    def faissmodel(self, length, settings):
-        """
-        Builds a seeded Faiss model.
-
-        Args:
-            length: number of rows to generate
-            settings: Faiss settings
-
-        Returns:
-            Faiss model and random number generator
-        """
-
-        rng = np.random.default_rng(0)
-        data = rng.random((length, 8), dtype=np.float32)
-        self.normalize(data)
-
-        model = ANNFactory.create({"backend": "faiss", "dimensions": data.shape[1], "faiss": settings})
-        model.index(data)
-
-        return model, rng
-
-    def append(self, name, params=None, length=500):
-        """
-        Appends new data to index.
-
-        Args:
-            name: backend name
-            params: additional config parameters
-            length: number of rows to generate
-
-        Returns:
-            ANN model
-        """
-
-        # Initial model
-        model = self.backend(name, params)
-
-        # Generate test data
-        data = np.random.rand(length, 240).astype(np.float32)
-        self.normalize(data)
-
-        model.append(data)
-
-        return model
-
-    def delete(self, name, params=None, ids=None):
-        """
-        Deletes data from index.
-
-        Args:
-            name: backend name
-            params: additional config parameters
-            ids: ids to delete
-
-        Returns:
-            ANN model
-        """
-
-        # Initial model
-        model = self.backend(name, params)
-        model.delete(ids)
-
-        return model
-
-    def save(self, name, params=None):
-        """
-        Test save/load.
-
-        Args:
-            name: backend name
-            params: additional config parameters
-
-        Returns:
-            ANN model
-        """
-
-        model = self.backend(name, params)
-
-        # Generate temp file path
-        index = os.path.join(tempfile.gettempdir(), "ann")
-
-        # Save and close index
-        model.save(index)
-        model.close()
-
-        # Reload index
-        model.load(index)
-
-        return model
-
-    def search(self, name, params=None):
-        """
-        Test ANN search.
-
-        Args:
-            name: backend name
-            params: additional config parameters
-
-        Returns:
-            search results
-        """
-
-        # Generate ANN index
-        model = self.backend(name, params)
-
-        # Generate query vector
-        query = np.random.rand(240).astype(np.float32)
-        self.normalize(query)
-
-        # Ensure top result has similarity > 0
-        return model.search(np.array([query]), 1)[0][0][1]
-
-    def limit(self, name, params=None):
-        """
-        Test ANN limit search.
-
-        Args:
-            name: backend name
-            params: additional config parameters
-
-        Returns:
-            search results
-        """
-
-        # Generate ANN index
-        model = self.backend(name, params, 50)
-
-        # Generate query vector
-        query = np.random.rand(240).astype(np.float32)
-        self.normalize(query)
-
-        # Ensure limit being > count doesn't throw an error
-        return model.search(np.array([query]), 100)[0][0][1]
-
-    def normalize(self, embeddings):
-        """
-        Normalizes embeddings using L2 normalization. Operation applied directly on array.
-
-        Args:
-            embeddings: input embeddings matrix
-        """
-
-        # Calculation is different for matrices vs vectors
-        if len(embeddings.shape) > 1:
-            embeddings /= np.linalg.norm(embeddings, axis=1)[:, np.newaxis]
-        else:
-            embeddings /= np.linalg.norm(embeddings)

test/python/testann/testdense/base.py

@@ -0,0 +1,240 @@
+"""
+Base class for ANN Dense tests
+"""
+
+import os
+import tempfile
+import time
+import unittest
+
+import numpy as np
+
+from txtai.ann import ANNFactory
+
+
+class DenseTest(unittest.TestCase):
+    """
+    Base class for ANN Dense tests.
+    """
+
+    def runTests(self, name, params=None, update=True):
+        """
+        Runs a series of standard backend tests.
+
+        Args:
+            name: backend name
+            params: additional config parameters
+            update: If append/delete options should be tested
+        """
+
+        self.assertEqual(self.backend(name, params).config["backend"], name)
+        self.assertEqual(self.save(name, params).count(), 10000)
+
+        if update:
+            self.assertEqual(self.append(name, params, 500).count(), 10500)
+            self.assertEqual(self.delete(name, params, [0, 1]).count(), 9998)
+            self.assertEqual(self.delete(name, params, [100000]).count(), 10000)
+
+        self.assertGreater(self.search(name, params), 0)
+        self.assertGreater(self.limit(name, params), 0)
+
+    def backend(self, name, params=None, length=10000):
+        """
+        Test a backend.
+
+        Args:
+            name: backend name
+            params: additional config parameters
+            length: number of rows to generate
+
+        Returns:
+            ANN model
+        """
+
+        # Generate test data
+        data = np.random.rand(length, 240).astype(np.float32)
+        self.normalize(data)
+
+        config = {"backend": name, "dimensions": data.shape[1]}
+        if params:
+            config.update(params)
+
+        model = ANNFactory.create(config)
+        model.index(data)
+
+        return model
+
+    def faissmodel(self, length, settings):
+        """
+        Builds a seeded Faiss model.
+
+        Args:
+            length: number of rows to generate
+            settings: Faiss settings
+
+        Returns:
+            Faiss model and random number generator
+        """
+
+        rng = np.random.default_rng(0)
+        data = rng.random((length, 8), dtype=np.float32)
+        self.normalize(data)
+
+        model = ANNFactory.create({"backend": "faiss", "dimensions": data.shape[1], "faiss": settings})
+        model.index(data)
+
+        return model, rng
+
+    def append(self, name, params=None, length=500):
+        """
+        Appends new data to index.
+
+        Args:
+            name: backend name
+            params: additional config parameters
+            length: number of rows to generate
+
+        Returns:
+            ANN model
+        """
+
+        # Initial model
+        model = self.backend(name, params)
+
+        # Generate test data
+        data = np.random.rand(length, 240).astype(np.float32)
+        self.normalize(data)
+
+        model.append(data)
+
+        return model
+
+    def delete(self, name, params=None, ids=None):
+        """
+        Deletes data from index.
+
+        Args:
+            name: backend name
+            params: additional config parameters
+            ids: ids to delete
+
+        Returns:
+            ANN model
+        """
+
+        # Initial model
+        model = self.backend(name, params)
+        model.delete(ids)
+
+        return model
+
+    def deletenumpy(self, name, params=None):
+        """
+        Test deleting ids passed as a NumPy array. Only the requested rows should be deleted.
+
+        Args:
+            name: backend name
+            params: additional config parameters
+        """
+
+        model = self.backend(name, params, 10)
+        model.delete(np.array([3, 4]))
+
+        # Generate query vector
+        query = np.random.rand(240).astype(np.float32)
+        self.normalize(query)
+
+        self.assertEqual(model.count(), 8)
+        self.assertEqual(sorted(uid for uid, _ in model.search(np.array([query]), 10)[0]), [0, 1, 2, 5, 6, 7, 8, 9])
+
+    def save(self, name, params=None):
+        """
+        Test save/load.
+
+        Args:
+            name: backend name
+            params: additional config parameters
+
+        Returns:
+            ANN model
+        """
+
+        model = self.backend(name, params)
+
+        # Generate temp file path
+        index = os.path.join(tempfile.gettempdir(), f"ann.{name}.{round(time.time() * 1000)}")
+
+        # Generate query vector
+        query = np.random.rand(240).astype(np.float32)
+        self.normalize(query)
+
+        # Save index and ensure it's still searchable
+        model.save(index)
+        self.assertGreater(model.search(np.array([query]), 1)[0][0][1], 0)
+
+        # Close and reload index
+        model.close()
+        model.load(index)
+
+        # Ensure reloaded index is searchable
+        self.assertGreater(model.search(np.array([query]), 1)[0][0][1], 0)
+
+        return model
+
+    def search(self, name, params=None):
+        """
+        Test ANN search.
+
+        Args:
+            name: backend name
+            params: additional config parameters
+
+        Returns:
+            search results
+        """
+
+        # Generate ANN index
+        model = self.backend(name, params)
+
+        # Generate query vector
+        query = np.random.rand(240).astype(np.float32)
+        self.normalize(query)
+
+        # Ensure top result has similarity > 0
+        return model.search(np.array([query]), 1)[0][0][1]
+
+    def limit(self, name, params=None):
+        """
+        Test ANN limit search.
+
+        Args:
+            name: backend name
+            params: additional config parameters
+
+        Returns:
+            search results
+        """
+
+        # Generate ANN index
+        model = self.backend(name, params, 50)
+
+        # Generate query vector
+        query = np.random.rand(240).astype(np.float32)
+        self.normalize(query)
+
+        # Ensure limit being > count doesn't throw an error
+        return model.search(np.array([query]), 100)[0][0][1]
+
+    def normalize(self, embeddings):
+        """
+        Normalizes embeddings using L2 normalization. Operation applied directly on array.
+
+        Args:
+            embeddings: input embeddings matrix
+        """
+
+        # Calculation is different for matrices vs vectors
+        if len(embeddings.shape) > 1:
+            embeddings /= np.linalg.norm(embeddings, axis=1)[:, np.newaxis]
+        else:
+            embeddings /= np.linalg.norm(embeddings)

test/python/testann/testdense/testannoy.py

@@ -0,0 +1,26 @@
+"""
+Annoy ANN module tests
+"""
+
+from .base import DenseTest
+
+
+class TestAnnoy(DenseTest):
+    """
+    Annoy ANN tests.
+    """
+
+    def testAnnoy(self):
+        """
+        Test Annoy backend
+        """
+
+        self.runTests("annoy", None, False)
+
+    def testAnnoyCustom(self):
+        """
+        Test Annoy backend with custom settings
+        """
+
+        # Test with custom settings
+        self.runTests("annoy", {"annoy": {"ntrees": 2, "searchk": 1}}, False)

test/python/testann/testdense/testcustom.py

@@ -0,0 +1,51 @@
+"""
+Custom ANN module tests
+"""
+
+from txtai.ann import ANNFactory, ANN
+
+from .base import DenseTest
+
+
+class TestCustom(DenseTest):
+    """
+    Custom ANN tests.
+    """
+
+    def testCustomBackend(self):
+        """
+        Test resolving a custom backend
+        """
+
+        self.runTests("txtai.ann.Faiss")
+
+    def testCustomBackendInvalid(self):
+        """
+        Test resolving an invalid backend
+        """
+
+        with self.assertRaises(ImportError):
+            ANNFactory.create({"backend": "pprint.pprint"})
+
+    def testCustomBackendNotFound(self):
+        """
+        Test resolving an unresolvable backend
+        """
+
+        with self.assertRaises(ImportError):
+            ANNFactory.create({"backend": "notfound.ann"})
+
+    def testNotImplemented(self):
+        """
+        Test exceptions for non-implemented methods
+        """
+
+        ann = ANN({})
+
+        self.assertRaises(NotImplementedError, ann.load, None)
+        self.assertRaises(NotImplementedError, ann.index, None)
+        self.assertRaises(NotImplementedError, ann.append, None)
+        self.assertRaises(NotImplementedError, ann.delete, None)
+        self.assertRaises(NotImplementedError, ann.search, None, None)
+        self.assertRaises(NotImplementedError, ann.count)
+        self.assertRaises(NotImplementedError, ann.save, None)