import { expect, test } from "@jest/globals";
import { PromptTemplate } from "../prompt.js";
import { Document } from "../../documents/document.js";

test("Test using partial", async () => {
  const prompt = new PromptTemplate({
    template: "{foo}{bar}",
    inputVariables: ["foo"],
    partialVariables: { bar: "baz" },
  });
  expect(await prompt.format({ foo: "foo" })).toBe("foobaz");
});

test("Test using partial with an extra variable", async () => {
  const prompt = new PromptTemplate({
    template: "{foo}{bar}",
    inputVariables: ["foo"],
    partialVariables: { bar: "baz" },
  });
  expect(await prompt.format({ foo: "foo", unused: "nada" })).toBe("foobaz");
});

test("Test fromTemplate", async () => {
  const prompt = PromptTemplate.fromTemplate("{foo}{bar}");
  expect(
    (await prompt.invoke({ foo: "foo", bar: "baz", unused: "eee" })).value
  ).toBe("foobaz");
});

test("Test fromTemplate with a non-string value", async () => {
  const prompt = PromptTemplate.fromTemplate("{foo}{bar}");
  expect(
    (
      await prompt.invoke({
        foo: ["barbar"],
        bar: [new Document({ pageContent: "bar" })],
      })
    ).value
  ).toBe(`["barbar"][{"pageContent":"bar","metadata":{}}]`);
});

test("Test fromTemplate with escaped strings", async () => {
  const prompt = PromptTemplate.fromTemplate("{{foo}}{{bar}}");
  expect(await prompt.format({ unused: "eee" })).toBe("{foo}{bar}");
});

test("Test fromTemplate with type parameter", async () => {
  const prompt = PromptTemplate.fromTemplate<{ foo: string }>("test");
  // @ts-expect-error TS compiler should flag
  expect(await prompt.format({ unused: "eee" })).toBe("test");
});

test("Test fromTemplate with missing variable should raise compiler error", async () => {
  const prompt = PromptTemplate.fromTemplate("{foo}");
  await expect(async () => {
    // @ts-expect-error TS compiler should flag missing variable
    await prompt.format({ unused: "eee" });
  }).rejects.toThrow();
  await expect(async () => {
    // @ts-expect-error TS compiler should flag missing variable
    await prompt.invoke({ unused: "eee" });
  }).rejects.toThrow();
});

test("Test fromTemplate with extra variable should work", async () => {
  const prompt = PromptTemplate.fromTemplate("{foo}");
  expect(await prompt.format({ foo: "test", unused: "eee" })).toBe("test");
  expect((await prompt.invoke({ foo: "test", unused: "eee" })).value).toBe(
    "test"
  );
});

test("Test using full partial", async () => {
  const prompt = new PromptTemplate({
    template: "{foo}{bar}",
    inputVariables: [],
    partialVariables: { bar: "baz", foo: "boo" },
  });
  expect(await prompt.format({})).toBe("boobaz");
});

test("Test partial", async () => {
  const prompt = new PromptTemplate({
    template: "{foo}{bar}",
    inputVariables: ["foo", "bar"],
  });
  expect(prompt.inputVariables).toEqual(["foo", "bar"]);
  const partialPrompt = await prompt.partial({ foo: "foo" });
  // original prompt is not modified
  expect(prompt.inputVariables).toEqual(["foo", "bar"]);
  // partial prompt has only remaining variables
  expect(partialPrompt.inputVariables).toEqual(["bar"]);
  expect(await partialPrompt.format({ bar: "baz" })).toBe("foobaz");
});

test("Test partial with function", async () => {
  const prompt = new PromptTemplate({
    template: "{foo}{bar}",
    inputVariables: ["foo", "bar"],
  });
  const partialPrompt = await prompt.partial({
    foo: () => Promise.resolve("boo"),
  });
  expect(await partialPrompt.format({ bar: "baz" })).toBe("boobaz");
});


/* eslint-disable @typescript-eslint/no-explicit-any */

import { expect, test } from "@jest/globals";
import {
  AIMessagePromptTemplate,
  ChatPromptTemplate,
  ChatMessagePromptTemplate,
  HumanMessagePromptTemplate,
  SystemMessagePromptTemplate,
  MessagesPlaceholder,
} from "../chat.js";
import { PromptTemplate } from "../prompt.js";
import {
  SystemMessage,
  HumanMessage,
  AIMessage,
  ChatMessage,
  FunctionMessage,
} from "../../messages/index.js";
import { Document } from "../../documents/document.js";

function createChatPromptTemplate() {
  const systemPrompt = new PromptTemplate({
    template: "Here's some context: {context}",
    inputVariables: ["context"],
  });
  const userPrompt = new PromptTemplate({
    template: "Hello {foo}, I'm {bar}. Thanks for the {context}",
    inputVariables: ["foo", "bar", "context"],
  });
  const aiPrompt = new PromptTemplate({
    template: "I'm an AI. I'm {foo}. I'm {bar}.",
    inputVariables: ["foo", "bar"],
  });
  const genericPrompt = new PromptTemplate({
    template: "I'm a generic message. I'm {foo}. I'm {bar}.",
    inputVariables: ["foo", "bar"],
  });
  // return new ChatPromptTemplate({
  //   promptMessages: [
  //     new SystemMessagePromptTemplate(systemPrompt),
  //     new HumanMessagePromptTemplate(userPrompt),
  //     new AIMessagePromptTemplate({ prompt: aiPrompt }),
  //     new ChatMessagePromptTemplate(genericPrompt, "test"),
  //   ],
  //   inputVariables: ["context", "foo", "bar"],
  // });
  return ChatPromptTemplate.fromMessages<{
    foo: string;
    bar: string;
    context: string;
  }>([
    new SystemMessagePromptTemplate(systemPrompt),
    new HumanMessagePromptTemplate(userPrompt),
    new AIMessagePromptTemplate({ prompt: aiPrompt }),
    new ChatMessagePromptTemplate(genericPrompt, "test"),
  ]);
}

test("Test format", async () => {
  const chatPrompt = createChatPromptTemplate();
  const messages = await chatPrompt.formatPromptValue({
    context: "This is a context",
    foo: "Foo",
    bar: "Bar",
    unused: "extra",
  });
  expect(messages.toChatMessages()).toEqual([
    new SystemMessage("Here's some context: This is a context"),
    new HumanMessage("Hello Foo, I'm Bar. Thanks for the This is a context"),
    new AIMessage("I'm an AI. I'm Foo. I'm Bar."),
    new ChatMessage("I'm a generic message. I'm Foo. I'm Bar.", "test"),
  ]);
});

test("Test format with invalid input values", async () => {
  const chatPrompt = createChatPromptTemplate();
  let error: any | undefined;
  try {
    // @ts-expect-error TS compiler should flag missing input variables
    await chatPrompt.formatPromptValue({
      context: "This is a context",
      foo: "Foo",
    });
  } catch (e) {
    error = e;
  }
  expect(error?.message).toContain("Missing value for input variable `bar`");
  expect(error?.lc_error_code).toEqual("INVALID_PROMPT_INPUT");
});

test("Test format with invalid input variables", async () => {
  const systemPrompt = new PromptTemplate({
    template: "Here's some context: {context}",
    inputVariables: ["context"],
  });
  const userPrompt = new PromptTemplate({
    template: "Hello {foo}, I'm {bar}",
    inputVariables: ["foo", "bar"],
  });
  expect(
    () =>
      new ChatPromptTemplate({
        promptMessages: [
          new SystemMessagePromptTemplate(systemPrompt),
          new HumanMessagePromptTemplate(userPrompt),
        ],
        inputVariables: ["context", "foo", "bar", "baz"],
      })
  ).toThrow(
    "Input variables `baz` are not used in any of the prompt messages."
  );

  expect(
    () =>
      new ChatPromptTemplate({
        promptMessages: [
          new SystemMessagePromptTemplate(systemPrompt),
          new HumanMessagePromptTemplate(userPrompt),
        ],
        inputVariables: ["context", "foo"],
      })
  ).toThrow(
    "Input variables `bar` are used in prompt messages but not in the prompt template."
  );
});

test("Test fromTemplate", async () => {
  const chatPrompt = ChatPromptTemplate.fromTemplate("Hello {foo}, I'm {bar}");
  expect(chatPrompt.inputVariables).toEqual(["foo", "bar"]);
  const messages = await chatPrompt.formatPromptValue({
    foo: "Foo",
    bar: "Bar",
  });
  expect(messages.toChatMessages()).toEqual([
    new HumanMessage("Hello Foo, I'm Bar"),
  ]);
});

test("Test fromTemplate", async () => {
  const chatPrompt = ChatPromptTemplate.fromTemplate("Hello {foo}, I'm {bar}");
  expect(chatPrompt.inputVariables).toEqual(["foo", "bar"]);
  expect(
    (
      await chatPrompt.invoke({
        foo: ["barbar"],
        bar: [new Document({ pageContent: "bar" })],
      })
    ).toChatMessages()
  ).toEqual([
    new HumanMessage(
      `Hello ["barbar"], I'm [{"pageContent":"bar","metadata":{}}]`
    ),
  ]);
});

test("Test fromMessages", async () => {
  const systemPrompt = new PromptTemplate({
    template: "Here's some context: {context}",
    inputVariables: ["context"],
  });
  const userPrompt = new PromptTemplate({
    template: "Hello {foo}, I'm {bar}",
    inputVariables: ["foo", "bar"],
  });
  // TODO: Fix autocomplete for the fromMessages method
  const chatPrompt = ChatPromptTemplate.fromMessages([
    new SystemMessagePromptTemplate(systemPrompt),
    new HumanMessagePromptTemplate(userPrompt),
  ]);
  expect(chatPrompt.inputVariables).toEqual(["context", "foo", "bar"]);
  const messages = await chatPrompt.formatPromptValue({
    context: "This is a context",
    foo: "Foo",
    bar: "Bar",
  });
  expect(messages.toChatMessages()).toEqual([
    new SystemMessage("Here's some context: This is a context"),
    new HumanMessage("Hello Foo, I'm Bar"),
  ]);
});

test("Test fromMessages with non-string inputs", async () => {
  const systemPrompt = new PromptTemplate({
    template: "Here's some context: {context}",
    inputVariables: ["context"],
  });
  const userPrompt = new PromptTemplate({
    template: "Hello {foo}, I'm {bar}",
    inputVariables: ["foo", "bar"],
  });
  // TODO: Fix autocomplete for the fromMessages method
  const chatPrompt = ChatPromptTemplate.fromMessages([
    new SystemMessagePromptTemplate(systemPrompt),
    new HumanMessagePromptTemplate(userPrompt),
  ]);
  expect(chatPrompt.inputVariables).toEqual(["context", "foo", "bar"]);
  const messages = await chatPrompt.formatPromptValue({
    context: [new Document({ pageContent: "bar" })],
    foo: "Foo",
    bar: "Bar",
  });
  expect(messages.toChatMessages()).toEqual([
    new SystemMessage(
      `Here's some context: [{"pageContent":"bar","metadata":{}}]`
    ),
    new HumanMessage("Hello Foo, I'm Bar"),
  ]);
});

test("Test fromMessages with a variety of ways to declare prompt messages", async () => {
  const systemPrompt = new PromptTemplate({
    template: "Here's some context: {context}",
    inputVariables: ["context"],
  });
  // TODO: Fix autocomplete for the fromMessages method
  const chatPrompt = ChatPromptTemplate.fromMessages([
    new SystemMessagePromptTemplate(systemPrompt),
    "Hello {foo}, I'm {bar}",
    ["assistant", "Nice to meet you, {bar}!"],
    ["human", "Thanks {foo}!!"],
  ]);
  const messages = await chatPrompt.formatPromptValue({
    context: "This is a context",
    foo: "Foo",
    bar: "Bar",
  });
  expect(messages.toChatMessages()).toEqual([
    new SystemMessage("Here's some context: This is a context"),
    new HumanMessage("Hello Foo, I'm Bar"),
    new AIMessage("Nice to meet you, Bar!"),
    new HumanMessage("Thanks Foo!!"),
  ]);
});


import { Document } from "../documents/document.js";
import { BaseDocumentTransformer } from "../documents/transformers.js";

/**
 * Interface that defines the methods for loading and splitting documents.
 */
export interface DocumentLoader {
  load(): Promise<Document[]>;
  loadAndSplit(textSplitter?: BaseDocumentTransformer): Promise<Document[]>;
}

