budibase/packages/server/src/api/routes/tests/ai.spec.ts

952 lines
27 KiB
TypeScript
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import { z } from "zod"
import { zodResponseFormat } from "openai/helpers/zod"
import { mockChatGPTResponse } from "../../../tests/utilities/mocks/ai/openai"
import TestConfiguration from "../../../tests/utilities/TestConfiguration"
import nock from "nock"
import { configs, env, features, setEnv } from "@budibase/backend-core"
import {
AIInnerConfig,
AIOperationEnum,
AttachmentSubType,
ConfigType,
Feature,
FieldType,
License,
PlanModel,
PlanType,
ProviderConfig,
RelationshipType,
} from "@budibase/types"
import { context } from "@budibase/backend-core"
import { generator } from "@budibase/backend-core/tests"
import { quotas, ai } from "@budibase/pro"
import { MockLLMResponseFn } from "../../../tests/utilities/mocks/ai"
import { mockAnthropicResponse } from "../../../tests/utilities/mocks/ai/anthropic"
function dedent(str: string) {
return str
.split("\n")
.map(line => line.trim())
.join("\n")
}
type SetupFn = (
config: TestConfiguration
) => Promise<() => Promise<void> | void>
interface TestSetup {
name: string
setup: SetupFn
mockLLMResponse: MockLLMResponseFn
}
function budibaseAI(): SetupFn {
return async (config: TestConfiguration) => {
await config.doInTenant(async () => {
await configs.save({
type: ConfigType.AI,
config: {
budibaseAI: {
provider: "BudibaseAI",
name: "Budibase AI",
active: true,
isDefault: true,
},
},
})
})
return setEnv({ OPENAI_API_KEY: "test-key", SELF_HOSTED: false })
}
}
function customAIConfig(providerConfig: Partial<ProviderConfig>): SetupFn {
return async (config: TestConfiguration) => {
const innerConfig: AIInnerConfig = {
myaiconfig: {
provider: "OpenAI",
name: "OpenAI",
apiKey: "test-key",
defaultModel: "gpt-4o-mini",
active: true,
isDefault: true,
...providerConfig,
},
}
const { id, rev } = await config.doInTenant(
async () =>
await configs.save({
type: ConfigType.AI,
config: innerConfig,
})
)
return async () => {
await config.doInTenant(async () => {
const db = context.getGlobalDB()
await db.remove(id, rev)
})
}
}
}
const allProviders: TestSetup[] = [
{
name: "OpenAI API key",
setup: async () => {
return setEnv({
OPENAI_API_KEY: "test-key",
})
},
mockLLMResponse: mockChatGPTResponse,
},
{
name: "OpenAI API key with custom config",
setup: customAIConfig({ provider: "OpenAI", defaultModel: "gpt-4o-mini" }),
mockLLMResponse: mockChatGPTResponse,
},
{
name: "Anthropic API key with custom config",
setup: customAIConfig({
provider: "Anthropic",
defaultModel: "claude-3-5-sonnet-20240620",
}),
mockLLMResponse: mockAnthropicResponse,
},
{
name: "BudibaseAI",
setup: budibaseAI(),
mockLLMResponse: mockChatGPTResponse,
},
]
describe("AI", () => {
const config = new TestConfiguration()
beforeAll(async () => {
await config.init()
})
afterAll(() => {
config.end()
})
beforeEach(() => {
nock.cleanAll()
})
describe.each(allProviders)(
"provider: $name",
({ setup, mockLLMResponse }: TestSetup) => {
let cleanup: () => Promise<void> | void
beforeAll(async () => {
cleanup = await setup(config)
})
afterAll(async () => {
const maybePromise = cleanup()
if (maybePromise) {
await maybePromise
}
})
describe("POST /api/ai/js", () => {
let cleanup: () => void
beforeAll(() => {
cleanup = features.testutils.setFeatureFlags("*", {
AI_JS_GENERATION: true,
})
})
afterAll(() => {
cleanup()
})
it("handles correct plain code response", async () => {
mockLLMResponse(`return 42`)
const { code } = await config.api.ai.generateJs({ prompt: "test" })
expect(code).toBe("return 42")
})
it("handles correct markdown code response", async () => {
mockLLMResponse(
dedent(`
\`\`\`js
return 42
\`\`\`
`)
)
const { code } = await config.api.ai.generateJs({ prompt: "test" })
expect(code).toBe("return 42")
})
it("handles multiple markdown code blocks returned", async () => {
mockLLMResponse(
dedent(`
This:
\`\`\`js
return 42
\`\`\`
Or this:
\`\`\`js
return 10
\`\`\`
`)
)
const { code } = await config.api.ai.generateJs({ prompt: "test" })
expect(code).toBe("return 42")
})
// TODO: handle when this happens
it.skip("handles no code response", async () => {
mockLLMResponse("I'm sorry, you're quite right, etc.")
const { code } = await config.api.ai.generateJs({ prompt: "test" })
expect(code).toBe("")
})
it("handles LLM errors", async () => {
mockLLMResponse(() => {
throw new Error("LLM error")
})
await config.api.ai.generateJs({ prompt: "test" }, { status: 500 })
})
})
describe("POST /api/ai/cron", () => {
it("handles correct cron response", async () => {
mockLLMResponse("0 0 * * *")
const { message } = await config.api.ai.generateCron({
prompt: "test",
})
expect(message).toBe("0 0 * * *")
})
it("handles expected LLM error", async () => {
mockLLMResponse("Error generating cron: skill issue")
await config.api.ai.generateCron(
{
prompt: "test",
},
{ status: 400 }
)
})
it("handles unexpected LLM error", async () => {
mockLLMResponse(() => {
throw new Error("LLM error")
})
await config.api.ai.generateCron(
{
prompt: "test",
},
{ status: 500 }
)
})
})
}
)
})
describe("BudibaseAI", () => {
const config = new TestConfiguration()
let cleanup: () => void | Promise<void>
beforeAll(async () => {
await config.init()
cleanup = await budibaseAI()(config)
})
afterAll(async () => {
if ("then" in cleanup) {
await cleanup()
} else {
cleanup()
}
config.end()
})
describe("POST /api/ai/chat", () => {
let licenseKey = "test-key"
let internalApiKey = "api-key"
let envCleanup: () => void
let featureCleanup: () => void
beforeAll(() => {
envCleanup = setEnv({
SELF_HOSTED: false,
INTERNAL_API_KEY: internalApiKey,
})
featureCleanup = features.testutils.setFeatureFlags("*", {
AI_JS_GENERATION: true,
})
})
afterAll(() => {
featureCleanup()
envCleanup()
})
beforeEach(async () => {
await config.newTenant()
nock.cleanAll()
const license: License = {
plan: {
type: PlanType.FREE,
model: PlanModel.PER_USER,
usesInvoicing: false,
},
features: [Feature.BUDIBASE_AI],
quotas: {} as any,
tenantId: config.tenantId,
}
nock(env.ACCOUNT_PORTAL_URL)
.get(`/api/license/${licenseKey}`)
.reply(200, license)
})
async function getQuotaUsage() {
return await context.doInSelfHostTenantUsingCloud(
config.getTenantId(),
async () => {
return await quotas.getQuotaUsage()
}
)
}
it("handles correct chat response", async () => {
let usage = await getQuotaUsage()
expect(usage._id).toBe(`quota_usage_${config.getTenantId()}`)
expect(usage.monthly.current.budibaseAICredits).toBe(0)
mockChatGPTResponse("Hi there!")
const { message } = await config.api.ai.chat({
messages: [{ role: "user", content: "Hello!" }],
licenseKey: licenseKey,
})
expect(message).toBe("Hi there!")
usage = await getQuotaUsage()
expect(usage.monthly.current.budibaseAICredits).toBeGreaterThan(0)
})
it("handles chat response error", async () => {
mockChatGPTResponse(() => {
throw new Error("LLM error")
})
await config.api.ai.chat(
{
messages: [{ role: "user", content: "Hello!" }],
licenseKey: "test-key",
},
{ status: 500 }
)
})
it("handles no license", async () => {
nock.cleanAll()
nock(env.ACCOUNT_PORTAL_URL).get(`/api/license/${licenseKey}`).reply(404)
await config.api.ai.chat(
{
messages: [{ role: "user", content: "Hello!" }],
licenseKey: "test-key",
},
{
status: 403,
}
)
})
it("handles no license key", async () => {
await config.api.ai.chat(
{
messages: [{ role: "user", content: "Hello!" }],
// @ts-expect-error - intentionally wrong
licenseKey: undefined,
},
{
status: 403,
}
)
})
it("handles text format", async () => {
let usage = await getQuotaUsage()
expect(usage._id).toBe(`quota_usage_${config.getTenantId()}`)
expect(usage.monthly.current.budibaseAICredits).toBe(0)
const gptResponse = generator.word()
mockChatGPTResponse(gptResponse, { format: "text" })
const { message } = await config.api.ai.chat({
messages: [{ role: "user", content: "Hello!" }],
format: "text",
licenseKey: licenseKey,
})
expect(message).toBe(gptResponse)
usage = await getQuotaUsage()
expect(usage.monthly.current.budibaseAICredits).toBeGreaterThan(0)
})
it("handles json format", async () => {
let usage = await getQuotaUsage()
expect(usage._id).toBe(`quota_usage_${config.getTenantId()}`)
expect(usage.monthly.current.budibaseAICredits).toBe(0)
const gptResponse = JSON.stringify({
[generator.word()]: generator.word(),
})
mockChatGPTResponse(gptResponse, { format: "json" })
const { message } = await config.api.ai.chat({
messages: [{ role: "user", content: "Hello!" }],
format: "json",
licenseKey: licenseKey,
})
expect(message).toBe(gptResponse)
usage = await getQuotaUsage()
expect(usage.monthly.current.budibaseAICredits).toBeGreaterThan(0)
})
it("handles structured outputs", async () => {
let usage = await getQuotaUsage()
expect(usage._id).toBe(`quota_usage_${config.getTenantId()}`)
expect(usage.monthly.current.budibaseAICredits).toBe(0)
const gptResponse = generator.guid()
const structuredOutput = zodResponseFormat(
z.object({
[generator.word()]: z.string(),
}),
"key"
)
mockChatGPTResponse(gptResponse, { format: structuredOutput })
const { message } = await config.api.ai.chat({
messages: [{ role: "user", content: "Hello!" }],
format: structuredOutput,
licenseKey: licenseKey,
})
expect(message).toBe(gptResponse)
usage = await getQuotaUsage()
expect(usage.monthly.current.budibaseAICredits).toBeGreaterThan(0)
})
})
describe("POST /api/ai/tables", () => {
let featureCleanup: () => void
beforeAll(() => {
featureCleanup = features.testutils.setFeatureFlags("*", {
AI_TABLE_GENERATION: true,
})
})
afterAll(() => {
featureCleanup()
})
beforeEach(async () => {
await config.newTenant()
nock.cleanAll()
})
const mockAIGenerationStructure = (
generationStructure: ai.GenerationStructure
) =>
mockChatGPTResponse(JSON.stringify(generationStructure), {
format: zodResponseFormat(ai.generationStructure, "key"),
})
const mockAIColumnGeneration = (
generationStructure: ai.GenerationStructure,
aiColumnGeneration: ai.AIColumnSchemas
) =>
mockChatGPTResponse(JSON.stringify(aiColumnGeneration), {
format: zodResponseFormat(
ai.aiColumnSchemas(
ai.aiTableResponseToTableSchema(generationStructure)
),
"key"
),
})
const mockDataGeneration = (
dataGeneration: Record<string, Record<string, any>[]>
) =>
mockChatGPTResponse(JSON.stringify(dataGeneration), {
format: zodResponseFormat(ai.tableDataStructuredOutput([]), "key"),
})
const mockProcessAIColumn = (response: string) =>
mockChatGPTResponse(response)
it("handles correct chat response", async () => {
const prompt = "Create me a table for managing IT tickets"
const generationStructure: ai.GenerationStructure = {
tables: [
{
name: "Tickets",
primaryDisplay: "Title",
schema: [
{
name: "Title",
type: FieldType.STRING,
constraints: {
presence: true,
},
},
{
name: "Description",
type: FieldType.LONGFORM,
constraints: {
presence: true,
},
},
{
name: "Priority",
type: FieldType.OPTIONS,
constraints: {
inclusion: ["Low", "Medium", "High"],
presence: true,
},
},
{
name: "Status",
type: FieldType.OPTIONS,
constraints: {
inclusion: ["Open", "In Progress", "Closed"],
presence: true,
},
},
{
name: "Assignee",
type: FieldType.LINK,
tableId: "Employees",
relationshipType: RelationshipType.MANY_TO_ONE,
reverseFieldName: "AssignedTickets",
relationshipId: "TicketUser",
},
{
name: "Created Date",
type: FieldType.DATETIME,
ignoreTimezones: false,
dateOnly: true,
},
{
name: "Resolution Time (Days)",
type: FieldType.FORMULA,
formula:
'return (new Date() - new Date($("Created Date"))) / (1000 * 60 * 60 * 24);',
responseType: FieldType.NUMBER,
},
{
name: "Attachment",
type: FieldType.ATTACHMENT_SINGLE,
},
],
},
{
name: "Employees",
primaryDisplay: "First Name",
schema: [
{
name: "First Name",
type: FieldType.STRING,
constraints: {
presence: true,
},
},
{
name: "Last Name",
type: FieldType.STRING,
constraints: {
presence: true,
},
},
{
name: "Position",
type: FieldType.STRING,
constraints: {
presence: true,
},
},
{
name: "Photo",
type: FieldType.ATTACHMENT_SINGLE,
subtype: AttachmentSubType.IMAGE,
},
{
name: "Documents",
type: FieldType.ATTACHMENTS,
},
{
name: "AssignedTickets",
type: FieldType.LINK,
tableId: "Tickets",
relationshipType: RelationshipType.ONE_TO_MANY,
reverseFieldName: "Assignee",
relationshipId: "TicketUser",
},
],
},
],
}
mockAIGenerationStructure(generationStructure)
const aiColumnGeneration: ai.AIColumnSchemas = {
Tickets: [
{
name: "Ticket Summary",
type: FieldType.AI,
operation: AIOperationEnum.SUMMARISE_TEXT,
columns: ["Title", "Description"],
},
{
name: "Translated Description",
type: FieldType.AI,
operation: AIOperationEnum.TRANSLATE,
column: "Description",
language: "es",
},
],
Employees: [
{
name: "Role Category",
type: FieldType.AI,
operation: AIOperationEnum.CATEGORISE_TEXT,
columns: ["Position"],
categories: "Manager,Staff,Intern,Contractor",
},
],
}
mockAIColumnGeneration(generationStructure, aiColumnGeneration)
nock("https://photourl.com").get("/any.png").reply(200).persist()
const dataGeneration: Record<string, Record<string, any>[]> = {
Tickets: [
{
Title: "System slow performance",
Description:
"User reports significant slowdowns when using multiple applications simultaneously on their PC.",
Priority: "Medium",
Status: "Closed",
"Created Date": "2025-04-17",
Attachment: {
name: "performance_logs.txt",
extension: ".txt",
content: "performance logs",
},
},
{
Title: "Email delivery failure",
Description:
"Emails sent to external clients are bouncing back. Bounce back message: '550: Recipient address rejected'.",
Priority: "Medium",
Status: "In Progress",
"Created Date": "2025-04-19",
Attachment: {
name: "email_bounce_back.txt",
extension: ".txt",
content: "Email delivery failure",
},
},
{
Title: "Software installation request",
Description:
"Request to install Adobe Photoshop on users workstation for design work.",
Priority: "Low",
Status: "In Progress",
"Created Date": "2025-04-18",
Attachment: {
name: "software_request_form.pdf",
extension: ".pdf",
content: "Software installation request",
},
},
{
Title: "Unable to connect to VPN",
Description:
"User is experiencing issues when trying to connect to the VPN. Error message: 'VPN connection failed due to incorrect credentials'.",
Priority: "High",
Status: "Open",
"Created Date": "2025-04-20",
Attachment: {
name: "vpn_error_screenshot.pdf",
extension: ".pdf",
content: "vpn error",
},
},
],
Employees: [
{
"First Name": "Joshua",
"Last Name": "Lee",
Position: "Application Developer",
Photo: "https://photourl.com/any.png",
Documents: [
{
name: "development_guidelines.pdf",
extension: ".pdf",
content: "any content",
},
{
name: "project_documents.txt",
extension: ".txt",
content: "any content",
},
],
},
{
"First Name": "Emily",
"Last Name": "Davis",
Position: "Software Deployment Technician",
Photo: "https://photourl.com/any.png",
Documents: [
{
name: "software_license_list.txt",
extension: ".txt",
content: "any content",
},
{
name: "deployment_guide.pdf",
extension: ".pdf",
content: "any content",
},
{
name: "installation_logs.txt",
extension: ".txt",
content: "any content",
},
],
},
{
"First Name": "James",
"Last Name": "Smith",
Position: "IT Support Specialist",
Photo: "https://photourl.com/any.png",
Documents: [
{
name: "certificates.pdf",
extension: ".pdf",
content: "any content",
},
{
name: "employment_contract.pdf",
extension: ".pdf",
content: "any content",
},
],
},
{
"First Name": "Jessica",
"Last Name": "Taylor",
Position: "Cybersecurity Analyst",
Photo: "https://photourl.com/any.png",
Documents: [
{
name: "security_audit_report.pdf",
extension: ".pdf",
content: "any content",
},
{
name: "incident_response_plan.pdf",
extension: ".pdf",
content: "any content",
},
],
},
{
"First Name": "Ashley",
"Last Name": "Harris",
Position: "Database Administrator",
Photo: "https://photourl.com/any.png",
Documents: [
{
name: "database_backup.txt",
extension: ".txt",
content: "any content",
},
{
name: "permission_settings.pdf",
extension: ".pdf",
content: "any content",
},
],
},
],
}
mockDataGeneration(dataGeneration)
mockProcessAIColumn("Mock LLM Response")
const { createdTables } = await config.api.ai.generateTables({ prompt })
expect(createdTables).toEqual([
{ id: expect.stringMatching(/ta_\w+/), name: "Tickets" },
{ id: expect.stringMatching(/ta_\w+/), name: "Employees" },
])
const tables = [
await config.api.table.get(createdTables[0].id),
await config.api.table.get(createdTables[1].id),
]
expect(tables).toEqual([
expect.objectContaining({
name: "Tickets",
schema: {
Title: {
name: "Title",
type: "string",
constraints: {
presence: true,
},
aiGenerated: true,
},
Description: {
name: "Description",
type: "longform",
constraints: {
presence: true,
},
aiGenerated: true,
},
Priority: {
name: "Priority",
type: "options",
constraints: {
inclusion: ["Low", "Medium", "High"],
presence: true,
},
aiGenerated: true,
},
Status: {
name: "Status",
type: "options",
constraints: {
inclusion: ["Open", "In Progress", "Closed"],
presence: true,
},
aiGenerated: true,
},
Assignee: {
name: "Assignee",
type: "link",
tableId: createdTables[1].id,
fieldName: "AssignedTickets",
relationshipType: "one-to-many",
aiGenerated: true,
},
"Created Date": {
name: "Created Date",
type: "datetime",
ignoreTimezones: false,
dateOnly: true,
aiGenerated: true,
},
"Resolution Time (Days)": {
name: "Resolution Time (Days)",
type: "formula",
formula:
'{{ js "cmV0dXJuIChuZXcgRGF0ZSgpIC0gbmV3IERhdGUoJCgiQ3JlYXRlZCBEYXRlIikpKSAvICgxMDAwICogNjAgKiA2MCAqIDI0KTs=" }}',
responseType: "number",
aiGenerated: true,
},
Attachment: {
name: "Attachment",
type: "attachment_single",
aiGenerated: true,
},
"Ticket Summary": {
name: "Ticket Summary",
type: "ai",
operation: "SUMMARISE_TEXT",
columns: ["Title", "Description"],
aiGenerated: true,
},
"Translated Description": {
name: "Translated Description",
type: "ai",
operation: "TRANSLATE",
column: "Description",
language: "es",
aiGenerated: true,
},
},
aiGenerated: true,
}),
expect.objectContaining({
name: "Employees 2",
schema: {
"First Name": {
constraints: {
presence: true,
},
name: "First Name",
type: "string",
aiGenerated: true,
},
"Last Name": {
constraints: {
presence: true,
},
name: "Last Name",
type: "string",
aiGenerated: true,
},
Photo: {
name: "Photo",
subtype: "image",
type: "attachment_single",
aiGenerated: true,
},
Position: {
constraints: {
presence: true,
},
name: "Position",
type: "string",
aiGenerated: true,
},
AssignedTickets: {
fieldName: "Assignee",
name: "AssignedTickets",
relationshipType: "many-to-one",
tableId: createdTables[0].id,
type: "link",
aiGenerated: true,
},
Documents: {
name: "Documents",
type: "attachment",
aiGenerated: true,
},
"Role Category": {
categories: "Manager,Staff,Intern,Contractor",
columns: ["Position"],
name: "Role Category",
operation: "CATEGORISE_TEXT",
type: "ai",
aiGenerated: true,
},
},
aiGenerated: true,
}),
])
const tickets = await config.api.row.fetch(createdTables[0].id)
expect(tickets).toHaveLength(4)
const employees = await config.api.row.fetch(createdTables[1].id)
expect(employees).toHaveLength(5)
})
})
})