feat(ai): refactor adapter and add growing schema

Author: gwardwell

packages/ai/src/mocking/AIAdapter.ts

@@ -1,6 +1,4 @@
 import { ApolloLink } from "@apollo/client";
-import { print } from "graphql";
-import { BASE_SYSTEM_PROMPT } from "./consts.js";
 
 export declare namespace AIAdapter {
   export interface Options {
@@ -11,56 +9,14 @@ export declare namespace AIAdapter {
 }
 
 export abstract class AIAdapter {
-  public providedSystemPrompt: string | undefined;
-  protected systemPrompt: string;
+  public systemPrompt?: string;
 
-  constructor(options: AIAdapter.Options = {}) {
-    this.systemPrompt = BASE_SYSTEM_PROMPT;
-    if (options.systemPrompt) {
-      this.providedSystemPrompt = options.systemPrompt;
-      this.systemPrompt += `\n\n${this.providedSystemPrompt}`;
-    }
+  constructor(options?: AIAdapter.Options) {
+    this.systemPrompt = options?.systemPrompt;
   }
 
-  public generateResponseForOperation(
-    operation: ApolloLink.Operation,
-    prompt: string
-  ): Promise<AIAdapter.Result> {
-    return Promise.resolve({ data: null });
-  }
-
-  protected createPrompt(
-    { query, variables }: ApolloLink.Operation,
-    prompt: string
-  ): string {
-    // Try to get the GraphQL query document string
-    // from the AST location if available, otherwise
-    // use the `print` function to get the query string.
-    //
-    // The AST location may not be available if the query
-    // was parsed with the `noLocation: true` option.
-    //
-    // If the query document string is available through
-    // the AST location, it will save some processing time
-    // over the `print` function.
-    const queryString = query?.loc?.source?.body ?? print(query);
-
-    let promptVariables = "";
-    if (variables) {
-      promptVariables = `
-
-    With variables:
-    \`\`\`json
-    ${JSON.stringify(variables, null, 2)}
-    \`\`\``;
-    }
-
-    return `Give me mock data that fulfills this query:
-    \`\`\`graphql
-    ${queryString}
-    \`\`\`
-    ${promptVariables}
-    ${prompt ? `\nAdditional instructions:\n${prompt}` : ""}
-    `;
-  }
+  public abstract generateObject(
+    prompt: string,
+    systemPrompt: string
+  ): Promise<AIAdapter.Result>;
 }

packages/ai/src/mocking/AIMockLink.ts

@@ -1,5 +1,6 @@
 import { ApolloLink, Observable } from "@apollo/client";
 import { AIAdapter } from "./AIAdapter.js";
+import { BaseAIAdapter } from "./BaseAIAdapter.js";
 
 export declare namespace AIMockLink {
   export type DefaultOptions = {};
@@ -12,32 +13,28 @@ export declare namespace AIMockLink {
 }
 
 export class AIMockLink extends ApolloLink {
-  private adapter: AIAdapter;
+  private adapter: BaseAIAdapter;
   public showWarnings: boolean = true;
 
   public static defaultOptions: AIMockLink.DefaultOptions = {};
 
   constructor(options: AIMockLink.Options) {
     super();
 
-    this.adapter = options.adapter;
+    this.adapter = new BaseAIAdapter(options.adapter);
     this.showWarnings = options.showWarnings ?? true;
   }
 
   public request(
     operation: ApolloLink.Operation
   ): Observable<ApolloLink.Result> {
-    const prompt = operation.getContext().prompt;
-
     return new Observable((observer) => {
       try {
-        this.adapter
-          .generateResponseForOperation(operation, prompt)
-          .then((result) => {
-            // Notify the observer with the generated response
-            observer.next(result);
-            observer.complete();
-          });
+        this.adapter.performQuery(operation).then((result) => {
+          // Notify the observer with the generated response
+          observer.next(result);
+          observer.complete();
+        });
       } catch (error) {
         observer.error(error);
         observer.complete();

packages/ai/src/mocking/BaseAIAdapter.ts

@@ -0,0 +1,93 @@
+import { ApolloLink } from "@apollo/client";
+import { AIAdapter } from "./AIAdapter.js";
+import { print } from "graphql";
+import { BASE_SYSTEM_PROMPT } from "./consts.js";
+import { GrowingSchema } from "./GrowingSchema.js";
+
+export class BaseAIAdapter {
+  private static baseSystemPrompt = BASE_SYSTEM_PROMPT;
+  private schema: GrowingSchema;
+
+  constructor(private implementation: AIAdapter) {
+    this.schema = new GrowingSchema();
+  }
+
+  /**
+   * Performs a query using the implementation adapter.
+   * @param operation - The operation to perform.
+   * @returns The result of the query.
+   */
+  public async performQuery(
+    operation: ApolloLink.Operation
+  ): Promise<AIAdapter.Result> {
+    const systemPrompt = BaseAIAdapter.createSystemPrompt(
+      this.implementation.systemPrompt
+    );
+
+    const prompt = BaseAIAdapter.createPrompt(operation);
+
+    const result = await this.implementation.generateObject(
+      prompt,
+      systemPrompt
+    );
+
+    this.schema.add(operation, result);
+
+    return result;
+  }
+
+  /**
+   * Creates a system prompt from the base system prompt and the provided prompt.
+   * @param prompt - The prompt to add to the base system prompt.
+   * @returns The system prompt.
+   */
+  private static createSystemPrompt(prompt?: string) {
+    return [BaseAIAdapter.baseSystemPrompt, prompt]
+      .filter(Boolean)
+      .join("\n\n");
+  }
+
+  /**
+   * Creates a prompt from the operation.
+   * @param operation - The operation to create a prompt from.
+   * @returns The prompt.
+   */
+  private static createPrompt(operation: ApolloLink.Operation): string {
+    const { query, variables } = operation;
+    const providedPrompt = operation.getContext().prompt;
+
+    // Try to get the GraphQL query document string
+    // from the AST location if available, otherwise
+    // use the `print` function to get the query string.
+    //
+    // The AST location may not be available if the query
+    // was parsed with the `noLocation: true` option.
+    //
+    // If the query document string is available through
+    // the AST location, it will save some processing time
+    // over the `print` function.
+    const queryString = query?.loc?.source?.body ?? print(query);
+
+    const promptParts = [
+      "Give me mock data that fulfills this query:",
+      "```graphql",
+      queryString,
+      "```",
+    ];
+
+    if (variables) {
+      promptParts.push(
+        "\n",
+        "```json",
+        JSON.stringify(variables, null, 2),
+        "```"
+      );
+    }
+
+    if (providedPrompt) {
+      promptParts.push("\n", "Additional instructions:", providedPrompt);
+    }
+
+    return promptParts.join("\n");
+  }
+}

packages/ai/src/mocking/__tests__/AIAdapter.test.ts

@@ -2,7 +2,7 @@ import { AIAdapter } from "../AIAdapter.js";
 
 class DerivedAdapter extends AIAdapter {
   constructor() {
-    super();
+    super({});
   }
 
   public generateObject(prompt: string): Promise<any> {