/**
 * Abstract class that provides a default implementation for the
 * loadAndSplit() method from the DocumentLoader interface. The load()
 * method is left abstract and needs to be implemented by subclasses.
 */
export abstract class BaseDocumentLoader implements DocumentLoader {
  /**
   * Loads the documents.
   * @returns A Promise that resolves with an array of Document instances.
   */
  abstract load(): Promise<Document[]>;

  /**
   * @deprecated Use `this.load()` and `splitter.splitDocuments()` individually.
   * Loads the documents and splits them using a specified text splitter.
   * @param textSplitter The TextSplitter instance to use for splitting the loaded documents. Defaults to a RecursiveCharacterTextSplitter instance.
   * @returns A Promise that resolves with an array of Document instances, each split according to the provided TextSplitter.
   */
  async loadAndSplit(splitter?: BaseDocumentTransformer): Promise<Document[]> {
    if (splitter === undefined) {
      throw new Error("You must pass a text splitter to use this method.");
    }
    const docs = await this.load();
    return splitter.invoke(docs);
  }
}


export interface DocumentInput<
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  Metadata extends Record<string, any> = Record<string, any>
> {
  pageContent: string;

  metadata?: Metadata;

  /**
   * An optional identifier for the document.
   *
   * Ideally this should be unique across the document collection and formatted
   * as a UUID, but this will not be enforced.
   */
  id?: string;
}

export interface DocumentInterface<
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  Metadata extends Record<string, any> = Record<string, any>
> {
  pageContent: string;

  metadata: Metadata;

  /**
   * An optional identifier for the document.
   *
   * Ideally this should be unique across the document collection and formatted
   * as a UUID, but this will not be enforced.
   */
  id?: string;
}

/**
 * Interface for interacting with a document.
 */
export class Document<
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  Metadata extends Record<string, any> = Record<string, any>
> implements DocumentInput, DocumentInterface
{
  pageContent: string;

  metadata: Metadata;

  // The ID field is optional at the moment.
  // It will likely become required in a future major release after
  // it has been adopted by enough vectorstore implementations.
  /**
   * An optional identifier for the document.
   *
   * Ideally this should be unique across the document collection and formatted
   * as a UUID, but this will not be enforced.
   */
  id?: string;

  constructor(fields: DocumentInput<Metadata>) {
    this.pageContent =
      fields.pageContent !== undefined ? fields.pageContent.toString() : "";
    this.metadata = fields.metadata ?? ({} as Metadata);
    this.id = fields.id;
  }
}


import { DocumentInterface } from "@langchain/core/documents";
import { Embeddings } from "@langchain/core/embeddings";
import {
  cosineSimilarity,
  euclideanDistance,
  innerProduct,
} from "@langchain/core/utils/math";
import {
  VectorStore,
  VectorStoreRetriever,
  VectorStoreRetrieverInput,
} from "@langchain/core/vectorstores";

/**
 * Type for options when adding a document to the VectorStore.
 */
// eslint-disable-next-line @typescript-eslint/no-explicit-any
type AddDocumentOptions = Record<string, any>;

export interface MatryoshkaRetrieverFields {
  /**
   * The number of documents to retrieve from the small store.
   * @default 50
   */
  smallK?: number;
  /**
   * The number of documents to retrieve from the large store.
   * @default 8
   */
  largeK?: number;
  /**
   * The metadata key to store the larger embeddings.
   * @default "lc_large_embedding"
   */
  largeEmbeddingKey?: string;
  /**
   * The embedding model to use when generating the large
   * embeddings.
   */
  largeEmbeddingModel: Embeddings;
  /**
   * The type of search to perform using the large embeddings.
   * @default "cosine"
   */
  searchType?: "cosine" | "innerProduct" | "euclidean";
}

/**
 * A retriever that uses two sets of embeddings to perform adaptive retrieval. Based
 * off of the "Matryoshka embeddings: faster OpenAI vector search using Adaptive Retrieval"
 * blog post {@link https://supabase.com/blog/matryoshka-embeddings}.
 *
 *
 * This class performs "Adaptive Retrieval" for searching text embeddings efficiently using the
 * Matryoshka Representation Learning (MRL) technique. It retrieves documents similar to a query
 * embedding in two steps:
 *
 * First-pass: Uses a lower dimensional sub-vector from the MRL embedding for an initial, fast,
 * but less accurate search.
 *
 * Second-pass: Re-ranks the top results from the first pass using the full, high-dimensional
 * embedding for higher accuracy.
 *
 *
 * This code implements MRL embeddings for efficient vector search by combining faster,
 * lower-dimensional initial search with accurate, high-dimensional re-ranking.
 */
export class MatryoshkaRetriever<
  Store extends VectorStore = VectorStore
> extends VectorStoreRetriever<Store> {
  smallK = 50;

  largeK = 8;

  largeEmbeddingKey = "lc_large_embedding";

  largeEmbeddingModel: Embeddings;

  searchType: "cosine" | "innerProduct" | "euclidean" = "cosine";

  constructor(
    fields: MatryoshkaRetrieverFields & VectorStoreRetrieverInput<Store>
  ) {
    super(fields);
    this.smallK = fields.smallK ?? this.smallK;
    this.largeK = fields.largeK ?? this.largeK;
    this.largeEmbeddingKey = fields.largeEmbeddingKey ?? this.largeEmbeddingKey;
    this.largeEmbeddingModel = fields.largeEmbeddingModel;
    this.searchType = fields.searchType ?? this.searchType;
  }

  /**
   * Ranks documents based on their similarity to a query embedding using larger embeddings.
   *
   * This method takes a query embedding and a list of documents (smallResults) as input. Each document
   * in the smallResults array has previously been associated with a large embedding stored in its metadata.
   * Depending on the `searchType` (cosine, innerProduct, or euclidean), it calculates the similarity scores
   * between the query embedding and each document's large embedding. It then ranks the documents based on
   * these similarity scores, from the most similar to the least similar.
   *
   * The method returns a promise that resolves to an array of the top `largeK` documents, where `largeK`
   * is a class property defining the number of documents to return. This subset of documents is determined
   * by sorting the entire list of documents based on their similarity scores and then selecting the top
   * `largeK` documents.
   *
   * @param {number[]} embeddedQuery The embedding of the query, represented as an array of numbers.
   * @param {DocumentInterface[]} smallResults An array of documents, each with metadata that includes a large embedding for similarity comparison.
   * @returns {Promise<DocumentInterface[]>} A promise that resolves to an array of the top `largeK` ranked documents based on their similarity to the query embedding.
   */
  private _rankByLargeEmbeddings(
    embeddedQuery: number[],
    smallResults: DocumentInterface[]
  ): DocumentInterface[] {
    const largeEmbeddings: Array<number[]> = smallResults.map((doc) =>
      JSON.parse(doc.metadata[this.largeEmbeddingKey])
    );
    let func: () => Array<number[]>;

    switch (this.searchType) {
      case "cosine":
        func = () => cosineSimilarity([embeddedQuery], largeEmbeddings);
        break;
      case "innerProduct":
        func = () => innerProduct([embeddedQuery], largeEmbeddings);
        break;
      case "euclidean":
        func = () => euclideanDistance([embeddedQuery], largeEmbeddings);
        break;
      default:
        throw new Error(`Unknown search type: ${this.searchType}`);
    }

    // Calculate the similarity scores between the query embedding and the large embeddings
    const [similarityScores] = func();

    // Create an array of indices from 0 to N-1, where N is the number of documents
    let indices = Array.from(
      { length: smallResults.length },
      (_, index) => index
    );

    indices = indices
      .map((v, i) => [similarityScores[i], v])
      .sort(([a], [b]) => b - a)
      .slice(0, this.largeK)
      .map(([, i]) => i);

    return indices.map((i) => smallResults[i]);
  }

  async _getRelevantDocuments(query: string): Promise<DocumentInterface[]> {
    const [embeddedQuery, smallResults] = await Promise.all([
      this.largeEmbeddingModel.embedQuery(query),
      this.vectorStore.similaritySearch(query, this.smallK, this.filter),
    ]);

    return this._rankByLargeEmbeddings(embeddedQuery, smallResults);
  }

  /**
   * Override the default `addDocuments` method to embed the documents twice,
   * once using the larger embeddings model, and then again using the default
   * embedding model linked to the vector store.
   *
   * @param {DocumentInterface[]} documents - An array of documents to add to the vector store.
   * @param {AddDocumentOptions} options - An optional object containing additional options for adding documents.
   * @returns {Promise<string[] | void>} A promise that resolves to an array of the document IDs that were added to the vector store.
   */
  override addDocuments = async (
    documents: DocumentInterface[],
    options?: AddDocumentOptions
  ): Promise<string[] | void> => {
    // Insure documents metadata does not contain the large embedding key
    if (documents.some((doc) => this.largeEmbeddingKey in doc.metadata)) {
      throw new Error(
        `All documents must not contain the large embedding key: ${this.largeEmbeddingKey} in their metadata.`
      );
    }

    const allDocPageContent = documents.map((doc) => doc.pageContent);
    const allDocLargeEmbeddings = await this.largeEmbeddingModel.embedDocuments(
      allDocPageContent
    );

    const newDocuments: Array<DocumentInterface> = documents.map(
      (doc, idx) => ({
        ...doc,
        metadata: {
          ...doc.metadata,
          [this.largeEmbeddingKey]: JSON.stringify(allDocLargeEmbeddings[idx]),
        },
      })
    );

    return this.vectorStore.addDocuments(newDocuments, options);
  };
}


import {
  BaseRetriever,
  type BaseRetrieverInput,
  type BaseRetrieverInterface,
} from "@langchain/core/retrievers";
import type { DocumentInterface } from "@langchain/core/documents";
import { CallbackManagerForRetrieverRun } from "@langchain/core/callbacks/manager";
import { BaseDocumentCompressor } from "./document_compressors/index.js";

/**
 * Interface for the arguments required to construct a
 * ContextualCompressionRetriever. It extends the BaseRetrieverInput
 * interface with two additional fields: baseCompressor and baseRetriever.
 */
export interface ContextualCompressionRetrieverArgs extends BaseRetrieverInput {
  baseCompressor: BaseDocumentCompressor;
  baseRetriever: BaseRetrieverInterface;
}

/**
 * A retriever that wraps a base retriever and compresses the results. It
 * retrieves relevant documents based on a given query and then compresses
 * these documents using a specified document compressor.
 * @example
 * ```typescript
 * const retriever = new ContextualCompressionRetriever({
 *   baseCompressor: new LLMChainExtractor(),
 *   baseRetriever: new HNSWLib().asRetriever(),
 * });
 * const retrievedDocs = await retriever.getRelevantDocuments(
 *   "What did the speaker say about Justice Breyer?",
 * );
 * ```
 */
export class ContextualCompressionRetriever extends BaseRetriever {
  static lc_name() {
    return "ContextualCompressionRetriever";
  }

  lc_namespace = ["langchain", "retrievers", "contextual_compression"];

  baseCompressor: BaseDocumentCompressor;

  baseRetriever: BaseRetrieverInterface;

  constructor(fields: ContextualCompressionRetrieverArgs) {
    super(fields);

    this.baseCompressor = fields.baseCompressor;
    this.baseRetriever = fields.baseRetriever;
  }

  async _getRelevantDocuments(
    query: string,
    runManager?: CallbackManagerForRetrieverRun
  ): Promise<DocumentInterface[]> {
    const docs = await this.baseRetriever.getRelevantDocuments(
      query,
      runManager?.getChild("base_retriever")
    );
    const compressedDocs = await this.baseCompressor.compressDocuments(
      docs,
      query,
      runManager?.getChild("base_compressor")
    );
    return compressedDocs;
  }
}


import type { EmbeddingsInterface } from "@langchain/core/embeddings";
import type { DocumentInterface } from "@langchain/core/documents";
import { cosineSimilarity } from "@langchain/core/utils/math";
import { BaseDocumentCompressor } from "./index.js";

/**
 * Interface for the parameters of the `EmbeddingsFilter` class.
 */
export interface EmbeddingsFilterParams {
  embeddings: EmbeddingsInterface;
  similarityFn?: (x: number[][], y: number[][]) => number[][];
  similarityThreshold?: number;
  k?: number;
}

