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);
  });
});


// eslint-disable-next-line import/no-extraneous-dependencies
import { loadPyodide, type PyodideInterface } from "pyodide";
import { Tool, ToolParams } from "@langchain/core/tools";

export type PythonInterpreterToolParams = Parameters<typeof loadPyodide>[0] &
  ToolParams & {
    instance: PyodideInterface;
  };

export class PythonInterpreterTool extends Tool {
  static lc_name() {
    return "PythonInterpreterTool";
  }

  name = "python_interpreter";

  description = `Evaluates python code in a sandbox environment. The environment resets on every execution. You must send the whole script every time and print your outputs. Script should be pure python code that can be evaluated. Packages available:
${this.availableDefaultPackages}`;

  pyodideInstance: PyodideInterface;

  stdout = "";

  stderr = "";

  constructor(options: PythonInterpreterToolParams) {
    super(options);
    this.pyodideInstance = options.instance;
    this.pyodideInstance.setStderr({
      batched: (text: string) => {
        this.stderr += text;
      },
    });

    this.pyodideInstance.setStdout({
      batched: (text: string) => {
        this.stdout += text;
      },
    });
  }

  async addPackage(packageName: string) {
    await this.pyodideInstance.loadPackage(packageName);
    this.description += `, ${packageName}`;
  }

  get availableDefaultPackages(): string {
    return [
      "asciitree",
      "astropy",
      "atomicwrites",
      "attrs",
      "autograd",
      "awkward-cpp",
      "bcrypt",
      "beautifulsoup4",
      "biopython",
      "bitarray",
      "bitstring",
      "bleach",
      "bokeh",
      "boost-histogram",
      "brotli",
      "cachetools",
      "Cartopy",
      "cbor-diag",
      "certifi",
      "cffi",
      "cffi_example",
      "cftime",
      "click",
      "cligj",
      "cloudpickle",
      "cmyt",
      "colorspacious",
      "contourpy",
      "coolprop",
      "coverage",
      "cramjam",
      "cryptography",
      "cssselect",
      "cycler",
      "cytoolz",
      "decorator",
      "demes",
      "deprecation",
      "distlib",
      "docutils",
      "exceptiongroup",
      "fastparquet",
      "fiona",
      "fonttools",
      "freesasa",
      "fsspec",
      "future",
      "galpy",
      "gensim",
      "geopandas",
      "gmpy2",
      "gsw",
      "h5py",
      "html5lib",
      "idna",
      "igraph",
      "imageio",
      "iniconfig",
      "jedi",
      "Jinja2",
      "joblib",
      "jsonschema",
      "kiwisolver",
      "lazy-object-proxy",
      "lazy_loader",
      "lightgbm",
      "logbook",
      "lxml",
      "MarkupSafe",
      "matplotlib",
      "matplotlib-pyodide",
      "micropip",
      "mne",
      "more-itertools",
      "mpmath",
      "msgpack",
      "msprime",
      "multidict",
      "munch",
      "mypy",
      "netcdf4",
      "networkx",
      "newick",
      "nlopt",
      "nltk",
      "nose",
      "numcodecs",
      "numpy",
      "opencv-python",
      "optlang",
      "orjson",
      "packaging",
      "pandas",
      "parso",
      "patsy",
      "peewee",
      "Pillow",
      "pillow_heif",
      "pkgconfig",
      "pluggy",
      "protobuf",
      "py",
      "pyb2d",
      "pyclipper",
      "pycparser",
      "pycryptodome",
      "pydantic",
      "pyerfa",
      "Pygments",
      "pyheif",
      "pyinstrument",
      "pynacl",
      "pyodide-http",
      "pyodide-tblib",
      "pyparsing",
      "pyproj",
      "pyrsistent",
      "pyshp",
      "pytest",
      "pytest-benchmark",
      "python-dateutil",
      "python-magic",
      "python-sat",
      "python_solvespace",
      "pytz",
      "pywavelets",
      "pyxel",
      "pyyaml",
      "rebound",
      "reboundx",
      "regex",
      "retrying",
      "RobotRaconteur",
      "ruamel.yaml",
      "rust-panic-test",
      "scikit-image",
      "scikit-learn",
      "scipy",
      "screed",
      "setuptools",
      "shapely",
      "simplejson",
      "six",
      "smart_open",
      "soupsieve",
      "sourmash",
      "sparseqr",
      "sqlalchemy",
      "statsmodels",
      "svgwrite",
      "swiglpk",
      "sympy",
      "termcolor",
      "texttable",
      "threadpoolctl",
      "tomli",
      "tomli-w",
      "toolz",
      "tqdm",
      "traits",
      "tskit",
      "typing-extensions",
      "uncertainties",
      "unyt",
      "webencodings",
      "wordcloud",
      "wrapt",
      "xarray",
      "xgboost",
      "xlrd",
      "xyzservices",
      "yarl",
      "yt",
      "zarr",
    ].join(", ");
  }

  static async initialize(
    options: Omit<PythonInterpreterToolParams, "instance">
  ) {
    const instance = await loadPyodide(options);
    return new this({ ...options, instance });
  }

  async _call(script: string) {
    this.stdout = "";
    this.stderr = "";

    await this.pyodideInstance.runPythonAsync(script);
    return JSON.stringify({ stdout: this.stdout, stderr: this.stderr });
  }
}


import { TextServiceClient, protos } from "@google-ai/generativelanguage";
import { GoogleAuth } from "google-auth-library";
import { type BaseLLMParams, LLM } from "@langchain/core/language_models/llms";
import { getEnvironmentVariable } from "@langchain/core/utils/env";

/**
 * @deprecated - Deprecated by Google. Will be removed in 0.3.0
 *
 * Input for Text generation for Google Palm
 */
export interface GooglePaLMTextInput extends BaseLLMParams {
  /**
   * Model Name to use
   *
   * Alias for `model`
   *
   * Note: The format must follow the pattern - `models/{model}`
   */
  modelName?: string;
  /**
   * Model Name to use
   *
   * Note: The format must follow the pattern - `models/{model}`
   */
  model?: string;

  /**
   * Controls the randomness of the output.
   *
   * Values can range from [0.0,1.0], inclusive. A value closer to 1.0
   * will produce responses that are more varied and creative, while
   * a value closer to 0.0 will typically result in more straightforward
   * responses from the model.
   *
   * Note: The default value varies by model
   */
  temperature?: number;

  /**
   * Maximum number of tokens to generate in the completion.
   */
  maxOutputTokens?: number;

  /**
   * Top-p changes how the model selects tokens for output.
   *
   * Tokens are selected from most probable to least until the sum
   * of their probabilities equals the top-p value.
   *
   * For example, if tokens A, B, and C have a probability of
   * .3, .2, and .1 and the top-p value is .5, then the model will
   * select either A or B as the next token (using temperature).
   *
   * Note: The default value varies by model
   */
  topP?: number;

  /**
   * Top-k changes how the model selects tokens for output.
   *
   * A top-k of 1 means the selected token is the most probable among
   * all tokens in the model’s vocabulary (also called greedy decoding),
   * while a top-k of 3 means that the next token is selected from
   * among the 3 most probable tokens (using temperature).
   *
   * Note: The default value varies by model
   */
  topK?: number;

  /**
   * The set of character sequences (up to 5) that will stop output generation.
   * If specified, the API will stop at the first appearance of a stop
   * sequence.
   *
   * Note: The stop sequence will not be included as part of the response.
   */
  stopSequences?: string[];

  /**
   * A list of unique `SafetySetting` instances for blocking unsafe content. The API will block
   * any prompts and responses that fail to meet the thresholds set by these settings. If there
   * is no `SafetySetting` for a given `SafetyCategory` provided in the list, the API will use
   * the default safety setting for that category.
   */
  safetySettings?: protos.google.ai.generativelanguage.v1beta2.ISafetySetting[];

  /**
   * Google Palm API key to use
   */
  apiKey?: string;
}

/**
 * @deprecated - Deprecated by Google. Will be removed in 0.3.0
 *
 * Google Palm 2 Language Model Wrapper to generate texts
 */
export class GooglePaLM extends LLM implements GooglePaLMTextInput {
  lc_serializable = true;

  get lc_secrets(): { [key: string]: string } | undefined {
    return {
      apiKey: "GOOGLE_PALM_API_KEY",
    };
  }

  modelName = "models/text-bison-001";

  model = "models/text-bison-001";

  temperature?: number; // default value chosen based on model

  maxOutputTokens?: number; // defaults to 64

  topP?: number; // default value chosen based on model

  topK?: number; // default value chosen based on model

  stopSequences: string[] = [];

  safetySettings?: protos.google.ai.generativelanguage.v1beta2.ISafetySetting[]; // default safety setting for that category

  apiKey?: string;

  private client: TextServiceClient;

  constructor(fields?: GooglePaLMTextInput) {
    super(fields ?? {});

    this.modelName = fields?.model ?? fields?.modelName ?? this.model;
    this.model = this.modelName;

    this.temperature = fields?.temperature ?? this.temperature;
    if (this.temperature && (this.temperature < 0 || this.temperature > 1)) {
      throw new Error("`temperature` must be in the range of [0.0,1.0]");
    }

    this.maxOutputTokens = fields?.maxOutputTokens ?? this.maxOutputTokens;
    if (this.maxOutputTokens && this.maxOutputTokens < 0) {
      throw new Error("`maxOutputTokens` must be a positive integer");
    }

    this.topP = fields?.topP ?? this.topP;
    if (this.topP && this.topP < 0) {
      throw new Error("`topP` must be a positive integer");
    }

    if (this.topP && this.topP > 1) {
      throw new Error("Google PaLM `topP` must in the range of [0,1]");
    }

    this.topK = fields?.topK ?? this.topK;
    if (this.topK && this.topK < 0) {
      throw new Error("`topK` must be a positive integer");
    }

    this.stopSequences = fields?.stopSequences ?? this.stopSequences;

    this.safetySettings = fields?.safetySettings ?? this.safetySettings;
    if (this.safetySettings && this.safetySettings.length > 0) {
      const safetySettingsSet = new Set(
        this.safetySettings.map((s) => s.category)
      );
      if (safetySettingsSet.size !== this.safetySettings.length) {
        throw new Error(
          "The categories in `safetySettings` array must be unique"
        );
      }
    }

    this.apiKey =
      fields?.apiKey ?? getEnvironmentVariable("GOOGLE_PALM_API_KEY");
    if (!this.apiKey) {
      throw new Error(
        "Please set an API key for Google Palm 2 in the environment variable GOOGLE_PALM_API_KEY or in the `apiKey` field of the GooglePalm constructor"
      );
    }

    this.client = new TextServiceClient({
      authClient: new GoogleAuth().fromAPIKey(this.apiKey),
    });
  }

  _llmType(): string {
    return "googlepalm";
  }

  async _call(
    prompt: string,
    options: this["ParsedCallOptions"]
  ): Promise<string> {
    const res = await this.caller.callWithOptions(
      { signal: options.signal },
      this._generateText.bind(this),
      prompt
    );
    return res ?? "";
  }

