Skip to content
This repository has been archived by the owner on Aug 30, 2024. It is now read-only.

Commit

Permalink
Fix document endpoints
Browse files Browse the repository at this point in the history
  • Loading branch information
NolanTrem committed Jul 5, 2024
1 parent f2d60cb commit 59dc0c7
Show file tree
Hide file tree
Showing 4 changed files with 119 additions and 19 deletions.
64 changes: 50 additions & 14 deletions __tests__/r2rClientIntegration.test.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
import { r2rClient } from "../src/index";
import { FilterCriteria, AnalysisTypes } from "../src/models";
const fs = require("fs");

const baseUrl = "http://localhost:8000";

Expand All @@ -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" },
Expand All @@ -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();
});
Expand All @@ -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();
Expand All @@ -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",
],
);
});
});
9 changes: 6 additions & 3 deletions src/models.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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 {
Expand Down
44 changes: 42 additions & 2 deletions src/r2rClient.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ if (typeof window === "undefined") {

import { feature, initializeTelemetry } from "./feature";
import {
Document,
R2RUpdatePromptRequest,
R2RIngestDocumentsRequest,
R2RIngestFilesRequest,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
21 changes: 21 additions & 0 deletions testr2r.js
Original file line number Diff line number Diff line change
@@ -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);

0 comments on commit 59dc0c7

Please sign in to comment.