chore: read endpoint from env vars

run-llama · Oct 25, 2024 · 5efeef5 · 5efeef5
1 parent a5dd21d
commit 5efeef5
Show file tree

Hide file tree

Showing 7 changed files with 408 additions and 404 deletions.
diff --git a/examples/azure-cosmosdb.ts b/examples/azure-cosmosdb.ts
@@ -1,7 +1,4 @@
-process.env.AZURE_OPENAI_ENDPOINT =
-  "https://cog-nadwwr3zcwntm.openai.azure.com";
-process.env.AZURE_DEPLOYMENT_NAME = "gpt-4o-mini";
-process.env.EMBEDDING_MODEL = "text-embedding-3-large";
+import "dotenv/config";
 
 import {
   DefaultAzureCredential,
@@ -38,25 +35,21 @@ import { AzureCosmosNoSqlIndexStore } from "llamaindex/storage/indexStore/AzureC
       deployment: process.env.EMBEDDING_MODEL,
     },
   });
-  const endpoint = "https://llamaindex.documents.azure.com:443/";
-  const docStore = AzureCosmosNoSqlDocumentStore.fromAadToken({
-    endpoint,
-  });
+  const docStore = AzureCosmosNoSqlDocumentStore.fromAadToken();
   console.log({ docStore });
-  const indexStore = AzureCosmosNoSqlIndexStore.fromAadToken({
-    endpoint,
-  });
+  const indexStore = AzureCosmosNoSqlIndexStore.fromAadToken();
   console.log({ indexStore });
-  const vectorStore = AzureCosmosDBNoSqlVectorStore.fromAadToken({
-    endpoint,
-  });
+  const vectorStore = AzureCosmosDBNoSqlVectorStore.fromAadToken();
   console.log({ vectorStore });
   const storageContext = await storageContextFromDefaults({
     docStore,
     indexStore,
-    // vectorStore,
+    vectorStores: {
+      ["TEXT"]: vectorStore,
+    },
   });
   console.log({ storageContext });
