tests: add tests for factor and config

README.md

@@ -1 +1 @@
-# SOON
+# SOON

index.ts

@@ -1 +1 @@
-//todo
+export * from "./src/rag-chat";

package.json

@@ -1,10 +1,9 @@
 {
   "name": "@upstash/rag-chat",
-  "version": "1.0.0",
+  "version": "0.0.6-alpha",
   "main": "./dist/index.js",
   "module": "./dist/index.mjs",
   "types": "./dist/index.d.ts",
-  "type": "module",
   "scripts": {
     "test": "bun test",
     "fmt": "prettier --write .",
@@ -52,7 +51,7 @@
     "@langchain/community": "^0.0.50",
     "@langchain/core": "^0.1.58",
     "@langchain/openai": "^0.0.28",
-    "@upstash/sdk": "0.0.19-alpha",
+    "@upstash/sdk": "0.0.23-alpha",
     "ai": "^3.0.35"
   }
 }

src/client-factory.test.ts

@@ -0,0 +1,41 @@
+/* eslint-disable @typescript-eslint/no-non-null-assertion */
+import { Upstash } from "@upstash/sdk";
+import { describe, expect, test } from "bun:test";
+import type { ClientFactoryConfig } from "./client-factory";
+import { ClientFactory } from "./client-factory";
+
+describe("ClientFactory Tests", () => {
+  const config: ClientFactoryConfig = {
+    email: process.env.UPSTASH_EMAIL!,
+    token: process.env.UPSTASH_TOKEN!,
+    region: "us-east-1",
+    redis: "test-rag-chat-client-factor-redis",
+    vector: "test-rag-chat-client-factor-vector",
+  };
+  const upstash = new Upstash(config);
+  const clientFactory = new ClientFactory(config);
+
+  test(
+    "Redis client initialization",
+    async () => {
+      const initialized = await clientFactory.init({ redis: true });
+      expect(initialized).toHaveProperty("redis");
+      expect(initialized.vector).toBeUndefined();
+
+      await upstash.deleteRedisDatabase("test-rag-chat-client-factor-redis");
+    },
+    { timeout: 20_000 }
+  );
+
+  test(
+    "Vector client initialization",
+    async () => {
+      const initialized = await clientFactory.init({ vector: true });
+      expect(initialized).toHaveProperty("vector");
+      expect(initialized.redis).toBeUndefined();
+
+      await upstash.deleteVectorIndex("test-rag-chat-client-factor-vector");
+    },
+    { timeout: 20_000 }
+  );
+});

src/client-factory.ts

@@ -0,0 +1,69 @@
+import type { Index, Redis } from "@upstash/sdk";
+import type { Config } from "./config";
+
+import { Upstash } from "@upstash/sdk";
+import { RedisClient } from "./clients/redis";
+import { VectorClient } from "./clients/vector";
+import { InternalUpstashError } from "./error/internal";
+
+export type ClientFactoryConfig = Pick<Config, "email" | "token" | "vector" | "redis" | "region">;
+export class ClientFactory {
+  private upstashSDK: Upstash;
+  private config: ClientFactoryConfig;
+
+  constructor(config: ClientFactoryConfig) {
+    this.upstashSDK = new Upstash(config);
+    this.config = config;
+  }
+  private createVectorClient(): Promise<Index | undefined> {
+    return new VectorClient({
+      upstashSDK: this.upstashSDK,
+      indexNameOrInstance: this.config.vector,
+      region: this.config.region,
+    }).getVectorClient();
+  }
+
+  private createRedisClient(): Promise<Redis | undefined> {
+    return new RedisClient({
+      upstashSDK: this.upstashSDK,
+      redisDbNameOrInstance: this.config.redis,
+      region: this.config.region,
+    }).getRedisClient();
+  }
+
+  /** Allows initting only desired clients. Reason is that in some cases like `HistoryService` we only require
+   * redis client, but in `RAGChat` we need all three clients to work efficiently.
+   *
+   * @default redis: false
+   * @default vector: false
+   * @default ratelimit: false
+   */
+  async init<TInit extends { redis?: boolean; vector?: boolean }>(
+    options: TInit = { redis: false, vector: false } as TInit
+  ): Promise<{
+    redis: TInit["redis"] extends true ? Redis : undefined;
+    vector: TInit["vector"] extends true ? Index : undefined;
+  }> {
+    let redis: Redis | undefined;
+    let vector: Index | undefined;
+
+    if (options.redis) {
+      redis = await this.createRedisClient();
+      if (!redis) {
+        throw new InternalUpstashError("Couldn't initialize Redis client");
+      }
+    }
+
+    if (options.vector) {
+      vector = await this.createVectorClient();
+      if (!vector) {
+        throw new InternalUpstashError("Couldn't initialize Vector client");
+      }
+    }
+
+    return { redis, vector } as {
+      redis: TInit["redis"] extends true ? Redis : undefined;
+      vector: TInit["vector"] extends true ? Index : undefined;
+    };
+  }
+}

