chore: import upstream snapshot with attribution

test/python/testpipeline/testllm/testgenerator.py

@@ -0,0 +1,24 @@
+"""
+Generator module tests
+"""
+
+import unittest
+
+from txtai.pipeline import Generator
+
+
+class TestGenerator(unittest.TestCase):
+    """
+    Sequences tests.
+    """
+
+    def testGeneration(self):
+        """
+        Test text pipeline generation
+        """
+
+        model = Generator("hf-internal-testing/tiny-random-gpt2")
+        start = "Hello, how are"
+
+        # Test that text is generated
+        self.assertIsNotNone(model(start))

test/python/testpipeline/testllm/testlitellm.py

@@ -0,0 +1,115 @@
+"""
+LiteLLM module tests
+"""
+
+import json
+import os
+import time
+import unittest
+import uuid
+
+from unittest.mock import patch
+
+from http.server import HTTPServer, BaseHTTPRequestHandler
+from threading import Thread
+
+from txtai.pipeline import LLM
+
+
+class RequestHandler(BaseHTTPRequestHandler):
+    """
+    Test HTTP handler.
+    """
+
+    def do_POST(self):
+        """
+        POST request handler.
+        """
+
+        # Parse input headers
+        length = int(self.headers["content-length"])
+        data = json.loads(self.rfile.read(length))
+
+        if data.get("stream"):
+            # Mock streaming response
+            content = "application/octet-stream"
+            response = (
+                "data: "
+                + json.dumps(
+                    {
+                        "id": str(uuid.uuid4()),
+                        "object": "chat.completion.chunk",
+                        "created": int(time.time() * 1000),
+                        "model": "test",
+                        "choices": [{"id": 0, "delta": {"content": "blue"}}],
+                    }
+                )
+                + "\n\ndata: [DONE]\n\n"
+            )
+        else:
+            # Mock standard response
+            content = "application/json"
+            response = json.dumps(
+                {
+                    "id": str(uuid.uuid4()),
+                    "object": "chat.completion",
+                    "created": int(time.time() * 1000),
+                    "model": "test",
+                    "choices": [{"id": 0, "message": {"role": "assistant", "content": "blue"}, "finish_reason": "stop"}],
+                }
+            )
+
+        # Encode response as bytes
+        response = response.encode("utf-8")
+
+        self.send_response(200)
+        self.send_header("content-type", content)
+        self.send_header("content-length", len(response))
+        self.end_headers()
+
+        self.wfile.write(response)
+        self.wfile.flush()
+
+
+class TestLiteLLM(unittest.TestCase):
+    """
+    LiteLLM tests.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create mock http server.
+        """
+
+        cls.httpd = HTTPServer(("127.0.0.1", 8000), RequestHandler)
+
+        server = Thread(target=cls.httpd.serve_forever, daemon=True)
+        server.start()
+
+    @classmethod
+    def tearDownClass(cls):
+        """
+        Shutdown mock http server.
+        """
+
+        cls.httpd.shutdown()
+
+    @patch.dict(os.environ, {"OPENAI_API_KEY": "test"})
+    def testGeneration(self):
+        """
+        Test generation with LiteLLM
+        """
+
+        # Test model generation with LiteLLM
+        model = LLM("openai/gpt-4o", api_base="http://127.0.0.1:8000")
+        self.assertEqual(model("The sky is"), "blue")
+
+        # Test default role
+        self.assertEqual(model("The sky is", defaultrole="user"), "blue")
+
+        # Test streaming
+        self.assertEqual(" ".join(x for x in model("The sky is", stream=True)), "blue")
+
+        # Test vision
+        self.assertEqual(model.isvision(), False)

test/python/testpipeline/testllm/testlitert.py

@@ -0,0 +1,31 @@
+"""
+LiteRT module tests
+"""
+
+import unittest
+
+from txtai.pipeline import LLM
+
+
+class TestLiteRT(unittest.TestCase):
+    """
+    LiteRT tests.
+    """
+
+    def testGeneration(self):
+        """
+        Test generation with LiteRT
+        """
+
+        # Test model generation with LiteRT
+        model = LLM("neuml/gemma-4-tiny-random-litert-lm/gemma-4-tiny-random.litertlm", mtp=False, maxlength=25)
+
+        # Test standard
+        self.assertIsNotNone(model("Hello"))
+
+        # Test streaming
+        self.assertIsNotNone(list(model("Hello", stream=True)))
+
+        # Test CPU fallback
+        model = LLM("neuml/gemma-4-tiny-random-litert-lm/gemma-4-tiny-random.litertlm", mtp=True, maxlength=25)
+        self.assertIsNotNone(model("Hello"))

