Typesafe Chat Builder

ChatBuilder.ts

class Prompt<
  TPromptTemplate extends string,
  TSuppliedInputArgs extends ExtractArgs<TPromptTemplate, {}>
> {
  constructor(
    public template: TPromptTemplate,
    public args: TSuppliedInputArgs
  ) {}

  toString() {
    return Object.keys(this.args).reduce((acc, x) => {
      return acc.replace(`{{${x}}}`, this.args[x as keyof typeof this.args]);
    }, this.template) as ReplaceArgs<TPromptTemplate, TSuppliedInputArgs>;
  }
}

class PromptBuilder<
  TPromptTemplate extends string,
  TExpectedInput extends ExtractArgs<TPromptTemplate, {}>
> {
  constructor(protected template: TPromptTemplate) {}

  addInputValidation<
    TSTypeValidator extends ExtractArgs<TPromptTemplate, TSTypeValidator>
  >(): PromptBuilder<TPromptTemplate, TSTypeValidator> {
    return new PromptBuilder(this.template) as any;
  }

  build<const TSuppliedInputArgs extends TExpectedInput>(
    args: TSuppliedInputArgs
  ) {
    return new Prompt<TPromptTemplate, TSuppliedInputArgs>(this.template, args)
      .toString();
  }
}


// Basic Example Usage (not type safe; will take arguments of any value)
const basicPromptBuilder = new PromptBuilder(
  "Tell {{me}} {{num}} {{jokeType}} joke"
);
const basicPrompt = basicPromptBuilder.build({
  //        ^?
  jokeType: "funny",
  me: "Brett",
  num: 1,
});

// Input Validation Example (Will throw type error if you call build without the right arguments)
const validatedPromptBuilder = basicPromptBuilder.addInputValidation<{
  jokeType: "funny" | "silly";
  me: "Brett" | "Liana";
  num: number;
}>();

const validatedPrompt = validatedPromptBuilder.build({
  //        ^?
  jokeType: "funny",
  me: "Brett",
  num: 1,
});

const invalidPrompt = validatedPromptBuilder.build({
  // @ts-expect-error Type '"error"' is not assignable to type '"funny" | "silly"'.
  jokeType: "error",
  // @ts-expect-error Type '"error"' is not assignable to type '"Brett" | "Liana"'.
  me: "error",
  // @ts-expect-error Type 'string' is not assignable to type 'number'.
  num: "error",
});


//-----------------------------------------
// Chat
//-----------------------------------------
type ChatMessageTemplate<TPromptTemplate extends string, TRole = "user" | "assistant" | "system",> = {
    role: TRole,
    content: TPromptTemplate
}

class Chat<
  TMessages extends [ChatMessageTemplate<any>, ...ChatMessageTemplate<any>[]],
  const TArgs extends Record<string, any>
> {
  constructor(
    protected messages: TMessages,
    protected args: TArgs
  ) {}

  toArray() {
    return this.messages
        .map(m => new PromptBuilder(m.content)
            .addInputValidation<ExtractArgs<typeof m.content, TArgs>>()
            .build(this.args)
        ) as ReplaceChatArgs<TMessages, TArgs>;
  }

  toString() {
    return JSON.stringify(this.toArray());
  }
}

class ChatBuilder<
  TMessages extends [ChatMessageTemplate<any>, ...ChatMessageTemplate<any>[]],
  const TArgs extends Record<string, any>
> {
  constructor(protected messages: TMessages) {}

  addInputValidation<
    TSTypeValidator extends ExtractChatArgs<TMessages, TSTypeValidator>
  >(): ChatBuilder<TMessages, TSTypeValidator> {
    return new ChatBuilder(this.messages) as any;
  }

  build<const TSuppliedInputArgs extends TArgs>(
    args: TSuppliedInputArgs
  ) {
    return new Chat<TMessages, TSuppliedInputArgs>(this.messages, args)
      .toArray();
  }
}

// Basic Example Usage (not type safe; will take arguments of any value)
const basicChatBuilder = new ChatBuilder([
//      ^?
    system(`You are a joke generator you only tell {{jokeType}} jokes`),
    user( "Tell me {{num}} Jokes.")
])
const basicChat = basicChatBuilder.build({
//      ^?
    num: 1,
    jokeType: "dumb",
});

// Input Validation Example (Will throw type error if you call build without the right arguments)
const validatedChatBuilder = basicChatBuilder.addInputValidation<{
//      ^?
  jokeType: "funny" | "silly";
  me: "Brett" | "Liana";
  num: number;
}>();

const validatedChat = validatedChatBuilder.build({
  //        ^?
  jokeType: "funny",
  me: "Brett",
  num: 1,
});

const invalidChat = validatedChatBuilder.build({
  // @ts-expect-error Type '"error"' is not assignable to type '"funny" | "silly"'.
  jokeType: "error",
  // @ts-expect-error Type '"error"' is not assignable to type '"Brett" | "Liana"'.
  me: "error",
  // @ts-expect-error Type 'string' is not assignable to type 'number'.
  num: "error",
});

