chore: import upstream snapshot with attribution

test/python/testapi/testapiagent.py

@@ -0,0 +1,91 @@
+"""
+Agent API module tests
+"""
+
+import os
+import tempfile
+import unittest
+
+from unittest.mock import patch
+
+from fastapi.testclient import TestClient
+
+from txtai.api import API, application
+
+# Configuration for agents
+AGENTS = """
+agent:
+    test:
+        max_iterations: 1
+        tools:
+            - name: testtool
+              description: Test tool
+              target: testapi.testapiagent.TestTool
+
+llm:
+    path: hf-internal-testing/tiny-random-LlamaForCausalLM
+"""
+
+
+# pylint: disable=R0904
+class TestAgent(unittest.TestCase):
+    """
+    API tests for agents.
+    """
+
+    @staticmethod
+    @patch.dict(os.environ, {"CONFIG": os.path.join(tempfile.gettempdir(), "testapi.yml"), "API_CLASS": "txtai.api.API"})
+    def start():
+        """
+        Starts a mock FastAPI client.
+        """
+
+        config = os.path.join(tempfile.gettempdir(), "testapi.yml")
+
+        with open(config, "w", encoding="utf-8") as output:
+            output.write(AGENTS)
+
+        # Create new application and set on client
+        application.app = application.create()
+        client = TestClient(application.app)
+        application.start()
+
+        # Patch LLM to generate answer
+        agent = application.get().agents["test"]
+        agent.process.model.llm = lambda *args, **kwargs: 'Action:\n{"name": "final_answer", "arguments": "Hi"}'
+
+        return client
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create API client on creation of class.
+        """
+
+        cls.client = TestAgent.start()
+
+    def testAgent(self):
+        """
+        Test agent via API
+        """
+
+        results = self.client.post("agent", json={"name": "test", "text": "Hello"}).json()
+        self.assertEqual(results, "Hi")
+
+    def testEmpty(self):
+        """
+        Test empty API configuration
+        """
+
+        api = API({})
+
+        self.assertIsNone(api.agent("junk", "test"))
+
+
+class TestTool:
+    """
+    Class to test agent tools
+    """
+
+    def __call__(self):
+        pass

test/python/testapi/testapiembeddings.py

@@ -0,0 +1,502 @@
+"""
+Embeddings API module tests
+"""
+
+import os
+import tempfile
+import unittest
+import urllib.parse
+
+from unittest.mock import patch
+
+from fastapi.testclient import TestClient
+
+from txtai.api import API, application
+
+# Configuration for a read/write embeddings index
+INDEX = """
+# Index file path
+path: %s
+
+# Allow indexing of documents
+writable: True
+
+# Allow reindexing
+reindex: True
+
+# Questions settings
+questions:
+    path: distilbert-base-cased-distilled-squad
+
+# Embeddings settings
+embeddings:
+    path: sentence-transformers/nli-mpnet-base-v2
+
+# Extractor settings
+extractor:
+    path: questions
+"""
+
+# Configuration for a read-only embeddings index
+READONLY = """
+# Index file path
+path: %s
+
+# Allow indexing of documents
+writable: False
+
+# Embeddings settings
+embeddings:
+"""
+
+# Configuration for an index with custom functions
+FUNCTIONS = """
+# Ignore existing index
+pathignore: %s
+
+# Allow indexing of documents
+writable: True
+
+# Embeddings settings
+embeddings:
+    path: sentence-transformers/nli-mpnet-base-v2
+    content: True
+    functions:
+        - testapi.testapiembeddings.Elements
+        - name: length
+          argcount: 1
+          function: testapi.testapiembeddings.length
+        - name: ann
+          function: ann
+    transform: testapi.testapiembeddings.transform
+"""
+
+# Configuration for RAG
+RAG = """
+# Ignore existing index
+pathignore: %s
+
+# Allow indexing of documents
+writable: True
+
+# Embeddings settings
+embeddings:
+    path: sentence-transformers/nli-mpnet-base-v2
+    content: True
+
+# LLM
+llm:
+    path: hf-internal-testing/tiny-random-gpt2
+    task: language-generation
+
+# RAG settings
+rag:
+    path: llm
+    output: flatten
+"""
+
+# Configuration for reindexing disabled
+REINDEXDISABLED = """
+# Index file path
+path: %s
+
+# Allow indexing of documents
+writable: True
+
+# Embeddings settings
+embeddings:
+    path: sentence-transformers/nli-mpnet-base-v2
+    content: True
+"""
+
+# Configuration for reranker
+RERANK = """
+# Index file path
+path: %s
+
+# Allow indexing of documents
+writable: True
+
+# Embeddings settings
+embeddings:
+    path: sentence-transformers/nli-mpnet-base-v2
+    content: True
+
+# Similarity and Reranking settings
+similarity:
+    path: neuml/colbert-bert-tiny
+    lateencode: True
+
+reranker:
+"""
+
+
+# pylint: disable=R0904
+class TestEmbeddings(unittest.TestCase):
+    """
+    API tests for embeddings indices.
+    """
+
+    @staticmethod
+    @patch.dict(os.environ, {"CONFIG": os.path.join(tempfile.gettempdir(), "testapi.yml"), "API_CLASS": "txtai.api.API"})
+    def start(yaml, path="testapi"):
+        """
+        Starts a mock FastAPI client.
+
+        Args:
+            yaml: input configuration
+            path: output path
+        """
+
+        config = os.path.join(tempfile.gettempdir(), "testapi.yml")
+        index = os.path.join(tempfile.gettempdir(), path)
+
+        with open(config, "w", encoding="utf-8") as output:
+            output.write(yaml % index)
+
+        # Create new application and set on client
+        application.app = application.create()
+        client = TestClient(application.app)
+        application.start()
+
+        return client
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create API client on creation of class.
+        """
+
+        cls.client = TestEmbeddings.start(INDEX, "testapi")
+
+        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",
+        ]
+
+        # Index data
+        cls.client.post("add", json=[{"id": x, "text": row} for x, row in enumerate(cls.data)])
+        cls.client.get("index")
+
+    def testCount(self):
+        """
+        Test count via API
+        """
+
+        self.assertEqual(self.client.get("count").json(), 6)
+
+    def testDelete(self):
+        """
+        Test delete via API
+        """
+
+        # Delete best match
+        ids = self.client.post("delete", json=[4]).json()
+        self.assertEqual(ids, [4])
+
+        # Search for best match
+        query = urllib.parse.quote("feel good story")
+        uid = self.client.get(f"search?query={query}&limit=1").json()[0]["id"]
+
+        self.assertEqual(self.client.get("count").json(), 5)
+        self.assertEqual(uid, 5)
+
+        # Reset data
+        self.client.post("add", json=[{"id": x, "text": row} for x, row in enumerate(self.data)])
+        self.client.get("index")
+
+    def testEmpty(self):
+        """
+        Test empty API configuration
+        """
+
+        api = API({"writable": True})
+
+        self.assertIsNone(api.search("test", None))
+        self.assertIsNone(api.batchsearch(["test"], None))
+        self.assertIsNone(api.delete(["test"]))
+        self.assertIsNone(api.count())
+        self.assertIsNone(api.similarity("test", ["test"]))
+        self.assertIsNone(api.batchsimilarity(["test"], ["test"]))
+        self.assertIsNone(api.explain("test"))
+        self.assertIsNone(api.batchexplain(["test"]))
+        self.assertIsNone(api.transform("test"))
+        self.assertIsNone(api.batchtransform(["test"]))
+        self.assertIsNone(api.extract(["test"], ["test"]))
+
+    def testExtractor(self):
+        """
+        Test qa extraction via API
+        """
+
+        data = [
+            "Giants hit 3 HRs to down Dodgers",
+            "Giants 5 Dodgers 4 final",
+            "Dodgers drop Game 2 against the Giants, 5-4",
+            "Blue Jays beat Red Sox final score 2-1",
+            "Red Sox lost to the Blue Jays, 2-1",
+            "Blue Jays at Red Sox is over. Score: 2-1",
+            "Phillies win over the Braves, 5-0",
+            "Phillies 5 Braves 0 final",
+            "Final: Braves lose to the Phillies in the series opener, 5-0",
+            "Lightning goaltender pulled, lose to Flyers 4-1",
+            "Flyers 4 Lightning 1 final",
+            "Flyers win 4-1",
+        ]
+
+        questions = ["What team won the game?", "What was score?"]
+
+        # pylint: disable=C3001
+        execute = lambda query: self.client.post(
+            "extract",
+            json={"queue": [{"name": question, "query": query, "question": question, "snippet": False} for question in questions], "texts": data},
+        ).json()
+
+        answers = execute("Red Sox - Blue Jays")
+        self.assertEqual("Blue Jays", answers[0]["answer"])
+        self.assertEqual("2-1", answers[1]["answer"])
+
+        # Ad-hoc questions
+        question = "What hockey team won?"
+
+        answers = self.client.post(
+            "extract", json={"queue": [{"name": question, "query": question, "question": question, "snippet": False}], "texts": data}
+        ).json()
+        self.assertEqual("Flyers", answers[0]["answer"])
+
+    def testReindex(self):
+        """
+        Test reindex via API
+        """
+
+        # Reindex data
+        self.client.post("reindex", json={"config": {"path": "sentence-transformers/nli-mpnet-base-v2"}})
+
+        # Search for best match
+        query = urllib.parse.quote("feel good story")
+        uid = self.client.get(f"search?query={query}&limit=1").json()[0]["id"]
+
+        self.assertEqual(uid, 4)
+
+        # Reset data
+        self.client.post("add", json=[{"id": x, "text": row} for x, row in enumerate(self.data)])
+        self.client.get("index")
+
+    def testSearch(self):
+        """
+        Test search via API
+        """
+
+        query = urllib.parse.quote("feel good story")
+        uid = self.client.get(f"search?query={query}&limit=1").json()[0]["id"]
+        self.assertEqual(uid, 4)
+
+    def testSearchBatch(self):
+        """
+        Test batch search via API
+        """
+
+        results = self.client.post("batchsearch", json={"queries": ["feel good story", "climate change"], "limit": 1}).json()
+
+        uids = [result[0]["id"] for result in results]
+        self.assertEqual(uids, [4, 1])
+
+    def testSimilarity(self):
+        """
+        Test similarity via API
+        """
+
+        uid = self.client.post("similarity", json={"query": "feel good story", "texts": self.data}).json()[0]["id"]
+
+        self.assertEqual(uid, 4)
+
+    def testSimilarityBatch(self):
+        """
+        Test batch similarity via API
+        """
+
+        results = self.client.post("batchsimilarity", json={"queries": ["feel good story", "climate change"], "texts": self.data}).json()
+
+        uids = [result[0]["id"] for result in results]
+        self.assertEqual(uids, [4, 1])
+
+    def testTransform(self):
+        """
+        Test embeddings transform via API
+        """
+
+        self.assertEqual(len(self.client.get("transform?text=testembed").json()), 768)
+
+    def testTransformBatch(self):
+        """
+        Test batch embeddings transform via API
+        """
+
+        embeddings = self.client.post("batchtransform", json=self.data).json()
+
+        self.assertEqual(len(embeddings), len(self.data))
+        self.assertEqual(len(embeddings[0]), 768)
+
+    def testUpsert(self):
+        """
+        Test upsert via API
+        """
+
+        # Update data
+        self.client.post("add", json=[{"id": 0, "text": "Feel good story: baby panda born"}])
+        self.client.get("upsert")
+
+        # Search for best match
+        query = urllib.parse.quote("feel good story")
+        uid = self.client.get(f"search?query={query}&limit=1").json()[0]["id"]
+
+        self.assertEqual(uid, 0)
+
+        # Reset data
+        self.client.post("add", json=[{"id": x, "text": row} for x, row in enumerate(self.data)])
+        self.client.get("index")
+
+    def testViewOnly(self):
+        """
+        Test read-only API instance
+        """
+
+        # Re-create application with a read-only index
+        self.client = TestEmbeddings.start(READONLY)
+
+        # Test search
+        query = urllib.parse.quote("feel good story")
+        uid = self.client.get(f"search?query={query}&limit=1").json()[0]["id"]
+        self.assertEqual(uid, 4)
+
+        # Test similarity
+        uid = self.client.post("similarity", json={"query": "feel good story", "texts": self.data}).json()[0]["id"]
+        self.assertEqual(uid, 4)
+
+        # Test errors raised for write operations
+        self.assertEqual(self.client.post("add", json=[{"id": 0, "text": "test"}]).status_code, 403)
+        self.assertEqual(self.client.get("index").status_code, 403)
+        self.assertEqual(self.client.get("upsert").status_code, 403)
+        self.assertEqual(self.client.post("delete", json=[0]).status_code, 403)
+        self.assertEqual(self.client.post("reindex", json={"config": {"path": "sentence-transformers/nli-mpnet-base-v2"}}).status_code, 403)
+
+    def testXFunctions(self):
+        """
+        Test API instance with custom functions
+        """
+
+        # Re-create application with custom functions
+        self.client = TestEmbeddings.start(FUNCTIONS)
+
+        # Index data
+        self.client.post("add", json=[{"id": x, "text": row} for x, row in enumerate(self.data)])
+        self.client.get("index")
+
+        query = urllib.parse.quote("select elements('text') length from txtai limit 1")
+        self.assertEqual(self.client.get(f"search?query={query}").json()[0]["length"], 4)
+
+        query = urllib.parse.quote("select length('text') length from txtai limit 1")
+        self.assertEqual(self.client.get(f"search?query={query}").json()[0]["length"], 4)
+
+    def testXPlain(self):
+        """
+        Test API instance with explain methods
+        """
+
+        results = self.client.post("explain", json={"query": "feel good story", "limit": 1}).json()
+
+        self.assertEqual(results[0]["text"], self.data[4])
+        self.assertIsNotNone(results[0].get("tokens"))
+
+    def testXPlainBatch(self):
+        """
+        Test batch query explain via API
+        """
+
+        results = self.client.post("batchexplain", json={"queries": ["feel good story", "climate change"], "limit": 1}).json()
+
+        text = [result[0]["text"] for result in results]
+        self.assertEqual(text, [self.data[4], self.data[1]])
+        self.assertIsNotNone(results[0][0].get("tokens"))
+
+    def testXRAG(self):
+        """
+        Test RAG via API
+        """
+
+        # Re-create application with a RAG pipeline
+        self.client = TestEmbeddings.start(RAG)
+
+        # Index data
+        self.client.post("add", json=[{"id": x, "text": row} for x, row in enumerate(self.data)])
+        self.client.get("index")
+
+        response = self.client.get("rag?query=bear").json()
+        self.assertIsInstance(response, str)
+
+        response = self.client.post("batchrag", json={"queries": ["bear", "bear"]}).json()
+        self.assertEqual(len(response), 2)
+
+    def testXReindexDisabled(self):
+        """
+        Test reindexing is disabled
+        """
+
+        # Re-create application with reindexing disabled
+        self.client = TestEmbeddings.start(REINDEXDISABLED, "testapi-reindex")
+
+        # Index data
+        self.client.post("add", json=[{"id": x, "text": row} for x, row in enumerate(self.data)])
+        self.client.get("index")
+
+        # Assert error raised
+        self.assertEqual(self.client.post("reindex", json={"config": {"path": "sentence-transformers/nli-mpnet-base-v2"}}).status_code, 403)
+
+    def testXRerank(self):
+        """
+        Test rerank via API
+        """
+
+        # Re-create application with a reranker pipeline
+        self.client = TestEmbeddings.start(RERANK, "testapi-rerank")
+
+        # Index data
+        self.client.post("add", json=[{"id": x, "text": row} for x, row in enumerate(self.data)])
+        self.client.get("index")
+
+        uid = self.client.get("rerank?query=bear").json()[0]["id"]
+        self.assertEqual(uid, "3")
+
+        results = self.client.post("batchrerank", json={"queries": ["bear", "bear"]}).json()
+
+        uids = [result[0]["id"] for result in results]
+        self.assertEqual(uids, ["3", "3"])
+
+
+class Elements:
+    """
+    Custom SQL function as callable object.
+    """
+
+    def __call__(self, text):
+        return length(text)
+
+
+def transform(document):
+    """
+    Custom transform function.
+    """
+
+    return document
+
+
+def length(text):
+    """
+    Custom SQL function.
+    """
+
+    return len(text)

