Extract collections request to separate API

elpassion · Feb 23, 2024 · c0024bd · c0024bd
1 parent b457ce7
commit c0024bd
Show file tree

Hide file tree

Showing 14 changed files with 182 additions and 129 deletions.
diff --git a/apps/web-remix/app/api/knowledgeBase/KnowledgeBaseApi.ts b/apps/web-remix/app/api/knowledgeBase/KnowledgeBaseApi.ts
@@ -1,11 +1,79 @@
 import { fetchTyped } from "~/utils/fetch.server";
 import { buildUrlWithParams } from "~/utils/url";
-import { MemoryChunksResponse } from "~/components/pages/knowledgeBase/contracts";
+import {
+  KnowledgeBaseCollectionFromListResponse,
+  KnowledgeBaseCollectionListResponse,
+  KnowledgeBaseFileListResponse,
+  MemoryChunksResponse,
+  CreateCollectionSchema,
+} from "./knowledgeApi.contracts";
 import { PaginationQueryParams } from "~/components/pagination/usePagination";
+import { z } from "zod";
 
 export class KnowledgeBaseApi {
   constructor(private client: typeof fetchTyped) {}
 
+  async getCollections(organizationId: string | number) {
+    return this.client(
+      KnowledgeBaseCollectionListResponse,
+      `/organizations/${organizationId}/memory_collections`
+    );
+  }
+
+  async getCollectionByName(
+    organizationId: string | number,
+    collectionName: string | number
+  ) {
+    return this.client(
+      KnowledgeBaseCollectionFromListResponse,
+      `/organizations/${organizationId}/memory_collections?collection_name=${collectionName}`
+    );
+  }
+
+  async deleteCollection(
+    organizationId: string | number,
+    collectionId: string | number
+  ) {
+    return this.client(
+      z.any(),
+      `/organizations/${organizationId}/memory_collections/${collectionId}`,
+      { method: "DELETE" }
+    );
+  }
+
+  async getCollectionMemories(
+    organizationId: string | number,
+    collectionId: string | number
+  ) {
+    return this.client(
+      KnowledgeBaseFileListResponse,
+      `/organizations/${organizationId}/memory_collections/${collectionId}/memories`
+    );
+  }
+
+  async deleteCollectionMemory(
+    organizationId: string | number,
+    collectionId: string | number,
+    memoryId: string | number
+  ) {
+    return this.client(
+      z.any(),
+      `/organizations/${organizationId}/memory_collections/${collectionId}/memories/${memoryId}`,
+      { method: "DELETE" }
+    );
+  }
+
+  async createCollection(
+    organizationId: string | number,
+    data: z.TypeOf<typeof CreateCollectionSchema>
+  ) {
+    return this.client(
+      z.any(),
+      `/organizations/${organizationId}/memory_collections`,
+      { method: "POST", body: JSON.stringify(data) }
+    );
+  }
+
   async getMemoryChunk(
     organizationId: string | number,
     collectionId: string | number,

diff --git a/apps/web-remix/app/api/knowledgeBase/knowledgeApi.contracts.ts b/apps/web-remix/app/api/knowledgeBase/knowledgeApi.contracts.ts
@@ -1,5 +1,6 @@
 import { z } from "zod";
 import { PaginationMeta } from "~/components/pagination/pagination.types";
+import { NotFoundError } from "~/utils/errors";
 
 export const MemoryChunk = z.object({
   id: z.string(),
@@ -16,3 +17,58 @@ export const MemoryChunksResponse = z.object({
 });
 
 export type IMemoryChunksResponse = z.TypeOf<typeof MemoryChunksResponse>;
+
+export const CreateCollectionSchema = z.object({
+  collection_name: z.string().min(2),
+  embeddings: z.object({
+    api_type: z.string().min(2),
+    model: z.string().min(2),
+    secret_name: z.string().min(2),
+  }),
+});
+
+export const KnowledgeBaseFile = z.object({
+  id: z.number(),
+  file_name: z.string(),
+  file_size: z.number(),
+  file_type: z.string(),
+});
+
+export const KnowledgeBaseCollection = z.object({
+  id: z.number(),
+  name: z.string(),
+  embeddings: z.object({
+    api_type: z.string(),
+    model: z.string(),
+    secret_name: z.string(),
+  }),
+});
+
+export const KnowledgeBaseFileResponse = z
+  .object({ data: KnowledgeBaseFile })
+  .transform((res) => res.data);
+
+export const KnowledgeBaseFileList = z.array(KnowledgeBaseFile);
+
+export const KnowledgeBaseCollectionList = z.array(KnowledgeBaseCollection);
+export const KnowledgeBaseFileListResponse = z
+  .object({
+    data: KnowledgeBaseFileList,
+  })
+  .transform((res) => res.data);
+
+export const KnowledgeBaseCollectionListResponse = z
+  .object({
+    data: KnowledgeBaseCollectionList,
+  })
+  .transform((res) => res.data);
+
+export const KnowledgeBaseCollectionFromListResponse =
+  KnowledgeBaseCollectionListResponse.transform((res) => {
+    const collection = res[0];
+
+    if (!collection) {
+      throw new NotFoundError();
+    }
+    return collection;
+  });
diff --git a/apps/web-remix/app/components/pages/knowledgeBase/collection/action.server.tsx b/apps/web-remix/app/components/pages/knowledgeBase/collection/action.server.tsx
@@ -2,8 +2,8 @@ import { ActionFunctionArgs, json } from "@remix-run/node";
 import { actionBuilder } from "~/utils.server";
 import { requireLogin } from "~/session.server";
 import invariant from "tiny-invariant";
-import { z } from "zod";
 import { setServerToast } from "~/utils/toast.server";
+import { KnowledgeBaseApi } from "~/api/knowledgeBase/KnowledgeBaseApi";
 
 export async function action(actionArgs: ActionFunctionArgs) {
   return actionBuilder({
@@ -15,18 +15,21 @@ export async function action(actionArgs: ActionFunctionArgs) {
       const memoryId = (await request.formData()).get("memoryId");
       const collectionName = params.collectionName;
 
+      const knowledgeBaseApi = new KnowledgeBaseApi(fetch);
+
       const {
         data: { id: collectionId },
-      } = await fetch(
-        z.any(),
-        `/organizations/${params.organizationId}/memory_collections?collection_name=${collectionName}`
+      } = await knowledgeBaseApi.getCollectionByName(
+        params.organizationId,
+        collectionName
       );
 
-      await fetch(
-        z.any(),
-        `/organizations/${params.organizationId}/memory_collections/${collectionId}/memories/${memoryId}`,
-        { method: "DELETE" }
+      await knowledgeBaseApi.deleteCollectionMemory(
+        params.organizationId,
+        collectionId,
+        memoryId as string
       );
+
       return json(
         {},
         {

diff --git a/apps/web-remix/app/components/pages/knowledgeBase/collection/loader.server.tsx b/apps/web-remix/app/components/pages/knowledgeBase/collection/loader.server.tsx
@@ -2,30 +2,26 @@ import { json, LoaderFunctionArgs } from "@remix-run/node";
 import invariant from "tiny-invariant";
 import { loaderBuilder } from "~/utils.server";
 import { requireLogin } from "~/session.server";
-import {
-  KnowledgeBaseCollectionFromListResponse,
-  KnowledgeBaseFileListResponse,
-} from "../contracts";
+import { KnowledgeBaseApi } from "~/api/knowledgeBase/KnowledgeBaseApi";
 
 export async function loader(args: LoaderFunctionArgs) {
   return loaderBuilder(async ({ request, params }, { fetch }) => {
     await requireLogin(request);
     invariant(params.organizationId, "organizationId not found");
     invariant(params.collectionName, "collectionName not found");
 
+    const knowledgeBaseApi = new KnowledgeBaseApi(fetch);
+
     const {
       data: { id: collectionId },
-    } = await fetch(
-      KnowledgeBaseCollectionFromListResponse,
-      `/organizations/${
-        params.organizationId
-      }/memory_collections?collection_name=${encodeURIComponent(
-        params.collectionName
-      )}`
+    } = await knowledgeBaseApi.getCollectionByName(
+      params.organizationId,
+      params.collectionName
     );
-    const knowledgeBase = await fetch(
-      KnowledgeBaseFileListResponse,
-      `/organizations/${params.organizationId}/memory_collections/${collectionId}/memories`
+
+    const knowledgeBase = await knowledgeBaseApi.getCollectionMemories(
+      params.organizationId,
+      collectionId
     );
 
     return json({

diff --git a/apps/web-remix/app/components/pages/knowledgeBase/contracts.ts b/apps/web-remix/app/components/pages/knowledgeBase/contracts.ts
diff --git a/apps/web-remix/app/components/pages/knowledgeBase/knowledgeBase.types.ts b/apps/web-remix/app/components/pages/knowledgeBase/knowledgeBase.types.ts
@@ -3,7 +3,7 @@ import {
   KnowledgeBaseCollection,
   KnowledgeBaseFile,
   KnowledgeBaseFileList,
-} from "./contracts";
+} from "~/api/knowledgeBase/knowledgeApi.contracts";
 
 export type IKnowledgeBaseFile = z.TypeOf<typeof KnowledgeBaseFile>;
 

diff --git a/apps/web-remix/app/components/pages/knowledgeBase/list/action.server.tsx b/apps/web-remix/app/components/pages/knowledgeBase/list/action.server.tsx
@@ -1,29 +1,31 @@
 import { ActionFunctionArgs, json } from "@remix-run/node";
 import invariant from "tiny-invariant";
-import { z } from "zod";
 import { requireLogin } from "~/session.server";
 import { actionBuilder } from "~/utils.server";
 import { setServerToast } from "~/utils/toast.server";
-import { KnowledgeBaseCollectionFromListResponse } from "../contracts";
+import { KnowledgeBaseApi } from "~/api/knowledgeBase/KnowledgeBaseApi";
 
 export async function action(actionArgs: ActionFunctionArgs) {
   return actionBuilder({
     delete: async ({ params, request }, { fetch }) => {
       await requireLogin(request);
       invariant(params.organizationId, "Missing organizationId");
 
+      const knowledgeBaseApi = new KnowledgeBaseApi(fetch);
+
       const collectionName = (await request.formData()).get("collectionName");
       const {
         data: { id: collectionId },
-      } = await fetch(
-        KnowledgeBaseCollectionFromListResponse,
-        `/organizations/${params.organizationId}/memory_collections?collection_name=${collectionName}`
+      } = await knowledgeBaseApi.getCollectionByName(
+        params.organizationId,
+        collectionName as string
       );
-      await fetch(
-        z.any(),
-        `/organizations/${params.organizationId}/memory_collections/${collectionId}`,
-        { method: "DELETE" }
+
+      await knowledgeBaseApi.deleteCollection(
+        params.organizationId,
+        collectionId
       );
+
       return json(
         {},
         {

diff --git a/apps/web-remix/app/components/pages/knowledgeBase/list/loader.server.tsx b/apps/web-remix/app/components/pages/knowledgeBase/list/loader.server.tsx
@@ -2,16 +2,17 @@ import { json, LoaderFunctionArgs } from "@remix-run/node";
 import invariant from "tiny-invariant";
 import { loaderBuilder } from "~/utils.server";
 import { requireLogin } from "~/session.server";
-import { KnowledgeBaseCollectionListResponse } from "../contracts";
+import { KnowledgeBaseApi } from "~/api/knowledgeBase/KnowledgeBaseApi";
 
 export async function loader(args: LoaderFunctionArgs) {
   return loaderBuilder(async ({ request, params }, { fetch }) => {
     await requireLogin(request);
     invariant(params.organizationId, "organizationId not found");
 
-    const collections = await fetch(
-      KnowledgeBaseCollectionListResponse,
-      `/organizations/${params.organizationId}/memory_collections`
+    const knowledgeBaseApi = new KnowledgeBaseApi(fetch);
+
+    const collections = await knowledgeBaseApi.getCollections(
+      params.organizationId
     );
 
     return json({

diff --git a/apps/web-remix/app/components/pages/knowledgeBase/memory/loader.server.tsx b/apps/web-remix/app/components/pages/knowledgeBase/memory/loader.server.tsx
@@ -2,14 +2,8 @@ import { json, LoaderFunctionArgs } from "@remix-run/node";
 import invariant from "tiny-invariant";
 import { loaderBuilder } from "~/utils.server";
 import { requireLogin } from "~/session.server";
-import {
-  KnowledgeBaseCollectionFromListResponse,
-  KnowledgeBaseFileListResponse,
-  MemoryChunksResponse,
-} from "~/components/pages/knowledgeBase/contracts";
-import z from "zod";
+import { KnowledgeBaseCollectionFromListResponse } from "~/api/knowledgeBase/knowledgeApi.contracts";
 import { getParamsPagination } from "~/components/pagination/usePagination";
-import { buildUrlWithParams } from "~/utils/url";
 import { KnowledgeBaseApi } from "~/api/knowledgeBase/KnowledgeBaseApi";
 
 export async function loader(args: LoaderFunctionArgs) {