/**
 * Class that represents a document compressor that uses embeddings to
 * drop documents unrelated to the query.
 * @example
 * ```typescript
 * const embeddingsFilter = new EmbeddingsFilter({
 *   embeddings: new OpenAIEmbeddings(),
 *   similarityThreshold: 0.8,
 *   k: 5,
 * });
 * const retrievedDocs = await embeddingsFilter.filterDocuments(
 *   getDocuments(),
 *   "What did the speaker say about Justice Breyer in the 2022 State of the Union?",
 * );
 * console.log({ retrievedDocs });
 * ```
 */
export class EmbeddingsFilter extends BaseDocumentCompressor {
  /**
   * Embeddings to use for embedding document contents and queries.
   */
  embeddings: EmbeddingsInterface;

  /**
   * Similarity function for comparing documents.
   */
  similarityFn = cosineSimilarity;

  /**
   * Threshold for determining when two documents are similar enough
   * to be considered redundant. Must be specified if `k` is not set.
   */
  similarityThreshold?: number;

  /**
   * The number of relevant documents to return. Can be explicitly set to undefined, in which case
   * similarity_threshold` must be specified. Defaults to 20
   */
  k? = 20;

  constructor(params: EmbeddingsFilterParams) {
    super();
    this.embeddings = params.embeddings;
    this.similarityFn = params.similarityFn ?? this.similarityFn;
    this.similarityThreshold = params.similarityThreshold;
    this.k = "k" in params ? params.k : this.k;
    if (this.k === undefined && this.similarityThreshold === undefined) {
      throw new Error(`Must specify one of "k" or "similarity_threshold".`);
    }
  }

  async compressDocuments(
    documents: DocumentInterface[],
    query: string
  ): Promise<DocumentInterface[]> {
    const embeddedDocuments = await this.embeddings.embedDocuments(
      documents.map((doc) => doc.pageContent)
    );
    const embeddedQuery = await this.embeddings.embedQuery(query);
    const similarity = this.similarityFn([embeddedQuery], embeddedDocuments)[0];
    let includedIdxs = Array.from(
      { length: embeddedDocuments.length },
      (_, i) => i
    );

    if (this.k !== undefined) {
      includedIdxs = includedIdxs
        .map((v, i) => [similarity[i], v])
        .sort(([a], [b]) => b - a)
        .slice(0, this.k)
        .map(([, i]) => i);
    }

    if (this.similarityThreshold !== undefined) {
      const threshold = this.similarityThreshold;
      includedIdxs = includedIdxs.filter((i) => similarity[i] > threshold);
    }

    return includedIdxs.map((i) => documents[i]);
  }
}


import type { BaseLanguageModelInterface } from "@langchain/core/language_models/base";
import {
  BaseMessage,
  SystemMessage,
  getBufferString,
} from "@langchain/core/messages";
import { BasePromptTemplate } from "@langchain/core/prompts";
import {
  InputValues,
  MemoryVariables,
  OutputValues,
} from "@langchain/core/memory";
import { LLMChain } from "../chains/llm_chain.js";
import { SUMMARY_PROMPT } from "./prompt.js";
import { BaseChatMemory, BaseChatMemoryInput } from "./chat_memory.js";

/**
 * Interface for the input parameters of the ConversationSummaryMemory
 * class.
 */
export interface ConversationSummaryMemoryInput
  extends BaseConversationSummaryMemoryInput {}

/**
 * Interface for the input parameters of the BaseConversationSummaryMemory
 * class.
 */
export interface BaseConversationSummaryMemoryInput
  extends BaseChatMemoryInput {
  llm: BaseLanguageModelInterface;
  memoryKey?: string;
  humanPrefix?: string;
  aiPrefix?: string;
  prompt?: BasePromptTemplate;
  summaryChatMessageClass?: new (content: string) => BaseMessage;
}

/**
 * Abstract class that provides a structure for storing and managing the
 * memory of a conversation. It includes methods for predicting a new
 * summary for the conversation given the existing messages and summary.
 */
export abstract class BaseConversationSummaryMemory extends BaseChatMemory {
  memoryKey = "history";

  humanPrefix = "Human";

  aiPrefix = "AI";

  llm: BaseLanguageModelInterface;

  prompt: BasePromptTemplate = SUMMARY_PROMPT;

  summaryChatMessageClass: new (content: string) => BaseMessage = SystemMessage;

  constructor(fields: BaseConversationSummaryMemoryInput) {
    const {
      returnMessages,
      inputKey,
      outputKey,
      chatHistory,
      humanPrefix,
      aiPrefix,
      llm,
      prompt,
      summaryChatMessageClass,
    } = fields;

    super({ returnMessages, inputKey, outputKey, chatHistory });

    this.memoryKey = fields?.memoryKey ?? this.memoryKey;
    this.humanPrefix = humanPrefix ?? this.humanPrefix;
    this.aiPrefix = aiPrefix ?? this.aiPrefix;
    this.llm = llm;
    this.prompt = prompt ?? this.prompt;
    this.summaryChatMessageClass =
      summaryChatMessageClass ?? this.summaryChatMessageClass;
  }

  /**
   * Predicts a new summary for the conversation given the existing messages
   * and summary.
   * @param messages Existing messages in the conversation.
   * @param existingSummary Current summary of the conversation.
   * @returns A promise that resolves to a new summary string.
   */
  async predictNewSummary(
    messages: BaseMessage[],
    existingSummary: string
  ): Promise<string> {
    const newLines = getBufferString(messages, this.humanPrefix, this.aiPrefix);
    const chain = new LLMChain({ llm: this.llm, prompt: this.prompt });
    return await chain.predict({
      summary: existingSummary,
      new_lines: newLines,
    });
  }
}

/**
 * Class that provides a concrete implementation of the conversation
 * memory. It includes methods for loading memory variables, saving
 * context, and clearing the memory.
 * @example
 * ```typescript
 * const memory = new ConversationSummaryMemory({
 *   memoryKey: "chat_history",
 *   llm: new ChatOpenAI({ modelName: "gpt-3.5-turbo", temperature: 0 }),
 * });
 *
 * const model = new ChatOpenAI();
 * const prompt =
 *   PromptTemplate.fromTemplate(`The following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.
 *
 * Current conversation:
 * {chat_history}
 * Human: {input}
 * AI:`);
 * const chain = new LLMChain({ llm: model, prompt, memory });
 *
 * const res1 = await chain.call({ input: "Hi! I'm Jim." });
 * console.log({ res1, memory: await memory.loadMemoryVariables({}) });
 *
 * const res2 = await chain.call({ input: "What's my name?" });
 * console.log({ res2, memory: await memory.loadMemoryVariables({}) });
 *
 * ```
 */
export class ConversationSummaryMemory extends BaseConversationSummaryMemory {
  buffer = "";

  constructor(fields: ConversationSummaryMemoryInput) {
    super(fields);
  }

  get memoryKeys() {
    return [this.memoryKey];
  }

  /**
   * Loads the memory variables for the conversation memory.
   * @returns A promise that resolves to an object containing the memory variables.
   */
  async loadMemoryVariables(_: InputValues): Promise<MemoryVariables> {
    if (this.returnMessages) {
      const result = {
        [this.memoryKey]: [new this.summaryChatMessageClass(this.buffer)],
      };
      return result;
    }
    const result = { [this.memoryKey]: this.buffer };
    return result;
  }

  /**
   * Saves the context of the conversation memory.
   * @param inputValues Input values for the conversation.
   * @param outputValues Output values from the conversation.
   * @returns A promise that resolves when the context has been saved.
   */
  async saveContext(
    inputValues: InputValues,
    outputValues: OutputValues
  ): Promise<void> {
    await super.saveContext(inputValues, outputValues);
    const messages = await this.chatHistory.getMessages();
    this.buffer = await this.predictNewSummary(messages.slice(-2), this.buffer);
  }

  /**
   * Clears the conversation memory.
   * @returns A promise that resolves when the memory has been cleared.
   */
  async clear() {
    await super.clear();
    this.buffer = "";
  }
}


import { getBufferString } from "@langchain/core/messages";
import {
  InputValues,
  MemoryVariables,
  OutputValues,
} from "@langchain/core/memory";
import {
  BaseConversationSummaryMemory,
  BaseConversationSummaryMemoryInput,
} from "./summary.js";

/**
 * Interface for the input parameters of the
 * ConversationSummaryBufferMemory class.
 */
export interface ConversationSummaryBufferMemoryInput
  extends BaseConversationSummaryMemoryInput {
  maxTokenLimit?: number;
}

/**
 * Class that extends BaseConversationSummaryMemory and implements
 * ConversationSummaryBufferMemoryInput. It manages the conversation
 * history in a LangChain application by maintaining a buffer of chat
 * messages and providing methods to load, save, prune, and clear the
 * memory.
 * @example
 * ```typescript
 * // Initialize the memory with a specific model and token limit
 * const memory = new ConversationSummaryBufferMemory({
 *   llm: new ChatOpenAI({ modelName: "gpt-3.5-turbo-instruct", temperature: 0 }),
 *   maxTokenLimit: 10,
 * });
 *
 * // Save conversation context to memory
 * await memory.saveContext({ input: "hi" }, { output: "whats up" });
 * await memory.saveContext({ input: "not much you" }, { output: "not much" });
 *
 * // Load the conversation history from memory
 * const history = await memory.loadMemoryVariables({});
 * console.log({ history });
 *
 * // Create a chat prompt using the conversation history
 * const chatPrompt = ChatPromptTemplate.fromMessages([
 *   SystemMessagePromptTemplate.fromTemplate(
 *     "The following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.",
 *   ),
 *   new MessagesPlaceholder("history"),
 *   HumanMessagePromptTemplate.fromTemplate("{input}"),
 * ]);
 *
 * // Initialize the conversation chain with the model, memory, and prompt
 * const chain = new ConversationChain({
 *   llm: new ChatOpenAI({ temperature: 0.9, verbose: true }),
 *   memory: memory,
 *   prompt: chatPrompt,
 * });
 * ```
 */
export class ConversationSummaryBufferMemory
  extends BaseConversationSummaryMemory
  implements ConversationSummaryBufferMemoryInput
{
  movingSummaryBuffer = "";

  maxTokenLimit = 2000;

  constructor(fields: ConversationSummaryBufferMemoryInput) {
    super(fields);

    this.maxTokenLimit = fields?.maxTokenLimit ?? this.maxTokenLimit;
  }

  get memoryKeys() {
    return [this.memoryKey];
  }

  /**
   * Method that loads the chat messages from the memory and returns them as
   * a string or as a list of messages, depending on the returnMessages
   * property.
   * @param _ InputValues object, not used in this method.
   * @returns Promise that resolves with MemoryVariables object containing the loaded chat messages.
   */
  async loadMemoryVariables(_?: InputValues): Promise<MemoryVariables> {
    let buffer = await this.chatHistory.getMessages();
    if (this.movingSummaryBuffer) {
      buffer = [
        new this.summaryChatMessageClass(this.movingSummaryBuffer),
        ...buffer,
      ];
    }

    let finalBuffer;
    if (this.returnMessages) {
      finalBuffer = buffer;
    } else {
      finalBuffer = getBufferString(buffer, this.humanPrefix, this.aiPrefix);
    }

    return { [this.memoryKey]: finalBuffer };
  }

  /**
   * Method that saves the context of the conversation, including the input
   * and output values, and prunes the memory if it exceeds the maximum
   * token limit.
   * @param inputValues InputValues object containing the input values of the conversation.
   * @param outputValues OutputValues object containing the output values of the conversation.
   * @returns Promise that resolves when the context is saved and the memory is pruned.
   */
  async saveContext(
    inputValues: InputValues,
    outputValues: OutputValues
  ): Promise<void> {
    await super.saveContext(inputValues, outputValues);
    await this.prune();
  }

  /**
   * Method that prunes the memory if the total number of tokens in the
   * buffer exceeds the maxTokenLimit. It removes messages from the
   * beginning of the buffer until the total number of tokens is within the
   * limit.
   * @returns Promise that resolves when the memory is pruned.
   */
  async prune() {
    // Prune buffer if it exceeds max token limit
    let buffer = await this.chatHistory.getMessages();
    if (this.movingSummaryBuffer) {
      buffer = [
        new this.summaryChatMessageClass(this.movingSummaryBuffer),
        ...buffer,
      ];
    }

    let currBufferLength = await this.llm.getNumTokens(
      getBufferString(buffer, this.humanPrefix, this.aiPrefix)
    );

    if (currBufferLength > this.maxTokenLimit) {
      const prunedMemory = [];
      while (currBufferLength > this.maxTokenLimit) {
        const poppedMessage = buffer.shift();
        if (poppedMessage) {
          prunedMemory.push(poppedMessage);
          currBufferLength = await this.llm.getNumTokens(
            getBufferString(buffer, this.humanPrefix, this.aiPrefix)
          );
        }
      }
      this.movingSummaryBuffer = await this.predictNewSummary(
        prunedMemory,
        this.movingSummaryBuffer
      );
    }
  }

  /**
   * Method that clears the memory and resets the movingSummaryBuffer.
   * @returns Promise that resolves when the memory is cleared.
   */
  async clear() {
    await super.clear();
    this.movingSummaryBuffer = "";
  }
}