  protected async _generateText(
    prompt: string
  ): Promise<string | null | undefined> {
    const res = await this.client.generateText({
      model: this.model,
      temperature: this.temperature,
      candidateCount: 1,
      topK: this.topK,
      topP: this.topP,
      maxOutputTokens: this.maxOutputTokens,
      stopSequences: this.stopSequences,
      safetySettings: this.safetySettings,
      prompt: {
        text: prompt,
      },
    });
    return res[0].candidates && res[0].candidates.length > 0
      ? res[0].candidates[0].output
      : undefined;
  }
}


import { z } from "zod";
import { zodToJsonSchema } from "zod-to-json-schema";
import { BaseLanguageModel } from "@langchain/core/language_models/base";
import { ChatPromptTemplate } from "@langchain/core/prompts";
import { Document } from "@langchain/core/documents";
import {
  Node,
  Relationship,
  GraphDocument,
} from "../../graphs/graph_document.js";

export const SYSTEM_PROMPT = `
# Knowledge Graph Instructions for GPT-4\n
## 1. Overview\n
You are a top-tier algorithm designed for extracting information in structured formats to build a knowledge graph.\n
Try to capture as much information from the text as possible without sacrifing accuracy. Do not add any information that is not explicitly mentioned in the text\n"
- **Nodes** represent entities and concepts.\n"
- The aim is to achieve simplicity and clarity in the knowledge graph, making it\n
accessible for a vast audience.\n
## 2. Labeling Nodes\n
- **Consistency**: Ensure you use available types for node labels.\n
Ensure you use basic or elementary types for node labels.\n
- For example, when you identify an entity representing a person, always label it as **'person'**. Avoid using more specific terms like 'mathematician' or 'scientist'
- **Node IDs**: Never utilize integers as node IDs. Node IDs should be names or human-readable identifiers found in the text.\n
- **Relationships** represent connections between entities or concepts.\n
Ensure consistency and generality in relationship types when constructing knowledge graphs. Instead of using specific and momentary types such as 'BECAME_PROFESSOR', use more general and timeless relationship types like 'PROFESSOR'. Make sure to use general and timeless relationship types!\n
## 3. Coreference Resolution\n
- **Maintain Entity Consistency**: When extracting entities, it's vital to ensure consistency.\n
If an entity, such as "John Doe", is mentioned multiple times in the text but is referred to by different names or pronouns (e.g., "Joe", "he"), always use the most complete identifier for that entity throughout the knowledge graph. In this example, use "John Doe" as the entity ID.\n
Remember, the knowledge graph should be coherent and easily understandable, so maintaining consistency in entity references is crucial.\n
## 4. Strict Compliance\n
Adhere to the rules strictly. Non-compliance will result in termination.
`;

const DEFAULT_PROMPT = /* #__PURE__ */ ChatPromptTemplate.fromMessages([
  ["system", SYSTEM_PROMPT],
  [
    "human",
    "Tip: Make sure to answer in the correct format and do not include any explanations. Use the given format to extract information from the following input: {input}",
  ],
]);

interface OptionalEnumFieldProps {
  enumValues?: string[];
  description: string;
  isRel?: boolean;
  fieldKwargs?: object;
}

function toTitleCase(str: string): string {
  return str
    .split(" ")
    .map((w) => w[0].toUpperCase() + w.substring(1).toLowerCase())
    .join("");
}

function createOptionalEnumType({
  enumValues = undefined,
  description = "",
  isRel = false,
}: OptionalEnumFieldProps): z.ZodTypeAny {
  let schema;

  if (enumValues && enumValues.length) {
    schema = z
      .enum(enumValues as [string, ...string[]])
      .describe(
        `${description} Available options are: ${enumValues.join(", ")}.`
      );
  } else {
    const nodeInfo =
      "Ensure you use basic or elementary types for node labels.\n" +
      "For example, when you identify an entity representing a person, " +
      "always label it as **'Person'**. Avoid using more specific terms " +
      "like 'Mathematician' or 'Scientist'";
    const relInfo =
      "Instead of using specific and momentary types such as " +
      "'BECAME_PROFESSOR', use more general and timeless relationship types like " +
      "'PROFESSOR'. However, do not sacrifice any accuracy for generality";

    const additionalInfo = isRel ? relInfo : nodeInfo;

    schema = z.string().describe(description + additionalInfo);
  }

  return schema;
}

function createSchema(allowedNodes: string[], allowedRelationships: string[]) {
  const dynamicGraphSchema = z.object({
    nodes: z
      .array(
        z.object({
          id: z.string(),
          type: createOptionalEnumType({
            enumValues: allowedNodes,
            description: "The type or label of the node.",
          }),
        })
      )
      .describe("List of nodes"),
    relationships: z
      .array(
        z.object({
          sourceNodeId: z.string(),
          sourceNodeType: createOptionalEnumType({
            enumValues: allowedNodes,
            description: "The source node of the relationship.",
          }),
          relationshipType: createOptionalEnumType({
            enumValues: allowedRelationships,
            description: "The type of the relationship.",
            isRel: true,
          }),
          targetNodeId: z.string(),
          targetNodeType: createOptionalEnumType({
            enumValues: allowedNodes,
            description: "The target node of the relationship.",
          }),
        })
      )
      .describe("List of relationships."),
  });

  return dynamicGraphSchema;
}

// eslint-disable-next-line @typescript-eslint/no-explicit-any
function mapToBaseNode(node: any): Node {
  return new Node({
    id: node.id,
    type: node.type ? toTitleCase(node.type) : "",
  });
}

// eslint-disable-next-line @typescript-eslint/no-explicit-any
function mapToBaseRelationship(relationship: any): Relationship {
  return new Relationship({
    source: new Node({
      id: relationship.sourceNodeId,
      type: relationship.sourceNodeType
        ? toTitleCase(relationship.sourceNodeType)
        : "",
    }),
    target: new Node({
      id: relationship.targetNodeId,
      type: relationship.targetNodeType
        ? toTitleCase(relationship.targetNodeType)
        : "",
    }),
    type: relationship.relationshipType.replace(" ", "_").toUpperCase(),
  });
}

export interface LLMGraphTransformerProps {
  llm: BaseLanguageModel;
  allowedNodes?: string[];
  allowedRelationships?: string[];
  prompt?: ChatPromptTemplate;
  strictMode?: boolean;
}


import { test } from "@jest/globals";
import { Document } from "@langchain/core/documents";
import { OpenAIEmbeddings, OpenAI } from "@langchain/openai";
import { AttributeInfo } from "langchain/chains/query_constructor";
import {
  FunctionalTranslator,
  SelfQueryRetriever,
} from "langchain/retrievers/self_query";
import { HNSWLib } from "../../vectorstores/hnswlib.js";

test("HNSWLib Store Self Query Retriever Test", async () => {
  const docs = [
    new Document({
      pageContent:
        "A bunch of scientists bring back dinosaurs and mayhem breaks loose",
      metadata: { year: 1993, rating: 7.7, genre: "science fiction" },
    }),
    new Document({
      pageContent:
        "Leo DiCaprio gets lost in a dream within a dream within a dream within a ...",
      metadata: { year: 2010, director: "Christopher Nolan", rating: 8.2 },
    }),
    new Document({
      pageContent:
        "A psychologist / detective gets lost in a series of dreams within dreams within dreams and Inception reused the idea",
      metadata: { year: 2006, director: "Satoshi Kon", rating: 8.6 },
    }),
    new Document({
      pageContent:
        "A bunch of normal-sized women are supremely wholesome and some men pine after them",
      metadata: { year: 2019, director: "Greta Gerwig", rating: 8.3 },
    }),
    new Document({
      pageContent: "Toys come alive and have a blast doing so",
      metadata: { year: 1995, genre: "animated" },
    }),
    new Document({
      pageContent:
        "Three men walk into the Zone, three men walk out of the Zone",
      metadata: {
        year: 1979,
        director: "Andrei Tarkovsky",
        genre: "science fiction",
        rating: 9.9,
      },
    }),
  ];

  const attributeInfo: AttributeInfo[] = [
    {
      name: "genre",
      description: "The genre of the movie",
      type: "string or array of strings",
    },
    {
      name: "year",
      description: "The year the movie was released",
      type: "number",
    },
    {
      name: "director",
      description: "The director of the movie",
      type: "string",
    },
    {
      name: "rating",
      description: "The rating of the movie (1-10)",
      type: "number",
    },
    {
      name: "length",
      description: "The length of the movie in minutes",
      type: "number",
    },
  ];

  const embeddings = new OpenAIEmbeddings();
  const llm = new OpenAI({
    modelName: "gpt-3.5-turbo",
    temperature: 0.01,
  });
  const documentContents = "Brief summary of a movie";
  const vectorStore = await HNSWLib.fromDocuments(docs, embeddings);
  const selfQueryRetriever = SelfQueryRetriever.fromLLM({
    llm,
    vectorStore,
    documentContents,
    attributeInfo,
    structuredQueryTranslator: new FunctionalTranslator(),
  });

  const query1 = await selfQueryRetriever.getRelevantDocuments(
    "Which movies are less than 90 minutes?"
  );
  // console.log(query1);
  expect(query1.length).toEqual(0);

  const query2 = await selfQueryRetriever.getRelevantDocuments(
    "Which movies are rated higher than 8.5?"
  );
  // console.log(query2);
  expect(query2.length).toEqual(2);

  const query3 = await selfQueryRetriever.getRelevantDocuments(
    "Which movies are directed by Greta Gerwig?"
  );
  // console.log(query3);
  expect(query3.length).toEqual(1);
});

test("HNSWLib shouldn't throw an error if a filter can't be generated, but should return no items", async () => {
  const docs = [
    new Document({
      pageContent:
        "A bunch of scientists bring back dinosaurs and mayhem breaks loose",
      metadata: { year: 1993, rating: 7.7, genre: "science fiction" },
    }),
    new Document({
      pageContent:
        "Leo DiCaprio gets lost in a dream within a dream within a dream within a ...",
      metadata: { year: 2010, director: "Christopher Nolan", rating: 8.2 },
    }),
    new Document({
      pageContent:
        "A psychologist / detective gets lost in a series of dreams within dreams within dreams and Inception reused the idea",
      metadata: { year: 2006, director: "Satoshi Kon", rating: 8.6 },
    }),
    new Document({
      pageContent:
        "A bunch of normal-sized women are supremely wholesome and some men pine after them",
      metadata: { year: 2019, director: "Greta Gerwig", rating: 8.3 },
    }),
    new Document({
      pageContent: "Toys come alive and have a blast doing so",
      metadata: { year: 1995, genre: "animated" },
    }),
    new Document({
      pageContent:
        "Three men walk into the Zone, three men walk out of the Zone",
      metadata: {
        year: 1979,
        director: "Andrei Tarkovsky",
        genre: "science fiction",
        rating: 9.9,
      },
    }),
  ];

  const attributeInfo = [
    {
      name: "sectionNumber",
      description: "The section number of the rule",
      type: "number",
    },
    {
      name: "sectionTitle",
      description: "The section title of the rule",
      type: "string",
    },
    {
      name: "sectionScope",
      description: "The section scope of the rule",
      type: "string",
    },
    {
      name: "codeRule",
      description: "The code rule of the rule",
      type: "string",
    },
  ];

  const embeddings = new OpenAIEmbeddings();
  const llm = new OpenAI({
    modelName: "gpt-3.5-turbo",
    temperature: 0.01,
  });
  const documentContents = "Brief summary of a movie";
  const vectorStore = await HNSWLib.fromDocuments(docs, embeddings);
  const selfQueryRetriever = SelfQueryRetriever.fromLLM({
    llm,
    vectorStore,
    documentContents,
    attributeInfo,
    structuredQueryTranslator: new FunctionalTranslator(),
  });

  const query1 = await selfQueryRetriever.getRelevantDocuments(
    "Which sectionTitle talks about pools?"
  );
  // console.log(query1);
  expect(query1.length).toEqual(0);
});


