Add txtai console, closes #252

.coveragerc

@@ -2,6 +2,7 @@
 source = src/python
 concurrency = multiprocessing,thread
 disable_warnings = no-data-collected
+omit = **/__main__.py
 
 [combine]
 disable_warnings = no-data-collected

examples/console.py

@@ -1,83 +0,0 @@
-"""
-txtai console module.
-
-Requires tabulate to be installed.
-  pip install tabulate
-"""
-
-import shutil
-import sys
-import textwrap
-
-from cmd import Cmd
-
-from tabulate import tabulate
-from txtai.embeddings import Embeddings
-
-
-class Console(Cmd):
-    """
-    txtai console.
-    """
-
-    def __init__(self, path):
-        super().__init__()
-
-        self.intro = "txtai console"
-        self.prompt = ">>> "
-
-        self.embeddings = None
-        self.path = path
-
-    def preloop(self):
-        print(f"Loading model {self.path}")
-
-        # Load model
-        self.embeddings = Embeddings()
-        self.embeddings.load(self.path)
-
-    def default(self, line):
-        # pylint: disable=W0703
-        try:
-            results = []
-            maxwidth = None
-            for result in self.embeddings.search(line):
-                # Calculate max width using current terminal width and number of columns
-                if not maxwidth:
-                    maxwidth = int(shutil.get_terminal_size()[0] / len(result))
-
-                # Wrap each value at maxwidth, if necessary
-                results.append({key: self.wrap(value, maxwidth) for key, value in result.items()})
-
-            print(tabulate(results, headers="keys", tablefmt="psql"))
-        except Exception as ex:
-            print(ex)
-
-    def wrap(self, value, maxwidth):
-        """
-        Wraps value at maxwidth if value is a string.
-
-        Args:
-            value: input value
-            maxwidth: maximum number of characters before splitting text
-
-        Returns:
-            newline wrapped text at maxwidth
-        """
-
-        return "\n".join(textwrap.wrap(value, maxwidth)) if isinstance(value, str) else value
-
-
-def main(path=None):
-    """
-    Console execution loop.
-
-    Args:
-        path: model path
-    """
-
-    Console(path).cmdloop()
-
-
-if __name__ == "__main__":
-    main(sys.argv[1] if len(sys.argv) > 1 else None)

setup.py

@@ -21,6 +21,8 @@ extras["api"] = [
 
 extras["cloud"] = ["apache-libcloud>=3.3.1"]
 
+extras["console"] = ["rich>=12.0.1"]
+
 extras["database"] = ["pillow>=9.0.1"]
 
 extras["model"] = ["onnxruntime>=1.8.1"]
@@ -61,7 +63,14 @@ extras["workflow"] = [
 ]
 
 extras["all"] = (
-    extras["api"] + extras["cloud"] + extras["database"] + extras["model"] + extras["pipeline"] + extras["similarity"] + extras["workflow"]
+    extras["api"]
+    + extras["cloud"]
+    + extras["console"]
+    + extras["database"]
+    + extras["model"]
+    + extras["pipeline"]
+    + extras["similarity"]
+    + extras["workflow"]
 )
 
 setup(

src/python/txtai/console/__init__.py

@@ -0,0 +1,5 @@
+"""
+Console imports
+"""
+
+from .base import Console

src/python/txtai/console/__main__.py

@@ -0,0 +1,22 @@
+"""
+Main module.
+"""
+
+import sys
+
+from .base import Console
+
+
+def main(path=None):
+    """
+    Console execution loop.
+
+    Args:
+        path: model path
+    """
+
+    Console(path).cmdloop()
+
+
+if __name__ == "__main__":
+    main(sys.argv[1] if len(sys.argv) > 1 else None)

src/python/txtai/console/base.py

@@ -0,0 +1,235 @@
+"""
+Console module
+"""
+
+import os
+import shlex
+
+from cmd import Cmd
+
+from rich import box
+from rich.console import Console as RichConsole
+from rich.table import Table
+
+from txtai.app import Application
+from txtai.embeddings import Embeddings
+
+
+class Console(Cmd):
+    """
+    txtai console.
+    """
+
+    def __init__(self, path=None):
+        """
+        Creates a new command line console.
+
+        Args:
+            path: path to initial configuration, if any
+        """
+
+        super().__init__()
+
+        self.prompt = ">>> "
+
+        # Rich console
+        self.console = RichConsole()
+
+        # App parameters
+        self.app = None
+        self.path = path
+
+        # Parameters
+        self.vhighlight = None
+        self.vlimit = None
+
+    def preloop(self):
+        """
+        Loads initial configuration.
+        """
+
+        self.console.print("txtai console", style="#03a9f4")
+
+        # Load default path
+        if self.path:
+            self.load(self.path)
+
+    def default(self, line):
+        """
+        Default event loop.
+
+        Args:
+            line: command line
+        """
+
+        # pylint: disable=W0703
+        try:
+            command = line.lower()
+            if command.startswith(".config"):
+                self.config()
+            elif command.startswith(".highlight"):
+                self.highlight(command)
+            elif command.startswith(".limit"):
+                self.limit(command)
+            elif command.startswith(".load"):
+                command = shlex.split(line)
+                self.path = command[1]
+                self.load(self.path)
+            elif command.startswith(".workflow"):
+                self.workflow(line)
+            else:
+                # Search is default action
+                self.search(line)
+        except Exception:
+            self.console.print_exception()
+
+    def config(self):
+        """
+        Processes .config command.
+        """
+
+        self.console.print(self.app.config)
+
+    def load(self, path):
+        """
+        Processes .load command.
+
+        Args:
+            path: path to configuration
+        """
+
+        if os.path.isfile(path):
+            self.console.print(f"Loading application {path}")
+            self.app = Application(path)
+        else:
+            self.console.print(f"Loading index {path}")
+
+            # Load embeddings index
+            self.app = Embeddings()
+            self.app.load(path)
+
+    def highlight(self, command):
+        """
+        Processes .highlight command.
+
+        Args:
+            command: command line
+        """
+
+        _, action = self.split(command, "#ffff00")
+        self.vhighlight = action
+        self.console.print(f"Set highlight to {self.vhighlight}")
+
+    def limit(self, command):
+        """
+        Processes .limit command.
+
+        Args:
+            command: command line
+        """
+
+        _, action = self.split(command, 10)
+        self.vlimit = int(action)
+        self.console.print(f"Set limit to {self.vlimit}")
+
+    def workflow(self, command):
+        """
+        Processes .workflow command.
+
+        Args:
+            command: command line
+        """
+
+        command = shlex.split(command)
+        if isinstance(self.app, Application):
+            self.console.print(list(self.app.workflow(command[1], command[2:])))
+
+    def split(self, command, default):
+        """
+        Splits command by whitespace.
+
+        Args:
+            command: command line
+            default: default command action
+
+        Returns:
+            command action
+        """
+
+        values = command.split(" ", 1)
+        return values if len(values) > 1 else (command, default)
+
+    def search(self, query):
+        """
+        Runs a search query.
+
+        Args:
+            query: query to run
+        """
+
+        if self.vhighlight:
+            results = self.app.explain(query, limit=self.vlimit)
+        else:
+            results = self.app.search(query, limit=self.vlimit)
+
+        columns, table = {}, Table(box=box.SQUARE, style="#03a9f4")
+
+        # Build column list
+        result = results[0]
+        if isinstance(result, tuple):
+            columns = dict.fromkeys(["id", "score"])
+        else:
+            columns = dict(result)
+
+        # Add columns to table
+        columns = list(x for x in columns if x != "tokens")
+        for column in columns:
+            table.add_column(column)
+
+        # Add rows to table
+        for result in results:
+            if isinstance(result, tuple):
+                table.add_row(*(self.render(result, None, x) for x in result))
+            else:
+                table.add_row(*(self.render(result, column, result.get(column)) for column in columns))
+
+        # Print table to console
+        self.console.print(table)
+
+    def render(self, result, column, value):
+        """
+        Renders a search result column value.
+
+        Args:
+            result: result row
+            column: column name
+            value: column value
+        """
+
+        if isinstance(value, float):
+            return f"{value:.4f}"
+
+        # Explain highlighting
+        if column == "text" and "tokens" in result:
+            spans = []
+            for token, score in result["tokens"]:
+                color = None
+                if score >= 0.02:
+                    color = f"b {self.vhighlight}"
+
+                spans.append((token, score, color))
+
+            if result["score"] >= 0.05 and not [color for _, _, color in spans if color]:
+                mscore = max([score for _, score, _ in spans])
+                spans = [(token, score, f"b {self.vhighlight}" if score == mscore else color) for token, score, color in spans]
+
+            output = ""
+            for token, _, color in spans:
+                if color:
+                    output += f"[{color}]{token}[/{color}] "
+                else:
+                    output += f"{token} "
+
+            return output
+
+        return str(value)

test/python/testconsole.py

@@ -0,0 +1,174 @@
+"""
+Console module tests
+"""
+
+import contextlib
+import io
+import os
+import tempfile
+import unittest
+
+from txtai.console import Console
+from txtai.embeddings import Embeddings
+
+
+class TestConsole(unittest.TestCase):
+    """
+    Console 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",
+        ]
+
+        # Create embeddings model, backed by sentence-transformers & transformers
+        cls.embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "content": True})
+
+        # Create an index for the list of text
+        cls.embeddings.index([(uid, text, None) for uid, text in enumerate(cls.data)])
+
+        # Create app paths
+        cls.apppath = os.path.join(tempfile.gettempdir(), "console.yml")
+        cls.embedpath = os.path.join(tempfile.gettempdir(), "embeddings")
+
+        # Create app.yml
+        with open(cls.apppath, "w", encoding="utf-8") as out:
+            out.write(
+                """
+                path: /tmp/embeddings
+
+                workflow:
+                    test:
+                        tasks:
+                            - task: console
+            """
+            )
+
+        # Save index
+        cls.embeddings.save(cls.embedpath)
+
+        # Create console
+        cls.app = os.path.join(cls.apppath)
+        cls.console = Console(cls.embedpath)
+
+    def testApplication(self):
+        """
+        Test application
+        """
+
+        self.assertIn("console.yml", self.command(f".load {self.apppath}"))
+        self.assertIn("1", self.command(".limit 1"))
+        self.assertIn("Maine man wins", self.command("feel good story"))
+
+    def testConfig(self):
+        """
+        Test .config command
+        """
+
+        self.assertIn("tasks", self.command(".config"))
+
+    def testEmbeddings(self):
+        """
+        Test embeddings index
+        """
+
+        self.assertIn("embeddings", self.command(f".load {self.embedpath}"))
+        self.assertIn("1", self.command(".limit 1"))
+        self.assertIn("Maine man wins", self.command("feel good story"))
+
+    def testEmbeddingsNoDatabase(self):
+        """
+        Test embeddings with no database/content
+        """
+
+        console = Console()
+
+        # Create embeddings model, backed by sentence-transformers & transformers
+        embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2"})
+
+        # Create an index for the list of text
+        embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
+
+        # Set embeddings on console
+        console.app = embeddings
+        self.assertIn("4", self.command("feel good story", console))
+
+    def testEmpty(self):
+        """
+        Test empty console instance
+        """
+
+        console = Console()
+        self.assertIn("AttributeError", self.command("search", console))
+
+    def testHighlight(self):
+        """
+        Test .highlight command
+        """
+
+        self.assertIn("highlight", self.command(".highlight"))
+        self.assertIn("wins", self.command("feel good story"))
+        self.assertIn("Taiwan", self.command("asia"))
+
+    def testPreloop(self):
+        """
+        Test preloop
+        """
+
+        self.assertIn("txtai console", self.preloop())
+
+    def testWorkflow(self):
+        """
+        Test .workflow command
+        """
+
+        self.command(f".load {self.apppath}")
+        self.assertIn("echo", self.command(".workflow test echo"))
+
+    def command(self, command, console=None):
+        """
+        Runs a console command.
+
+        Args:
+            command: command to run
+            console: console instance, defaults to self.console
+
+        Returns:
+            command output
+        """
+
+        # Run info
+        output = io.StringIO()
+        with contextlib.redirect_stdout(output):
+            if not console:
+                console = self.console
+
+            console.onecmd(command)
+
+        return output.getvalue()
+
+    def preloop(self):
+        """
+        Runs console.preloop and redirects stdout.
+
+        Returns:
+            preloop output
+        """
+
+        # Run info
+        output = io.StringIO()
+        with contextlib.redirect_stdout(output):
+            self.console.preloop()
+
+        return output.getvalue()

test/python/testdatabase/testembeddings.py

@@ -205,6 +205,13 @@ class TestEmbeddings(unittest.TestCase):
         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 testGenerator(self):
         """
         Test index with a generator