test/python/testpipeline/testllm/testllama.py

@@ -0,0 +1,76 @@
+"""
+Llama module tests
+"""
+
+import unittest
+
+from unittest.mock import patch
+
+from txtai.pipeline import LLM
+
+
+class TestLlama(unittest.TestCase):
+    """
+    llama.cpp tests.
+    """
+
+    @patch("llama_cpp.Llama")
+    def testContext(self, llama):
+        """
+        Test n_ctx with llama.cpp
+        """
+
+        class Llama:
+            """
+            Mock llama.cpp instance to test invalid context
+            """
+
+            def __init__(self, **kwargs):
+                if kwargs.get("n_ctx") == 0 or kwargs.get("n_ctx", 0) >= 10000:
+                    raise ValueError("Failed to create context")
+
+                # Save parameters
+                self.params = kwargs
+
+        # Mock llama.cpp instance
+        llama.side_effect = Llama
+
+        # Model to test
+        path = "TheBloke/TinyLlama-1.1B-Chat-v0.3-GGUF/tinyllama-1.1b-chat-v0.3.Q2_K.gguf"
+
+        # Test omitting n_ctx falls back to default settings
+        llm = LLM(path)
+        self.assertNotIn("n_ctx", llm.generator.llm.params)
+
+        # Test n_ctx=0 falls back to default settings
+        llm = LLM(path, n_ctx=0)
+        self.assertNotIn("n_ctx", llm.generator.llm.params)
+
+        # Test n_ctx manually set
+        llm = LLM(path, n_ctx=1024)
+        self.assertEqual(llm.generator.llm.params["n_ctx"], 1024)
+
+        # Mock a value for n_ctx that's too big
+        with self.assertRaises(ValueError):
+            llm = LLM(path, n_ctx=10000)
+
+    def testGeneration(self):
+        """
+        Test generation with llama.cpp
+        """
+
+        # Test model generation with llama.cpp
+        model = LLM("TheBloke/TinyLlama-1.1B-Chat-v0.3-GGUF/tinyllama-1.1b-chat-v0.3.Q2_K.gguf", chat_format="chatml")
+
+        # Test with prompt
+        self.assertEqual(model("2 + 2 = ", maxlength=10, seed=0, stop=["."], defaultrole="prompt")[0], "4")
+
+        # Test with list of messages
+        messages = [{"role": "system", "content": "You are a helpful assistant. You answer math problems."}, {"role": "user", "content": "2+2?"}]
+        self.assertIsNotNone(model(messages, maxlength=10, seed=0, stop=["."]))
+
+        # Test default role
+        self.assertIsNotNone(model("2 + 2 = ", maxlength=10, seed=0, stop=["."], defaultrole="user"))
+
+        # Test streaming
+        self.assertEqual(" ".join(x for x in model("2 + 2 = ", maxlength=10, stream=True, seed=0, stop=["."], defaultrole="prompt"))[0], "4")

test/python/testpipeline/testllm/testllm.py

