Notebook.tsx•17.1 kB
"use client";
import { v4 as uuid4 } from "uuid";
import classNames from "classnames";
import { Fragment, MouseEvent, MutableRefObject, useCallback, useEffect, useRef, useState } from "react";
import { useModal } from "@/ui/elements/Modal";
import { CaretIcon, CloseIcon, PlusIcon } from "@/ui/Icons";
import { IconButton, PopupMenu, TextArea, Modal, GhostButton, CTAButton } from "@/ui/elements";
import { GraphControlsAPI } from "@/app/(graph)/GraphControls";
import GraphVisualization, { GraphVisualizationAPI } from "@/app/(graph)/GraphVisualization";
import NotebookCellHeader from "./NotebookCellHeader";
import { Cell, Notebook as NotebookType } from "./types";
interface NotebookProps {
notebook: NotebookType;
runCell: (notebook: NotebookType, cell: Cell, cogneeInstance: string) => Promise<void>;
updateNotebook: (updatedNotebook: NotebookType) => void;
}
export default function Notebook({ notebook, updateNotebook, runCell }: NotebookProps) {
useEffect(() => {
if (notebook.cells.length === 0) {
const newCell: Cell = {
id: uuid4(),
name: "first cell",
type: "code",
content: "",
};
updateNotebook({
...notebook,
cells: [newCell],
});
toggleCellOpen(newCell.id)
}
}, [notebook, updateNotebook]);
const handleCellRun = useCallback((cell: Cell, cogneeInstance: string) => {
return runCell(notebook, cell, cogneeInstance);
}, [notebook, runCell]);
const handleCellAdd = useCallback((afterCellIndex: number, cellType: "markdown" | "code") => {
const newCell: Cell = {
id: uuid4(),
name: "new cell",
type: cellType,
content: "",
};
const newNotebook = {
...notebook,
cells: [
...notebook.cells.slice(0, afterCellIndex + 1),
newCell,
...notebook.cells.slice(afterCellIndex + 1),
],
};
toggleCellOpen(newCell.id);
updateNotebook(newNotebook);
}, [notebook, updateNotebook]);
const removeCell = useCallback((cell: Cell, event?: MouseEvent) => {
event?.preventDefault();
updateNotebook({
...notebook,
cells: notebook.cells.filter((c: Cell) => c.id !== cell.id),
});
}, [notebook, updateNotebook]);
const {
isModalOpen: isRemoveCellConfirmModalOpen,
openModal: openCellRemoveConfirmModal,
closeModal: closeCellRemoveConfirmModal,
confirmAction: handleCellRemoveConfirm,
} = useModal<Cell, MouseEvent>(false, removeCell);
const handleCellRemove = useCallback((cell: Cell) => {
openCellRemoveConfirmModal(cell);
}, [openCellRemoveConfirmModal]);
const handleCellInputChange = useCallback((notebook: NotebookType, cell: Cell, value: string) => {
const newCell = {...cell, content: value };
updateNotebook({
...notebook,
cells: notebook.cells.map((cell: Cell) => (cell.id === newCell.id ? newCell : cell)),
});
}, [updateNotebook]);
const handleCellUp = useCallback((cell: Cell) => {
const index = notebook.cells.indexOf(cell);
if (index > 0) {
const newCells = [...notebook.cells];
newCells[index] = notebook.cells[index - 1];
newCells[index - 1] = cell;
updateNotebook({
...notebook,
cells: newCells,
});
}
}, [notebook, updateNotebook]);
const handleCellDown = useCallback((cell: Cell) => {
const index = notebook.cells.indexOf(cell);
if (index < notebook.cells.length - 1) {
const newCells = [...notebook.cells];
newCells[index] = notebook.cells[index + 1];
newCells[index + 1] = cell;
updateNotebook({
...notebook,
cells: newCells,
});
}
}, [notebook, updateNotebook]);
const handleCellRename = useCallback((cell: Cell) => {
const newName = prompt("Enter a new name for the cell:");
if (newName) {
updateNotebook({
...notebook,
cells: notebook.cells.map((c: Cell) => (c.id === cell.id ? {...c, name: newName } : c)),
});
}
}, [notebook, updateNotebook]);
const [openCells, setOpenCells] = useState(new Set(notebook.cells.map((c: Cell) => c.id)));
const toggleCellOpen = (id: string) => {
setOpenCells((prev) => {
const newState = new Set(prev);
if (newState.has(id)) {
newState.delete(id)
} else {
newState.add(id);
}
return newState;
});
};
return (
<>
<div className="bg-white rounded-xl flex flex-col gap-0.5 px-7 py-5 flex-1">
<div className="mb-5">{notebook.name}</div>
{notebook.cells.map((cell: Cell, index) => (
<Fragment key={cell.id}>
<div key={cell.id} className="flex flex-row rounded-xl border-1 border-gray-100">
<div className="flex flex-col flex-1 relative">
{cell.type === "code" ? (
<>
<div className="absolute left-[-1.35rem] top-2.5">
<IconButton className="p-[0.25rem] m-[-0.25rem]" onClick={toggleCellOpen.bind(null, cell.id)}>
<CaretIcon className={classNames("transition-transform", openCells.has(cell.id) ? "rotate-0" : "rotate-180")} />
</IconButton>
</div>
<NotebookCellHeader
cell={cell}
runCell={handleCellRun}
renameCell={handleCellRename}
removeCell={handleCellRemove}
moveCellUp={handleCellUp}
moveCellDown={handleCellDown}
className="rounded-tl-xl rounded-tr-xl"
/>
{openCells.has(cell.id) && (
<>
<TextArea
value={cell.content}
onChange={handleCellInputChange.bind(null, notebook, cell)}
// onKeyUp={handleCellRunOnEnter}
isAutoExpanding
name="cellInput"
placeholder="Type your code here..."
contentEditable={true}
className="resize-none min-h-36 max-h-96 overflow-y-auto rounded-tl-none rounded-tr-none rounded-bl-xl rounded-br-xl border-0 !outline-0"
/>
<div className="flex flex-col bg-gray-100 overflow-x-auto max-w-full">
{cell.result && (
<div className="px-2 py-2">
output: <CellResult content={cell.result} />
</div>
)}
{!!cell.error?.length && (
<div className="px-2 py-2">
error: {cell.error}
</div>
)}
</div>
</>
)}
</>
) : (
<>
<div className="absolute left-[-1.35rem] top-2.5">
<IconButton className="p-[0.25rem] m-[-0.25rem]" onClick={toggleCellOpen.bind(null, cell.id)}>
<CaretIcon className={classNames("transition-transform", openCells.has(cell.id) ? "rotate-0" : "rotate-180")} />
</IconButton>
</div>
<NotebookCellHeader
cell={cell}
renameCell={handleCellRename}
removeCell={handleCellRemove}
moveCellUp={handleCellUp}
moveCellDown={handleCellDown}
className="rounded-tl-xl rounded-tr-xl"
/>
{openCells.has(cell.id) && (
<TextArea
value={cell.content}
onChange={handleCellInputChange.bind(null, notebook, cell)}
// onKeyUp={handleCellRunOnEnter}
isAutoExpanding
name="cellInput"
placeholder="Type your text here..."
contentEditable={true}
className="resize-none min-h-24 max-h-96 overflow-y-auto rounded-tl-none rounded-tr-none rounded-bl-xl rounded-br-xl border-0 !outline-0"
/>
)}
</>
)}
</div>
</div>
<div className="ml-[-1.35rem]">
<PopupMenu
openToRight={true}
triggerElement={<PlusIcon />}
triggerClassName="p-[0.25rem] m-[-0.25rem]"
>
<div className="flex flex-col gap-0.5">
<button
onClick={() => handleCellAdd(index, "markdown")}
className="hover:bg-gray-100 w-full text-left px-2 cursor-pointer"
>
<span>text</span>
</button>
</div>
<div
onClick={() => handleCellAdd(index, "code")}
className="hover:bg-gray-100 w-full text-left px-2 cursor-pointer"
>
<span>code</span>
</div>
</PopupMenu>
</div>
</Fragment>
))}
</div>
<Modal isOpen={isRemoveCellConfirmModalOpen}>
<div className="w-full max-w-2xl">
<div className="flex flex-row items-center justify-between">
<span className="text-2xl">Delete notebook cell?</span>
<IconButton onClick={closeCellRemoveConfirmModal}><CloseIcon /></IconButton>
</div>
<div className="mt-8 mb-6">Are you sure you want to delete a notebook cell? This action cannot be undone.</div>
<div className="flex flex-row gap-4 mt-4 justify-end">
<GhostButton type="button" onClick={closeCellRemoveConfirmModal}>cancel</GhostButton>
<CTAButton onClick={handleCellRemoveConfirm} type="submit">delete</CTAButton>
</div>
</div>
</Modal>
</>
);
}
function CellResult({ content }: { content: [] }) {
const parsedContent = [];
const graphRef = useRef<GraphVisualizationAPI>();
const graphControls = useRef<GraphControlsAPI>({
setSelectedNode: () => {},
getSelectedNode: () => null,
});
for (const line of content) {
try {
if (Array.isArray(line)) {
// Insights search returns uncommon graph data structure
if (Array.from(line).length > 0 && Array.isArray(line[0]) && line[0][1]["relationship_name"]) {
parsedContent.push(
<div key={line[0][1]["relationship_name"]} className="w-full h-full bg-white">
<span className="text-sm pl-2 mb-4">reasoning graph</span>
<GraphVisualization
data={transformInsightsGraphData(line)}
ref={graphRef as MutableRefObject<GraphVisualizationAPI>}
graphControls={graphControls}
className="min-h-80"
/>
</div>
);
continue;
}
// @ts-expect-error line can be Array or string
for (const item of line) {
if (
typeof item === "object" && item["search_result"] && (typeof(item["search_result"]) === "string"
|| (Array.isArray(item["search_result"]) && typeof(item["search_result"][0]) === "string"))
) {
parsedContent.push(
<div key={String(item["search_result"])} className="w-full h-full bg-white">
<span className="text-sm pl-2 mb-4">query response (dataset: {item["dataset_name"]})</span>
<span className="block px-2 py-2 whitespace-normal">{item["search_result"]}</span>
</div>
);
} else if (typeof(item) === "object" && item["search_result"] && typeof(item["search_result"]) === "object") {
parsedContent.push(
<pre className="px-2 w-full h-full bg-white text-sm" key={String(item).slice(0, -10)}>
{JSON.stringify(item, null, 2)}
</pre>
)
} else if (typeof(item) === "string") {
parsedContent.push(
<pre className="px-2 w-full h-full bg-white text-sm whitespace-normal" key={item.slice(0, -10)}>
{item}
</pre>
);
} else if (typeof(item) === "object" && !(item["search_result"] || item["graphs"])) {
parsedContent.push(
<pre className="px-2 w-full h-full bg-white text-sm" key={String(item).slice(0, -10)}>
{JSON.stringify(item, null, 2)}
</pre>
)
}
if (typeof item === "object" && item["graphs"] && typeof item["graphs"] === "object") {
Object.entries<{ nodes: []; edges: []; }>(item["graphs"]).forEach(([datasetName, graph]) => {
parsedContent.push(
<div key={datasetName} className="w-full h-full bg-white">
<span className="text-sm pl-2 mb-4">reasoning graph (datasets: {datasetName})</span>
<GraphVisualization
data={transformToVisualizationData(graph)}
ref={graphRef as MutableRefObject<GraphVisualizationAPI>}
graphControls={graphControls}
className="min-h-80"
/>
</div>
);
});
}
}
}
if (typeof(line) === "object" && line["result"] && typeof(line["result"]) === "string") {
const datasets = Array.from(
// eslint-disable-next-line @typescript-eslint/no-explicit-any
new Set(Object.values(line["datasets"]).map((dataset: any) => dataset.name))
).join(", ");
parsedContent.push(
<div key={line["result"]} className="w-full h-full bg-white">
<span className="text-sm pl-2 mb-4">query response (datasets: {datasets})</span>
<span className="block px-2 py-2 whitespace-normal">{line["result"]}</span>
</div>
);
}
if (typeof(line) === "object" && line["graphs"]) {
Object.entries<{ nodes: []; edges: []; }>(line["graphs"]).forEach(([datasetName, graph]) => {
parsedContent.push(
<div key={datasetName} className="w-full h-full bg-white">
<span className="text-sm pl-2 mb-4">reasoning graph (datasets: {datasetName})</span>
<GraphVisualization
data={transformToVisualizationData(graph)}
ref={graphRef as MutableRefObject<GraphVisualizationAPI>}
graphControls={graphControls}
className="min-h-80"
/>
</div>
);
});
}
if (typeof(line) === "object" && line["result"] && typeof(line["result"]) === "object") {
parsedContent.push(
<pre className="px-2 w-full h-full bg-white text-sm" key={String(line).slice(0, -10)}>
{JSON.stringify(line["result"], null, 2)}
</pre>
)
}
if (typeof(line) === "string") {
parsedContent.push(
<pre className="px-2 w-full h-full bg-white text-sm whitespace-normal" key={String(line).slice(0, -10)}>
{line}
</pre>
)
}
} catch (error) {
console.error(error);
parsedContent.push(
<pre className="px-2 w-full h-full bg-white text-sm whitespace-normal" key={String(line).slice(0, -10)}>
{line}
</pre>
);
}
}
return parsedContent.map((item, index) => (
<div key={index} className="px-2 py-1">
{item}
</div>
));
};
function transformToVisualizationData(graph: { nodes: [], edges: [] }) {
return {
nodes: graph.nodes,
links: graph.edges,
};
}
type Triplet = [{
id: string,
name: string,
type: string,
}, {
relationship_name: string,
}, {
id: string,
name: string,
type: string,
}]
function transformInsightsGraphData(triplets: Triplet[]) {
const nodes: {
[key: string]: {
id: string,
label: string,
type: string,
}
} = {};
const links: {
[key: string]: {
source: string,
target: string,
label: string,
}
} = {};
for (const triplet of triplets) {
nodes[triplet[0].id] = {
id: triplet[0].id,
label: triplet[0].name || triplet[0].id,
type: triplet[0].type,
};
nodes[triplet[2].id] = {
id: triplet[2].id,
label: triplet[2].name || triplet[2].id,
type: triplet[2].type,
};
const linkKey = `${triplet[0]["id"]}_${triplet[1]["relationship_name"]}_${triplet[2]["id"]}`;
links[linkKey] = {
source: triplet[0].id,
target: triplet[2].id,
label: triplet[1]["relationship_name"],
};
}
return {
nodes: Object.values(nodes),
links: Object.values(links),
};
}