Fix document endpoints

Author: NolanTrem

__tests__/r2rClientIntegration.test.ts

@@ -1,4 +1,6 @@
 import { r2rClient } from "../src/index";
+import { FilterCriteria, AnalysisTypes } from "../src/models";
+const fs = require("fs");
 
 const baseUrl = "http://localhost:8000";
 
@@ -13,6 +15,33 @@ describe("r2rClient Integration Tests", () => {
     await expect(client.healthCheck()).resolves.not.toThrow();
   });
 
+  test("Ingest documents", async () => {
+    const documentsToIngest = [
+      {
+        type: "txt",
+        data: fs.readFileSync("examples/data/myshkin.txt", "utf8"),
+        metadata: { title: "myshkin.txt" },
+      },
+    ];
+    await expect(
+      client.ingestDocuments(documentsToIngest),
+    ).resolves.not.toThrow();
+  });
+
+  test("Update documents", async () => {
+    const documentsToUpdate = [
+      {
+        id: "4430ddbf-4323-5a14-9205-c092920e9321",
+        type: "txt",
+        data: fs.readFileSync("examples/data/raskolnikov.txt", "utf8"),
+        metadata: { title: "updated_myshkin.txt" },
+      },
+    ];
+    await expect(
+      client.updateDocuments(documentsToUpdate),
+    ).resolves.not.toThrow();
+  });
+
   test("Ingest files", async () => {
     const files = [
       { path: "examples/data/raskolnikov.txt", name: "raskolnikov.txt" },
@@ -37,7 +66,7 @@ describe("r2rClient Integration Tests", () => {
     await expect(
       client.updateFiles(updated_file, {
         document_ids: ["48e29904-3010-54fe-abe5-a4f3fba59110"],
-        metadatas: [{ title: "myshkin.txt" }],
+        metadatas: [{ title: "updated_karamozov.txt" }],
       }),
     ).resolves.not.toThrow();
   });
@@ -64,19 +93,23 @@ describe("r2rClient Integration Tests", () => {
     await expect(client.appSettings()).resolves.not.toThrow();
   });
 
-  // test('Get analytics', async () => {
-  //   const filterCriteria = {
-  //     filters: {
-  //       "search_latencies": "search_latency"
-  //     }
-  //   };
+  test("Get analytics", async () => {
+    const filterCriteria: FilterCriteria = {
+      filters: {
+        search_latencies: "search_latency",
+      },
+    };
 
-  //   const analysisTypes = {
-  //     "search_latencies": ["basic_statistics", "search_latency"]
-  //   };
+    const analysisTypes: AnalysisTypes = {
+      analysis_types: {
+        search_latencies: ["basic_statistics", "search_latency"],
+      },
+    };
 
-  //   await expect(client.analytics(filterCriteria, analysisTypes)).resolves.not.toThrow();
-  // });
+    await expect(
+      client.analytics(filterCriteria, analysisTypes),
+    ).resolves.not.toThrow();
+  });
 
   test("Get users overview", async () => {
     await expect(client.usersOverview()).resolves.not.toThrow();
@@ -95,8 +128,11 @@ describe("r2rClient Integration Tests", () => {
   afterAll(async () => {
     // Clean up
     await client.delete(
-      ["document_id"],
-      ["48e29904-3010-54fe-abe5-a4f3fba59110"],
+      ["document_id", "document_id"],
+      [
+        "48e29904-3010-54fe-abe5-a4f3fba59110",
+        "4430ddbf-4323-5a14-9205-c092920e9321",
+      ],
     );
   });
 });

src/models.tsx

@@ -37,7 +37,10 @@ export interface KGSearchSettings {
 }
 
 export interface Document {
-  // Define the properties of Document
+  id?: string;
+  type: string;
+  data: string;
+  metadata: Record<string, any>;
 }
 
 export interface R2RUpdatePromptRequest {
@@ -53,8 +56,8 @@ export interface R2RIngestDocumentsRequest {
 
 export interface R2RUpdateDocumentsRequest {
   documents: Document[];
-  versions?: string[];
-  metadatas?: Record<string, any>[];
+  versions?: string[] | null;
+  metadatas?: Record<string, any>[] | null;
 }
 
 export interface R2RIngestFilesRequest {

src/r2rClient.ts

@@ -8,6 +8,7 @@ if (typeof window === "undefined") {
 
 import { feature, initializeTelemetry } from "./feature";
 import {
+  Document,
   R2RUpdatePromptRequest,
   R2RIngestDocumentsRequest,
   R2RIngestFilesRequest,
@@ -85,7 +86,22 @@ export class r2rClient {
   }
 
   @feature("ingestDocuments")
-  async ingestDocuments(request: R2RIngestDocumentsRequest): Promise<any> {
+  async ingestDocuments(
+    documents: Document[],
+    versions?: string[],
+  ): Promise<any> {
+    const processedDocuments = documents.map((doc) => ({
+      type: doc.type,
+      data: Buffer.from(doc.data).toString("base64"),
+      metadata: doc.metadata,
+      ...(doc.id && { id: doc.id }),
+    }));
+
+    const request: R2RIngestDocumentsRequest = {
+      documents: processedDocuments,
+      ...(versions && { versions }),
+    };
+
     const response = await this.axiosInstance.post(
       "/ingest_documents",
       request,
@@ -149,7 +165,31 @@ export class r2rClient {
   }
 
   @feature("updateDocuments")
-  async updateDocuments(request: R2RUpdateDocumentsRequest): Promise<any> {
+  async updateDocuments(
+    documents: Document[],
+    versions?: string[] | null,
+    metadatas?: Record<string, any>[] | null,
+  ): Promise<any> {
+    const processedDocuments = documents.map((doc) => ({
+      id: doc.id,
+      type: doc.type,
+      data: Buffer.from(doc.data).toString("base64"),
+      metadata: doc.metadata,
+    }));
+
+    const request: R2RUpdateDocumentsRequest = {
+      documents: processedDocuments,
+    };
+
+    // Only include versions and metadatas if they're explicitly provided
+    if (versions !== undefined && versions !== null) {
+      request.versions = versions;
+    }
+
+    if (metadatas !== undefined && metadatas !== null) {
+      request.metadatas = metadatas;
+    }
+
     const response = await this.axiosInstance.post(
       "/update_documents",
       request,

testr2r.js

@@ -0,0 +1,21 @@
+const { r2rClient } = require("./dist/r2rClient");
+const fs = require("fs");
+
+const client = new r2rClient("http://localhost:8000");
+
+async function main() {
+  const filterCriteria = {
+    filters: {
+      search_latencies: "search_latency",
+    },
+  };
+
+  const analysisTypes = {
+    search_latencies: ["basic_statistics", "search_latency"],
+  };
+
+  result = await client.analytics(filterCriteria, analysisTypes);
+  console.log(result);
+}
+
+main().catch(console.error);