Handle default models for embed and prompt

Author: Florents-Tselai

tests/test_tsellm.py

@@ -1,10 +1,12 @@
+import llm.cli
 from sqlite_utils import Database
 from tsellm.cli import cli
 import unittest
 from test.support import captured_stdout, captured_stderr, captured_stdin, os_helper
 from test.support.os_helper import TESTFN, unlink
-
+from llm import models
 import sqlite3
+from llm import cli as llm_cli
 
 
 class CommandLineInterface(unittest.TestCase):
@@ -25,7 +27,7 @@ def expect_success(self, *args):
         # This makes DeprecationWarning and other warnings cause a failure.
         # Let's not be that harsh yet.
         # See https://github.com/Florents-Tselai/llm/tree/fix-utc-warning-312
-        #self.assertEqual(err, "")
+        # self.assertEqual(err, "")
         return out
 
     def expect_failure(self, *args):
@@ -65,17 +67,26 @@ def test_cli_on_disk_db(self):
         self.assertIn("(0,)", out)
 
 
-class PromptFunction(CommandLineInterface):
-    """Testing the SELECT prompt(...) function"""
+class SQLiteLLMFunction(CommandLineInterface):
 
-    def test_prompt_markov(self):
-        out = self.expect_success(":memory:", "select prompt('hello world', 'markov')")
+    def setUp(self):
+        super().setUp()
+        llm_cli.set_default_model("markov")
+        llm_cli.set_default_embedding_model("hazo")
+
+    def assertMarkovResult(self, prompt, generated):
         # Every word should be one of the original prompt (see https://github.com/simonw/llm-markov/blob/657ca504bcf9f0bfc1c6ee5fe838cde9a8976381/tests/test_llm_markov.py#L20)
-        self.assertIn("hello", out)
-        self.assertIn("world", out)
+        for w in prompt.split(" "):
+            self.assertIn(w, generated)
 
+    def test_prompt_markov(self):
+        out = self.expect_success(":memory:", "select prompt('hello world', 'markov')")
+        self.assertMarkovResult("hello world", out)
 
-class EmbedFunction(CommandLineInterface):
+    def test_prompt_default_markov(self):
+        self.assertEquals(llm_cli.get_default_model(), "markov")
+        out = self.expect_success(":memory:", "select prompt('hello world')")
+        self.assertMarkovResult("hello world", out)
 
     def test_embed_hazo(self):
         out = self.expect_success(":memory:", "select embed('hello world', 'hazo')")
@@ -84,6 +95,14 @@ def test_embed_hazo(self):
             out,
         )
 
+    def test_embed_default_hazo(self):
+        self.assertEquals(llm_cli.get_default_embedding_model(), "hazo")
+        out = self.expect_success(":memory:", "select embed('hello world')")
+        self.assertEqual(
+            "('[5.0, 5.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]',)\n",
+            out,
+        )
+
 #
 # def test_cli_prompt(existing_db_path):
 #     db = Database(existing_db_path)

tsellm/core.py

@@ -1,6 +1,8 @@
 import llm
 import json
 
+from llm import cli as llm_cli
+
 TSELLM_CONFIG_SQL = """
 -- tsellm configuration table
 -- need to be taken care of accross migrations and versions.
@@ -15,13 +17,19 @@
 def _prompt_model(prompt, model):
     return llm.get_model(model).prompt(prompt).text()
 
+def _prompt_model_default(prompt):
+    return llm.get_model("markov").prompt(prompt).text()
 
 def _embed_model(text, model):
     return json.dumps(llm.get_embedding_model(model).embed(text))
 
+def _embed_model_default(text):
+    return json.dumps(llm.get_embedding_model(llm_cli.get_default_embedding_model()).embed(text))
 
 def _tsellm_init(con):
     """Entry-point for tsellm initialization."""
     con.execute(TSELLM_CONFIG_SQL)
     con.create_function("prompt", 2, _prompt_model)
+    con.create_function("prompt", 1, _prompt_model_default)
     con.create_function("embed", 2, _embed_model)
+    con.create_function("embed", 1, _embed_model_default)