//-----------------------------------------
// Message Creation Helpers
//-----------------------------------------
export function system<T extends string>(
  literals: TemplateStringsArray | T,
  ...placeholders: unknown[]
) {
  return {
    role: "system" as const,
    content: dedent(literals, ...placeholders),
  } as ChatMessageTemplate<T, "system">;
}
export function user<T extends string>(
  literals: TemplateStringsArray | T,
  ...placeholders: unknown[]
) {
  return {
    role: "user" as const,
    content: dedent(literals, ...placeholders),
  } as ChatMessageTemplate<T, "user">;
}
export function assistant<T extends string>(
  literals: TemplateStringsArray | T,
  ...placeholders: unknown[]
) {
  return {
    role: "assistant" as const,
    content: dedent(literals, ...placeholders),
  } as ChatMessageTemplate<T, "assistant">;
}

export function dedent<T extends string>(
  templ: TemplateStringsArray | T,
  ...values: unknown[]
): typeof templ extends TemplateStringsArray ? string : T {
  let strings = Array.from(typeof templ === "string" ? [templ] : templ);

  // 1. Remove trailing whitespace.
  strings[strings.length - 1] = strings[strings.length - 1].replace(
    /\r?\n([\t ]*)$/,
    ""
  );

  // 2. Find all line breaks to determine the highest common indentation level.
  const indentLengths = strings.reduce<number[]>((arr, str) => {
    const matches = str.match(/\n([\t ]+|(?!\s).)/g);
    if (matches) {
      return arr.concat(
        matches.map((match) => match.match(/[\t ]/g)?.length ?? 0)
      );
    }
    return arr;
  }, []);

  // 3. Remove the common indentation from all strings.
  if (indentLengths.length) {
    const pattern = new RegExp(`\n[\t ]{${Math.min(...indentLengths)}}`, "g");

    strings = strings.map((str) => str.replace(pattern, "\n"));
  }

  // 4. Remove leading whitespace.
  strings[0] = strings[0].replace(/^\r?\n/, "");

  // 5. Perform interpolation.
  let string = strings[0];

  values.forEach((value, i) => {
    // 5.1 Read current indentation level
    const endentations = string.match(/(?:^|\n)( *)$/);
    const endentation = endentations ? endentations[1] : "";
    let indentedValue = value;
    // 5.2 Add indentation to values with multiline strings
    if (typeof value === "string" && value.includes("\n")) {
      indentedValue = String(value)
        .split("\n")
        .map((str, i) => {
          return i === 0 ? str : `${endentation}${str}`;
        })
        .join("\n");
    }

    string += indentedValue + strings[i + 1];
  });

  return string as any;
}


//-----------------------------------------
// TS-Helpers
//-----------------------------------------
type ReplaceArgs<
  TPromptTemplate extends string,
  TArgs extends Record<string, any>
> = TPromptTemplate extends `${infer TStart}{{${infer TDataType}}}${infer TRest}`
  ? TRest extends `${string}{{${string}}}` | `${string}{{${string}}}${string}`
    ? `${TStart}${TArgs[TDataType]}${ReplaceArgs<TRest, TArgs>}`
    : `${TStart}${TArgs[TDataType]}${TRest}`
  : TPromptTemplate;

type ExtractArgsAsTuple<TPromptTemplate extends string> =
  TPromptTemplate extends `${string}{{${infer TDataType}}}${infer TRest}`
    ? TRest extends `${string}{{${string}}}` | `${string}{{${string}}}${string}`
      ? [TDataType, ...ExtractArgsAsTuple<TRest>]
      : [TDataType]
    : [];

type ExtractArgs<
  TPromptTemplate extends string,
  TSTypeValidator = ExtractArgs<TPromptTemplate, {}>
> = {
  [K in ExtractArgsAsTuple<TPromptTemplate>[number] as K]: K extends keyof TSTypeValidator
    ? TSTypeValidator[K]
    : any;
};

type ExtractArgsAsUnion<
  TPromptTemplate extends string,
  TSTypeValidator = ExtractArgs<TPromptTemplate, {}>
> = keyof ExtractArgs<TPromptTemplate, TSTypeValidator>;

type ReplaceChatArgs< 
  TMessages,
  TArgs extends Record<string, any>
> = {
  [K in keyof TMessages]: TMessages[K] extends ChatMessageTemplate<string> ? {
    role: TMessages[K]["role"],
    content: ReplaceArgs<TMessages[K]["content"], TArgs>
  } : never
};

type ExtractChatArgs< 
  TMessages,
  TSTypeValidator = ExtractChatArgs<TMessages, {}>
> = ExtractArgs<TMessages extends ChatMessageTemplate<string>[] ? TMessages[number]["content"] : never, TSTypeValidator>


//-----------------------------------------
// TS-test
//-----------------------------------------
type Expect<T extends true> = T;
type Equal<X, Y> = (<T>() => T extends X ? 1 : 2) extends <T>() => T extends Y
  ? 1
  : 2
  ? true
  : false;

type testReplaceArgs = [
  Expect<
    Equal<
      ReplaceArgs<
        "Tell {{person}} a {{jokeType}} joke",
        { jokeType: "funny"; person: "Brett" }
      >,
      "Tell Brett a funny joke"
    >
  >,
  Expect<
    Equal<
      ReplaceArgs<"Tell me a {{jokeType}} joke", { jokeType: "funny" }>,
      "Tell me a funny joke"
    >
  >,
  Expect<
    Equal<
      ReplaceArgs<
        "Tell me a {{jokeType}} {{joke}}",
        { jokeType: "funny"; joke: "poem" }
      >,
      "Tell me a funny poem"
    >
  >
];

type testExtractArgsAsTuple = [
  Expect<
    Equal<
      ExtractArgsAsTuple<"Tell {{person}} a {{jokeType}} joke">,
      ["person", "jokeType"]
    >
  >,
  Expect<
    Equal<ExtractArgsAsTuple<"Tell me a {{jokeType}} joke">, ["jokeType"]>
  >,
  Expect<
    Equal<
      ExtractArgsAsTuple<"Tell me a {{jokeType}} {{joke}}">,
      ["jokeType", "joke"]
    >
  >
];
type fadsfsa = ExtractArgs<
  "Tell {{person}} a {{jokeType}} joke",
  { jokeType: number }
>;
//      ^?

type testExtractArgs = [
  Expect<
    Equal<
      ExtractArgs<
        "Tell {{person}} a {{jokeType}} joke",
        {
          person: string;
          jokeType: string;
        }
      >,
      { jokeType: string; person: string }
    >
  >,
  Expect<
    Equal<
      ExtractArgs<"Tell me a {{jokeType}} joke", { jokeType: "funny" | "dad" }>,
      { jokeType: "funny" | "dad" }
    >
  >,
  Expect<
    Equal<
      ExtractArgs<
        "Tell me a {{jokeType}} {{num}} {{joke}}",
        { jokeType: string; num: number; joke: string }
      >,
      { jokeType: string; num: number; joke: string }
    >
  >
];

type testExtractArgsAsUnion = [
  Expect<
    Equal<
      ExtractArgsAsUnion<"Tell {{person}} a {{jokeType}} joke">,
      "jokeType" | "person"
    >
  >,
  Expect<Equal<ExtractArgsAsUnion<"Tell me a {{jokeType}} joke">, "jokeType">>,
  Expect<
    Equal<
      ExtractArgsAsUnion<"Tell me a {{jokeType}} {{joke}}">,
      "jokeType" | "joke"
    >
  >
];

type testExtractChatArgs = [
    Expect<
        Equal<
            ExtractChatArgs<[
            //    ^?
                { role: "system", content: "foo {{bar}} test"},
                { role: "user", content: "foo {{buzz}} test"}
            ]>,
            {
                bar: any,
                buzz: any
            }
        >
    >,
    Expect<
        Equal<
            ExtractChatArgs<[
            //    ^?
                { role: "system", content: "foo {{bar}} test"},
                { role: "user", content: "foo {{buzz}} test"}
            ], {
                bar: "fizz" | "buzz",
                buzz: number
            }>,
            {
                bar: "fizz" | "buzz",
                buzz: number
            }
        >
    >,
    Expect<
        Equal<
            ExtractChatArgs<[
            //    ^?
                { role: "system", content: "foo {{bar}} test"},
                { role: "user", content: "foo {{buzz}} test"}
            ], {
                bar: "fizz" | "buzz"
            }>,
            {
                bar: "fizz" | "buzz",
                buzz: any
            }
        >
    >,
]

type testReplaceChatArgs = [
  Expect<
    Equal<
      ReplaceChatArgs<
        Array<
            { role: "system", content: "Tell {{person}} a {{jokeType}} joke" }
        >,
        { jokeType: "funny"; person: "Brett" }
      >,
      Array<{ role: "system", content: "Tell Brett a funny joke" }>
    >
  >,
  Expect<
    Equal<
      ReplaceChatArgs<
        Array<
            { role: "system", content: "Tell {{person}} a {{jokeType}} joke" }
        >,
        { jokeType: "funny" | "dad"; person: "Brett" }
      >,
      Array<{ role: "system", content: "Tell Brett a funny joke" | "Tell Brett a dad joke" }>
    >
  >,
  Expect<
    Equal<
      ReplaceChatArgs<[
            { role: "system", content: "foo {{bar}} test"},
            { role: "user", content: "foo {{buzz}} test"}
        ], {
            bar: "test",
            buzz: "test2" | "test3"
        }>,
      [
        { role: "system", content: "foo test test" },
        { role: "user", content: "foo test2 test" | "foo test3 test" }
      ]
    >
  >
];