@@ -0,0 +1,185 @@
+"""
+LLM module tests
+"""
+
+import unittest
+
+import torch
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from txtai.pipeline import LLM, Generation
+
+# pylint: disable=C0411
+from utils import Utils
+
+
+class TestLLM(unittest.TestCase):
+    """
+    LLM tests.
+    """
+
+    def testArguments(self):
+        """
+        Test pipeline keyword arguments
+        """
+
+        start = "Hello, how are"
+
+        # Test that text is generated with custom parameters
+        model = LLM("hf-internal-testing/tiny-random-gpt2", task="language-generation", dtype="torch.float32")
+        self.assertIsNotNone(model(start))
+
+        model = LLM("hf-internal-testing/tiny-random-gpt2", task="language-generation", dtype=torch.float32)
+        self.assertIsNotNone(model(start))
+
+    def testBatchSize(self):
+        """
+        Test batch size
+        """
+
+        model = LLM("sshleifer/tiny-gpt2")
+        self.assertIsNotNone(model(["Hello, how are"] * 2, batch_size=2))
+
+    def testCustom(self):
+        """
+        Test custom LLM framework
+        """
+
+        model = LLM("hf-internal-testing/tiny-random-gpt2", task="language-generation", method="txtai.pipeline.HFGeneration")
+        self.assertIsNotNone(model("Hello, how are"))
+
+    def testCustomNotFound(self):
+        """
+        Test resolving an unresolvable LLM framework
+        """
+
+        with self.assertRaises(ImportError):
+            LLM("hf-internal-testing/tiny-random-gpt2", method="notfound.generation")
+
+    def testDefaultRole(self):
+        """
+        Test default role
+        """
+
+        model = LLM("hf-internal-testing/tiny-random-LlamaForCausalLM")
+        generator = model.generator
+
+        # Validate that the LLM supports chat messages
+        self.assertEqual(model.ischat(), True)
+
+        messages = [
+            ("Hello", list),
+            ("\n<|im_start|>Hello<|im_end|>", str),
+            ("<|start|>Hello<|end|>", str),
+            ("<|start_of_role|>system<|end_of_role|>", str),
+            ("[INST]Hello[/INST]", str),
+        ]
+
+        for message, expected in messages:
+            # Test auto detection of formats
+            self.assertEqual(type(generator.format([message], "auto")[0]), expected)
+
+            # Test always setting user chat messages
+            self.assertEqual(type(generator.format([message], "user")[0]), list)
+
+            # Test always keeping as prompt text
+            self.assertEqual(type(generator.format([message], "prompt")[0]), str)
+
+    def testExternal(self):
+        """
+        Test externally loaded model
+        """
+
+        model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2")
+        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
+
+        model = LLM((model, tokenizer), template="{text}")
+        start = "Hello, how are"
+
+        # Test that text is generated
+        self.assertIsNotNone(model(start))
+
+    def testMaxLength(self):
+        """
+        Test max length
+        """
+
+        model = LLM("sshleifer/tiny-gpt2")
+        self.assertIsInstance(model("Hello, how are", maxlength=10), str)
+
+    def testNotImplemented(self):
+        """
+        Test exceptions for non-implemented methods
+        """
+
+        generation = Generation()
+        self.assertRaises(NotImplementedError, generation.stream, None, None, None, None)
+
+    def testStop(self):
+        """
+        Test stop strings
+        """
+
+        model = LLM("sshleifer/tiny-gpt2")
+        self.assertIsNotNone(model("Hello, how are", stop=["you"]))
+
+    def testStream(self):
+        """
+        Test streaming generation
+        """
+
+        model = LLM("sshleifer/tiny-gpt2")
+        self.assertIsInstance(" ".join(x for x in model("Hello, how are", stream=True)), str)
+
+    def testStripThink(self):
+        """
+        Test stripthink parameter
+        """
+
+        # pylint: disable=W0613
+        def execute1(*args, **kwargs):
+            return ["<think>test</think>you"]
+
+        def execute2(*args, **kwargs):
+            return ["<|channel|>final<|message|> you"]
+
+        model = LLM("hf-internal-testing/tiny-random-LlamaForCausalLM")
+
+        for method in [execute1, execute2]:
+            # Override execute method
+            model.generator.execute = method
+            self.assertEqual(model("Hello, how are", stripthink=True), "you")
+            self.assertEqual(model("Hello, how are", stripthink=False), method()[0])
+
+    def testStripThinkStream(self):
+        """
+        Test stripthink parameter with streaming output
+        """
+
+        # pylint: disable=W0613
+        def execute1(*args, **kwargs):
+            yield from "<think>test</think>you"
+
+        def execute2(*args, **kwargs):
+            yield from "<|channel|>final<|message|>you"
+
+        model = LLM("hf-internal-testing/tiny-random-LlamaForCausalLM")
+
+        for method in [execute1, execute2]:
+            # Override execute method
+            model.generator.execute = method
+            self.assertEqual("".join(model("Hello, how are", stripthink=True, stream=True)), "you")
+            self.assertEqual("".join(model("Hello, how are", stripthink=False, stream=True)), "".join(list(method())))
+
+    def testVision(self):
+        """
+        Test vision LLM
+        """
+
+        model = LLM("neuml/tiny-random-qwen2vl")
+        result = model(
+            [{"role": "user", "content": [{"type": "text", "text": "What is in this image?"}, {"type": "image", "image": Utils.PATH + "/books.jpg"}]}]
+        )
+
+        self.assertIsNotNone(result)

