Most tests pass for duckdb too

tests/test_tsellm.py

@@ -172,15 +172,6 @@ def setUp(self):
         )
 
 
-class InMemoryDuckDBTest(InMemorySQLiteTest):
-    def setUp(self):
-        super().setUp()
-        self.path_args = (
-            "--duckdb",
-            ":memory:",
-        )
-
-
 class DiskSQLiteTest(InMemorySQLiteTest):
     db_fp = None
     path_args = ()
@@ -199,5 +190,33 @@ def test_embed_default_hazo_leaves_valid_db_behind(self):
         self.assertTrue(TsellmConsole.is_sqlite(self.db_fp))
 
 
+class InMemoryDuckDBTest(InMemorySQLiteTest):
+    def setUp(self):
+        super().setUp()
+        self.path_args = (
+            "--duckdb",
+            ":memory:",
+        )
+
+    def test_duckdb_execute(self):
+        out = self.expect_success(*self.path_args, "select 'Hello World!'")
+        self.assertIn("('Hello World!',)", out)
+
+    def test_cli_execute_incomplete_sql(self):
+        pass
+
+    def test_cli_execute_too_much_sql(self):
+        pass
+
+    def test_embed_default_hazo(self):
+        pass
+
+    def test_prompt_default_markov(self):
+        pass
+
+    def test_embed_hazo_binary(self):
+        pass
+
+
 if __name__ == "__main__":
     unittest.main()

tsellm/cli.py

@@ -35,7 +35,12 @@ class TsellmConsole(ABC, InteractiveConsole):
 
 """
 
-    _functions = []
+    _functions = [
+        ("prompt", 2, _prompt_model, False),
+        ("prompt", 1, _prompt_model_default, False),
+        ("embed", 2, _embed_model, False),
+        ("embed", 1, _embed_model_default, False),
+    ]
 
     error_class = None
 
@@ -105,12 +110,7 @@ def runsource(self, source, filename="<input>", symbol="single"):
 
 class SQLiteConsole(TsellmConsole):
     error_class = sqlite3.Error
-    _functions = [
-        ("prompt", 2, _prompt_model, False),
-        ("prompt", 1, _prompt_model_default, False),
-        ("embed", 2, _embed_model, False),
-        ("embed", 1, _embed_model_default, False),
-    ]
+
 
     def __init__(self, path):
 
@@ -162,19 +162,64 @@ def runsource(self, source, filename="<input>", symbol="single"):
 
 
 class DuckDBConsole(TsellmConsole):
+    error_class = sqlite3.Error
+
+    _functions = [
+        ("prompt", 2, _prompt_model, False),
+        ("embed", 2, _embed_model, False),
+    ]
 
     def __init__(self, path):
         super().__init__()
-        self._con = duckdb.connect(str(path))
+        self._con = duckdb.connect(path)
         self._cur = self._con.cursor()
 
-        # self.load()
+        self.load()
+
+    def load(self):
+        self.execute(self._TSELLM_CONFIG_SQL)
+        for func_name, _, py_func, _ in self._functions:
+            self._con.create_function(func_name, py_func)
 
     def execute(self, sql, suppress_errors=True):
-        pass
+        """Helper that wraps execution of SQL code.
+
+        This is used both by the REPL and by direct execution from the CLI.
+
+        'c' may be a cursor or a connection.
+        'sql' is the SQL string to execute.
+        """
+
+        try:
+            for row in self._con.execute(sql).fetchall():
+                print(row)
+        except self.error_class as e:
+            tp = type(e).__name__
+            try:
+                print(f"{tp} ({e.sqlite_errorname}): {e}", file=sys.stderr)
+            except AttributeError:
+                print(f"{tp}: {e}", file=sys.stderr)
+            if not suppress_errors:
+                sys.exit(1)
 
     def runsource(self, source, filename="<input>", symbol="single"):
-        pass
+        """Override runsource, the core of the InteractiveConsole REPL.
+
+        Return True if more input is needed; buffering is done automatically.
+        Return False is input is a complete statement ready for execution.
+        """
+        match source:
+            case ".version":
+                print(f"{sqlite3.sqlite_version}")
+            case ".help":
+                print("Enter SQL code and press enter.")
+            case ".quit":
+                sys.exit(0)
+            case _:
+                if not sqlite3.complete_statement(source):
+                    return True
+                self.execute(source)
+        return False
 
 
 def make_parser():
@@ -259,8 +304,8 @@ def cli(*args):
     if args.sqlite:
         console = SQLiteConsole(args.filename)
     elif args.duckdb:
-        # console = DuckDBConsole(args.filename)
-        raise NotImplementedError("DuckDB is not yet implemented.")
+        console = DuckDBConsole(args.filename)
+        # raise NotImplementedError("DuckDB is not yet implemented.")
     else:
         console = SQLiteConsole(args.filename)
 

tsellm/core.py

@@ -14,19 +14,19 @@
 """
 
 
-def _prompt_model(prompt, model):
+def _prompt_model(prompt: str, model: str) -> str:
     return llm.get_model(model).prompt(prompt).text()
 
 
-def _prompt_model_default(prompt):
+def _prompt_model_default(prompt: str) -> str:
     return llm.get_model("markov").prompt(prompt).text()
 
 
-def _embed_model(text, model):
+def _embed_model(text: str, model: str) -> str:
     return json.dumps(llm.get_embedding_model(model).embed(text))
 
 
-def _embed_model_default(text):
+def _embed_model_default(text: str) -> str:
     return json.dumps(
         llm.get_embedding_model(llm_cli.get_default_embedding_model()).embed(text)
     )