test/python/testapi/testapipipeline.py

@@ -0,0 +1,391 @@
+"""
+Pipeline API module tests
+"""
+
+import os
+import tempfile
+import unittest
+import urllib
+
+from unittest.mock import patch
+
+from fastapi.testclient import TestClient
+
+from txtai.api import API, application
+
+# pylint: disable=C0411
+from utils import Utils
+
+# Configuration for pipelines
+PIPELINES = """
+# Image captions
+caption:
+
+# Entity extraction
+entity:
+    path: dslim/bert-base-NER
+
+# Extractor settings
+extractor:
+    similarity: similarity
+    path: llm
+
+# Label settings
+labels:
+    path: prajjwal1/bert-medium-mnli
+
+# LLM settings
+llm:
+    path: hf-internal-testing/tiny-random-gpt2
+    task: language-generation
+
+# Image objects
+objects:
+
+# Text segmentation
+segmentation:
+    sentences: true
+
+# Enable pipeline similarity backed by zero shot classifier
+similarity:
+
+# Summarization
+summary:
+    path: t5-small
+
+# Tabular
+tabular:
+
+# Text extraction
+textractor:
+    safeopen: /tmp/txtai
+
+# Text to speech
+texttospeech:
+
+# Transcription
+transcription:
+
+# Translation:
+translation:
+
+# Enable file uploads
+upload:
+"""
+
+
+# pylint: disable=R0904
+class TestPipeline(unittest.TestCase):
+    """
+    API tests for pipelines.
+    """
+
+    @staticmethod
+    @patch.dict(os.environ, {"CONFIG": os.path.join(tempfile.gettempdir(), "testapi.yml"), "API_CLASS": "txtai.api.API"})
+    def start():
+        """
+        Starts a mock FastAPI client.
+        """
+
+        config = os.path.join(tempfile.gettempdir(), "testapi.yml")
+
+        with open(config, "w", encoding="utf-8") as output:
+            output.write(PIPELINES)
+
+        # Create new application and set on client
+        application.app = application.create()
+        client = TestClient(application.app)
+        application.start()
+
+        return client
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create API client on creation of class.
+        """
+
+        cls.client = TestPipeline.start()
+
+        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",
+        ]
+
+        cls.text = (
+            "Search is the base of many applications. Once data starts to pile up, users want to be able to find it. It's the foundation "
+            "of the internet and an ever-growing challenge that is never solved or done. The field of Natural Language Processing (NLP) is "
+            "rapidly evolving with a number of new developments. Large-scale general language models are an exciting new capability "
+            "allowing us to add amazing functionality quickly with limited compute and people. Innovation continues with new models "
+            "and advancements coming in at what seems a weekly basis. This article introduces txtai, an AI-powered search engine "
+            "that enables Natural Language Understanding (NLU) based search in any application."
+        )
+
+        # Create invalid test file
+        directory = os.path.join(tempfile.gettempdir(), "txtai-1")
+        os.makedirs(directory, exist_ok=True)
+        with open(os.path.join(directory, "test"), "w", encoding="utf-8") as output:
+            output.write("123")
+
+    def testCaption(self):
+        """
+        Test caption via API
+        """
+
+        caption = self.client.get(f"caption?file={Utils.PATH}/books.jpg").json()
+
+        self.assertEqual(caption, "a book shelf filled with books and a stack of books")
+
+    def testCaptionBatch(self):
+        """
+        Test batch caption via API
+        """
+
+        path = Utils.PATH + "/books.jpg"
+
+        captions = self.client.post("batchcaption", json=[path, path]).json()
+        self.assertEqual(captions, ["a book shelf filled with books and a stack of books"] * 2)
+
+    def testEntity(self):
+        """
+        Test entity extraction via API
+        """
+
+        entities = self.client.get(f"entity?text={self.data[1]}").json()
+        self.assertEqual([e[0] for e in entities], ["Canada", "Manhattan"])
+
+    def testEntityBatch(self):
+        """
+        Test batch entity via API
+        """
+
+        entities = self.client.post("batchentity", json=[self.data[1]]).json()
+        self.assertEqual([e[0] for e in entities[0]], ["Canada", "Manhattan"])
+
+    def testEmpty(self):
+        """
+        Test empty API configuration
+        """
+
+        api = API({})
+
+        self.assertIsNone(api.label("test", ["test"]))
+        self.assertIsNone(api.pipeline("junk", "test"))
+
+    def testLabel(self):
+        """
+        Test label via API
+        """
+
+        labels = self.client.post("label", json={"text": "this is the best sentence ever", "labels": ["positive", "negative"]}).json()
+
+        self.assertEqual(labels[0]["id"], 0)
+
+    def testLabelBatch(self):
+        """
+        Test batch label via API
+        """
+
+        labels = self.client.post(
+            "batchlabel", json={"texts": ["this is the best sentence ever", "This is terrible"], "labels": ["positive", "negative"]}
+        ).json()
+
+        results = [l[0]["id"] for l in labels]
+        self.assertEqual(results, [0, 1])
+
+    def testLLM(self):
+        """
+        Test LLM inference via API
+        """
+
+        response = self.client.get("llm?text=test").json()
+        self.assertIsInstance(response, str)
+
+    def testLLMBatch(self):
+        """
+        Test batch LLM inference via API
+        """
+
+        response = self.client.post("batchllm", json={"texts": ["test", "test"]}).json()
+        self.assertEqual(len(response), 2)
+
+    def testObjects(self):
+        """
+        Test objects via API
+        """
+
+        objects = self.client.get(f"objects?file={Utils.PATH}/books.jpg").json()
+
+        self.assertEqual(objects[0][0], "book")
+
+    def testObjectsBatch(self):
+        """
+        Test batch objects via API
+        """
+
+        path = Utils.PATH + "/books.jpg"
+
+        objects = self.client.post("batchobjects", json=[path, path]).json()
+        self.assertEqual([o[0][0] for o in objects], ["book"] * 2)
+
+    def testSegment(self):
+        """
+        Test segmentation via API
+        """
+
+        text = self.client.get("segment?text=This is a test. And another test.").json()
+
+        # Check array length is 2
+        self.assertEqual(len(text), 2)
+
+    def testSegmentBatch(self):
+        """
+        Test batch segmentation via API
+        """
+
+        text = "This is a test. And another test."
+        texts = self.client.post("batchsegment", json=[text, text]).json()
+
+        # Check array length is 2 and first element length is 2
+        self.assertEqual(len(texts), 2)
+        self.assertEqual(len(texts[0]), 2)
+
+    def testSimilarity(self):
+        """
+        Test similarity via API
+        """
+
+        uid = self.client.post("similarity", json={"query": "feel good story", "texts": self.data}).json()[0]["id"]
+
+        self.assertEqual(self.data[uid], self.data[4])
+
+    def testSimilarityBatch(self):
+        """
+        Test batch similarity via API
+        """
+
+        results = self.client.post("batchsimilarity", json={"queries": ["feel good story", "climate change"], "texts": self.data}).json()
+
+        uids = [result[0]["id"] for result in results]
+        self.assertEqual(uids, [4, 1])
+
+    def testSummary(self):
+        """
+        Test summary via API
+        """
+
+        summary = self.client.get(f"summary?text={urllib.parse.quote(self.text)}&minlength=15&maxlength=15").json()
+        self.assertEqual(summary, "the field of natural language processing (NLP) is rapidly evolving")
+
+    def testSummaryBatch(self):
+        """
+        Test batch summary via API
+        """
+
+        summaries = self.client.post("batchsummary", json={"texts": [self.text, self.text], "minlength": 15, "maxlength": 15}).json()
+        self.assertEqual(summaries, ["the field of natural language processing (NLP) is rapidly evolving"] * 2)
+
+    def testTabular(self):
+        """
+        Test tabular via API
+        """
+
+        results = self.client.get(f"tabular?file={Utils.PATH}/tabular.csv").json()
+
+        # Check length of results is as expected
+        self.assertEqual(len(results), 6)
+
+    def testTabularBatch(self):
+        """
+        Test batch tabular via API
+        """
+
+        path = Utils.PATH + "/tabular.csv"
+
+        results = self.client.post("batchtabular", json=[path, path]).json()
+        self.assertEqual((len(results[0]), len(results[1])), (6, 6))
+
+    def testTextractor(self):
+        """
+        Test textractor via API
+        """
+
+        text = self.client.get(f"textract?file={Utils.PATH}/article.pdf").json()
+
+        # Check length of text is as expected
+        self.assertEqual(len(text), 2471)
+
+        # Check invalid URLs
+        for url in ["http://192.168.1.1/path", "http://127.0.0.1/path", "http://invalid", "/etc/config", "/tmp/txtai-1/test"]:
+            with self.assertRaises(IOError):
+                self.client.get(f"textract?file={url}").json()
+
+    def testTextractorBatch(self):
+        """
+        Test batch textractor via API
+        """
+
+        path = Utils.PATH + "/article.pdf"
+
+        texts = self.client.post("batchtextract", json=[path, path]).json()
+        self.assertEqual((len(texts[0]), len(texts[1])), (2471, 2471))
+
+    def testTextToSpeech(self):
+        """
+        Test text to speech
+        """
+
+        # Generate audio and check for WAV signature
+        audio = self.client.get("texttospeech?text=hello&encoding=wav").content
+        self.assertTrue(audio[0:4] == b"RIFF")
+
+    def testTranscribe(self):
+        """
+        Test transcribe via API
+        """
+
+        text = self.client.get(f"transcribe?file={Utils.PATH}/Make_huge_profits.wav").json()
+
+        # Check length of text is as expected
+        self.assertEqual(text, "Make huge profits without working make up to one hundred thousand dollars a day")
+
+    def testTranscribeBatch(self):
+        """
+        Test batch transcribe via API
+        """
+
+        path = Utils.PATH + "/Make_huge_profits.wav"
+
+        texts = self.client.post("batchtranscribe", json=[path, path]).json()
+        self.assertEqual(texts, ["Make huge profits without working make up to one hundred thousand dollars a day"] * 2)
+
+    def testTranslate(self):
+        """
+        Test translate via API
+        """
+
+        translation = self.client.get(f"translate?text={urllib.parse.quote('This is a test translation into Spanish')}&target=es").json()
+        self.assertEqual(translation, "Esta es una traducción de prueba al español")
+
+    def testTranslateBatch(self):
+        """
+        Test batch translate via API
+        """
+
+        text = "This is a test translation into Spanish"
+        translations = self.client.post("batchtranslate", json={"texts": [text, text], "target": "es"}).json()
+        self.assertEqual(translations, ["Esta es una traducción de prueba al español"] * 2)
+
+    def testUpload(self):
+        """
+        Test file upload
+        """
+
+        path = Utils.PATH + "/article.pdf"
+        with open(path, "rb") as f:
+            path = self.client.post("upload", files={"files": f}).json()[0]
+            self.assertTrue(os.path.exists(path))

test/python/testapi/testapiworkflow.py

@@ -0,0 +1,259 @@
+"""
+Workflow API module tests
+"""
+
+import json
+import os
+import tempfile
+import unittest
+
+from http.server import HTTPServer, BaseHTTPRequestHandler
+from multiprocessing.pool import ThreadPool
+from threading import Thread
+from unittest.mock import patch
+
+from fastapi.testclient import TestClient
+
+from txtai.api import API, application
+
+# Configuration for workflows
+WORKFLOWS = """
+# Embeddings index
+writable: true
+embeddings:
+    path: sentence-transformers/nli-mpnet-base-v2
+
+# Labels
+labels:
+    path: prajjwal1/bert-medium-mnli
+
+nop:
+
+# Text segmentation
+segmentation:
+    sentences: true
+
+# Workflow definitions
+workflow:
+    labels:
+        tasks:
+            - action: labels
+              args: [[positive, negative]]
+    multiaction:
+        tasks:
+            - action:
+                - labels
+                - nop
+              initialize: testapi.testapiworkflow.TestInitFinal
+              finalize: testapi.testapiworkflow.TestInitFinal
+              merge: concat
+              args:
+                - [[positive, negative], false, True]
+                - null
+    schedule:
+        schedule:
+            cron: '* * * * * *'
+            elements:
+                - This is a test sentence. And another sentence to split.
+            iterations: 1
+        tasks:
+            - action: segmentation
+    segment:
+        tasks:
+            - action: segmentation
+            - action: index
+    get:
+        tasks:
+            - task: service
+              url: http://127.0.0.1:8001/testget
+              method: get
+              params:
+                text:
+    post:
+        tasks:
+            - task: service
+              url: http://127.0.0.1:8001/testpost
+              params:
+
+    xml:
+        tasks:
+            - task: service
+              url: http://127.0.0.1:8001/xml
+              method: get
+              batch: false
+              extract: row
+              params:
+                text:
+"""
+
+
+class RequestHandler(BaseHTTPRequestHandler):
+    """
+    Test HTTP handler.
+    """
+
+    def do_GET(self):
+        """
+        GET request handler.
+        """
+
+        self.send_response(200)
+
+        if self.path.startswith("/xml"):
+            response = "<row><text>test</text></row>".encode("utf-8")
+            mime = "application/xml"
+        else:
+            response = '[{"text": "test"}]'.encode("utf-8")
+            mime = "application/json"
+
+        self.send_header("content-type", mime)
+        self.send_header("content-length", len(response))
+        self.end_headers()
+
+        self.wfile.write(response)
+        self.wfile.flush()
+
+    def do_POST(self):
+        """
+        POST request handler.
+        """
+
+        length = int(self.headers["content-length"])
+        data = json.loads(self.rfile.read(length))
+
+        response = json.dumps([[y for y in x.split(".") if y] for x in data]).encode("utf-8")
+
+        self.send_response(200)
+        self.send_header("content-type", "application/json")
+        self.send_header("content-length", len(response))
+        self.end_headers()
+
+        self.wfile.write(response)
+        self.wfile.flush()
+
+
+class TestWorkflow(unittest.TestCase):
+    """
+    API tests for workflows.
+    """
+
+    @staticmethod
+    @patch.dict(os.environ, {"CONFIG": os.path.join(tempfile.gettempdir(), "testapi.yml"), "API_CLASS": "txtai.api.API"})
+    def start():
+        """
+        Starts a mock FastAPI client.
+        """
+
+        config = os.path.join(tempfile.gettempdir(), "testapi.yml")
+
+        with open(config, "w", encoding="utf-8") as output:
+            output.write(WORKFLOWS)
+
+        # Create new application and set on client
+        application.app = application.create()
+        client = TestClient(application.app)
+        application.start()
+
+        return client
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create API client on creation of class.
+        """
+
+        cls.client = TestWorkflow.start()
+
+        cls.httpd = HTTPServer(("127.0.0.1", 8001), RequestHandler)
+
+        server = Thread(target=cls.httpd.serve_forever, daemon=True)
+        server.start()
+
+    @classmethod
+    def tearDownClass(cls):
+        """
+        Shutdown mock http server.
+        """
+
+        cls.httpd.shutdown()
+
+    def testAPICleanup(self):
+        """
+        Test API threadpool closed when __del__ called.
+        """
+
+        api = API({})
+        api.pool = ThreadPool()
+
+        # pylint: disable=C2801
+        api.__del__()
+
+        self.assertIsNone(api.pool)
+
+    def testServiceGet(self):
+        """
+        Test workflow with ServiceTask GET via API
+        """
+
+        text = "This is a test sentence. And another sentence to split."
+        results = self.client.post("workflow", json={"name": "get", "elements": [text]}).json()
+
+        self.assertEqual(len(results), 1)
+        self.assertEqual(len(results[0]), 1)
+
+    def testServicePost(self):
+        """
+        Test workflow with ServiceTask POST via API
+        """
+
+        text = "This is a test sentence. And another sentence to split."
+        results = self.client.post("workflow", json={"name": "post", "elements": [text]}).json()
+
+        self.assertEqual(len(results), 1)
+        self.assertEqual(len(results[0]), 2)
+
+    def testServiceXml(self):
+        """
+        Test workflow with ServiceTask GET via API and XML response
+        """
+
+        text = "This is a test sentence. And another sentence to split."
+        results = self.client.post("workflow", json={"name": "xml", "elements": [text]}).json()
+
+        self.assertEqual(len(results), 1)
+        self.assertEqual(len(results[0]), 1)
+
+    def testWorkflowLabels(self):
+        """
+        Test workflow with labels via API
+        """
+
+        text = "This is the best"
+
+        results = self.client.post("workflow", json={"name": "labels", "elements": [text]}).json()
+        self.assertEqual(results[0][0], 0)
+
+        results = self.client.post("workflow", json={"name": "multiaction", "elements": [text]}).json()
+        self.assertEqual(results[0], "['positive']. This is the best")
+
+    def testWorkflowSegment(self):
+        """
+        Test workflow with segmentation via API
+        """
+
+        text = "This is a test sentence. And another sentence to split."
+
+        results = self.client.post("workflow", json={"name": "segment", "elements": [text]}).json()
+        self.assertEqual(len(results), 2)
+
+        results = self.client.post("workflow", json={"name": "segment", "elements": [[0, text]]}).json()
+        self.assertEqual(len(results), 2)
+
+
+class TestInitFinal:
+    """
+    Class to test task initialize and finalize calls.
+    """
+
+    def __call__(self):
+        pass

