Skip to content

Commit

Permalink
Export graph entities and relationships
Browse files Browse the repository at this point in the history
  • Loading branch information
@NolanTrem
NolanTrem committed Dec 31, 2024
1 parent ba39a88 commit 6b4944f
Show file tree
Hide file tree
Showing 6 changed files with 684 additions and 10 deletions.
187 changes: 178 additions & 9 deletions js/sdk/__tests__/GraphsIntegrationSuperUser.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,7 @@ describe("r2rClient V3 Graphs Integration Tests", () => {
);
await client.documents.exportEntities({
id: documentId,
outputPath,
outputPath: outputPath,
columns: ["id", "name", "created_at"],
includeHeader: true,
});
Expand Down Expand Up @@ -193,9 +193,10 @@ describe("r2rClient V3 Graphs Integration Tests", () => {
TEST_OUTPUT_DIR,
"document_entities_empty.csv",
);
await client.documents.export({
outputPath,
filters: { document_type: { $eq: "non_existent_type" } },
await client.documents.exportEntities({
id: documentId,
outputPath: outputPath,
filters: { name: { $eq: "non_existent_name" } },
});

expect(fs.existsSync(outputPath)).toBe(true);
Expand Down Expand Up @@ -226,7 +227,7 @@ describe("r2rClient V3 Graphs Integration Tests", () => {
);
await client.documents.exportRelationships({
id: documentId,
outputPath,
outputPath: outputPath,
columns: ["subject", "object", "created_at"],
includeHeader: true,
});
Expand Down Expand Up @@ -275,14 +276,15 @@ describe("r2rClient V3 Graphs Integration Tests", () => {
const content = fs.readFileSync(outputPath, "utf-8");
});

test("Handle empty document entity export result", async () => {
test("Handle empty document relationships export result", async () => {
const outputPath = path.join(
TEST_OUTPUT_DIR,
"document_relationships_empty.csv",
);
await client.documents.export({
outputPath,
filters: { document_type: { $eq: "non_existent_type" } },
await client.documents.exportRelationships({
id: documentId,
outputPath: outputPath,
filters: { subject: { $eq: "non_existent_subject" } },
});

expect(fs.existsSync(outputPath)).toBe(true);
Expand Down Expand Up @@ -321,6 +323,173 @@ describe("r2rClient V3 Graphs Integration Tests", () => {
expect(response.totalEntries).toBeGreaterThanOrEqual(1);
});

test("Export graph entities to CSV with default options", async () => {
const outputPath = path.join(TEST_OUTPUT_DIR, "graph_entities_default.csv");
await client.graphs.exportEntities({
collectionId: collectionId,
outputPath: outputPath,
});

expect(fs.existsSync(outputPath)).toBe(true);
const content = fs.readFileSync(outputPath, "utf-8");
expect(content).toBeTruthy();
expect(content.split("\n").length).toBeGreaterThan(1);
});

test("Export graph entities to CSV with custom columns", async () => {
const outputPath = path.join(TEST_OUTPUT_DIR, "graph_entities_custom.csv");
await client.graphs.exportEntities({
collectionId: collectionId,
outputPath: outputPath,
columns: ["id", "name", "created_at"],
includeHeader: true,
});

expect(fs.existsSync(outputPath)).toBe(true);
const content = fs.readFileSync(outputPath, "utf-8");
const headers = content
.split("\n")[0]
.split(",")
.map((h) => h.trim());

expect(headers).toContain('"id"');
expect(headers).toContain('"name"');
expect(headers).toContain('"created_at"');
});

test("Export filtered graph entities to CSV", async () => {
const outputPath = path.join(
TEST_OUTPUT_DIR,
"graph_entities_filtered.csv",
);
await client.graphs.exportEntities({
collectionId: collectionId,
outputPath: outputPath,
filters: { document_type: { $eq: "txt" } },
includeHeader: true,
});

expect(fs.existsSync(outputPath)).toBe(true);
const content = fs.readFileSync(outputPath, "utf-8");
expect(content).toBeTruthy();
});

test("Export graph entities without headers", async () => {
const outputPath = path.join(
TEST_OUTPUT_DIR,
"graph_entities_no_header.csv",
);
await client.graphs.exportEntities({
collectionId: collectionId,
outputPath: outputPath,
includeHeader: false,
});

expect(fs.existsSync(outputPath)).toBe(true);
const content = fs.readFileSync(outputPath, "utf-8");
});

test("Handle empty graph entity export result", async () => {
const outputPath = path.join(TEST_OUTPUT_DIR, "graph_entities_empty.csv");
await client.graphs.exportEntities({
collectionId: collectionId,
outputPath: outputPath,
filters: { document_type: { $eq: "non_existent_type" } },
});

expect(fs.existsSync(outputPath)).toBe(true);
const content = fs.readFileSync(outputPath, "utf-8");
expect(content.split("\n").filter((line) => line.trim()).length).toBe(1);
});

test("Export graph relationships to CSV with default options", async () => {
const outputPath = path.join(
TEST_OUTPUT_DIR,
"graphs_relationships_default.csv",
);
await client.graphs.exportRelationships({
collectionId: collectionId,
outputPath: outputPath,
});

expect(fs.existsSync(outputPath)).toBe(true);
const content = fs.readFileSync(outputPath, "utf-8");
expect(content).toBeTruthy();
expect(content.split("\n").length).toBeGreaterThan(1);
});

test("Export graph relationships to CSV with custom columns", async () => {
const outputPath = path.join(
TEST_OUTPUT_DIR,
"graph_relationships_custom.csv",
);
await client.graphs.exportRelationships({
collectionId: collectionId,
outputPath: outputPath,
columns: ["subject", "object", "created_at"],
includeHeader: true,
});

expect(fs.existsSync(outputPath)).toBe(true);
const content = fs.readFileSync(outputPath, "utf-8");
const headers = content
.split("\n")[0]
.split(",")
.map((h) => h.trim());

expect(headers).toContain('"subject"');
expect(headers).toContain('"object"');
expect(headers).toContain('"created_at"');
});

test("Export filtered graphs entities to CSV", async () => {
const outputPath = path.join(
TEST_OUTPUT_DIR,
"graphs_entities_filtered.csv",
);
await client.graphs.exportEntities({
collectionId: collectionId,
outputPath: outputPath,
filters: { document_type: { $eq: "txt" } },
includeHeader: true,
});

expect(fs.existsSync(outputPath)).toBe(true);
const content = fs.readFileSync(outputPath, "utf-8");
expect(content).toBeTruthy();
});

test("Export document relationships without headers", async () => {
const outputPath = path.join(
TEST_OUTPUT_DIR,
"document_relationships_no_header.csv",
);
await client.documents.exportRelationships({
id: documentId,
outputPath: outputPath,
includeHeader: false,
});

expect(fs.existsSync(outputPath)).toBe(true);
const content = fs.readFileSync(outputPath, "utf-8");
});

test("Handle empty graphs entity export result", async () => {
const outputPath = path.join(
TEST_OUTPUT_DIR,
"document_relationships_empty.csv",
);
await client.graphs.exportEntities({
collectionId: collectionId,
outputPath: outputPath,
filters: { document_type: { $eq: "non_existent_type" } },
});

expect(fs.existsSync(outputPath)).toBe(true);
const content = fs.readFileSync(outputPath, "utf-8");
expect(content.split("\n").filter((line) => line.trim()).length).toBe(1);
});

test("Check that there are no communities in the graph prior to building", async () => {
const response = await client.graphs.listCommunities({
collectionId: collectionId,
Expand Down
145 changes: 145 additions & 0 deletions js/sdk/src/v3/clients/graphs.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,14 @@ import {
WrappedCommunitiesResponse,
WrappedCommunityResponse,
} from "../../types";
import { downloadBlob } from "../../utils";

let fs: any;
if (typeof window === "undefined") {
import("fs").then((module) => {
fs = module;
});
}

export class GraphsClient {
constructor(private client: r2rClient) {}
Expand Down Expand Up @@ -358,6 +366,143 @@ export class GraphsClient {
);
}

/**
* Export document entities as a CSV file with support for filtering and column selection.
*
* @param options Export configuration options
* @param options.outputPath Path where the CSV file should be saved (Node.js only)
* @param options.columns Optional list of specific columns to include
* @param options.filters Optional filters to limit which documents are exported
* @param options.includeHeader Whether to include column headers (default: true)
* @returns Promise<Blob> in browser environments, Promise<void> in Node.js
*/
@feature("documents.exportEntities")
async exportEntities(options: {
collectionId: string;
outputPath?: string;
columns?: string[];
filters?: Record<string, any>;
includeHeader?: boolean;
}): Promise<Blob | void> {
const data: Record<string, any> = {
id: options.collectionId,
include_header: options.includeHeader ?? true,
};

if (options.columns) {
data.columns = options.columns;
}
if (options.filters) {
data.filters = options.filters;
}

const response = await this.client.makeRequest(
"POST",
`documents/${options.collectionId}/entities/export`,
{
data,
responseType: "arraybuffer",
headers: { Accept: "text/csv" },
},
);

// Node environment
if (options.outputPath && typeof process !== "undefined") {
await fs.promises.writeFile(options.outputPath, Buffer.from(response));
return;
}

// Browser
return new Blob([response], { type: "text/csv" });
}

/**
* Export documents as a CSV file and save it to the user's device.
* @param filename
* @param options
*/
@feature("documents.exportEntitiesToFile")
async exportEntitiesToFile(options: {
filename: string;
collectionId: string;
columns?: string[];
filters?: Record<string, any>;
includeHeader?: boolean;
}): Promise<void> {
const blob = await this.exportEntities(options);
if (blob instanceof Blob) {
downloadBlob(blob, options.filename);
}
}

/**
* Export document relationships as a CSV file with support for filtering and column selection.
*
* @param options Export configuration options
* @param options.outputPath Path where the CSV file should be saved (Node.js only)
* @param options.columns Optional list of specific columns to include
* @param options.filters Optional filters to limit which documents are exported
* @param options.includeHeader Whether to include column headers (default: true)
* @returns Promise<Blob> in browser environments, Promise<void> in Node.js
*/
@feature("documents.exportRelationships")
async exportRelationships(options: {
collectionId: string;
outputPath?: string;
columns?: string[];
filters?: Record<string, any>;
includeHeader?: boolean;
}): Promise<Blob | void> {
const data: Record<string, any> = {
include_header: options.includeHeader ?? true,
};

if (options.columns) {
data.columns = options.columns;
}
if (options.filters) {
data.filters = options.filters;
}

const response = await this.client.makeRequest(
"POST",
`documents/${options.collectionId}/relationships/export`,
{
data,
responseType: "arraybuffer",
headers: { Accept: "text/csv" },
},
);

// Node environment
if (options.outputPath && typeof process !== "undefined") {
await fs.promises.writeFile(options.outputPath, Buffer.from(response));
return;
}

// Browser
return new Blob([response], { type: "text/csv" });
}

/**
* Export document relationships as a CSV file and save it to the user's device.
* @param filename
* @param options
*/
@feature("documents.exportRelationshipsToFile")
async exportRelationshipsToFile(options: {
filename: string;
collectionId: string;
columns?: string[];
filters?: Record<string, any>;
includeHeader?: boolean;
}): Promise<void> {
const blob = await this.exportRelationships(options);
if (blob instanceof Blob) {
downloadBlob(blob, options.filename);
}
}

/**
* Creates a new community in the graph.
*
Expand Down
Loading

0 comments on commit 6b4944f

Please sign in to comment.