import { InputValues, MemoryVariables } from "@langchain/core/memory";
import { getBufferString } from "@langchain/core/messages";
import { BaseChatMemory, BaseChatMemoryInput } from "./chat_memory.js";

/**
 * Interface for the input parameters of the `BufferMemory` class.
 */
export interface BufferMemoryInput extends BaseChatMemoryInput {
  humanPrefix?: string;
  aiPrefix?: string;
  memoryKey?: string;
}

/**
 * The `BufferMemory` class is a type of memory component used for storing
 * and managing previous chat messages. It is a wrapper around
 * `ChatMessageHistory` that extracts the messages into an input variable.
 * This class is particularly useful in applications like chatbots where
 * it is essential to remember previous interactions. Note: The memory
 * instance represents the history of a single conversation. Therefore, it
 * is not recommended to share the same history or memory instance between
 * two different chains. If you deploy your LangChain app on a serverless
 * environment, do not store memory instances in a variable, as your
 * hosting provider may reset it by the next time the function is called.
 * @example
 * ```typescript
 * // Initialize the memory to store chat history and set up the language model with a specific temperature.
 * const memory = new BufferMemory({ memoryKey: "chat_history" });
 * const model = new ChatOpenAI({ temperature: 0.9 });
 *
 * // Create a prompt template for a friendly conversation between a human and an AI.
 * const prompt =
 *   PromptTemplate.fromTemplate(`The following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.
 *
 * Current conversation:
 * {chat_history}
 * Human: {input}
 * AI:`);
 *
 * // Set up the chain with the language model, prompt, and memory.
 * const chain = new LLMChain({ llm: model, prompt, memory });
 *
 * // Example usage of the chain to continue the conversation.
 * // The `call` method sends the input to the model and returns the AI's response.
 * const res = await chain.call({ input: "Hi! I'm Jim." });
 * console.log({ res });
 *
 * ```
 */
export class BufferMemory extends BaseChatMemory implements BufferMemoryInput {
  humanPrefix = "Human";

  aiPrefix = "AI";

  memoryKey = "history";

  constructor(fields?: BufferMemoryInput) {
    super({
      chatHistory: fields?.chatHistory,
      returnMessages: fields?.returnMessages ?? false,
      inputKey: fields?.inputKey,
      outputKey: fields?.outputKey,
    });
    this.humanPrefix = fields?.humanPrefix ?? this.humanPrefix;
    this.aiPrefix = fields?.aiPrefix ?? this.aiPrefix;
    this.memoryKey = fields?.memoryKey ?? this.memoryKey;
  }

  get memoryKeys() {
    return [this.memoryKey];
  }

  /**
   * Loads the memory variables. It takes an `InputValues` object as a
   * parameter and returns a `Promise` that resolves with a
   * `MemoryVariables` object.
   * @param _values `InputValues` object.
   * @returns A `Promise` that resolves with a `MemoryVariables` object.
   */
  async loadMemoryVariables(_values: InputValues): Promise<MemoryVariables> {
    const messages = await this.chatHistory.getMessages();
    if (this.returnMessages) {
      const result = {
        [this.memoryKey]: messages,
      };
      return result;
    }
    const result = {
      [this.memoryKey]: getBufferString(
        messages,
        this.humanPrefix,
        this.aiPrefix
      ),
    };
    return result;
  }
}


import { test, expect } from "@jest/globals";
import { ChatOpenAI, OpenAIChat } from "@langchain/openai";
import { SystemMessage } from "@langchain/core/messages";
import { ConversationSummaryMemory } from "../summary.js";

test("Test summary memory", async () => {
  const memory = new ConversationSummaryMemory({
    llm: new OpenAIChat({ modelName: "gpt-3.5-turbo", temperature: 0 }),
  });
  expect(await memory.loadMemoryVariables({})).toEqual({
    history: "",
  });

  await memory.saveContext(
    { input: "How's it going?" },
    { response: "Hello! I'm doing fine. and you?" }
  );
  // @eslint-disable-next-line/@typescript-eslint/ban-ts-comment
  // @ts-expect-error unused var
  const result2 = await memory.loadMemoryVariables({});
  // console.log("result2", result2);

  await memory.clear();
  expect(await memory.loadMemoryVariables({})).toEqual({
    history: "",
  });
});

test("Test summary memory with chat model", async () => {
  const memory = new ConversationSummaryMemory({
    llm: new ChatOpenAI({ temperature: 0 }),
  });
  expect(await memory.loadMemoryVariables({})).toEqual({
    history: "",
  });

  await memory.saveContext(
    { input: "How's it going?" },
    { response: "Hello! I'm doing fine. and you?" }
  );
  // @eslint-disable-next-line/@typescript-eslint/ban-ts-comment
  // @ts-expect-error unused var
  const result2 = await memory.loadMemoryVariables({});
  // console.log("result2", result2);

  await memory.clear();
  expect(await memory.loadMemoryVariables({})).toEqual({
    history: "",
  });
});

test("Test summary memory return messages", async () => {
  const memory = new ConversationSummaryMemory({
    llm: new OpenAIChat({ modelName: "gpt-3.5-turbo", temperature: 0 }),
    returnMessages: true,
  });
  expect(await memory.loadMemoryVariables({})).toEqual({
    history: [new SystemMessage("")],
  });

  await memory.saveContext(
    { input: "How's it going?" },
    { response: "Hello! I'm doing fine. and you?" }
  );
  // @eslint-disable-next-line/@typescript-eslint/ban-ts-comment
  // @ts-expect-error unused var
  const result2 = await memory.loadMemoryVariables({});
  // console.log("result2", result2);

  await memory.clear();
  expect(await memory.loadMemoryVariables({})).toEqual({
    history: [new SystemMessage("")],
  });
});


import { test, expect } from "@jest/globals";
import { OpenAI, ChatOpenAI } from "@langchain/openai";
import { SystemMessage } from "@langchain/core/messages";
import { ConversationSummaryBufferMemory } from "../summary_buffer.js";

test("Test summary buffer memory", async () => {
  const memory = new ConversationSummaryBufferMemory({
    llm: new OpenAI({ modelName: "gpt-3.5-turbo-instruct", temperature: 0 }),
    maxTokenLimit: 10,
  });
  expect(await memory.loadMemoryVariables({})).toEqual({
    history: "",
  });

  await memory.saveContext(
    { input: "How's it going?" },
    { response: "Hello! I'm doing fine. and you?" }
  );

  // @eslint-disable-next-line/@typescript-eslint/ban-ts-comment
  // @ts-expect-error unused var
  const result = await memory.loadMemoryVariables({});
  // console.log("result", result);

  await memory.clear();
  expect(await memory.loadMemoryVariables({})).toEqual({
    history: "",
  });
});

test("Test summary buffer memory with chat model", async () => {
  const memory = new ConversationSummaryBufferMemory({
    llm: new ChatOpenAI({ temperature: 0 }),
    maxTokenLimit: 10,
  });
  expect(await memory.loadMemoryVariables({})).toEqual({
    history: "",
  });

  await memory.saveContext(
    { input: "How's it going?" },
    { response: "Hello! I'm doing fine. and you?" }
  );
  // @eslint-disable-next-line/@typescript-eslint/ban-ts-comment
  // @ts-expect-error unused var
  const result = await memory.loadMemoryVariables({});
  // console.log("result", result);

  await memory.clear();
  expect(await memory.loadMemoryVariables({})).toEqual({
    history: "",
  });
});

test("Test summary buffer memory return messages", async () => {
  const memory = new ConversationSummaryBufferMemory({
    llm: new OpenAI({ modelName: "gpt-3.5-turbo-instruct", temperature: 0 }),
    returnMessages: true,
    maxTokenLimit: 10,
  });
  const exampleBuffer = "hello summary buffer";
  memory.movingSummaryBuffer = exampleBuffer;
  expect(await memory.loadMemoryVariables({})).toEqual({
    history: [new SystemMessage(exampleBuffer)],
  });

  await memory.saveContext(
    { input: "How's it going?" },
    { response: "Hello! I'm doing fine. and you?" }
  );

  // @eslint-disable-next-line/@typescript-eslint/ban-ts-comment
  // @ts-expect-error unused var
  const result = await memory.loadMemoryVariables({});
  // console.log("result", result);

  await memory.clear();
  expect(await memory.loadMemoryVariables({})).toEqual({
    history: [],
  });
});


import { Document } from "@langchain/core/documents";

/**
 * Given a list of documents, this util formats their contents
 * into a string, separated by newlines.
 *
 * @param documents
 * @returns A string of the documents page content, separated by newlines.
 */
export const formatDocumentsAsString = (documents: Document[]): string =>
  documents.map((doc) => doc.pageContent).join("\n\n");


import { describe, expect, test } from "@jest/globals";
import { Document } from "@langchain/core/documents";
import {
  CharacterTextSplitter,
  LatexTextSplitter,
  MarkdownTextSplitter,
  RecursiveCharacterTextSplitter,
  TokenTextSplitter,
} from "../text_splitter.js";

function textLineGenerator(char: string, length: number) {
  const line = new Array(length).join(char);
  return `${line}\n`;
}

describe("Character text splitter", () => {
  test("Test splitting by character count.", async () => {
    const text = "foo bar baz 123";
    const splitter = new CharacterTextSplitter({
      separator: " ",
      chunkSize: 7,
      chunkOverlap: 3,
    });
    const output = await splitter.splitText(text);
    const expectedOutput = ["foo bar", "bar baz", "baz 123"];
    expect(output).toEqual(expectedOutput);
  });

  test("Test splitting by character count doesn't create empty documents.", async () => {
    const text = "foo  bar";
    const splitter = new CharacterTextSplitter({
      separator: " ",
      chunkSize: 2,
      chunkOverlap: 0,
    });
    const output = await splitter.splitText(text);
    const expectedOutput = ["foo", "bar"];
    expect(output).toEqual(expectedOutput);
  });

  test("Test splitting by character count on long words.", async () => {
    const text = "foo bar baz a a";
    const splitter = new CharacterTextSplitter({
      separator: " ",
      chunkSize: 3,
      chunkOverlap: 1,
    });
    const output = await splitter.splitText(text);
    const expectedOutput = ["foo", "bar", "baz", "a a"];
    expect(output).toEqual(expectedOutput);
  });

  test("Test splitting by character count when shorter words are first.", async () => {
    const text = "a a foo bar baz";
    const splitter = new CharacterTextSplitter({
      separator: " ",
      chunkSize: 3,
      chunkOverlap: 1,
    });
    const output = await splitter.splitText(text);
    const expectedOutput = ["a a", "foo", "bar", "baz"];
    expect(output).toEqual(expectedOutput);
  });

  test("Test splitting by characters when splits not found easily.", async () => {
    const text = "foo bar baz 123";
    const splitter = new CharacterTextSplitter({
      separator: " ",
      chunkSize: 1,
      chunkOverlap: 0,
    });
    const output = await splitter.splitText(text);
    const expectedOutput = ["foo", "bar", "baz", "123"];
    expect(output).toEqual(expectedOutput);
  });

  test("Test invalid arguments.", () => {
    expect(() => {
      // @eslint-disable-next-line/@typescript-eslint/ban-ts-comment
      // @ts-expect-error unused var
      const res = new CharacterTextSplitter({ chunkSize: 2, chunkOverlap: 4 });
      // console.log(res);
    }).toThrow();
  });

  test("Test create documents method.", async () => {
    const texts = ["foo bar", "baz"];
    const splitter = new CharacterTextSplitter({
      separator: " ",
      chunkSize: 3,
      chunkOverlap: 0,
    });
    const docs = await splitter.createDocuments(texts);
    const metadata = { loc: { lines: { from: 1, to: 1 } } };
    const expectedDocs = [
      new Document({ pageContent: "foo", metadata }),
      new Document({ pageContent: "bar", metadata }),
      new Document({ pageContent: "baz", metadata }),
    ];
    expect(docs).toEqual(expectedDocs);
  });

  test("Test create documents with metadata method.", async () => {
    const texts = ["foo bar", "baz"];
    const splitter = new CharacterTextSplitter({
      separator: " ",
      chunkSize: 3,
      chunkOverlap: 0,
    });
    const docs = await splitter.createDocuments(texts, [
      { source: "1" },
      { source: "2" },
    ]);
    const loc = { lines: { from: 1, to: 1 } };
    const expectedDocs = [
      new Document({ pageContent: "foo", metadata: { source: "1", loc } }),
      new Document({
        pageContent: "bar",
        metadata: { source: "1", loc },
      }),
      new Document({ pageContent: "baz", metadata: { source: "2", loc } }),
    ];
    expect(docs).toEqual(expectedDocs);
  });

  test("Test create documents method with metadata and an added chunk header.", async () => {
    const texts = ["foo bar", "baz"];
    const splitter = new CharacterTextSplitter({
      separator: " ",
      chunkSize: 3,
      chunkOverlap: 0,
    });
    const docs = await splitter.createDocuments(
      texts,
      [{ source: "1" }, { source: "2" }],
      {
        chunkHeader: `SOURCE NAME: testing\n-----\n`,
        appendChunkOverlapHeader: true,
      }
    );
    const loc = { lines: { from: 1, to: 1 } };
    const expectedDocs = [
      new Document({
        pageContent: "SOURCE NAME: testing\n-----\nfoo",
        metadata: { source: "1", loc },
      }),
      new Document({
        pageContent: "SOURCE NAME: testing\n-----\n(cont'd) bar",
        metadata: { source: "1", loc },
      }),
      new Document({
        pageContent: "SOURCE NAME: testing\n-----\nbaz",
        metadata: { source: "2", loc },
      }),
    ];
    expect(docs).toEqual(expectedDocs);
  });
});


