-
Notifications
You must be signed in to change notification settings - Fork 45
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Bridge between ChatCraft DB and DuckDB (#783)
* Bridge between ChatCraft DB and DuckDB * feat: add error handling and retry logic for missing chatcraft tables * style: Remove trailing whitespace in duckdb-chatcraft.ts * fix: add proper error typing and type guard in chatCraftQuery * refactor: add stricter validation for chatcraft table sync in queries * feat: sync all referenced chatcraft tables before query retry * refactor: improve error handling flow in chatCraftQuery function * lint * todos * Refactor * Review fixes --------- Co-authored-by: Taras Glek (aider) <[email protected]>
- Loading branch information
Showing
5 changed files
with
208 additions
and
87 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,45 +1,33 @@ | ||
| import { ChatCraftCommand } from "../ChatCraftCommand"; | ||
| import { ChatCraftChat } from "../ChatCraftChat"; | ||
| import { ChatCraftHumanMessage } from "../ChatCraftMessage"; | ||
| import db from "../../lib/db"; | ||
| import { query, queryResultToJson, queryToMarkdown } from "../duckdb"; | ||
| import { getTables, queryToMarkdown } from "../duckdb-chatcraft"; | ||
|
|
||
| export class DuckCommand extends ChatCraftCommand { | ||
| constructor() { | ||
| super("duck", "/duck", "Do some SQL queries"); | ||
| } | ||
|
|
||
| async execute(chat: ChatCraftChat, _user: User | undefined, args?: string[]) { | ||
| const exportResult = await db.exportToDuckDB(); | ||
| let sql: string; | ||
| let markdown: string; | ||
|
|
||
| if (!args?.length) { | ||
| // Get a list of all tables and describe each one | ||
| const message: string[] = ["## DuckDB Tables"]; | ||
|
|
||
| const tables = await query("SHOW TABLES"); | ||
| const tableNames = queryResultToJson(tables).map((row: any) => row.name); | ||
|
|
||
| await Promise.all( | ||
| tableNames.map(async (name: string) => { | ||
| const rowCount = exportResult.tables.find((table) => table.name === name)?.rowCount || 0; | ||
| const tableDescription = await queryToMarkdown(`DESCRIBE ${name}`); | ||
| message.push(`### ${name} (${rowCount} rows)`, tableDescription); | ||
| }) | ||
| ); | ||
|
|
||
| return chat.addMessage(new ChatCraftHumanMessage({ text: message.join("\n\n") })); | ||
| sql = "SHOW TABLES;"; | ||
| markdown = await getTables(); | ||
| } else { | ||
| sql = args.join(" "); | ||
| markdown = await queryToMarkdown(sql); | ||
| } | ||
|
|
||
| const sql = args.join(" "); | ||
| const results = await queryToMarkdown(sql); | ||
| const message = [ | ||
| // show query | ||
| "```sql", | ||
| sql, | ||
| "```", | ||
| // show results | ||
| results, | ||
| ].join("\n\n"); | ||
| markdown, | ||
| ].join("\n"); | ||
| return chat.addMessage(new ChatCraftHumanMessage({ text: message })); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,117 @@ | ||
| import { DataType } from "apache-arrow"; | ||
| import db, { CHATCRAFT_TABLES, ChatCraftTableName, isChatCraftTableName } from "./db"; | ||
| import { | ||
| withConnection, | ||
| insertJSON, | ||
| QueryResult, | ||
| query, | ||
| DuckDBCatalogError, | ||
| queryResultToJson, | ||
| } from "./duckdb"; | ||
| import { jsonToMarkdownTable } from "./utils"; | ||
|
|
||
| /** | ||
| * Extracts chatcraft schema table references from a SQL query | ||
| * TODO: Implement this function using json_serialize_sql and AST traversal | ||
| * @param sql The SQL query to analyze | ||
| * @returns Array of table names referenced in the chatcraft schema | ||
| */ | ||
| function extractChatCraftTables(sql: string): string[] { | ||
| // Match "chatcraft.table_name" or "chatcraft.table_name;" | ||
| const regex = /chatcraft\.(\w+)(?:\s|;|$)/g; | ||
| const matches = [...sql.matchAll(regex)]; | ||
| return [...new Set(matches.map((match) => match[1]))]; | ||
| } | ||
|
|
||
| /** | ||
| * Synchronizes a single ChatCraft table to DuckDB | ||
| * @param tableName The name of the table to synchronize | ||
| * @returns Promise resolving when sync is complete | ||
| */ | ||
| async function syncChatCraftTable(tableName: ChatCraftTableName): Promise<void> { | ||
| const data = await db.byTableName(tableName).toArray(); | ||
|
|
||
| // Convert dates to ISO strings for JSON serialization | ||
| const jsonData = data.map((record) => ({ | ||
| ...record, | ||
| date: record.date instanceof Date ? record.date.toISOString() : record.date, | ||
| })); | ||
|
|
||
| await insertJSON(tableName, JSON.stringify(jsonData), { schema: "chatcraft" }); | ||
| } | ||
|
|
||
| /** | ||
| * Enhanced query function that handles ChatCraft db data synchronization silently | ||
| * @param sql The SQL query to execute | ||
| * @param params Optional parameters for prepared statement | ||
| * @returns Query results as an Arrow Table | ||
| */ | ||
| async function chatCraftQuery<T extends { [key: string]: DataType } = any>( | ||
| sql: string, | ||
| params?: any[] | ||
| ): Promise<QueryResult<T>> { | ||
| try { | ||
| // First, attempt to execute the query assuming everything is already created | ||
| return await query<T>(sql, params); | ||
| } catch (error: unknown) { | ||
| // If the query fails, see if the error is due to a missing table that we can provide | ||
| // by injecting ChatCraft data from Dexie. NOTE: if a user happens to create a tabled | ||
| // that shares the same name as our injected tables (e.g., chatcraft.messages), we'll | ||
| // let them use theirs instead of generating a new one. This also reduces the risk of | ||
| // overhead from premature table creation. | ||
| if (error instanceof DuckDBCatalogError) { | ||
| const referencedTables = extractChatCraftTables(sql); | ||
|
|
||
| // If we have referenced tables, sync them | ||
| if (referencedTables.length > 0) { | ||
| // Create schema if needed | ||
| await withConnection(async (conn) => { | ||
| await conn.query(`CREATE SCHEMA IF NOT EXISTS chatcraft`); | ||
| }); | ||
|
|
||
| // Sync all referenced chatcraft tables | ||
| for (const tableName of referencedTables) { | ||
| if (isChatCraftTableName(tableName)) { | ||
| await syncChatCraftTable(tableName); | ||
| } | ||
| } | ||
|
|
||
| // Retry the query after syncing | ||
| return await query<T>(sql, params); | ||
| } | ||
| } | ||
|
|
||
| // If not a catalog error or no referenced tables, rethrow | ||
| throw error; | ||
| } | ||
| } | ||
|
|
||
| // Replace the original query export with the enhanced, ChatCraft version | ||
| export { chatCraftQuery as query }; | ||
|
|
||
| /** | ||
| * Executes a SQL query and returns the results as a Markdown table | ||
| * @param sql The SQL query to execute | ||
| * @param params Optional parameters for prepared statement | ||
| * @returns Promise resolving to a Markdown formatted table string | ||
| * @throws {Error} If the query fails | ||
| */ | ||
| export async function queryToMarkdown(sql: string, params?: any[]): Promise<string> { | ||
| const result = await chatCraftQuery(sql, params); | ||
| const json = queryResultToJson(result); | ||
| return jsonToMarkdownTable(json); | ||
| } | ||
|
|
||
| /** | ||
| * Get a list of all available tables, including the "virtual" | ||
| * chatcraft.* tables we can sync into duckdb on demand. | ||
| */ | ||
| export async function getTables() { | ||
| const result = await query("show tables"); | ||
| const json = queryResultToJson(result); | ||
| // TODO: this isn't really accurate, since `show tables` only shows what's in the current schema (main) | ||
| CHATCRAFT_TABLES.forEach((table) => { | ||
| json.push({ name: `chatcraft.${table}` }); | ||
| }); | ||
| return jsonToMarkdownTable(json); | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters