chore: fix llm

examples/prompts/promptMixin.ts

@@ -37,10 +37,6 @@ async function main() {
     responseSynthesizer,
   });
 
-  console.log({
-    promptsToUse: queryEngine.getPrompts(),
-  });
-
   queryEngine.updatePrompts({
     "responseSynthesizer:summaryTemplate": treeSummarizePrompt,
   });

packages/core/package.json

@@ -30,6 +30,7 @@
     "mammoth": "^1.6.0",
     "md-utils-ts": "^2.0.0",
     "mongodb": "^6.3.0",
+    "mustache": "^4.2.0",
     "notion-md-crawler": "^0.0.2",
     "openai": "^4.26.1",
     "papaparse": "^5.4.1",
@@ -46,6 +47,7 @@
   "devDependencies": {
     "@swc/cli": "^0.3.9",
     "@swc/core": "^1.4.2",
+    "@types/mustache": "^4.2.5",
     "concurrently": "^8.2.2",
     "glob": "^10.3.10",
     "madge": "^6.1.0",

packages/core/src/Prompt.ts

@@ -1,5 +1,6 @@
 import type { SubQuestion } from "./engines/query/types.js";
 import type { ChatMessage } from "./llm/types.js";
+import { PromptTemplate } from "./prompts/types.js";
 import type { ToolMetadata } from "./types.js";
 
 /**
@@ -24,28 +25,30 @@ DEFAULT_TEXT_QA_PROMPT_TMPL = (
 )
 */
 
-export const defaultTextQaPrompt = ({ context = "", query = "" }) => {
-  return `Context information is below.
+export const defaultTextQaPrompt = () => `Context information is below.
 ---------------------
-${context}
+{{context}}
 ---------------------
 Given the context information and not prior knowledge, answer the query.
-Query: ${query}
+Query: {{query}}
 Answer:`;
-};
+
+export const defaultTextQaTemplate = new PromptTemplate<{
+  context: string;
+  query: string;
+}>(defaultTextQaPrompt);
+
+export type TextQaPromptTemplate = typeof defaultTextQaTemplate;
 
 export type TextQaPrompt = typeof defaultTextQaPrompt;
 
-export const anthropicTextQaPrompt: TextQaPrompt = ({
-  context = "",
-  query = "",
-}) => {
+export const anthropicTextQaPrompt = () => {
   return `Context information:
 <context>
-${context}
+{{context}}
 </context>
 Given the context information and not prior knowledge, answer the query.
-Query: ${query}`;
+Query: {{query}}`;
 };
 
 /*
@@ -101,21 +104,25 @@ DEFAULT_REFINE_PROMPT_TMPL = (
 )
 */
 
-export const defaultRefinePrompt = ({
-  query = "",
-  existingAnswer = "",
-  context = "",
-}) => {
-  return `The original query is as follows: ${query}
-We have provided an existing answer: ${existingAnswer}
+export const defaultRefinePrompt = () => {
+  return `The original query is as follows: {{query}}
+We have provided an existing answer: {{existingAnswer}}
 We have the opportunity to refine the existing answer (only if needed) with some more context below.
 ------------
-${context}
+{{context}}
 ------------
 Given the new context, refine the original answer to better answer the query. If the context isn't useful, return the original answer.
 Refined Answer:`;
 };
 
+export const defaultRefinePromptTemplate = new PromptTemplate<{
+  query: string;
+  existingAnswer: string;
+  context: string;
+}>(defaultRefinePrompt);
+
+export type RefinePromptTemplate = typeof defaultRefinePromptTemplate;
+
 export type RefinePrompt = typeof defaultRefinePrompt;
 
 /*

packages/core/src/llm/base.ts

@@ -1,3 +1,4 @@
+import { ChatPromptTemplate, PromptTemplate } from "llamaindex";
 import type {
   ChatMessage,
   ChatResponse,
@@ -15,6 +16,69 @@ import { streamConverter } from "./utils.js";
 export abstract class BaseLLM implements LLM {
   abstract metadata: LLMMetadata;
 
+  predict(
+    params: LLMCompletionParamsStreaming,
+  ): Promise<AsyncIterable<CompletionResponse>>;
+  predict(params: LLMCompletionParamsNonStreaming): Promise<CompletionResponse>;
+  async predict(
+    params: LLMCompletionParamsStreaming | LLMCompletionParamsNonStreaming,
+  ): Promise<CompletionResponse | AsyncIterable<CompletionResponse>> {
+    const { prompt, parentEvent, stream } = params;
+
+    if (prompt instanceof ChatPromptTemplate) {
+      if (stream) {
+        const stream = await this.chat({
+          messages: prompt.formatMessages(),
+          parentEvent,
+          stream: true,
+        });
+        return streamConverter(stream, (chunk) => {
+          return {
+            text: chunk.delta,
+          };
+        });
+      }
+
+      const chatResponse = await this.chat({
+        messages: prompt.formatMessages(),
+        parentEvent,
+      });
+      return { text: chatResponse.message.content as string };
+    }
+
+    if (prompt instanceof PromptTemplate) {
+      if (stream) {
+        const stream = await this.complete({
+          prompt: prompt.format(),
+          parentEvent,
+          stream: true,
+        });
+
+        return stream;
+      }
+
+      return this.complete({
+        prompt: prompt.format(),
+        parentEvent,
+      });
+    }
+
+    if (stream) {
+      const stream = await this.complete({
+        prompt,
+        parentEvent,
+        stream: true,
+      });
+
+      return stream;
+    }
+
+    return this.complete({
+      prompt,
+      parentEvent,
+    });
+  }
+
   complete(
     params: LLMCompletionParamsStreaming,
   ): Promise<AsyncIterable<CompletionResponse>>;

packages/core/src/llm/ollama.ts

@@ -64,6 +64,16 @@ export class Ollama extends BaseEmbedding implements LLM {
     };
   }
 
+  predict(
+    params: LLMCompletionParamsStreaming,
+  ): Promise<AsyncIterable<CompletionResponse>>;
+  predict(params: LLMCompletionParamsNonStreaming): Promise<CompletionResponse>;
+  predict(
+    params: unknown,
+  ): Promise<AsyncIterable<CompletionResponse>> | Promise<CompletionResponse> {
+    throw new Error("TODO: method not implemented for OLLAMA.");
+  }
+
   chat(
     params: LLMChatParamsStreaming,
   ): Promise<AsyncIterable<ChatResponseChunk>>;

packages/core/src/llm/types.ts

@@ -1,11 +1,23 @@
 import type { Tokenizers } from "../GlobalsHelper.js";
 import type { Event } from "../callbacks/CallbackManager.js";
+import type { BasePromptTemplate } from "../prompts/types.js";
 
 /**
  * Unified language model interface
  */
 export interface LLM {
   metadata: LLMMetadata;
+
+  /**
+   * Predict the next completion from the LLM
+   * *
+   * @param params
+   */
+  predict(
+    params: LLMCompletionParamsStreaming,
+  ): Promise<AsyncIterable<CompletionResponse>>;
+  predict(params: LLMCompletionParamsNonStreaming): Promise<CompletionResponse>;
+
   /**
    * Get a chat response from the LLM
    *
@@ -91,7 +103,7 @@ export interface LLMChatParamsNonStreaming extends LLMChatParamsBase {
 }
 
 export interface LLMCompletionParamsBase {
-  prompt: any;
+  prompt: BasePromptTemplate | string;
   parentEvent?: Event;
 }
 

packages/core/src/prompts/index.ts

@@ -1 +1,3 @@
 export * from "./Mixin.js";
+export * from "./types.js";
+export * from "./utils.js";

packages/core/src/prompts/types.ts

@@ -0,0 +1,120 @@
+import type { ChatMessage } from "../llm/types.js";
+
+import mustache from "mustache";
+
+import { mapTemplateVars, messagesToPrompt } from "./utils.js";
+
+export interface BasePromptTemplate<T extends object = {}> {
+  templateVars: Partial<T>;
+
+  partialFormat: (vars: Partial<T>) => PromptTemplate | ChatPromptTemplate;
+
+  format: (extraVars?: T) => string;
+  formatMessages: (extraVars?: T) => ChatMessage[];
+}
+
+export class PromptTemplate<T extends object = {}>
+  implements BasePromptTemplate<T>
+{
+  template: (...args: any) => string;
+
+  templateVars: Partial<T> = {} as T;
+
+  constructor(template: (...args: any) => string, templateVars?: T) {
+    this.template = template;
+    this.templateVars = templateVars ?? ({} as T);
+  }
+
+  mapTemplateVars(): Array<string> {
+    return mapTemplateVars(this.template());
+  }
+
+  partialFormat(vars: Partial<T>): PromptTemplate {
+    const templateVars = {
+      ...vars,
+    };
+
+    this.templateVars = templateVars;
+
+    return new PromptTemplate(this.template, templateVars);
+  }
+
+  format(extraVars?: Partial<T>): string {
+    const prompt = mustache.render(this.template(), {
+      ...this.templateVars,
+      ...extraVars,
+    });
+    return prompt;
+  }
+
+  formatMessages(extraVars?: Partial<T>): ChatMessage[] {
+    const prompt = this.format(extraVars);
+
+    return [
+      {
+        content: prompt,
+        role: "user",
+      },
+    ];
+  }
+}
+
+export class ChatPromptTemplate<T extends object = {}>
+  implements BasePromptTemplate<T>
+{
+  messageTemplates: (...args: any) => ChatMessage[];
+  templateVars: Partial<T> = {} as T;
+
+  constructor(
+    messageTemplates: (...args: any) => ChatMessage[],
+    templateVars?: T,
+  ) {
+    this.messageTemplates = messageTemplates;
+    this.templateVars = templateVars ?? ({} as T);
+  }
+
+  mapTemplateVars(): Array<string> {
+    const allVars: Array<string> = [];
+
+    const messages = this.messageTemplates();
+
+    for (const message of messages) {
+      allVars.push(...mapTemplateVars(message.content));
+    }
+
+    return allVars;
+  }
+
+  partialFormat(vars: Partial<T>): ChatPromptTemplate {
+    const templateVars = {
+      ...vars,
+    };
+
+    this.templateVars = templateVars;
+
+    return new ChatPromptTemplate(this.messageTemplates, templateVars);
+  }
+
+  format(extraVars?: T): string {
+    const messages = this.formatMessages(extraVars);
+    return messagesToPrompt(messages);
+  }
+
+  formatMessages(extraVars?: Partial<T>): ChatMessage[] {
+    const messages: ChatMessage[] = [];
+
+    for (const message of this.messageTemplates()) {
+      const formattedMessage = mustache.render(message.content, {
+        ...this.templateVars,
+        ...extraVars,
+      });
+
+      messages.push({
+        content: formattedMessage,
+        role: message.role,
+      });
+    }
+
+    return messages;
+  }
+}

packages/core/src/prompts/utils.ts

@@ -0,0 +1,44 @@
+import type { ChatMessage } from "../llm/types.js";
+
+export const messagesToPrompt = (messages: ChatMessage[]): string => {
+  const stringMessages = [];
+  for (const message of messages) {
+    const role = message.role;
+    const content = message.content;
+    let stringMessage = `${role}: ${content}`;
+
+    const additionalKwargs = message.additionalKwargs;
+    if (additionalKwargs) {
+      stringMessage += `\n${additionalKwargs}`;
+    }
+    stringMessages.push(stringMessage);
+  }
+
+  stringMessages.push(`assistant: `);
+  return stringMessages.join("\n");
+};
+
+export const mapTemplateVars = (template: string): Array<string> => {
+  const placeholders: Array<string> = [];
+  let index = 0;
+
+  while (index < template.length) {
+    const startIndex = template.indexOf("{{", index);
+    const endIndex = template.indexOf("}}", startIndex) + 2; // +2 to include the '}}'
+
+    if (startIndex === -1 || endIndex === 1) {
+      // No more placeholders
+      break;
+    }
+
+    const placeholderName = template
+      .substring(startIndex + 2, endIndex - 2)
+      .trim(); // Extract name
+
+    placeholders.push(placeholderName);
+
+    index = endIndex; // Move index to end of current placeholder
+  }
+
+  return placeholders;
+};

packages/core/src/selectors/llmSelectors.ts

@@ -9,10 +9,13 @@ import type {
 } from "../types.js";
 import type { SelectorResult } from "./base.js";
 import { BaseSelector } from "./base.js";
-import type { MultiSelectPrompt, SingleSelectPrompt } from "./prompts.js";
+import type {
+  MultiSelectPromptTemplate,
+  SingleSelectPromptTemplate,
+} from "./prompts.js";
 import {
-  defaultMultiSelectPrompt,
-  defaultSingleSelectPrompt,
+  defaultMultiSelectPromptTemplate,
+  defaultSingleSelectPromptTemplate,
 } from "./prompts.js";
 
 function buildChoicesText(choices: ToolMetadataOnlyDescription[]): string {
@@ -46,29 +49,29 @@ type LLMPredictorType = LLM;
  */
 export class LLMMultiSelector extends BaseSelector {
   llm: LLMPredictorType;
-  prompt: MultiSelectPrompt;
+  prompt: MultiSelectPromptTemplate;
   maxOutputs: number;
   outputParser: BaseOutputParser<StructuredOutput<Answer[]>> | null;
 
   constructor(init: {
     llm: LLMPredictorType;
-    prompt?: MultiSelectPrompt;
+    prompt?: MultiSelectPromptTemplate;
     maxOutputs?: number;
     outputParser?: BaseOutputParser<StructuredOutput<Answer[]>>;
   }) {
     super();
     this.llm = init.llm;
-    this.prompt = init.prompt ?? defaultMultiSelectPrompt;
+    this.prompt = init.prompt ?? defaultMultiSelectPromptTemplate;
     this.maxOutputs = init.maxOutputs ?? 10;
 
     this.outputParser = init.outputParser ?? new SelectionOutputParser();
   }
 
-  _getPrompts(): Record<string, MultiSelectPrompt> {
+  _getPrompts(): Record<string, MultiSelectPromptTemplate> {
     return { prompt: this.prompt };
   }
 
-  _updatePrompts(prompts: Record<string, MultiSelectPrompt>): void {
+  _updatePrompts(prompts: Record<string, MultiSelectPromptTemplate>): void {
     if ("prompt" in prompts) {
       this.prompt = prompts.prompt;
     }
@@ -85,15 +88,19 @@ export class LLMMultiSelector extends BaseSelector {
   ): Promise<SelectorResult> {
     const choicesText = buildChoicesText(choices);
 
-    const prompt = this.prompt(
-      choicesText.length,
-      choicesText,
-      query.queryStr,
-      this.maxOutputs,
-    );
+    const prompt = this.prompt.format({
+      numChoices: choicesText.length,
+      contextList: choicesText,
+      queryStr: query.queryStr,
+      maxOutputs: this.maxOutputs,
+    });
 
     const formattedPrompt = this.outputParser?.format(prompt);
 
+    if (!formattedPrompt) {
+      throw new Error("Formatted prompt is undefined");
+    }
+
     const prediction = await this.llm.complete({
       prompt: formattedPrompt,
     });
@@ -117,25 +124,25 @@ export class LLMMultiSelector extends BaseSelector {
  */
 export class LLMSingleSelector extends BaseSelector {
   llm: LLMPredictorType;
-  prompt: SingleSelectPrompt;
+  prompt: SingleSelectPromptTemplate;
   outputParser: BaseOutputParser<StructuredOutput<Answer[]>> | null;
 
   constructor(init: {
     llm: LLMPredictorType;
-    prompt?: SingleSelectPrompt;
+    prompt?: SingleSelectPromptTemplate;
     outputParser?: BaseOutputParser<StructuredOutput<Answer[]>>;
   }) {
     super();
     this.llm = init.llm;
-    this.prompt = init.prompt ?? defaultSingleSelectPrompt;
+    this.prompt = init.prompt ?? defaultSingleSelectPromptTemplate;
     this.outputParser = init.outputParser ?? new SelectionOutputParser();
   }
 
-  _getPrompts(): Record<string, SingleSelectPrompt> {
+  _getPrompts(): Record<string, SingleSelectPromptTemplate> {
     return { prompt: this.prompt };
   }
 
-  _updatePrompts(prompts: Record<string, SingleSelectPrompt>): void {
+  _updatePrompts(prompts: Record<string, SingleSelectPromptTemplate>): void {
     if ("prompt" in prompts) {
       this.prompt = prompts.prompt;
     }
@@ -152,10 +159,18 @@ export class LLMSingleSelector extends BaseSelector {
   ): Promise<SelectorResult> {
     const choicesText = buildChoicesText(choices);
 
-    const prompt = this.prompt(choicesText.length, choicesText, query.queryStr);
+    const prompt = this.prompt.format({
+      numChoices: choicesText.length,
+      contextList: choicesText,
+      queryStr: query.queryStr,
+    });
 
     const formattedPrompt = this.outputParser?.format(prompt);
 
+    if (!formattedPrompt) {
+      throw new Error("Formatted prompt is undefined");
+    }
+
     const prediction = await this.llm.complete({
       prompt: formattedPrompt,
     });

packages/core/src/selectors/prompts.ts

@@ -1,16 +1,23 @@
-export const defaultSingleSelectPrompt = (
-  numChoices: number,
-  contextList: string,
-  queryStr: string,
-): string => {
-  return `Some choices are given below. It is provided in a numbered list (1 to ${numChoices}), where each item in the list corresponds to a summary.
+import { PromptTemplate } from "../index.js";
+
+export const defaultSingleSelectPrompt = (): string => {
+  return `Some choices are given below. It is provided in a numbered list (1 to {{numChoices}), where each item in the list corresponds to a summary.
 ---------------------
-${contextList}
+{{contextList}}
 ---------------------
-Using only the choices above and not prior knowledge, return the choice that is most relevant to the question: '${queryStr}'
+Using only the choices above and not prior knowledge, return the choice that is most relevant to the question: '{{queryStr}}'
 `;
 };
 
+export const defaultSingleSelectPromptTemplate = new PromptTemplate<{
+  numChoices: number;
+  contextList: string;
+  queryStr: string;
+}>(defaultSingleSelectPrompt);
+
+export type SingleSelectPromptTemplate =
+  typeof defaultSingleSelectPromptTemplate;
+
 export type SingleSelectPrompt = typeof defaultSingleSelectPrompt;
 
 export const defaultMultiSelectPrompt = (
@@ -27,4 +34,13 @@ Using only the choices above and not prior knowledge, return the top choices (no
 `;
 };
 
+export const defaultMultiSelectPromptTemplate = new PromptTemplate<{
+  numChoices: number;
+  contextList: string;
+  queryStr: string;
+  maxOutputs: number;
+}>(defaultMultiSelectPrompt);
+
+export type MultiSelectPromptTemplate = typeof defaultMultiSelectPromptTemplate;
+
 export type MultiSelectPrompt = typeof defaultMultiSelectPrompt;

packages/core/src/synthesizers/MultiModalResponseSynthesizer.ts

@@ -6,8 +6,8 @@ import { serviceContextFromDefaults } from "../ServiceContext.js";
 import { imageToDataUrl } from "../embeddings/index.js";
 import type { MessageContentDetail } from "../llm/types.js";
 import { PromptMixin } from "../prompts/Mixin.js";
-import type { TextQaPrompt } from "./../Prompt.js";
-import { defaultTextQaPrompt } from "./../Prompt.js";
+import type { TextQaPromptTemplate } from "./../Prompt.js";
+import { defaultTextQaTemplate } from "./../Prompt.js";
 import type {
   BaseSynthesizer,
   SynthesizeParamsNonStreaming,
@@ -20,7 +20,7 @@ export class MultiModalResponseSynthesizer
 {
   serviceContext: ServiceContext;
   metadataMode: MetadataMode;
-  textQATemplate: TextQaPrompt;
+  textQATemplate: TextQaPromptTemplate;
 
   constructor({
     serviceContext,
@@ -31,17 +31,17 @@ export class MultiModalResponseSynthesizer
 
     this.serviceContext = serviceContext ?? serviceContextFromDefaults();
     this.metadataMode = metadataMode ?? MetadataMode.NONE;
-    this.textQATemplate = textQATemplate ?? defaultTextQaPrompt;
+    this.textQATemplate = textQATemplate ?? defaultTextQaTemplate;
   }
 
-  protected _getPrompts(): { textQATemplate: TextQaPrompt } {
+  protected _getPrompts(): { textQATemplate: TextQaPromptTemplate } {
     return {
       textQATemplate: this.textQATemplate,
     };
   }
 
   protected _updatePrompts(promptsDict: {
-    textQATemplate: TextQaPrompt;
+    textQATemplate: TextQaPromptTemplate;
   }): void {
     if (promptsDict.textQATemplate) {
       this.textQATemplate = promptsDict.textQATemplate;
@@ -70,7 +70,7 @@ export class MultiModalResponseSynthesizer
     );
     // TODO: use builders to generate context
     const context = textChunks.join("\n\n");
-    const textPrompt = this.textQATemplate({ context, query });
+    const textPrompt = this.textQATemplate.format({ context, query });
     const images = await Promise.all(
       imageNodes.map(async (node: ImageNode) => {
         return {
@@ -85,8 +85,10 @@ export class MultiModalResponseSynthesizer
       { type: "text", text: textPrompt },
       ...images,
     ];
+
+    // TODO: Fix this
     const response = await this.serviceContext.llm.complete({
-      prompt,
+      prompt: textPrompt,
       parentEvent,
     });
     return new Response(response.text, nodes);

packages/core/src/synthesizers/builders.ts

@@ -1,21 +1,24 @@
-import type { Event } from "../callbacks/CallbackManager.js";
-import type { LLM } from "../llm/index.js";
-import { streamConverter } from "../llm/utils.js";
+import type { BasePromptTemplate } from "llamaindex";
 import type {
   RefinePrompt,
+  RefinePromptTemplate,
   SimplePrompt,
   TextQaPrompt,
+  TextQaPromptTemplate,
   TreeSummarizePrompt,
 } from "../Prompt.js";
 import {
-  defaultRefinePrompt,
-  defaultTextQaPrompt,
+  defaultRefinePromptTemplate,
+  defaultTextQaTemplate,
   defaultTreeSummarizePrompt,
 } from "../Prompt.js";
 import type { PromptHelper } from "../PromptHelper.js";
 import { getBiggestPrompt } from "../PromptHelper.js";
-import { PromptMixin } from "../prompts/Mixin.js";
 import type { ServiceContext } from "../ServiceContext.js";
+import type { Event } from "../callbacks/CallbackManager.js";
+import type { LLM } from "../llm/index.js";
+import { streamConverter } from "../llm/utils.js";
+import { PromptMixin } from "../prompts/Mixin.js";
 import type {
   ResponseBuilder,
   ResponseBuilderParamsNonStreaming,
@@ -37,11 +40,14 @@ enum ResponseMode {
  */
 export class SimpleResponseBuilder implements ResponseBuilder {
   llm: LLM;
-  textQATemplate: TextQaPrompt;
+  textQATemplate: TextQaPromptTemplate;
 
-  constructor(serviceContext: ServiceContext, textQATemplate?: TextQaPrompt) {
+  constructor(
+    serviceContext: ServiceContext,
+    textQATemplate?: TextQaPromptTemplate,
+  ) {
     this.llm = serviceContext.llm;
-    this.textQATemplate = textQATemplate ?? defaultTextQaPrompt;
+    this.textQATemplate = textQATemplate ?? defaultTextQaTemplate;
   }
 
   getResponse(
@@ -63,7 +69,8 @@ export class SimpleResponseBuilder implements ResponseBuilder {
       context: textChunks.join("\n\n"),
     };
 
-    const prompt = this.textQATemplate(input);
+    const prompt = this.textQATemplate.format(input);
+
     if (stream) {
       const response = await this.llm.complete({ prompt, parentEvent, stream });
       return streamConverter(response, (chunk) => chunk.text);
@@ -80,25 +87,26 @@ export class SimpleResponseBuilder implements ResponseBuilder {
 export class Refine extends PromptMixin implements ResponseBuilder {
   llm: LLM;
   promptHelper: PromptHelper;
-  textQATemplate: TextQaPrompt;
-  refineTemplate: RefinePrompt;
+
+  textQATemplate: TextQaPromptTemplate;
+  refineTemplate: RefinePromptTemplate;
 
   constructor(
     serviceContext: ServiceContext,
-    textQATemplate?: TextQaPrompt,
-    refineTemplate?: RefinePrompt,
+    textQATemplate?: TextQaPromptTemplate,
+    refineTemplate?: RefinePromptTemplate,
   ) {
     super();
 
     this.llm = serviceContext.llm;
     this.promptHelper = serviceContext.promptHelper;
-    this.textQATemplate = textQATemplate ?? defaultTextQaPrompt;
-    this.refineTemplate = refineTemplate ?? defaultRefinePrompt;
+    this.textQATemplate = textQATemplate ?? defaultTextQaTemplate;
+    this.refineTemplate = refineTemplate ?? defaultRefinePromptTemplate;
   }
 
   protected _getPrompts(): {
-    textQATemplate: RefinePrompt;
-    refineTemplate: RefinePrompt;
+    textQATemplate: TextQaPromptTemplate;
+    refineTemplate: RefinePromptTemplate;
   } {
     return {
       textQATemplate: this.textQATemplate,
@@ -107,8 +115,8 @@ export class Refine extends PromptMixin implements ResponseBuilder {
   }
 
   protected _updatePrompts(prompts: {
-    textQATemplate: RefinePrompt;
-    refineTemplate: RefinePrompt;
+    textQATemplate: TextQaPromptTemplate;
+    refineTemplate: RefinePromptTemplate;
   }): void {
     if (prompts.textQATemplate) {
       this.textQATemplate = prompts.textQATemplate;
@@ -166,20 +174,24 @@ export class Refine extends PromptMixin implements ResponseBuilder {
     stream: boolean,
     parentEvent?: Event,
   ) {
-    const textQATemplate: SimplePrompt = (input) =>
-      this.textQATemplate({ ...input, query: queryStr });
+    const textQATemplate: SimplePrompt = (input) => {
+      this.textQATemplate.partialFormat({ ...input, query: queryStr });
+      return this.textQATemplate.format();
+    };
+
     const textChunks = this.promptHelper.repack(textQATemplate, [textChunk]);
 
     let response: AsyncIterable<string> | string | undefined = undefined;
 
     for (let i = 0; i < textChunks.length; i++) {
       const chunk = textChunks[i];
+
+      this.textQATemplate.partialFormat({ context: chunk, query: queryStr });
+
       const lastChunk = i === textChunks.length - 1;
       if (!response) {
         response = await this.complete({
-          prompt: textQATemplate({
-            context: chunk,
-          }),
+          prompt: this.textQATemplate,
           parentEvent,
           stream: stream && lastChunk,
         });
@@ -205,8 +217,10 @@ export class Refine extends PromptMixin implements ResponseBuilder {
     stream: boolean,
     parentEvent?: Event,
   ) {
-    const refineTemplate: SimplePrompt = (input) =>
-      this.refineTemplate({ ...input, query: queryStr });
+    const refineTemplate: SimplePrompt = (input) => {
+      this.refineTemplate.partialFormat({ ...input, query: queryStr });
+      return this.refineTemplate.format();
+    };
 
     const textChunks = this.promptHelper.repack(refineTemplate, [textChunk]);
 
@@ -214,12 +228,16 @@ export class Refine extends PromptMixin implements ResponseBuilder {
 
     for (let i = 0; i < textChunks.length; i++) {
       const chunk = textChunks[i];
+
+      this.refineTemplate.partialFormat({
+        context: chunk,
+        existingAnswer: response as string,
+      });
+
       const lastChunk = i === textChunks.length - 1;
+
       response = await this.complete({
-        prompt: refineTemplate({
-          context: chunk,
-          existingAnswer: response as string,
-        }),
+        prompt: this.refineTemplate,
         parentEvent,
         stream: stream && lastChunk,
       });
@@ -228,15 +246,15 @@ export class Refine extends PromptMixin implements ResponseBuilder {
   }
 
   async complete(params: {
-    prompt: string;
+    prompt: BasePromptTemplate | string;
     stream: boolean;
     parentEvent?: Event;
   }): Promise<AsyncIterable<string> | string> {
     if (params.stream) {
-      const response = await this.llm.complete({ ...params, stream: true });
+      const response = await this.llm.predict({ ...params, stream: true });
       return streamConverter(response, (chunk) => chunk.text);
     } else {
-      const response = await this.llm.complete({ ...params, stream: false });
+      const response = await this.llm.predict({ ...params, stream: false });
       return response.text;
     }
   }
@@ -261,10 +279,8 @@ export class CompactAndRefine extends Refine {
     | ResponseBuilderParamsNonStreaming): Promise<
     AsyncIterable<string> | string
   > {
-    const textQATemplate: SimplePrompt = (input) =>
-      this.textQATemplate({ ...input, query: query });
-    const refineTemplate: SimplePrompt = (input) =>
-      this.refineTemplate({ ...input, query: query });
+    const textQATemplate: SimplePrompt = () => this.textQATemplate.format();
+    const refineTemplate: SimplePrompt = () => this.refineTemplate.format();
 
     const maxPrompt = getBiggestPrompt([textQATemplate, refineTemplate]);
     const newTexts = this.promptHelper.repack(maxPrompt, textChunks);

packages/core/src/test.ts

@@ -0,0 +1,71 @@
+import {
+  ChatPromptTemplate,
+  PromptTemplate,
+  SimpleDirectoryReader,
+  VectorStoreIndex,
+  type ChatMessage,
+} from "./index.js";
+
+// Define a default prompt template
+const defaultTextQaPromptTemplate = () => `Context information is below.
+---------------------
+{{context}}
+---------------------
+Given the context information and not prior knowledge, answer the query.
+Query: {{query}}
+Answer:`;
+
+// Define a default chat prompt template
+const defaultChatPromtTemplate = (): ChatMessage[] => [
+  {
+    content: "Always answer the question, even if you don't know the answer.",
+    role: "system",
+  },
+  {
+    content: defaultTextQaPromptTemplate(),
+    role: "user",
+  },
+];
+
+// Instantiate the prompt templates
+const textQaPromptTemplate = new PromptTemplate(defaultTextQaPromptTemplate);
+const chatMessageTemplate = new ChatPromptTemplate(defaultChatPromtTemplate);
+
+// Map the template variables (get from the prompt templates all text that includes {{text}})
+console.log({
+  textQaVars: textQaPromptTemplate.mapTemplateVars(),
+  chatTemplateVars: chatMessageTemplate.mapTemplateVars(),
+});
+
+async function main() {
+  const documents = await new SimpleDirectoryReader().loadData({
+    directoryPath: "../examples",
+  });
+
+  const index = await VectorStoreIndex.fromDocuments(documents);
+
+  const retriever = await index.asRetriever({});
+  retriever.similarityTopK = 10;
+
+  const queryEngine = index.asQueryEngine({
+    retriever,
+  });
+
+  // Update the prompts
+  queryEngine.updatePrompts({
+    "responseSynthesizer:textQATemplate": chatMessageTemplate,
+  });
+
+  // Query the engine
+  const query = "Tell me about abramov";
+
+  const response = await queryEngine.query({
+    query,
+  });
+
+  console.log({
+    response,
+  });
+}
+
+main();

packages/edge/package.json

@@ -27,6 +27,7 @@
     "lodash": "^4.17.21",
     "magic-bytes.js": "^1.10.0",
     "mammoth": "^1.6.0",
+    "mustache": "^4.2.0",
     "md-utils-ts": "^2.0.0",
     "mongodb": "^6.3.0",
     "notion-md-crawler": "^0.0.2",

pnpm-lock.yaml

@@ -243,6 +243,9 @@ importers:
       mongodb:
         specifier: ^6.3.0
         version: 6.3.0
+      mustache:
+        specifier: ^4.2.0
+        version: 4.2.0
       notion-md-crawler:
         specifier: ^0.0.2
         version: 0.0.2
@@ -286,6 +289,9 @@ importers:
       '@swc/core':
         specifier: ^1.4.2
         version: 1.4.2
+      '@types/mustache':
+        specifier: ^4.2.5
+        version: 4.2.5
       concurrently:
         specifier: ^8.2.2
         version: 8.2.2
@@ -4568,6 +4574,9 @@ packages:
   /@types/ms@0.7.34:
     resolution: {integrity: sha512-nG96G3Wp6acyAgJqGasjODb+acrI7KltPiRxzHPXnP3NgI28bpQDRv53olbqGXbfcgF5aiiHmO3xpwEpS5Ld9g==}
 
+  /@types/mustache@4.2.5:
+    resolution: {integrity: sha512-PLwiVvTBg59tGFL/8VpcGvqOu3L4OuveNvPi0EYbWchRdEVP++yRUXJPFl+CApKEq13017/4Nf7aQ5lTtHUNsA==}
+    dev: true
   /@types/node-fetch@2.6.11:
     resolution: {integrity: sha512-24xFj9R5+rfQJLRyM56qh+wnVSYhyXC2tkoBndtY0U+vubqNsYXGjufB2nn8Q6gt0LrARwL6UBtMCSVCwl4B1g==}
     dependencies:
@@ -11063,6 +11072,11 @@ packages:
       dns-packet: 5.6.1
       thunky: 1.1.0
 
+  /mustache@4.2.0:
+    resolution: {integrity: sha512-71ippSywq5Yb7/tVYyGbkBggbU8H3u5Rz56fH60jGFgr8uHwxs+aSKeqmluIVzM0m0kB7xQjKS6qPfd0b2ZoqQ==}
+    hasBin: true
+    dev: false
+
   /mz@2.7.0:
     resolution: {integrity: sha512-z81GNO7nnYMEhrGh9LeymoE4+Yr0Wn5McHIZMK5cfQCl+NDX08sCZgUc9/6MHni9IWuFLm1Z3HTCXu2z9fN62Q==}
     dependencies: