opensource/sqlchat
1
0
Fork 0
mirror of https://github.com/sqlchat/sqlchat.git synced 2025-09-23 18:43:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

chore: better formatting schema

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

4
src/lib/connectors/mssql/index.ts
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" : ""}` `${row["COLUMN_NAME"]} ${row["DATA_TYPE"].toUpperCase()} ${String(row["IS_NULLABLE"]).toUpperCase() === "NO" ? "NOT NULL" : ""}`
); );
} }
table.structure = `CREATE TABLE [${table.name}] ( table.structure = `CREATE TABLE [${table.name}] (${columnList.join(",\n")});`;
${columnList.join(",\n")}
);`;
} }
} }
await pool.close(); await pool.close();

4
src/lib/connectors/postgres/index.ts
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" : ""}` `${row["column_name"]} ${row["data_type"].toUpperCase()} ${String(row["is_nullable"]).toUpperCase() === "NO" ? "NOT NULL" : ""}`
); );
} }
table.structure = `CREATE TABLE \`${table.name}\` ( table.structure = `CREATE TABLE \`${table.name}\` (${columnList.join(",\n")});`;
${columnList.join(",\n")}
);`;
} }
} }

1
src/store/assistant.ts
View File

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

11
src/utils/openai.ts
View File

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