Add hosted indexers support (#50)

- Move local indexers into `indexers/local` folder - Various test fixes - Add web scraper and Github indexer API and tests
cortexclick · Nov 3, 2024 · 99922ee · 99922ee
1 parent c4d3063
commit 99922ee
Show file tree

Hide file tree

Showing 14 changed files with 414 additions and 147 deletions.
diff --git a/src/catalog.test.ts b/src/catalog.test.ts
@@ -2,20 +2,16 @@ import { expect, test, afterEach } from "vitest";
 import { Catalog, CatalogConfig } from "./catalog";
 import { testClient } from "./vitest-test-client";
 
-let catalog: Catalog;
+let catalog: Catalog | undefined;
 
 afterEach(async () => {
-  try {
-    if (catalog) {
-      await catalog.delete();
-    }
-  } catch {
-    // probably catalog has already been deleted by a test, so ignore
+  if (catalog) {
+    await catalog.delete();
   }
-});
+}, 20000);
 
-test("Catalog CRUD", { timeout: 10000 }, async () => {
-  const catalogName = `catalog-${Math.floor(Math.random() * 10000)}`;
+test("Catalog CRUD", { timeout: 60000 }, async () => {
+  const catalogName = `catalog-sdk-test-${Date.now()}`;
 
   const config: CatalogConfig = {
     description: "foo bar",
@@ -58,6 +54,8 @@ test("Catalog CRUD", { timeout: 10000 }, async () => {
 
   // delete
   await catalog.delete();
+  catalog = undefined;
+
   // assert that the get fails
   await expect(async () => {
     await testClient.getCatalog(catalogName);

diff --git a/src/catalog.ts b/src/catalog.ts
@@ -7,13 +7,16 @@ import {
   DocumentListItem,
 } from "./document";
 import * as fs from "node:fs";
-import { JSONIndexer, JSONIndexerOpts } from "./indexers/json-indexer";
+import { JSONIndexer, JSONIndexerOpts } from "./indexers/local/json-indexer";
 import {
   DirectoryIndexer,
   DirectoryIndexerOpts,
-} from "./indexers/directory-indexer";
-import { TSVIndexer, TSVIndexerOpts } from "./indexers/tsv-indexer";
-import { ShopifyIndexer, ShopifyIndexerOpts } from "./indexers/shopify-indexer";
+} from "./indexers/local/directory-indexer";
+import { TSVIndexer, TSVIndexerOpts } from "./indexers/local/tsv-indexer";
+import {
+  ShopifyIndexer,
+  ShopifyIndexerOpts,
+} from "./indexers/local/shopify-indexer";
 
 export type CatalogConfig = {
   description: string;
@@ -202,25 +205,25 @@ export class Catalog {
     return { warnings: body?.warnings ?? [] };
   }
 
-  async delete() {
+  public async delete() {
     this.checkDeleted();
     this.deleted = true;
     await this.apiClient.DELETE(`/catalogs/${this.name}`);
     return;
   }
 
-  async getDocument(documentId: string): Promise<Document> {
+  public async getDocument(documentId: string): Promise<Document> {
     this.checkDeleted();
     return Document.get(this.apiClient, this, documentId);
   }
 
-  async deleteDocument(documentId: string) {
+  public async deleteDocument(documentId: string) {
     this.checkDeleted();
     const doc = await this.getDocument(documentId);
     await doc.delete();
   }
 
-  async listDocuments(
+  public async listDocuments(
     paginationOpts?: DocumentPaginationOpts,
   ): Promise<DocumentListResult> {
     const { page, pageSize } = paginationOpts || { page: 1, pageSize: 50 };
@@ -246,19 +249,19 @@ export class Catalog {
   }
 
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  jsonIndexer(documents: any[], opts?: JSONIndexerOpts) {
+  public jsonIndexer(documents: any[], opts?: JSONIndexerOpts) {
     return new JSONIndexer(this, documents, opts);
   }
 
-  directoryIndexer(opts: DirectoryIndexerOpts) {
+  public directoryIndexer(opts: DirectoryIndexerOpts) {
     return new DirectoryIndexer(this, opts);
   }
 
-  tsvIndexer(file: string, opts?: TSVIndexerOpts) {
+  public tsvIndexer(file: string, opts?: TSVIndexerOpts) {
     return new TSVIndexer(this, file, opts);
   }
 
-  shopifyIndexer(opts: ShopifyIndexerOpts) {
+  public shopifyIndexer(opts: ShopifyIndexerOpts) {
     return new ShopifyIndexer(this, opts);
   }
 

diff --git a/src/client.ts b/src/client.ts
@@ -5,6 +5,12 @@ import { CortexApiClient } from "./api-client";
 import { Chat, StreamingChatResult } from "./chat";
 import { Content, ContentStatus, StreamingContentResult } from "./content";
 import { Readable } from "stream";
+import {
+  GithubDataSourceConfig,
+  Indexer,
+  IndexerScheduleFrequency,
+  WebScraperDataSourceConfig,
+} from "./indexers/hosted/indexer";
 
 export type CortexClientArgs = {
   org: string;
@@ -163,4 +169,39 @@ export class CortexClient {
   async listCatalogs() {
     return Catalog.list(this.apiClient);
   }
+
+  async createWebScraperIndexer(
+    name: string,
+    catalogName: string,
+    schedule: IndexerScheduleFrequency,
+    config: WebScraperDataSourceConfig,
+  ): Promise<Indexer> {
+    return Indexer.create(this.apiClient, name, catalogName, schedule, {
+      type: "webScraper",
+      config,
+    });
+  }
+
+  async createGithubIndexer(
+    name: string,
+    catalogName: string,
+    schedule: IndexerScheduleFrequency,
+    config: GithubDataSourceConfig,
+  ): Promise<Indexer> {
+    if (!config.code) {
+      config.code = {};
+    }
+    return Indexer.create(this.apiClient, name, catalogName, schedule, {
+      type: "github",
+      config,
+    });
+  }
+
+  async getIndexer(name: string): Promise<Indexer> {
+    return Indexer.get(this.apiClient, name);
+  }
+
+  async listIndexers(): Promise<Indexer[]> {
+    return Indexer.list(this.apiClient);
+  }
 }
diff --git a/src/cortex.test.ts b/src/cortex.test.ts
@@ -1,7 +1,22 @@
-import { expect, test } from "vitest";
+import { afterEach, expect, test } from "vitest";
 import { OrgConfigOpts } from "./org";
 import { CortexConfig } from "./cortex";
 import { testClient } from "./vitest-test-client";
+import { Catalog, CatalogConfig } from "./catalog";
+
+let catalog: Catalog | undefined;
+let catalog2: Catalog | undefined;
+
+afterEach(async () => {
+  if (catalog) {
+    await catalog.delete();
+    catalog = undefined;
+  }
+  if (catalog2) {
+    await catalog2.delete();
+    catalog2 = undefined;
+  }
+});
 
 test("can get and set OrgConfig", async () => {
   const orgConfigOpts: OrgConfigOpts = {
@@ -26,42 +41,57 @@ test("can get and set OrgConfig", async () => {
   expect(getOrgConfig.companyInfo).toBe(orgConfigOpts.companyInfo);
 });
 
-test("can configure, get, and delete and Cortexes", async () => {
-  const cortexName = `cortex-${Math.floor(Math.random() * 10000)}`;
+test(
+  "can configure, get, and delete and Cortexes",
+  { timeout: 30000 },
+  async () => {
+    const cortexName = `cortex-sdk-test-${Date.now()}`;
 
-  const cortexConfig: CortexConfig = {
-    friendlyName: "Cortex AI",
-    catalogs: ["cat-1", "cat-2"],
-    instructions: ["do your job", "do it well"],
-    public: false,
-    customizations: {
-      rules: ["be nice to the user"],
-      personality: ["saucy"],
-      chatVerbosity: "concise",
-      writeVerbosity: "long-form",
-    },
-    chatConfig: {
-      intro: "hello world",
-      examples: ["q1", "q2"],
-      greeting: "who lives in a pineapple under the sea? CORTEX AI.",
-    },
-    overrides: {
-      companyInfo: "a very good company that does AI stuff",
-      companyName: "Cortex Click, Inc. --test",
-    },
-  };
+    const catalogName = `catalog-sdk-test-${Date.now()}`;
+    const catalogName2 = `catalog-sdk-test-${Date.now() + 1}`;
 
-  let cortex = await testClient.configureCortex(cortexName, cortexConfig);
+    const config: CatalogConfig = {
+      description: "foo bar",
+      instructions: ["a", "b"],
+    };
 
-  cortex = await testClient.getCortex(cortexName);
-  expect(cortex.config.catalogs).toStrictEqual(cortexConfig.catalogs);
-  expect(cortex.config.overrides?.inheritRules).toBe(true); // test input doesn't specify `inheritRules`, should be true by default
-  // TODO - check all the properties
+    catalog = await testClient.configureCatalog(catalogName, config);
+    catalog2 = await testClient.configureCatalog(catalogName2, config);
 
-  // delete the cortex
-  await cortex.delete();
-  // assert that the get fails
-  await expect(async () => {
-    await testClient.getCortex(cortexName);
-  }).rejects.toThrowError("Failed to get cortex: Not Found");
-});
+    const cortexConfig: CortexConfig = {
+      friendlyName: "Cortex AI",
+      catalogs: [catalogName, catalogName2],
+      instructions: ["do your job", "do it well"],
+      public: false,
+      customizations: {
+        rules: ["be nice to the user"],
+        personality: ["saucy"],
+        chatVerbosity: "concise",
+        writeVerbosity: "long-form",
+      },
+      chatConfig: {
+        intro: "hello world",
+        examples: ["q1", "q2"],
+        greeting: "who lives in a pineapple under the sea? CORTEX AI.",
+      },
+      overrides: {
+        companyInfo: "a very good company that does AI stuff",
+        companyName: "Cortex Click, Inc. --test",
+      },
+    };
+
+    let cortex = await testClient.configureCortex(cortexName, cortexConfig);
+
+    cortex = await testClient.getCortex(cortexName);
+    expect(cortex.config.catalogs).toStrictEqual(cortexConfig.catalogs);
+    expect(cortex.config.overrides?.inheritRules).toBe(true); // test input doesn't specify `inheritRules`, should be true by default
+    // TODO - check all the properties
+
+    // delete the cortex
+    await cortex.delete();
+    // assert that the get fails
+    await expect(async () => {
+      await testClient.getCortex(cortexName);
+    }).rejects.toThrowError("Failed to get cortex: Not Found");
+  },
+);
diff --git a/src/cortex.ts b/src/cortex.ts
@@ -171,7 +171,8 @@ export class Cortex {
     }
 
     if (res.status > 201) {
-      throw new Error(`Failed to configure cortex: ${res.statusText}`);
+      const message = res.status === 400 ? await res.text() : res.statusText;
+      throw new Error(`Failed to configure cortex: ${message}`);
     }
     return new Cortex(config, apiClient, name);
   }