test/python/testpipeline/testllm/testopencode.py

@@ -0,0 +1,71 @@
+"""
+OpenCode module tests
+"""
+
+import json
+import unittest
+
+from http.server import HTTPServer, BaseHTTPRequestHandler
+from threading import Thread
+
+from txtai.pipeline import LLM
+
+
+class RequestHandler(BaseHTTPRequestHandler):
+    """
+    Test HTTP handler.
+    """
+
+    def do_POST(self):
+        """
+        POST request handler.
+        """
+
+        # Mock response
+        content = "application/json"
+        response = json.dumps({"id": "0", "parts": [{"type": "text", "text": "blue"}]})
+
+        # Encode response as bytes
+        response = response.encode("utf-8")
+
+        self.send_response(200)
+        self.send_header("content-type", content)
+        self.send_header("content-length", len(response))
+        self.end_headers()
+
+        self.wfile.write(response)
+        self.wfile.flush()
+
+
+class TestOpenCode(unittest.TestCase):
+    """
+    OpenCode tests.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create mock http server.
+        """
+
+        cls.httpd = HTTPServer(("127.0.0.1", 8005), RequestHandler)
+
+        server = Thread(target=cls.httpd.serve_forever, daemon=True)
+        server.start()
+
+    @classmethod
+    def tearDownClass(cls):
+        """
+        Shutdown mock http server.
+        """
+
+        cls.httpd.shutdown()
+
+    def testGeneration(self):
+        """
+        Test generation with OpenCode
+        """
+
+        # Test model generation with LiteLLM
+        model = LLM("opencode/big-pickle", url="http://127.0.0.1:8005")
+        self.assertEqual(model("The sky is"), "blue")

test/python/testpipeline/testllm/testrag.py

