import { type WriterOptions } from "https://esm.town/v/nbbaier/WriterOptions";
import { sqlite } from "https://esm.town/v/std/sqlite";
import OpenAI from "npm:openai";
interface QueryWriterOptons extends WriterOptions {
table: string;
}
export class QueryWriter {
table: string;
model: string;
apiKey: string;
openai: OpenAI;
constructor(options: QueryWriterOptons) {
const { table, model, ...openaiOptions } = options;
this.table = table;
this.model = model;
this.openai = new OpenAI(openaiOptions);
}
private async getSchema() {
const tableCols = (await sqlite.execute(`PRAGMA table_info(${this.table})`)).rows.map(column => {
return `${column[1]} ${column[2]}`;
}).join(", ");
return `${this.table}(${tableCols})`;
}
private async executeQuery(query: string) {
try {
return await sqlite.execute(query);
} catch (error) {
throw new Error("Error executing query: " + error.message);
}
}
private createPrompt(schema: string, str: string) {
return `
You are an AI assistant that returns raw SQL queries using natural language.
You only output raw SQLite queries. Never return anything other than raw SQLite.
Always begin the query with SELECT. You will be given the following schema:
${schema}
Take the below query and return raw SQLite:
${str}
`;
}
async writeQuery(str: string) {
const schema = await this.getSchema();
const prompt = this.createPrompt(schema, str);
try {
const response = await this.openai.chat.completions.create({
messages: [{ role: "system", content: prompt }],
model: this.model,
});
if (!response.choices || response.choices.length === 0) {
throw new Error("No response from OpenAI");
}
const query = response.choices[0].message?.content;
if (!query) {
throw new Error("No SQL returned from OpenAI. Try again.");
}
return query;
} catch (error) {
throw new Error("Error generating query: " + error.message);
}
}
async gptQuery(str: string) {
const schema = await this.getSchema();
const prompt = this.createPrompt(schema, str);
try {
const response = await this.openai.chat.completions.create({
messages: [{ role: "system", content: prompt }],
model: this.model,
});
if (!response.choices || response.choices.length === 0) {
throw new Error("No response from OpenAI");
}
const query = response.choices[0].message?.content;
if (!query) {
throw new Error("No SQL returned from OpenAI. Try again.");
}
return this.executeQuery(query);
} catch (error) {
throw new Error("Error generating and executing query: " + error.message);
}