describe("RecursiveCharacter text splitter", () => {
  test("One unique chunk", async () => {
    const splitter = new RecursiveCharacterTextSplitter({
      chunkSize: 100,
      chunkOverlap: 0,
    });
    const content = textLineGenerator("A", 70);

    const docs = await splitter.createDocuments([content]);

    const expectedDocs = [
      new Document({
        pageContent: content.trim(),
        metadata: { loc: { lines: { from: 1, to: 1 } } },
      }),
    ];

    expect(docs).toEqual(expectedDocs);
  });

  test("Test iterative text splitter.", async () => {
    const text = `Hi.\n\nI'm Harrison.\n\nHow? Are? You?\nOkay then f f f f.
This is a weird text to write, but gotta test the splittingggg some how.\n\n
Bye!\n\n-H.`;
    const splitter = new RecursiveCharacterTextSplitter({
      chunkSize: 10,
      chunkOverlap: 1,
    });
    const output = await splitter.splitText(text);
    const expectedOutput = [
      "Hi.",
      "I'm",
      "Harrison.",
      "How? Are?",
      "You?",
      "Okay then",
      "f f f f.",
      "This is a",
      "weird",
      "text to",
      "write,",
      "but gotta",
      "test the",
      "splitting",
      "gggg",
      "some how.",
      "Bye!",
      "-H.",
    ];
    expect(output).toEqual(expectedOutput);
  });

  test("A basic chunked document", async () => {
    const splitter = new RecursiveCharacterTextSplitter({
      chunkSize: 100,
      chunkOverlap: 0,
    });
    const line1 = textLineGenerator("A", 70);
    const line2 = textLineGenerator("B", 70);
    const content = line1 + line2;

    const docs = await splitter.createDocuments([content]);

    const expectedDocs = [
      new Document({
        pageContent: line1.trim(),
        metadata: { loc: { lines: { from: 1, to: 1 } } },
      }),
      new Document({
        pageContent: line2.trim(),
        metadata: { loc: { lines: { from: 2, to: 2 } } },
      }),
    ];

    expect(docs).toEqual(expectedDocs);
  });

  test("A chunked document with similar text", async () => {
    const splitter = new RecursiveCharacterTextSplitter({
      chunkSize: 100,
      chunkOverlap: 0,
    });
    const line = textLineGenerator("A", 70);
    const content = line + line;

    const docs = await splitter.createDocuments([content]);

    const expectedDocs = [
      new Document({
        pageContent: line.trim(),
        metadata: { loc: { lines: { from: 1, to: 1 } } },
      }),
      new Document({
        pageContent: line.trim(),
        metadata: { loc: { lines: { from: 2, to: 2 } } },
      }),
    ];

    expect(docs).toEqual(expectedDocs);
  });

  test("A chunked document starting with new lines", async () => {
    const splitter = new RecursiveCharacterTextSplitter({
      chunkSize: 100,
      chunkOverlap: 0,
    });
    const line1 = textLineGenerator("\n", 2);
    const line2 = textLineGenerator("A", 70);
    const line3 = textLineGenerator("\n", 4);
    const line4 = textLineGenerator("B", 70);
    const line5 = textLineGenerator("\n", 4);
    const content = line1 + line2 + line3 + line4 + line5;

    const docs = await splitter.createDocuments([content]);

    const expectedDocs = [
      new Document({
        pageContent: line2.trim(),
        metadata: { loc: { lines: { from: 3, to: 3 } } },
      }),
      new Document({
        pageContent: line4.trim(),
        metadata: { loc: { lines: { from: 8, to: 8 } } },
      }),
    ];

    expect(docs).toEqual(expectedDocs);
  });

  test("A chunked with overlap", async () => {
    const splitter = new RecursiveCharacterTextSplitter({
      chunkSize: 100,
      chunkOverlap: 30,
    });
    const line1 = textLineGenerator("A", 70);
    const line2 = textLineGenerator("B", 20);
    const line3 = textLineGenerator("C", 70);
    const content = line1 + line2 + line3;

    const docs = await splitter.createDocuments([content]);

    const expectedDocs = [
      new Document({
        pageContent: line1 + line2.trim(),
        metadata: { loc: { lines: { from: 1, to: 2 } } },
      }),
      new Document({
        pageContent: line2 + line3.trim(),
        metadata: { loc: { lines: { from: 2, to: 3 } } },
      }),
    ];

    expect(docs).toEqual(expectedDocs);
  });

  test("Chunks with overlap that contains new lines", async () => {
    const splitter = new RecursiveCharacterTextSplitter({
      chunkSize: 100,
      chunkOverlap: 30,
    });
    const line1 = textLineGenerator("A", 70);
    const line2 = textLineGenerator("B", 10);
    const line3 = textLineGenerator("C", 10);
    const line4 = textLineGenerator("D", 70);
    const content = line1 + line2 + line3 + line4;

    const docs = await splitter.createDocuments([content]);

    const expectedDocs = [
      new Document({
        pageContent: line1 + line2 + line3.trim(),
        metadata: { loc: { lines: { from: 1, to: 3 } } },
      }),
      new Document({
        pageContent: line2 + line3 + line4.trim(),
        metadata: { loc: { lines: { from: 2, to: 4 } } },
      }),
    ];
    expect(docs).toEqual(expectedDocs);
  });
});

test("Separator length is considered correctly for chunk size", async () => {
  const text = "aa ab ac ba bb";
  const splitter = new RecursiveCharacterTextSplitter({
    keepSeparator: false,
    chunkSize: 7,
    chunkOverlap: 3,
  });
  const output = await splitter.splitText(text);
  const expectedOutput = ["aa ab", "ab ac", "ac ba", "ba bb"];

  expect(output).toEqual(expectedOutput);
});

test("Token text splitter", async () => {
  const text = "foo bar baz a a";
  const splitter = new TokenTextSplitter({
    encodingName: "r50k_base",
    chunkSize: 3,
    chunkOverlap: 0,
  });
  const output = await splitter.splitText(text);
  const expectedOutput = ["foo bar b", "az a a"];

  expect(output).toEqual(expectedOutput);
});

test("Token text splitter overlap when last chunk is large", async () => {
  const text = "foo bar baz a a";
  const splitter = new TokenTextSplitter({
    encodingName: "r50k_base",
    chunkSize: 5,
    chunkOverlap: 3,
  });
  const output = await splitter.splitText(text);
  const expectedOutput = ["foo bar baz a", " baz a a"];

  expect(output).toEqual(expectedOutput);
});

test("Test markdown text splitter", async () => {
  const text =
    "# 🦜🔗 LangChain\n" +
    "\n" +
    "⚡ Building applications with LLMs through composability ⚡\n" +
    "\n" +
    "## Quick Install\n" +
    "\n" +
    "```bash\n" +
    "# Hopefully this code block isn't split\n" +
    "pip install langchain\n" +
    "```\n" +
    "\n" +
    "As an open source project in a rapidly developing field, we are extremely open to contributions.";
  const splitter = new MarkdownTextSplitter({
    chunkSize: 100,
    chunkOverlap: 0,
  });
  const output = await splitter.splitText(text);

  const expectedOutput = [
    "# 🦜🔗 LangChain\n\n⚡ Building applications with LLMs through composability ⚡",
    "## Quick Install\n\n```bash\n# Hopefully this code block isn't split\npip install langchain",
    "```",
    "As an open source project in a rapidly developing field, we are extremely open to contributions.",
  ];
  expect(output).toEqual(expectedOutput);
});

test("Test latex text


litter.", async () => {
  const text = `\\begin{document}
\\title{🦜🔗 LangChain}
⚡ Building applications with LLMs through composability ⚡

\\section{Quick Install}

\\begin{verbatim}
Hopefully this code block isn't split
yarn add langchain
\\end{verbatim}

As an open source project in a rapidly developing field, we are extremely open to contributions.

\\end{document}`;
  const splitter = new LatexTextSplitter({
    chunkSize: 100,
    chunkOverlap: 0,
  });
  const output = await splitter.splitText(text);

  const expectedOutput = [
    "\\begin{document}\n\\title{🦜🔗 LangChain}\n⚡ Building applications with LLMs through composability ⚡",
    "\\section{Quick Install}",
    "\\begin{verbatim}\nHopefully this code block isn't split\nyarn add langchain\n\\end{verbatim}",
    "As an open source project in a rapidly developing field, we are extremely open to contributions.",
    "\\end{document}",
  ];
  expect(output).toEqual(expectedOutput);
});

test("Test HTML text splitter", async () => {
  const text = `<!DOCTYPE html>

    <title>🦜🔗 LangChain</title>

🦜🔗 LangChain

⚡ Building applications with LLMs through composability ⚡

      As an open source project in a rapidly developing field, we are extremely open to contributions.

`;
  const splitter = RecursiveCharacterTextSplitter.fromLanguage("html", {
    chunkSize: 175,
    chunkOverlap: 20,
  });
  const output = await splitter.splitText(text);

  const expectedOutput = [
    "<!DOCTYPE html>\n",
    "\n    <title>🦜🔗 LangChain</title>",
    `
  `,
    `

🦜🔗 LangChain

⚡ Building applications with LLMs through composability ⚡

`,
    `

      As an open source project in a rapidly developing field, we are extremely open to contributions.

`,
  ];
  expect(output).toEqual(expectedOutput);
});

test("Test lines loc on iterative text splitter.", async () => {
  const text = `Hi.\nI'm Harrison.\n\nHow?\na\nb`;
  const splitter = new RecursiveCharacterTextSplitter({
    chunkSize: 20,
    chunkOverlap: 1,
  });
  const docs = await splitter.createDocuments([text]);

  const expectedDocs = [
    new Document({
      pageContent: "Hi.\nI'm Harrison.",
      metadata: { loc: { lines: { from: 1, to: 2 } } },
    }),
    new Document({
      pageContent: "How?\na\nb",
      metadata: { loc: { lines: { from: 4, to: 6 } } },
    }),
  ];

  expect(docs).toEqual(expectedDocs);
});


export {
  mapStoredMessagesToChatMessages,
  mapChatMessagesToStoredMessages,
} from "@langchain/core/messages";


import type { BaseLanguageModelInterface } from "@langchain/core/language_models/base";
import type {
  StructuredToolInterface,
  ToolInterface,
} from "@langchain/core/tools";
import { CallbackManager } from "@langchain/core/callbacks/manager";
import { BufferMemory } from "../memory/buffer_memory.js";
import { ChatAgent } from "./chat/index.js";
import { ChatConversationalAgent } from "./chat_convo/index.js";
import { StructuredChatAgent } from "./structured_chat/index.js";
import { AgentExecutor, AgentExecutorInput } from "./executor.js";
import { ZeroShotAgent } from "./mrkl/index.js";
import { OpenAIAgent } from "./openai_functions/index.js";
import { XMLAgent } from "./xml/index.js";

