feat: Enhance tool mapping and output rendering with animations (#4481)

* Enhance tool block mapping by using unique tool keys with name and run_id

* Enhance tool output rendering with Markdown and JSON formatting in ContentDisplay component

* Add animations for block title and content separators in ContentBlockDisplay component

* Allow 'size' prop to accept string values and update styling in BorderTrail component

* Adjust BorderTrail animation size and duration based on expansion state

* fix both borders trailing at the same time

* [autofix.ci] apply automated fixes

* fix text sizing

* fix spacing issues

* Adjust header title and text styling in ContentBlockDisplay and DurationDisplay components

* Refactor header title in ContentBlockDisplay component

* [autofix.ci] apply automated fixes

* Convert `test_handle_on_tool_start` to an async function and update tool content key logic

* Handle logger without 'opt' method in code parsing error handling

* Update test duration values in .test_durations file

---------

Co-authored-by: anovazzi1 <otavio2204@gmail.com>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-11-11 13:32:08 -03:00 • committed by GitHub
commit 06b7450652
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 882 additions and 702 deletions

View file

@ -105,6 +105,7 @@ def handle_on_tool_start(
tool_name = event["name"] tool_name = event["name"]
tool_input = event["data"].get("input") tool_input = event["data"].get("input")
run_id = event.get("run_id", "") run_id = event.get("run_id", "")
tool_key = f"{tool_name}_{run_id}"
# Create content blocks if they don't exist # Create content blocks if they don't exist
if not agent_message.content_blocks: if not agent_message.content_blocks:
@ -122,11 +123,11 @@ def handle_on_tool_start(
) )
# Store in map and append to message # Store in map and append to message
tool_blocks_map[run_id] = tool_content tool_blocks_map[tool_key] = tool_content
agent_message.content_blocks[0].contents.append(tool_content) agent_message.content_blocks[0].contents.append(tool_content)
agent_message = send_message_method(message=agent_message) agent_message = send_message_method(message=agent_message)
tool_blocks_map[run_id] = agent_message.content_blocks[0].contents[-1] tool_blocks_map[tool_key] = agent_message.content_blocks[0].contents[-1]
return agent_message, start_time return agent_message, start_time
@ -138,7 +139,9 @@ def handle_on_tool_end(
start_time: float, start_time: float,
) -> tuple[Message, float]: ) -> tuple[Message, float]:
run_id = event.get("run_id", "") run_id = event.get("run_id", "")
tool_content = tool_blocks_map.get(run_id) tool_name = event.get("name", "")
tool_key = f"{tool_name}_{run_id}"
tool_content = tool_blocks_map.get(tool_key)
if tool_content and isinstance(tool_content, ToolContent): if tool_content and isinstance(tool_content, ToolContent):
tool_content.output = event["data"].get("output") tool_content.output = event["data"].get("output")
@ -159,7 +162,9 @@ def handle_on_tool_error(
start_time: float, start_time: float,
) -> tuple[Message, float]: ) -> tuple[Message, float]:
run_id = event.get("run_id", "") run_id = event.get("run_id", "")
tool_content = tool_blocks_map.get(run_id) tool_name = event.get("name", "")
tool_key = f"{tool_name}_{run_id}"
tool_content = tool_blocks_map.get(tool_key)
if tool_content and isinstance(tool_content, ToolContent): if tool_content and isinstance(tool_content, ToolContent):
tool_content.error = event["data"].get("error", "Unknown error") tool_content.error = event["data"].get("error", "Unknown error")

View file

@ -29,7 +29,10 @@ def validate_code(code):
try: try:
tree = ast.parse(code) tree = ast.parse(code)
except Exception as e: # noqa: BLE001 except Exception as e: # noqa: BLE001
logger.opt(exception=True).debug("Error parsing code") if hasattr(logger, "opt"):
logger.opt(exception=True).debug("Error parsing code")
else:
logger.debug("Error parsing code")
errors["function"]["errors"].append(str(e)) errors["function"]["errors"].append(str(e))
return errors return errors

File diff suppressed because it is too large Load diff

View file

