chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,71 @@
|
||||
"""
|
||||
Client module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
import tempfile
|
||||
|
||||
from txtai.embeddings import Embeddings
|
||||
|
||||
from .testrdbms import Common
|
||||
|
||||
|
||||
# pylint: disable=R0904
|
||||
class TestClient(Common.TestRDBMS):
|
||||
"""
|
||||
Embeddings with content stored in a client RDBMS.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Initialize test data.
|
||||
"""
|
||||
|
||||
cls.data = [
|
||||
"US tops 5 million confirmed virus cases",
|
||||
"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
|
||||
"Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
|
||||
"The National Park Service warns against sacrificing slower friends in a bear attack",
|
||||
"Maine man wins $1M from $25 lottery ticket",
|
||||
"Make huge profits without work, earn up to $100,000 a day",
|
||||
]
|
||||
|
||||
# Content backend
|
||||
cls.backend = None
|
||||
|
||||
# Create embeddings model, backed by sentence-transformers & transformers
|
||||
cls.embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2"})
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""
|
||||
Cleanup data.
|
||||
"""
|
||||
|
||||
if cls.embeddings:
|
||||
cls.embeddings.close()
|
||||
|
||||
def setUp(self):
|
||||
"""
|
||||
Set unique database path for each test.
|
||||
"""
|
||||
|
||||
# Generate unique database path and set on embeddings
|
||||
path = os.path.join(tempfile.gettempdir(), f"{int(time.time() * 1000)}.sqlite")
|
||||
self.backend = f"sqlite:///{path}"
|
||||
|
||||
self.embeddings.config["content"] = self.backend
|
||||
|
||||
def testSchema(self):
|
||||
"""
|
||||
Test database creation with a specified schema
|
||||
"""
|
||||
|
||||
# Default sequence id
|
||||
embeddings = Embeddings(path="sentence-transformers/nli-mpnet-base-v2", content=self.backend, schema="txtai")
|
||||
embeddings.index(self.data)
|
||||
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
@@ -0,0 +1,29 @@
|
||||
"""
|
||||
Custom database tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from txtai.database import DatabaseFactory
|
||||
|
||||
|
||||
class TestCustom(unittest.TestCase):
|
||||
"""
|
||||
Custom database backend tests.
|
||||
"""
|
||||
|
||||
def testCustomBackend(self):
|
||||
"""
|
||||
Test resolving a custom backend
|
||||
"""
|
||||
|
||||
database = DatabaseFactory.create({"content": "txtai.database.SQLite"})
|
||||
self.assertIsNotNone(database)
|
||||
|
||||
def testCustomBackendNotFound(self):
|
||||
"""
|
||||
Test resolving an unresolvable backend
|
||||
"""
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
DatabaseFactory.create({"content": "notfound.database"})
|
||||
@@ -0,0 +1,32 @@
|
||||
"""
|
||||
Database tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from txtai.database import Database
|
||||
|
||||
|
||||
class TestDatabase(unittest.TestCase):
|
||||
"""
|
||||
Base database tests.
|
||||
"""
|
||||
|
||||
def testNotImplemented(self):
|
||||
"""
|
||||
Test exceptions for non-implemented methods
|
||||
"""
|
||||
|
||||
database = Database({})
|
||||
|
||||
self.assertRaises(NotImplementedError, database.load, None)
|
||||
self.assertRaises(NotImplementedError, database.insert, None)
|
||||
self.assertRaises(NotImplementedError, database.delete, None)
|
||||
self.assertRaises(NotImplementedError, database.reindex, None)
|
||||
self.assertRaises(NotImplementedError, database.save, None)
|
||||
self.assertRaises(NotImplementedError, database.close)
|
||||
self.assertRaises(NotImplementedError, database.ids, None)
|
||||
self.assertRaises(NotImplementedError, database.count)
|
||||
self.assertRaises(NotImplementedError, database.resolve, None, None)
|
||||
self.assertRaises(NotImplementedError, database.embed, None, None)
|
||||
self.assertRaises(NotImplementedError, database.query, None, None, None, None)
|
||||
@@ -0,0 +1,84 @@
|
||||
"""
|
||||
DuckDB module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
|
||||
from txtai.embeddings import Embeddings
|
||||
|
||||
from .testrdbms import Common
|
||||
|
||||
|
||||
# pylint: disable=R0904
|
||||
class TestDuckDB(Common.TestRDBMS):
|
||||
"""
|
||||
Embeddings with content stored in DuckDB.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Initialize test data.
|
||||
"""
|
||||
|
||||
cls.data = [
|
||||
"US tops 5 million confirmed virus cases",
|
||||
"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
|
||||
"Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
|
||||
"The National Park Service warns against sacrificing slower friends in a bear attack",
|
||||
"Maine man wins $1M from $25 lottery ticket",
|
||||
"Make huge profits without work, earn up to $100,000 a day",
|
||||
]
|
||||
|
||||
# Content backend
|
||||
cls.backend = "duckdb"
|
||||
|
||||
# Create embeddings model, backed by sentence-transformers & transformers
|
||||
cls.embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "content": cls.backend})
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""
|
||||
Cleanup data.
|
||||
"""
|
||||
|
||||
if cls.embeddings:
|
||||
cls.embeddings.close()
|
||||
|
||||
@unittest.skipIf(os.name == "nt", "testArchive skipped on Windows")
|
||||
def testArchive(self):
|
||||
"""
|
||||
Test embeddings index archiving
|
||||
"""
|
||||
|
||||
super().testArchive()
|
||||
|
||||
def testFunction(self):
|
||||
"""
|
||||
Test custom functions
|
||||
"""
|
||||
|
||||
embeddings = Embeddings(
|
||||
{
|
||||
"path": "sentence-transformers/nli-mpnet-base-v2",
|
||||
"content": self.backend,
|
||||
"functions": [{"name": "textlength", "function": "testdatabase.testduckdb.length"}],
|
||||
}
|
||||
)
|
||||
|
||||
# Create an index for the list of text
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Search for best match
|
||||
result = embeddings.search("select textlength(text) length from txtai where id = 0", 1)[0]
|
||||
|
||||
self.assertEqual(int(result["length"]), 39)
|
||||
|
||||
|
||||
def length(text):
|
||||
"""
|
||||
Custom SQL function.
|
||||
"""
|
||||
|
||||
return len(text)
|
||||
@@ -0,0 +1,155 @@
|
||||
"""
|
||||
Test encoding/decoding database objects
|
||||
"""
|
||||
|
||||
import glob
|
||||
import os
|
||||
import unittest
|
||||
import tempfile
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from io import BytesIO
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from txtai.embeddings import Embeddings
|
||||
|
||||
# pylint: disable=C0411
|
||||
from utils import Utils
|
||||
|
||||
|
||||
class TestEncoder(unittest.TestCase):
|
||||
"""
|
||||
Encoder tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Initialize test data.
|
||||
"""
|
||||
|
||||
cls.data = []
|
||||
for path in glob.glob(Utils.PATH + "/*jpg"):
|
||||
cls.data.append((path, {"object": Image.open(path)}, None))
|
||||
|
||||
# Create embeddings model, backed by sentence-transformers & transformers
|
||||
cls.embeddings = Embeddings(
|
||||
{"method": "sentence-transformers", "path": "sentence-transformers/clip-ViT-B-32", "content": True, "objects": "image"}
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""
|
||||
Cleanup data.
|
||||
"""
|
||||
|
||||
if cls.embeddings:
|
||||
cls.embeddings.close()
|
||||
|
||||
def testDefault(self):
|
||||
"""
|
||||
Test an index with default encoder
|
||||
"""
|
||||
|
||||
try:
|
||||
# Set default encoder
|
||||
self.embeddings.config["objects"] = True
|
||||
|
||||
# Test all database providers
|
||||
for content in ["duckdb", "sqlite"]:
|
||||
self.embeddings.config["content"] = content
|
||||
|
||||
data = [(0, {"object": bytearray([1, 2, 3]), "text": "default test"}, None)]
|
||||
|
||||
# Create an index
|
||||
self.embeddings.index(data)
|
||||
|
||||
result = self.embeddings.search("select object from txtai limit 1")[0]
|
||||
|
||||
self.assertEqual(result["object"].getvalue(), bytearray([1, 2, 3]))
|
||||
finally:
|
||||
self.embeddings.config["objects"] = "image"
|
||||
self.embeddings.config["content"] = True
|
||||
|
||||
def testImages(self):
|
||||
"""
|
||||
Test an index with image encoder
|
||||
"""
|
||||
|
||||
# Create an index for the list of images
|
||||
self.embeddings.index(self.data)
|
||||
|
||||
result = self.embeddings.search("select id, object from txtai where similar('universe') limit 1")[0]
|
||||
|
||||
self.assertTrue(result["id"].endswith("stars.jpg"))
|
||||
self.assertTrue(isinstance(result["object"], Image.Image))
|
||||
|
||||
@patch.dict(os.environ, {"ALLOW_PICKLE": "True"})
|
||||
def testPickle(self):
|
||||
"""
|
||||
Test an index with pickle encoder
|
||||
"""
|
||||
|
||||
try:
|
||||
# Set pickle encoder
|
||||
self.embeddings.config["objects"] = "pickle"
|
||||
data = [(0, {"object": [1, 2, 3, 4, 5], "text": "default test"}, None)]
|
||||
|
||||
# Create an index
|
||||
self.embeddings.index(data)
|
||||
|
||||
result = self.embeddings.search("select object from txtai limit 1")[0]
|
||||
|
||||
self.assertEqual(result["object"], [1, 2, 3, 4, 5])
|
||||
finally:
|
||||
self.embeddings.config["objects"] = "image"
|
||||
|
||||
def testReindex(self):
|
||||
"""
|
||||
Test reindex with objects
|
||||
"""
|
||||
|
||||
# Create an index for the list of images
|
||||
self.embeddings.index(self.data)
|
||||
|
||||
# Reindex images
|
||||
self.embeddings.reindex({"method": "sentence-transformers", "path": "sentence-transformers/clip-ViT-B-32"})
|
||||
|
||||
result = self.embeddings.search("select id, object from txtai where similar('universe') limit 1")[0]
|
||||
|
||||
self.assertTrue(result["id"].endswith("stars.jpg"))
|
||||
self.assertTrue(isinstance(result["object"], Image.Image))
|
||||
|
||||
def testReindexFunction(self):
|
||||
"""
|
||||
Test reindex with objects and a function
|
||||
"""
|
||||
|
||||
try:
|
||||
# Streaming function that loads images on the fly
|
||||
def prepare(documents):
|
||||
for uid, data, tags in documents:
|
||||
yield (uid, Image.open(data), tags)
|
||||
|
||||
# Create an index for the list of images
|
||||
self.embeddings.index(self.data)
|
||||
|
||||
# Set default encoder and use function to load images
|
||||
self.embeddings.config["objects"] = True
|
||||
|
||||
# Save and load index to force default encoder
|
||||
index = os.path.join(tempfile.gettempdir(), "objects")
|
||||
self.embeddings.save(index)
|
||||
self.embeddings.load(index)
|
||||
|
||||
# Reindex images
|
||||
self.embeddings.reindex({"method": "sentence-transformers", "path": "sentence-transformers/clip-ViT-B-32"}, function=prepare)
|
||||
|
||||
result = self.embeddings.search("select id, object from txtai where similar('universe') limit 1")[0]
|
||||
|
||||
self.assertTrue(result["id"].endswith("stars.jpg"))
|
||||
self.assertTrue(isinstance(result["object"], BytesIO))
|
||||
finally:
|
||||
self.embeddings.config["objects"] = "image"
|
||||
@@ -0,0 +1,935 @@
|
||||
"""
|
||||
Common file database module tests
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import io
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from txtai.embeddings import Embeddings, IndexNotFoundError
|
||||
from txtai.database import Embedded, RDBMS, SQLError
|
||||
|
||||
|
||||
class Common:
|
||||
"""
|
||||
Wraps common file database tests to prevent unit test discovery for this class.
|
||||
"""
|
||||
|
||||
# pylint: disable=R0904
|
||||
class TestRDBMS(unittest.TestCase):
|
||||
"""
|
||||
Embeddings with content stored in a file database tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Initialize test data.
|
||||
"""
|
||||
|
||||
cls.data = [
|
||||
"US tops 5 million confirmed virus cases",
|
||||
"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
|
||||
"Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
|
||||
"The National Park Service warns against sacrificing slower friends in a bear attack",
|
||||
"Maine man wins $1M from $25 lottery ticket",
|
||||
"Make huge profits without work, earn up to $100,000 a day",
|
||||
]
|
||||
|
||||
# Content backend
|
||||
cls.backend = None
|
||||
|
||||
# Create embeddings model, backed by sentence-transformers & transformers
|
||||
cls.embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "content": cls.backend})
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""
|
||||
Cleanup data.
|
||||
"""
|
||||
|
||||
if cls.embeddings:
|
||||
cls.embeddings.close()
|
||||
|
||||
def testArchive(self):
|
||||
"""
|
||||
Test embeddings index archiving
|
||||
"""
|
||||
|
||||
for extension in ["tar.bz2", "tar.gz", "tar.xz", "zip"]:
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), f"embeddings.{self.category()}.{extension}")
|
||||
|
||||
self.embeddings.save(index)
|
||||
self.embeddings.load(index)
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
# Test offsets still work after save/load
|
||||
self.embeddings.upsert([(0, "Looking out into the dreadful abyss", None)])
|
||||
self.assertEqual(self.embeddings.count(), len(self.data))
|
||||
|
||||
def testAutoId(self):
|
||||
"""
|
||||
Test auto id generation
|
||||
"""
|
||||
|
||||
# Default sequence id
|
||||
embeddings = Embeddings(path="sentence-transformers/nli-mpnet-base-v2", content=self.backend)
|
||||
embeddings.index(self.data)
|
||||
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
# UUID
|
||||
embeddings.config["autoid"] = "uuid4"
|
||||
embeddings.index(self.data)
|
||||
|
||||
result = embeddings.search(self.data[4], 1)[0]
|
||||
self.assertEqual(len(result["id"]), 36)
|
||||
|
||||
def testCheckpoint(self):
|
||||
"""
|
||||
Test embeddings index checkpoints
|
||||
"""
|
||||
|
||||
# Checkpoint directory
|
||||
checkpoint = os.path.join(tempfile.gettempdir(), f"embeddings.{self.category()}.checkpoint")
|
||||
|
||||
# Save embeddings checkpoint
|
||||
self.embeddings.index(self.data, checkpoint=checkpoint)
|
||||
|
||||
# Reindex with checkpoint
|
||||
self.embeddings.index(self.data, checkpoint=checkpoint)
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
def testColumns(self):
|
||||
"""
|
||||
Test custom text/object columns
|
||||
"""
|
||||
|
||||
embeddings = Embeddings({"keyword": True, "content": self.backend, "columns": {"text": "value"}})
|
||||
data = [{"value": x} for x in self.data]
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(data)])
|
||||
|
||||
# Run search
|
||||
result = embeddings.search("lottery", 1)[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
def testClose(self):
|
||||
"""
|
||||
Test embeddings close
|
||||
"""
|
||||
|
||||
embeddings = None
|
||||
|
||||
# Create index twice to test open/close and ensure resources are freed
|
||||
for _ in range(2):
|
||||
embeddings = Embeddings(
|
||||
{"path": "sentence-transformers/nli-mpnet-base-v2", "scoring": {"method": "bm25", "terms": True}, "content": self.backend}
|
||||
)
|
||||
|
||||
# Add record to index
|
||||
embeddings.index([(0, "Close test", None)])
|
||||
|
||||
# Save index
|
||||
index = os.path.join(tempfile.gettempdir(), f"embeddings.{self.category()}.close")
|
||||
embeddings.save(index)
|
||||
|
||||
# Close index
|
||||
embeddings.close()
|
||||
|
||||
# Test embeddings is empty
|
||||
self.assertIsNone(embeddings.ann)
|
||||
self.assertIsNone(embeddings.database)
|
||||
|
||||
def testData(self):
|
||||
"""
|
||||
Test content storage and retrieval
|
||||
"""
|
||||
|
||||
data = self.data + [{"date": "2021-01-01", "text": "Baby panda", "flag": 1}]
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(data)])
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], data[-1]["text"])
|
||||
|
||||
def testDelete(self):
|
||||
"""
|
||||
Test delete
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Delete best match
|
||||
self.embeddings.delete([4])
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
|
||||
self.assertEqual(self.embeddings.count(), 5)
|
||||
self.assertEqual(result["text"], self.data[5])
|
||||
|
||||
def testEmpty(self):
|
||||
"""
|
||||
Test empty index
|
||||
"""
|
||||
|
||||
# Test search against empty index
|
||||
embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "content": self.backend})
|
||||
self.assertEqual(embeddings.search("test"), [])
|
||||
|
||||
# Test index with no data
|
||||
embeddings.index([])
|
||||
self.assertIsNone(embeddings.ann)
|
||||
|
||||
# Test upsert with no data
|
||||
embeddings.index([(0, "this is a test", None)])
|
||||
embeddings.upsert([])
|
||||
self.assertIsNotNone(embeddings.ann)
|
||||
|
||||
def testEmptyString(self):
|
||||
"""
|
||||
Test empty string indexing
|
||||
"""
|
||||
|
||||
# Test empty string
|
||||
self.embeddings.index([(0, "", None)])
|
||||
self.assertTrue(self.embeddings.search("test"))
|
||||
|
||||
# Test empty string with dict
|
||||
self.embeddings.index([(0, {"text": ""}, None)])
|
||||
self.assertTrue(self.embeddings.search("test"))
|
||||
|
||||
def testExplain(self):
|
||||
"""
|
||||
Test query explain
|
||||
"""
|
||||
|
||||
# Test explain with similarity
|
||||
result = self.embeddings.explain("feel good story", self.data)[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
self.assertEqual(len(result.get("tokens")), 8)
|
||||
|
||||
def testExplainBatch(self):
|
||||
"""
|
||||
Test query explain batch
|
||||
"""
|
||||
|
||||
# Test explain with query
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
result = self.embeddings.batchexplain(["feel good story"], limit=1)[0][0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
self.assertEqual(len(result.get("tokens")), 8)
|
||||
|
||||
def testExplainEmpty(self):
|
||||
"""
|
||||
Test query explain with no filtering criteria
|
||||
"""
|
||||
|
||||
self.assertEqual(self.embeddings.explain("select * from txtai limit 1")[0]["id"], "0")
|
||||
|
||||
def testExpressions(self):
|
||||
"""
|
||||
Test expressions
|
||||
"""
|
||||
|
||||
# Test indexed expressions
|
||||
embeddings = Embeddings(
|
||||
path="sentence-transformers/nli-mpnet-base-v2",
|
||||
content=self.backend,
|
||||
expressions=[{"name": "textlength", "expression": "length(text)", "index": True}],
|
||||
)
|
||||
embeddings.index(self.data)
|
||||
|
||||
result = embeddings.search("SELECT textlength FROM txtai WHERE id = 0", 1)[0]
|
||||
self.assertEqual(result["textlength"], len(self.data[0]))
|
||||
|
||||
def testGenerator(self):
|
||||
"""
|
||||
Test index with a generator
|
||||
"""
|
||||
|
||||
def documents():
|
||||
for uid, text in enumerate(self.data):
|
||||
yield (uid, text, None)
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index(documents())
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
def testHybrid(self):
|
||||
"""
|
||||
Test hybrid search
|
||||
"""
|
||||
|
||||
# Build data array
|
||||
data = [(uid, text, None) for uid, text in enumerate(self.data)]
|
||||
|
||||
# Index data with sparse + dense vectors.
|
||||
embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "hybrid": True, "content": self.backend})
|
||||
embeddings.index(data)
|
||||
|
||||
# Run search
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], data[4][1])
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), f"embeddings.{self.category()}.hybrid")
|
||||
|
||||
# Test load/save
|
||||
embeddings.save(index)
|
||||
embeddings.load(index)
|
||||
|
||||
# Run search
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], data[4][1])
|
||||
|
||||
# Index data with sparse + dense vectors and unnormalized scores.
|
||||
embeddings.config["scoring"]["normalize"] = False
|
||||
embeddings.index(data)
|
||||
|
||||
# Run search
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], data[4][1])
|
||||
|
||||
# Index data with sparse + dense vectors and bb25 normalized scores
|
||||
embeddings.config["scoring"]["normalize"] = "bb25"
|
||||
embeddings.index(data)
|
||||
|
||||
# Run search
|
||||
result = embeddings.search("canada intact iceberg a", 1)[0]
|
||||
self.assertEqual(result["text"], data[1][1])
|
||||
|
||||
# Test upsert
|
||||
data[0] = (0, "Feel good story: baby panda born", None)
|
||||
embeddings.upsert([data[0]])
|
||||
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], data[0][1])
|
||||
|
||||
def testIndex(self):
|
||||
"""
|
||||
Test index
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
def testIndexTokens(self):
|
||||
"""
|
||||
Test index with tokens
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text.split(), None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
def testInfo(self):
|
||||
"""
|
||||
Test info
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
output = io.StringIO()
|
||||
with contextlib.redirect_stdout(output):
|
||||
self.embeddings.info()
|
||||
|
||||
self.assertIn("txtai", output.getvalue())
|
||||
|
||||
def testInstructions(self):
|
||||
"""
|
||||
Test indexing with instruction prefixes.
|
||||
"""
|
||||
|
||||
embeddings = Embeddings(
|
||||
{
|
||||
"path": "sentence-transformers/nli-mpnet-base-v2",
|
||||
"content": self.backend,
|
||||
"instructions": {"query": "query: ", "data": "passage: "},
|
||||
}
|
||||
)
|
||||
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Search for best match
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
def testInvalidData(self):
|
||||
"""
|
||||
Test invalid JSON data
|
||||
"""
|
||||
|
||||
# Test invalid JSON value
|
||||
with self.assertRaises(ValueError):
|
||||
self.embeddings.index([(0, {"text": "This is a test", "flag": float("NaN")}, None)])
|
||||
|
||||
def testKeyword(self):
|
||||
"""
|
||||
Test keyword only (sparse) search
|
||||
"""
|
||||
|
||||
# Build data array
|
||||
data = [(uid, text, None) for uid, text in enumerate(self.data)]
|
||||
|
||||
# Index data with sparse keyword vectors
|
||||
embeddings = Embeddings({"keyword": True, "content": self.backend})
|
||||
embeddings.index(data)
|
||||
|
||||
# Run search
|
||||
result = embeddings.search("lottery ticket", 1)[0]
|
||||
self.assertEqual(result["text"], data[4][1])
|
||||
|
||||
# Test count method
|
||||
self.assertEqual(embeddings.count(), len(data))
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), f"embeddings.{self.category()}.keyword")
|
||||
|
||||
# Test load/save
|
||||
embeddings.save(index)
|
||||
embeddings.load(index)
|
||||
|
||||
# Run search
|
||||
result = embeddings.search("lottery ticket", 1)[0]
|
||||
self.assertEqual(result["text"], data[4][1])
|
||||
|
||||
# Update data
|
||||
data[0] = (0, "Feel good story: baby panda born", None)
|
||||
embeddings.upsert([data[0]])
|
||||
|
||||
# Search for best match
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], data[0][1])
|
||||
|
||||
def testMultiData(self):
|
||||
"""
|
||||
Test indexing with multiple data types (text, documents)
|
||||
"""
|
||||
|
||||
embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "content": self.backend, "batch": len(self.data)})
|
||||
|
||||
# Create an index using mixed data (text and documents)
|
||||
data = []
|
||||
for uid, text in enumerate(self.data):
|
||||
data.append((uid, text, None))
|
||||
data.append((uid, {"content": text}, None))
|
||||
|
||||
embeddings.index(data)
|
||||
|
||||
# Search for best match
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
def testMultiSave(self):
|
||||
"""
|
||||
Test multiple successive saves
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Save original index
|
||||
index = os.path.join(tempfile.gettempdir(), f"embeddings.{self.category()}.insert")
|
||||
self.embeddings.save(index)
|
||||
|
||||
# Modify index
|
||||
self.embeddings.upsert([(0, "Looking out into the dreadful abyss", None)])
|
||||
|
||||
# Save to a different location
|
||||
indexupdate = os.path.join(tempfile.gettempdir(), f"embeddings.{self.category()}.update")
|
||||
self.embeddings.save(indexupdate)
|
||||
|
||||
# Save to same location
|
||||
self.embeddings.save(index)
|
||||
|
||||
# Test all indexes match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
self.embeddings.load(index)
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
self.embeddings.load(indexupdate)
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
def testNoIndex(self):
|
||||
"""
|
||||
Test an embeddings instance with no available indexes
|
||||
"""
|
||||
|
||||
# Disable top-level indexing
|
||||
embeddings = Embeddings(
|
||||
{
|
||||
"content": self.backend,
|
||||
"defaults": False,
|
||||
}
|
||||
)
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
with self.assertRaises(IndexNotFoundError):
|
||||
embeddings.search("select id, text, score from txtai where similar('feel good story')")
|
||||
|
||||
def testNotImplemented(self):
|
||||
"""
|
||||
Test exceptions for non-implemented methods
|
||||
"""
|
||||
|
||||
db = RDBMS({})
|
||||
|
||||
self.assertRaises(NotImplementedError, db.connect, None)
|
||||
self.assertRaises(NotImplementedError, db.getcursor)
|
||||
self.assertRaises(NotImplementedError, db.jsonprefix)
|
||||
self.assertRaises(NotImplementedError, db.jsoncolumn, None)
|
||||
self.assertRaises(NotImplementedError, db.rows)
|
||||
self.assertRaises(NotImplementedError, db.addfunctions)
|
||||
|
||||
db = Embedded({})
|
||||
self.assertRaises(NotImplementedError, db.copy, None)
|
||||
|
||||
def testObject(self):
|
||||
"""
|
||||
Test object field
|
||||
"""
|
||||
|
||||
# Encode object
|
||||
embeddings = Embeddings({"defaults": False, "content": self.backend, "objects": True})
|
||||
embeddings.index([{"object": "binary data".encode("utf-8")}])
|
||||
|
||||
# Decode and test extracted object
|
||||
obj = embeddings.search("select object from txtai where id = 0")[0]["object"]
|
||||
self.assertEqual(str(obj.getvalue(), "utf-8"), "binary data")
|
||||
|
||||
@patch.dict(os.environ, {"ALLOW_PICKLE": "True"})
|
||||
def testPickle(self):
|
||||
"""
|
||||
Test pickle configuration
|
||||
"""
|
||||
|
||||
embeddings = Embeddings(
|
||||
{
|
||||
"format": "pickle",
|
||||
"path": "sentence-transformers/nli-mpnet-base-v2",
|
||||
"content": self.backend,
|
||||
}
|
||||
)
|
||||
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), f"embeddings.{self.category()}.pickle")
|
||||
|
||||
embeddings.save(index)
|
||||
|
||||
# Check that config exists
|
||||
self.assertTrue(os.path.exists(os.path.join(index, "config")))
|
||||
|
||||
# Check that index can be reloaded
|
||||
embeddings.load(index)
|
||||
self.assertEqual(embeddings.count(), 6)
|
||||
|
||||
def testQuantize(self):
|
||||
"""
|
||||
Test scalar quantization
|
||||
"""
|
||||
|
||||
# Index data with 1-bit scalar quantization
|
||||
embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "quantize": 1, "content": self.backend})
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
def testQueryModel(self):
|
||||
"""
|
||||
Test index
|
||||
"""
|
||||
|
||||
embeddings = Embeddings(
|
||||
{"path": "sentence-transformers/nli-mpnet-base-v2", "content": self.backend, "query": {"path": "neuml/t5-small-txtsql"}}
|
||||
)
|
||||
|
||||
# Create an index for the list of text
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Search for best match
|
||||
result = embeddings.search("feel good story with win in text", 1)[0]
|
||||
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
def testReindex(self):
|
||||
"""
|
||||
Test reindex
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Delete records to test indexids still match
|
||||
self.embeddings.delete(([0, 1]))
|
||||
|
||||
# Reindex
|
||||
self.embeddings.reindex({"path": "sentence-transformers/nli-mpnet-base-v2"})
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
def testSave(self):
|
||||
"""
|
||||
Test save
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), f"embeddings.{self.category()}")
|
||||
|
||||
self.embeddings.save(index)
|
||||
self.embeddings.load(index)
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
# Test offsets still work after save/load
|
||||
self.embeddings.upsert([(0, "Looking out into the dreadful abyss", None)])
|
||||
self.assertEqual(self.embeddings.count(), len(self.data))
|
||||
|
||||
def testSettings(self):
|
||||
"""
|
||||
Test custom SQLite settings
|
||||
"""
|
||||
|
||||
# Index with write-ahead logging enabled
|
||||
embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "content": self.backend, "sqlite": {"wal": True}})
|
||||
|
||||
# Create an index for the list of text
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Search for best match
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
def testSQL(self):
|
||||
"""
|
||||
Test running a SQL query
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, {"text": text, "length": len(text), "attribute": f"ID{uid}"}, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Test similar
|
||||
result = self.embeddings.search(
|
||||
"select text, score from txtai where similar('feel good story') group by text, score having count(*) > 0 order by score desc", 1
|
||||
)[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
# Test similar with limits
|
||||
result = self.embeddings.search("select * from txtai where similar('feel good story', 1) limit 1")[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
# Test similar with offset
|
||||
result = self.embeddings.search("select * from txtai where similar('feel good story') offset 1")[0]
|
||||
self.assertEqual(result["text"], self.data[5])
|
||||
|
||||
# Test where
|
||||
result = self.embeddings.search("select * from txtai where text like '%iceberg%'", 1)[0]
|
||||
self.assertEqual(result["text"], self.data[1])
|
||||
|
||||
# Test count
|
||||
result = self.embeddings.search("select count(*) from txtai")[0]
|
||||
self.assertEqual(list(result.values())[0], len(self.data))
|
||||
|
||||
# Test columns
|
||||
result = self.embeddings.search("select id, text, length, data, entry from txtai")[0]
|
||||
self.assertEqual(sorted(result.keys()), ["data", "entry", "id", "length", "text"])
|
||||
|
||||
# Test column filtering
|
||||
result = self.embeddings.search("select text from txtai where attribute = 'ID4'", 1)[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
# Test SQL parse error
|
||||
with self.assertRaises(SQLError):
|
||||
self.embeddings.search("select * from txtai where bad,query")
|
||||
|
||||
def testSQLBind(self):
|
||||
"""
|
||||
Test SQL statements with bind parameters
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Test similar clause bind parameters
|
||||
result = self.embeddings.search("select id, text, score from txtai where similar(:x)", parameters={"x": "feel good story"})[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
# Test similar clause bind and non-bind parameters
|
||||
result = self.embeddings.search("select id, text, score from txtai where similar(:x, 0.5)", parameters={"x": "feel good story"})[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
# Test where filtering with bind parameters
|
||||
result = self.embeddings.search("select * from txtai where text like :x", parameters={"x": "%iceberg%"})[0]
|
||||
self.assertEqual(result["text"], self.data[1])
|
||||
|
||||
def testSparse(self):
|
||||
"""
|
||||
Test sparse vector search
|
||||
"""
|
||||
|
||||
# Build data array
|
||||
data = [(uid, text, None) for uid, text in enumerate(self.data)]
|
||||
|
||||
# Index data with sparse vectors
|
||||
embeddings = Embeddings({"sparse": "sparse-encoder-testing/splade-bert-tiny-nq", "content": self.backend})
|
||||
embeddings.index(data)
|
||||
|
||||
# Run search
|
||||
result = embeddings.search("lottery ticket", 1)[0]
|
||||
self.assertEqual(result["text"], data[4][1])
|
||||
|
||||
# Test count method
|
||||
self.assertEqual(embeddings.count(), len(data))
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), f"embeddings.{self.category()}.sparse")
|
||||
|
||||
# Test load/save
|
||||
embeddings.save(index)
|
||||
embeddings.load(index)
|
||||
|
||||
# Run search
|
||||
result = embeddings.search("lottery ticket", 1)[0]
|
||||
self.assertEqual(result["text"], data[4][1])
|
||||
|
||||
# Update data
|
||||
data[0] = (0, "Feel good story: baby panda born", None)
|
||||
embeddings.upsert([data[0]])
|
||||
|
||||
# Search for best match
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], data[0][1])
|
||||
|
||||
def testSubindex(self):
|
||||
"""
|
||||
Test subindex
|
||||
"""
|
||||
|
||||
# Build data array
|
||||
data = [(uid, text, None) for uid, text in enumerate(self.data)]
|
||||
|
||||
# Disable top-level indexing and create subindex
|
||||
embeddings = Embeddings(
|
||||
{"content": self.backend, "defaults": False, "indexes": {"index1": {"path": "sentence-transformers/nli-mpnet-base-v2"}}}
|
||||
)
|
||||
embeddings.index(data)
|
||||
|
||||
# Test transform
|
||||
self.assertEqual(embeddings.transform("feel good story").shape, (768,))
|
||||
|
||||
# Run search
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], data[4][1])
|
||||
|
||||
# Run SQL search
|
||||
result = embeddings.search("select id, text, score from txtai where similar('feel good story', 10, 0.5)")[0]
|
||||
self.assertEqual(result["text"], data[4][1])
|
||||
|
||||
# Test missing index
|
||||
with self.assertRaises(IndexNotFoundError):
|
||||
embeddings.search("select id, text, score from txtai where similar('feel good story', 'notindex')")
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), f"embeddings.{self.category()}.subindex")
|
||||
|
||||
# Test load/save
|
||||
embeddings.save(index)
|
||||
embeddings.load(index)
|
||||
|
||||
# Run search
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], data[4][1])
|
||||
|
||||
# Update data
|
||||
data[0] = (0, "Feel good story: baby panda born", None)
|
||||
embeddings.upsert([data[0]])
|
||||
|
||||
# Search for best match
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], data[0][1])
|
||||
|
||||
# Check missing text is set to id when top-level indexing is disabled
|
||||
embeddings.upsert([(embeddings.count(), {"content": "empty text"}, None)])
|
||||
result = embeddings.search(f"{embeddings.count() - 1}", 1)[0]
|
||||
self.assertEqual(result["text"], str(embeddings.count() - 1))
|
||||
|
||||
# Close embeddings
|
||||
embeddings.close()
|
||||
|
||||
def testSubindexEmpty(self):
|
||||
"""
|
||||
Test loading an empty subindex
|
||||
"""
|
||||
|
||||
# Build data array
|
||||
data = [(uid, {"column1": text}, None) for uid, text in enumerate(self.data)]
|
||||
|
||||
# Disable top-level indexing and create subindexes
|
||||
embeddings = Embeddings(
|
||||
{
|
||||
"content": self.backend,
|
||||
"defaults": False,
|
||||
"indexes": {
|
||||
"index1": {"path": "sentence-transformers/nli-mpnet-base-v2", "columns": {"text": "column1"}},
|
||||
"index2": {"path": "sentence-transformers/nli-mpnet-base-v2", "columns": {"text": "column2"}},
|
||||
},
|
||||
}
|
||||
)
|
||||
embeddings.index(data)
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), f"embeddings.{self.category()}.subindexempty")
|
||||
|
||||
# Save index
|
||||
embeddings.save(index)
|
||||
|
||||
# Test exists
|
||||
self.assertTrue(embeddings.exists(index))
|
||||
|
||||
# Load index
|
||||
embeddings.load(index)
|
||||
|
||||
# Test search
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], data[4][1]["text"])
|
||||
|
||||
def testTerms(self):
|
||||
"""
|
||||
Test extracting keyword terms from query
|
||||
"""
|
||||
|
||||
result = self.embeddings.terms("select * from txtai where similar('keyword terms')")
|
||||
self.assertEqual(result, "keyword terms")
|
||||
|
||||
def testTruncate(self):
|
||||
"""
|
||||
Test dimensionality truncation
|
||||
"""
|
||||
|
||||
# Truncate vectors to a specified number of dimensions
|
||||
embeddings = Embeddings(
|
||||
{"path": "sentence-transformers/nli-mpnet-base-v2", "dimensionality": 750, "content": self.backend, "vectors": {"revision": "main"}}
|
||||
)
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
def testUpsert(self):
|
||||
"""
|
||||
Test upsert
|
||||
"""
|
||||
|
||||
# Build data array
|
||||
data = [(uid, text, None) for uid, text in enumerate(self.data)]
|
||||
|
||||
# Reset embeddings for test
|
||||
self.embeddings.ann = None
|
||||
self.embeddings.database = None
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.upsert(data)
|
||||
|
||||
# Update data
|
||||
data[0] = (0, "Feel good story: baby panda born", None)
|
||||
self.embeddings.upsert([data[0]])
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], data[0][1])
|
||||
|
||||
def testUpsertBatch(self):
|
||||
"""
|
||||
Test upsert batch
|
||||
"""
|
||||
|
||||
try:
|
||||
# Build data array
|
||||
data = [(uid, text, None) for uid, text in enumerate(self.data)]
|
||||
|
||||
# Reset embeddings for test
|
||||
self.embeddings.ann = None
|
||||
self.embeddings.database = None
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.upsert(data)
|
||||
|
||||
# Set batch size to 1
|
||||
self.embeddings.config["batch"] = 1
|
||||
|
||||
# Update data
|
||||
data[0] = (0, "Feel good story: baby panda born", None)
|
||||
data[1] = (0, "Not good news", None)
|
||||
self.embeddings.upsert([data[0], data[1]])
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
|
||||
self.assertEqual(result["text"], data[0][1])
|
||||
finally:
|
||||
del self.embeddings.config["batch"]
|
||||
|
||||
def category(self):
|
||||
"""
|
||||
Content backend category.
|
||||
|
||||
Returns:
|
||||
category
|
||||
"""
|
||||
|
||||
return self.__class__.__name__.lower().replace("test", "")
|
||||
@@ -0,0 +1,325 @@
|
||||
"""
|
||||
SQL module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from txtai.database import DatabaseFactory, SQL, SQLError
|
||||
|
||||
|
||||
class TestSQL(unittest.TestCase):
|
||||
"""
|
||||
Test SQL parsing and generation.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Initialize test data.
|
||||
"""
|
||||
|
||||
# Create SQL parser for SQLite
|
||||
cls.db = DatabaseFactory.create({"content": True})
|
||||
cls.db.initialize()
|
||||
|
||||
cls.sql = SQL(cls.db)
|
||||
|
||||
def testAlias(self):
|
||||
"""
|
||||
Test alias clauses
|
||||
"""
|
||||
|
||||
self.assertSql("select", "select a as a1 from txtai", "json_extract(data, '$.a') as a1")
|
||||
self.assertSql("select", "select a 'a1' from txtai", "json_extract(data, '$.a') 'a1'")
|
||||
self.assertSql("select", 'select a "a1" from txtai', "json_extract(data, '$.a') \"a1\"")
|
||||
self.assertSql("select", "select a a1 from txtai", "json_extract(data, '$.a') a1")
|
||||
self.assertSql(
|
||||
"select",
|
||||
"select a, b as b1, c, d + 1 as 'd1' from txtai",
|
||||
"json_extract(data, '$.a') as \"a\", json_extract(data, '$.b') as b1, "
|
||||
+ "json_extract(data, '$.c') as \"c\", json_extract(data, '$.d') + 1 as 'd1'",
|
||||
)
|
||||
self.assertSql("select", "select id as myid from txtai", "s.id as myid")
|
||||
self.assertSql("select", "select length(a) t from txtai", "length(json_extract(data, '$.a')) t")
|
||||
|
||||
self.assertSql("where", "select id as myid from txtai where myid != 3 and a != 1", "myid != 3 and json_extract(data, '$.a') != 1")
|
||||
self.assertSql("where", "select txt T from txtai where t LIKE '%abc%'", "t LIKE '%abc%'")
|
||||
self.assertSql("where", "select txt 'T' from txtai where t LIKE '%abc%'", "t LIKE '%abc%'")
|
||||
self.assertSql("where", "select txt \"T\" from txtai where t LIKE '%abc%'", "t LIKE '%abc%'")
|
||||
self.assertSql("where", "select txt as T from txtai where t LIKE '%abc%'", "t LIKE '%abc%'")
|
||||
self.assertSql("where", "select txt as 'T' from txtai where t LIKE '%abc%'", "t LIKE '%abc%'")
|
||||
self.assertSql("where", "select txt as \"T\" from txtai where t LIKE '%abc%'", "t LIKE '%abc%'")
|
||||
|
||||
self.assertSql("groupby", "select id as myid, count(*) from txtai group by myid, a", "myid, json_extract(data, '$.a')")
|
||||
self.assertSql("orderby", "select id as myid from txtai order by myid, a", "myid, json_extract(data, '$.a')")
|
||||
|
||||
def testBadSQL(self):
|
||||
"""
|
||||
Test invalid SQL
|
||||
"""
|
||||
|
||||
with self.assertRaises(SQLError):
|
||||
self.db.search("select * from txtai where order by")
|
||||
|
||||
with self.assertRaises(SQLError):
|
||||
self.db.search("select * from txtai where groupby order by")
|
||||
|
||||
with self.assertRaises(SQLError):
|
||||
self.db.search("select * from txtai where a(1)")
|
||||
|
||||
with self.assertRaises(SQLError):
|
||||
self.db.search("select a b c from txtai where id match id")
|
||||
|
||||
def testBracket(self):
|
||||
"""
|
||||
Test bracket expressions
|
||||
"""
|
||||
|
||||
self.assertSql("select", "select [a] from txtai", "json_extract(data, '$.a') as \"a\"")
|
||||
self.assertSql("select", "select [a] ab from txtai", "json_extract(data, '$.a') ab")
|
||||
self.assertSql("select", "select [abc] from txtai", "json_extract(data, '$.abc') as \"abc\"")
|
||||
self.assertSql("select", "select [id], text, score from txtai", "s.id, text, score")
|
||||
self.assertSql("select", "select [ab cd], text, score from txtai", "json_extract(data, '$.ab cd') as \"ab cd\", text, score")
|
||||
self.assertSql("select", "select [a[0]] from txtai", "json_extract(data, '$.a[0]') as \"a[0]\"")
|
||||
self.assertSql("select", "select [a[0].ab] from txtai", "json_extract(data, '$.a[0].ab') as \"a[0].ab\"")
|
||||
self.assertSql("select", "select [a[0].c[0]] from txtai", "json_extract(data, '$.a[0].c[0]') as \"a[0].c[0]\"")
|
||||
self.assertSql("select", "select avg([a]) from txtai", "avg(json_extract(data, '$.a')) as \"avg([a])\"")
|
||||
|
||||
# Test single quote escaping in bracket expressions
|
||||
self.assertSql("select", "select [field'] from txtai", "json_extract(data, '$.field''') as \"field'\"")
|
||||
|
||||
self.assertSql("where", "select * from txtai where [a b] < 1 or a > 1", "json_extract(data, '$.a b') < 1 or json_extract(data, '$.a') > 1")
|
||||
self.assertSql("where", "select [a[0].c[0]] a from txtai where a < 1", "a < 1")
|
||||
self.assertSql("groupby", "select * from txtai group by [a]", "json_extract(data, '$.a')")
|
||||
self.assertSql("orderby", "select * from txtai where order by [a]", "json_extract(data, '$.a')")
|
||||
|
||||
def testDistinct(self):
|
||||
"""
|
||||
Test distinct expressions
|
||||
"""
|
||||
|
||||
# Attributes
|
||||
self.assertSql("select", "select distinct id from txtai", "distinct s.id")
|
||||
self.assertSql("select", "select distinct id as myid from txtai", "distinct s.id as myid")
|
||||
self.assertSql("select", "select distinct a from txtai", "distinct json_extract(data, '$.a') as \"a\"")
|
||||
self.assertSql("select", "select distinct a.b from txtai", "distinct json_extract(data, '$.a.b') as \"a.b\"")
|
||||
|
||||
# Bracket expression
|
||||
self.assertSql("select", "select distinct [ab cd] from txtai", "distinct json_extract(data, '$.ab cd') as \"distinct[ab cd]\"")
|
||||
|
||||
# Function expression
|
||||
self.assertSql("select", "select distinct(id) from txtai", 'distinct(s.id) as "distinct(id)"')
|
||||
self.assertSql("select", "select count(distinct id) from txtai", 'count(distinct s.id) as "count(distinct id)"')
|
||||
self.assertSql("select", "select count(distinct a) from txtai", "count(distinct json_extract(data, '$.a')) as \"count(distinct a)\"")
|
||||
self.assertSql("select", "select count(distinct avg(id)) from txtai", 'count(distinct avg(s.id)) as "count(distinct avg(id))"')
|
||||
self.assertSql(
|
||||
"select", "select count(distinct avg(a)) from txtai", "count(distinct avg(json_extract(data, '$.a'))) as \"count(distinct avg(a))\""
|
||||
)
|
||||
|
||||
# Compound expression
|
||||
self.assertSql("select", "select distinct a/1 from txtai", "distinct json_extract(data, '$.a') / 1 as \"a / 1\"")
|
||||
self.assertSql("select", "select distinct(a/1) from txtai", "distinct(json_extract(data, '$.a') / 1) as \"distinct(a / 1)\"")
|
||||
|
||||
def testGroupby(self):
|
||||
"""
|
||||
Test group by clauses
|
||||
"""
|
||||
|
||||
prefix = "select count(*), flag from txtai "
|
||||
|
||||
self.assertSql("groupby", prefix + "group by text", "text")
|
||||
self.assertSql("groupby", prefix + "group by distinct(a)", "distinct(json_extract(data, '$.a'))")
|
||||
self.assertSql("groupby", prefix + "where a > 1 group by text", "text")
|
||||
|
||||
def testHaving(self):
|
||||
"""
|
||||
Test having clauses
|
||||
"""
|
||||
|
||||
prefix = "select count(*), flag from txtai "
|
||||
|
||||
self.assertSql("having", prefix + "group by text having count(*) > 1", "count(*) > 1")
|
||||
self.assertSql("having", prefix + "where flag = 1 group by text having count(*) > 1", "count(*) > 1")
|
||||
|
||||
def testIsSQL(self):
|
||||
"""
|
||||
Test SQL detection method.
|
||||
"""
|
||||
|
||||
self.assertTrue(self.sql.issql("select text from txtai where id = 1"))
|
||||
self.assertFalse(self.sql.issql(1234))
|
||||
|
||||
def testLimit(self):
|
||||
"""
|
||||
Test limit clauses
|
||||
"""
|
||||
|
||||
prefix = "select count(*) from txtai "
|
||||
|
||||
self.assertSql("limit", prefix + "limit 100", "100")
|
||||
|
||||
def testOffset(self):
|
||||
"""
|
||||
Test offset clauses
|
||||
"""
|
||||
|
||||
prefix = "select count(*) from txtai "
|
||||
|
||||
self.assertSql("offset", prefix + "limit 100 offset 50", "50")
|
||||
self.assertSql("offset", prefix + "offset 50", "50")
|
||||
|
||||
def testOrderby(self):
|
||||
"""
|
||||
Test order by clauses
|
||||
"""
|
||||
|
||||
prefix = "select * from txtai "
|
||||
|
||||
self.assertSql("orderby", prefix + "order by id", "s.id")
|
||||
self.assertSql("orderby", prefix + "order by id, text", "s.id, text")
|
||||
self.assertSql("orderby", prefix + "order by id asc", "s.id asc")
|
||||
self.assertSql("orderby", prefix + "order by id desc", "s.id desc")
|
||||
self.assertSql("orderby", prefix + "order by id asc, text desc", "s.id asc, text desc")
|
||||
|
||||
def testSelectBasic(self):
|
||||
"""
|
||||
Test basic select clauses
|
||||
"""
|
||||
|
||||
self.assertSql("select", "select id, indexid, tags from txtai", "s.id, s.indexid, s.tags")
|
||||
self.assertSql("select", "select id, indexid, flag from txtai", "s.id, s.indexid, json_extract(data, '$.flag') as \"flag\"")
|
||||
self.assertSql("select", "select id, indexid, a.b.c from txtai", "s.id, s.indexid, json_extract(data, '$.a.b.c') as \"a.b.c\"")
|
||||
self.assertSql("select", "select 'id', [id], (id) from txtai", "'id', s.id, (s.id)")
|
||||
self.assertSql("select", "select * from txtai", "*")
|
||||
|
||||
def testSelectCompound(self):
|
||||
"""
|
||||
Test compound select clauses
|
||||
"""
|
||||
|
||||
self.assertSql("select", "select a + 1 from txtai", "json_extract(data, '$.a') + 1 as \"a + 1\"")
|
||||
self.assertSql("select", "select 1 * a from txtai", "1 * json_extract(data, '$.a') as \"1 * a\"")
|
||||
self.assertSql("select", "select a/1 from txtai", "json_extract(data, '$.a') / 1 as \"a / 1\"")
|
||||
self.assertSql("select", "select avg(a-b) from txtai", "avg(json_extract(data, '$.a') - json_extract(data, '$.b')) as \"avg(a - b)\"")
|
||||
self.assertSql("select", "select distinct(text) from txtai", "distinct(text)")
|
||||
self.assertSql("select", "select id, score, (a/2)*3 from txtai", "s.id, score, (json_extract(data, '$.a') / 2) * 3 as \"(a / 2) * 3\"")
|
||||
self.assertSql("select", "select id, score, (a/2*3) from txtai", "s.id, score, (json_extract(data, '$.a') / 2 * 3) as \"(a / 2 * 3)\"")
|
||||
self.assertSql(
|
||||
"select",
|
||||
"select func(func2(indexid + 1), a) from txtai",
|
||||
"func(func2(s.indexid + 1), json_extract(data, '$.a')) as \"func(func2(indexid + 1), a)\"",
|
||||
)
|
||||
self.assertSql("select", "select func(func2(indexid + 1), a) a from txtai", "func(func2(s.indexid + 1), json_extract(data, '$.a')) a")
|
||||
self.assertSql("select", "select 'prefix' || id from txtai", "'prefix' || s.id as \"'prefix' || id\"")
|
||||
self.assertSql("select", "select 'prefix' || id id from txtai", "'prefix' || s.id id")
|
||||
self.assertSql("select", "select 'prefix' || a a from txtai", "'prefix' || json_extract(data, '$.a') a")
|
||||
|
||||
def testSimilar(self):
|
||||
"""
|
||||
Test similar functions
|
||||
"""
|
||||
|
||||
prefix = "select * from txtai "
|
||||
|
||||
self.assertSql("where", prefix + "where similar('abc')", "__SIMILAR__0")
|
||||
self.assertSql("similar", prefix + "where similar('abc')", [["abc"]])
|
||||
|
||||
self.assertSql("where", prefix + "where similar('abc') AND id = 1", "__SIMILAR__0 AND s.id = 1")
|
||||
self.assertSql("similar", prefix + "where similar('abc')", [["abc"]])
|
||||
|
||||
self.assertSql("where", prefix + "where similar('abc') and similar('def')", "__SIMILAR__0 and __SIMILAR__1")
|
||||
self.assertSql("similar", prefix + "where similar('abc') and similar('def')", [["abc"], ["def"]])
|
||||
|
||||
self.assertSql("where", prefix + "where similar('abc', 1000)", "__SIMILAR__0")
|
||||
self.assertSql("similar", prefix + "where similar('abc', 1000)", [["abc", "1000"]])
|
||||
|
||||
self.assertSql("where", prefix + "where similar('abc', 1000) and similar('def', 10)", "__SIMILAR__0 and __SIMILAR__1")
|
||||
self.assertSql("similar", prefix + "where similar('abc', 1000) and similar('def', 10)", [["abc", "1000"], ["def", "10"]])
|
||||
|
||||
self.assertSql("where", prefix + "where coalesce(similar('abc'), similar('abc'))", "coalesce(__SIMILAR__0, __SIMILAR__1)")
|
||||
self.assertSql("similar", prefix + "where coalesce(similar('abc'), similar('abc'))", [["abc"], ["abc"]])
|
||||
|
||||
def testUnterminated(self):
|
||||
"""
|
||||
Test unterminated clauses
|
||||
"""
|
||||
|
||||
# Unterminated bracket expressions
|
||||
with self.assertRaises(SQLError):
|
||||
self.db.search("select [a from txtai")
|
||||
|
||||
with self.assertRaises(SQLError):
|
||||
self.db.search("select avg([a) from txtai")
|
||||
|
||||
with self.assertRaises(SQLError):
|
||||
self.db.search("select [a[0] from txtai")
|
||||
|
||||
# Unterminated function expressions
|
||||
with self.assertRaises(SQLError):
|
||||
self.db.search("select func(a from txtai")
|
||||
|
||||
with self.assertRaises(SQLError):
|
||||
self.db.search("select * from txtai where coalesce(a")
|
||||
|
||||
# Unterminated similar clause
|
||||
with self.assertRaises(SQLError):
|
||||
self.db.search("select * from txtai where similar('abc'")
|
||||
|
||||
def testUpper(self):
|
||||
"""
|
||||
Test SQL statements are case insensitive.
|
||||
"""
|
||||
|
||||
self.assertSql("groupby", "SELECT * FROM TXTAI WHERE a = 1 GROUP BY id", "s.id")
|
||||
self.assertSql("orderby", "SELECT * FROM TXTAI WHERE a = 1 ORDER BY id", "s.id")
|
||||
|
||||
def testWhereBasic(self):
|
||||
"""
|
||||
Test basic where clauses
|
||||
"""
|
||||
|
||||
prefix = "select * from txtai "
|
||||
|
||||
self.assertSql("where", prefix + "where a = b", "json_extract(data, '$.a') = json_extract(data, '$.b')")
|
||||
self.assertSql("where", prefix + "where abc = def", "json_extract(data, '$.abc') = json_extract(data, '$.def')")
|
||||
self.assertSql("where", prefix + "where a = b.value", "json_extract(data, '$.a') = json_extract(data, '$.b.value')")
|
||||
self.assertSql("where", prefix + "where a = 1", "json_extract(data, '$.a') = 1")
|
||||
self.assertSql("where", prefix + "WHERE 1 = a", "1 = json_extract(data, '$.a')")
|
||||
self.assertSql("where", prefix + "WHERE a LIKE 'abc'", "json_extract(data, '$.a') LIKE 'abc'")
|
||||
self.assertSql("where", prefix + "WHERE a NOT LIKE 'abc'", "json_extract(data, '$.a') NOT LIKE 'abc'")
|
||||
self.assertSql("where", prefix + "WHERE a IN (1, 2, 3, b)", "json_extract(data, '$.a') IN (1, 2, 3, json_extract(data, '$.b'))")
|
||||
self.assertSql("where", prefix + "WHERE a is not null", "json_extract(data, '$.a') is not null")
|
||||
self.assertSql("where", prefix + "WHERE score >= 0.15", "score >= 0.15")
|
||||
|
||||
def testWhereCompound(self):
|
||||
"""
|
||||
Test compound where clauses
|
||||
"""
|
||||
|
||||
prefix = "select * from txtai "
|
||||
|
||||
self.assertSql("where", prefix + "where a > (b + 1)", "json_extract(data, '$.a') > (json_extract(data, '$.b') + 1)")
|
||||
self.assertSql("where", prefix + "where a > func('abc')", "json_extract(data, '$.a') > func('abc')")
|
||||
self.assertSql(
|
||||
"where", prefix + "where (id = 1 or id = 2) and a like 'abc'", "(s.id = 1 or s.id = 2) and json_extract(data, '$.a') like 'abc'"
|
||||
)
|
||||
self.assertSql(
|
||||
"where",
|
||||
prefix + "where a > f(d(b, c, 1),1)",
|
||||
"json_extract(data, '$.a') > f(d(json_extract(data, '$.b'), json_extract(data, '$.c'), 1), 1)",
|
||||
)
|
||||
self.assertSql("where", prefix + "where (id = 1 AND id = 2) OR indexid = 3", "(s.id = 1 AND s.id = 2) OR s.indexid = 3")
|
||||
self.assertSql("where", prefix + "where f(id) = b(id)", "f(s.id) = b(s.id)")
|
||||
self.assertSql("where", prefix + "WHERE f(id)", "f(s.id)")
|
||||
|
||||
def assertSql(self, clause, query, expected):
|
||||
"""
|
||||
Helper method to assert a query clause is as expected.
|
||||
|
||||
Args:
|
||||
clause: clause to select
|
||||
query: input query
|
||||
expected: expected transformed query value
|
||||
"""
|
||||
|
||||
self.assertEqual(self.sql(query)[clause], expected)
|
||||
@@ -0,0 +1,73 @@
|
||||
"""
|
||||
SQLite module tests
|
||||
"""
|
||||
|
||||
from txtai.embeddings import Embeddings
|
||||
|
||||
from .testrdbms import Common
|
||||
|
||||
|
||||
# pylint: disable=R0904
|
||||
class TestSQLite(Common.TestRDBMS):
|
||||
"""
|
||||
Embeddings with content stored in SQLite tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Initialize test data.
|
||||
"""
|
||||
|
||||
cls.data = [
|
||||
"US tops 5 million confirmed virus cases",
|
||||
"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
|
||||
"Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
|
||||
"The National Park Service warns against sacrificing slower friends in a bear attack",
|
||||
"Maine man wins $1M from $25 lottery ticket",
|
||||
"Make huge profits without work, earn up to $100,000 a day",
|
||||
]
|
||||
|
||||
# Content backend
|
||||
cls.backend = "sqlite"
|
||||
|
||||
# Create embeddings model, backed by sentence-transformers & transformers
|
||||
cls.embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "content": cls.backend})
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""
|
||||
Cleanup data.
|
||||
"""
|
||||
|
||||
if cls.embeddings:
|
||||
cls.embeddings.close()
|
||||
|
||||
def testFunction(self):
|
||||
"""
|
||||
Test custom functions
|
||||
"""
|
||||
|
||||
embeddings = Embeddings(
|
||||
{
|
||||
"path": "sentence-transformers/nli-mpnet-base-v2",
|
||||
"content": self.backend,
|
||||
"functions": [{"name": "textlength", "function": "testdatabase.testsqlite.length"}],
|
||||
}
|
||||
)
|
||||
|
||||
# Create an index for the list of text
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Search for best match
|
||||
result = embeddings.search("select textlength(text) length from txtai where id = 0", 1)[0]
|
||||
|
||||
self.assertEqual(int(result["length"]), 39)
|
||||
|
||||
|
||||
def length(text):
|
||||
"""
|
||||
Custom SQL function.
|
||||
"""
|
||||
|
||||
return len(text)
|
||||
Reference in New Issue
Block a user