/**
 * Represents the type of an agent in LangChain. It can be
 * "zero-shot-react-description", "chat-zero-shot-react-description", or
 * "chat-conversational-react-description".
 */
type AgentType =
  | "zero-shot-react-description"
  | "chat-zero-shot-react-description"
  | "chat-conversational-react-description";

/**
 * @deprecated See {@link https://js.langchain.com/docs/modules/agents/agent_types/ | new agent creation docs}.
 */
export const initializeAgentExecutor = async (
  tools: ToolInterface[],
  llm: BaseLanguageModelInterface,
  _agentType?: AgentType,
  _verbose?: boolean,
  _callbackManager?: CallbackManager
): Promise<AgentExecutor> => {
  const agentType = _agentType ?? "zero-shot-react-description";
  const verbose = _verbose;
  const callbackManager = _callbackManager;
  switch (agentType) {
    case "zero-shot-react-description":
      return AgentExecutor.fromAgentAndTools({
        agent: ZeroShotAgent.fromLLMAndTools(llm, tools),
        tools,
        returnIntermediateSteps: true,
        verbose,
        callbackManager,
      });
    case "chat-zero-shot-react-description":
      return AgentExecutor.fromAgentAndTools({
        agent: ChatAgent.fromLLMAndTools(llm, tools),
        tools,
        returnIntermediateSteps: true,
        verbose,
        callbackManager,
      });
    case "chat-conversational-react-description":
      return AgentExecutor.fromAgentAndTools({
        agent: ChatConversationalAgent.fromLLMAndTools(llm, tools),
        tools,
        verbose,
        callbackManager,
      });
    default:
      throw new Error("Unknown agent type");
  }
};

/**
 * @interface
 */
export type InitializeAgentExecutorOptions =
  | ({
      agentType: "zero-shot-react-description";
      agentArgs?: Parameters<typeof ZeroShotAgent.fromLLMAndTools>[2];
      memory?: never;
    } & Omit<AgentExecutorInput, "agent" | "tools">)
  | ({
      agentType: "chat-zero-shot-react-description";
      agentArgs?: Parameters<typeof ChatAgent.fromLLMAndTools>[2];
      memory?: never;
    } & Omit<AgentExecutorInput, "agent" | "tools">)
  | ({
      agentType: "chat-conversational-react-description";
      agentArgs?: Parameters<typeof ChatConversationalAgent.fromLLMAndTools>[2];
    } & Omit<AgentExecutorInput, "agent" | "tools">)
  | ({
      agentType: "xml";
      agentArgs?: Parameters<typeof XMLAgent.fromLLMAndTools>[2];
    } & Omit<AgentExecutorInput, "agent" | "tools">);

/**
 * @interface
 */
export type InitializeAgentExecutorOptionsStructured =
  | ({
      agentType: "structured-chat-zero-shot-react-description";
      agentArgs?: Parameters<typeof StructuredChatAgent.fromLLMAndTools>[2];
    } & Omit<AgentExecutorInput, "agent" | "tools">)
  | ({
      agentType: "openai-functions";
      agentArgs?: Parameters<typeof OpenAIAgent.fromLLMAndTools>[2];
    } & Omit<AgentExecutorInput, "agent" | "tools">);

/**
 * Initialize an agent executor with options.
 * @deprecated See {@link https://js.langchain.com/docs/modules/agents/agent_types/ | new agent creation docs}.
 * @param tools Array of tools to use in the agent
 * @param llm LLM or ChatModel to use in the agent
 * @param options Options for the agent, including agentType, agentArgs, and other options for AgentExecutor.fromAgentAndTools
 * @returns AgentExecutor
 */
export async function initializeAgentExecutorWithOptions(
  tools: StructuredToolInterface[],
  llm: BaseLanguageModelInterface,
  options: InitializeAgentExecutorOptionsStructured
): Promise<AgentExecutor>;
/** @deprecated See {@link https://js.langchain.com/docs/modules/agents/agent_types/ | new agent creation docs}. */
export async function initializeAgentExecutorWithOptions(
  tools: ToolInterface[],
  llm: BaseLanguageModelInterface,
  options?: InitializeAgentExecutorOptions
): Promise<AgentExecutor>;
/** @deprecated See {@link https://js.langchain.com/docs/modules/agents/agent_types/ | new agent creation docs}. */
export async function initializeAgentExecutorWithOptions(
  tools: StructuredToolInterface[] | ToolInterface[],
  llm: BaseLanguageModelInterface,
  options:
    | InitializeAgentExecutorOptions
    | InitializeAgentExecutorOptionsStructured = {
    agentType:
      llm._modelType() === "base_chat_model"
        ? "chat-zero-shot-react-description"
        : "zero-shot-react-description",
  }
): Promise<AgentExecutor> {
  // Note this tools cast is safe as the overload signatures prevent
  // the function from being called with a StructuredTool[] when
  // the agentType is not in InitializeAgentExecutorOptionsStructured
  switch (options.agentType) {
    case "zero-shot-react-description": {
      const { agentArgs, tags, ...rest } = options;
      return AgentExecutor.fromAgentAndTools({
        tags: [...(tags ?? []), "zero-shot-react-description"],
        agent: ZeroShotAgent.fromLLMAndTools(
          llm,
          tools as ToolInterface[],
          agentArgs
        ),
        tools,
        ...rest,
      });
    }
    case "chat-zero-shot-react-description": {
      const { agentArgs, tags, ...rest } = options;
      return AgentExecutor.fromAgentAndTools({
        tags: [...(tags ?? []), "chat-zero-shot-react-description"],
        agent: ChatAgent.fromLLMAndTools(
          llm,
          tools as ToolInterface[],
          agentArgs
        ),
        tools,
        ...rest,
      });
    }
    case "chat-conversational-react-description": {
      const { agentArgs, memory, tags, ...rest } = options;
      const executor = AgentExecutor.fromAgentAndTools({
        tags: [...(tags ?? []), "chat-conversational-react-description"],
        agent: ChatConversationalAgent.fromLLMAndTools(
          llm,
          tools as ToolInterface[],
          agentArgs
        ),
        tools,
        memory:
          memory ??
          new BufferMemory({
            returnMessages: true,
            memoryKey: "chat_history",
            inputKey: "input",
            outputKey: "output",
          }),
        ...rest,
      });
      return executor;
    }
    case "xml": {
      const { agentArgs, tags, ...rest } = options;
      const executor = AgentExecutor.fromAgentAndTools({
        tags: [...(tags ?? []), "xml"],
        agent: XMLAgent.fromLLMAndTools(
          llm,
          tools as ToolInterface[],
          agentArgs
        ),
        tools,
        ...rest,
      });
      return executor;
    }
    case "structured-chat-zero-shot-react-description": {
      const { agentArgs, memory, tags, ...rest } = options;
      const executor = AgentExecutor.fromAgentAndTools({
        tags: [...(tags ?? []), "structured-chat-zero-shot-react-description"],
        agent: StructuredChatAgent.fromLLMAndTools(llm, tools, agentArgs),
        tools,
        memory,
        ...rest,
      });
      return executor;
    }
    case "openai-functions": {
      const { agentArgs, memory, tags, ...rest } = options;
      const executor = AgentExecutor.fromAgentAndTools({
        tags: [...(tags ?? []), "openai-functions"],
        agent: OpenAIAgent.fromLLMAndTools(llm, tools, agentArgs),
        tools,
        memory:
          memory ??
          new BufferMemory({
            returnMessages: true,
            memoryKey: "chat_history",
            inputKey: "input",
            outputKey: "output",
          }),
        ...rest,
      });
      return executor;
    }
    default: {
      throw new Error("Unknown agent type");
    }
  }
}


import { OutputParserException } from "@langchain/core/output_parsers";
import { OutputParserArgs } from "../agent.js";
import { AgentActionOutputParser } from "../types.js";

import { FORMAT_INSTRUCTIONS } from "./prompt.js";

export const FINAL_ANSWER_ACTION = "Final Answer:";
/**
 * A class that extends `AgentActionOutputParser` to provide a custom
 * implementation for parsing the output of a ZeroShotAgent action.
 */
export class ZeroShotAgentOutputParser extends AgentActionOutputParser {
  lc_namespace = ["langchain", "agents", "mrkl"];

  finishToolName: string;

  constructor(fields?: OutputParserArgs) {
    super(fields);
    this.finishToolName = fields?.finishToolName || FINAL_ANSWER_ACTION;
  }