test/python/testapi/testauthorization.py

@@ -0,0 +1,73 @@
+"""
+Authorization module tests
+"""
+
+import hashlib
+import os
+import tempfile
+import unittest
+
+from unittest.mock import patch
+
+from fastapi.testclient import TestClient
+
+from txtai.api import application
+
+
+class TestAuthorization(unittest.TestCase):
+    """
+    API tests for token authorization.
+    """
+
+    @staticmethod
+    @patch.dict(
+        os.environ,
+        {
+            "CONFIG": os.path.join(tempfile.gettempdir(), "testapi.yml"),
+            "DEPENDENCIES": "txtai.api.Authorization",
+            "TOKEN": hashlib.sha256("token".encode("utf-8")).hexdigest(),
+        },
+    )
+    def start():
+        """
+        Starts a mock FastAPI client.
+        """
+
+        config = os.path.join(tempfile.gettempdir(), "testapi.yml")
+
+        with open(config, "w", encoding="utf-8") as output:
+            output.write("embeddings:\n")
+
+        # Create new application and set on client
+        application.app = application.create()
+        client = TestClient(application.app)
+        application.start()
+
+        return client
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create API client on creation of class.
+        """
+
+        cls.client = TestAuthorization.start()
+
+    def testInvalid(self):
+        """
+        Test invalid authorization
+        """
+
+        response = self.client.get("search?query=test")
+        self.assertEqual(response.status_code, 401)
+
+        response = self.client.get("search?query=test", headers={"Authorization": "Bearer invalid"})
+        self.assertEqual(response.status_code, 401)
+
+    def testValid(self):
+        """
+        Test valid authorization
+        """
+
+        results = self.client.get("search?query=test", headers={"Authorization": "Bearer token"}).json()
+        self.assertEqual(results, [])

test/python/testapi/testcluster.py

@@ -0,0 +1,261 @@
+"""
+Cluster API module tests
+"""
+
+import json
+import os
+import tempfile
+import unittest
+import urllib.parse
+
+from http.server import HTTPServer, BaseHTTPRequestHandler
+from threading import Thread
+from unittest.mock import patch
+
+from fastapi.testclient import TestClient
+
+from txtai.api import application
+
+# Configuration for an embeddings cluster
+CLUSTER = """
+cluster:
+    shards:
+        - http://127.0.0.1:8002
+        - http://127.0.0.1:8003
+"""
+
+
+class RequestHandler(BaseHTTPRequestHandler):
+    """
+    Test HTTP handler.
+    """
+
+    def do_GET(self):
+        """
+        GET request handler.
+        """
+
+        if self.path == "/count":
+            response = 26
+        elif self.path.startswith("/search?query=select"):
+            if "group+by+id" in self.path:
+                response = [{"count(*)": 26}]
+            elif "group+by+text" in self.path:
+                response = [{"count(*)": 12, "text": "This is a test"}, {"count(*)": 14, "text": "And another test"}]
+            elif "group+by+txt" in self.path:
+                response = [{"count(*)": 12, "txt": "This is a test"}, {"count(*)": 14, "txt": "And another test"}]
+            else:
+                if self.server.server_port == 8002:
+                    response = [{"count(*)": 12, "min(indexid)": 0, "max(indexid)": 11, "avg(indexid)": 6.3}]
+                else:
+                    response = [{"count(*)": 16, "min(indexid)": 2, "max(indexid)": 14, "avg(indexid)": 6.7}]
+        elif self.path.startswith("/search"):
+            response = [{"id": 4, "score": 0.40}]
+        else:
+            response = {"result": "ok"}
+
+        # Convert response to string
+        response = json.dumps(response).encode("utf-8")
+
+        self.send_response(200)
+        self.send_header("content-type", "application/json")
+        self.send_header("content-length", len(response))
+        self.end_headers()
+
+        self.wfile.write(response)
+        self.wfile.flush()
+
+    def do_POST(self):
+        """
+        POST request handler.
+        """
+
+        if self.path.startswith("/batchsearch"):
+            response = [[{"id": 4, "score": 0.40}], [{"id": 1, "score": 0.40}]]
+        elif self.path.startswith("/delete"):
+            if self.server.server_port == 8002:
+                response = [0]
+            else:
+                response = []
+        else:
+            response = {"result": "ok"}
+
+        response = json.dumps(response).encode("utf-8")
+
+        self.send_response(200)
+        self.send_header("content-type", "application/json")
+        self.send_header("content-length", len(response))
+        self.end_headers()
+
+        self.wfile.write(response)
+        self.wfile.flush()
+
+
+@unittest.skipIf(os.name == "nt", "TestCluster skipped on Windows")
+class TestCluster(unittest.TestCase):
+    """
+    API tests for embeddings clusters
+    """
+
+    @staticmethod
+    @patch.dict(os.environ, {"CONFIG": os.path.join(tempfile.gettempdir(), "testapi.yml"), "API_CLASS": "txtai.api.API"})
+    def start():
+        """
+        Starts a mock FastAPI client.
+        """
+
+        config = os.path.join(tempfile.gettempdir(), "testapi.yml")
+
+        with open(config, "w", encoding="utf-8") as output:
+            output.write(CLUSTER)
+
+        # Create new application and set on client
+        application.app = application.create()
+        client = TestClient(application.app)
+        application.start()
+
+        return client
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create API client on creation of class.
+        """
+
+        cls.client = TestCluster.start()
+
+        cls.httpd1 = HTTPServer(("127.0.0.1", 8002), RequestHandler)
+
+        server1 = Thread(target=cls.httpd1.serve_forever, daemon=True)
+        server1.start()
+
+        cls.httpd2 = HTTPServer(("127.0.0.1", 8003), RequestHandler)
+
+        server2 = Thread(target=cls.httpd2.serve_forever, daemon=True)
+        server2.start()
+
+        # Index data
+        cls.client.post("add", json=[{"id": 0, "text": "test"}])
+        cls.client.get("index")
+
+    @classmethod
+    def tearDownClass(cls):
+        """
+        Shutdown mock http server.
+        """
+
+        cls.httpd1.shutdown()
+        cls.httpd2.shutdown()
+
+    def testCount(self):
+        """
+        Test cluster count
+        """
+
+        self.assertEqual(self.client.get("count").json(), 52)
+
+    def testDelete(self):
+        """
+        Test cluster delete
+        """
+
+        self.assertEqual(self.client.post("delete", json=[0]).json(), [0])
+
+    def testDeleteString(self):
+        """
+        Test cluster delete with string id
+        """
+
+        self.assertEqual(self.client.post("delete", json=["0"]).json(), [0])
+
+    def testIds(self):
+        """
+        Test id configurations
+        """
+
+        # String ids
+        self.client.post("add", json=[{"id": "0", "text": "test"}])
+        self.assertEqual(self.client.get("index").status_code, 200)
+
+        # Auto ids
+        self.client.post("add", json=[{"text": "test"}])
+        self.assertEqual(self.client.get("index").status_code, 200)
+
+    def testReindex(self):
+        """
+        Test cluster reindex
+        """
+
+        self.assertEqual(self.client.post("reindex", json={"config": {"path": "sentence-transformers/nli-mpnet-base-v2"}}).status_code, 200)
+
+    def testSearch(self):
+        """
+        Test cluster search
+        """
+
+        # Encode parameters
+        params = json.dumps({"x": 1})
+
+        query = urllib.parse.quote("feel good story")
+        uid = self.client.get(f"search?query={query}&limit=1&weights=0.5&index=default&parameters={params}&graph=False").json()[0]["id"]
+        self.assertEqual(uid, 4)
+
+    def testSearchBatch(self):
+        """
+        Test cluster batch search
+        """
+
+        results = self.client.post(
+            "batchsearch",
+            json={
+                "queries": ["feel good story", "climate change"],
+                "limit": 1,
+                "weights": 0.5,
+                "index": "default",
+                "parameters": [{"x": 1}, {"x": 2}],
+                "graph": False,
+            },
+        ).json()
+
+        uids = [result[0]["id"] for result in results]
+        self.assertEqual(uids, [4, 1])
+
+    def testSQL(self):
+        """
+        Test cluster SQL statement
+        """
+
+        query = urllib.parse.quote("select count(*), min(indexid), max(indexid), avg(indexid) from txtai where text='This is a test'")
+        self.assertEqual(
+            self.client.get(f"search?query={query}").json(), [{"count(*)": 28, "min(indexid)": 0, "max(indexid)": 14, "avg(indexid)": 6.5}]
+        )
+
+        query = urllib.parse.quote("select count(*), text txt from txtai group by txt order by count(*) desc")
+        self.assertEqual(
+            self.client.get(f"search?query={query}").json(),
+            [{"count(*)": 28, "txt": "And another test"}, {"count(*)": 24, "txt": "This is a test"}],
+        )
+
+        query = urllib.parse.quote("select count(*), text from txtai group by text order by count(*) asc")
+        self.assertEqual(
+            self.client.get(f"search?query={query}").json(),
+            [{"count(*)": 24, "text": "This is a test"}, {"count(*)": 28, "text": "And another test"}],
+        )
+
+        query = urllib.parse.quote("select count(*) from txtai group by id order by count(*)")
+        self.assertEqual(self.client.get(f"search?query={query}").json(), [{"count(*)": 52}])
+
+    def testUpsert(self):
+        """
+        Test cluster upsert
+        """
+
+        # Update data
+        self.client.post("add", json=[{"id": 4, "text": "Feel good story: baby panda born"}])
+        self.client.get("upsert")
+
+        # Search for best match
+        query = "feel good story"
+        uid = self.client.get(f"search?query={query}&limit=1").json()[0]["id"]
+
+        self.assertEqual(uid, 4)

