diff --git a/CHANGELOG.md b/CHANGELOG.md index 6c1db58..7256595 100644 --- a/CHANGELOG.md +++ b/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)) diff --git a/litecli/main.py b/litecli/main.py index 8895980..7646bf8 100644 --- a/litecli/main.py +++ b/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 diff --git a/litecli/packages/special/llm.py b/litecli/packages/special/llm.py index 4473af3..67d116c 100644 --- a/litecli/packages/special/llm.py +++ b/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,7 +254,10 @@ 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() @@ -262,17 +265,21 @@ def handle_llm(text, cur) -> Tuple[str, Optional[str]]: 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() diff --git a/tests/test_llm_special.py b/tests/test_llm_special.py index 25bcc9a..14eb82b 100644 --- a/tests/test_llm_special.py +++ b/tests/test_llm_special.py @@ -59,7 +59,7 @@ 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? @@ -67,6 +67,7 @@ def test_llm_command_with_c_flag_and_fenced_sql(mock_run_cmd, mock_llm, executor # 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)