  /**
   * Parses the text output of an agent action, extracting the tool, tool
   * input, and output.
   * @param text The text output of an agent action.
   * @returns An object containing the tool, tool input, and output extracted from the text, along with the original text as a log.
   */
  async parse(text: string) {
    if (text.includes(this.finishToolName)) {
      const parts = text.split(this.finishToolName);
      const output = parts[parts.length - 1].trim();
      return {
        returnValues: { output },
        log: text,
      };
    }

    const match = /Action:([\s\S]*?)(?:\nAction Input:([\s\S]*?))?$/.exec(text);
    if (!match) {
      throw new OutputParserException(`Could not parse LLM output: ${text}`);
    }

    return {
      tool: match[1].trim(),
      toolInput: match[2]
        ? match[2].trim().replace(/^("+)(.*?)(\1)$/, "$2")
        : "",
      log: text,
    };
  }

  /**
   * Returns the format instructions for parsing the output of an agent
   * action in the style of the ZeroShotAgent.
   * @returns The format instructions for parsing the output.
   */
  getFormatInstructions(): string {
    return FORMAT_INSTRUCTIONS;
  }
}


export type CreateStructuredChatAgentParams = {
  /** LLM to use as the agent. */
  llm: BaseLanguageModelInterface;
  /** Tools this agent has access to. */
  tools: (StructuredToolInterface | ToolDefinition)[];
  /**
   * The prompt to use. Must have input keys for
   * `tools`, `tool_names`, and `agent_scratchpad`.
   */
  prompt: BasePromptTemplate;
  /**
   * Whether to invoke the underlying model in streaming mode,
   * allowing streaming of intermediate steps. Defaults to true.
   */
  streamRunnable?: boolean;
};

/**
 * Create an agent aimed at supporting tools with multiple inputs.
 * @param params Params required to create the agent. Includes an LLM, tools, and prompt.
 * @returns A runnable sequence representing an agent. It takes as input all the same input
 *     variables as the prompt passed in does. It returns as output either an
 *     AgentAction or AgentFinish.
 *
 * @example
 * ```typescript
 * import { AgentExecutor, createStructuredChatAgent } from "langchain/agents";
 * import { pull } from "langchain/hub";
 * import type { ChatPromptTemplate } from "@langchain/core/prompts";
 * import { AIMessage, HumanMessage } from "@langchain/core/messages";
 *
 * import { ChatOpenAI } from "@langchain/openai";
 *
 * // Define the tools the agent will have access to.
 * const tools = [...];
 *
 * // Get the prompt to use - you can modify this!
 * // If you want to see the prompt in full, you can at:
 * // https://smith.langchain.com/hub/hwchase17/structured-chat-agent
 * const prompt = await pull<ChatPromptTemplate>(
 *   "hwchase17/structured-chat-agent"
 * );
 *
 * const llm = new ChatOpenAI({
 *   temperature: 0,
 *   modelName: "gpt-3.5-turbo-1106",
 * });
 *
 * const agent = await createStructuredChatAgent({
 *   llm,
 *   tools,
 *   prompt,
 * });
 *
 * const agentExecutor = new AgentExecutor({
 *   agent,
 *   tools,
 * });
 *
 * const result = await agentExecutor.invoke({
 *   input: "what is LangChain?",
 * });
 *
 * // With chat history
 * const result2 = await agentExecutor.invoke({
 *   input: "what's my name?",
 *   chat_history: [
 *     new HumanMessage("hi! my name is cob"),
 *     new AIMessage("Hello Cob! How can I assist you today?"),
 *   ],
 * });
 * ```
 */
export async function createStructuredChatAgent({
  llm,
  tools,
  prompt,
  streamRunnable,
}: CreateStructuredChatAgentParams) {
  const missingVariables = ["tools", "tool_names", "agent_scratchpad"].filter(
    (v) => !prompt.inputVariables.includes(v)
  );
  if (missingVariables.length > 0) {
    throw new Error(
      `Provided prompt is missing required input variables: ${JSON.stringify(
        missingVariables
      )}`
    );
  }
  let toolNames: string[] = [];
  if (tools.every(isOpenAITool)) {
    toolNames = tools.map((tool) => tool.function.name);
  } else if (tools.every(isStructuredTool)) {
    toolNames = tools.map((tool) => tool.name);
  } else {
    throw new Error(
      "All tools must be either OpenAI or Structured tools, not a mix."
    );
  }
  const partialedPrompt = await prompt.partial({
    tools: renderTextDescriptionAndArgs(tools),
    tool_names: toolNames.join(", "),
  });
  // TODO: Add .bind to core runnable interface.
  const llmWithStop = (llm as BaseLanguageModel).bind({
    stop: ["Observation"],
  });
  const agent = AgentRunnableSequence.fromRunnables(
    [
      RunnablePassthrough.assign({
        agent_scratchpad: (input: { steps: AgentStep[] }) =>
          formatLogToString(input.steps),
      }),
      partialedPrompt,
      llmWithStop,
      StructuredChatOutputParserWithRetries.fromLLM(llm, {
        toolNames,
      }),
    ],
    {
      name: "StructuredChatAgent",
      streamRunnable,
      singleAction: true,
    }
  );
  return agent;
}


import type { BaseLanguageModelInterface } from "@langchain/core/language_models/base";
import { Callbacks } from "@langchain/core/callbacks/manager";
import { AgentAction, AgentFinish } from "@langchain/core/agents";
import { OutputParserException } from "@langchain/core/output_parsers";
import { renderTemplate } from "@langchain/core/prompts";
import { AgentActionOutputParser } from "../types.js";
import {
  AGENT_ACTION_FORMAT_INSTRUCTIONS,
  FORMAT_INSTRUCTIONS,
} from "./prompt.js";
import { OutputFixingParser } from "../../output_parsers/fix.js";

/**
 * A class that provides a custom implementation for parsing the output of
 * a StructuredChatAgent action. It extends the `AgentActionOutputParser`
 * class and extracts the action and action input from the text output,
 * returning an `AgentAction` or `AgentFinish` object.
 */
export class StructuredChatOutputParser extends AgentActionOutputParser {
  lc_namespace = ["langchain", "agents", "structured_chat"];

  private toolNames: string[];

  constructor(fields: { toolNames: string[] }) {
    super(...arguments);
    this.toolNames = fields.toolNames;
  }

  /**
   * Parses the given text and returns an `AgentAction` or `AgentFinish`
   * object. If an `OutputFixingParser` is provided, it is used for parsing;
   * otherwise, the base parser is used.
   * @param text The text to parse.
   * @param callbacks Optional callbacks for asynchronous operations.
   * @returns A Promise that resolves to an `AgentAction` or `AgentFinish` object.
   */
  async parse(text: string): Promise<AgentAction | AgentFinish> {
    try {
      const regex = /```(?:json)?(.*)(```)/gs;
      const actionMatch = regex.exec(text);
      if (actionMatch === null) {
        throw new OutputParserException(
          `Could not parse an action. The agent action must be within a markdown code block, and "action" must be a provided tool or "Final Answer"`
        );
      }
      const response = JSON.parse(actionMatch[1].trim());
      const { action, action_input } = response;

      if (action === "Final Answer") {
        return { returnValues: { output: action_input }, log: text };
      }
      return { tool: action, toolInput: action_input || {}, log: text };
    } catch (e) {
      throw new OutputParserException(
        `Failed to parse. Text: "${text}". Error: ${e}`
      );
    }
  }

  /**
   * Returns the format instructions for parsing the output of an agent
   * action in the style of the StructuredChatAgent.
   * @returns A string representing the format instructions.
   */
  getFormatInstructions(): string {
    return renderTemplate(AGENT_ACTION_FORMAT_INSTRUCTIONS, "f-string", {
      tool_names: this.toolNames.join(", "),
    });
  }
}

/**
 * An interface for the arguments used to construct a
 * `StructuredChatOutputParserWithRetries` instance.
 */
export interface StructuredChatOutputParserArgs {
  baseParser?: StructuredChatOutputParser;
  outputFixingParser?: OutputFixingParser<AgentAction | AgentFinish>;
  toolNames?: string[];
}

/**
 * A class that provides a wrapper around the `StructuredChatOutputParser`
 * and `OutputFixingParser` classes. It extends the
 * `AgentActionOutputParser` class and allows for retrying the output
 * parsing using the `OutputFixingParser` if it is provided.
 * @example
 * ```typescript
 * const outputParser = new StructuredChatOutputParserWithRetries.fromLLM(
 *   new ChatOpenAI({ temperature: 0 }),
 *   {
 *     toolNames: ["calculator", "random-number-generator"],
 *   },
 * );
 * const result = await outputParser.parse(
 *  "What is a random number between 5 and 10 raised to the second power?"
 * );
 * ```
 */
export class StructuredChatOutputParserWithRetries extends AgentActionOutputParser {
  lc_namespace = ["langchain", "agents", "structured_chat"];

  private baseParser: StructuredChatOutputParser;

  private outputFixingParser?: OutputFixingParser<AgentAction | AgentFinish>;

  private toolNames: string[] = [];

  constructor(fields: StructuredChatOutputParserArgs) {
    super(fields);
    this.toolNames = fields.toolNames ?? this.toolNames;
    this.baseParser =
      fields?.baseParser ??
      new StructuredChatOutputParser({ toolNames: this.toolNames });
    this.outputFixingParser = fields?.outputFixingParser;
  }

  /**
   * Parses the given text and returns an `AgentAction` or `AgentFinish`
   * object. Throws an `OutputParserException` if the parsing fails.
   * @param text The text to parse.
   * @returns A Promise that resolves to an `AgentAction` or `AgentFinish` object.
   */
  async parse(
    text: string,
    callbacks?: Callbacks
  ): Promise<AgentAction | AgentFinish> {
    if (this.outputFixingParser !== undefined) {
      return this.outputFixingParser.parse(text, callbacks);
    }
    return this.baseParser.parse(text);
  }

  /**
   * Returns the format instructions for parsing the output of an agent
   * action in the style of the StructuredChatAgent.
   * @returns A string representing the format instructions.
   */
  getFormatInstructions(): string {
    return renderTemplate(FORMAT_INSTRUCTIONS, "f-string", {
      tool_names: this.toolNames.join(", "),
    });
  }

