diff --git a/.coveragerc b/.coveragerc index 24834123..0cb653fa 100644 --- a/.coveragerc +++ b/.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 diff --git a/examples/console.py b/examples/console.py deleted file mode 100644 index f3644bb4..00000000 --- a/examples/console.py +++ /dev/null @@ -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) diff --git a/setup.py b/setup.py index ec2745b2..4b81cb15 100644 --- a/setup.py +++ b/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( diff --git a/src/python/txtai/console/__init__.py b/src/python/txtai/console/__init__.py new file mode 100644 index 00000000..890e1e2b --- /dev/null +++ b/src/python/txtai/console/__init__.py @@ -0,0 +1,5 @@ +""" +Console imports +""" + +from .base import Console diff --git a/src/python/txtai/console/__main__.py b/src/python/txtai/console/__main__.py new file mode 100644 index 00000000..c0a474d6 --- /dev/null +++ b/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) diff --git a/src/python/txtai/console/base.py b/src/python/txtai/console/base.py new file mode 100644 index 00000000..855ff482 --- /dev/null +++ b/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) diff --git a/test/python/testconsole.py b/test/python/testconsole.py new file mode 100644 index 00000000..dbf36f1f --- /dev/null +++ b/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() diff --git a/test/python/testdatabase/testembeddings.py b/test/python/testdatabase/testembeddings.py index 191276aa..d8a65393 100644 --- a/test/python/testdatabase/testembeddings.py +++ b/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