test/python/testapi/testencoding.py

@@ -0,0 +1,163 @@
+"""
+Encoding module tests
+"""
+
+import base64
+import os
+import tempfile
+import unittest
+import urllib.parse
+
+from io import BytesIO
+from unittest.mock import patch
+
+import msgpack
+import numpy as np
+import PIL
+
+from fastapi.testclient import TestClient
+
+from txtai.api import application
+from txtai.api.responses import JSONEncoder
+
+# pylint: disable=C0411
+from utils import Utils
+
+# Configuration for image storage
+INDEX = """
+# Allow indexing of documents
+writable: %s
+
+embeddings:
+    defaults: False
+    content: True
+    objects: %s
+"""
+
+
+class TestEncoding(unittest.TestCase):
+    """
+    API tests for response encoding
+    """
+
+    @staticmethod
+    @patch.dict(os.environ, {"CONFIG": os.path.join(tempfile.gettempdir(), "testapi.yml"), "API_CLASS": "txtai.api.API"})
+    def start(yaml):
+        """
+        Starts a mock FastAPI client.
+
+        Args:
+            yaml: input configuration
+        """
+
+        config = os.path.join(tempfile.gettempdir(), "testapi.yml")
+
+        with open(config, "w", encoding="utf-8") as output:
+            output.write(yaml)
+
+        # Create new application and set on client
+        application.app = application.create()
+        client = TestClient(application.app)
+        application.start()
+
+        return client
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create API client on creation of class.
+        """
+
+        cls.client = TestEncoding.start(INDEX % ("True", "image"))
+
+    def testImages(self):
+        """
+        Test image encoding
+        """
+
+        with open(Utils.PATH + "/books.jpg", "rb") as f:
+            self.client.post("addimage", data={"uid": 0}, files={"data": f})
+            self.client.get("index")
+
+        query = urllib.parse.quote_plus("select id, object from txtai limit 1")
+        results = self.client.get(f"search?query={query}").json()
+
+        # Test reading image
+        self.assertIsInstance(PIL.Image.open(BytesIO(base64.b64decode(results[0]["object"]))), PIL.Image.Image)
+
+    def testInvalidInputs(self):
+        """
+        Test invalid parameter inputs
+        """
+
+        response = self.client.post("addimage", data={"uid": [0, 1]}, files={"data": b"123"})
+        self.assertEqual(response.status_code, 422)
+
+        response = self.client.post("addobject", data={"uid": [0, 1]}, files={"data": b"123"})
+        self.assertEqual(response.status_code, 422)
+
+    def testInvalidJSON(self):
+        """
+        Test that invalid JSON raises an exception
+        """
+
+        with self.assertRaises(TypeError):
+            JSONEncoder().encode(np.random.rand(1, 1))
+
+    def testMessagePack(self):
+        """
+        Test message pack encoding
+        """
+
+        # Validate binary encoding
+        results = self.client.get("count", headers={"Accept": "application/msgpack"}).content
+        self.assertEqual(results, b"\x01")
+
+        # Validate query result
+        query = urllib.parse.quote_plus("select id, object from txtai limit 1")
+        results = self.client.get(f"search?query={query}", headers={"Accept": "application/msgpack"}).content
+        results = msgpack.unpackb(results)
+
+        # Test reading image
+        self.assertIsInstance(PIL.Image.open(BytesIO(results[0]["object"])), PIL.Image.Image)
+
+    def testObjects(self):
+        """
+        Test object encoding
+        """
+
+        # Recreate model with standard object encoding
+        self.client = TestEncoding.start(INDEX % ("True", "True"))
+
+        # Test various formats
+        self.client.post("addobject", data={"uid": "id0"}, files={"data": b"1234"})
+        self.client.post("addobject", files={"data": b"ABC"})
+        self.client.post("addobject", data={"uid": "id1", "field": "object"}, files={"data": b"A1234"})
+        self.client.get("index")
+
+        query = urllib.parse.quote_plus("select id, object from txtai where id = 'id0' limit 1")
+        results = self.client.get(f"search?query={query}").json()
+        self.assertEqual(base64.b64decode(results[0]["object"]), b"1234")
+
+        # Test with messagepack encoding
+        results = self.client.get(f"search?query={query}", headers={"Accept": "application/msgpack"}).content
+        results = msgpack.unpackb(results)
+        self.assertEqual(results[0]["object"], b"1234")
+
+        count = self.client.get("count").json()
+        self.assertEqual(count, 3)
+
+    def testReadOnly(self):
+        """
+        Test read only indexes
+        """
+
+        # Recreate model with standard object encoding
+        self.client = TestEncoding.start(INDEX % ("False", "True"))
+
+        # Test errors raised for write operations
+        with open(Utils.PATH + "/books.jpg", "rb") as f:
+            response = self.client.post("addimage", data={"uid": 0}, files={"data": f})
+            self.assertEqual(response.status_code, 403)
+
+        self.assertEqual(self.client.post("addobject", data={"uid": 0}, files={"data": b"1234"}).status_code, 403)

test/python/testapi/testextension.py

@@ -0,0 +1,117 @@
+"""
+Extension module tests
+"""
+
+import os
+import tempfile
+import unittest
+
+from unittest.mock import patch
+
+from fastapi import APIRouter
+from fastapi.testclient import TestClient
+
+from txtai.api import application, Extension
+from txtai.pipeline import Pipeline
+
+# Example pipeline extension
+PIPELINES = """
+testapi.testextension.SamplePipeline:
+"""
+
+
+class SampleRouter:
+    """
+    Sample API router.
+    """
+
+    router = APIRouter()
+
+    @staticmethod
+    @router.get("/sample")
+    def sample(text: str):
+        """
+        Calls sample pipeline.
+
+        Args:
+            text: input text
+
+        Returns:
+            formatted text
+        """
+
+        return application.get().pipeline("testapi.testextension.SamplePipeline", (text,))
+
+
+class SampleExtension(Extension):
+    """
+    Sample API extension.
+    """
+
+    def __call__(self, app):
+        app.include_router(SampleRouter().router)
+
+
+class SamplePipeline(Pipeline):
+    """
+    Sample pipeline.
+    """
+
+    def __call__(self, text):
+        return text.lower()
+
+
+class TestExtension(unittest.TestCase):
+    """
+    API tests for extensions.
+    """
+
+    @staticmethod
+    @patch.dict(
+        os.environ,
+        {
+            "CONFIG": os.path.join(tempfile.gettempdir(), "testapi.yml"),
+            "API_CLASS": "txtai.api.API",
+            "EXTENSIONS": "testapi.testextension.SampleExtension",
+        },
+    )
+    def start():
+        """
+        Starts a mock FastAPI client.
+        """
+
+        config = os.path.join(tempfile.gettempdir(), "testapi.yml")
+
+        with open(config, "w", encoding="utf-8") as output:
+            output.write(PIPELINES)
+
+        # Create new application and set on client
+        application.app = application.create()
+        client = TestClient(application.app)
+        application.start()
+
+        return client
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create API client on creation of class.
+        """
+
+        cls.client = TestExtension.start()
+
+    def testEmpty(self):
+        """
+        Test an empty extension
+        """
+
+        extension = Extension()
+        self.assertIsNone(extension(None))
+
+    def testExtension(self):
+        """
+        Test a pipeline extension
+        """
+
+        text = self.client.get("sample?text=Test%20String").json()
+        self.assertEqual(text, "test string")