  /**
   * Creates a new `StructuredChatOutputParserWithRetries` instance from a
   * `BaseLanguageModel` and options. The options can include a base parser
   * and tool names.
   * @param llm A `BaseLanguageModel` instance.
   * @param options Options for creating a `StructuredChatOutputParserWithRetries` instance.
   * @returns A new `StructuredChatOutputParserWithRetries` instance.
   */
  static fromLLM(
    llm: BaseLanguageModelInterface,
    options: Omit<StructuredChatOutputParserArgs, "outputFixingParser">
  ): StructuredChatOutputParserWithRetries {
    const baseParser =
      options.baseParser ??
      new StructuredChatOutputParser({ toolNames: options.toolNames ?? [] });
    const outputFixingParser = OutputFixingParser.fromLLM(llm, baseParser);
    return new StructuredChatOutputParserWithRetries({
      baseParser,
      outputFixingParser,
      toolNames: options.toolNames,
    });
  }
}


/* eslint-disable no-process-env */
import { expect, test } from "@jest/globals";
import { OpenAI, OpenAIEmbeddings, ChatOpenAI } from "@langchain/openai";
import { Tool } from "@langchain/core/tools";
import { RunnableSequence } from "@langchain/core/runnables";
import { OutputParserException } from "@langchain/core/output_parsers";
import { AIMessage } from "@langchain/core/messages";
import { AgentStep } from "@langchain/core/agents";
import { ChatMessageHistory } from "../../stores/message/in_memory.js";
import { AgentExecutor, ZeroShotAgent } from "../index.js";
import { SerpAPI } from "../../util/testing/tools/serpapi.js";
import { Calculator } from "../../util/testing/tools/calculator.js";
import { initializeAgentExecutorWithOptions } from "../initialize.js";
import { WebBrowser } from "../../tools/webbrowser.js";
import { BufferMemory } from "../../memory/buffer_memory.js";

test("Pass runnable to agent executor", async () => {
  const model = new ChatOpenAI({ temperature: 0, modelName: "gpt-3.5-turbo" });
  const tools: Tool[] = [
    new SerpAPI(undefined, {
      location: "Austin,Texas,United States",
      hl: "en",
      gl: "us",
    }),
    new Calculator(),
  ];

  const prompt = ZeroShotAgent.createPrompt(tools);
  const outputParser = ZeroShotAgent.getDefaultOutputParser();

  const runnable = RunnableSequence.from([
    {
      input: (i: { input: string }) => i.input,
      agent_scratchpad: (i: { input: string }) => i.input,
    },
    prompt,
    model,
    outputParser,
  ]);

  const executor = AgentExecutor.fromAgentAndTools({
    agent: runnable,
    tools,
  });
  const res = await executor.invoke({
    input:
      "Who is Olivia Wilde's boyfriend? What is his current age raised to the 0.23 power?",
  });
  // console.log(
  //   {
  //     res,
  //   },
  //   "Pass runnable to agent executor"
  // );
  expect(res.output).not.toEqual("");
  expect(res.output).not.toEqual("Agent stopped due to max iterations.");
});

test("Custom output parser", async () => {
  const model = new ChatOpenAI({ temperature: 0, modelName: "gpt-3.5-turbo" });
  const tools: Tool[] = [
    new SerpAPI(undefined, {
      location: "Austin,Texas,United States",
      hl: "en",
      gl: "us",
    }),
    new Calculator(),
  ];

  const parser = (output: AIMessage) => {
    const text = output.content;
    if (typeof text !== "string") {
      throw new Error("Cannot parse non-string output.");
    }
    if (text.includes("Final Answer:")) {
      return {
        returnValues: {
          output: "We did it!",
        },
        log: text,
      };
    }

    const match = /Action:([\s\S]*?)(?:\nAction Input:([\s\S]*?))?$/.exec(text);
    if (!match) {
      throw new OutputParserException(`Could not parse LLM output: ${text}`);
    }

    return {
      tool: match[1].trim(),
      toolInput: match[2]
        ? match[2].trim().replace(/^("+)(.*?)(\1)$/, "$2")
        : "",
      log: text,
    };
  };

  const prompt = ZeroShotAgent.createPrompt(tools);

  const runnable = RunnableSequence.from([
    {
      input: (i: { input: string }) => i.input,
      agent_scratchpad: (i: { input: string }) => i.input,
    },
    prompt,
    model,
    parser,
  ]);

  const executor = AgentExecutor.fromAgentAndTools({
    agent: runnable,
    tools,
  });
  const res = await executor.invoke({
    input:
      "Who is Olivia Wilde's boyfriend? What is his current age raised to the 0.23 power?",
  });
  // console.log(
  //   {
  //     res,
  //   },
  //   "Custom output parser"
  // );
  expect(res.output).toEqual("We did it!");
});

test("Add a fallback method", async () => {
  // Model should always fail since the model name passed does not exist.
  const modelBase = new ChatOpenAI({
    modelName: "fake-model",
    temperature: 10,
  });

  const modelLarge = new ChatOpenAI({
    modelName: "gpt-3.5-turbo-16k",
    temperature: 0.6,
  });

  const model = modelBase.withFallbacks({
    fallbacks: [modelLarge],
  });

  const prompt = ZeroShotAgent.createPrompt([]);
  const outputParser = ZeroShotAgent.getDefaultOutputParser();

  const runnable = RunnableSequence.from([
    {
      input: (i: { input: string }) => i.input,
      agent_scratchpad: (i: { input: string }) => i.input,
    },
    prompt,
    model,
    outputParser,
  ]);

  const executor = AgentExecutor.fromAgentAndTools({
    agent: runnable,
    tools: [],
  });
  const res = await executor.invoke({
    input: "Is the sky blue? Response with a concise answer",
  });
  // console.log(
  //   {
  //     res,
  //   },
  //   "Pass runnable to agent executor"
  // );
  expect(res.output).not.toEqual("");
  expect(res.output).not.toEqual("Agent stopped due to max iterations.");
});

test("Run agent with an abort signal", async () => {
  const model = new OpenAI({ temperature: 0, modelName: "text-babbage-001" });
  const tools = [new Calculator()];

  const executor = await initializeAgentExecutorWithOptions(tools, model, {
    agentType: "zero-shot-react-description",
  });
  // console.log("Loaded agent.");

  const input = `What is 3 to the fourth power?`;
  // console.log(`Executing with input "${input}"...`);

  const controller = new AbortController();
  await expect(() => {
    const result = executor.call({ input, signal: controller.signal });
    controller.abort();
    return result;
  }).rejects.toThrow();
});

test("Run agent with incorrect api key should throw error", async () => {
  const model = new OpenAI({
    temperature: 0,
    modelName: "text-babbage-001",
    openAIApiKey: "invalid",
  });
  const tools = [
    new SerpAPI(undefined, {
      location: "Austin,Texas,United States",
      hl: "en",
      gl: "us",
    }),
    new Calculator(),
  ];

  const executor = await initializeAgentExecutorWithOptions(tools, model, {
    agentType: "zero-shot-react-description",
  });
  // console.log("Loaded agent.");

  const input = `Who is Olivia Wilde's boyfriend? What is his current age raised to the 0.23 power?`;

  let error;
  // Test that the model throws an error
  await expect(async () => {
    try {
      await model.invoke(input);
    } catch (e) {
      error = e;
      throw e;
    }
  }).rejects.toThrowError();

  // Test that the agent throws the same error
  await expect(() => executor.call({ input })).rejects.toThrowError(
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    (error as any).message
  );
}, 10000);

test("Run tool web-browser", async () => {
  const model = new OpenAI({ temperature: 0 });
  const tools = [
    new SerpAPI(process.env.SERPAPI_API_KEY, {
      location: "Austin,Texas,United States",
      hl: "en",
      gl: "us",
    }),
    new Calculator(),
    new WebBrowser({ model, embeddings: new OpenAIEmbeddings() }),
  ];

  const executor = await initializeAgentExecutorWithOptions(tools, model, {
    agentType: "zero-shot-react-description",
    returnIntermediateSteps: true,
  });
  // console.log("Loaded agent.");

  const input = `What is the word of the day on merriam webster`;
  // console.log(`Executing with input "${input}"...`);

  const result = await executor.call({ input });
  // console.log(
  //   {
  //     result,
  //   },
  //   "Run tool web-browser"
  // );
  expect(result.intermediateSteps.length).toBeGreaterThanOrEqual(1);
  expect(result.intermediateSteps[0].action.tool).toEqual("search");
  expect(result.intermediateSteps[1].action.tool).toEqual("web-browser");
  expect(result.output).not.toEqual("");
  expect(result.output).not.toEqual("Agent stopped due to max iterations.");
});


import { zodToJsonSchema } from "zod-to-json-schema";
import fs from "fs";
import { z } from "zod";
import { AgentAction, AgentFinish, AgentStep } from "@langchain/core/agents";
import { AIMessage } from "@langchain/core/messages";
import { OpenAIEmbeddings, ChatOpenAI } from "@langchain/openai";
import { convertToOpenAIFunction } from "@langchain/core/utils/function_calling";
import { RunnableSequence } from "@langchain/core/runnables";
import {
  ChatPromptTemplate,
  MessagesPlaceholder,
} from "@langchain/core/prompts";
import { createRetrieverTool } from "../toolkits/index.js";
import { RecursiveCharacterTextSplitter } from "../../text_splitter.js";
import { MemoryVectorStore } from "../../vectorstores/memory.js";
import { AgentExecutor } from "../executor.js";
import { formatForOpenAIFunctions } from "../format_scratchpad/openai_functions.js";

/** Define a custom structured output parser. */
const structuredOutputParser = (
  output: AIMessage
): AgentAction | AgentFinish => {
  if (typeof output.content !== "string") {
    throw new Error("Cannot parse non-string output.");
  }
  if (output.additional_kwargs.function_call === undefined) {
    return { returnValues: { output: output.content }, log: output.content };
  }

  const functionCall = output.additional_kwargs.function_call;
  const name = functionCall?.name as string;
  const inputs = functionCall?.arguments as string;
  // console.log(functionCall);

  const jsonInput = JSON.parse(inputs);

  if (name === "response") {
    return { returnValues: { ...jsonInput }, log: output.content };
  }

  return {
    tool: name,
    toolInput: jsonInput,
    log: output.content,
  };
};

test("Pass custom structured output parsers", async () => {
  /** Read text file & embed documents */
  const text = fs.readFileSync("../examples/state_of_the_union.txt", "utf8");
  const textSplitter = new RecursiveCharacterTextSplitter({ chunkSize: 1000 });
  let docs = await textSplitter.createDocuments([text]);
  // Add fake source information
  docs = docs.map((doc, i) => ({
    ...doc,
    metadata: {
      page_chunk: i,
    },
  }));
  /** Initialize docs & create retriever */
  const vectorStore = await MemoryVectorStore.fromDocuments(
    docs,
    new OpenAIEmbeddings()
  );
  const retriever = vectorStore.asRetriever();
  /** Instantiate the LLM */
  const llm = new ChatOpenAI({});
  /** Define the prompt template */
  const prompt = ChatPromptTemplate.fromMessages([
    ["system", "You are a helpful assistant"],
    new MessagesPlaceholder("agent_scratchpad"),
    ["user", "{input}"],
  ]);
  /** Define the response schema */
  const responseSchema = z.object({
    answer: z.string().describe("The final answer to respond to the user"),
    sources: z
      .array(z.string())
      .describe(
        "List of page chunks that contain answer to the question. Only include a page chunk if it contains relevant information"
      ),
  });
  /** Create the response function */
  const responseOpenAIFunction = {
    name: "response",
    description: "Return the response to the user",
    parameters: zodToJsonSchema(responseSchema),
  };
  /** Convert retriever into a tool */
  const retrieverTool = createRetrieverTool(retriever, {
    name: "state-of-union-retriever",
    description:
      "Query a retriever to get information about state of the union address",
  });
  /** Bind both retriever and response functions to LLM */
  const llmWithTools = llm.bind({
    functions: [convertToOpenAIFunction(retrieverTool), responseOpenAIFunction],
  });
  /** Create the runnable */
  const runnableAgent = RunnableSequence.from([
    {
      input: (i: { input: string; steps: Array<AgentStep> }) => i.input,
      agent_scratchpad: (i: { input: string; steps: Array<AgentStep> }) =>
        formatForOpenAIFunctions(i.steps),
    },
    prompt,
    llmWithTools,
    structuredOutputParser,
  ]);
  /** Create the agent by passing in the runnable & tools */
  const executor = AgentExecutor.fromAgentAndTools({
    agent: runnableAgent,
    tools: [retrieverTool],
  });
  /** Call invoke on the agent */
  // @eslint-disable-next-line/@typescript-eslint/ban-ts-comment
  // @ts-expect-error unused var
  const res = await executor.invoke({
    input: "what did the president say about kentaji brown jackson",
  });
  // console.log({
  //   res,
  // });
  /**
    {
      res: {
        answer: 'President mentioned that he nominated Circuit Court of Appeals Judge Ketanji Brown Jackson. He described her as one of our nation’s top legal minds and stated that she will continue Justice Breyer’s legacy of excellence.',
        sources: [
          'And I did that 4 days ago, when I nominated Circuit Court of Appeals Judge Ketanji Brown Jackson. One of our nation’s top legal minds, who will continue Justice Breyer’s legacy of excellence. A former top litigator in private practice. A former federal public defender. And from a family of public school educators and police officers. A consensus builder. Since she’s been nominated, she’s received a broad range of support—from the Fraternal Order of Police to former judges appointed by Democrats and Republicans.'
        ]
      }
    }
   */
});


import type { BaseLanguageModelInterface } from "@langchain/core/language_models/base";
import { type ToolInterface, BaseToolkit } from "@langchain/core/tools";
import { renderTemplate } from "@langchain/core/prompts";
import {
  InfoSqlTool,
  ListTablesSqlTool,
  QueryCheckerTool,
  QuerySqlTool,
} from "../../../tools/sql.js";
import { SQL_PREFIX, SQL_SUFFIX } from "./prompt.js";
import { LLMChain } from "../../../chains/llm_chain.js";
import { ZeroShotAgent, ZeroShotCreatePromptArgs } from "../../mrkl/index.js";
import { AgentExecutor } from "../../executor.js";
import { SqlDatabase } from "../../../sql_db.js";

/**
 * Interface that extends ZeroShotCreatePromptArgs and adds an optional
 * topK parameter for specifying the number of results to return.
 */
export interface SqlCreatePromptArgs extends ZeroShotCreatePromptArgs {
  /** Number of results to return. */
  topK?: number;
}

/**
 * Class that represents a toolkit for working with SQL databases. It
 * initializes SQL tools based on the provided SQL database.
 * @example
 * ```typescript
 * const model = new ChatOpenAI({});
 * const toolkit = new SqlToolkit(sqlDb, model);
 * const executor = createSqlAgent(model, toolkit);
 * const result = await executor.invoke({ input: 'List the total sales per country. Which country's customers spent the most?' });
 * console.log(`Got output ${result.output}`);
 * ```
 */
export class SqlToolkit extends BaseToolkit {
  tools: ToolInterface[];

  db: SqlDatabase;

  dialect = "sqlite";

  constructor(db: SqlDatabase, llm?: BaseLanguageModelInterface) {
    super();
    this.db = db;
    this.tools = [
      new QuerySqlTool(db),
      new InfoSqlTool(db),
      new ListTablesSqlTool(db),
      new QueryCheckerTool({ llm }),
    ];
  }
}

export function createSqlAgent(
  llm: BaseLanguageModelInterface,
  toolkit: SqlToolkit,
  args?: SqlCreatePromptArgs
) {
  const {
    prefix = SQL_PREFIX,
    suffix = SQL_SUFFIX,
    inputVariables = ["input", "agent_scratchpad"],
    topK = 10,
  } = args ?? {};
  const { tools } = toolkit;
  const formattedPrefix = renderTemplate(prefix, "f-string", {
    dialect: toolkit.dialect,
    top_k: topK,
  });

  const prompt = ZeroShotAgent.createPrompt(tools, {
    prefix: formattedPrefix,
    suffix,
    inputVariables,
  });
  const chain = new LLMChain({ prompt, llm });
  const agent = new ZeroShotAgent({
    llmChain: chain,
    allowedTools: tools.map((t) => t.name),
  });
  return AgentExecutor.fromAgentAndTools({
    agent,
    tools,
    returnIntermediateSteps: true,
  });
}


import { AgentAction, AgentFinish } from "@langchain/core/agents";
import { OutputParserException } from "@langchain/core/output_parsers";
import { AgentActionOutputParser } from "../types.js";

/**
 * @example
 * ```typescript
 * const prompt = ChatPromptTemplate.fromMessages([
 *   HumanMessagePromptTemplate.fromTemplate(AGENT_INSTRUCTIONS),
 *   new MessagesPlaceholder("agent_scratchpad"),
 * ]);
 * const runnableAgent = RunnableSequence.from([
 *   ...rest of runnable
 *   prompt,
 *   new ChatAnthropic({ modelName: "claude-2", temperature: 0 }).bind({
 *     stop: ["</tool_input>", "</final_answer>"],
 *   }),
 *   new XMLAgentOutputParser(),
 * ]);
 * const result = await executor.invoke({
 *   input: "What is the weather in Honolulu?",
 *   tools: [],
 * });
 * ```
 */
export class XMLAgentOutputParser extends AgentActionOutputParser {
  lc_namespace = ["langchain", "agents", "xml"];

  static lc_name() {
    return "XMLAgentOutputParser";
  }

  /**
   * Parses the output text from the agent and returns an AgentAction or
   * AgentFinish object.
   * @param text The output text from the agent.
   * @returns An AgentAction or AgentFinish object.
   */
  async parse(text: string): Promise<AgentAction | AgentFinish> {
    if (text.includes("</tool>")) {
      const [tool, toolInput] = text.split("</tool>");
      const _tool = tool.split("<tool>")[1];
      const _toolInput = toolInput.split("<tool_input>")[1];
      return { tool: _tool, toolInput: _toolInput, log: text };
    } else if (text.includes("<final_answer>")) {
      const [, answer] = text.split("<final_answer>");
      return { returnValues: { output: answer }, log: text };
    } else {
      throw new OutputParserException(`Could not parse LLM output: ${text}`);
    }
  }

  getFormatInstructions(): string {
    throw new Error(
      "getFormatInstructions not implemented inside OpenAIFunctionsAgentOutputParser."
    );
  }
}