chore: better formatting schema

This commit is contained in:
Tianzhou Chen
2023-06-08 00:36:18 +08:00
parent 3f31392dc1
commit 7c6129e9ff
4 changed files with 8 additions and 12 deletions

View File

@ -92,9 +92,7 @@ const getTableSchema = async (connection: Connection, databaseName: string): Pro
`${row["COLUMN_NAME"]} ${row["DATA_TYPE"].toUpperCase()} ${String(row["IS_NULLABLE"]).toUpperCase() === "NO" ? "NOT NULL" : ""}`
);
}
table.structure = `CREATE TABLE [${table.name}] (
${columnList.join(",\n")}
);`;
table.structure = `CREATE TABLE [${table.name}] (${columnList.join(",\n")});`;
}
}
await pool.close();

View File

@ -128,9 +128,7 @@ const getTableSchema = async (connection: Connection, databaseName: string): Pro
`${row["column_name"]} ${row["data_type"].toUpperCase()} ${String(row["is_nullable"]).toUpperCase() === "NO" ? "NOT NULL" : ""}`
);
}
table.structure = `CREATE TABLE \`${table.name}\` (
${columnList.join(",\n")}
);`;
table.structure = `CREATE TABLE \`${table.name}\` (${columnList.join(",\n")});`;
}
}

View File

@ -2,7 +2,6 @@ import { first } from "lodash-es";
import { Assistant, Id } from "@/types";
import * as customAssistantList from "../../assistants";
export const GeneralBotId = "general-bot";
export const SQLChatBotId = "sql-chat-bot";
export const assistantList: Assistant[] = Object.keys(customAssistantList).map((name) => {

View File

@ -13,13 +13,12 @@ export const countTextTokens = (text: string) => {
export function generateDbPromptFromContext(
promptGenerator: (input: string | undefined) => string,
schemaList: any,
schemaList: Schema[],
selectedSchemaName: string,
selectedTablesName: string[],
maxToken: number,
userPrompt?: string
): string {
let schema = "";
// userPrompt is the message that user want to send to bot. When to look prompt in drawer, userPrompt is undefined.
let tokens = countTextTokens(userPrompt || "");
@ -38,13 +37,15 @@ export function generateDbPromptFromContext(
tableList.push(table!.structure);
}
}
let finalTableList = [];
if (tableList) {
for (const table of tableList) {
if (tokens < maxToken / 2) {
tokens += countTextTokens(table);
schema += table;
tokens += countTextTokens(table + "\n\n");
finalTableList.push(table);
}
}
}
return promptGenerator(schema);
return promptGenerator(finalTableList.join("\n\n"));
}