test/python/testapi/testmcp.py

@@ -0,0 +1,59 @@
+"""
+Agent API module tests
+"""
+
+import os
+import tempfile
+import unittest
+
+from unittest.mock import patch
+
+from fastapi.testclient import TestClient
+
+from txtai.api import application
+
+# Configuration for agents
+MCP = """
+mcp: True
+"""
+
+
+# pylint: disable=R0904
+class TestMCP(unittest.TestCase):
+    """
+    API tests for model context protocol (MCP)
+    """
+
+    @staticmethod
+    @patch.dict(os.environ, {"CONFIG": os.path.join(tempfile.gettempdir(), "testapi.yml"), "API_CLASS": "txtai.api.API"})
+    def start():
+        """
+        Starts a mock FastAPI client.
+        """
+
+        config = os.path.join(tempfile.gettempdir(), "testapi.yml")
+
+        with open(config, "w", encoding="utf-8") as output:
+            output.write(MCP)
+
+        # Create new application and set on client
+        application.app = application.create()
+        client = TestClient(application.app)
+        application.start()
+
+        return client
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create API client on creation of class.
+        """
+
+        cls.client = TestMCP.start()
+
+    def testMCP(self):
+        """
+        Test that application a /mcp route
+        """
+
+        self.assertTrue(any(route.path == "/mcp" for route in self.client.app.routes))

test/python/testapi/testopenai.py

@@ -0,0 +1,201 @@
+"""
+OpenAI API module tests
+"""
+
+import os
+import tempfile
+import unittest
+
+from unittest.mock import patch
+
+from fastapi.testclient import TestClient
+
+from txtai.api import application
+
+# pylint: disable=C0411
+from utils import Utils
+
+# API Configuration
+CONFIG = """
+# Enable OpenAI-compatible API
+openai: True
+
+# Allow indexing of documents
+writable: True
+
+# Agent configuration
+agent:
+    hello:
+        max_iterations: 1
+
+# Embeddings settings
+embeddings:
+    path: sentence-transformers/nli-mpnet-base-v2
+    content: True        
+
+# LLM configuration
+llm:
+    path: hf-internal-testing/tiny-random-LlamaForCausalLM
+
+# Text segmentation
+segmentation:
+
+# Text to speech
+texttospeech:
+
+# Transcription
+transcription:
+
+# Workflow
+workflow:
+    echo:
+        tasks:
+            - task: console
+"""
+
+
+# pylint: disable=R0904
+class TestOpenAI(unittest.TestCase):
+    """
+    Tests for OpenAI-compatible API endpoint for txtai.
+    """
+
+    @staticmethod
+    @patch.dict(os.environ, {"CONFIG": os.path.join(tempfile.gettempdir(), "testopenai.yml"), "API_CLASS": "txtai.api.API"})
+    def start():
+        """
+        Starts a mock FastAPI client.
+        """
+
+        config = os.path.join(tempfile.gettempdir(), "testopenai.yml")
+
+        with open(config, "w", encoding="utf-8") as output:
+            output.write(CONFIG)
+
+        # Create new application and set on client
+        application.app = application.create()
+        client = TestClient(application.app)
+        application.start()
+
+        # Patch LLM to generate answer
+        agent = application.get().agents["hello"]
+        agent.process.model.llm = lambda *args, **kwargs: 'Action:\n{"name": "final_answer", "arguments": "Hi"}'
+
+        return client
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create API client on creation of class.
+        """
+
+        cls.client = TestOpenAI.start()
+
+        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",
+        ]
+
+        # Index data
+        cls.client.post("add", json=[{"id": x, "text": row} for x, row in enumerate(cls.data)])
+        cls.client.get("index")
+
+    def testChatAgent(self):
+        """
+        Test a chat completion with an agent
+        """
+
+        response = self.client.post("/v1/chat/completions", json={"messages": [{"role": "user", "content": "Hello"}], "model": "hello"}).json()
+
+        self.assertEqual(response["choices"][0]["message"]["content"], "Hi")
+
+    def testChatLLM(self):
+        """
+        Test a chat completion with a LLM
+        """
+
+        response = self.client.post("/v1/chat/completions", json={"messages": [{"role": "user", "content": "Hello"}], "model": "llm"}).json()
+
+        self.assertIsNotNone(response["choices"][0]["message"]["content"])
+
+    def testChatPipeline(self):
+        """
+        Test a chat completion with a pipeline
+        """
+
+        response = self.client.post("/v1/chat/completions", json={"messages": [{"role": "user", "content": "Hello"}], "model": "segmentation"}).json()
+
+        self.assertEqual(response["choices"][0]["message"]["content"], "Hello")
+
+    def testChatSearch(self):
+        """
+        Test a chat completion with an embeddings search
+        """
+
+        response = self.client.post(
+            "/v1/chat/completions", json={"messages": [{"role": "user", "content": "feel good story"}], "model": "embeddings"}
+        ).json()
+
+        self.assertEqual(response["choices"][0]["message"]["content"], self.data[4])
+
+    def testChatStream(self):
+        """
+        Test a chat completion with a LLM
+        """
+
+        response = self.client.post("/v1/chat/completions", json={"messages": [{"role": "user", "content": "Hello"}], "model": "llm", "stream": True})
+
+        self.assertGreater(len(response.text.split("\n\n")), 0)
+
+    def testChatWorkflow(self):
+        """
+        Test a chat completion with a workflow
+        """
+
+        response = self.client.post("/v1/chat/completions", json={"messages": [{"role": "user", "content": "Hello"}], "model": "echo"}).json()
+
+        self.assertEqual(response["choices"][0]["message"]["content"], "Hello")
+
+    def testEmbeddings(self):
+        """
+        Test generating embeddings vectors
+        """
+
+        response = self.client.post("/v1/embeddings", json={"input": "text to embed", "model": "nli-mpnet-base-v2"}).json()
+
+        self.assertEqual(len(response["data"][0]["embedding"]), 768)
+
+    def testSpeech(self):
+        """
+        Test generating speech for input text
+        """
+
+        response = self.client.post(
+            "/v1/audio/speech", json={"model": "tts", "input": "text to speak", "voice": "default", "response_format": "wav"}
+        ).content
+
+        self.assertTrue(response[0:4] == b"RIFF")
+
+    def testTranscribe(self):
+        """
+        Test audio to text transcription
+        """
+
+        path = Utils.PATH + "/Make_huge_profits.wav"
+        with open(path, "rb") as f:
+            text = self.client.post("/v1/audio/transcriptions", files={"file": f}).json()["text"]
+            self.assertEqual(text, "Make huge profits without working make up to one hundred thousand dollars a day")
+
+    def testTranslate(self):
+        """
+        Test audio translation
+        """
+
+        path = Utils.PATH + "/Make_huge_profits.wav"
+        with open(path, "rb") as f:
+            text = self.client.post("/v1/audio/translations", files={"file": f}).json()["text"]
+            self.assertEqual(text, "Make huge profits without working make up to one hundred thousand dollars a day")