add code documentation (thanks @EmanuelCampos)

.changeset/famous-singers-shake.md

@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+add new OpenAI embeddings (with dimension reduction support)

examples/vectorIndexEmbed3.ts

@@ -0,0 +1,41 @@
+import fs from "node:fs/promises";
+
+import {
+  Document,
+  OpenAIEmbedding,
+  VectorStoreIndex,
+  serviceContextFromDefaults,
+} from "llamaindex";
+
+async function main() {
+  // Load essay from abramov.txt in Node
+  const path = "node_modules/llamaindex/examples/abramov.txt";
+
+  const essay = await fs.readFile(path, "utf-8");
+
+  // Create Document object with essay
+  const document = new Document({ text: essay, id_: path });
+
+  // Create service context and specify text-embedding-3-large
+  const embedModel = new OpenAIEmbedding({
+    model: "text-embedding-3-large",
+    dimensions: 1024,
+  });
+  const serviceContext = serviceContextFromDefaults({ embedModel });
+
+  // Split text and create embeddings. Store them in a VectorStoreIndex
+  const index = await VectorStoreIndex.fromDocuments([document], {
+    serviceContext,
+  });
+
+  // Query the index
+  const queryEngine = index.asQueryEngine();
+  const response = await queryEngine.query({
+    query: "What did the author do in college?",
+  });
+
+  // Output response
+  console.log(response.toString());
+}
+
+main().catch(console.error);

package.json

@@ -18,7 +18,7 @@
   },
   "devDependencies": {
     "@changesets/cli": "^2.27.1",
-    "@turbo/gen": "^1.11.2",
+    "@turbo/gen": "^1.11.3",
     "@types/jest": "^29.5.11",
     "eslint": "^8.56.0",
     "eslint-config-custom": "workspace:*",
@@ -27,8 +27,8 @@
     "lint-staged": "^15.2.0",
     "prettier": "^3.2.4",
     "prettier-plugin-organize-imports": "^3.2.4",
-    "ts-jest": "^29.1.1",
-    "turbo": "^1.11.2",
+    "ts-jest": "^29.1.2",
+    "turbo": "^1.11.3",
     "typescript": "^5.3.3"
   },
   "packageManager": "pnpm@8.10.5+sha256.a4bd9bb7b48214bbfcd95f264bd75bb70d100e5d4b58808f5cd6ab40c6ac21c5",

packages/core/package.json

@@ -19,7 +19,7 @@
     "md-utils-ts": "^2.0.0",
     "mongodb": "^6.3.0",
     "notion-md-crawler": "^0.0.2",
-    "openai": "^4.20.1",
+    "openai": "^4.26.0",
     "papaparse": "^5.4.1",
     "pathe": "^1.1.2",
     "pdfjs-dist": "4.0.269",

packages/core/src/constants.ts

@@ -6,6 +6,4 @@ export const DEFAULT_CHUNK_OVERLAP = 20;
 export const DEFAULT_CHUNK_OVERLAP_RATIO = 0.1;
 export const DEFAULT_SIMILARITY_TOP_K = 2;
 
-// NOTE: for text-embedding-ada-002
-export const DEFAULT_EMBEDDING_DIM = 1536;
 export const DEFAULT_PADDING = 5;

packages/core/src/embeddings/OpenAIEmbedding.ts

