refactor: architecture — boxes left, descriptions right, 4-layer layout

architecture.py

@@ -1,76 +1,135 @@
-"""AstrAI promo: Full architecture overview."""
+"""AstrAI promo: 4-layer architecture — boxes left, explanations right."""
 
 from manim import *
 
 
 class Architecture(Scene):
-    """Animates the full AstrAI system stack layer by layer."""
+    """Boxes on left, description text on right for each layer."""
 
     def construct(self):
-        title = Text("AstrAI Architecture", font_size=48, color=BLUE)
+        title = Text("AstrAI Architecture", font_size=42, color=BLUE)
+        title.to_edge(UP, buff=0.25)
         self.play(Write(title))
-        self.wait(0.2)
-        self.play(title.animate.to_edge(UP))
 
-        layers_data = [
+        W, BH = 5.2, 1.15
-            (0.9, GREEN, "API Layer", ["FastAPI Server  •  OpenAI-Compatible API"]),
+        BX = -3.6
-            (0.9, BLUE, "Inference Engine", ["Streaming  •  Async  •  Batch Modes"]),
+        TX = 3.6
-            (1.6, YELLOW, "Continuous Batching Scheduler",
+
-             ["Cleanup → Refill → Prefill → Decode",
+        def make_box(header, color, bits, src):
-              "Position-Grouped Decode",
+            b = Rectangle(width=W, height=BH, color=color, fill_opacity=0.1, stroke_width=1.5)
-              "Bitmask O(1) Slot Allocation"]),
+            h = Text(header, font_size=16, color=color, weight=BOLD)
-            (1.2, ORANGE, "Prefix Cache + KV Cache",
+            items = [h]
-             ["Radix Tree  •  Slot Versioning",
+            for line in bits:
-              "GPU copy_()  →  Zero-Copy Reuse"]),
+                items.append(Text(line, font_size=10, color=WHITE))
-            (1.2, PURPLE, "Transformer Model (1B params)",
+            items.append(Text(src, font_size=9, color=GRAY))
-             ["24-layer GQA  •  RoPE  •  SwiGLU",
+            c = VGroup(*items).arrange(DOWN, buff=0.04)
-              "bfloat16  •  100K vocab"]),
+            c.move_to(b.get_center())
+            return VGroup(b, c)
+
+        L1 = make_box("HTTP API Server", GREEN,
+                      ["FastAPI  ·  OpenAI-Compatible",
+                       "/v1/chat/completions  ·  SSE streaming"],
+                      "astrai/inference/server.py")
+
+        L2 = make_box("Inference Engine", BLUE,
+                      ["generate()  ·  batch mode  ·  streaming",
+                       "4-phase daemon: Cleanup → Refill → Prefill → Decode",
+                       "Position-grouped decode  ·  Bitmask O(1) slots"],
+                      "astrai/inference/engine.py  ·  scheduler.py")
+
+        L3 = make_box("Prefix Cache  +  KV Cache", ORANGE,
+                      ["Radix Tree prefix matching  ·  LRU eviction",
+                       "Slot versioning  ·  GPU copy_() zero-copy reuse"],
+                      "astrai/inference/scheduler.py")
+
+        L4 = make_box("Transformer Model", PURPLE,
+                      ["24× DecoderBlock  ·  GQA 6:1  ·  RoPE",
+                       "SwiGLU MLP  ·  Dim 1536  ·  bfloat16"],
+                      "astrai/model/transformer.py")
+
+        layers = VGroup(L1, L2, L3, L4)
+        layers.arrange(DOWN, buff=0.08)
+        layers.move_to([BX, 0, 0])
+        layers.next_to(title, DOWN, buff=0.25)
+
+        # Description panels (right side)
+        descs_text = [
+            ["HTTP API Server",
+             "Receives chat requests via",
+             "OpenAI-compatible endpoints.",
+             "Streams generated tokens back",
+             "through Server-Sent Events."],
+            ["Inference Engine",
+             "Orchestrates the full generation",
+             "pipeline with a background daemon.",
+             "4-phase loop: Cleanup tasks,",
+             "Refill batch, Prefill prompts,",
+             "Decode tokens one by one."],
+            ["Prefix Cache + KV Cache",
+             "Caches key-value states using",
+             "a Radix Tree for O(n) prefix lookup.",
+             "Reuses matched prefixes via GPU",
+             "memcpy — zero recomputation."],
+            ["Transformer Model (1B params)",
+             "Decoder-only Transformer with",
+             "Grouped-Query Attention (GQA 6:1).",
+             "RoPE rotary encoding, SwiGLU",
+             "activation, 100K vocabulary."],
         ]
 
-        layers = VGroup()
+        def make_desc(lines, color):
-        for height, color, label, subs in layers_data:
+            els = [Text(lines[0], font_size=20, color=color, weight=BOLD)]
-            box = Rectangle(width=7.5, height=height, color=color, fill_opacity=0.1)
+            for ln in lines[1:]:
-            lbl = Text(label, font_size=18, color=color)
+                els.append(Text(ln, font_size=14, color=WHITE))
-            items = [lbl] + [Text(s, font_size=11, color=WHITE) for s in subs]
+            grp = VGroup(*els).arrange(DOWN, buff=0.1, aligned_edge=LEFT)
-            content = VGroup(*items)
+            return grp
-            content.arrange(DOWN, buff=0.22)
-            content.move_to(box.get_center())
-            layers.add(VGroup(box, content))
 
-        layers.arrange(DOWN, buff=0.18)
+        COLORS = [GREEN, BLUE, ORANGE, PURPLE]
-        layers.next_to(title, DOWN, buff=0.3)
+        descs = [make_desc(lns, c) for lns, c in zip(descs_text, COLORS)]
 
-        for i in range(len(layers)):
+        arrows = VGroup()
-            self.play(Create(layers[i]), run_time=0.35)
+        for i, (layer, desc) in enumerate(zip(layers, descs)):
-            if i > 0:
+            b = layer[0]
-                # Use box-to-box for arrow endpoints (not content)
+            self.play(Create(layer), run_time=0.35)
-                prev_box = layers[i - 1][0]
+            desc.next_to(b, RIGHT, buff=1.0)
-                curr_box = layers[i][0]
+            desc.align_to(b, UP)
+            self.play(Write(desc), run_time=0.3)
+            self.wait(2.0 if i == 0 else 1.8)
+            if i < len(layers) - 1:
+                self.play(FadeOut(desc))
+                nxt = layers[i + 1][0]
                 arrow = Arrow(
-                    prev_box.get_bottom(),
+                    b.get_bottom(), nxt.get_top(),
-                    curr_box.get_top(),
+                    color=GRAY, buff=0.04,
-                    color=GRAY,
+                    max_tip_length_to_length_ratio=0.18,
-                    buff=0.1,
-                    max_tip_length_to_length_ratio=0.15,
                 )
-                self.play(Create(arrow), run_time=0.15)
+                self.play(Create(arrow), run_time=0.12)
-
+                arrows.add(arrow)
+            else:
                 self.wait(0.5)
+                self.play(FadeOut(desc))
 
-        hl = SurroundingRectangle(layers[3], color=GREEN, buff=0.12)
+        # Show all boxes + arrows together briefly
-        hl_note = Text("Zero-Copy Prefix Reuse", font_size=18, color=GREEN)
+        self.wait(0.3)
-        hl_note.next_to(hl, LEFT, buff=0.4)
-        self.play(Create(hl), Write(hl_note))
-        self.wait(1.5)
-        self.play(FadeOut(hl), FadeOut(hl_note))
 
-        self.play(FadeOut(layers))
+        # Highlight innovation layers
+        hl2 = SurroundingRectangle(L2, color=BLUE, buff=0.1, stroke_width=2)
+        hl3 = SurroundingRectangle(L3, color=ORANGE, buff=0.1, stroke_width=2)
+        hl_note = Text("Key Innovations", font_size=20, color=GOLD)
+        hl_note.next_to(VGroup(hl2, hl3), RIGHT, buff=1.5)
+        hl_note.align_to(hl2, UP)
+        self.play(Create(hl2), Create(hl3), Write(hl_note))
+        self.wait(2.0)
+        self.play(FadeOut(hl2), FadeOut(hl3), FadeOut(hl_note))
+
+        self.play(FadeOut(VGroup(*layers)), FadeOut(arrows))
 
         cta = VGroup(
             Text("AstrAI", font_size=52, color=BLUE),
-            Text("Single GPU  •  Open Source  •  1B params", font_size=24, color=GRAY),
+            Text("Single GPU  ·  Open Source  ·  1B params", font_size=22, color=GRAY),
-            Text("github.com/ViperEkura/AstrAI", font_size=20, color=YELLOW),
+            Text("github.com/ViperEkura/AstrAI", font_size=18, color=YELLOW),
         ).arrange(DOWN, buff=0.35)
+        cta.move_to(ORIGIN)
         self.play(Write(cta))
-        self.wait(2)
+        self.wait(2.5)
         self.play(FadeOut(cta), FadeOut(title))