119 lines
4.9 KiB
Python
119 lines
4.9 KiB
Python
"""AstrAI promo: Prefix Cache animation (Radix tree with branches)."""
|
|
|
|
from manim import *
|
|
|
|
Text.set_default(font="Times New Roman")
|
|
|
|
|
|
class PrefixCache(Scene):
|
|
"""Animates the radix-tree prefix cache with multiple distinct branches."""
|
|
|
|
def _add_node(self, parent_pos, label, color, dx, dy):
|
|
pos = parent_pos + np.array([dx, dy, 0])
|
|
dot = Dot(point=pos, color=color, radius=0.1)
|
|
txt = Text(label, font_size=13, color=color)
|
|
txt.next_to(dot, UP, buff=0.1)
|
|
grp = VGroup(dot, txt)
|
|
edge = Line(parent_pos, pos, color=GRAY, stroke_width=1.5)
|
|
return grp, edge, pos
|
|
|
|
def _add_leaf(self, parent_pos, color, tag):
|
|
leaf = Square(side_length=0.25, color=color, fill_opacity=0.4)
|
|
leaf.move_to(parent_pos + DOWN * 0.7)
|
|
edge = Line(parent_pos, leaf.get_top(), color=color, stroke_width=1.5)
|
|
lbl = Text(tag, font_size=10, color=color).next_to(leaf, DOWN, buff=0.1)
|
|
return VGroup(leaf, edge, lbl)
|
|
|
|
def construct(self):
|
|
title = Text("Prefix Cache", font_size=48, color=BLUE)
|
|
self.play(Write(title))
|
|
self.wait(0.2)
|
|
self.play(title.animate.to_edge(UP).scale(0.6))
|
|
|
|
# Root at top-left, tree stays visible throughout
|
|
root_pos = np.array([-4.5, 2.0, 0])
|
|
root = Circle(radius=0.25, color=BLUE, fill_opacity=0.2)
|
|
root.move_to(root_pos)
|
|
root_lbl = Text("root", font_size=10, color=GRAY).move_to(root)
|
|
root_grp = VGroup(root, root_lbl)
|
|
self.play(FadeIn(root_grp, scale=0.5), run_time=0.3)
|
|
|
|
# Labels accumulate on the right side
|
|
right_x = 3.5
|
|
label_y = 2.5
|
|
label_step = 0.5
|
|
|
|
def show_label(text, color):
|
|
nonlocal label_y
|
|
lbl = Text(text, font_size=14, color=color)
|
|
lbl.move_to([right_x, label_y, 0])
|
|
label_y -= label_step
|
|
self.play(Write(lbl))
|
|
return lbl
|
|
|
|
# ── R1: A → B → C ──
|
|
r1_lbl = show_label('R1: "A B C"', GREEN)
|
|
|
|
a_grp, a_edge, a_pos = self._add_node(root_pos, "A", GREEN, 0.6, -0.9)
|
|
self.play(Create(a_edge), FadeIn(a_grp, scale=0.5), run_time=0.2)
|
|
b_grp, b_edge, b_pos = self._add_node(a_pos, "B", GREEN, 0.6, -0.9)
|
|
self.play(Create(b_edge), FadeIn(b_grp, scale=0.5), run_time=0.2)
|
|
c_grp, c_edge, c_pos = self._add_node(b_pos, "C", GREEN, 0.6, -0.9)
|
|
self.play(Create(c_edge), FadeIn(c_grp, scale=0.5), run_time=0.2)
|
|
self.play(FadeIn(self._add_leaf(c_pos, GREEN, "slot 0"), scale=0.8), run_time=0.3)
|
|
self.wait(0.3)
|
|
|
|
# ── R2: shares A B, branches D E ──
|
|
r2_lbl = show_label('R2: "A B D E"', ORANGE)
|
|
|
|
for g in [a_grp, b_grp]:
|
|
flash = SurroundingRectangle(g, color=YELLOW, buff=0.12)
|
|
self.play(Create(flash), run_time=0.1)
|
|
self.play(FadeOut(flash), run_time=0.08)
|
|
|
|
d_grp, d_edge, d_pos = self._add_node(b_pos, "D", ORANGE, -0.6, -0.9)
|
|
self.play(Create(d_edge), FadeIn(d_grp, scale=0.5), run_time=0.2)
|
|
e_grp, e_edge, e_pos = self._add_node(d_pos, "E", ORANGE, -0.6, -0.9)
|
|
self.play(Create(e_edge), FadeIn(e_grp, scale=0.5), run_time=0.2)
|
|
self.play(FadeIn(self._add_leaf(e_pos, ORANGE, "slot 1"), scale=0.8), run_time=0.3)
|
|
self.wait(0.3)
|
|
|
|
# ── R3: shares A B, single F ──
|
|
r3_lbl = show_label('R3: "A B F"', PINK)
|
|
|
|
f_grp, f_edge, f_pos = self._add_node(b_pos, "F", PINK, 0.0, -1.2)
|
|
self.play(Create(f_edge), FadeIn(f_grp, scale=0.5), run_time=0.2)
|
|
self.play(FadeIn(self._add_leaf(f_pos, PINK, "slot 2"), scale=0.8), run_time=0.3)
|
|
self.wait(0.3)
|
|
|
|
# ── R4: new prefix from root ──
|
|
r4_lbl = show_label('R4: "X Y"', TEAL)
|
|
|
|
x_grp, x_edge, x_pos = self._add_node(root_pos, "X", TEAL, -1.0, -0.9)
|
|
self.play(Create(x_edge), FadeIn(x_grp, scale=0.5), run_time=0.2)
|
|
y_grp, y_edge, y_pos = self._add_node(x_pos, "Y", TEAL, -0.6, -0.9)
|
|
self.play(Create(y_edge), FadeIn(y_grp, scale=0.5), run_time=0.2)
|
|
self.play(FadeIn(self._add_leaf(y_pos, TEAL, "slot 3"), scale=0.8), run_time=0.3)
|
|
self.wait(0.5)
|
|
|
|
# ── highlight shared prefix (tree stays) ──
|
|
reuse_box = SurroundingRectangle(VGroup(a_grp, b_grp), color=YELLOW, buff=0.15)
|
|
reuse_note = Text(
|
|
'Prefix "A B" shared\nby 3 requests — 0 copy',
|
|
font_size=16,
|
|
color=YELLOW,
|
|
)
|
|
reuse_note.next_to(reuse_box, LEFT, buff=1.0)
|
|
self.play(Create(reuse_box), Write(reuse_note))
|
|
self.wait(2)
|
|
self.play(FadeOut(reuse_box), FadeOut(reuse_note))
|
|
|
|
# ── summary below tree (tree stays visible) ──
|
|
summary = VGroup(
|
|
Text("KV cache reuse across requests", font_size=26, color=GREEN)
|
|
).arrange(DOWN, buff=0.2)
|
|
summary.to_edge(DOWN, buff=1.0)
|
|
self.play(Write(summary))
|
|
self.wait(2)
|
|
self.play(*[FadeOut(m) for m in self.mobjects])
|