feat(platform): Add delete agent functionality (#8273)

pull/8243/head^2
Zamil Majdy 2024-10-08 19:03:26 +03:00 committed by GitHub
parent d42ed088dd
commit 2a74381ae8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 130 additions and 10 deletions

View File

@ -500,6 +500,15 @@ async def get_graph_all_versions(graph_id: str, user_id: str) -> list[Graph]:
return [Graph.from_db(graph) for graph in graph_versions]
async def delete_graph(graph_id: str, user_id: str) -> int:
entries_count = await AgentGraph.prisma().delete_many(
where={"id": graph_id, "userId": user_id}
)
if entries_count:
logger.info(f"Deleted {entries_count} graph entries for Graph #{graph_id}")
return entries_count
async def create_graph(graph: Graph, user_id: str) -> Graph:
async with transaction() as tx:
await __create_graph(tx, graph, user_id)

View File

@ -10,6 +10,7 @@ from autogpt_libs.auth.middleware import auth_middleware
from fastapi import APIRouter, Body, Depends, FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from typing_extensions import TypedDict
from backend.data import block, db
from backend.data import execution as execution_db
@ -168,6 +169,12 @@ class AgentServer(AppService):
methods=["PUT"],
tags=["templates", "graphs"],
)
api_router.add_api_route(
path="/graphs/{graph_id}",
endpoint=self.delete_graph,
methods=["DELETE"],
tags=["graphs"],
)
api_router.add_api_route(
path="/graphs/{graph_id}/versions",
endpoint=self.get_graph_all_versions,
@ -395,6 +402,17 @@ class AgentServer(AppService):
) -> graph_db.Graph:
return await cls.create_graph(create_graph, is_template=True, user_id=user_id)
class DeleteGraphResponse(TypedDict):
version_counts: int
@classmethod
async def delete_graph(
cls, graph_id: str, user_id: Annotated[str, Depends(get_user_id)]
) -> DeleteGraphResponse:
return {
"version_counts": await graph_db.delete_graph(graph_id, user_id=user_id)
}
@classmethod
async def create_graph(
cls,

View File

@ -0,0 +1,5 @@
-- DropForeignKey
ALTER TABLE "AgentGraph" DROP CONSTRAINT "AgentGraph_agentGraphParentId_version_fkey";
-- AddForeignKey
ALTER TABLE "AgentGraph" ADD CONSTRAINT "AgentGraph_agentGraphParentId_version_fkey" FOREIGN KEY ("agentGraphParentId", "version") REFERENCES "AgentGraph"("id", "version") ON DELETE CASCADE ON UPDATE CASCADE;

View File

@ -53,7 +53,7 @@ model AgentGraph {
// All sub-graphs are defined within this 1-level depth list (even if it's a nested graph).
AgentSubGraphs AgentGraph[] @relation("AgentSubGraph")
agentGraphParentId String?
AgentGraphParent AgentGraph? @relation("AgentSubGraph", fields: [agentGraphParentId, version], references: [id, version])
AgentGraphParent AgentGraph? @relation("AgentSubGraph", fields: [agentGraphParentId, version], references: [id, version], onDelete: Cascade)
@@id(name: "graphVersionId", [id, version])
}
@ -63,7 +63,7 @@ model AgentNode {
id String @id @default(uuid())
agentBlockId String
AgentBlock AgentBlock @relation(fields: [agentBlockId], references: [id])
AgentBlock AgentBlock @relation(fields: [agentBlockId], references: [id], onUpdate: Cascade)
agentGraphId String
agentGraphVersion Int @default(1)

View File

@ -7,3 +7,28 @@ from backend.util.test import SpinTestServer
async def server():
async with SpinTestServer() as server:
yield server
@pytest.fixture(scope="session", autouse=True)
async def graph_cleanup(server):
created_graph_ids = []
original_create_graph = server.agent_server.create_graph
async def create_graph_wrapper(*args, **kwargs):
created_graph = await original_create_graph(*args, **kwargs)
# Extract user_id correctly
user_id = kwargs.get("user_id", args[2] if len(args) > 2 else None)
created_graph_ids.append((created_graph.id, user_id))
return created_graph
try:
server.agent_server.create_graph = create_graph_wrapper
yield # This runs the test function
finally:
server.agent_server.create_graph = original_create_graph
# Delete the created graphs and assert they were deleted
for graph_id, user_id in created_graph_ids:
resp = await server.agent_server.delete_graph(graph_id, user_id)
num_deleted = resp["version_counts"]
assert num_deleted > 0, f"Graph {graph_id} was not deleted."

View File

@ -5,10 +5,15 @@ from backend.blocks.basic import FindInDictionaryBlock, StoreValueBlock
from backend.blocks.maths import CalculatorBlock, Operation
from backend.data import execution, graph
from backend.server import AgentServer
from backend.server.model import CreateGraph
from backend.usecases.sample import create_test_graph, create_test_user
from backend.util.test import SpinTestServer, wait_execution
async def create_graph(s: SpinTestServer, g: graph.Graph, u: User) -> graph.Graph:
return await s.agent_server.create_graph(CreateGraph(graph=g), False, u.id)
async def execute_graph(
agent_server: AgentServer,
test_graph: graph.Graph,
@ -99,9 +104,8 @@ async def assert_sample_graph_executions(
@pytest.mark.asyncio(scope="session")
async def test_agent_execution(server: SpinTestServer):
test_graph = create_test_graph()
test_user = await create_test_user()
await graph.create_graph(test_graph, user_id=test_user.id)
test_graph = await create_graph(server, create_test_graph(), test_user)
data = {"input_1": "Hello", "input_2": "World"}
graph_exec_id = await execute_graph(
server.agent_server,
@ -163,7 +167,7 @@ async def test_input_pin_always_waited(server: SpinTestServer):
links=links,
)
test_user = await create_test_user()
test_graph = await graph.create_graph(test_graph, user_id=test_user.id)
test_graph = await create_graph(server, test_graph, test_user)
graph_exec_id = await execute_graph(
server.agent_server, test_graph, test_user, {}, 3
)
@ -244,7 +248,7 @@ async def test_static_input_link_on_graph(server: SpinTestServer):
links=links,
)
test_user = await create_test_user()
test_graph = await graph.create_graph(test_graph, user_id=test_user.id)
test_graph = await create_graph(server, test_graph, test_user)
graph_exec_id = await execute_graph(
server.agent_server, test_graph, test_user, {}, 8
)

View File

@ -1,7 +1,8 @@
import pytest
from backend.data import db, graph
from backend.data import db
from backend.executor import ExecutionScheduler
from backend.server.model import CreateGraph
from backend.usecases.sample import create_test_graph, create_test_user
from backend.util.service import get_service_client
from backend.util.settings import Config
@ -12,7 +13,11 @@ from backend.util.test import SpinTestServer
async def test_agent_schedule(server: SpinTestServer):
await db.connect()
test_user = await create_test_user()
test_graph = await graph.create_graph(create_test_graph(), user_id=test_user.id)
test_graph = await server.agent_server.create_graph(
create_graph=CreateGraph(graph=create_test_graph()),
is_template=False,
user_id=test_user.id,
)
scheduler = get_service_client(
ExecutionScheduler, Config().execution_scheduler_port

View File

@ -90,6 +90,11 @@ const Monitor = () => {
flow={selectedFlow}
flowRuns={flowRuns.filter((r) => r.graphID == selectedFlow.id)}
className={column3}
refresh={() => {
fetchAgents();
setSelectedFlow(null);
setSelectedRun(null);
}}
/>
)) || (
<Card className={`p-6 ${column3}`}>

View File

@ -20,14 +20,24 @@ import { ClockIcon, ExitIcon, Pencil2Icon } from "@radix-ui/react-icons";
import Link from "next/link";
import { exportAsJSONFile } from "@/lib/utils";
import { FlowRunsStats } from "@/components/monitor/index";
import { Trash2Icon } from "lucide-react";
import {
Dialog,
DialogContent,
DialogHeader,
DialogTitle,
DialogDescription,
DialogFooter,
} from "@/components/ui/dialog";
export const FlowInfo: React.FC<
React.HTMLAttributes<HTMLDivElement> & {
flow: GraphMeta;
flowRuns: FlowRun[];
flowVersion?: number | "all";
refresh: () => void;
}
> = ({ flow, flowRuns, flowVersion, ...props }) => {
> = ({ flow, flowRuns, flowVersion, refresh, ...props }) => {
const api = useMemo(() => new AutoGPTServerAPI(), []);
const [flowVersions, setFlowVersions] = useState<Graph[] | null>(null);
@ -39,6 +49,8 @@ export const FlowInfo: React.FC<
v.version == (selectedVersion == "all" ? flow.version : selectedVersion),
);
const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false);
useEffect(() => {
api.getGraphAllVersions(flow.id).then((result) => setFlowVersions(result));
}, [flow.id, api]);
@ -96,7 +108,7 @@ export const FlowInfo: React.FC<
className={buttonVariants({ variant: "outline" })}
href={`/build?flowID=${flow.id}`}
>
<Pencil2Icon className="mr-2" /> Edit
<Pencil2Icon />
</Link>
<Button
variant="outline"
@ -116,6 +128,9 @@ export const FlowInfo: React.FC<
>
<ExitIcon />
</Button>
<Button variant="outline" onClick={() => setIsDeleteModalOpen(true)}>
<Trash2Icon className="h-full" />
</Button>
</div>
</CardHeader>
<CardContent>
@ -128,6 +143,36 @@ export const FlowInfo: React.FC<
)}
/>
</CardContent>
<Dialog open={isDeleteModalOpen} onOpenChange={setIsDeleteModalOpen}>
<DialogContent>
<DialogHeader>
<DialogTitle>Delete Agent</DialogTitle>
<DialogDescription>
Are you sure you want to delete this agent? <br />
This action cannot be undone.
</DialogDescription>
</DialogHeader>
<DialogFooter>
<Button
variant="outline"
onClick={() => setIsDeleteModalOpen(false)}
>
Cancel
</Button>
<Button
variant="destructive"
onClick={() => {
api.deleteGraph(flow.id).then(() => {
setIsDeleteModalOpen(false);
refresh();
});
}}
>
Delete
</Button>
</DialogFooter>
</DialogContent>
</Dialog>
</Card>
);
};

View File

@ -124,6 +124,10 @@ export default class BaseAutoGPTServerAPI {
return this._request("PUT", `/templates/${id}`, template);
}
deleteGraph(id: string): Promise<void> {
return this._request("DELETE", `/graphs/${id}`);
}
setGraphActiveVersion(id: string, version: number): Promise<Graph> {
return this._request("PUT", `/graphs/${id}/versions/active`, {
active_graph_version: version,