sea-dev 1.0.0

package/libs/core/src/data-extraction/structuring/custom.ts

@@ -0,0 +1,495 @@
+import type { StructuredData, StructuredResponse } from "../types";
+import { documentExtractionPromptTemplate } from "./prompts";
+import type { FormGroupHydrated, FormWithGroups } from "@sea/db/types";
+import { z } from "zod/v4";
+import { getModel, type ExtractionModel } from "../models";
+import {
+  jsonSchema,
+  generateText,
+  Output,
+  zodSchema,
+  type JSONSchema7,
+  type Schema,
+} from "ai";
+
+// import { writeFileSync } from "fs";
+
+export function addAdditionalPropertiesToJsonSchema(
+  jsonSchema: JSONSchema7,
+): JSONSchema7 {
+  if (jsonSchema.type === "object") {
+    jsonSchema.additionalProperties = false;
+    const properties = jsonSchema.properties;
+    if (properties != null) {
+      for (const property in properties) {
+        properties[property] = addAdditionalPropertiesToJsonSchema(
+          properties[property] as JSONSchema7,
+        );
+      }
+    }
+  }
+  if (jsonSchema.type === "array" && jsonSchema.items != null) {
+    if (Array.isArray(jsonSchema.items)) {
+      jsonSchema.items = jsonSchema.items.map((item) =>
+        addAdditionalPropertiesToJsonSchema(item as JSONSchema7),
+      );
+    } else {
+      jsonSchema.items = addAdditionalPropertiesToJsonSchema(
+        jsonSchema.items as JSONSchema7,
+      );
+    }
+  }
+  if (jsonSchema.anyOf != null) {
+    jsonSchema.anyOf = jsonSchema.anyOf.map((schema) =>
+      addAdditionalPropertiesToJsonSchema(schema as JSONSchema7),
+    );
+  }
+  if (jsonSchema.allOf != null) {
+    jsonSchema.allOf = jsonSchema.allOf.map((schema) =>
+      addAdditionalPropertiesToJsonSchema(schema as JSONSchema7),
+    );
+  }
+  if (jsonSchema.oneOf != null) {
+    jsonSchema.oneOf = jsonSchema.oneOf.map((schema) =>
+      addAdditionalPropertiesToJsonSchema(schema as JSONSchema7),
+    );
+  }
+  return jsonSchema;
+}
+
+function asStrictZodSchema<T>(schema: Schema<T>): Schema<T> {
+  addAdditionalPropertiesToJsonSchema(schema.jsonSchema as JSONSchema7);
+  return jsonSchema(schema.jsonSchema);
+}
+
+const BANK_TYPES: readonly string[] = [
+  "ABN AMRO Bank Commercial",
+  "Airwallex",
+  "Allied Irish Banks",
+  "Allica Bank",
+  "Alpha FX",
+  "American Express",
+  "Arbuthnot Latham",
+  "Bank of Ireland 365 Online",
+  "Bank of Ireland Business On Line",
+  "Bank of Scotland Business",
+  "Bank of Scotland Commercial",
+  "Bank of Scotland Personal",
+  "Barclaycard Commercial Payments",
+  "Barclaycard UK",
+  "Barclays Business",
+  "Barclays Corporate",
+  "Barclays Personal",
+  "Barclays Wealth",
+  "BBVA",
+  "C. Hoare & Co",
+  "Capital One",
+  "Cater Allen",
+  "Caxton",
+  "Chase Bank",
+  "Chelsea Building Society",
+  "Citi Private Bank",
+  "Coutts&Co",
+  "Coutts&Co ClearSpend",
+  "Credit Suisse",
+  "Cynergy Bank",
+  "Danske Bank Business",
+  "Danske Bank Private",
+  "Ebury",
+  "First Direct",
+  "Ghana International Bank",
+  "GoHenry",
+  "Halifax Personal",
+  "Handelsbanken Corporate",
+  "Handelsbanken Private",
+  "Hargreaves Lansdown Saving Limited",
+  "HSBC Business",
+  "HSBC Kinetic",
+  "HSBC Personal",
+  "HSBCnet",
+  "Investec",
+  "Lloyds Bank Business",
+  "Lloyds Bank Commercial",
+  "Lloyds Bank Personal",
+  "Lombard Odier",
+  "Marks & Spencer",
+  "MBNA Personal",
+  "Metro Bank",
+  "Mettle",
+  "Monese",
+  "Monzo Bank Limited",
+  "Nationwide Building Society",
+  "Natwest",
+  "Natwest Bankline",
+  "Natwest ClearSpend",
+  "Natwest International",
+  "Revolut",
+  "Royal Bank of Scotland",
+  "Royal Bank of Scotland Bankline",
+  "Royal Bank of Scotland ClearSpend",
+  "Santander",
+  "smile",
+  "Soldo",
+  "Starling Bank",
+  "Sumitomo Mitsui Banking Corporation",
+  "Tesco Bank",
+  "The Co-Operative Bank Business",
+  "The Co-Operative Bank Retail",
+  "thinkmoney",
+  "Tide",
+  "Triodos Bank",
+  "TSB Bank",
+  "UBS",
+  "Ulster Bank",
+  "Union Bancaire Privée London",
+  "Unity Trust Bank",
+  "Vanquis Bank",
+  "Virgin Money",
+  "Wise",
+  "Yorkshire Building Society",
+  "Zempler Bank",
+] as const;
+
+export type CustomExtractionOptions = {
+  model: ExtractionModel;
+  userPrompt?: string;
+  context?: string;
+  reasoningEffort?: "none" | "minimal" | "low" | "medium" | "high";
+};
+
+type BuildExtractionPromptOptions = {
+  schema: string;
+  userPrompt?: string;
+  context?: string;
+};
+
+export async function structureCustomFormData(
+  serializedLayout: string,
+  form: FormWithGroups,
+  extractionOptions: CustomExtractionOptions,
+): Promise<StructuredResponse> {
+  try {
+    const extractionSchema = form.groups.reduce((acc, group) => {
+      const groupSchema = buildExtractionSchemaFromGroup(group);
+      return z.object({
+        ...acc.shape,
+        ...groupSchema.shape,
+      });
+    }, z.object({}));
+
+    const { model, userPrompt, context, reasoningEffort } = extractionOptions;
+    const schemaJson = JSON.stringify(zodSchema(extractionSchema));
+    const prompt = buildExtractionPrompt({
+      schema: schemaJson,
+      userPrompt,
+      context,
+    });
+
+    const requestParams = {
+      model: getModel(model, extractionOptions),
+      system: prompt,
+      prompt: serializedLayout,
+      output: Output.object({
+        schema: asStrictZodSchema(zodSchema(extractionSchema)),
+      }),
+      maxRetries: 1,
+      temperature: 0,
+      providerOptions: {
+        openai: {
+          ...(reasoningEffort && reasoningEffort !== "none"
+            ? { reasoningEffort }
+            : {}),
+        },
+      },
+    };
+    const result = await generateText(requestParams);
+    const structuredData = extractionSchema.parse(
+      result.output,
+    ) as StructuredData;
+    if (!structuredData) return { data: null };
+
+    return { data: structuredData };
+  } catch (error) {
+    console.error("Error structuring custom form data:", error);
+    throw error;
+  }
+}
+
+export async function structureCustomGroupData(
+  serializedLayout: string,
+  formGroup: FormGroupHydrated,
+  extractionOptions: CustomExtractionOptions,
+): Promise<StructuredResponse> {
+  try {
+    const extractionSchema = buildExtractionSchemaFromGroup(formGroup);
+
+    const { model, userPrompt, context, reasoningEffort } = extractionOptions;
+    const schemaJson = JSON.stringify(zodSchema(extractionSchema));
+    const prompt = buildExtractionPrompt({
+      schema: schemaJson,
+      userPrompt,
+      context,
+    });
+
+    // // Intercept fetch to log OpenAI requests
+    // const originalFetch = global.fetch;
+    // global.fetch = async (url, options) => {
+    //   if (
+    //     url.toString().includes("openai.com") ||
+    //     url.toString().includes("api.openai.com")
+    //   ) {
+    //     const logData = {
+    //       timestamp: new Date().toISOString(),
+    //       url: url.toString(),
+    //       body: options?.body ? JSON.parse(options.body as string) : null,
+    //     };
+    //     writeFileSync(
+    //       ".claude/debug/openai-request.json",
+    //       JSON.stringify(logData, null, 2),
+    //     );
+    //     console.log(
+    //       "OpenAI request saved to .claude/debug/openai-request.json",
+    //     );
+    //   }
+    //   return originalFetch(url, options);
+    // };
+
+    const requestParams = {
+      model: getModel(model, extractionOptions),
+      system: prompt,
+      prompt: serializedLayout,
+      output: Output.object({
+        schema: asStrictZodSchema(zodSchema(extractionSchema)),
+      }),
+      maxRetries: 1,
+      temperature: 0,
+      providerOptions: {
+        openai: {
+          ...(reasoningEffort && reasoningEffort !== "none"
+            ? { reasoningEffort }
+            : {}),
+        },
+      },
+    };
+    const result = await generateText(requestParams);
+
+    // // Restore original fetch
+    // global.fetch = originalFetch;
+
+    // // Save LLM output
+    // writeFileSync(
+    //   ".claude/debug/llm-output.json",
+    //   JSON.stringify(result.output, null, 2),
+    // );
+    // console.log("LLM output saved to .claude/debug/llm-output.json");
+
+    const structuredData = extractionSchema.parse(
+      result.output,
+    ) as StructuredData;
+    if (!structuredData) return { data: null };
+
+    // Parse for bank statement data
+    // TODO: Remove this once we have a proper way to handle this.
+    if (
+      structuredData["Bank Statement"] &&
+      "BankName" in structuredData["Bank Statement"] &&
+      (!structuredData["Bank Statement"].BankName ||
+        structuredData["Bank Statement"].BankName.value === "unknown")
+    ) {
+      if (structuredData["Bank Statement"].BankNameFree) {
+        structuredData["Bank Statement"].BankName = {
+          value: structuredData["Bank Statement"].BankNameFree.value,
+          refId: structuredData["Bank Statement"].BankNameFree.refId,
+        };
+      }
+    }
+
+    return { data: structuredData };
+  } catch (error) {
+    console.error(
+      "Error structuring data for form group:",
+      formGroup.name,
+      error,
+    );
+    throw new Error(
+      `Failed to structure data for form group: ${formGroup.name}: ${error}`,
+    );
+  }
+}
+
+/**
+ * Builds a Zod validation schema from a form group definition.
+ * Converts form groups and fields into a structured schema with proper types and constraints.
+ *
+ * @param group - Form group to convert into a schema
+ * @returns Zod schema object representing the form structure
+ */
+export function buildExtractionSchemaFromGroup(
+  group: FormGroupHydrated,
+): z.ZodObject {
+  const groupFields: Record<string, z.ZodType> = {};
+  for (const field of group.fields) {
+    // TODO: This is a hack to handle the bank name field, as we don't have a
+    // multiple choice field type. We should remove this once we have a way
+    // to handle this.
+    if (
+      group.name === "Bank Statement" &&
+      (field.name === "BankName" || field.name === "Bank Name")
+    ) {
+      groupFields[field.name] = z
+        .object({
+          value: z.enum([...BANK_TYPES, "unknown"]),
+          refId: z.string(),
+        })
+        .nullable()
+        .describe(
+          "Bank name from the available list. If the bank name is not in the list, return 'unknown'.",
+        );
+
+      groupFields["BankNameFree"] = z
+        .object({
+          value: z.string(),
+          refId: z.string(),
+        })
+        .nullable()
+        .describe(
+          "Bank name. If you cannot determine the bank name, return 'Other'.",
+        );
+
+      continue;
+    }
+
+    switch (field.type) {
+      case "string":
+      case "string-long":
+        groupFields[field.name] = z
+          .object({
+            value: z.string(),
+            refId: z.string(),
+          })
+          .nullable()
+          .describe(field.hint);
+        break;
+      case "number":
+        groupFields[field.name] = z
+          .object({
+            value: z.number(),
+            refId: z.string(),
+          })
+          .nullable()
+          .describe(field.hint);
+        break;
+      case "boolean":
+        groupFields[field.name] = z
+          .object({
+            value: z.boolean(),
+            refId: z.string(),
+          })
+          .nullable()
+          .describe(field.hint);
+        break;
+      case "date":
+        groupFields[field.name] = z
+          .object({
+            value: z.iso.date(),
+            refId: z.string(),
+          })
+          .nullable()
+          .describe(field.hint);
+        break;
+      case "email":
+        groupFields[field.name] = z
+          .object({
+            value: z.email(),
+            refId: z.string(),
+          })
+          .nullable()
+          .describe(field.hint);
+        break;
+      case "money":
+        groupFields[field.name] = z
+          .object({
+            value: z.number(),
+            refId: z.string(),
+          })
+          .nullable()
+          .describe(field.hint);
+        break;
+      case "url":
+        groupFields[field.name] = z
+          .object({
+            value: z.string().describe("Valid URL"),
+            refId: z.string(),
+          })
+          .nullable();
+        break;
+      case "enum":
+        if (field.options && field.options.length > 0) {
+          groupFields[field.name] = z
+            .object({
+              value: z.enum(field.options as [string, ...string[]]),
+              refId: z.string(),
+            })
+            .nullable()
+            .describe(
+              `${field.hint}. Must be one of: ${field.options.join(", ")}`,
+            );
+        } else {
+          groupFields[field.name] = z
+            .object({
+              value: z.string(),
+              refId: z.string(),
+            })
+            .nullable()
+            .describe(field.hint);
+        }
+        break;
+      default:
+        groupFields[field.name] = z
+          .object({
+            value: z.string(),
+            refId: z.string(),
+          })
+          .nullable()
+          .describe(field.hint);
+    }
+  }
+
+  const groupSchema = z.object(groupFields).describe(group.hint);
+  return z.object({
+    [group.name]: group.multi ? z.array(groupSchema) : groupSchema,
+  });
+}
+
+/**
+ * Constructs the prompt for document extraction by combining schema, context, and custom instructions.
+ *
+ * @param options - Configuration for building the extraction prompt
+ * @param options.schema - JSON schema defining the expected extraction structure
+ * @param options.context - Optional contextual information to include in the prompt
+ * @param options.customInstructions - Optional custom instructions for the extraction model
+ * @returns Formatted extraction prompt string
+ */
+export function buildExtractionPrompt(
+  options: BuildExtractionPromptOptions,
+): string {
+  const { schema, context, userPrompt } = options;
+  const contextSection = `
+*** Important Context ***
+${context ? `Today is: ${new Date().toISOString().split("T")[0]}\n${context}` : `Today is: ${new Date().toISOString().split("T")[0]}`}`;
+
+  const customInstructionsStub = `
+*** VERY IMPORTANT ***
+
+{userPrompt}
+
+**********************`;
+  const customInstructionsSection = userPrompt
+    ? customInstructionsStub.replace("{userPrompt}", userPrompt)
+    : "";
+
+  return documentExtractionPromptTemplate
+    .replace("{schema}", schema)
+    .replace("{contextSection}", contextSection)
+    .replace("{customInstructionsSection}", customInstructionsSection);
+}