@ -392,7 +392,7 @@ async def test_handle_on_chain_end_with_empty_return_values():
assert isinstance(start_time, float) assert isinstance(start_time, float)
def test_handle_on_tool_start(): async def test_handle_on_tool_start():
"""Test handle_on_tool_start event.""" """Test handle_on_tool_start event."""
send_message = MagicMock(side_effect=lambda message: message) send_message = MagicMock(side_effect=lambda message: message)
tool_blocks_map = {} tool_blocks_map = {}
@ -414,8 +414,9 @@ def test_handle_on_tool_start():
assert len(updated_message.content_blocks) == 1 assert len(updated_message.content_blocks) == 1
assert len(updated_message.content_blocks[0].contents) > 0 assert len(updated_message.content_blocks[0].contents) > 0
tool_key = f"{event['name']}_{event['run_id']}"
tool_content = updated_message.content_blocks[0].contents[-1] tool_content = updated_message.content_blocks[0].contents[-1]
assert tool_content == tool_blocks_map.get("test_run") assert tool_content == tool_blocks_map.get(tool_key)
assert isinstance(tool_content, ToolContent) assert isinstance(tool_content, ToolContent)
assert tool_content.name == "test_tool" assert tool_content.name == "test_tool"
assert tool_content.tool_input == {"query": "tool input"} assert tool_content.tool_input == {"query": "tool input"}
@ -452,6 +453,7 @@ async def test_handle_on_tool_end():
updated_message, start_time = handle_on_tool_end(end_event, agent_message, tool_blocks_map, send_message, 0.0) updated_message, start_time = handle_on_tool_end(end_event, agent_message, tool_blocks_map, send_message, 0.0)
f"{end_event['name']}_{end_event['run_id']}"
tool_content = updated_message.content_blocks[0].contents[-1] tool_content = updated_message.content_blocks[0].contents[-1]
assert tool_content.name == "test_tool" assert tool_content.name == "test_tool"
assert tool_content.output == "tool output" assert tool_content.output == "tool output"

View file

