Skip to content
This repository has been archived by the owner on Nov 23, 2023. It is now read-only.

Commit

Permalink
Merge pull request #221 from COS301-SE-2023/fix/ai-lang-test
Browse files Browse the repository at this point in the history
Fix some tests
  • Loading branch information
ArmandKrynauw authored Sep 27, 2023
2 parents 2a4b37e + 79009bb commit ca98ded
Show file tree
Hide file tree
Showing 6 changed files with 313 additions and 408 deletions.
2 changes: 1 addition & 1 deletion blix-plugins/blink/src/main.cjs
Original file line number Diff line number Diff line change
Expand Up @@ -262,7 +262,7 @@ const nodes = {
componentId: "imagePicker",
label: "Pick an image",
defaultValue: "",
triggerUpdate: true,
triggerUpdate: true,
}, {});
// ui.addCachePicker({
// componentId: "cachePicker",
Expand Down
6 changes: 3 additions & 3 deletions src/electron/lib/ai/AiManagerv2.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ import { NodeInstance, ToolboxRegistry } from "../registries/ToolboxRegistry";
import { CoreGraphManager } from "../core-graph/CoreGraphManager";
import type { MainWindow } from "../api/apis/WindowApi";
import {
BlypescriptExportStrategyV2,
BlypescriptExportStrategy,
CoreGraphExporter,
} from "../../lib/core-graph/CoreGraphExporter";
import { readFileSync, writeFileSync } from "fs";
Expand Down Expand Up @@ -55,15 +55,15 @@ export class AiManager {
) {
this.blypescriptToolbox = BlypescriptToolbox.fromToolbox(this.toolbox);
this.graphExporter = new CoreGraphExporter(
new BlypescriptExportStrategyV2(this.blypescriptToolbox)
new BlypescriptExportStrategy(this.blypescriptToolbox)
);
this.blypescriptInterpreter = new BlypescriptInterpreter(toolbox, graphManager);
}