+
   const document = new Document({ text: "Test Text" });
   const index = await VectorStoreIndex.fromDocuments([document], {
     storageContext,

diff --git a/packages/llamaindex/src/storage/docStore/AzureCosmosNoSqlDocumentStore.ts b/packages/llamaindex/src/storage/docStore/AzureCosmosNoSqlDocumentStore.ts
@@ -21,18 +21,12 @@ interface AccountAndKeyOptions extends CosmosClientCommonOptions {
   key: string;
 }
 
-interface AadTokenOptions extends CosmosClientCommonOptions {
-  endpoint: string;
-}
-
 export class AzureCosmosNoSqlDocumentStore extends KVDocumentStore {
   constructor({ azureCosmosNoSqlKVStore, namespace }: DocumentStoreArgs) {
     super(azureCosmosNoSqlKVStore, namespace);
   }
 
-  static fromConnectionString(
-    options: ConnectionStringOptions,
-  ): AzureCosmosNoSqlDocumentStore {
+  static fromConnectionString(options: ConnectionStringOptions) {
     const { dbName = DEFAULT_DATABASE, containerName = DEFAULT_CONTAINER } =
       options;
 
@@ -45,9 +39,7 @@ export class AzureCosmosNoSqlDocumentStore extends KVDocumentStore {
     });
   }
 
-  static fromAccountAndKey(
-    options: AccountAndKeyOptions,
-  ): AzureCosmosNoSqlDocumentStore {
+  static fromAccountAndKey(options: AccountAndKeyOptions) {
     const { dbName = DEFAULT_DATABASE, containerName = DEFAULT_CONTAINER } =
       options;
 
@@ -60,9 +52,15 @@ export class AzureCosmosNoSqlDocumentStore extends KVDocumentStore {
     });
   }
 
-  static fromAadToken(options: AadTokenOptions): AzureCosmosNoSqlDocumentStore {
-    const { dbName = DEFAULT_DATABASE, containerName = DEFAULT_CONTAINER } =
-      options;
+  static fromAadToken(
+    options?: {
+      endpoint?: string;
+      containerName?: string;
+      dbName?: string;
+    } & CosmosClientCommonOptions,
+  ) {
+    const dbName = options?.dbName || DEFAULT_DATABASE;
+    const containerName = options?.containerName || DEFAULT_CONTAINER;
 
     const azureCosmosNoSqlKVStore =
       AzureCosmosNoSqlKVStore.fromAadToken(options);

diff --git a/packages/llamaindex/src/storage/docStore/utils.ts b/packages/llamaindex/src/storage/docStore/utils.ts
@@ -59,8 +59,6 @@ export function jsonToDoc<Data>(
   const dataDict = serializer.fromPersistence(docDict[DATA_KEY]) as any;
   let doc: BaseNode;
 
-  console.log({ docType, dataDict });
-
   if (docType === ObjectType.DOCUMENT) {
     doc = new Document({
       text: dataDict.text,

diff --git a/packages/llamaindex/src/storage/indexStore/AzureCosmosNoSqlIndexStore.ts b/packages/llamaindex/src/storage/indexStore/AzureCosmosNoSqlIndexStore.ts
@@ -31,9 +31,7 @@ export class AzureCosmosNoSqlIndexStore extends KVIndexStore {
   }
 
   // Static method for creating an instance using a connection string
-  static fromConnectionString(
-    options: ConnectionStringOptions,
-  ): AzureCosmosNoSqlIndexStore {
+  static fromConnectionString(options: ConnectionStringOptions) {
     const { dbName = DEFAULT_DATABASE, containerName = DEFAULT_CONTAINER } =
       options;
     const azureCosmosNoSqlKVStore =
@@ -46,9 +44,7 @@ export class AzureCosmosNoSqlIndexStore extends KVIndexStore {
   }
 
   // Static method for creating an instance using account and key
-  static fromAccountAndKey(
-    options: AccountAndKeyOptions,
-  ): AzureCosmosNoSqlIndexStore {
+  static fromAccountAndKey(options: AccountAndKeyOptions) {
     const { dbName = DEFAULT_DATABASE, containerName = DEFAULT_CONTAINER } =
       options;
 
@@ -62,9 +58,15 @@ export class AzureCosmosNoSqlIndexStore extends KVIndexStore {
   }
 
   // Static method for creating an instance using AAD token
-  static fromAadToken(options: AadTokenOptions): AzureCosmosNoSqlIndexStore {
-    const { dbName = DEFAULT_DATABASE, containerName = DEFAULT_CONTAINER } =
-      options;
+  static fromAadToken(
+    options?: {
+      endpoint?: string;
+      containerName?: string;
+      dbName?: string;
+    } & CosmosClientCommonOptions,
+  ) {
+    const dbName = options?.dbName || DEFAULT_DATABASE;
+    const containerName = options?.containerName || DEFAULT_CONTAINER;
 
     const azureCosmosNoSqlKVStore =
       AzureCosmosNoSqlKVStore.fromAadToken(options);

diff --git a/packages/llamaindex/src/storage/kvStore/AzureCosmosNoSqlKVStore.ts b/packages/llamaindex/src/storage/kvStore/AzureCosmosNoSqlKVStore.ts
@@ -5,9 +5,11 @@ import {
   type CosmosClientOptions,
 } from "@azure/cosmos";
 import { DefaultAzureCredential } from "@azure/identity";
+import { IndexStructType } from "@llamaindex/core/data-structs";
+import { getEnv } from "@llamaindex/env";
 import { BaseKVStore } from "./types.js";
 
-const USER_AGENT_PREFIX = "LlamaIndex-CDBNoSQL-VectorStore-JavaScript";
+const USER_AGENT_PREFIX = "LlamaIndex-CDBNoSQL-KVStore-JavaScript";
 const DEFAULT_CHAT_DATABASE = "KVStoreDB";
 const DEFAULT_CHAT_CONTAINER = "KVStoreContainer";
 const DEFAULT_OFFER_THROUGHPUT = 400;
@@ -71,6 +73,10 @@ export class AzureCosmosNoSqlKVStore extends BaseKVStore {
     this.cosmosDatabaseProperties = cosmosDatabaseProperties;
   }
 
+  client(): CosmosClient {
+    return this.cosmosClient;
+  }
+
   // Asynchronous initialization method to create database and container
   async init(): Promise<void> {
     // Set default throughput if not provided
@@ -132,10 +138,22 @@ export class AzureCosmosNoSqlKVStore extends BaseKVStore {
   }
 
   static fromAadToken(
-    options: {
-      endpoint: string;
+    options?: {
+      endpoint?: string;
     } & CosmosClientCommonOptions,
   ): AzureCosmosNoSqlKVStore {
+    if (!options) {
+      options = {
+        endpoint: getEnv("AZURE_COSMOSDB_NOSQL_ENDPOINT") ?? "",
+      };
+    }
+
+    if (!options.endpoint) {
+      throw new Error(
+        "AZURE_COSMOSDB_NOSQL_ENDPOINT is required for AzureCosmosNoSqlKVStore",
+      );
+    }
+
     const aadCredentials = new DefaultAzureCredential();
     const cosmosClient = new CosmosClient({
       ...options,
@@ -180,6 +198,10 @@ export class AzureCosmosNoSqlKVStore extends BaseKVStore {
       .fetchAll();
     const output: Record<string, Record<string, any>> = {};
     resources.forEach((item) => {
+      item = {
+        ...item,
+        type: IndexStructType.LIST, // TODO: how can we automatically determine this?
+      };
       output[item.id] = item;
     });
     return output;

diff --git a/packages/llamaindex/src/vector-store/AzureCosmosDBNoSqlVectorStore.ts b/packages/llamaindex/src/vector-store/AzureCosmosDBNoSqlVectorStore.ts
@@ -5,15 +5,13 @@ import {
   VectorEmbeddingDistanceFunction,
   VectorIndexType,
   type ContainerRequest,
-  type CosmosClientOptions,
   type DatabaseRequest,
   type IndexingPolicy,
   type VectorEmbeddingPolicy,
   type VectorIndex,
 } from "@azure/cosmos";
-import { DefaultAzureCredential, type TokenCredential } from "@azure/identity";
+import { type TokenCredential } from "@azure/identity";
 import { BaseNode, MetadataMode } from "@llamaindex/core/schema";
-import { getEnv } from "@llamaindex/env";
 import { metadataDictToNode, nodeToMetadata } from "./utils.js";
 
 import {
@@ -64,8 +62,6 @@ export interface AzureCosmosDBNoSQLConfig
   readonly idKey?: string;
 }
 
-const USER_AGENT_PREFIX = "LlamaIndex-CDBNoSQL-VectorStore-JavaScript";
-
 const DEFAULT_VECTOR_EMBEDDING_POLICY = {
   vectorEmbeddings: [
     {
@@ -81,32 +77,6 @@ const DEFAULT_VECTOR_INDEXING_POLICY: VectorIndex[] = [
   { path: "/embedding", type: VectorIndexType.QuantizedFlat },
 ];
 
-function parseConnectionString(connectionString: string): {
-  endpoint: string;
-  key: string;
-} {
-  const parts = connectionString.split(";");
-  let endpoint = "";
-  let accountKey = "";
-
-  parts.forEach((part) => {
-    const [key, value] = part.split("=");
-    if (key && key.trim() === "AccountEndpoint") {
-      endpoint = value?.trim() ?? "";
-    } else if ((key ?? "").trim() === "AccountKey") {
-      accountKey = value?.trim() ?? "";
-    }
-  });
-
-  if (!endpoint || !accountKey) {
-    throw new Error(
-      "Invalid connection string: missing AccountEndpoint or AccountKey.",
-    );
-  }
-
-  return { endpoint, key: accountKey };
-}
-
 interface AadTokenOptions extends CosmosClientCommonOptions {
   endpoint: string;
 }
@@ -154,60 +124,32 @@ export class AzureCosmosDBNoSqlVectorStore extends BaseVectorStore {
     return this.cosmosClient;
   }
 
-  static fromAadToken(options: AadTokenOptions): AzureCosmosDBNoSqlVectorStore {
-    const { dbName, containerName } = options;
-
+  static fromAadToken(options?: AadTokenOptions) {
     const azureCosmosNoSqlKVStore =
       AzureCosmosNoSqlKVStore.fromAadToken(options);
-    const namespace = `${dbName}.${containerName}`;
     return new AzureCosmosDBNoSqlVectorStore(azureCosmosNoSqlKVStore, options);
   }
+
+  kvStore: AzureCosmosNoSqlKVStore;
+
   constructor(
     azureCosmosNoSqlKVStore: AzureCosmosNoSqlKVStore,
-    dbConfig: AzureCosmosDBNoSQLConfig,
+    dbConfig?: AzureCosmosDBNoSQLConfig,
     embedModel?: VectorStoreBaseParams,
   ) {
     super(embedModel);
-    const connectionString =
-      dbConfig.connectionString ??
-      getEnv("AZURE_COSMOSDB_NOSQL_CONNECTION_STRING");
-
-    const endpoint =
-      dbConfig.endpoint ?? getEnv("AZURE_COSMOSDB_NOSQL_ENDPOINT");
-
-    if (!dbConfig.client && !connectionString && !endpoint) {
-      throw new Error(
-        "AzureCosmosDBNoSQLVectorStore client, connection string or endpoint must be set.",
-      );
-    }
-
-    if (!dbConfig.client) {
-      if (connectionString) {
-        const { endpoint, key } = parseConnectionString(connectionString);
-        this.cosmosClient = new CosmosClient({
-          endpoint,
-          key,
-          userAgentSuffix: USER_AGENT_PREFIX,
-        } as CosmosClientOptions);
-      } else {
-        // Use managed identity
-        this.cosmosClient = new CosmosClient({
-          endpoint,
-          aadCredentials: dbConfig.credentials ?? new DefaultAzureCredential(),
-          userAgentSuffix: USER_AGENT_PREFIX,
-        } as CosmosClientOptions);
-      }
-    } else {
-      this.cosmosClient = dbConfig.client;
-    }
 
-    const client = this.cosmosClient;
-    const databaseName = dbConfig.databaseName ?? "vectorSearchDB";
-    const containerName = dbConfig.containerName ?? "vectorSearchContainer";
+    dbConfig = dbConfig ?? {};
+    this.kvStore = azureCosmosNoSqlKVStore;
+    this.cosmosClient = this.kvStore.client();
     this.idKey = dbConfig.idKey ?? "id";
     this.textKey = dbConfig.textKey ?? "text";
     this.flatMetadata = dbConfig.flatMetadata ?? true;
     this.metadataKey = dbConfig.metadataKey ?? "metadata";
+
+    const client = this.cosmosClient;
+    const databaseName = dbConfig.databaseName ?? "vectorSearchDB";
+    const containerName = dbConfig.containerName ?? "vectorSearchContainer";
     const vectorEmbeddingPolicy =
       dbConfig.vectorEmbeddingPolicy ?? DEFAULT_VECTOR_EMBEDDING_POLICY;
     const indexingPolicy: IndexingPolicy = dbConfig.indexingPolicy ?? {