Update Agent (#215)

* Install missing better-auth dep

* Add working better auth secret

* Delete AI Elements

* Add new AI Elements

* Remove avatars, fix reasoning

* Remove group

* Remove conversation copy

* Update agent-page-content.tsx

* Update globals.css

* Fix sizing

* Rework tool calls

Author: haydenbleasel

.env.example

@@ -6,7 +6,7 @@ REDIS_URL="redis://localhost:6379"
 AI_API_KEY=""
 
 BETTER_AUTH_URL="http://localhost:3000"
-BETTER_AUTH_SECRET="your_bcrypt_secret"
+BETTER_AUTH_SECRET="wSl7AHwHRxm6HVi0rSLcvFnn0SiZG+thg9IAF/vBkHs="
 
 # OpenPageRank key, not needed for most operations
 OPR_API_KEY=""

apps/dashboard/app/(main)/websites/[id]/agent/_components/agent-input.tsx

@@ -52,7 +52,7 @@ export function AgentInput() {
 
 	return (
 		<div className="shrink-0 border-t bg-sidebar/30 backdrop-blur-sm">
-			<div className="mx-auto max-w-2xl p-4">
+			<div className="mx-auto max-w-4xl p-4">
 				<div className="relative">
 					<AgentCommandMenu />
 

apps/dashboard/app/(main)/websites/[id]/agent/_components/agent-messages.tsx

@@ -1,28 +1,24 @@
 "use client";
 
-import { authClient } from "@databuddy/auth/client";
 import type { UIMessage } from "ai";
 import { useEffect, useState } from "react";
+import {
+	ChainOfThought,
+	ChainOfThoughtContent,
+	ChainOfThoughtHeader,
+	ChainOfThoughtStep,
+} from "@/components/ai-elements/chain-of-thought";
 import {
 	Message,
-	MessageAvatar,
 	MessageContent,
+	MessageResponse,
 } from "@/components/ai-elements/message";
 import {
 	Reasoning,
 	ReasoningContent,
 	ReasoningTrigger,
 } from "@/components/ai-elements/reasoning";
-import { Response } from "@/components/ai-elements/response";
-import {
-	Tool,
-	ToolContent,
-	ToolHeader,
-	ToolInput,
-	ToolOutput,
-} from "@/components/ai-elements/tool";
 import { Avatar, AvatarFallback, AvatarImage } from "@/components/ui/avatar";
-import { Skeleton } from "@/components/ui/skeleton";
 import { cn } from "@/lib/utils";
 
 type AgentMessagesProps = {
@@ -34,24 +30,6 @@ type AgentMessagesProps = {
 
 type MessagePart = UIMessage["parts"][number];
 
-function isReasoningPart(part: MessagePart): boolean {
-	return (
-		part.type === "reasoning" ||
-		part.type === "step-start" ||
-		part.type === "data-step-start" ||
-		part.type?.includes("reasoning") ||
-		part.type?.includes("step")
-	);
-}
-
-function isTextPart(part: MessagePart): boolean {
-	return part.type === "text";
-}
-
-function isToolPart(part: MessagePart): boolean {
-	return part.type?.startsWith("tool") ?? false;
-}
-
 function getReasoningText(part: MessagePart): string {
 	const reasoning = part as {
 		text?: string;
@@ -66,43 +44,44 @@ function getReasoningText(part: MessagePart): string {
 	);
 }
 
-function getToolState(part: MessagePart) {
-	const tool = part as { errorText?: string; output?: unknown };
-	if (tool.errorText) {
-		return "output-error";
+function formatToolOutput(output: unknown) {
+	if (output === undefined) {
+		return null;
 	}
-	if (tool.output !== undefined) {
-		return "output-available";
+
+	console.log(output);
+
+	if (typeof output === "object" && "data" in output) {
+		return <p>Found {output.data.length} results.</p>;
 	}
-	return "input-available";
-}
 
-function formatToolOutput(output: unknown, toolName?: string) {
-	if (output === undefined) {
-		return null;
+	if (typeof output === "object" && "pages" in output) {
+		return <p>Found {output.pages.length} results.</p>;
 	}
 
-	if (
-		toolName === "web_search" &&
-		typeof output === "object" &&
-		output !== null
-	) {
-		const webData = output as { data?: unknown[] };
-		if (Array.isArray(webData.data)) {
-			return {
-				summary: `Scraped ${webData.data.length} page(s)`,
-				results: webData.data.map((page, index) => ({
-					page: index + 1,
-					...(typeof page === "object" ? page : { content: page }),
-				})),
-			};
-		}
+	if (typeof output === "object" && "errorText" in output) {
+		return <p>Error: {output.errorText}</p>;
 	}
 
-	if (typeof output === "string" || typeof output === "object") {
-		return output as string | Record<string, unknown>;
+	if (typeof output === "string") {
+		const obj = JSON.parse(output);
+
+		if ("data" in obj) {
+			return <p>Found {obj.data.length} results.</p>;
+		}
+
+		if ("pages" in obj) {
+			return <p>Found {obj.pages.length} results.</p>;
+		}
+
+		if ("errorText" in obj) {
+			return <p>Error: {obj.errorText}</p>;
+		}
+
+		return <p>Found 0 results.</p>;
 	}
-	return String(output);
+
+	return <p>Found 0 results.</p>;
 }
 
 function ReasoningMessage({
@@ -128,61 +107,36 @@ function ReasoningMessage({
 	);
 }
 
-function ToolMessage({
-	part,
-	partIndex,
-	isStreaming,
-}: {
-	part: MessagePart;
-	partIndex: number;
-	isStreaming: boolean;
-}) {
-	const toolPart = part as {
-		toolCallId?: string;
-		toolName?: string;
-		input?: unknown;
-		output?: unknown;
-		errorText?: string;
-	};
-
-	const state = getToolState(part);
-	const isRunning = state === "input-available";
-	const hasCompleted = state === "output-available" || state === "output-error";
-
-	const [hasBeenStreaming, setHasBeenStreaming] = useState(
-		isStreaming && isRunning
-	);
-
-	useEffect(() => {
-		if (isStreaming && isRunning) {
-			setHasBeenStreaming(true);
+function groupConsecutiveToolCalls(parts: MessagePart[]) {
+	const grouped: Array<MessagePart | MessagePart[]> = [];
+	let currentToolGroup: MessagePart[] = [];
+
+	for (const part of parts) {
+		if (part.type?.includes("tool")) {
+			currentToolGroup.push(part);
+		} else {
+			if (currentToolGroup.length > 0) {
+				grouped.push(
+					currentToolGroup.length === 1 ? currentToolGroup[0] : currentToolGroup
+				);
+				currentToolGroup = [];
+			}
+			grouped.push(part);
 		}
-	}, [isStreaming, isRunning]);
+	}
 
-	const shouldBeOpen = hasBeenStreaming || hasCompleted;
+	// Don't forget the last group
+	if (currentToolGroup.length > 0) {
+		grouped.push(
+			currentToolGroup.length === 1 ? currentToolGroup[0] : currentToolGroup
+		);
+	}
 
-	return (
-		<Tool defaultOpen={shouldBeOpen}>
-			<ToolHeader
-				state={state}
-				type={
-					(toolPart.toolName as `tool-${string}`) ??
-					(`tool-${partIndex}` as const)
-				}
-			/>
-			<ToolContent>
-				{toolPart.input !== undefined && <ToolInput input={toolPart.input} />}
-				<ToolOutput
-					errorText={toolPart.errorText}
-					output={formatToolOutput(toolPart.output, toolPart.toolName)}
-				/>
-			</ToolContent>
-		</Tool>
-	);
+	return grouped;
 }
 
 function renderMessagePart(
-	part: MessagePart,
+	part: MessagePart | MessagePart[],
 	partIndex: number,
 	messageId: string,
 	isLastMessage: boolean,
@@ -191,7 +145,27 @@ function renderMessagePart(
 	const key = `${messageId}-${partIndex}`;
 	const isCurrentlyStreaming = isLastMessage && isStreaming;
 
-	if (isReasoningPart(part)) {
+	// Handle grouped tool calls
+	if (Array.isArray(part)) {
+		return (
+			<ChainOfThought className="my-4" defaultOpen key={key}>
+				<ChainOfThoughtHeader>Running {part.length} tools</ChainOfThoughtHeader>
+				<ChainOfThoughtContent>
+					{part.map((toolPart, idx) => (
+						<ChainOfThoughtStep
+							key={`${key}-tool-${idx}`}
+							label={`Running ${toolPart.type}`}
+							status="complete"
+						>
+							{formatToolOutput(toolPart.output)}
+						</ChainOfThoughtStep>
+					))}
+				</ChainOfThoughtContent>
+			</ChainOfThought>
+		);
+	}
+
+	if (part.type === "reasoning") {
 		return (
 			<ReasoningMessage
 				isStreaming={isCurrentlyStreaming}
@@ -201,27 +175,32 @@ function renderMessagePart(
 		);
 	}
 
-	if (isTextPart(part)) {
+	if (part.type === "text") {
 		const textPart = part as { text: string };
 		if (!textPart.text?.trim()) {
 			return null;
 		}
 
 		return (
-			<Response isAnimating={isCurrentlyStreaming} key={key}>
+			<MessageResponse isAnimating={isCurrentlyStreaming} key={key}>
 				{textPart.text}
-			</Response>
+			</MessageResponse>
 		);
 	}
 
-	if (isToolPart(part)) {
+	console.log(part);
+
+	if (part.type?.includes("tool")) {
 		return (
-			<ToolMessage
-				isStreaming={isCurrentlyStreaming}
-				key={key}
-				part={part}
-				partIndex={partIndex}
-			/>
+			<ChainOfThought defaultOpen key={key}>
+				<ChainOfThoughtHeader />
+				<ChainOfThoughtContent>
+					<ChainOfThoughtStep
+						label={`Running ${part.type}`}
+						status="complete"
+					/>
+				</ChainOfThoughtContent>
+			</ChainOfThought>
 		);
 	}
 
@@ -245,29 +224,28 @@ export function AgentMessages({
 				const showError =
 					isLastMessage && hasError && message.role === "assistant";
 
+				const groupedParts = message.parts
+					? groupConsecutiveToolCalls(message.parts)
+					: [];
+
 				return (
-					<div className="group" key={message.id}>
-						<Message from={message.role}>
-							<MessageContent className="max-w-[80%]" variant="flat">
-								{message.parts?.map((part, partIndex) =>
-									renderMessagePart(
-										part,
-										partIndex,
-										message.id,
-										isLastMessage,
-										isStreaming
-									)
-								)}
-
-								{showError && <ErrorMessage />}
-							</MessageContent>
-
-							{message.role === "user" && <UserAvatar />}
-							{message.role === "assistant" && (
-								<AssistantAvatar hasError={hasError} />
+					<Message from={message.role} key={message.id}>
+						<MessageContent
+							className={cn(message.role === "assistant" ? "w-full" : "")}
+						>
+							{groupedParts.map((part, partIndex) =>
+								renderMessagePart(
+									part,
+									partIndex,
+									message.id,
+									isLastMessage,
+									isStreaming
+								)
 							)}
-						</Message>
-					</div>
+
+							{showError && <ErrorMessage />}
+						</MessageContent>
+					</Message>
 				);
 			})}
 
@@ -319,35 +297,3 @@ function StreamingIndicator({ statusText }: { statusText?: string }) {
 		</div>
 	);
 }
-
-function UserAvatar() {
-	const { data: session, isPending } = authClient.useSession();
-	const user = session?.user;
-
-	if (isPending) {
-		return <Skeleton className="size-8 shrink-0 rounded-full" />;
-	}
-
-	return (
-		<MessageAvatar
-			name={user?.name || user?.email || "User"}
-			src={user?.image || ""}
-		/>
-	);
-}
-
-function AssistantAvatar({ hasError = false }: { hasError?: boolean }) {
-	return (
-		<Avatar className="size-8 shrink-0 ring-1 ring-border">
-			<AvatarImage alt="Databunny" src="/databunny.webp" />
-			<AvatarFallback
-				className={cn(
-					"bg-primary/10 font-semibold text-primary",
-					hasError && "bg-destructive/10 text-destructive"
-				)}
-			>
-				DB
-			</AvatarFallback>
-		</Avatar>
-	);
-}