async executePrompt({ prompt, graphId, model, apiKey, chatId, verbose }: PromptOptions) {
this.blypescriptToolbox = BlypescriptToolbox.fromToolbox(this.toolbox);
this.graphExporter = new CoreGraphExporter(
new BlypescriptExportStrategyV2(this.blypescriptToolbox)
new BlypescriptExportStrategy(this.blypescriptToolbox)
);
// console.log(this.blypescriptToolbox.toString());
const blypescriptProgram = this.graphExporter.exportGraph(this.graphManager.getGraph(graphId));
Expand Down
5 changes: 2 additions & 3 deletions src/electron/lib/ai/ai-profiler-v2.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@ import { PackageData } from "../../lib/plugins/PluginManager";
import { NodePluginContext, Plugin } from "../../lib/plugins/Plugin";
import {
BlypescriptExportStrategy,
BlypescriptExportStrategyV2,
CoreGraphExporter,
} from "../../lib/core-graph/CoreGraphExporter";
import { BlypescriptProgram, BlypescriptToolbox, colorString } from "./AiLang";
Expand Down Expand Up @@ -47,7 +46,7 @@ export class Profiler {
const blypescriptToolbox = BlypescriptToolbox.fromToolbox(this.toolboxRegistry);

// const blypescriptExporter = new BlypescriptExportStrategy(this.toolboxRegistry);
const blypescriptExporter = new BlypescriptExportStrategyV2(blypescriptToolbox);
const blypescriptExporter = new BlypescriptExportStrategy(blypescriptToolbox);
const coreGraphExporter = new CoreGraphExporter(blypescriptExporter);
const blypescriptProgram = coreGraphExporter.exportGraph(coreGraph);

Expand Down Expand Up @@ -219,7 +218,7 @@ export class Profiler {
const blypescriptToolbox = BlypescriptToolbox.fromToolbox(this.toolboxRegistry);

const coreGraph = this.graphManager.getGraph(graphId);
const blypescriptExporter = new BlypescriptExportStrategyV2(blypescriptToolbox);
const blypescriptExporter = new BlypescriptExportStrategy(blypescriptToolbox);
const coreGraphExporter = new CoreGraphExporter(blypescriptExporter);
const blypescriptProgram = coreGraphExporter.exportGraph(coreGraph);

Expand Down
214 changes: 1 addition & 213 deletions src/electron/lib/core-graph/CoreGraphExporter.ts
Original file line number Diff line number Diff line change
Expand Up @@ -151,226 +151,14 @@ export class LLMExportStrategy implements ExportStrategy<LLMGraph> {
}

// ==================================================================
// Blypescript Exporter v1
// Blypescript Exporter
// ==================================================================

export class BlypescriptExportStrategy implements ExportStrategy<BlypescriptProgram> {
private nodeOccurrenceMap!: Map<string, number>;
private nodeIdNameMap!: Map<UUID, string>;
private graph!: CoreGraph;

constructor(private readonly toolbox: ToolboxRegistry) {}

public export(graph: CoreGraph) {
this.nodeOccurrenceMap = new Map<string, number>();
this.nodeIdNameMap = new Map<UUID, string>();
this.graph = graph;

const statements: BlypescriptStatement[] = [];
const nodeStack: Node[] = [];
const graphNodes = Object.values(graph.getNodes);
let remainingNodes = [...graphNodes];
const numNodes = graphNodes.length;

// Traverse each graph
while (this.nodeIdNameMap.size !== numNodes) {
const remainingNode = remainingNodes.pop();
const sourceStatements: BlypescriptStatement[] = [];
const internalStatements: BlypescriptStatement[] = [];
const terminalStatements: BlypescriptStatement[] = [];

nodeStack.push(...this.findSubgraphSources(remainingNode, new Set()));

// Traverse each subgraph
while (nodeStack.length !== 0) {
const node = nodeStack.shift()!;

if (this.nodeIdNameMap.has(node.uuid)) {
continue;
}

const statement = this.createStatement(node);
const nodeNeighbors = this.findNodeNeighbors(node);

if (nodeNeighbors.inputNeighbors.length && nodeNeighbors.outputNeighbors.length) {
internalStatements.push(statement);
} else if (nodeNeighbors.inputNeighbors.length) {
terminalStatements.push(statement);
} else {
sourceStatements.push(statement);
}

nodeStack.push(...nodeNeighbors.outputNeighbors);
remainingNodes = remainingNodes.filter((n) => n.uuid !== node.uuid);
}

statements.push(...sourceStatements, ...internalStatements, ...terminalStatements);
}

const nodeNameIdMap = new Map<string, UUID>();

for (const [key, value] of this.nodeIdNameMap.entries()) {
nodeNameIdMap.set(value, key);
}

return new BlypescriptProgram(statements, nodeNameIdMap);
}

private getOutputNodes(): Node[] {
const outputIds = Object.keys(this.graph.getOutputNodes);
return outputIds.map((id) => this.graph.getNodes[id]);
}

/**
* Finds the terminal nodes within a subgraph starting from the given node in a core graph.
* Terminal nodes are nodes that have no outgoing connections.
*
* @param node - The starting node.
* @param graph - The core graph.
* @returns An array of terminal nodes within the subgraph.
*/
// private findSubgraphTerminals(node: Node | undefined, graph: CoreGraph): Node[] {
// if (!node) {
// return [];
// }

// const outputNeighbors: Node[] = [];
// const nodeAnchors = Object.values(node.getAnchors);
// const nodeOutputAnchors = nodeAnchors.filter((a) => a.ioType === AnchorIO.output);

// for (const anchor of nodeOutputAnchors) {
// outputNeighbors.push(...this.findSubgraphTerminals(anchor.parent, graph));
// }

// if (outputNeighbors.length === 0) {
// return [node];
// }

// return outputNeighbors;
// }

/**
* Finds the source nodes within a subgraph starting from the given node in a core graph.
* Source nodes are nodes at the start of the directed graph.
*
* @param node - The starting node.
* @param graph - The core graph.
* @returns An array of source nodes within the subgraph.
*/
private findSubgraphSources(node: Node | undefined, visited: Set<UUID>): Node[] {
if (!node) {
return [];
}

visited.add(node.uuid);

const { inputNeighbors, outputNeighbors } = this.findNodeNeighbors(node);
const neighbors = [...inputNeighbors, ...outputNeighbors];
const sourceNodes: Node[] = [];

if (inputNeighbors.length === 0) {
sourceNodes.push(node);
}

for (const neighbor of neighbors) {
if (!visited.has(neighbor.uuid)) {
sourceNodes.push(...this.findSubgraphSources(neighbor, visited));
}
}

return sourceNodes;
}

private findNodeNeighbors(node: Node) {
const inputNeighbors: Record<UUID, Node> = {};
const outputNeighbors: Record<UUID, Node> = {};
const anchors = Object.values(node.getAnchors);

for (const anchor of anchors) {
if (anchor.ioType === AnchorIO.input) {
const outputAnchorId = this.graph.getEdgeDest[anchor.uuid]?.getAnchorFrom;

if (outputAnchorId) {
const outputAnchor = this.graph.getAnchors[outputAnchorId];
inputNeighbors[outputAnchor.parent.uuid] = outputAnchor.parent;
}
} else {
const inputAnchorIds = this.graph.getEdgeSrc[anchor.uuid];

if (inputAnchorIds) {
for (const inputAnchorId of inputAnchorIds) {
const inputAnchor = this.graph.getAnchors[inputAnchorId];
outputNeighbors[inputAnchor.parent.uuid] = inputAnchor.parent;
}
}
}
}

const neighbors = {
inputNeighbors: Object.values(inputNeighbors),
outputNeighbors: Object.values(outputNeighbors),
};

return neighbors;
}

private createStatement(node: Node): BlypescriptStatement {
let name = node.getName;
const signature = `${node.getPlugin}.${node.getName}` as const;
const inputs: string[] = [];
// Potential look into using node signature instead of node name
const counter = this.nodeOccurrenceMap.get(name) || 0;

this.nodeOccurrenceMap.set(name, counter + 1);
name += `${counter + 1}`;
this.nodeIdNameMap.set(node.uuid, name);

const nodeInstance = this.toolbox.getNodeInstance(signature);
const inputIdentifiers = nodeInstance.inputs.map((input) => input.id);

// Add node inputs first
for (const identifier of inputIdentifiers) {
const anchor = Object.values(node.getAnchors).find((a) => a.anchorId === identifier);
const inputEdge = this.graph.getEdgeDest[anchor?.uuid || ""];

if (inputEdge) {
const outputAnchor = this.graph.getAnchors[inputEdge.getAnchorFrom];
const outputNode = outputAnchor.parent;
inputs.push(`${this.nodeIdNameMap.get(outputNode.uuid)!}['${outputAnchor.anchorId}']`);
} else {
inputs.push("null");
}
}

const nodeUiInputs = this.graph.getUIInputs(node.uuid);

if (!nodeUiInputs) {
return new BlypescriptStatement(name, signature, inputs);
}

const uiConfigs = Object.values(nodeInstance.uiConfigs);
const uiInputIdentifiers = uiConfigs.map((config) => config.componentId);

// Then add UI inputs
for (const identifier of uiInputIdentifiers) {
const uiValue = nodeUiInputs[identifier];
// Improve to work with more than just string an ints
inputs.push(typeof uiValue === "string" ? `'${uiValue}'` : `${uiValue as string}`);
}

return new BlypescriptStatement(name, signature, inputs);
}
}

// ==================================================================
// Blypescript Exporter v2
// ==================================================================

export class BlypescriptExportStrategyV2 implements ExportStrategy<BlypescriptProgram> {
private nodeOccurrenceMap!: Map<string, number>;
private nodeIdNameMap!: Map<UUID, string>;
private graph!: CoreGraph;

constructor(private readonly toolbox: BlypescriptToolbox) {}

public export(graph: CoreGraph) {
Expand Down
87 changes: 86 additions & 1 deletion tests/unit-tests/electron/lib/ai/AiLang.spec.ts
Original file line number Diff line number Diff line change
@@ -1,10 +1,20 @@
import { BlypescriptProgram, BlypescriptToolbox } from "../../../../../src/electron/lib/ai/AiLang";
import {
BlypescriptInterpreter,
BlypescriptProgram,
BlypescriptToolbox,
} from "../../../../../src/electron/lib/ai/AiLang";
import { ToolboxRegistry } from "../../../../../src/electron/lib/registries/ToolboxRegistry";
import { createOutputNode } from "../../../../../src/electron/lib/Blix";
import { join } from "path";
import { readFileSync, readdirSync } from "fs";
import { PackageData } from "../../../../../src/electron/lib/plugins/PluginManager";
import { Plugin, NodePluginContext } from "../../../../../src/electron/lib/plugins/Plugin";
import {
BlypescriptExportStrategy,
CoreGraphExporter,
} from "../../../../../src/electron/lib/core-graph/CoreGraphExporter";
import { CoreGraph } from "../../../../../src/electron/lib/core-graph/CoreGraph";
import { CoreGraphManager } from "../../../../../src/electron/lib/core-graph/CoreGraphManager";

describe("Blypescript parser", () => {
test("should parse an empty graph", () => {
Expand Down Expand Up @@ -33,6 +43,81 @@ describe("Blypescript parser", () => {
});
});

describe("Blypescript interpreter", () => {
test("add nodes to empty graph", () => {
const toolbox = generateToolbox(["glfx", "input", "blink"]);
const blypescriptToolbox = BlypescriptToolbox.fromToolbox(toolbox);
const graph = new CoreGraph();
const coreGraphManager = new CoreGraphManager(toolbox);
const blypescriptInterpreter = new BlypescriptInterpreter(toolbox, coreGraphManager);
const exporter = new CoreGraphExporter(new BlypescriptExportStrategy(blypescriptToolbox));

coreGraphManager.addGraph(graph);

let response1 = graph.addNode(toolbox.getNodeInstance("input.number"), { x: 0, y: 0 });
const uuid1 = response1.data!.nodeId;
response1 = graph.addNode(toolbox.getNodeInstance("blix.output"), { x: 0, y: 0 });
const uuid2 = response1.data!.nodeId;
response1 = graph.addNode(toolbox.getNodeInstance("blink.inputImage"), { x: 0, y: 0 });
const uuid3 = response1.data!.nodeId;

const node1 = graph.getNodes[uuid1];
const node2 = graph.getNodes[uuid2];
const node3 = graph.getNodes[uuid3];

const anchorFrom = Object.keys(node3.getAnchors)[1];
const anchorTo = Object.keys(node2.getAnchors)[0];
let response2 = graph.addEdge(anchorFrom, anchorTo);

const result2 = BlypescriptProgram.fromString(
`graph() {
const inputGLFXImage1 = glfx.GLFXImage();
const output = blix.output(inputGLFXImage1['res'], 'output');
}`,
blypescriptToolbox
);

expect(result2.success).toBe(true);
const program1 = exporter.exportGraph(graph);
const program2 = result2.data as BlypescriptProgram;

blypescriptInterpreter.run(graph.uuid, program1, program2);
});
});

describe("BlypescriptExporter", () => {
let toolbox: ToolboxRegistry;
let blypescriptToolbox: BlypescriptToolbox;
let exporter: CoreGraphExporter<BlypescriptProgram>;
let graph: CoreGraph;

beforeEach(() => {
toolbox = generateToolbox(["glfx", "input"]);
blypescriptToolbox = BlypescriptToolbox.fromToolbox(toolbox);
exporter = new CoreGraphExporter(new BlypescriptExportStrategy(blypescriptToolbox));

graph = new CoreGraph();
let response1 = graph.addNode(toolbox.getNodeInstance("input.number"), { x: 0, y: 0 });
const uuid1 = response1.data!.nodeId;
response1 = graph.addNode(toolbox.getNodeInstance("blix.output"), { x: 0, y: 0 });
const uuid2 = response1.data!.nodeId;
response1 = graph.addNode(toolbox.getNodeInstance("glfx.brightnessContrast"), { x: 0, y: 0 });
const uuid3 = response1.data!.nodeId;

const node1 = graph.getNodes[uuid1];
const node2 = graph.getNodes[uuid2];

const anchorFrom = Object.keys(node1.getAnchors)[0];
const anchorTo = Object.keys(node2.getAnchors)[0];
let response2 = graph.addEdge(anchorFrom, anchorTo);
});

test("should export a graph to blypescript", () => {
const blypescriptProgram = exporter.exportGraph(graph);
expect(blypescriptProgram).toBeDefined();
});
});

// ===================================================================
// HELPERS
// ===================================================================
Expand Down
Loading

0 comments on commit ca98ded

Please sign in to comment.