@autobe/agent/src/orchestrate/realize/orchestrateRealizeCollectorPlan.ts

AI backend server code generator

import {
  AutoBeEventSource,
  AutoBeInterfaceHistory,
  AutoBeOpenApi,
  AutoBeProgressEventBase,
  AutoBeRealizeCollectorPlan,
  AutoBeRealizePlanEvent,
} from "@autobe/interface";
import { StringUtil } from "@autobe/utils";
import { IPointer, Singleton } from "tstl";
import typia, { ILlmApplication, ILlmController, IValidation } from "typia";
import { v4 } from "uuid";

import { AutoBeContext } from "../../context/AutoBeContext";
import { buildAnalysisContextSections } from "../../utils/RAGRetrieval";
import { executeCachedBatch } from "../../utils/executeCachedBatch";
import { forceRetry } from "../../utils/forceRetry";
import { getEmbedder } from "../../utils/getEmbedder";
import { AutoBePreliminaryController } from "../common/AutoBePreliminaryController";
import { convertToSectionEntries } from "../common/internal/convertToSectionEntries";
import { IAnalysisSectionEntry } from "../common/structures/IAnalysisSectionEntry";
import { transformRealizeCollectorPlanHistory } from "./histories/transformRealizeCollectorPlanHistory";
import { AutoBeRealizeCollectorProgrammer } from "./programmers/AutoBeRealizeCollectorProgrammer";
import { IAutoBeRealizeCollectorPlanApplication } from "./structures/IAutoBeRealizeCollectorPlanApplication";

export async function orchestrateRealizeCollectorPlan(
  ctx: AutoBeContext,
  props: {
    progress: AutoBeProgressEventBase;
  },
): Promise<AutoBeRealizeCollectorPlan[]> {
  const history: AutoBeInterfaceHistory | null = ctx.state().interface;
  if (history === null)
    throw new Error("Cannot realize collector plan without interface.");

  const document: AutoBeOpenApi.IDocument = history.document;
  const dtoTypeNames: string[] = Object.keys(
    document.components.schemas,
  ).filter(AutoBeRealizeCollectorProgrammer.filter);
  const prismaSchemaNames: Set<string> = new Set(
    ctx
      .state()
      .database!.result.data.files.map((f) => f.models)
      .flat()
      .map((m) => m.name),
  );

  const result: AutoBeRealizeCollectorPlan[][] = await executeCachedBatch(
    ctx,
    Array.from(dtoTypeNames).map((it) => async (promptCacheKey) => {
      const counter = new Singleton(() => ++props.progress.completed);
      try {
        return await forceRetry(() =>
          process(ctx, {
            document,
            dtoTypeName: it,
            prismaSchemaNames,
            promptCacheKey,
            progress: props.progress,
            counter,
          }),
        );
      } catch (error) {
        counter.get();
        throw error;
      }
    }),
  );
  return result.flat();
}

