From 0186af06dbbad150e2a9174efdba826538deb432 Mon Sep 17 00:00:00 2001 From: satoshi toyama Date: Wed, 11 Dec 2024 20:12:57 +0900 Subject: [PATCH 1/2] refactor(playground): Update action handler to fetch graph URL from DB Add agent ID lookup and remove direct graph URL passing for better security and data consistency. Centralizes graph URL storage in database instead of passing it through the client. - Replace graphUrl parameter with agentId in action handler - Add DB query to fetch agent's graph URL - Add AgentId type definition - Update page component to pass agentId instead of graphUrl --- app/(playground)/p/[agentId]/canary/actions.ts | 13 +++++++++++-- app/(playground)/p/[agentId]/canary/page.tsx | 5 ++--- app/(playground)/p/[agentId]/canary/types.ts | 2 ++ 3 files changed, 15 insertions(+), 5 deletions(-) diff --git a/app/(playground)/p/[agentId]/canary/actions.ts b/app/(playground)/p/[agentId]/canary/actions.ts index 47b942e6..a7881bd8 100644 --- a/app/(playground)/p/[agentId]/canary/actions.ts +++ b/app/(playground)/p/[agentId]/canary/actions.ts @@ -1,5 +1,6 @@ "use server"; +import { db } from "@/drizzle"; import { anthropic } from "@ai-sdk/anthropic"; import { google } from "@ai-sdk/google"; import { openai } from "@ai-sdk/openai"; @@ -16,6 +17,7 @@ import * as v from "valibot"; import { vercelBlobFileFolder, vercelBlobGraphFolder } from "./constants"; import { textGenerationPrompt } from "./prompts"; import type { + AgentId, ArtifactId, FileData, FileId, @@ -106,7 +108,7 @@ type ActionSource = TextSource | TextGenerationSource | FileSource; export async function action( artifactId: ArtifactId, - graphUrl: string, + agentId: AgentId, nodeId: NodeId, ) { const lf = new Langfuse(); @@ -114,7 +116,14 @@ export async function action( sessionId: artifactId, }); - const graph = await fetch(graphUrl).then( + const agent = await db.query.agents.findFirst({ + where: (agents, { eq }) => eq(agents.id, agentId), + }); + if (agent === undefined || agent.graphUrl === null) { + throw new Error(`Agent with id ${agentId} not found`); + } + + const graph = await fetch(agent.graphUrl).then( (res) => res.json() as unknown as Graph, ); const node = graph.nodes.find((node) => node.id === nodeId); diff --git a/app/(playground)/p/[agentId]/canary/page.tsx b/app/(playground)/p/[agentId]/canary/page.tsx index d6b3d562..cfba4e59 100644 --- a/app/(playground)/p/[agentId]/canary/page.tsx +++ b/app/(playground)/p/[agentId]/canary/page.tsx @@ -4,7 +4,6 @@ import { del, list } from "@vercel/blob"; import { ReactFlowProvider } from "@xyflow/react"; import { eq } from "drizzle-orm"; import { notFound } from "next/navigation"; -import type { AgentId } from "../beta-proto/types"; import { action, putGraph } from "./actions"; import { Editor } from "./components/editor"; import { AgentNameProvider } from "./contexts/agent-name"; @@ -16,7 +15,7 @@ import { PropertiesPanelProvider } from "./contexts/properties-panel"; import { ToastProvider } from "./contexts/toast"; import { ToolbarContextProvider } from "./contexts/toolbar"; import { isLatestVersion, migrateGraph } from "./graph"; -import type { ArtifactId, Graph, NodeId } from "./types"; +import type { AgentId, ArtifactId, Graph, NodeId } from "./types"; import { buildGraphFolderPath } from "./utils"; // Extend the max duration of the server actions from this page to 5 minutes @@ -94,7 +93,7 @@ export default async function Page({ nodeId: NodeId, ) { "use server"; - return await action(artifactId, graphUrl, nodeId); + return await action(artifactId, agentId, nodeId); } return ( diff --git a/app/(playground)/p/[agentId]/canary/types.ts b/app/(playground)/p/[agentId]/canary/types.ts index 2b26c241..3a23a5c4 100644 --- a/app/(playground)/p/[agentId]/canary/types.ts +++ b/app/(playground)/p/[agentId]/canary/types.ts @@ -212,3 +212,5 @@ export interface SubGraph { nodes: Set; connections: Set; } + +export type AgentId = `agnt_${string}`; From cce94aa476c30265fc62fdf5d3ae1f38a2d308e9 Mon Sep 17 00:00:00 2001 From: satoshi toyama Date: Wed, 11 Dec 2024 20:19:19 +0900 Subject: [PATCH 2/2] refactor(playground): Simplify execution API by removing graphUrl param Removes the graphUrl parameter from execution flow since it's now fetched directly from the database using agentId. This simplifies the API and prevents potential URL manipulation. - Remove graphUrl parameter from ExecutionProvider interface - Update execute function signature in page component - Remove unnecessary graphUrl passing in execution context --- app/(playground)/p/[agentId]/canary/contexts/execution.tsx | 5 ++--- app/(playground)/p/[agentId]/canary/page.tsx | 6 +----- 2 files changed, 3 insertions(+), 8 deletions(-) diff --git a/app/(playground)/p/[agentId]/canary/contexts/execution.tsx b/app/(playground)/p/[agentId]/canary/contexts/execution.tsx index 36195b12..438930f0 100644 --- a/app/(playground)/p/[agentId]/canary/contexts/execution.tsx +++ b/app/(playground)/p/[agentId]/canary/contexts/execution.tsx @@ -20,7 +20,6 @@ interface ExecutionProviderProps { children: ReactNode; executeAction: ( artifactId: ArtifactId, - graphUrl: string, nodeId: NodeId, ) => Promise>; } @@ -56,9 +55,9 @@ export function ExecutionProvider({ }, }); setTab("Result"); - const latestGraphUrl = await flush(); + await flush(); try { - const stream = await executeAction(artifactId, latestGraphUrl, nodeId); + const stream = await executeAction(artifactId, nodeId); let textArtifactObject: TextArtifactObject = { type: "text", diff --git a/app/(playground)/p/[agentId]/canary/page.tsx b/app/(playground)/p/[agentId]/canary/page.tsx index cfba4e59..96958ce5 100644 --- a/app/(playground)/p/[agentId]/canary/page.tsx +++ b/app/(playground)/p/[agentId]/canary/page.tsx @@ -87,11 +87,7 @@ export default async function Page({ return agentName; } - async function execute( - artifactId: ArtifactId, - graphUrl: string, - nodeId: NodeId, - ) { + async function execute(artifactId: ArtifactId, nodeId: NodeId) { "use server"; return await action(artifactId, agentId, nodeId); }