Skip to content

Commit

Permalink
Export communities
Browse files Browse the repository at this point in the history
  • Loading branch information
@NolanTrem
NolanTrem committed Dec 31, 2024
1 parent 6b4944f commit 8226c15
Show file tree
Hide file tree
Showing 5 changed files with 418 additions and 18 deletions.
89 changes: 87 additions & 2 deletions js/sdk/__tests__/GraphsIntegrationSuperUser.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -394,7 +394,7 @@ describe("r2rClient V3 Graphs Integration Tests", () => {
await client.graphs.exportEntities({
collectionId: collectionId,
outputPath: outputPath,
filters: { document_type: { $eq: "non_existent_type" } },
filters: { name: { $eq: "non_existent_name" } },
});

expect(fs.existsSync(outputPath)).toBe(true);
Expand Down Expand Up @@ -482,7 +482,7 @@ describe("r2rClient V3 Graphs Integration Tests", () => {
await client.graphs.exportEntities({
collectionId: collectionId,
outputPath: outputPath,
filters: { document_type: { $eq: "non_existent_type" } },
filters: { name: { $eq: "non_existent_name" } },
});

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

test("Export graph communities to CSV with default options", async () => {
const outputPath = path.join(
TEST_OUTPUT_DIR,
"graph_communities_default.csv",
);
await client.graphs.exportCommunities({
collectionId: documentId,
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 communities to CSV with custom columns", async () => {
const outputPath = path.join(TEST_OUTPUT_DIR, "graph_entities_custom.csv");
await client.graphs.exportCommunities({
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 communities to CSV", async () => {
const outputPath = path.join(
TEST_OUTPUT_DIR,
"graph_communities_filtered.csv",
);
await client.graphs.exportCommunities({
collectionId: collectionId,
outputPath: outputPath,
filters: { name: { $eq: "txt" } },
includeHeader: true,
});

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

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

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

test("Handle empty graph communities export result", async () => {
const outputPath = path.join(
TEST_OUTPUT_DIR,
"graph_communities_empty.csv",
);
await client.graphs.exportCommunities({
collectionId: collectionId,
outputPath: outputPath,
filters: { name: { $eq: "non_existent_name" } },
});

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("Create a new entity", async () => {
const response = await client.graphs.createEntity({
collectionId: collectionId,
Expand Down
89 changes: 78 additions & 11 deletions js/sdk/src/v3/clients/graphs.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -367,7 +367,7 @@ export class GraphsClient {
}

/**
* Export document entities as a CSV file with support for filtering and column selection.
* Export graph 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)
Expand All @@ -376,7 +376,7 @@ export class GraphsClient {
* @param options.includeHeader Whether to include column headers (default: true)
* @returns Promise<Blob> in browser environments, Promise<void> in Node.js
*/
@feature("documents.exportEntities")
@feature("graphs.exportEntities")
async exportEntities(options: {
collectionId: string;
outputPath?: string;
Expand All @@ -385,7 +385,6 @@ export class GraphsClient {
includeHeader?: boolean;
}): Promise<Blob | void> {
const data: Record<string, any> = {
id: options.collectionId,
include_header: options.includeHeader ?? true,
};

Expand All @@ -398,7 +397,7 @@ export class GraphsClient {

const response = await this.client.makeRequest(
"POST",
`documents/${options.collectionId}/entities/export`,
`graphs/${options.collectionId}/entities/export`,
{
data,
responseType: "arraybuffer",
Expand All @@ -417,11 +416,11 @@ export class GraphsClient {
}

/**
* Export documents as a CSV file and save it to the user's device.
* Export graph entities as a CSV file and save it to the user's device.
* @param filename
* @param options
*/
@feature("documents.exportEntitiesToFile")
@feature("graphs.exportEntitiesToFile")
async exportEntitiesToFile(options: {
filename: string;
collectionId: string;
Expand All @@ -436,7 +435,7 @@ export class GraphsClient {
}

/**
* Export document relationships as a CSV file with support for filtering and column selection.
* Export graph 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)
Expand All @@ -445,7 +444,7 @@ export class GraphsClient {
* @param options.includeHeader Whether to include column headers (default: true)
* @returns Promise<Blob> in browser environments, Promise<void> in Node.js
*/
@feature("documents.exportRelationships")
@feature("graphs.exportRelationships")
async exportRelationships(options: {
collectionId: string;
outputPath?: string;
Expand All @@ -466,7 +465,7 @@ export class GraphsClient {

const response = await this.client.makeRequest(
"POST",
`documents/${options.collectionId}/relationships/export`,
`graphs/${options.collectionId}/relationships/export`,
{
data,
responseType: "arraybuffer",
Expand All @@ -485,11 +484,11 @@ export class GraphsClient {
}

/**
* Export document relationships as a CSV file and save it to the user's device.
* Export graph relationships as a CSV file and save it to the user's device.
* @param filename
* @param options
*/
@feature("documents.exportRelationshipsToFile")
@feature("graphs.exportRelationshipsToFile")
async exportRelationshipsToFile(options: {
filename: string;
collectionId: string;
Expand All @@ -503,6 +502,74 @@ export class GraphsClient {
}
}

/**
* Export graph communities 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("graphs.exportCommunities")
async exportCommunities(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",
`graphs/${options.collectionId}/communities/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 graph communities as a CSV file and save it to the user's device.
* @param filename
* @param options
*/
@feature("graphs.exportCommunitiesToFile")
async exportCommunitiesToFile(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 8226c15

Please sign in to comment.