@@ -0,0 +1,225 @@
+"""
+RAG module tests
+"""
+
+import platform
+import unittest
+
+from txtai.embeddings import Embeddings
+from txtai.pipeline import Questions, RAG, Similarity
+
+
+class TestRAG(unittest.TestCase):
+    """
+    RAG tests.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create single rag instance.
+        """
+
+        cls.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",
+        ]
+
+        # Create embeddings model, backed by sentence-transformers & transformers
+        cls.embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2"})
+
+        # Create rag instance
+        cls.rag = RAG(cls.embeddings, "distilbert-base-cased-distilled-squad")
+
+    @classmethod
+    def tearDownClass(cls):
+        """
+        Cleanup data.
+        """
+
+        if cls.embeddings:
+            cls.embeddings.close()
+
+    def testAnswer(self):
+        """
+        Test qa extraction with an answer
+        """
+
+        questions = ["What team won the game?", "What was score?"]
+
+        # pylint: disable=C3001
+        execute = lambda query: self.rag([(question, query, question, False) for question in questions], self.data)
+
+        answers = execute("Red Sox - Blue Jays")
+        self.assertEqual("Blue Jays", answers[0][1])
+        self.assertEqual("2-1", answers[1][1])
+
+        # Ad-hoc questions
+        question = "What hockey team won?"
+
+        answers = self.rag([(question, question, question, False)], self.data)
+        self.assertEqual("Flyers", answers[0][1])
+
+    def testEmptyQuery(self):
+        """
+        Test an empty queries list
+        """
+
+        self.assertEqual(self.rag.query(None, None), [])
+
+    def testNoAnswer(self):
+        """
+        Test qa extraction with no answer
+        """
+
+        question = ""
+
+        answers = self.rag([(question, question, question, False)], self.data)
+        self.assertIsNone(answers[0][1])
+
+        question = "abcdef"
+        answers = self.rag([(question, question, question, False)], self.data)
+        self.assertIsNone(answers[0][1])
+
+    @unittest.skipIf(platform.system() == "Darwin", "Quantized models not supported on macOS")
+    def testQuantize(self):
+        """
+        Test qa extraction backed by a quantized model
+        """
+
+        rag = RAG(self.embeddings, "distilbert-base-cased-distilled-squad", True)
+
+        question = "How many home runs?"
+
+        answers = rag([(question, question, question, True)], self.data)
+        self.assertIsNotNone(answers[0][1])
+
+    def testOutputs(self):
+        """
+        Test output formatting rules
+        """
+
+        question = "How many home runs?"
+
+        # Test flatten to list of answers
+        rag = RAG(self.embeddings, "distilbert-base-cased-distilled-squad", output="flatten")
+        answers = rag([(question, question, question, True)], self.data)
+        self.assertTrue(answers[0].startswith("Giants hit 3 HRs"))
+
+        # Test reference field
+        rag = RAG(self.embeddings, "distilbert-base-cased-distilled-squad", output="reference")
+        answers = rag([(question, question, question, True)], self.data)
+        self.assertTrue(self.data[answers[0][2]].startswith("Giants hit 3 HRs"))
+
+    def testPrompt(self):
+        """
+        Test a user prompt with templating
+        """
+
+        embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "content": True})
+        embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
+
+        rag = RAG(
+            embeddings,
+            "google/flan-t5-small",
+            template="""
+              Answer the following question and return a number.
+              Question: {question}
+              Context:{context}""",
+            output="flatten",
+        )
+
+        self.assertEqual(rag("How many HRs"), "3")
+
+    def testPromptTemplates(self):
+        """
+        Test system and user prompt templates
+        """
+
+        rag = RAG(
+            self.embeddings,
+            "sshleifer/tiny-gpt2",
+            system="You are a friendly assistant",
+            template="""
+              Answer the following question and return a number.
+              Question: {question}
+              Context:{context}""",
+        )
+
+        prompts = rag.prompts(["How many HRs?"], [self.data])[0]
+        self.assertEqual([x["role"] for x in prompts], ["system", "user"])
+
+    def testSearch(self):
+        """
+        Test qa extraction with an embeddings search for context
+        """
+
+        embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "content": True})
+        embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
+
+        rag = RAG(embeddings, "distilbert-base-cased-distilled-squad")
+
+        question = "How many home runs?"
+
+        answers = rag([(question, question, question, True)])
+        self.assertTrue(answers[0][1].startswith("Giants hit 3 HRs"))
+
+    def testSimilarity(self):
+        """
+        Test qa extraction using a Similarity pipeline to build context
+        """
+
+        # Create rag instance
+        rag = RAG(Similarity("prajjwal1/bert-medium-mnli"), Questions("distilbert-base-cased-distilled-squad"))
+
+        question = "How many home runs?"
+
+        answers = rag([(question, "HRs", question, True)], self.data)
+        self.assertTrue(answers[0][1].startswith("Giants hit 3 HRs"))
+
+    def testSnippet(self):
+        """
+        Test qa extraction with a full answer snippet
+        """
+
+        question = "How many home runs?"
+
+        answers = self.rag([(question, question, question, True)], self.data)
+        self.assertTrue(answers[0][1].startswith("Giants hit 3 HRs"))
+
+    def testSnippetEmpty(self):
+        """
+        Test snippet method can handle empty parameters
+        """
+
+        self.assertEqual(self.rag.snippets(["name"], [None], [None], [None]), [("name", None)])
+
+    def testStringInput(self):
+        """
+        Test with single string input
+        """
+
+        result = self.rag("How many home runs?", self.data)
+        self.assertEqual(result["answer"], "3")
+
+    def testTasks(self):
+        """
+        Test loading models with task parameter
+        """
+
+        for task, model in [
+            ("language-generation", "hf-internal-testing/tiny-random-gpt2"),
+            ("sequence-sequence", "hf-internal-testing/tiny-random-t5"),
+        ]:
+            rag = RAG(self.embeddings, model, task=task)
+            self.assertIsNotNone(rag)

test/python/testpipeline/testllm/testsequences.py

@@ -0,0 +1,21 @@
+"""
+Sequences module tests
+"""
+
+import unittest
+
+from txtai.pipeline import Sequences
+
+
+class TestSequences(unittest.TestCase):
+    """
+    Sequences tests.
+    """
+
+    def testGeneration(self):
+        """
+        Test text2text pipeline generation
+        """
+
+        model = Sequences("t5-small")
+        self.assertEqual(model("Testing the model", prefix="translate English to German: "), "Das Modell zu testen")