Skip to content

Commit

Permalink
Flav/viz draft (#6486)
Browse files Browse the repository at this point in the history
* Working iframe

Saving progress, dirty code,

File transfer working.

Working. Need UI work + code clean up.

Build working

* Fix file upload on front edge. to be reverted.

* Fix UX

* Tab manually changed.

* Fix UI and rendering.

* Better naming

* Clean up

* Adding front-viz

* Adding comments.

* Clean up console.log

* Remove unused image.

* Draft

* Create .nvmrc

* Add foundations for viz

* ✨

* 🙈

* 🙈

* Fix

* Fix useFile hook

* ✂️

* ✂️

* Remove Sparkle from viz

* ✨

* 🔙

* Use NEXT_PUBLIC_VIZ_URL at build time

* 🙈

* ✂️

* 🙈

* ✂️

* 😑

* 🐵

* 🔨

* 💥

* Only listen to current action

* ✂️

* ✂️

* ✨

* ✂️

---------

Co-authored-by: Aric Lasry <lasry.aric@gmail.com>
  • Loading branch information
flvndvd and lasryaric authored Jul 25, 2024
1 parent 83edeb4 commit 4ab90aa
Show file tree
Hide file tree
Showing 12 changed files with 457 additions and 82 deletions.
5 changes: 4 additions & 1 deletion front/Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,10 @@ RUN npm ci
COPY /front .

ARG COMMIT_HASH
ENV NEXT_PUBLIC_COMMIT_HASH=${COMMIT_HASH}
ARG NEXT_PUBLIC_VIZ_URL

ENV NEXT_PUBLIC_COMMIT_HASH=$COMMIT_HASH
ENV NEXT_PUBLIC_VIZ_URL=$NEXT_PUBLIC_VIZ_URL

# fake database URIs are needed because Sequelize will throw if the `url` parameter
# is undefined, and `next build` imports the `models.ts` file while "Collecting page data"
Expand Down
56 changes: 25 additions & 31 deletions front/components/assistant/conversation/AgentMessage.tsx
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
import {
ArrowPathIcon,
BracesIcon,
Button,
Chip,
Citation,
Expand All @@ -22,7 +21,6 @@ import type {
LightAgentConfigurationType,
RetrievalActionType,
UserType,
VisualizationActionType,
WebsearchActionType,
WebsearchResultType,
WorkspaceType,
Expand All @@ -39,13 +37,15 @@ import {
isWebsearchActionType,
removeNulls,
} from "@dust-tt/types";
import assert from "assert";
import Link from "next/link";
import { useRouter } from "next/router";
import { useCallback, useContext, useEffect, useRef, useState } from "react";

import { makeDocumentCitations } from "@app/components/actions/retrieval/utils";
import { AssistantDetailsDropdownMenu } from "@app/components/assistant/AssistantDetailsDropdownMenu";
import { AgentMessageActions } from "@app/components/assistant/conversation/actions/AgentMessageActions";
import { VisualizationActionIframe } from "@app/components/assistant/conversation/actions/VisualizationActionIframe";
import type { MessageSizeType } from "@app/components/assistant/conversation/ConversationMessage";
import { ConversationMessage } from "@app/components/assistant/conversation/ConversationMessage";
import { GenerationContext } from "@app/components/assistant/conversation/GenerationContextProvider";
Expand Down Expand Up @@ -88,19 +88,9 @@ export function AgentMessage({
const [streamedAgentMessage, setStreamedAgentMessage] =
useState<AgentMessageType>(message);

const defaultVisualizations: VisualizationActionType[] =
message.actions.filter((a): a is VisualizationActionType =>
isVisualizationActionType(a)
) as VisualizationActionType[];

const [streamedVisualizations, setStreamedVisualizations] = useState<
{ actionId: number; visualization: string }[]
>(
defaultVisualizations.map((v) => ({
actionId: v.id,
visualization: v.generation ?? "",
}))
);
>([]);

const [isRetryHandlerProcessing, setIsRetryHandlerProcessing] =
useState<boolean>(false);
Expand Down Expand Up @@ -213,6 +203,7 @@ export function AgentMessage({
...event.message,
};
});
setStreamedVisualizations([]);
break;
}

Expand Down Expand Up @@ -535,24 +526,6 @@ export function AgentMessage({
</div>
) : null}

{/* This is where we will we plug Aric's work to render the graph in an iframe. */}
{streamedVisualizations.map(({ actionId, visualization }) => {
return (
<div key={actionId}>
<div className="flex flex-row gap-2">
<Icon size="sm" visual={BracesIcon} />
<div className="font-semibold">Visualization</div>
</div>
<div>
<RenderMessageMarkdown
content={"```js" + visualization + "\n```"}
isStreaming={true}
/>
</div>
</div>
);
})}

{agentMessage.content !== null && (
<div>
{lastTokenClassification !== "chain_of_thought" &&
Expand Down Expand Up @@ -583,6 +556,27 @@ export function AgentMessage({
)}
</div>
)}
<>
{agentMessage.actions
.filter((a) => isVisualizationActionType(a))
.map((a, i) => {
const streamingViz = streamedVisualizations.find(
(sv) => sv.actionId === a.id
);
assert(isVisualizationActionType(a));
return (
<VisualizationActionIframe
action={a}
conversationId={conversationId}
isStreaming={!!streamingViz}
key={i}
onRetry={() => retryHandler(agentMessage)}
owner={owner}
streamedCode={streamingViz?.visualization || null}
/>
);
})}
</>
{agentMessage.status === "cancelled" && (
<Chip
label="Message generation was interrupted"
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,171 @@
import { BracesIcon, PlayIcon, Tab } from "@dust-tt/sparkle";
import type {
VisualizationActionType,
VisualizationRPCRequest,
WorkspaceType,
} from "@dust-tt/types";
import {
isGetCodeToExecuteRequest,
isGetFileRequest,
isVisualizationRPCRequest,
visualizationExtractCodeNonStreaming,
visualizationExtractCodeStreaming,
} from "@dust-tt/types";
import { useCallback, useEffect, useState } from "react";

import { RenderMessageMarkdown } from "@app/components/assistant/RenderMessageMarkdown";

const sendResponseToIframe = (
request: VisualizationRPCRequest,
response: unknown,
target: MessageEventSource
) => {
target.postMessage(
{
command: "answer",
messageUniqueId: request.messageUniqueId,
actionId: request.actionId,
result: response,
},
// TODO(2024-07-24 flav) Restrict origin.
{ targetOrigin: "*" }
);
};

// Custom hook to encapsulate the logic for handling visualization messages.
function useVisualizationDataHandler(
action: VisualizationActionType,
workspaceId: string,
onRetry: () => void
) {
const getFile = useCallback(
async (fileId: string) => {
const response = await fetch(
`/api/w/${workspaceId}/files/${fileId}?action=view`
);
if (!response.ok) {
// TODO(2024-07-24 flav) Propagate the error to the iframe.
throw new Error(`Failed to fetch file ${fileId}`);
}

const resBuffer = await response.arrayBuffer();
return new File([resBuffer], fileId, {
type: response.headers.get("Content-Type") || undefined,
});
},
[workspaceId]
);

useEffect(() => {
const listener = async (event: MessageEvent) => {
const { data } = event;

// TODO(2024-07-24 flav) Check origin.
if (
!isVisualizationRPCRequest(data) ||
!event.source ||
data.actionId !== action.id
) {
return;
}

if (isGetFileRequest(data)) {
const file = await getFile(data.params.fileId);

sendResponseToIframe(data, { file }, event.source);
} else if (isGetCodeToExecuteRequest(data)) {
const code = action.generation;

sendResponseToIframe(data, { code }, event.source);
} else {
// TODO(2024-07-24 flav) Pass the error message to the host window.
onRetry();
}

// TODO: Types above are not accurate, as it can pass the first check but won't enter any if block.
};

window.addEventListener("message", listener);
return () => window.removeEventListener("message", listener);
}, [action.generation, action.id, onRetry, getFile]);

return { getFile };
}

export function VisualizationActionIframe({
owner,
action,
isStreaming,
streamedCode,
onRetry,
}: {
conversationId: string;
owner: WorkspaceType;
action: VisualizationActionType;
streamedCode: string | null;
isStreaming: boolean;
onRetry: () => void;
}) {
const [activeTab, setActiveTab] = useState<"code" | "runtime">("code");
const [tabManuallyChanged, setTabManuallyChanged] = useState(false);

const workspaceId = owner.sId;

useVisualizationDataHandler(action, workspaceId, onRetry);

useEffect(() => {
if (activeTab === "code" && action.generation && !tabManuallyChanged) {
setActiveTab("runtime");
setTabManuallyChanged(true);
}
}, [action.generation, activeTab, tabManuallyChanged]);

let extractedCode: string | null = null;

if (action.generation) {
extractedCode = visualizationExtractCodeNonStreaming(action.generation);
} else {
extractedCode = visualizationExtractCodeStreaming(streamedCode || "");
}

return (
<>
<Tab
tabs={[
{
label: "Code",
id: "code",
current: activeTab === "code",
icon: BracesIcon,
sizing: "expand",
},
{
label: "Run",
id: "runtime",
current: activeTab === "runtime",
icon: PlayIcon,
sizing: "expand",
hasSeparator: true,
},
]}
setCurrentTab={(tabId, event) => {
event.preventDefault();
setActiveTab(tabId as "code" | "runtime");
}}
/>
{activeTab === "code" && extractedCode && extractedCode.length > 0 && (
<RenderMessageMarkdown
content={"```javascript\n" + extractedCode + "\n```"}
isStreaming={isStreaming}
/>
)}
{activeTab === "runtime" && (
<iframe
style={{ width: "100%", height: "600px" }}
src={`${process.env.NEXT_PUBLIC_VIZ_URL}/content?aId=${action.id}`}
sandbox="allow-scripts"
/>
)}
</>
);
}
Original file line number Diff line number Diff line change
Expand Up @@ -178,7 +178,7 @@ export async function submitAssistantBuilderForm({
{
type: "visualization_configuration",
name: DEFAULT_VISUALIZATION_ACTION_NAME,
description: "Generates graphs from data.",
description: "Generate client side javascript react code.",
},
];

Expand Down
Loading

0 comments on commit 4ab90aa

Please sign in to comment.