json_embed(json,model) for DuckDB

Florents-Tselai · Aug 15, 2024 · c5bb51a · c5bb51a
1 parent b45100c
commit c5bb51a
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 19 deletions.
diff --git a/tests/test_tsellm.py b/tests/test_tsellm.py
@@ -174,6 +174,15 @@ def test_interact_valid_multiline_sql(self):
 
 class InMemorySQLiteTest(TsellmConsoleTest):
     path_args = None
+    alice_json = """{
+            \"name\": \"Alice\",
+            \"details\": {
+                \"age\": 30,
+                \"hobbies\": [\"reading\", \"cycling\"],
+                \"location\": \"Wonderland\"
+            },
+            \"greeting\": \"Hello, World!\"
+        }"""
 
     def setUp(self):
         super().setUp()
@@ -225,18 +234,9 @@ def test_embed_hazo_binary(self):
         self.expect_success(*self.path_args, "select embed(randomblob(16), 'hazo')")
 
     def test_embed_json_recursive(self):
-        example_json = """{
-            \"name\": \"Alice\",
-            \"details\": {
-                \"age\": 30,
-                \"hobbies\": [\"reading\", \"cycling\"],
-                \"location\": \"Wonderland\"
-            },
-            \"greeting\": \"Hello, World!\"
-        }"""
         out = self.expect_success(
             *self.path_args,
-            f"select json_extract('{example_json}', '$.name')",
+            f"select json_extract('{self.alice_json}', '$.name')",
         )
         self.assertEqual(
             "('Alice',)\n",
@@ -245,16 +245,18 @@ def test_embed_json_recursive(self):
 
         out = self.expect_success(
             *self.path_args,
-            f"select json_embed('{example_json}', 'hazo')",
+            f"select json_embed('{self.alice_json}', 'hazo')",
         )
         self.assertEqual(
-            ('(\'{"name": [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, 0.0], "details": {"age": 30, "hobbies": [[7.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, 0.0], [7.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, 0.0]], '
-             '"location": [10.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, 0.0]}, "greeting": [6.0, 6.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"),
+            (
+                '(\'{"name": [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, 0.0], "details": {"age": 30, "hobbies": [[7.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, 0.0], [7.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, 0.0]], "
+                '"location": [10.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, 0.0]}, "greeting": [6.0, 6.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,
         )
 
@@ -323,6 +325,33 @@ def test_embed_hazo_binary(self):
         # See https://github.com/Florents-Tselai/tsellm/issues/25
         pass
 
+    def test_embed_json_recursive(self):
+        out = self.expect_success(
+            *self.path_args,
+            f"select '{self.alice_json}'::json -> 'name'",
+        )
+        self.assertEqual(
+            "('\"Alice\"',)\n",
+            out,
+        )
+
+        out = self.expect_success(
+            *self.path_args,
+            f"select json_embed('{self.alice_json}'::json, 'hazo')",
+        )
+        self.assertEqual(
+            (
+                '(\'{"name": [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, 0.0], "details": {"age": 30, "hobbies": [[7.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, 0.0], [7.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, 0.0]], "
+                '"location": [10.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, 0.0]}, "greeting": [6.0, 6.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,
+        )
+
 
 class DiskDuckDBTest(InMemoryDuckDBTest):
     db_fp = None

diff --git a/tsellm/cli.py b/tsellm/cli.py
@@ -276,6 +276,7 @@ def is_valid_db(self) -> bool:
     _functions = [
         ("prompt", 2, _prompt_model, False),
         ("embed", 2, _embed_model, False),
+        ("json_embed", 2, _json_embed_model, False),
     ]
 
     def connect(self):

diff --git a/tsellm/core.py b/tsellm/core.py
@@ -14,7 +14,6 @@
 """
 
 
-
 def json_recurse_apply(json_obj, f):
     if isinstance(json_obj, dict):
         # Recursively apply the function to dictionary values
@@ -29,6 +28,7 @@ def json_recurse_apply(json_obj, f):
         # Return the object as is if it's neither a dictionary, list, or string
         return json_obj
 
+
 def _prompt_model(prompt: str, model: str) -> str:
     return llm.get_model(model).prompt(prompt).text()