import Metal from "@getmetal/metal-sdk";

import { BaseRetriever, BaseRetrieverInput } from "@langchain/core/retrievers";
import { Document } from "@langchain/core/documents";

/**
 * Interface for the fields required during the initialization of a
 * `MetalRetriever` instance. It extends the `BaseRetrieverInput`
 * interface and adds a `client` field of type `Metal`.
 */
export interface MetalRetrieverFields extends BaseRetrieverInput {
  client: Metal;
}

/**
 * Interface to represent a response item from the Metal service. It
 * contains a `text` field and an index signature to allow for additional
 * unknown properties.
 */
interface ResponseItem {
  text: string;
  [key: string]: unknown;
}

/**
 * Class used to interact with the Metal service, a managed retrieval &
 * memory platform. It allows you to index your data into Metal and run
 * semantic search and retrieval on it. It extends the `BaseRetriever`
 * class and requires a `Metal` instance and a dictionary of parameters to
 * pass to the Metal API during its initialization.
 * @example
 * ```typescript
 * const retriever = new MetalRetriever({
 *   client: new Metal(
 *     process.env.METAL_API_KEY,
 *     process.env.METAL_CLIENT_ID,
 *     process.env.METAL_INDEX_ID,
 *   ),
 * });
 * const docs = await retriever.getRelevantDocuments("hello");
 * ```
 */
export class MetalRetriever extends BaseRetriever {
  static lc_name() {
    return "MetalRetriever";
  }

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

  private client: Metal;

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

    this.client = fields.client;
  }

  async _getRelevantDocuments(query: string): Promise<Document[]> {
    const res = await this.client.search({ text: query });

    const items = ("data" in res ? res.data : res) as ResponseItem[];
    return items.map(
      ({ text, metadata }) =>
        new Document({
          pageContent: text,
          metadata: metadata as Record<string, unknown>,
        })
    );
  }
}


import {
  MemorySearchPayload,
  MemorySearchResult,
  NotFoundError,
  ZepClient,
} from "@getzep/zep-js";
import { BaseRetriever, BaseRetrieverInput } from "@langchain/core/retrievers";
import { Document } from "@langchain/core/documents";

/**
 * Configuration interface for the ZepRetriever class. Extends the
 * BaseRetrieverInput interface.
 *
 * @argument {string} sessionId - The ID of the Zep session.
 * @argument {string} url - The URL of the Zep API.
 * @argument {number} [topK] - The number of results to return.
 * @argument {string} [apiKey] - The API key for the Zep API.
 * @argument [searchScope] [searchScope] - The scope of the search: "messages" or "summary".
 * @argument [searchType] [searchType] - The type of search to perform: "similarity" or "mmr".
 * @argument {number} [mmrLambda] - The lambda value for the MMR search.
 * @argument {Record<string, unknown>} [filter] - The metadata filter to apply to the search.
 */
export interface ZepRetrieverConfig extends BaseRetrieverInput {
  sessionId: string;
  url: string;
  topK?: number;
  apiKey?: string;
  searchScope?: "messages" | "summary";
  searchType?: "similarity" | "mmr";
  mmrLambda?: number;
  filter?: Record<string, unknown>;
}

/**
 * Class for retrieving information from a Zep long-term memory store.
 * Extends the BaseRetriever class.
 * @example
 * ```typescript
 * const retriever = new ZepRetriever({
 *   url: "http:
 *   sessionId: "session_exampleUUID",
 *   topK: 3,
 * });
 * const query = "Can I drive red cars in France?";
 * const docs = await retriever.getRelevantDocuments(query);
 * ```
 */
export class ZepRetriever extends BaseRetriever {
  static lc_name() {
    return "ZepRetriever";
  }

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

  get lc_secrets(): { [key: string]: string } | undefined {
    return {
      apiKey: "ZEP_API_KEY",
      url: "ZEP_API_URL",
    };
  }

  get lc_aliases(): { [key: string]: string } | undefined {
    return { apiKey: "api_key" };
  }

  zepClientPromise: Promise<ZepClient>;

  private sessionId: string;

  private topK?: number;

  private searchScope?: "messages" | "summary";

  private searchType?: "similarity" | "mmr";

  private mmrLambda?: number;

  private filter?: Record<string, unknown>;

  constructor(config: ZepRetrieverConfig) {
    super(config);
    this.sessionId = config.sessionId;
    this.topK = config.topK;
    this.searchScope = config.searchScope;
    this.searchType = config.searchType;
    this.mmrLambda = config.mmrLambda;
    this.filter = config.filter;
    this.zepClientPromise = ZepClient.init(config.url, config.apiKey);
  }

  /**
   *  Converts an array of message search results to an array of Document objects.
   *  @param {MemorySearchResult[]} results - The array of search results.
   *  @returns {Document[]} An array of Document objects representing the search results.
   */
  private searchMessageResultToDoc(results: MemorySearchResult[]): Document[] {
    return results
      .filter((r) => r.message)
      .map(
        ({
          message: { content, metadata: messageMetadata } = {},
          dist,
          ...rest
        }) =>
          new Document({
            pageContent: content ?? "",
            metadata: { score: dist, ...messageMetadata, ...rest },
          })
      );
  }

  /**
   *  Converts an array of summary search results to an array of Document objects.
   *  @param {MemorySearchResult[]} results - The array of search results.
   *  @returns {Document[]} An array of Document objects representing the search results.
   */
  private searchSummaryResultToDoc(results: MemorySearchResult[]): Document[] {
    return results
      .filter((r) => r.summary)
      .map(
        ({
          summary: { content, metadata: summaryMetadata } = {},
          dist,
          ...rest
        }) =>
          new Document({
            pageContent: content ?? "",
            metadata: { score: dist, ...summaryMetadata, ...rest },
          })
      );
  }

  /**
   *  Retrieves the relevant documents based on the given query.
   *  @param {string} query - The query string.
   *  @returns {Promise<Document[]>} A promise that resolves to an array of relevant Document objects.
   */
  async _getRelevantDocuments(query: string): Promise<Document[]> {
    const payload: MemorySearchPayload = {
      text: query,
      metadata: this.filter,
      search_scope: this.searchScope,
      search_type: this.searchType,
      mmr_lambda: this.mmrLambda,
    };
    // Wait for ZepClient to be initialized
    const zepClient = await this.zepClientPromise;
    if (!zepClient) {
      throw new Error("ZepClient is not initialized");
    }
    try {
      const results: MemorySearchResult[] = await zepClient.memory.searchMemory(
        this.sessionId,
        payload,
        this.topK
      );
      return this.searchScope === "summary"
        ? this.searchSummaryResultToDoc(results)
        : this.searchMessageResultToDoc(results);
    } catch (error) {
      // eslint-disable-next-line no-instanceof/no-instanceof
      if (error instanceof NotFoundError) {
        return Promise.resolve([]); // Return an empty Document array
      }
      // If it's not a NotFoundError, throw the error again
      throw error;
    }
  }
}


import {
  BaseRetriever,
  type BaseRetrieverInput,
} from "@langchain/core/retrievers";
import { Document } from "@langchain/core/documents";
import {
  AsyncCaller,
  AsyncCallerParams,
} from "@langchain/core/utils/async_caller";

/**
 * Interface for the arguments required to create a new instance of
 * DataberryRetriever.
 */
export interface DataberryRetrieverArgs
  extends AsyncCallerParams,
    BaseRetrieverInput {
  datastoreUrl: string;
  topK?: number;
  apiKey?: string;
}

/**
 * Interface for the structure of a Berry object returned by the Databerry
 * API.
 */
interface Berry {
  text: string;
  score: number;
  source?: string;
  [key: string]: unknown;
}

/**
 * A specific implementation of a document retriever for the Databerry
 * API. It extends the BaseRetriever class, which is an abstract base
 * class for a document retrieval system in LangChain.
 */
/** @deprecated Use "langchain/retrievers/chaindesk" instead */
export class DataberryRetriever extends BaseRetriever {
  static lc_name() {
    return "DataberryRetriever";
  }

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

  get lc_secrets() {
    return { apiKey: "DATABERRY_API_KEY" };
  }

  get lc_aliases() {
    return { apiKey: "api_key" };
  }

  caller: AsyncCaller;

  datastoreUrl: string;

  topK?: number;

  apiKey?: string;

  constructor(fields: DataberryRetrieverArgs) {
    super(fields);
    const { datastoreUrl, apiKey, topK, ...rest } = fields;

    this.caller = new AsyncCaller(rest);
    this.datastoreUrl = datastoreUrl;
    this.apiKey = apiKey;
    this.topK = topK;
  }

  async _getRelevantDocuments(query: string): Promise<Document[]> {
    const r = await this.caller.call(fetch, this.datastoreUrl, {
      method: "POST",
      body: JSON.stringify({
        query,
        ...(this.topK ? { topK: this.topK } : {}),
      }),
      headers: {
        "Content-Type": "application/json",
        ...(this.apiKey ? { Authorization: `Bearer ${this.apiKey}` } : {}),
      },
    });

    const { results } = (await r.json()) as { results: Berry[] };

    return results.map(
      ({ text, score, source, ...rest }) =>
        new Document({
          pageContent: text,
          metadata: {
            score,
            source,
            ...rest,
          },
        })
    );
  }
}


import {
  BaseRetriever,
  type BaseRetrieverInput,
} from "@langchain/core/retrievers";
import { Document } from "@langchain/core/documents";
import {
  AsyncCaller,
  type AsyncCallerParams,
} from "@langchain/core/utils/async_caller";

export interface ChaindeskRetrieverArgs
  extends AsyncCallerParams,
    BaseRetrieverInput {
  datastoreId: string;
  topK?: number;
  filter?: Record<string, unknown>;
  apiKey?: string;
}

interface Berry {
  text: string;
  score: number;
  source?: string;
  [key: string]: unknown;
}

/**
 * @example
 * ```typescript
 * const retriever = new ChaindeskRetriever({
 *   datastoreId: "DATASTORE_ID",
 *   apiKey: "CHAINDESK_API_KEY",
 *   topK: 8,
 * });
 * const docs = await retriever.getRelevantDocuments("hello");
 * ```
 */
export class ChaindeskRetriever extends BaseRetriever {
  static lc_name() {
    return "ChaindeskRetriever";
  }

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

  caller: AsyncCaller;

  datastoreId: string;

  topK?: number;

  filter?: Record<string, unknown>;

