chore: add data and caching to AI assist (#13750)

Co-authored-by: Ben Church <[email protected]>

Author: Brian Leonard

airbyte-connector-builder-server/src/main/kotlin/io/airbyte/connector_builder/requester/assist/AssistProxy.kt

@@ -22,6 +22,8 @@ class AssistProxy(private val proxyConfig: AssistConfiguration) {
       requestMethod = "POST"
       setRequestProperty("Content-Type", "application/json")
       doOutput = true
+      connectTimeout = 1 * 60 * 1000 // 1 minute
+      readTimeout = 10 * 60 * 1000 // 10 minutes
     }
 
     connection.outputStream.use { outputStream ->

airbyte-webapp/jest.config.ts

@@ -18,6 +18,7 @@ const jestConfig: Config = {
     "\\.(css|png|scss)$": "test-utils/mock-data/mockEmpty.js",
     "\\.svg$": "test-utils/mock-data/mockSvgString.js",
     "\\.svg\\?react$": "test-utils/mock-data/mockSvg.js",
+    "\\?worker$": "test-utils/mock-data/mockWorker.js",
   },
   setupFilesAfterEnv: ["./src/test-utils/setup-tests.ts"],
   globalSetup: "./src/test-utils/global-setup.js",

airbyte-webapp/src/components/connectorBuilder/Builder/AddStreamButton.tsx

@@ -14,9 +14,8 @@ import { Icon } from "components/ui/Icon";
 import { Modal, ModalBody, ModalFooter } from "components/ui/Modal";
 
 import {
-  BuilderAssistCreateStreamParams,
   BuilderAssistFindStreamsResponse,
-  BuilderAssistGlobalParams,
+  BuilderAssistInputStreamParams,
   BuilderAssistManifestResponse,
   useBuilderAssistCreateStream,
   useBuilderAssistFindStreams,
@@ -29,7 +28,7 @@ import { convertToAssistFormValuesSync } from "./Assist/assist";
 import { AssistWaiting } from "./Assist/AssistWaiting";
 import { BuilderField } from "./BuilderField";
 import { BuilderFieldWithInputs } from "./BuilderFieldWithInputs";
-import { AssistData, BuilderStream, DEFAULT_BUILDER_STREAM_VALUES, DEFAULT_SCHEMA, useBuilderWatch } from "../types";
+import { BuilderStream, DEFAULT_BUILDER_STREAM_VALUES, DEFAULT_SCHEMA, useBuilderWatch } from "../types";
 
 interface AddStreamResponse {
   streamName: string;
@@ -172,26 +171,12 @@ const AddStreamModal = ({
   streams: BuilderStream[];
   initialUrlPath?: string;
 }) => {
-  const baseUrl = useBuilderWatch("formValues.global.urlBase");
-  const appName = useBuilderWatch("name") || "Connector";
-  const assistData: AssistData = useBuilderWatch("formValues.assist");
   const { assistEnabled } = useConnectorBuilderFormState();
-
   const shouldAssist = assistEnabled && !isDuplicate;
 
   // TODO refactor to useMutation, as this is a bit of a hack
   const [assistFormValues, setAssistFormValues] = useState<AddStreamFormValues | null>(null);
 
-  const assistParams = useMemo(
-    () => ({
-      docs_url: assistData?.docsUrl,
-      openapi_spec_url: assistData?.openApiSpecUrl,
-      app_name: appName,
-      url_base: baseUrl,
-    }),
-    [assistData?.docsUrl, assistData?.openApiSpecUrl, appName, baseUrl]
-  );
-
   const submitResponse = useCallback(
     (values: AddStreamFormValues) => {
       const otherStreamValues = values.copyOtherStream
@@ -235,10 +220,9 @@ const AddStreamModal = ({
       return null;
     }
     return {
-      ...assistParams,
       stream_name: assistFormValues.streamName,
     };
-  }, [assistParams, assistFormValues]);
+  }, [assistFormValues]);
 
   return (
     <Modal
@@ -255,7 +239,6 @@ const AddStreamModal = ({
           isDuplicate={isDuplicate}
           streams={streams}
           initialUrlPath={initialUrlPath}
-          assistParams={assistParams}
           shouldAssist={shouldAssist}
         />
       )}
@@ -276,19 +259,21 @@ const AddStreamForm = ({
   isDuplicate,
   streams,
   initialUrlPath,
-  assistParams,
   shouldAssist,
 }: {
   onSubmit: (values: AddStreamFormValues) => void;
   onCancel: () => void;
   isDuplicate: boolean;
   streams: BuilderStream[];
   initialUrlPath?: string;
-  assistParams: BuilderAssistGlobalParams;
   shouldAssist: boolean;
 }) => {
   const { formatMessage } = useIntl();
 
+  const { data, isFetching } = useBuilderAssistFindStreams({
+    enabled: shouldAssist,
+  });
+
   const showCopyFromStream = !isDuplicate && streams.length > 0;
   const showUrlPath = !shouldAssist;
 
@@ -306,6 +291,7 @@ const AddStreamForm = ({
     validator.urlPath = yup.string().required("form.empty.error");
   }
 
+  // put the main form default values here so the API can use the context in the new form
   const methods = useForm({
     defaultValues: {
       streamName: "",
@@ -324,7 +310,7 @@ const AddStreamForm = ({
       <form onSubmit={methods.handleSubmit(onSubmit)}>
         <ModalBody className={styles.body}>
           {shouldAssist ? (
-            <AssistedStreamNameField path="streamName" streams={streams} assistParams={assistParams} />
+            <AssistedStreamNameField path="streamName" streams={streams} data={data} isFetching={isFetching} />
           ) : (
             <BuilderField
               path="streamName"
@@ -393,19 +379,16 @@ const AssistLoadingMessage = () => {
 const AssistedStreamNameField = ({
   path,
   streams,
-  assistParams,
+  data,
+  isFetching,
 }: {
   path: string;
   streams: BuilderStream[];
-  assistParams: BuilderAssistGlobalParams;
+  data: BuilderAssistFindStreamsResponse | undefined;
+  isFetching: boolean;
 }) => {
   const { formatMessage } = useIntl();
 
-  const { data, isFetching } = useBuilderAssistFindStreams({
-    ...assistParams,
-    enabled: true,
-  });
-
   const streamOptions = useMemo(() => {
     if (data) {
       return getStreamOptions(streams, data);
@@ -443,7 +426,7 @@ const AssistProcessing = ({
   onComplete,
   onSkip,
 }: {
-  input: BuilderAssistCreateStreamParams;
+  input: BuilderAssistInputStreamParams;
   onComplete: (values: AddStreamResponse) => void;
   onSkip: () => void;
 }) => {

airbyte-webapp/src/components/connectorBuilder/Builder/Assist/AssistButton.tsx

@@ -13,6 +13,7 @@ import { Tooltip } from "components/ui/Tooltip";
 
 import {
   BuilderAssistFindStreamsResponse,
+  BuilderAssistInputAllParams,
   BuilderAssistManifestResponse,
   useBuilderAssistFindAuth,
   useBuilderAssistFindUrlBase,
@@ -144,66 +145,54 @@ export interface AssistButtonProps {
 }
 
 interface AssistButtonConfig {
-  useHook: (params: object) => UseQueryResult<BuilderAssistManifestResponse, HttpError<KnownExceptionInfo>>;
+  useHook: (
+    input: BuilderAssistInputAllParams
+  ) => UseQueryResult<BuilderAssistManifestResponse, HttpError<KnownExceptionInfo>>;
+  useHookParams: Array<keyof BuilderAssistInputAllParams>;
   formPathToSet: string | ((streamNum: number) => string);
   propertiesToPluck?: string[];
 }
 
-const assistButtonConfigs = {
+const assistButtonConfigs: { [key in AssistKey]: AssistButtonConfig } = {
   urlbase: {
     useHook: useBuilderAssistFindUrlBase,
+    useHookParams: [],
     formPathToSet: "global.urlBase",
   },
   auth: {
     useHook: useBuilderAssistFindAuth,
+    useHookParams: [],
     formPathToSet: "global.authenticator",
   },
   metadata: {
     useHook: useBuilderAssistStreamMetadata,
+    useHookParams: ["stream_name"],
     formPathToSet: (streamNum: number) => `streams.${streamNum}`,
     propertiesToPluck: ["urlPath", "httpMethod", "primaryKey"],
   },
   record_selector: {
     useHook: useBuilderAssistStreamResponse,
+    useHookParams: ["stream_name"],
     formPathToSet: (streamNum: number) => `streams.${streamNum}.recordSelector`,
   },
   paginator: {
     useHook: useBuilderAssistFindStreamPaginator,
+    useHookParams: ["stream_name"],
     formPathToSet: (streamNum: number) => `streams.${streamNum}.paginator`,
   },
-} as const;
+};
 
 export const AssistButton: React.FC<AssistButtonProps> = ({ assistKey, streamNum }) => {
-  const assistData: AssistData = useBuilderWatch("formValues.assist");
-  const app_name = useWatch({ name: "name" }) || "Connector";
-  const url_base = useWatch({ name: "formValues.global.urlBase" });
   const stream_name = useWatch({ name: `formValues.streams.${streamNum}.name` });
 
-  const { assistEnabled } = useConnectorBuilderFormState();
-  if (!assistData || !assistEnabled) {
-    return null;
-  }
-
-  // TODO: this is a hack to get the docs_url and openapi_spec_url from the assistData
-  // Before we bring this to our customers we should hoist this error higher.
-  const docs_url = assistData?.docsUrl?.trim();
-  const openapi_spec_url = assistData?.openApiSpecUrl?.trim();
+  const config = assistButtonConfigs[assistKey];
+  const hookParams = useMemo(() => pick({ stream_name }, config.useHookParams), [stream_name, config.useHookParams]);
 
-  const config = assistButtonConfigs[assistKey] as AssistButtonConfig | undefined;
-  if (!config) {
-    console.error(`Unknown assist key: ${assistKey}`);
+  const { assistEnabled } = useConnectorBuilderFormState();
+  if (!assistEnabled) {
     return null;
   }
 
-  const hookParams = {
-    docs_url,
-    openapi_spec_url,
-    app_name,
-    url_base,
-    stream_name,
-    streamNum,
-  };
-
   return <InternalAssistButton assistKey={assistKey} hookParams={hookParams} streamNum={streamNum} {...config} />;
 };
 
@@ -278,8 +267,8 @@ export const AssistAddStreamButton: React.FC<AssistAddStreamProps> = ({
 
 interface InternalAssistButtonProps extends AssistButtonConfig {
   streamNum?: number;
-  hookParams: object;
   assistKey: AssistKey;
+  hookParams: BuilderAssistInputAllParams;
 }
 
 const InternalAssistButton: React.FC<InternalAssistButtonProps> = ({