Merge pull request #198 from dbcli/ai

Ai

Author: amjith

.github/workflows/ci.yml

@@ -12,7 +12,7 @@ jobs:
 
     strategy:
       matrix:
-        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
+        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
 
     steps:
       - uses: actions/checkout@v3

.github/workflows/publish.yml

@@ -13,7 +13,7 @@ jobs:
 
     strategy:
       matrix:
-        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
+        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
 
     steps:
     - uses: actions/checkout@v4

CHANGELOG.md

@@ -1,3 +1,15 @@
+## 1.14.0 - 2025-01-20
+
+### Features
+
+* Add LLM feature to ask an LLM to create a SQL query.
+    - This adds a new `\llm` special command
+    - eg: `\llm "Who is the largest customer based on revenue?"`
+
+### Internal
+
+* Change min required python version to 3.9+
+
 ## 1.13.2 - 2024-11-24
 
 ### Internal

litecli/main.py

@@ -385,6 +385,47 @@ def get_continuation(width, line_number, is_soft_wrap):
         def show_suggestion_tip():
             return iterations < 2
 
+        def output_res(res, start):
+            result_count = 0
+            mutating = False
+            for title, cur, headers, status in res:
+                logger.debug("headers: %r", headers)
+                logger.debug("rows: %r", cur)
+                logger.debug("status: %r", status)
+                threshold = 1000
+                if is_select(status) and cur and cur.rowcount > threshold:
+                    self.echo(
+                        "The result set has more than {} rows.".format(threshold),
+                        fg="red",
+                    )
+                    if not confirm("Do you want to continue?"):
+                        self.echo("Aborted!", err=True, fg="red")
+                        break
+
+                if self.auto_vertical_output:
+                    max_width = self.prompt_app.output.get_size().columns
+                else:
+                    max_width = None
+
+                formatted = self.format_output(title, cur, headers, special.is_expanded_output(), max_width)
+
+                t = time() - start
+                try:
+                    if result_count > 0:
+                        self.echo("")
+                    try:
+                        self.output(formatted, status)
+                    except KeyboardInterrupt:
+                        pass
+                    self.echo("Time: %0.03fs" % t)
+                except KeyboardInterrupt:
+                    pass
+
+                start = time()
+                result_count += 1
+                mutating = mutating or is_mutating(status)
+            return mutating
+
         def one_iteration(text=None):
             if text is None:
                 try:
@@ -402,6 +443,24 @@ def one_iteration(text=None):
                     self.echo(str(e), err=True, fg="red")
                     return
 
+                if special.is_llm_command(text):
+                    try:
+                        start = time()
+                        cur = self.sqlexecute.conn.cursor()
+                        context, sql = special.handle_llm(text, cur)
+                        if context:
+                            click.echo(context)
+                        text = self.prompt_app.prompt(default=sql)
+                    except KeyboardInterrupt:
+                        return
+                    except special.FinishIteration as e:
+                        return output_res(e.results, start) if e.results else None
+                    except RuntimeError as e:
+                        logger.error("sql: %r, error: %r", text, e)
+                        logger.error("traceback: %r", traceback.format_exc())
+                        self.echo(str(e), err=True, fg="red")
+                        return
+
             if not text.strip():
                 return
 
@@ -415,9 +474,6 @@ def one_iteration(text=None):
                     self.echo("Wise choice!")
                     return
 
-            # Keep track of whether or not the query is mutating. In case
-            # of a multi-statement query, the overall query is considered
-            # mutating if any one of the component statements is mutating
             mutating = False
 
             try:
@@ -434,44 +490,11 @@ def one_iteration(text=None):
                 res = sqlexecute.run(text)
                 self.formatter.query = text
                 successful = True
-                result_count = 0
-                for title, cur, headers, status in res:
-                    logger.debug("headers: %r", headers)
-                    logger.debug("rows: %r", cur)
-                    logger.debug("status: %r", status)
-                    threshold = 1000
-                    if is_select(status) and cur and cur.rowcount > threshold:
-                        self.echo(
-                            "The result set has more than {} rows.".format(threshold),
-                            fg="red",
-                        )
-                        if not confirm("Do you want to continue?"):
-                            self.echo("Aborted!", err=True, fg="red")
-                            break
-
-                    if self.auto_vertical_output:
-                        max_width = self.prompt_app.output.get_size().columns
-                    else:
-                        max_width = None
-
-                    formatted = self.format_output(title, cur, headers, special.is_expanded_output(), max_width)
-
-                    t = time() - start
-                    try:
-                        if result_count > 0:
-                            self.echo("")
-                        try:
-                            self.output(formatted, status)
-                        except KeyboardInterrupt:
-                            pass
-                        self.echo("Time: %0.03fs" % t)
-                    except KeyboardInterrupt:
-                        pass
-
-                    start = time()
-                    result_count += 1
-                    mutating = mutating or is_mutating(status)
                 special.unset_once_if_written()
+                # Keep track of whether or not the query is mutating. In case
+                # of a multi-statement query, the overall query is considered
+                # mutating if any one of the component statements is mutating
+                mutating = output_res(res, start)
                 special.unset_pipe_once_if_written()
             except EOFError as e:
                 raise e

litecli/packages/completion_engine.py

@@ -118,6 +118,9 @@ def suggest_special(text):
         else:
             return [{"type": "table", "schema": []}]
 
+    if cmd in [".llm", ".ai", "\\llm", "\\ai"]:
+        return [{"type": "llm"}]
+
     return [{"type": "keyword"}, {"type": "special"}]
 
 

litecli/packages/special/__init__.py

@@ -12,3 +12,4 @@ def export(defn):
 
 from . import dbcommands
 from . import iocommands
+from . import llm

litecli/packages/special/dbcommands.py

@@ -6,6 +6,7 @@
 import platform
 import shlex
 
+
 from litecli import __version__
 from litecli.packages.special import iocommands
 from .main import special_command, RAW_QUERY, PARSED_QUERY

litecli/packages/special/iocommands.py

@@ -1,22 +1,23 @@
 from __future__ import unicode_literals
-import os
-import re
+
 import locale
 import logging
-import subprocess
+import os
+import re
 import shlex
+import subprocess
 from io import open
 from time import sleep
 
 import click
 import sqlparse
 from configobj import ConfigObj
 
+from ..prompt_utils import confirm_destructive_query
 from . import export
-from .main import special_command, NO_QUERY, PARSED_QUERY
 from .favoritequeries import FavoriteQueries
+from .main import NO_QUERY, PARSED_QUERY, special_command
 from .utils import handle_cd_command
-from litecli.packages.prompt_utils import confirm_destructive_query
 
 use_expanded_output = False
 PAGER_ENABLED = True
@@ -27,6 +28,8 @@
 written_to_pipe_once_process = False
 favoritequeries = FavoriteQueries(ConfigObj())
 
+log = logging.getLogger(__name__)
+
 
 @export
 def set_favorite_queries(config):
@@ -95,9 +98,6 @@ def is_expanded_output():
     return use_expanded_output
 
 
-_logger = logging.getLogger(__name__)
-
-
 @export
 def editor_command(command):
     """