  apiKey?: string;

  constructor({
    datastoreId,
    apiKey,
    topK,
    filter,
    ...rest
  }: ChaindeskRetrieverArgs) {
    super();

    this.caller = new AsyncCaller(rest);
    this.datastoreId = datastoreId;
    this.apiKey = apiKey;
    this.topK = topK;
    this.filter = filter;
  }

  async getRelevantDocuments(query: string): Promise<Document[]> {
    const r = await this.caller.call(
      fetch,
      `https://app.chaindesk.ai/api/datastores/${this.datastoreId}/query`,
      {
        method: "POST",
        body: JSON.stringify({
          query,
          ...(this.topK ? { topK: this.topK } : {}),
          ...(this.filter ? { filters: this.filter } : {}),
        }),
        headers: {
          "Content-Type": "application/json",
          ...(this.apiKey ? { Authorization: `Bearer ${this.apiKey}` } : {}),
        },
      }
    );

    const { results } = (await r.json()) as { results: Berry[] };

    return results.map(
      ({ text, score, source, ...rest }) =>
        new Document({
          pageContent: text,
          metadata: {
            score,
            source,
            ...rest,
          },
        })
    );
  }
}


import { ZepClient } from "@getzep/zep-cloud";
import {
  SearchScope,
  SearchType,
  MemorySearchResult,
  NotFoundError,
} from "@getzep/zep-cloud/api";
import { BaseRetriever, BaseRetrieverInput } from "@langchain/core/retrievers";
import { Document } from "@langchain/core/documents";

/**
 * Configuration interface for the ZepRetriever class. Extends the
 * BaseRetrieverInput interface.
 *
 * @argument {string} sessionId - The ID of the Zep session.
 * @argument {string} [apiKey] - The Zep Cloud Project Key.
 * @argument {number} [topK] - The number of results to return.
 * @argument [searchScope] [searchScope] - The scope of the search: "messages" or "summary".
 * @argument [searchType] [searchType] - The type of search to perform: "similarity" or "mmr".
 * @argument {number} [mmrLambda] - The lambda value for the MMR search.
 * @argument {Record<string, unknown>} [filter] - The metadata filter to apply to the search.
 */
export interface ZepCloudRetrieverConfig extends BaseRetrieverInput {
  sessionId: string;
  topK?: number;
  apiKey: string;
  searchScope?: SearchScope;
  searchType?: SearchType;
  mmrLambda?: number;
  filter?: Record<string, unknown>;
}

/**
 * Class for retrieving information from a Zep Cloud long-term memory store.
 * Extends the BaseRetriever class.
 * @example
 * ```typescript
 * const retriever = new ZepCloudRetriever({
 *   apiKey: "<zep cloud project api key>",
 *   sessionId: "session_exampleUUID",
 *   topK: 3,
 * });
 * const query = "Can I drive red cars in France?";
 * const docs = await retriever.getRelevantDocuments(query);
 * ```
 */
export class ZepCloudRetriever extends BaseRetriever {
  static lc_name() {
    return "ZepRetriever";
  }

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

  get lc_secrets(): { [key: string]: string } | undefined {
    return {
      apiKey: "ZEP_API_KEY",
    };
  }

  get lc_aliases(): { [key: string]: string } | undefined {
    return { apiKey: "api_key" };
  }

  client: ZepClient;

  private sessionId: string;

  private topK?: number;

  private searchScope?: SearchScope;

  private searchType?: SearchType;

  private mmrLambda?: number;

  private filter?: Record<string, unknown>;

  constructor(config: ZepCloudRetrieverConfig) {
    super(config);
    this.sessionId = config.sessionId;
    this.topK = config.topK;
    this.searchScope = config.searchScope;
    this.searchType = config.searchType;
    this.mmrLambda = config.mmrLambda;
    this.filter = config.filter;
    this.client = new ZepClient({ apiKey: config.apiKey });
  }

  /**
   *  Converts an array of message search results to an array of Document objects.
   *  @param {MemorySearchResult[]} results - The array of search results.
   *  @returns {Document[]} An array of Document objects representing the search results.
   */
  private searchMessageResultToDoc(results: MemorySearchResult[]): Document[] {
    return results
      .filter((r) => r.message)
      .map(
        ({
          message: { content, metadata: messageMetadata } = {},
          score,
          ...rest
        }) =>
          new Document({
            pageContent: content ?? "",
            metadata: { score, ...messageMetadata, ...rest },
          })
      );
  }

  /**
   *  Converts an array of summary search results to an array of Document objects.
   *  @param {MemorySearchResult[]} results - The array of search results.
   *  @returns {Document[]} An array of Document objects representing the search results.
   */
  private searchSummaryResultToDoc(results: MemorySearchResult[]): Document[] {
    return results
      .filter((r) => r.summary)
      .map(
        ({
          summary: { content, metadata: summaryMetadata } = {},
          score,
          ...rest
        }) =>
          new Document({
            pageContent: content ?? "",
            metadata: { score, ...summaryMetadata, ...rest },
          })
      );
  }

  /**
   *  Retrieves the relevant documents based on the given query.
   *  @param {string} query - The query string.
   *  @returns {Promise<Document[]>} A promise that resolves to an array of relevant Document objects.
   */
  async _getRelevantDocuments(query: string): Promise<Document[]> {
    try {
      const results: MemorySearchResult[] = await this.client.memory.search(
        this.sessionId,
        {
          text: query,
          metadata: this.filter,
          searchScope: this.searchScope,
          searchType: this.searchType,
          mmrLambda: this.mmrLambda,
          limit: this.topK,
        }
      );
      return this.searchScope === "summary"
        ? this.searchSummaryResultToDoc(results)
        : this.searchMessageResultToDoc(results);
    } catch (error) {
      // eslint-disable-next-line no-instanceof/no-instanceof
      if (error instanceof NotFoundError) {
        return Promise.resolve([]); // Return an empty Document array
      }
      // If it's not a NotFoundError, throw the error again
      throw error;
    }
  }
}


import {
  BaseRetriever,
  type BaseRetrieverInput,
} from "@langchain/core/retrievers";
import {
  AsyncCaller,
  type AsyncCallerParams,
} from "@langchain/core/utils/async_caller";
import type { DocumentInterface } from "@langchain/core/documents";

/**
 * Type for the authentication method used by the RemoteRetriever. It can
 * either be false (no authentication) or an object with a bearer token.
 */
export type RemoteRetrieverAuth = false | { bearer: string };

/**
 * Type for the JSON response values from the remote server.
 */
// eslint-disable-next-line @typescript-eslint/no-explicit-any
export type RemoteRetrieverValues = Record<string, any>;

/**
 * Interface for the parameters required to initialize a RemoteRetriever
 * instance.
 */
export interface RemoteRetrieverParams
  extends AsyncCallerParams,
    BaseRetrieverInput {
  /**
   * The URL of the remote retriever server
   */
  url: string;

  /**
   * The authentication method to use, currently implemented is
   * - false: no authentication
   * - { bearer: string }: Bearer token authentication
   */
  auth: RemoteRetrieverAuth;
}

/**
 * Abstract class for interacting with a remote server to retrieve
 * relevant documents based on a given query.
 */
export abstract class RemoteRetriever
  extends BaseRetriever
  implements RemoteRetrieverParams
{
  get lc_secrets(): { [key: string]: string } | undefined {
    return {
      "auth.bearer": "REMOTE_RETRIEVER_AUTH_BEARER",
    };
  }

  url: string;

  auth: RemoteRetrieverAuth;

  headers: Record<string, string>;

  asyncCaller: AsyncCaller;

  constructor(fields: RemoteRetrieverParams) {
    super(fields);
    const { url, auth, ...rest } = fields;
    this.url = url;
    this.auth = auth;
    this.headers = {
      Accept: "application/json",
      "Content-Type": "application/json",
      ...(this.auth && this.auth.bearer
        ? { Authorization: `Bearer ${this.auth.bearer}` }
        : {}),
    };
    this.asyncCaller = new AsyncCaller(rest);
  }

  /**
   * Abstract method that should be implemented by subclasses to create the
   * JSON body of the request based on the given query.
   * @param query The query based on which the JSON body of the request is created.
   * @returns The JSON body of the request.
   */
  abstract createJsonBody(query: string): RemoteRetrieverValues;

  /**
   * Abstract method that should be implemented by subclasses to process the
   * JSON response from the server and convert it into an array of Document
   * instances.
   * @param json The JSON response from the server.
   * @returns An array of Document instances.
   */
  abstract processJsonResponse(
    json: RemoteRetrieverValues
  ): DocumentInterface[];

  async _getRelevantDocuments(query: string): Promise<DocumentInterface[]> {
    const body = this.createJsonBody(query);
    const response = await this.asyncCaller.call(() =>
      fetch(this.url, {
        method: "POST",
        headers: this.headers,
        body: JSON.stringify(body),
      })
    );
    if (!response.ok) {
      throw new Error(
        `Failed to retrieve documents from ${this.url}: ${response.status} ${response.statusText}`
      );
    }
    const json = await response.json();
    return this.processJsonResponse(json);
  }
}


import { Zep, ZepClient } from "@getzep/zep-cloud";
import { Memory, NotFoundError } from "@getzep/zep-cloud/api";
import {
  InputValues,
  OutputValues,
  MemoryVariables,
  getInputValue,
  getOutputValue,
} from "@langchain/core/memory";
import {
  AIMessage,
  BaseMessage,
  ChatMessage,
  getBufferString,
  HumanMessage,
  SystemMessage,
} from "@langchain/core/messages";
import { BaseChatMemory, BaseChatMemoryInput } from "./chat_memory.js";

// Extract Summary and Facts from Zep memory, if present and compose a system prompt
export const zepMemoryContextToSystemPrompt = (memory: Memory) => {
  let systemPrompt = "";

  // Extract conversation facts, if present
  if (memory.facts) {
    systemPrompt += memory.facts.join("\n");
  }

  // Extract summary, if present
  if (memory.summary && memory.summary?.content) {
    systemPrompt += memory.summary.content;
  }

  return systemPrompt;
};

// We are condensing the Zep context into a human message in order to satisfy
// some models' input requirements and allow more flexibility for devs.
// (for example, Anthropic only supports one system message, and does not support multiple user messages in a row)
export const condenseZepMemoryIntoHumanMessage = (memory: Memory) => {
  const systemPrompt = zepMemoryContextToSystemPrompt(memory);

  let concatMessages = "";

  // Add message history to the prompt, if present
  if (memory.messages) {
    concatMessages = memory.messages
      .map((msg) => `${msg.role ?? msg.roleType}: ${msg.content}`)
      .join("\n");
  }

  return new HumanMessage(`${systemPrompt}\n${concatMessages}`);
};

// Convert Zep Memory to a list of BaseMessages
export const zepMemoryToMessages = (memory: Memory) => {
  const systemPrompt = zepMemoryContextToSystemPrompt(memory);

  let messages: BaseMessage[] = systemPrompt
    ? [new SystemMessage(systemPrompt)]
    : [];

  if (memory && memory.messages) {
    messages = messages.concat(
      memory.messages
        .filter((m) => m.content)
        .map((message) => {
          const { content, role, roleType } = message;
          const messageContent = content as string;
          if (roleType === "user") {
            return new HumanMessage(messageContent);
          } else if (role === "assistant") {
            return new AIMessage(messageContent);
          } else {
            // default to generic ChatMessage
            return new ChatMessage(
              messageContent,
              (roleType ?? role) as string
            );
          }
        })
    );
  }

  return messages;
};