@@ -9,28 +9,55 @@ import {
 import { OpenAISession, getOpenAISession } from "../llm/openai";
 import { BaseEmbedding } from "./types";
 
-export enum OpenAIEmbeddingModelType {
-  TEXT_EMBED_ADA_002 = "text-embedding-ada-002",
-}
+export const ALL_OPENAI_EMBEDDING_MODELS = {
+  "text-embedding-ada-002": {
+    dimensions: 1536,
+    maxTokens: 8191,
+  },
+  "text-embedding-3-small": {
+    dimensions: 1536,
+    dimensionOptions: [512, 1536],
+    maxTokens: 8191,
+  },
+  "text-embedding-3-large": {
+    dimensions: 3072,
+    dimensionOptions: [256, 1024, 3072],
+    maxTokens: 8191,
+  },
+};
 
 export class OpenAIEmbedding extends BaseEmbedding {
-  model: OpenAIEmbeddingModelType | string;
+  /** embeddding model. defaults to "text-embedding-ada-002" */
+  model: string;
+  /** number of dimensions of the resulting vector, for models that support choosing fewer dimensions. undefined will default to model default */
+  dimensions: number | undefined;
 
   // OpenAI session params
+
+  /** api key */
   apiKey?: string = undefined;
+  /** maximum number of retries, default 10 */
   maxRetries: number;
+  /** timeout in ms, default 60 seconds  */
   timeout?: number;
+  /** other session options for OpenAI */
   additionalSessionOptions?: Omit<
     Partial<OpenAIClientOptions>,
     "apiKey" | "maxRetries" | "timeout"
   >;
 
+  /** session object */
   session: OpenAISession;
 
+  /**
+   * OpenAI Embedding
+   * @param init - initial parameters
+   */
   constructor(init?: Partial<OpenAIEmbedding> & { azure?: AzureOpenAIConfig }) {
     super();
 
-    this.model = init?.model ?? OpenAIEmbeddingModelType.TEXT_EMBED_ADA_002;
+    this.model = init?.model ?? "text-embedding-ada-002";
+    this.dimensions = init?.dimensions; // if no dimensions provided, will be undefined/not sent to OpenAI
 
     this.maxRetries = init?.maxRetries ?? 10;
     this.timeout = init?.timeout ?? 60 * 1000; // Default is 60 seconds
@@ -76,6 +103,7 @@ export class OpenAIEmbedding extends BaseEmbedding {
   private async getOpenAIEmbedding(input: string) {
     const { data } = await this.session.openai.embeddings.create({
       model: this.model,
+      dimensions: this.dimensions, // only sent to OpenAI if set by user
       input,
     });
 

packages/core/src/llm/LLM.ts

@@ -41,14 +41,21 @@ import {
 export const GPT4_MODELS = {
   "gpt-4": { contextWindow: 8192 },
   "gpt-4-32k": { contextWindow: 32768 },
+  "gpt-4-32k-0613": { contextWindow: 32768 },
+  "gpt-4-turbo-preview": { contextWindow: 128000 },
   "gpt-4-1106-preview": { contextWindow: 128000 },
-  "gpt-4-vision-preview": { contextWindow: 8192 },
+  "gpt-4-0125-preview": { contextWindow: 128000 },
+  "gpt-4-vision-preview": { contextWindow: 128000 },
 };
 
+// NOTE we don't currently support gpt-3.5-turbo-instruct and don't plan to in the near future
 export const GPT35_MODELS = {
   "gpt-3.5-turbo": { contextWindow: 4096 },
+  "gpt-3.5-turbo-0613": { contextWindow: 4096 },
   "gpt-3.5-turbo-16k": { contextWindow: 16384 },
+  "gpt-3.5-turbo-16k-0613": { contextWindow: 16384 },
   "gpt-3.5-turbo-1106": { contextWindow: 16384 },
+  "gpt-3.5-turbo-0125": { contextWindow: 16384 },
 };
 
 /**

packages/core/src/llm/azure.ts

@@ -17,6 +17,14 @@ const ALL_AZURE_OPENAI_CHAT_MODELS = {
   },
   "gpt-4": { contextWindow: 8192, openAIModel: "gpt-4" },
   "gpt-4-32k": { contextWindow: 32768, openAIModel: "gpt-4-32k" },
+  "gpt-4-vision-preview": {
+    contextWindow: 128000,
+    openAIModel: "gpt-4-vision-preview",
+  },
+  "gpt-4-1106-preview": {
+    contextWindow: 128000,
+    openAIModel: "gpt-4-1106-preview",
+  },
 };
 
 const ALL_AZURE_OPENAI_EMBEDDING_MODELS = {
@@ -25,13 +33,29 @@ const ALL_AZURE_OPENAI_EMBEDDING_MODELS = {
     openAIModel: "text-embedding-ada-002",
     maxTokens: 8191,
   },
+  "text-embedding-3-small": {
+    dimensions: 1536,
+    dimensionOptions: [512, 1536],
+    openAIModel: "text-embedding-3-small",
+    maxTokens: 8191,
+  },
+  "text-embedding-3-large": {
+    dimensions: 3072,
+    dimensionOptions: [256, 1024, 3072],
+    openAIModel: "text-embedding-3-large",
+    maxTokens: 8191,
+  },
 };
 
 const ALL_AZURE_API_VERSIONS = [
   "2022-12-01",
   "2023-05-15",
-  "2023-06-01-preview",
-  "2023-07-01-preview",
+  "2023-03-15-preview", // retiring 2024-04-02
+  "2023-06-01-preview", // retiring 2024-04-02
+  "2023-07-01-preview", // retiring 2024-04-02
+  "2023-08-01-preview", // retiring 2024-04-02
+  "2023-09-01-preview",
+  "2023-12-01-preview",
 ];
 
 const DEFAULT_API_VERSION = "2023-05-15";

packages/core/src/storage/vectorStore/PGVectorStore.ts

@@ -20,6 +20,7 @@ export class PGVectorStore implements VectorStore {
   private schemaName: string = PGVECTOR_SCHEMA;
   private tableName: string = PGVECTOR_TABLE;
   private connectionString: string | undefined = undefined;
+  private dimensions: number = 1536;
 
   private db?: pg.Client;
 
@@ -38,15 +39,18 @@ export class PGVectorStore implements VectorStore {
    * @param {string} config.schemaName - The name of the schema (optional). Defaults to PGVECTOR_SCHEMA.
    * @param {string} config.tableName - The name of the table (optional). Defaults to PGVECTOR_TABLE.
    * @param {string} config.connectionString - The connection string (optional).
+   * @param {number} config.dimensions - The dimensions of the embedding model.
    */
   constructor(config?: {
     schemaName?: string;
     tableName?: string;
     connectionString?: string;
+    dimensions?: number;
   }) {
     this.schemaName = config?.schemaName ?? PGVECTOR_SCHEMA;
     this.tableName = config?.tableName ?? PGVECTOR_TABLE;
     this.connectionString = config?.connectionString;
+    this.dimensions = config?.dimensions ?? 1536;
   }
 
   /**
@@ -108,7 +112,7 @@ export class PGVectorStore implements VectorStore {
       collection VARCHAR,
       document TEXT,
       metadata JSONB DEFAULT '{}',
-      embeddings VECTOR(1536)
+      embeddings VECTOR(${this.dimensions})
     )`;
     await db.query(tbl);