src/clients/redis/index.test.ts

@@ -1,7 +1,7 @@
 /* eslint-disable @typescript-eslint/no-non-null-assertion */
 import { Upstash } from "@upstash/sdk";
 import { describe, expect, test } from "bun:test";
-import { DEFAULT_REDIS_CONFIG, DEFAULT_REDIS_DB_NAME, RedisClientConstructor } from ".";
+import { DEFAULT_REDIS_CONFIG, DEFAULT_REDIS_DB_NAME, RedisClient } from ".";
 
 const upstashSDK = new Upstash({
   email: process.env.UPSTASH_EMAIL!,
@@ -12,23 +12,23 @@ describe("Redis Client", () => {
   test(
     "Initialize client without db name",
     async () => {
-      const constructor = new RedisClientConstructor({
-        sdkClient: upstashSDK,
+      const constructor = new RedisClient({
+        upstashSDK: upstashSDK,
       });
       const redisClient = await constructor.getRedisClient();
 
       expect(redisClient).toBeTruthy();
 
       await upstashSDK.deleteRedisDatabase(DEFAULT_REDIS_DB_NAME);
     },
-    { timeout: 10_000 }
+    { timeout: 20_000 }
   );
 
   test(
     "Initialize client with db name",
     async () => {
-      const constructor = new RedisClientConstructor({
-        sdkClient: upstashSDK,
+      const constructor = new RedisClient({
+        upstashSDK: upstashSDK,
         redisDbNameOrInstance: "test-name",
       });
       const redisClient = await constructor.getRedisClient();
@@ -37,7 +37,7 @@ describe("Redis Client", () => {
 
       await upstashSDK.deleteRedisDatabase("test-name");
     },
-    { timeout: 10_000 }
+    { timeout: 20_000 }
   );
 
   test(
@@ -50,8 +50,8 @@ describe("Redis Client", () => {
       });
       const existingRedisClient = await upstashSDK.newRedisClient(redisInstance.database_name);
 
-      const constructor = new RedisClientConstructor({
-        sdkClient: upstashSDK,
+      const constructor = new RedisClient({
+        upstashSDK: upstashSDK,
         redisDbNameOrInstance: existingRedisClient,
       });
       const redisClient = await constructor.getRedisClient();

src/clients/redis/index.ts

@@ -13,21 +13,21 @@ export const DEFAULT_REDIS_CONFIG: CreateCommandPayload = {
 };
 
 type Config = {
-  sdkClient: Upstash;
+  upstashSDK: Upstash;
   redisDbNameOrInstance?: string | Redis;
-  preferredRegion?: PreferredRegions;
+  region?: PreferredRegions;
 };
 
-export class RedisClientConstructor {
+export class RedisClient {
   private redisDbNameOrInstance?: string | Redis;
-  private preferredRegion?: PreferredRegions;
-  private sdkClient: Upstash;
+  private region?: PreferredRegions;
+  private upstashSDK: Upstash;
   private redisClient?: Redis;
 
-  constructor({ sdkClient, preferredRegion, redisDbNameOrInstance }: Config) {
+  constructor({ upstashSDK, region, redisDbNameOrInstance }: Config) {
     this.redisDbNameOrInstance = redisDbNameOrInstance;
-    this.sdkClient = sdkClient;
-    this.preferredRegion = preferredRegion ?? "us-east-1";
+    this.upstashSDK = upstashSDK;
+    this.region = region ?? "us-east-1";
   }
 
   public async getRedisClient(): Promise<Redis | undefined> {
@@ -63,22 +63,23 @@ export class RedisClientConstructor {
 
   private createRedisClientByName = async (redisDbName: string) => {
     try {
-      const redis = await this.sdkClient.getRedisDatabase(redisDbName);
-      this.redisClient = await this.sdkClient.newRedisClient(redis.database_name);
+      const redis = await this.upstashSDK.getRedisDatabase(redisDbName);
+      this.redisClient = await this.upstashSDK.newRedisClient(redis.database_name);
     } catch {
+      console.error(`Requested '${redisDbName}' is missing in DB list. Creating new one...`);
       await this.createRedisClientByDefaultConfig(redisDbName);
     }
   };
 
   private createRedisClientByDefaultConfig = async (redisDbName?: string) => {
-    const redisDatabase = await this.sdkClient.getOrCreateRedisDatabase({
+    const redisDatabase = await this.upstashSDK.getOrCreateRedisDatabase({
       ...DEFAULT_REDIS_CONFIG,
       name: redisDbName ?? DEFAULT_REDIS_CONFIG.name,
-      region: this.preferredRegion ?? DEFAULT_REDIS_CONFIG.region,
+      region: this.region ?? DEFAULT_REDIS_CONFIG.region,
     });
 
     if (redisDatabase?.database_name) {
-      this.redisClient = await this.sdkClient.newRedisClient(redisDatabase.database_name);
+      this.redisClient = await this.upstashSDK.newRedisClient(redisDatabase.database_name);
     }
   };
 }

src/clients/vector/index.test.ts

@@ -1,34 +1,34 @@
 /* eslint-disable @typescript-eslint/no-non-null-assertion */
 import { Upstash } from "@upstash/sdk";
 import { describe, expect, test } from "bun:test";
-import { DEFAULT_VECTOR_DB_NAME, VectorClientConstructor, DEFAULT_VECTOR_CONFIG } from ".";
+import { DEFAULT_VECTOR_DB_NAME, VectorClient, DEFAULT_VECTOR_CONFIG } from ".";
 
 const upstashSDK = new Upstash({
   email: process.env.UPSTASH_EMAIL!,
   token: process.env.UPSTASH_TOKEN!,
 });
 
-describe("Redis Client", () => {
+describe("Vector Client", () => {
   test(
     "Initialize client without index name",
     async () => {
-      const constructor = new VectorClientConstructor({
-        sdkClient: upstashSDK,
+      const constructor = new VectorClient({
+        upstashSDK: upstashSDK,
       });
       const vectorClient = await constructor.getVectorClient();
 
       expect(vectorClient).toBeTruthy();
 
       await upstashSDK.deleteVectorIndex(DEFAULT_VECTOR_DB_NAME);
     },
-    { timeout: 10_000 }
+    { timeout: 20_000 }
   );
 
   test(
     "Initialize client with db name",
     async () => {
-      const constructor = new VectorClientConstructor({
-        sdkClient: upstashSDK,
+      const constructor = new VectorClient({
+        upstashSDK: upstashSDK,
         indexNameOrInstance: "test-name",
       });
       const redisClient = await constructor.getVectorClient();
@@ -37,7 +37,7 @@ describe("Redis Client", () => {
 
       await upstashSDK.deleteVectorIndex("test-name");
     },
-    { timeout: 10_000 }
+    { timeout: 20_000 }
   );
 
   test(
@@ -47,11 +47,12 @@ describe("Redis Client", () => {
       const vectorInstance = await upstashSDK.createVectorIndex({
         ...DEFAULT_VECTOR_CONFIG,
         name: indexName,
+        type: "free",
       });
       const existingVectorClient = await upstashSDK.newVectorClient(vectorInstance.name);
 
-      const constructor = new VectorClientConstructor({
-        sdkClient: upstashSDK,
+      const constructor = new VectorClient({
+        upstashSDK: upstashSDK,
         indexNameOrInstance: existingVectorClient,
       });
       const vectorClient = await constructor.getVectorClient();
@@ -60,6 +61,6 @@ describe("Redis Client", () => {
 
       await upstashSDK.deleteVectorIndex(indexName);
     },
-    { timeout: 10_000 }
+    { timeout: 20_000 }
   );
 });