/**
 * Interface defining the structure of the input data for the ZepMemory
 * class. It includes properties like humanPrefix, aiPrefix, memoryKey, memoryType
 * sessionId, and apiKey.
 */
export interface ZepCloudMemoryInput extends BaseChatMemoryInput {
  humanPrefix?: string;

  aiPrefix?: string;

  memoryKey?: string;

  sessionId: string;

  apiKey: string;

  memoryType?: Zep.MemoryType;

  // Whether to return separate messages for chat history with a SystemMessage containing (facts and summary) or return a single HumanMessage with the entire memory context.
  // Defaults to false (return a single HumanMessage) in order to allow more flexibility with different models.
  separateMessages?: boolean;
}

/**
 * Class used to manage the memory of a chat session, including loading
 * and saving the chat history, and clearing the memory when needed. It
 * uses the ZepClient to interact with the Zep service for managing the
 * chat session's memory.
 * @example
 * ```typescript
 * const sessionId = randomUUID();
 *
 * // Initialize ZepCloudMemory with session ID and API key
 * const memory = new ZepCloudMemory({
 *   sessionId,
 *   apiKey: "<zep api key>",
 * });
 *
 * // Create a ChatOpenAI model instance with specific parameters
 * const model = new ChatOpenAI({
 *   modelName: "gpt-3.5-turbo",
 *   temperature: 0,
 * });
 *
 * // Create a ConversationChain with the model and memory
 * const chain = new ConversationChain({ llm: model, memory });
 *
 * // Example of calling the chain with an input
 * const res1 = await chain.call({ input: "Hi! I'm Jim." });
 * console.log({ res1 });
 *
 * // Follow-up call to the chain to demonstrate memory usage
 * const res2 = await chain.call({ input: "What did I just say my name was?" });
 * console.log({ res2 });
 *
 * // Output the session ID and the current state of memory
 * console.log("Session ID: ", sessionId);
 * console.log("Memory: ", await memory.loadMemoryVariables({}));
 *
 * ```
 */