async function process(
  ctx: AutoBeContext,
  props: {
    document: AutoBeOpenApi.IDocument;
    dtoTypeName: string;
    prismaSchemaNames: Set<string>;
    promptCacheKey: string;
    progress: AutoBeProgressEventBase;
    counter: Singleton<number>;
  },
): Promise<AutoBeRealizeCollectorPlan[]> {
  const allSections: IAnalysisSectionEntry[] = convertToSectionEntries(
    ctx.state().analyze?.files ?? [],
  );

  const queryText: string = [
    "collector",
    "plan",
    "dto",
    "prisma",
    props.dtoTypeName,
  ].join(" ");

  const ragSections: IAnalysisSectionEntry[] =
    await buildAnalysisContextSections(
      getEmbedder(),
      allSections,
      queryText,
      "TOPK",
      { log: false, logPrefix: "realizeCollectorPlan" },
    );

  const preliminary: AutoBePreliminaryController<
    | "analysisSections"
    | "databaseSchemas"
    | "interfaceSchemas"
    | "interfaceOperations"
    | "complete"
  > = new AutoBePreliminaryController({
    dispatch: (e) => ctx.dispatch(e),
    state: ctx.state(),
    source: SOURCE,
    application:
      typia.json.application<IAutoBeRealizeCollectorPlanApplication>(),
    kinds: [
      "analysisSections",
      "databaseSchemas",
      "interfaceSchemas",
      "interfaceOperations",
      "complete",
    ],
    local: {
      analysisSections: ragSections,
      interfaceOperations: props.document.operations.filter(
        (op) => op.requestBody?.typeName === props.dtoTypeName,
      ),
      interfaceSchemas: Object.fromEntries(
        Object.entries(props.document.components.schemas).filter(
          ([key]) => key === props.dtoTypeName,
        ),
      ),
    },
  });
  const event: AutoBeRealizePlanEvent = await preliminary.orchestrate(
    ctx,
    async (out) => {
      const pointer: IPointer<IAutoBeRealizeCollectorPlanApplication.IWrite | null> =
        {
          value: null,
        };
      const result: AutoBeContext.IResult = await ctx.conversate({
        source: "realizePlan",
        controller: createController({
          prismaSchemaNames: props.prismaSchemaNames,
          dtoTypeName: props.dtoTypeName,
          build: (next) => {
            pointer.value = next;
          },
          preliminary,
        }),
        enforceFunctionCall: true,
        promptCacheKey: props.promptCacheKey,
        ...transformRealizeCollectorPlanHistory({
          state: ctx.state(),
          preliminary,
          dtoTypeName: props.dtoTypeName,
        }),
      });
      if (pointer.value === null) return out(result)(null);

      const plans: AutoBeRealizeCollectorPlan[] = pointer.value.plans
        .filter((p) => p.databaseSchemaName !== null)
        .map((p) => ({
          type: "collector",
          dtoTypeName: p.dtoTypeName,
          thinking: p.thinking,
          databaseSchemaName: p.databaseSchemaName!,
          references: p.references,
        }));
      const event: AutoBeRealizePlanEvent = {
        type: "realizePlan",
        id: v4(),
        plans,
        acquisition: preliminary.getAcquisition(),
        metric: result.metric,
        tokenUsage: result.tokenUsage,
        completed: props.counter.get(),
        total: props.progress.total,
        step: ctx.state().analyze?.step ?? 0,
        created_at: new Date().toISOString(),
      };
      return out(result)(event);
    },
  );
  ctx.dispatch(event);
  return event.plans as AutoBeRealizeCollectorPlan[];
}

function createController(props: {
  prismaSchemaNames: Set<string>;
  dtoTypeName: string;
  build: (next: IAutoBeRealizeCollectorPlanApplication.IWrite) => void;
  preliminary: AutoBePreliminaryController<
    | "analysisSections"
    | "databaseSchemas"
    | "interfaceSchemas"
    | "interfaceOperations"
    | "complete"
  >;
}): ILlmController {
  const validate: Validator = (input) => {
    const result: IValidation<IAutoBeRealizeCollectorPlanApplication.IProps> =
      typia.validate<IAutoBeRealizeCollectorPlanApplication.IProps>(input);
    if (result.success === false) return result;
    else if (result.data.request.type !== "write")
      return props.preliminary.validate({
        thinking: result.data.thinking,
        request: result.data.request,
      });

    const errors: IValidation.IError[] = [];
    result.data.request.plans.map((plan, i) => {
      if (props.dtoTypeName !== plan.dtoTypeName)
        errors.push({
          path: `$input.request.plans[${i}].dtoTypeName`,
          value: plan.dtoTypeName,
          expected: JSON.stringify(props.dtoTypeName),
          description: StringUtil.trim`
            The DTO type name must be ${JSON.stringify(props.dtoTypeName)}.

            If you have planned other DTO type's collector, 
            please entirely remake the plan with ONLY the DTO type 
            ${JSON.stringify(props.dtoTypeName)}.
          `,
        });
      if (
        plan.databaseSchemaName !== null &&
        props.prismaSchemaNames.has(plan.databaseSchemaName) === false
      )
        errors.push({
          path: `$input.request.plans[${i}].databaseSchemaName`,
          value: plan.databaseSchemaName,
          expected: Array.from(props.prismaSchemaNames)
            .map((s) => JSON.stringify(s))
            .join(" | "),
          description: StringUtil.trim`
            The database schema name must be one of the available database schemas.

            ${Array.from(props.prismaSchemaNames)
              .map((s) => `- ${s}`)
              .join("\n")}
          `,
        });
    });
    return errors.length
      ? {
          success: false,
          errors,
          data: result.data,
        }
      : result;
  };

  const application: ILlmApplication = props.preliminary.fixApplication(
    typia.llm.application<IAutoBeRealizeCollectorPlanApplication>({
      validate: {
        process: validate,
      },
    }),
  );

  return {
    protocol: "class",
    name: SOURCE,
    application,
    execute: {
      process: (next) => {
        if (next.request.type === "write") props.build(next.request);
      },
    } satisfies IAutoBeRealizeCollectorPlanApplication,
  };
}

type Validator = (
  input: unknown,
) => IValidation<IAutoBeRealizeCollectorPlanApplication.IProps>;

const SOURCE = "realizePlan" satisfies AutoBeEventSource;