cognee/cognee-frontend/src/modules/notebooks/useNotebooks.ts

110 lines
3.1 KiB
TypeScript

import { useCallback, useState } from "react";
import { Cell, Notebook } from "@/ui/elements/Notebook/types";
import { CogneeInstance } from "@/modules/instances/types";
import createNotebook from "./createNotebook";
import deleteNotebook from "./deleteNotebook";
import getNotebooks from "./getNotebooks";
import runNotebookCell from "./runNotebookCell";
import { default as persistNotebook } from "./saveNotebook";
function useNotebooks(instance: CogneeInstance) {
const [notebooks, setNotebooks] = useState<Notebook[]>([]);
const addNotebook = useCallback((notebookName: string) => {
return createNotebook(notebookName, instance)
.then((notebook: Notebook) => {
setNotebooks((notebooks) => [
...notebooks,
notebook,
]);
return notebook;
});
}, [instance]);
const removeNotebook = useCallback((notebookId: string) => {
return deleteNotebook(notebookId, instance)
.then(() => {
setNotebooks((notebooks) =>
notebooks.filter((notebook) => notebook.id !== notebookId)
);
});
}, [instance]);
const fetchNotebooks = useCallback(() => {
return getNotebooks(instance)
.then((notebooks) => {
setNotebooks(notebooks);
return notebooks;
})
.catch((error) => {
console.error("Error fetching notebooks:", error.detail);
throw error
});
}, [instance]);
const updateNotebook = useCallback((updatedNotebook: Notebook) => {
setNotebooks((existingNotebooks) =>
existingNotebooks.map((notebook) =>
notebook.id === updatedNotebook.id
? updatedNotebook
: notebook
)
);
}, []);
const saveNotebook = useCallback((notebook: Notebook) => {
return persistNotebook(notebook.id, {
name: notebook.name,
cells: notebook.cells,
}, instance);
}, [instance]);
const runCell = useCallback((notebook: Notebook, cell: Cell) => {
setNotebooks((existingNotebooks) =>
existingNotebooks.map((existingNotebook) =>
existingNotebook.id === notebook.id ? {
...existingNotebook,
cells: existingNotebook.cells.map((existingCell) =>
existingCell.id === cell.id ? {
...existingCell,
result: undefined,
error: undefined,
} : existingCell
),
} : existingNotebook
)
);
return runNotebookCell(notebook.id, cell, instance)
.then((response) => {
setNotebooks((existingNotebooks) =>
existingNotebooks.map((existingNotebook) =>
existingNotebook.id === notebook.id ? {
...existingNotebook,
cells: existingNotebook.cells.map((existingCell) =>
existingCell.id === cell.id ? {
...existingCell,
result: response.result,
error: response.error,
} : existingCell
),
} : existingNotebook
)
);
});
}, [instance]);
return {
notebooks,
addNotebook,
saveNotebook,
updateNotebook,
removeNotebook,
refreshNotebooks: fetchNotebooks,
runCell,
};
};
export default useNotebooks;