feat: history and tokens usage

.changeset/short-lions-double.md

@@ -0,0 +1,5 @@
+---
+'chat-dbt': minor
+---
+
+improve history modes

.changeset/spicy-countries-occur.md

@@ -0,0 +1,5 @@
+---
+"chat-dbt": minor
+---
+
+add information about tokens used

README.md

@@ -93,9 +93,15 @@ Each attempt is iterative and builds upon the previous ones, so OpenAI is suppos
 chat-dbt --auto-correct 3
 ```
 
-<!-- ### Keep context between queries -->
+### Skip history between queries
 
-<!-- TODO -->
+By default, Chat-DBT keeps the history of the previous exchanges with the new prompt sent to OpenAI.It gives more context to OpenAI and allows queries using previous results. On the other hand, it uses more tokens and is therefore more costly. You can either disable the history with the `--history-mode=none` option, or only keep the previous queries without sending their database result with the `--history-mode=queries` option.
+
+```sh
+chat-dbt --history-mode=[all|none|queries]
+```
+
+Please note that the previous query will however always be sent when you asked to retry a query that failed.
 
 ### Change the format of the result
 

ROADMAP.md

@@ -1,7 +1,7 @@
 # Roadmap
 
 -   Complete the readme file:
-    -   explaination
+    -   explanation
     -   limitations
         -   security
         -   access to DB data
@@ -14,16 +14,15 @@
 -   Implement the options into the web interface
     -   confirm
     -   automatic corrections
-    -   keep-context
+    -   change history mode
     -   Editable SQL query before confirmation (and keep the modifications in the history)
     -   Editable SQL query after an error
     -   Editable prompt query after an error
+-   reset history
 -   Use the npm package as a library
 -   Keep the modifications of the edited SQL query in the history
--   Statistics e.g. time per request, tokens used, etc
 -   Dark/light mode toggle
 -   Add CI tests for both the CLI and the web interface
--   Nice to have (but complicated): when the context is preserved, allow queries from previous data e.g. "translate the previous result into french"
 -   `--hide-sql`
 -   `--hide-result`
 -   Input from stdin/output to stdout (combined with the above)

cli/cli.ts

@@ -1,7 +1,7 @@
 import chalk from 'chalk'
 import inquirer from 'inquirer'
 import { Configuration, OpenAIApi } from 'openai'
-import ora, { Ora } from 'ora'
+import ora from 'ora'
 import table from 'tty-table'
 
 import {
@@ -10,11 +10,11 @@ import {
     getSqlQuery,
     runSqlQuery
 } from '@/shared/chat-gpt'
-import { getErrorPrompt } from '@/shared/error'
+import { ERROR_PROMPT } from '@/shared/error'
 import { getIntrospection } from '@/shared/introspection'
 
-import InputHistoryPrompt from './input-history'
 import { CommonOptions } from './index'
+import InputHistoryPrompt from './input-history'
 import { editFile } from './utils'
 
 type YesNo = 'yes' | 'no'
@@ -26,7 +26,11 @@ export const startCLI = async (options: CommonOptions) => {
     const openai = new OpenAIApi(
         new Configuration({ apiKey: key, organization: org })
     )
-    const history: string[] = []
+    const history: GptSqlResponse[] = []
+    // * The history of queries is not calculated from the full his
+    // * as the history may not be saved between queries when using
+    // * the option --history=queries or --history=none
+    const promptHistory: string[] = []
 
     while (true) {
         const { query } = await inquirer.prompt([
@@ -35,13 +39,14 @@ export const startCLI = async (options: CommonOptions) => {
                 name: 'query',
                 message: 'Describe your query',
                 validate: value => !!value || 'Query cannot be empty',
-                history
+                history: promptHistory
             }
         ])
-        history.push(query)
+        promptHistory.push(query)
         await executeQueryAndShowResult({
             openai,
             query,
+            history,
             ...options
         })
     }
@@ -89,30 +94,32 @@ const printResult = async (result: GptSqlResultItem[], format: string) => {
 
 const executeQueryAndShowResult = async ({
     query,
-    context = [],
+    history = [],
     ...options
 }: CommonOptions & {
     /** @example Number of users who have a first name starting with 'A' */
     query: string
     openai: OpenAIApi
-    context?: GptSqlResponse[]
+    history?: GptSqlResponse[]
 }) => {
-    const { openai, database, format, keepContext, model } = options
     const spinner = ora()
     spinner.start()
     let sqlQuery: string = ''
     try {
         if (!sqlQuery) {
             spinner.text = 'Getting SQL introspection...'
-            const introspection = await getIntrospection(database)
-            spinner.text = 'Calling OpenAI...'
-            sqlQuery = await getSqlQuery({
+            const introspection = await getIntrospection(options.database)
+            spinner.text = 'Calling OpenAI... '
+            const result = await getSqlQuery({
                 query,
-                context,
-                openai,
-                model,
+                history,
+                historyMode: options.historyMode,
+                openai: options.openai,
+                model: options.model,
                 introspection
             })
+            sqlQuery = result.sqlQuery
+            console.log(`${result.usage?.total_tokens} tokens used`)
             if (options.confirm) {
                 // * Ask the user's confirmation before executing the SQL query
                 spinner.stop()
@@ -156,20 +163,22 @@ const executeQueryAndShowResult = async ({
             }
         }
         spinner.text = 'Running query...'
-        const result = await runSqlQuery({ sqlQuery, database })
+        const result = await runSqlQuery({
+            sqlQuery,
+            database: options.database
+        })
         spinner.stop()
         if (!options.confirm) {
             // * Print the SQL query, but only if it's not already printed
             console.log(chalk.dim(sqlQuery))
         }
         spinner.succeed('Success')
-
-        if (keepContext) {
-            context.push({ query, sqlQuery, result })
+        if (options.historyMode === 'none') {
+            history.length = 0
         } else {
-            context.length = 0
+            history.push({ query, sqlQuery, result })
         }
-        printResult(result, format)
+        printResult(result, options.format)
     } catch (e) {
         const error = e as Error
         spinner.stop()
@@ -205,16 +214,19 @@ const executeQueryAndShowResult = async ({
             ])
             retry = prompt.retry
         }
+        if (
+            (retry !== 'no' && sqlQuery !== '') ||
+            options.historyMode !== 'none'
+        ) {
+            history.push({
+                query,
+                sqlQuery,
+                result: undefined,
+                error: error.message
+            })
+        }
         if (retry !== 'no') {
-            if (sqlQuery !== '') {
-                context.push({
-                    query,
-                    sqlQuery,
-                    result: undefined,
-                    error: error.message
-                })
-            }
-            query = getErrorPrompt(error.message)
+            query = ERROR_PROMPT
             if (retry === 'editPrompt') {
                 query = await editFile({ contents: query })
             }
@@ -225,19 +237,22 @@ const executeQueryAndShowResult = async ({
                 })
                 try {
                     spinner.start()
-                    const result = await runSqlQuery({ sqlQuery, database })
+                    const result = await runSqlQuery({
+                        sqlQuery,
+                        database: options.database
+                    })
                     spinner.stop()
                     console.log(chalk.dim(sqlQuery))
                     spinner.succeed('Success')
-                    printResult(result, format)
+                    printResult(result, options.format)
                 } finally {
                     return
                 }
             }
             console.log(chalk.blue('!'), 'Retrying with:', chalk.bold(query))
             await executeQueryAndShowResult({
                 query,
-                context,
+                history,
                 ...options
             })
         }

cli/index.ts

@@ -43,9 +43,11 @@ const program = envProgram
     )
     .addOption(
         new Option(
-            '-k, --keep-context',
-            'keep context between queries'
-        ).default(false)
+            '-t, --history-mode <choice>',
+            'part of the previous queries to keep in the context: all, none, or only the queries without their results'
+        )
+            .choices(['all', 'none', 'queries'])
+            .default('all')
     )
     .addOption(
         new Option('-f, --format <format>', 'format of the result')

components/Dialogs.tsx

@@ -54,7 +54,7 @@ export const ResponseDialog: React.FC<{
     }
     return (
         <LeftDialog>
-            <SqlQuery query={message.sqlQuery} />
+            <SqlQuery query={message.sqlQuery} usage={message.usage} />
             <Error error={message.error} active={last} />
             <Result result={message.result} />
         </LeftDialog>

components/Error.tsx

@@ -2,7 +2,7 @@ import { Button, Container, Group, Title } from '@mantine/core'
 import { IconRefresh } from '@tabler/icons-react'
 import { Fragment, MouseEventHandler } from 'react'
 
-import { getErrorPrompt } from '@/shared/error'
+import { ERROR_PROMPT } from '@/shared/error'
 import { fetcher } from '@/utils/fetch'
 import { useAppContext } from '@/utils/state'
 
@@ -17,18 +17,18 @@ export const Error: React.FC<{
 
     const askCorrection: MouseEventHandler<HTMLButtonElement> = async () => {
         setLoading(true)
-        const query = getErrorPrompt(error)
+        const query = ERROR_PROMPT
         const currentHistory = [...history]
         setHistory([...currentHistory, { query }])
 
         // * Get number of errors after the last successful query + 1 (index starts at 0)
-        const contextSize =
+        const historySize =
             currentHistory.reverse().findIndex(m => !m.error) + 1
 
         const result = await fetcher({
             query,
             // * Get all the queries after the last successful query
-            context: currentHistory.slice(-contextSize)
+            history: currentHistory.slice(-historySize)
         })
         setLoading(false) // TODO this does not re-focus the input
         setHistory([...currentHistory, result])

components/SqlQuery.tsx

@@ -1,12 +1,14 @@
-import { Title } from '@mantine/core'
+import { Badge, Title } from '@mantine/core'
 import { Prism } from '@mantine/prism'
 import { Fragment } from 'react'
 
 import { useStyles } from '@/utils/styles'
+import { CreateCompletionResponseUsage } from 'openai'
 
 export const SqlQuery: React.FC<{
     query?: string
-}> = ({ query }) => {
+    usage?: CreateCompletionResponseUsage
+}> = ({ query, usage }) => {
     const {
         classes: { code }
     } = useStyles()
@@ -17,7 +19,14 @@ export const SqlQuery: React.FC<{
 
     return (
         <Fragment>
-            <Title order={4}>Query</Title>
+            <Title order={4}>
+                Query
+                {usage && (
+                    <sup style={{ paddingLeft: '10px' }}>
+                        <Badge>{usage.total_tokens} tokens</Badge>
+                    </sup>
+                )}
+            </Title>
             <Prism classNames={{ code }} language='sql'>
                 {query}
             </Prism>

next.config.mjs

@@ -16,6 +16,6 @@ export default {
         model: 'gpt-4',
         confirm: false,
         autoCorrect: 0,
-        keepContext: false
+        historyMode: 'all'
     }
 }

pages/_document.tsx

@@ -1,10 +1,8 @@
 import { createGetInitialProps } from '@mantine/next'
 import Document, { Head, Html, Main, NextScript } from 'next/document'
 
-const getInitialProps = createGetInitialProps()
-
 export default class _Document extends Document {
-    static getInitialProps = getInitialProps
+    static getInitialProps = createGetInitialProps()
 
     render() {
         return (

pages/api/gpt-sql-query.ts

@@ -15,9 +15,10 @@ export default async function handler(
     res: NextApiResponse<GptSqlResponse>
 ) {
     // * Get the query
-    const { query, context } = req.body as {
+    const { query, history, historyMode } = req.body as {
         query?: string
-        context?: GptSqlResponse[]
+        history?: GptSqlResponse[]
+        historyMode: string
     }
     if (!query) {
         return res.status(400).json({ error: 'no request', query: '' })
@@ -27,7 +28,8 @@ export default async function handler(
         model,
         query,
         database,
-        context
+        history,
+        historyMode
     })
 
     return res.status(response.error ? 500 : 200).json(response)

pages/index.tsx

@@ -10,10 +10,12 @@ import { IconSend } from '@tabler/icons-react'
 import { Fragment, useState } from 'react'
 
 import { LeftDialog, QueryDialog, ResponseDialog } from '@/components/Dialogs'
-import { useStyles } from '@/utils/styles'
-import { useAppContext } from '@/utils/state'
 import { fetcher } from '@/utils/fetch'
-import React from 'react'
+import { getOptions } from '@/utils/options'
+import { useAppContext } from '@/utils/state'
+import { useStyles } from '@/utils/styles'
+
+const options = getOptions()
 
 const Page: React.FC = () => {
     const { loading, setLoading, history, setHistory } = useAppContext()
@@ -40,8 +42,11 @@ const Page: React.FC = () => {
         setLoading(true)
         const currentHistory = [...history]
         setHistory([...currentHistory, { query }])
-
-        const result = await fetcher({ query })
+        const result = await fetcher({
+            query,
+            historyMode: options.historyMode,
+            history: options.historyMode === 'none' ? undefined : currentHistory
+        })
         setLoading(false)
         setHistory([...currentHistory, result])
         setQuery('')