@ -40,10 +40,11 @@ export function ContentBlockDisplay({
contentBlocks[0]?.contents[contentBlocks[0]?.contents.length - 1]; contentBlocks[0]?.contents[contentBlocks[0]?.contents.length - 1];
const headerIcon = const headerIcon =
state === "partial" ? lastContent?.header?.icon || "Bot" : "Bot"; state === "partial" ? lastContent?.header?.icon || "Bot" : "Bot";
const headerTitle = const headerTitle =
(state === "partial" state === "partial" ? (lastContent?.header?.title ?? "Steps") : "Finished";
? lastContent?.header?.title // show the block title only if state === "partial"
: contentBlocks[0]?.title) || "Steps"; const showBlockTitle = state === "partial";
return ( return (
<div className="relative py-3"> <div className="relative py-3">
@ -61,11 +62,10 @@ export function ContentBlockDisplay({
> >
{isLoading && ( {isLoading && (
<BorderTrail <BorderTrail
className="bg-zinc-600 opacity-50 dark:bg-zinc-400" size={100}
size={60}
transition={{ transition={{
repeat: Infinity, repeat: Infinity,
duration: 2, duration: 10,
ease: "linear", ease: "linear",
}} }}
/> />
@ -92,7 +92,7 @@ export function ContentBlockDisplay({
<Markdown <Markdown
remarkPlugins={[remarkGfm]} remarkPlugins={[remarkGfm]}
rehypePlugins={[rehypeMathjax]} rehypePlugins={[rehypeMathjax]}
className="inline-block w-fit max-w-full font-semibold text-primary" className="inline-block w-fit max-w-full text-[14px] font-semibold text-primary"
> >
{headerTitle} {headerTitle}
</Markdown> </Markdown>
@ -139,33 +139,59 @@ export function ContentBlockDisplay({
animate={{ opacity: 1 }} animate={{ opacity: 1 }}
transition={{ duration: 0.2, delay: 0.1 }} transition={{ duration: 0.2, delay: 0.1 }}
className={cn( className={cn(
"relative p-4", "relative",
index !== contentBlocks.length - 1 && index !== contentBlocks.length - 1 &&
"border-b border-border", "border-b border-border",
)} )}
> >
<div className="mb-2 font-medium"> <AnimatePresence>
<Markdown {showBlockTitle && (
remarkPlugins={[remarkGfm]} <motion.div
linkTarget="_blank" initial={{ opacity: 0, height: 0, marginBottom: 0 }}
rehypePlugins={[rehypeMathjax]} animate={{
components={{ opacity: 1,
p({ node, ...props }) { height: "auto",
return ( marginBottom: 8,
<span className="inline">{props.children}</span> }}
); exit={{ opacity: 0, height: 0, marginBottom: 0 }}
}, transition={{ duration: 0.2 }}
}} className="overflow-hidden font-medium"
> >
{block.title} <Markdown
</Markdown> className="text-[14px] font-semibold text-foreground"
</div> remarkPlugins={[remarkGfm]}
linkTarget="_blank"
rehypePlugins={[rehypeMathjax]}
components={{
p({ node, ...props }) {
return (
<span className="inline">{props.children}</span>
);
},
}}
>
{block.title}
</Markdown>
</motion.div>
)}
</AnimatePresence>
<div className="text-sm text-muted-foreground"> <div className="text-sm text-muted-foreground">
{block.contents.map((content, index) => ( {block.contents.map((content, index) => (
<> <motion.div key={index}>
<Separator orientation="horizontal" className="my-2" /> <AnimatePresence>
<ContentDisplay key={index} content={content} /> {index !== 0 && (
</> <motion.div
initial={{ opacity: 0 }}
animate={{ opacity: 1 }}
exit={{ opacity: 0 }}
transition={{ duration: 0.2 }}
>
<Separator orientation="horizontal" />
</motion.div>
)}
</AnimatePresence>
<ContentDisplay content={content} />
</motion.div>
))} ))}
</div> </div>
</motion.div> </motion.div>

View file

@ -12,7 +12,7 @@ export default function ContentDisplay({ content }: { content: ContentType }) {
// First render the common BaseContent elements if they exist // First render the common BaseContent elements if they exist
const renderHeader = content.header && ( const renderHeader = content.header && (
<> <>
<div className="flex items-center gap-2"> <div className="flex items-center gap-2 pb-[12px]">
{content.header.icon && ( {content.header.icon && (
<ForwardedIconComponent <ForwardedIconComponent
name={content.header.icon} name={content.header.icon}
@ -25,7 +25,7 @@ export default function ContentDisplay({ content }: { content: ContentType }) {
<Markdown <Markdown
remarkPlugins={[remarkGfm]} remarkPlugins={[remarkGfm]}
rehypePlugins={[rehypeMathjax]} rehypePlugins={[rehypeMathjax]}
className="inline-block w-fit max-w-full" className="inline-block w-fit max-w-full text-[14px] font-semibold text-foreground"
> >
{content.header.title} {content.header.title}
</Markdown> </Markdown>
@ -35,7 +35,7 @@ export default function ContentDisplay({ content }: { content: ContentType }) {
</> </>
); );
const renderDuration = content.duration !== undefined && ( const renderDuration = content.duration !== undefined && (
<div className="absolute right-2 top-0"> <div className="absolute right-2 top-4">
<DurationDisplay duration={content.duration} /> <DurationDisplay duration={content.duration} />
</div> </div>
); );
@ -54,7 +54,7 @@ export default function ContentDisplay({ content }: { content: ContentType }) {
components={{ components={{
p({ node, ...props }) { p({ node, ...props }) {
return ( return (
<span className="inline-block w-fit max-w-full"> <span className="block w-fit max-w-full">
{props.children} {props.children}
</span> </span>
); );
@ -135,16 +135,91 @@ export default function ContentDisplay({ content }: { content: ContentType }) {
break; break;
case "tool_use": case "tool_use":
const formatToolOutput = (output: any) => {
if (output === null || output === undefined) return "";
// If it's a string, render as markdown
if (typeof output === "string") {
return (
<Markdown
remarkPlugins={[remarkGfm]}
rehypePlugins={[rehypeMathjax]}
className="markdown prose max-w-full text-[14px] font-normal dark:prose-invert"
components={{
pre({ node, ...props }) {
return <>{props.children}</>;
},
code: ({ node, inline, className, children, ...props }) => {
const match = /language-(\w+)/.exec(className || "");
return !inline ? (
<SimplifiedCodeTabComponent
language={(match && match[1]) || ""}
code={String(children).replace(/\n$/, "")}
/>
) : (
<code className={className} {...props}>
{children}
</code>
);
},
}}
>
{output}
</Markdown>
);
}
// For objects/arrays, format as JSON
try {
return (
<CodeBlock
language="json"
value={JSON.stringify(output, null, 2)}
/>
);
} catch {
return String(output);
}
};
contentData = ( contentData = (
<div> <div className="flex flex-col gap-2">
{content.name && <div>Tool: {content.name}</div>} <Markdown
<div>Input: {JSON.stringify(content.tool_input, null, 2)}</div> remarkPlugins={[remarkGfm]}
{content.output && ( rehypePlugins={[rehypeMathjax]}
<div>Output: {JSON.stringify(content.output)}</div> className="markdown prose max-w-full text-[14px] font-normal dark:prose-invert"
>
{`${content.name ? `**Tool:** ${content.name}\n\n` : ""}**Input:**`}
</Markdown>
<CodeBlock
language="json"
value={JSON.stringify(content.tool_input, null, 2)}
/>
{content.output !== undefined && (
<>
<Markdown
remarkPlugins={[remarkGfm]}
rehypePlugins={[rehypeMathjax]}
className="markdown prose max-w-full text-[14px] font-normal dark:prose-invert"
>
**Output:**
</Markdown>
<div className="mt-1">{formatToolOutput(content.output)}</div>
</>
)} )}
{content.error && ( {content.error && (
<div className="text-red-500"> <div className="text-red-500">
Error: {JSON.stringify(content.error)} <Markdown
remarkPlugins={[remarkGfm]}
rehypePlugins={[rehypeMathjax]}
className="markdown prose max-w-full text-[14px] font-normal dark:prose-invert"
>
**Error:**
</Markdown>
<CodeBlock
language="json"
value={JSON.stringify(content.error, null, 2)}
/>
</div> </div>
)} )}
</div> </div>
@ -168,7 +243,7 @@ export default function ContentDisplay({ content }: { content: ContentType }) {
} }
return ( return (
<div className="relative"> <div className="relative p-[16px]">
{renderHeader} {renderHeader}
{renderDuration} {renderDuration}
{contentData} {contentData}

View file

@ -53,7 +53,7 @@ export default function DurationDisplay({ duration }: { duration?: number }) {
bounce: 0, bounce: 0,
duration: 300, duration: 300,
}} }}
className="tabular-nums" className="text-[11px] font-bold tabular-nums"
/> />
</div> </div>
</div> </div>

View file

@ -1,10 +1,9 @@
"use client";
import { cn } from "@/utils/utils"; import { cn } from "@/utils/utils";
import { motion, Transition } from "framer-motion"; import { motion, Transition } from "framer-motion";
type BorderTrailProps = { type BorderTrailProps = {
className?: string; className?: string;
size?: number; size?: number | string;
transition?: Transition; transition?: Transition;
delay?: number; delay?: number;
onAnimationComplete?: () => void; onAnimationComplete?: () => void;
@ -28,10 +27,12 @@ export function BorderTrail({
return ( return (
<div className="pointer-events-none absolute inset-0 rounded-[inherit] border border-transparent [mask-clip:padding-box,border-box] [mask-composite:intersect] [mask-image:linear-gradient(transparent,transparent),linear-gradient(#000,#000)]"> <div className="pointer-events-none absolute inset-0 rounded-[inherit] border border-transparent [mask-clip:padding-box,border-box] [mask-composite:intersect] [mask-image:linear-gradient(transparent,transparent),linear-gradient(#000,#000)]">
<motion.div <motion.div
className={cn("absolute aspect-square bg-zinc-500", className)} className={cn("absolute bg-zinc-500", className)}
style={{ style={{
width: size, width: size,
offsetPath: `rect(0 auto auto 0 round ${size}px)`, offsetPath: `rect(0 auto auto 0 round 18px)`,
boxShadow:
"0px 0px 20px 5px rgb(255 255 255 / 90%), 0 0 30px 10px rgb(0 0 0 / 90%)",
...style, ...style,
}} }}
animate={{ animate={{