Update the prompt.

Author: amjith

litecli/packages/special/llm.py

@@ -155,27 +155,32 @@ def __init__(self, results=None):
 """
 
 _SQL_CODE_FENCE = r"```sql\n(.*?)\n```"
-PROMPT = """A SQLite database has the following schema:
+PROMPT = """
+You are a helpful assistant who is a SQLite expert. You are embedded in a SQLite
+cli tool called litecli.
 
-$db_schema
+Answer this question:
+
+$question
 
-Here is a sample row of data from each table: $sample_data
+Use the following context if it is relevant to answering the question. If the
+question is not about the current database then ignore the context.
 
-Use the provided schema and the sample data to construct a SQL query that
-can be run in SQLite3 to answer
+You are connected to a SQLite database with the following schema:
 
-$question
+$db_schema
+
+Here is a sample row of data from each table:
 
-Explain the reason for choosing each table in the SQL query you have
-written. Keep the explanation concise.
-Finally include a sql query in a code fence such as this one:
+$sample_data
+
+If the answer can be found using a SQL query, include a sql query in a code
+fence such as this one:
 
 ```sql
 SELECT count(*) FROM table_name;
 ```
-
-If the question cannot be answered based on the database schema respond with "I
-cannot answer that question" in a sql code fence.
+Keep your explanation concise and focused on the question asked.
 """
 
 

litecli/packages/special/main.py

@@ -52,17 +52,30 @@ def parse_special_command(sql):
     an invocation mode (regular, verbose, or succinct), and the argument.
     """
     raw, _, arg = sql.partition(" ")
-    is_verbose = raw.endswith("+")
-    is_succinct = raw.endswith("-")
-    # strip out any + or - modifiers to get the actual command name
-    command = raw.strip().rstrip("+-")
-    if is_verbose:
-        mode = Verbosity.VERBOSE
-    elif is_succinct:
-        mode = Verbosity.SUCCINCT
+    raw = raw.strip()
+
+    suffix_count = 0
+    idx = len(raw) - 1
+    while idx >= 0 and raw[idx] in "+-":
+        suffix_count += 1
+        idx -= 1
+    if suffix_count > 1:
+        raise ValueError("Invalid special command: %s" % raw)
+
+    if suffix_count == 1:
+        suffix = raw[-1]
+        base_cmd = raw[:-1]
+        mode = Verbosity.VERBOSE if suffix == "+" else Verbosity.SUCCINCT
     else:
+        suffix = None
+        base_cmd = raw
         mode = Verbosity.REGULAR
-    return (command, mode, arg.strip())
+
+    last_char = base_cmd[-1] if base_cmd else None
+    if suffix is None and last_char and not (last_char.isalnum() or last_char in "?."):
+        raise ValueError("Invalid special command: %s" % raw)
+
+    return (base_cmd, mode, arg.strip())
 
 
 @export

tests/test_special_iocommands.py

@@ -4,6 +4,7 @@
 import pytest
 
 import litecli.packages.special
+from litecli.packages.special.main import parse_special_command, Verbosity
 
 
 def test_once_command():
@@ -57,3 +58,32 @@ def test_pipe_once_command():
             litecli.packages.special.unset_pipe_once_if_written()
             f.seek(0)
             assert f.read() == b"hello world\n"
+
+
+@pytest.mark.parametrize(
+    "sql,expected",
+    [
+        (r"\d table_name", ("\\d", Verbosity.REGULAR, "table_name")),
+        (r"\d+ table_name", ("\\d", Verbosity.VERBOSE, "table_name")),
+        (r"\?", ("\\?", Verbosity.REGULAR, "")),
+        (r"\llm Question", ("\\llm", Verbosity.REGULAR, "Question")),
+        (r"\llm-", ("\\llm", Verbosity.SUCCINCT, "")),
+        (r"\llm+", ("\\llm", Verbosity.VERBOSE, "")),
+    ],
+)
+def test_parse_special_command(sql, expected):
+    """
+    Ensure parse_special_command correctly splits the command and mode.
+    """
+    result = parse_special_command(sql)
+    assert result == expected
+
+
+def test_parse_special_command_error():
+    sql = r"\llm* Question"
+    with pytest.raises(ValueError):
+        parse_special_command(sql)
+
+    sql = r"\llm+- Question"
+    with pytest.raises(ValueError):
+        parse_special_command(sql)