Merge pull request #215 from kracekumar/add-logs-in-llm-command

[LLM] Add logs while executing llm command

Author: amjith

CHANGELOG.md

@@ -2,6 +2,7 @@
 ## TBD
 
 ### Features
+* Add logs while invoking `\llm`and `\\m+` command. [(#215)](https://github.com/dbcli/litecli/pull/215)
 * Support `--help` in the `\llm`and `\llm+` command. ([#214](https://github.com/dbcli/litecli/pull/214))
 * Make the history file location configurable. ([#206](https://github.com/dbcli/litecli/issues/206))
 

litecli/main.py

@@ -448,9 +448,12 @@ def one_iteration(text=None):
                     try:
                         start = time()
                         cur = self.sqlexecute.conn and self.sqlexecute.conn.cursor()
-                        context, sql = special.handle_llm(text, cur)
+                        context, sql, duration = special.handle_llm(text, cur)
                         if context:
+                            click.echo("LLM Reponse:")
                             click.echo(context)
+                            click.echo('---')
+                        click.echo(f"Time: {duration:.2f} seconds")
                         text = self.prompt_app.prompt(default=sql)
                     except KeyboardInterrupt:
                         return

litecli/packages/special/llm.py

@@ -7,6 +7,7 @@
 import sys
 from runpy import run_module
 from typing import Optional, Tuple
+from time import time
 
 import click
 
@@ -191,9 +192,8 @@ def ensure_litecli_template(replace=False):
     run_external_cmd("llm", PROMPT, "--save", "litecli")
     return
 
-
 @export
-def handle_llm(text, cur) -> Tuple[str, Optional[str]]:
+def handle_llm(text, cur) -> Tuple[str, Optional[str], float]:
     """This function handles the special command `\\llm`.
 
     If it deals with a question that results in a SQL query then it will return
@@ -254,25 +254,32 @@ def handle_llm(text, cur) -> Tuple[str, Optional[str]]:
     if not use_context:
         args = parts
         if capture_output:
+            click.echo("Calling llm command")
+            start = time()
             _, result = run_external_cmd("llm", *args, capture_output=capture_output)
+            end = time()
             match = re.search(_SQL_CODE_FENCE, result, re.DOTALL)
             if match:
                 sql = match.group(1).strip()
             else:
                 output = [(None, None, None, result)]
                 raise FinishIteration(output)
 
-            return result if verbose else "", sql
+            return result if verbose else "", sql, end - start
         else:
             run_external_cmd("llm", *args, restart_cli=restart)
             raise FinishIteration(None)
 
     try:
         ensure_litecli_template()
+        # Measure end to end llm command invocation.
+        # This measures the internal DB command to pull the schema and llm command
+        start = time()
         context, sql = sql_using_llm(cur=cur, question=arg, verbose=verbose)
+        end = time()
         if not verbose:
             context = ""
-        return context, sql
+        return context, sql, end - start
     except Exception as e:
         # Something went wrong. Raise an exception and bail.
         raise RuntimeError(e)
@@ -301,6 +308,7 @@ def sql_using_llm(cur, question=None, verbose=False) -> Tuple[str, Optional[str]
             WHERE type IN ('table','view') AND name NOT LIKE 'sqlite_%'
             ORDER BY 1
     """
+    click.echo("Preparing schema information to feed the llm")
     sample_row_query = "SELECT * FROM {table} LIMIT 1"
     log.debug(schema_query)
     cur.execute(schema_query)
@@ -332,7 +340,9 @@ def sql_using_llm(cur, question=None, verbose=False) -> Tuple[str, Optional[str]
         question,
         " ",  # Dummy argument to prevent llm from waiting on stdin
     ]
+    click.echo("Invoking llm command with schema information")
     _, result = run_external_cmd("llm", *args, capture_output=True)
+    click.echo("Received response from the llm command")
     match = re.search(_SQL_CODE_FENCE, result, re.DOTALL)
     if match:
         sql = match.group(1).strip()

tests/test_llm_special.py

@@ -59,14 +59,15 @@ def test_llm_command_with_c_flag_and_fenced_sql(mock_run_cmd, mock_llm, executor
 
     test_text = r"\llm -c 'Rewrite the SQL without CTE'"
 
-    result, sql = handle_llm(test_text, executor)
+    result, sql, duration = handle_llm(test_text, executor)
 
     # We expect the function to return (result, sql), but result might be "" if verbose is not set
     # By default, `verbose` is false unless text has something like \llm --verbose?
     # The function code: return result if verbose else "", sql
     # Our test_text doesn't set verbose => we expect "" for the returned context.
     assert result == ""
     assert sql == "SELECT * FROM table;"
+    assert isinstance(duration, float)
 
 
 @patch("litecli.packages.special.llm.llm")
@@ -133,7 +134,7 @@ def test_llm_command_with_prompt(mock_sql_using_llm, mock_ensure_template, mock_
     mock_sql_using_llm.return_value = ("context from LLM", "SELECT 1;")
 
     test_text = r"\llm prompt 'Magic happening here?'"
-    context, sql = handle_llm(test_text, executor)
+    context, sql, duration = handle_llm(test_text, executor)
 
     # ensure_litecli_template should be called
     mock_ensure_template.assert_called_once()
@@ -143,6 +144,7 @@ def test_llm_command_with_prompt(mock_sql_using_llm, mock_ensure_template, mock_
     mock_sql_using_llm.assert_called()
     assert context == ""
     assert sql == "SELECT 1;"
+    assert isinstance(duration, float)
 
 
 @patch("litecli.packages.special.llm.llm")
@@ -155,12 +157,13 @@ def test_llm_command_question_with_context(mock_sql_using_llm, mock_ensure_templ
     mock_sql_using_llm.return_value = ("You have context!", "SELECT 2;")
 
     test_text = r"\llm 'Top 10 downloads by size.'"
-    context, sql = handle_llm(test_text, executor)
+    context, sql, duration = handle_llm(test_text, executor)
 
     mock_ensure_template.assert_called_once()
     mock_sql_using_llm.assert_called()
     assert context == ""
     assert sql == "SELECT 2;"
+    assert isinstance(duration, float)
 
 
 @patch("litecli.packages.special.llm.llm")
@@ -173,7 +176,9 @@ def test_llm_command_question_verbose(mock_sql_using_llm, mock_ensure_template,
     mock_sql_using_llm.return_value = ("Verbose context, oh yeah!", "SELECT 42;")
 
     test_text = r"\llm+ 'Top 10 downloads by size.'"
-    context, sql = handle_llm(test_text, executor)
+    context, sql, duration = handle_llm(test_text, executor)
 
     assert context == "Verbose context, oh yeah!"
     assert sql == "SELECT 42;"
+
+    assert isinstance(duration, float)