export class ZepCloudMemory
  extends BaseChatMemory
  implements ZepCloudMemoryInput
{
  humanPrefix = "Human";

  aiPrefix = "AI";

  memoryKey = "history";

  apiKey: string;

  sessionId: string;

  zepClient: ZepClient;

  memoryType: Zep.MemoryType;

  separateMessages: boolean;

  constructor(fields: ZepCloudMemoryInput) {
    super({
      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;
    this.apiKey = fields.apiKey;
    this.sessionId = fields.sessionId;
    this.memoryType = fields.memoryType ?? "perpetual";
    this.separateMessages = fields.separateMessages ?? false;
    this.zepClient = new ZepClient({
      apiKey: this.apiKey,
    });
  }

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

  /**
   * Method that retrieves the chat history from the Zep service and formats
   * it into a list of messages.
   * @param values Input values for the method.
   * @returns Promise that resolves with the chat history formatted into a list of messages.
   */
  async loadMemoryVariables(values: InputValues): Promise<MemoryVariables> {
    const memoryType = values.memoryType ?? "perpetual";
    let memory: Memory | null = null;
    try {
      memory = await this.zepClient.memory.get(this.sessionId, {
        memoryType,
      });
    } catch (error) {
      // eslint-disable-next-line no-instanceof/no-instanceof
      if (error instanceof NotFoundError) {
        return this.returnMessages
          ? { [this.memoryKey]: [] }
          : { [this.memoryKey]: "" };
      }
      throw error;
    }

    if (this.returnMessages) {
      return {
        [this.memoryKey]: this.separateMessages
          ? zepMemoryToMessages(memory)
          : [condenseZepMemoryIntoHumanMessage(memory)],
      };
    }
    return {
      [this.memoryKey]: this.separateMessages
        ? getBufferString(
            zepMemoryToMessages(memory),
            this.humanPrefix,
            this.aiPrefix
          )
        : condenseZepMemoryIntoHumanMessage(memory).content,
    };
  }

  /**
   * Method that saves the input and output messages to the Zep service.
   * @param inputValues Input messages to be saved.
   * @param outputValues Output messages to be saved.
   * @returns Promise that resolves when the messages have been saved.
   */
  async saveContext(
    inputValues: InputValues,
    outputValues: OutputValues
  ): Promise<void> {
    const input = getInputValue(inputValues, this.inputKey);
    const output = getOutputValue(outputValues, this.outputKey);

    // Add the new memory to the session using the ZepClient
    if (this.sessionId) {
      try {
        await this.zepClient.memory.add(this.sessionId, {
          messages: [
            {
              role: this.humanPrefix,
              roleType: "user",
              content: `${input}`,
            },
            {
              role: this.aiPrefix,
              roleType: "assistant",
              content: `${output}`,
            },
          ],
        });
      } catch (error) {
        console.error("Error adding memory: ", error);
      }
    }

    // Call the superclass's saveContext method
    await super.saveContext(inputValues, outputValues);
  }

  /**
   * Method that deletes the chat history from the Zep service.
   * @returns Promise that resolves when the chat history has been deleted.
   */
  async clear(): Promise<void> {
    try {
      await this.zepClient.memory.delete(this.sessionId);
    } catch (error) {
      console.error("Error deleting session: ", error);
    }

    // Clear the superclass's chat history
    await super.clear();
  }
}


import { test } from "@jest/globals";
import * as fs from "node:fs/promises";
import { fileURLToPath } from "node:url";
import * as path from "node:path";
import { AIMessage, HumanMessage } from "@langchain/core/messages";
import { PromptTemplate } from "@langchain/core/prompts";
import {
  BytesOutputParser,
  StringOutputParser,
} from "@langchain/core/output_parsers";
import { ChatOllama } from "../ollama.js";

test.skip("test call", async () => {
  const ollama = new ChatOllama({});
  // @eslint-disable-next-line/@typescript-eslint/ban-ts-comment
  // @ts-expect-error unused var
  const result = await ollama.invoke(
    "What is a good name for a company that makes colorful socks?"
  );
  // console.log({ result });
});

test.skip("test call with callback", async () => {
  const ollama = new ChatOllama({
    baseUrl: "http://localhost:11434",
  });
  const tokens: string[] = [];
  const result = await ollama.invoke(
    "What is a good name for a company that makes colorful socks?",
    {
      callbacks: [
        {
          handleLLMNewToken(token: string) {
            tokens.push(token);
          },
        },
      ],
    }
  );
  expect(tokens.length).toBeGreaterThan(1);
  expect(result).toEqual(tokens.join(""));
});

test.skip("test streaming call", async () => {
  const ollama = new ChatOllama({
    baseUrl: "http://localhost:11434",
  });
  const stream = await ollama.stream(
    `Translate "I love programming" into German.`
  );
  const chunks = [];
  for await (const chunk of stream) {
    chunks.push(chunk);
  }
  expect(chunks.length).toBeGreaterThan(1);
});

test.skip("should abort the request", async () => {
  const ollama = new ChatOllama({
    baseUrl: "http://localhost:11434",
  });
  const controller = new AbortController();

  await expect(() => {
    const ret = ollama.invoke("Respond with an extremely verbose response", {
      signal: controller.signal,
    });
    controller.abort();
    return ret;
  }).rejects.toThrow("This operation was aborted");
});

test.skip("Test multiple messages", async () => {
  const model = new ChatOllama({ baseUrl: "http://localhost:11434" });
  // @eslint-disable-next-line/@typescript-eslint/ban-ts-comment
  // @ts-expect-error unused var
  const res = await model.invoke([
    new HumanMessage({ content: "My name is Jonas" }),
  ]);
  // console.log({ res });
  // @eslint-disable-next-line/@typescript-eslint/ban-ts-comment
  // @ts-expect-error unused var
  const res2 = await model.invoke([
    new HumanMessage("My name is Jonas"),
    new AIMessage(
      "Hello Jonas! It's nice to meet you. Is there anything I can help you with?"
    ),
    new HumanMessage("What did I say my name was?"),
  ]);
  // console.log({ res2 });
});

test.skip("should stream through with a bytes output parser", async () => {
  const TEMPLATE = `You are a pirate named Patchy. All responses must be extremely verbose and in pirate dialect.

  User: {input}
  AI:`;

  // Infer the input variables from the template
  const prompt = PromptTemplate.fromTemplate(TEMPLATE);

  const ollama = new ChatOllama({
    model: "llama2",
    baseUrl: "http://127.0.0.1:11434",
  });
  const outputParser = new BytesOutputParser();
  const chain = prompt.pipe(ollama).pipe(outputParser);
  const stream = await chain.stream({
    input: `Translate "I love programming" into German.`,
  });
  const chunks = [];
  for await (const chunk of stream) {
    chunks.push(chunk);
  }
  // console.log(chunks.join(""));
  expect(chunks.length).toBeGreaterThan(1);
});

test.skip("JSON mode", async () => {
  const TEMPLATE = `You are a pirate named Patchy. All responses must be in pirate dialect and in JSON format, with a property named "response" followed by the value.

  User: {input}
  AI:`;

  // Infer the input variables from the template
  const prompt = PromptTemplate.fromTemplate(TEMPLATE);

  const ollama = new ChatOllama({
    model: "llama2",
    baseUrl: "http://127.0.0.1:11434",
    format: "json",
  });
  const outputParser = new StringOutputParser();
  const chain = prompt.pipe(ollama).pipe(outputParser);
  const res = await chain.invoke({
    input: `Translate "I love programming" into German.`,
  });
  expect(JSON.parse(res).response).toBeDefined();
});

test.skip("Test ChatOllama with an image", async () => {
  const __filename = fileURLToPath(import.meta.url);
  const __dirname = path.dirname(__filename);
  const imageData = await fs.readFile(path.join(__dirname, "/data/hotdog.jpg"));
  const chat = new ChatOllama({
    model: "llava",
    baseUrl: "http://127.0.0.1:11434",
  });
  // @eslint-disable-next-line/@typescript-eslint/ban-ts-comment
  // @ts-expect-error unused var
  const res = await chat.invoke([
    new HumanMessage({
      content: [
        {
          type: "text",
          text: "What is in this image?",
        },
        {
          type: "image_url",
          image_url: `data:image/jpeg;base64,${imageData.toString("base64")}`,
        },
      ],
    }),
  ]);
  // console.log({ res });
});

test.skip("test max tokens (numPredict)", async () => {
  const ollama = new ChatOllama({
    numPredict: 10,
  }).pipe(new StringOutputParser());
  const stream = await ollama.stream(
    "explain quantum physics to me in as many words as possible"
  );
  let numTokens = 0;
  let response = "";
  for await (const s of stream) {
    numTokens += 1;
    response += s;
  }

  // console.log({ numTokens, response });
  // Ollama doesn't always stream back the exact number of tokens, so we
  // check for a number which is slightly above the `numPredict`.
  expect(numTokens).toBeLessThanOrEqual(12);
});


import { Storage, File } from "@google-cloud/storage";

import { Document } from "@langchain/core/documents";
import { Docstore } from "langchain/stores/doc/base";

/**
 * Interface that defines the configuration for the
 * GoogleCloudStorageDocstore. It includes the bucket name and an optional
 * prefix.
 */
export interface GoogleCloudStorageDocstoreConfiguration {
  /** The identifier for the GCS bucket */
  bucket: string;

  /**
   * An optional prefix to prepend to each object name.
   * Often used to create a pseudo-hierarchy.
   */
  prefix?: string;
}

/**
 * Class that provides an interface for interacting with Google Cloud
 * Storage (GCS) as a document store. It extends the Docstore class and
 * implements methods to search, add, and add a document to the GCS
 * bucket.
 */
export class GoogleCloudStorageDocstore extends Docstore {
  bucket: string;

  prefix = "";

  storage: Storage;

  constructor(config: GoogleCloudStorageDocstoreConfiguration) {
    super();

    this.bucket = config.bucket;
    this.prefix = config.prefix ?? this.prefix;

    this.storage = new Storage();
  }

  /**
   * Searches for a document in the GCS bucket and returns it as a Document
   * instance.
   * @param search The name of the document to search for in the GCS bucket
   * @returns A Promise that resolves to a Document instance representing the found document
   */
  async search(search: string): Promise<Document> {
    const file = this.getFile(search);

    const [fileMetadata] = await file.getMetadata();
    const metadata = fileMetadata?.metadata;

    const [dataBuffer] = await file.download();
    const pageContent = dataBuffer.toString();

    const ret = new Document({
      pageContent,
      metadata,
    });

    return ret;
  }

  /**
   * Adds multiple documents to the GCS bucket.
   * @param texts An object where each key is the name of a document and the value is the Document instance to be added
   * @returns A Promise that resolves when all documents have been added
   */
  async add(texts: Record<string, Document>): Promise<void> {
    await Promise.all(
      Object.keys(texts).map((key) => this.addDocument(key, texts[key]))
    );
  }

  /**
   * Adds a single document to the GCS bucket.
   * @param name The name of the document to be added
   * @param document The Document instance to be added
   * @returns A Promise that resolves when the document has been added
   */
  async addDocument(name: string, document: Document): Promise<void> {
    const file = this.getFile(name);
    await file.save(document.pageContent);
    await file.setMetadata({ metadata: document.metadata });
  }

  /**
   * Gets a file from the GCS bucket.
   * @param name The name of the file to get from the GCS bucket
   * @returns A File instance representing the fetched file
   */
  private getFile(name: string): File {
    const filename = this.prefix + name;
    const file = this.storage.bucket(this.bucket).file(filename);
    return file;
  }
}


function isVersionLessThan(v1: number[], v2: number[]): boolean {
  for (let i = 0; i < Math.min(v1.length, v2.length); i += 1) {
    if (v1[i] < v2[i]) {
      return true;
    } else if (v1[i] > v2[i]) {
      return false;
    }
  }
  // If all the corresponding parts are equal, the shorter version is less
  return v1.length < v2.length;
}

// Filter utils

const COMPARISONS_TO_NATIVE: Record<string, string> = {
  $eq: "=",
  $ne: "<>",
  $lt: "<",
  $lte: "<=",
  $gt: ">",
  $gte: ">=",
};

const COMPARISONS_TO_NATIVE_OPERATORS = new Set(
  Object.keys(COMPARISONS_TO_NATIVE)
);

const TEXT_OPERATORS = new Set(["$like", "$ilike"]);

const LOGICAL_OPERATORS = new Set(["$and", "$or"]);

const SPECIAL_CASED_OPERATORS = new Set(["$in", "$nin", "$between"]);

const SUPPORTED_OPERATORS = new Set([
  ...COMPARISONS_TO_NATIVE_OPERATORS,
  ...TEXT_OPERATORS,
  ...LOGICAL_OPERATORS,
  ...SPECIAL_CASED_OPERATORS,
]);

const IS_IDENTIFIER_REGEX = /^[a-zA-Z_][a-zA-Z0-9_]*$/;

function combineQueries(
  inputQueries: [string, Record<string, Any>][],
  operator: string
): [string, Record<string, Any>] {
  let combinedQuery = "";
  const combinedParams: Record<string, Any> = {};
  const paramCounter: Record<string, number> = {};

  for (const [query, params] of inputQueries) {
    let newQuery = query;
    for (const [param, value] of Object.entries(params)) {
      if (param in paramCounter) {
        paramCounter[param] += 1;
      } else {
        paramCounter[param] = 1;
      }
      const newParamName = `${param}_${paramCounter[param]}`;

      newQuery = newQuery.replace(`$${param}`, `$${newParamName}`);
      combinedParams[newParamName] = value;
    }

    if (combinedQuery) {
      combinedQuery += ` ${operator} `;
    }
    combinedQuery += `(${newQuery})`;
  }

  return [combinedQuery, combinedParams];
}

function collectParams(
  inputData: [string, Record<string, string>][]
): [string[], Record<string, Any>] {
  const queryParts: string[] = [];
  const params: Record<string, Any> = {};

  for (const [queryPart, param] of inputData) {
    queryParts.push(queryPart);
    Object.assign(params, param);
  }

  return [queryParts, params];
}

function handleFieldFilter(
  field: string,
  value: Any,
  paramNumber = 1
): [string, Record<string, Any>] {
  if (typeof field !== "string") {
    throw new Error(
      `field should be a string but got: ${typeof field} with value: ${field}`
    );
  }

  if (field.startsWith("$")) {
    throw new Error(
      `Invalid filter condition. Expected a field but got an operator: ${field}`
    );
  }

  // Allow [a - zA - Z0 -9_], disallow $ for now until we support escape characters
  if (!IS_IDENTIFIER_REGEX.test(field)) {
    throw new Error(
      `Invalid field name: ${field}. Expected a valid identifier.`
    );
  }

  let operator: string;
  let filterValue: Any;

  if (typeof value === "object" && value !== null && !Array.isArray(value)) {
    const keys = Object.keys(value);

    if (keys.length !== 1) {
      throw new Error(`Invalid filter condition. Expected a value which is a dictionary
        with a single key that corresponds to an operator but got a dictionary
        with ${keys.length} keys. The first few keys are: ${keys
        .slice(0, 3)
        .join(", ")}
      `);
    }

    // eslint-disable-next-line prefer-destructuring
    operator = keys[0];
    filterValue = value[operator];

    if (!SUPPORTED_OPERATORS.has(operator)) {
      throw new Error(
        `Invalid operator: ${operator}. Expected one of ${SUPPORTED_OPERATORS}`
      );
    }
  } else {
    operator = "$eq";
    filterValue = value;
  }

  if (COMPARISONS_TO_NATIVE_OPERATORS.has(operator)) {
    const native = COMPARISONS_TO_NATIVE[operator];
    const querySnippet = `n.${field} ${native} $param_${paramNumber}`;
    const queryParam = { [`param_${paramNumber}`]: filterValue };

    return [querySnippet, queryParam];
  } else if (operator === "$between") {
    const [low, high] = filterValue;
    const querySnippet = `$param_${paramNumber}_low <= n.${field} <= $param_${paramNumber}_high`;
    const queryParam = {
      [`param_${paramNumber}_low`]: low,
      [`param_${paramNumber}_high`]: high,
    };

    return [querySnippet, queryParam];
  } else if (["$in", "$nin", "$like", "$ilike"].includes(operator)) {
    if (["$in", "$nin"].includes(operator)) {
      filterValue.forEach((val: Any) => {
        if (
          typeof val !== "string" &&
          typeof val !== "number" &&
          typeof val !== "boolean"
        ) {
          throw new Error(`Unsupported type: ${typeof val} for value: ${val}`);
        }
      });
    }

    if (operator === "$in") {
      const querySnippet = `n.${field} IN $param_${paramNumber}`;
      const queryParam = { [`param_${paramNumber}`]: filterValue };
      return [querySnippet, queryParam];
    } else if (operator === "$nin") {
      const querySnippet = `n.${field} NOT IN $param_${paramNumber}`;
      const queryParam = { [`param_${paramNumber}`]: filterValue };
      return [querySnippet, queryParam];
    } else if (operator === "$like") {
      const querySnippet = `n.${field} CONTAINS $param_${paramNumber}`;
      const queryParam = { [`param_${paramNumber}`]: filterValue.slice(0, -1) };
      return [querySnippet, queryParam];
    } else if (operator === "$ilike") {
      const querySnippet = `toLower(n.${field}) CONTAINS $param_${paramNumber}`;
      const queryParam = { [`param_${paramNumber}`]: filterValue.slice(0, -1) };
      return [querySnippet, queryParam];
    } else {
      throw new Error("Not Implemented");
    }
  } else {
    throw new Error("Not Implemented");
  }
}

function constructMetadataFilter(
  filter: Record<string, Any>
): [string, Record<string, Any>] {
  if (typeof filter !== "object" || filter === null) {
    throw new Error("Expected a dictionary representing the filter condition.");
  }

  const entries = Object.entries(filter);

  if (entries.length === 1) {
    const [key, value] = entries[0];

    if (key.startsWith("$")) {
      if (!["$and", "$or"].includes(key.toLowerCase())) {
        throw new Error(
          `Invalid filter condition. Expected $and or $or but got: ${key}`
        );
      }

      if (!Array.isArray(value)) {
        throw new Error(
          `Expected an array for logical conditions, but got ${typeof value} for value: ${value}`
        );
      }

      const operation = key.toLowerCase() === "$and" ? "AND" : "OR";
      const combinedQueries = combineQueries(
        value.map((v) => constructMetadataFilter(v)),
        operation
      );

      return combinedQueries;
    } else {
      return handleFieldFilter(key, value);
    }
  } else if (entries.length > 1) {
    for (const [key] of entries) {
      if (key.startsWith("$")) {
        throw new Error(
          `Invalid filter condition. Expected a field but got an operator: ${key}`
        );
      }
    }

    const and_multiple = collectParams(
      entries.map(([field, val], index) =>
        handleFieldFilter(field, val, index + 1)
      )
    );

    if (and_multiple.length >= 1) {
      return [and_multiple[0].join(" AND "), and_multiple[1]];
    } else {
      throw Error(
        "Invalid filter condition. Expected a dictionary but got an empty dictionary"
      );
    }
  } else {
    throw new Error("Filter condition contains no entries.");
  }
}


import * as uuid from "uuid";
import type {
  ChromaClient as ChromaClientT,
  Collection,
  ChromaClientParams,
  CollectionMetadata,
  Where,
} from "chromadb";

import type { EmbeddingsInterface } from "@langchain/core/embeddings";
import { VectorStore } from "@langchain/core/vectorstores";
import { Document } from "@langchain/core/documents";

type SharedChromaLibArgs = {
  numDimensions?: number;
  collectionName?: string;
  filter?: object;
  collectionMetadata?: CollectionMetadata;
  clientParams?: Omit<ChromaClientParams, "path">;
};

/**
 * Defines the arguments that can be passed to the `Chroma` class
 * constructor. It can either contain a `url` for the Chroma database, the
 * number of dimensions for the vectors (`numDimensions`), a
 * `collectionName` for the collection to be used in the database, and a
 * `filter` object; or it can contain an `index` which is an instance of
 * `ChromaClientT`, along with the `numDimensions`, `collectionName`, and
 * `filter`.
 */
export type ChromaLibArgs =
  | ({
      url?: string;
    } & SharedChromaLibArgs)
  | ({
      index?: ChromaClientT;
    } & SharedChromaLibArgs);

/**
 * Defines the parameters for the `delete` method in the `Chroma` class.
 * It can either contain an array of `ids` of the documents to be deleted
 * or a `filter` object to specify the documents to be deleted.
 */
export interface ChromaDeleteParams<T> {
  ids?: string[];
  filter?: T;
}

/**
 * Chroma vector store integration.
 *
 * Setup:
 * Install `@langchain/community` and `chromadb`.
 *
 * ```bash
 * npm install @langchain/community chromadb
 * ```
 *
 * ## [Constructor args](https://api.js.langchain.com/classes/langchain_community_vectorstores_chroma.Chroma.html#constructor)
 *
 *
 * Instantiate
 *
 * ```typescript
 * import { Chroma } from '@langchain/community/vectorstores/chroma';
 * // Or other embeddings
 * import { OpenAIEmbeddings } from '@langchain/openai';
 *
 * const embeddings = new OpenAIEmbeddings({
 *   model: "text-embedding-3-small",
 * })
 *
 * const vectorStore = new Chroma(
 *   embeddings,
 *   {
 *     collectionName: "foo",
 *     url: "http://localhost:8000", // URL of the Chroma server
 *   }
 * );
 * ```
 *
 *
 *

 *
 *
 * Add documents
 *
 * ```typescript
 * import type { Document } from '@langchain/core/documents';
 *
 * const document1 = { pageContent: "foo", metadata: { baz: "bar" } };
 * const document2 = { pageContent: "thud", metadata: { bar: "baz" } };
 * const document3 = { pageContent: "i will be deleted :(", metadata: {} };
 *
 * const documents: Document[] = [document1, document2, document3];
 * const ids = ["1", "2", "3"];
 * await vectorStore.addDocuments(documents, { ids });
 * ```
 *
 *
 *

 *
 *
 * Delete documents
 *
 * ```typescript
 * await vectorStore.delete({ ids: ["3"] });
 * ```
 *
 *
 *

 *
 *
 * Similarity search
 *
 * ```typescript
 * const results = await vectorStore.similaritySearch("thud", 1);
 * for (const doc of results) {
 *   console.log(`* ${doc.pageContent} [${JSON.stringify(doc.metadata, null)}]`);
 * }
 * // Output: * thud [{"baz":"bar"}]
 * ```
 *
 *
 *

 *
 *
 *
 * Similarity search with filter
 *
 * ```typescript
 * const resultsWithFilter = await vectorStore.similaritySearch("thud", 1, { baz: "bar" });
 *
 * for (const doc of resultsWithFilter) {
 *   console.log(`* ${doc.pageContent} [${JSON.stringify(doc.metadata, null)}]`);
 * }
 * // Output: * foo [{"baz":"bar"}]
 * ```
 *
 *
 *

 *
 *
 *
 * Similarity search with score
 *
 * ```typescript
 * const resultsWithScore = await vectorStore.similaritySearchWithScore("qux", 1);
 * for (const [doc, score] of resultsWithScore) {
 *   console.log(`* [SIM=${score.toFixed(6)}] ${doc.pageContent} [${JSON.stringify(doc.metadata, null)}]`);
 * }
 * // Output: * [SIM=0.000000] qux [{"bar":"baz","baz":"bar"}]
 * ```
 *
 *
 *

 *
 *
 * As a retriever
 *
 * ```typescript
 * const retriever = vectorStore.asRetriever({
 *   searchType: "mmr", // Leave blank for standard similarity search
 *   k: 1,
 * });
 * const resultAsRetriever = await retriever.invoke("thud");
 * console.log(resultAsRetriever);
 *
 * // Output: [Document({ metadata: { "baz":"bar" }, pageContent: "thud" })]
 * ```
 *
 *
 *

 */


export class Chroma extends VectorStore {
  declare FilterType: Where;

  index?: ChromaClientT;

  collection?: Collection;

  collectionName: string;

  collectionMetadata?: CollectionMetadata;

  numDimensions?: number;

  clientParams?: Omit<ChromaClientParams, "path">;

  url: string;

  filter?: object;

  _vectorstoreType(): string {
    return "chroma";
  }

  constructor(embeddings: EmbeddingsInterface, args: ChromaLibArgs) {
    super(embeddings, args);
    this.numDimensions = args.numDimensions;
    this.embeddings = embeddings;
    this.collectionName = ensureCollectionName(args.collectionName);
    this.collectionMetadata = args.collectionMetadata;
    this.clientParams = args.clientParams;
    if ("index" in args) {
      this.index = args.index;
    } else if ("url" in args) {
      this.url = args.url || "http://localhost:8000";
    }

    this.filter = args.filter;
  }

  /**
   * Adds documents to the Chroma database. The documents are first
   * converted to vectors using the `embeddings` instance, and then added to
   * the database.
   * @param documents An array of `Document` instances to be added to the database.
   * @param options Optional. An object containing an array of `ids` for the documents.
   * @returns A promise that resolves when the documents have been added to the database.
   */
  async addDocuments(documents: Document[], options?: { ids?: string[] }) {
    const texts = documents.map(({ pageContent }) => pageContent);
    return this.addVectors(
      await this.embeddings.embedDocuments(texts),
      documents,
      options
    );
  }

  /**
   * Ensures that a collection exists in the Chroma database. If the
   * collection does not exist, it is created.
   * @returns A promise that resolves with the `Collection` instance.
   */
  async ensureCollection(): Promise<Collection> {
    if (!this.collection) {
      if (!this.index) {
        const chromaClient = new (await Chroma.imports()).ChromaClient({
          path: this.url,
          ...(this.clientParams ?? {}),
        });
        this.index = chromaClient;
      }
      try {
        this.collection = await this.index.getOrCreateCollection({
          name: this.collectionName,
          ...(this.collectionMetadata && { metadata: this.collectionMetadata }),
        });
      } catch (err) {
        throw new Error(`Chroma getOrCreateCollection error: ${err}`);
      }
    }

    return this.collection;
  }

  /**
   * Adds vectors to the Chroma database. The vectors are associated with
   * the provided documents.
   * @param vectors An array of vectors to be added to the database.
   * @param documents An array of `Document` instances associated with the vectors.
   * @param options Optional. An object containing an array of `ids` for the vectors.
   * @returns A promise that resolves with an array of document IDs when the vectors have been added to the database.
   */
  async addVectors(
    vectors: number[][],
    documents: Document[],
    options?: { ids?: string[] }
  ) {
    if (vectors.length === 0) {
      return [];
    }
    if (this.numDimensions === undefined) {
      this.numDimensions = vectors[0].length;
    }
    if (vectors.length !== documents.length) {
      throw new Error(`Vectors and metadatas must have the same length`);
    }
    if (vectors[0].length !== this.numDimensions) {
      throw new Error(
        `Vectors must have the same length as the number of dimensions (${this.numDimensions})`
      );
    }

    const documentIds =
      options?.ids ?? Array.from({ length: vectors.length }, () => uuid.v1());
    const collection = await this.ensureCollection();

    const mappedMetadatas = documents.map(({ metadata }) => {
      let locFrom;
      let locTo;

      if (metadata?.loc) {
        if (metadata.loc.lines?.from !== undefined)
          locFrom = metadata.loc.lines.from;
        if (metadata.loc.lines?.to !== undefined) locTo = metadata.loc.lines.to;
      }

      const newMetadata: Document["metadata"] = {
        ...metadata,
        ...(locFrom !== undefined && { locFrom }),
        ...(locTo !== undefined && { locTo }),
      };

      if (newMetadata.loc) delete newMetadata.loc;

      return newMetadata;
    });

    await collection.upsert({
      ids: documentIds,
      embeddings: vectors,
      metadatas: mappedMetadatas,
      documents: documents.map(({ pageContent }) => pageContent),
    });
    return documentIds;
  }

  /**
   * Deletes documents from the Chroma database. The documents to be deleted
   * can be specified by providing an array of `ids` or a `filter` object.
   * @param params An object containing either an array of `ids` of the documents to be deleted or a `filter` object to specify the documents to be deleted.
   * @returns A promise that resolves when the specified documents have been deleted from the database.
   */
  async delete(params: ChromaDeleteParams<this["FilterType"]>): Promise<void> {
    const collection = await this.ensureCollection();
    if (Array.isArray(params.ids)) {
      await collection.delete({ ids: params.ids });
    } else if (params.filter) {
      await collection.delete({
        where: { ...params.filter },
      });
    } else {
      throw new Error(`You must provide one of "ids or "filter".`);
    }
  }

  /**
   * Searches for vectors in the Chroma database that are similar to the
   * provided query vector. The search can be filtered using the provided
   * `filter` object or the `filter` property of the `Chroma` instance.
   * @param query The query vector.
   * @param k The number of similar vectors to return.
   * @param filter Optional. A `filter` object to filter the search results.
   * @returns A promise that resolves with an array of tuples, each containing a `Document` instance and a similarity score.
   */
  async similaritySearchVectorWithScore(
    query: number[],
    k: number,
    filter?: this["FilterType"]
  ) {
    if (filter && this.filter) {
      throw new Error("cannot provide both `filter` and `this.filter`");
    }
    const _filter = filter ?? this.filter;

    const collection = await this.ensureCollection();

    // similaritySearchVectorWithScore supports one query vector at a time
    // chroma supports multiple query vectors at a time
    const result = await collection.query({
      queryEmbeddings: query,
      nResults: k,
      where: { ..._filter },
    });

    const { ids, distances, documents, metadatas } = result;
    if (!ids || !distances || !documents || !metadatas) {
      return [];
    }
    // get the result data from the first and only query vector
    const [firstIds] = ids;
    const [firstDistances] = distances;
    const [firstDocuments] = documents;
    const [firstMetadatas] = metadatas;

    const results: [Document, number][] = [];
    for (let i = 0; i < firstIds.length; i += 1) {
      let metadata: Document["metadata"] = firstMetadatas?.[i] ?? {};

      if (metadata.locFrom && metadata.locTo) {
        metadata = {
          ...metadata,
          loc: {
            lines: {
              from: metadata.locFrom,
              to: metadata.locTo,
            },
          },
        };

        delete metadata.locFrom;
        delete metadata.locTo;
      }

      results.push([
        new Document({
          pageContent: firstDocuments?.[i] ?? "",
          metadata,
        }),
        firstDistances[i],
      ]);
    }
    return results;
  }

  /**
   * Creates a new `Chroma` instance from an array of text strings. The text
   * strings are converted to `Document` instances and added to the Chroma
   * database.
   * @param texts An array of text strings.
   * @param metadatas An array of metadata objects or a single metadata object. If an array is provided, it must have the same length as the `texts` array.
   * @param embeddings An `Embeddings` instance used to generate embeddings for the documents.
   * @param dbConfig A `ChromaLibArgs` object containing the configuration for the Chroma database.
   * @returns A promise that resolves with a new `Chroma` instance.
   */


static async fromTexts(
    texts: string[],
    metadatas: object[] | object,
    embeddings: EmbeddingsInterface,
    dbConfig: ChromaLibArgs
  ): Promise<Chroma> {
    const docs: Document[] = [];
    for (let i = 0; i < texts.length; i += 1) {
      const metadata = Array.isArray(metadatas) ? metadatas[i] : metadatas;
      const newDoc = new Document({
        pageContent: texts[i],
        metadata,
      });
      docs.push(newDoc);
    }
    return this.fromDocuments(docs, embeddings, dbConfig);
  }

  /**
   * Creates a new `Chroma` instance from an array of `Document` instances.
   * The documents are added to the Chroma database.
   * @param docs An array of `Document` instances.
   * @param embeddings An `Embeddings` instance used to generate embeddings for the documents.
   * @param dbConfig A `ChromaLibArgs` object containing the configuration for the Chroma database.
   * @returns A promise that resolves with a new `Chroma` instance.
   */
  static async fromDocuments(
    docs: Document[],
    embeddings: EmbeddingsInterface,
    dbConfig: ChromaLibArgs
  ): Promise<Chroma> {
    const instance = new this(embeddings, dbConfig);
    await instance.addDocuments(docs);
    return instance;
  }

  /**
   * Creates a new `Chroma` instance from an existing collection in the
   * Chroma database.
   * @param embeddings An `Embeddings` instance used to generate embeddings for the documents.
   * @param dbConfig A `ChromaLibArgs` object containing the configuration for the Chroma database.
   * @returns A promise that resolves with a new `Chroma` instance.
   */
  static async fromExistingCollection(
    embeddings: EmbeddingsInterface,
    dbConfig: ChromaLibArgs
  ): Promise<Chroma> {
    const instance = new this(embeddings, dbConfig);
    await instance.ensureCollection();
    return instance;
  }

  /** @ignore */
  static async imports(): Promise<{
    ChromaClient: typeof ChromaClientT;
  }> {
    try {
      const { ChromaClient } = await import("chromadb");
      return { ChromaClient };
    } catch (e) {
      throw new Error(
        "Please install chromadb as a dependency with, e.g. `npm install -S chromadb`"
      );
    }
  }
}

/**
 * Generates a unique collection name if none is provided.
 */
function ensureCollectionName(collectionName?: string) {
  if (!collectionName) {
    return `langchain-${uuid.v4()}`;
  }
  return collectionName;
}


import type {
  HierarchicalNSW as HierarchicalNSWT,
  SpaceName,
} from "hnswlib-node";
import type { EmbeddingsInterface } from "@langchain/core/embeddings";
import { SaveableVectorStore } from "@langchain/core/vectorstores";
import { Document } from "@langchain/core/documents";
import { SynchronousInMemoryDocstore } from "../stores/doc/in_memory.js";

/**
 * Interface for the base configuration of HNSWLib. It includes the space
 * name and the number of dimensions.
 */
export interface HNSWLibBase {
  space: SpaceName;
  numDimensions?: number;
}

/**
 * Interface for the arguments that can be passed to the HNSWLib
 * constructor. It extends HNSWLibBase and includes properties for the
 * document store and HNSW index.
 */
export interface HNSWLibArgs extends HNSWLibBase {
  docstore?: SynchronousInMemoryDocstore;
  index?: HierarchicalNSWT;
}

/**
 * Class that implements a vector store using Hierarchical Navigable Small
 * World (HNSW) graphs. It extends the SaveableVectorStore class and
 * provides methods for adding documents and vectors, performing
 * similarity searches, and saving and loading the vector store.
 */


export class HNSWLib extends SaveableVectorStore {
  declare FilterType: (doc: Document) => boolean;

  _index?: HierarchicalNSWT;

  docstore: SynchronousInMemoryDocstore;

  args: HNSWLibBase;

  _vectorstoreType(): string {
    return "hnswlib";
  }

  constructor(embeddings: EmbeddingsInterface, args: HNSWLibArgs) {
    super(embeddings, args);
    this._index = args.index;
    this.args = args;
    this.embeddings = embeddings;
    this.docstore = args?.docstore ?? new SynchronousInMemoryDocstore();
  }

  /**
   * Method to add documents to the vector store. It first converts the
   * documents to vectors using the embeddings, then adds the vectors to the
   * vector store.
   * @param documents The documents to be added to the vector store.
   * @returns A Promise that resolves when the documents have been added.
   */
  async addDocuments(documents: Document[]): Promise<void> {
    const texts = documents.map(({ pageContent }) => pageContent);
    return this.addVectors(
      await this.embeddings.embedDocuments(texts),
      documents
    );
  }

  private static async getHierarchicalNSW(args: HNSWLibBase) {
    const { HierarchicalNSW } = await HNSWLib.imports();
    if (!args.space) {
      throw new Error("hnswlib-node requires a space argument");
    }
    if (args.numDimensions === undefined) {
      throw new Error("hnswlib-node requires a numDimensions argument");
    }
    return new HierarchicalNSW(args.space, args.numDimensions);
  }

  private async initIndex(vectors: number[][]) {
    if (!this._index) {
      if (this.args.numDimensions === undefined) {
        this.args.numDimensions = vectors[0].length;
      }
      this.index = await HNSWLib.getHierarchicalNSW(this.args);
    }
    if (!this.index.getCurrentCount()) {
      this.index.initIndex(vectors.length);
    }
  }

  public get index(): HierarchicalNSWT {
    if (!this._index) {
      throw new Error(
        "Vector store not initialised yet. Try calling `addTexts` first."
      );
    }
    return this._index;
  }

  private set index(index: HierarchicalNSWT) {
    this._index = index;
  }

  /**
   * Method to add vectors to the vector store. It first initializes the
   * index if it hasn't been initialized yet, then adds the vectors to the
   * index and the documents to the document store.
   * @param vectors The vectors to be added to the vector store.
   * @param documents The documents corresponding to the vectors.
   * @returns A Promise that resolves when the vectors and documents have been added.
   */
  async addVectors(vectors: number[][], documents: Document[]) {
    if (vectors.length === 0) {
      return;
    }
    await this.initIndex(vectors);

    // TODO here we could optionally normalise the vectors to unit length
    // so that dot product is equivalent to cosine similarity, like this
    // https://github.com/nmslib/hnswlib/issues/384#issuecomment-1155737730
    // While we only support OpenAI embeddings this isn't necessary
    if (vectors.length !== documents.length) {
      throw new Error(`Vectors and metadatas must have the same length`);
    }
    if (vectors[0].length !== this.args.numDimensions) {
      throw new Error(
        `Vectors must have the same length as the number of dimensions (${this.args.numDimensions})`
      );
    }
    const capacity = this.index.getMaxElements();
    const needed = this.index.getCurrentCount() + vectors.length;
    if (needed > capacity) {
      this.index.resizeIndex(needed);
    }
    const docstoreSize = this.index.getCurrentCount();
    const toSave: Record<string, Document> = {};
    for (let i = 0; i < vectors.length; i += 1) {
      this.index.addPoint(vectors[i], docstoreSize + i);
      toSave[docstoreSize + i] = documents[i];
    }
    this.docstore.add(toSave);
  }

  /**
   * Method to perform a similarity search in the vector store using a query
   * vector. It returns the k most similar documents along with their
   * similarity scores. An optional filter function can be provided to
   * filter the documents.
   * @param query The query vector.
   * @param k The number of most similar documents to return.
   * @param filter An optional filter function to filter the documents.
   * @returns A Promise that resolves to an array of tuples, where each tuple contains a document and its similarity score.
   */
  async similaritySearchVectorWithScore(
    query: number[],
    k: number,
    filter?: this["FilterType"]
  ) {
    if (this.args.numDimensions && !this._index) {
      await this.initIndex([[]]);
    }
    if (query.length !== this.args.numDimensions) {
      throw new Error(
        `Query vector must have the same length as the number of dimensions (${this.args.numDimensions})`
      );
    }
    if (k > this.index.getCurrentCount()) {
      const total = this.index.getCurrentCount();
      console.warn(
        `k (${k}) is greater than the number of elements in the index (${total}), setting k to ${total}`
      );
      // eslint-disable-next-line no-param-reassign
      k = total;
    }
    const filterFunction = (label: number): boolean => {
      if (!filter) {
        return true;
      }
      const document = this.docstore.search(String(label));
      // eslint-disable-next-line no-instanceof/no-instanceof
      if (typeof document !== "string") {
        return filter(document);
      }
      return false;
    };
    const result = this.index.searchKnn(
      query,
      k,
      filter ? filterFunction : undefined
    );
    return result.neighbors.map(
      (docIndex, resultIndex) =>
        [
          this.docstore.search(String(docIndex)),
          result.distances[resultIndex],
        ] as [Document, number]
    );
  }

  /**
   * Method to delete the vector store from a directory. It deletes the
   * hnswlib.index file, the docstore.json file, and the args.json file from
   * the directory.
   * @param params An object with a directory property that specifies the directory from which to delete the vector store.
   * @returns A Promise that resolves when the vector store has been deleted.
   */
  async delete(params: { directory: string }) {
    const fs = await import("node:fs/promises");
    const path = await import("node:path");
    try {
      await fs.access(path.join(params.directory, "hnswlib.index"));
    } catch (err) {
      throw new Error(
        `Directory ${params.directory} does not contain a hnswlib.index file.`
      );
    }

    await Promise.all([
      await fs.rm(path.join(params.directory, "hnswlib.index"), {
        force: true,
      }),
      await fs.rm(path.join(params.directory, "docstore.json"), {
        force: true,
      }),
      await fs.rm(path.join(params.directory, "args.json"), { force: true }),
    ]);
  }

  /**
   * Method to save the vector store to a directory. It saves the HNSW
   * index, the arguments, and the document store to the directory.
   * @param directory The directory to which to save the vector store.
   * @returns A Promise that resolves when the vector store has been saved.
   */
  async save(directory: string) {
    const fs = await import("node:fs/promises");
    const path = await import("node:path");
    await fs.mkdir(directory, { recursive: true });
    await Promise.all([
      this.index.writeIndex(path.join(directory, "hnswlib.index")),
      await fs.writeFile(
        path.join(directory, "args.json"),
        JSON.stringify(this.args)
      ),
      await fs.writeFile(
        path.join(directory, "docstore.json"),
        JSON.stringify(Array.from(this.docstore._docs.entries()))
      ),
    ]);
  }

  /**
   * Static method to load a vector store from a directory. It reads the
   * HNSW index, the arguments, and the document store from the directory,
   * then creates a new HNSWLib instance with these values.
   * @param directory The directory from which to load the vector store.
   * @param embeddings The embeddings to be used by the HNSWLib instance.
   * @returns A Promise that resolves to a new HNSWLib instance.
   */