pygpt-net 2.6.61__py3-none-any.whl → 2.6.63__py3-none-any.whl

pygpt_net/core/agents/custom/runner.py

@@ -6,7 +6,7 @@
 # GitHub:  https://github.com/szczyglis-dev/py-gpt   #
 # MIT License                                        #
 # Created By  : Marcin Szczygliński                  #
-# Updated Date: 2025.09.25 14:00:00                  #
+# Updated Date: 2025.09.27 06:00:00                  #
 # ================================================== #
 
 from __future__ import annotations
@@ -62,13 +62,133 @@ class FlowOrchestrator:
       - First agent (in the whole flow) gets full initial messages from the app.
       - Next agent WITHOUT memory gets only last step's displayed content as a single 'user' message.
       - Agent WITH memory:
-          * if memory has items -> use them;
-          * if memory empty -> seed input from last displayed content (or initial messages as fallback).
+          * if memory has items -> use base history (items[:-1]) and pass last displayed content as a single 'user' baton;
+          * if memory empty -> seed baton from last displayed content (or initial messages as fallback).
     """
     def __init__(self, window, logger: Optional[Logger] = None) -> None:
         self.window = window
         self.logger = logger or NullLogger()
 
+    # ---------- Helpers (production-ready) ----------
+
+    def _extract_text_from_item(self, item: TResponseInputItem) -> str:
+        """Best-effort extract plain text from TResponseInputItem."""
+        if isinstance(item, dict):
+            content = item.get("content", "")
+            if isinstance(content, str):
+                return content
+            if isinstance(content, list):
+                parts = []
+                for p in content:
+                    if isinstance(p, dict):
+                        t = p.get("text")
+                        if isinstance(t, str):
+                            parts.append(t)
+                return "\n".join(parts)
+            return ""
+        if isinstance(item, str):
+            return item
+        return ""
+
+    def _build_baton_input(
+        self,
+        *,
+        node_id: str,
+        g: FlowGraph,
+        mem: MemoryManager,
+        initial_messages: List[TResponseInputItem],
+        first_dispatch_done: bool,
+        last_plain_output: str,
+        dbg: DebugConfig,
+    ) -> tuple[List[TResponseInputItem], str, Optional[str], Any, str]:
+        """
+        Returns: (prepared_items, baton_user_text, mem_id, mem_state, source_tag)
+        Mirrors LI baton/memory policy.
+        """
+        mem_id = g.agent_to_memory.get(node_id)
+        mem_state = mem.get(mem_id) if mem_id else None
+
+        baton_user_text = ""
+        source = ""
+
+        if mem_state and mem_state.items:
+            # memory with history -> base history + baton from last output (preferred)
+            base_items = list(mem_state.items[:-1]) if len(mem_state.items) >= 1 else []
+            if last_plain_output and last_plain_output.strip():
+                baton_user_text = last_plain_output
+                prepared = base_items + [{"role": "user", "content": baton_user_text}]
+                source = "memory:existing_to_user_baton"
+            else:
+                # fallback: use last assistant content as baton
+                last_ass = mem_state.items[-1] if isinstance(mem_state.items[-1], dict) else {}
+                if isinstance(last_ass.get("content"), str):
+                    baton_user_text = last_ass.get("content", "")
+                elif isinstance(last_ass.get("content"), list) and last_ass["content"]:
+                    baton_user_text = last_ass["content"][0].get("text", "") or ""
+                else:
+                    baton_user_text = ""
+                prepared = base_items + [{"role": "user", "content": baton_user_text}]
+                source = "memory:existing_to_last_assistant"
+            return sanitize_input_items(prepared), baton_user_text, mem_id, mem_state, source
+
+        if mem_state:
+            # memory attached but empty -> seed from last output else from initial (use last user msg as baton)
+            if last_plain_output and last_plain_output.strip():
+                baton_user_text = last_plain_output
+                prepared = [{"role": "user", "content": baton_user_text}]
+                source = "memory:seed_from_last_output"
+            else:
+                base_items = list(initial_messages[:-1]) if initial_messages else []
+                last_item = initial_messages[-1] if initial_messages else {"role": "user", "content": ""}
+                baton_user_text = self._extract_text_from_item(last_item)
+                prepared = base_items + [{"role": "user", "content": baton_user_text}]
+                source = "memory:seed_from_initial"
+            return sanitize_input_items(prepared), baton_user_text, mem_id, mem_state, source
+
+        # no memory attached
+        if not first_dispatch_done:
+            # first agent: pass initial messages as-is; baton is last user text (for potential external memory)
+            last_item = initial_messages[-1] if initial_messages else {"role": "user", "content": ""}
+            baton_user_text = self._extract_text_from_item(last_item)
+            return sanitize_input_items(list(initial_messages)), baton_user_text, None, None, "no-mem:first_initial"
+        else:
+            baton_user_text = last_plain_output if last_plain_output and last_plain_output.strip() else (
+                self._extract_text_from_item(initial_messages[-1]) if initial_messages else ""
+            )
+            prepared = [{"role": "user", "content": baton_user_text}]
+            return sanitize_input_items(prepared), baton_user_text, None, None, "no-mem:last_output"
+
+    def _update_memory_after_step(
+        self,
+        *,
+        node_id: str,
+        mem_state: Any,
+        baton_user_text: str,
+        display_text: str,
+        last_response_id: Optional[str],
+        dbg: DebugConfig,
+    ) -> None:
+        """Update memory strictly with [user baton, assistant display_text], mirroring LI semantics."""
+        if not mem_state:
+            return
+        base_items = list(mem_state.items[:-1]) if getattr(mem_state, "items", None) else []
+        new_mem = (base_items or []) + [
+            {"role": "user", "content": baton_user_text or ""},
+            {"role": "assistant", "content": [{"type": "output_text", "text": display_text or ""}]},
+        ]
+        try:
+            mem_state.set_from(new_mem, last_response_id)
+            if dbg.log_inputs:
+                self.logger.debug(
+                    f"[memory] {node_id} updated len {len(base_items)} -> {len(new_mem)} "
+                    f"user='{ellipsize(baton_user_text or '', dbg.preview_chars)}' "
+                    f"assist='{ellipsize(display_text or '', dbg.preview_chars)}'"
+                )
+        except Exception as e:
+            self.logger.error(f"[memory] update failed for {node_id}: {e}")
+
+    # ---------- Main flow ----------
+
     async def run_flow(
         self,
         schema: List[Dict[str, Any]],
@@ -174,55 +294,34 @@ class FlowOrchestrator:
                     f" role='{node_rt.role}'"
                 )
 
-            # Memory selection and INPUT BUILD (memory-first policy)
-            mem_id = g.agent_to_memory.get(current_id)
-            mem_state = mem.get(mem_id) if mem_id else None
-            if dbg.log_runtime:
-                mem_info = f"{mem_id} (len={len(mem_state.items)})" if mem_state and mem_state.items else (mem_id or "-")
-                self.logger.debug(f"[memory] attached={bool(mem_id)} mem_id={mem_info}")
-
-            input_items: List[TResponseInputItem]
-            input_source = ""
-
-            if mem_state and not mem_state.is_empty():
-                # memory present and already has history
-                input_items = list(mem_state.items)
-                input_source = "memory:existing"
-            elif mem_state:
-                # memory present but empty -> seed from last output or initial messages
-                if last_plain_output and last_plain_output.strip():
-                    input_items = [{"role": "user", "content": last_plain_output}]
-                    input_source = "memory:seeded_from_last"
-                else:
-                    input_items = list(initial_messages)
-                    input_source = "memory:seeded_from_initial"
-            else:
-                # no memory -> first agent gets initial messages; others get only last output
-                if not first_dispatch_done:
-                    input_items = list(initial_messages)
-                    input_source = "no-mem:initial"
-                else:
-                    if last_plain_output and last_plain_output.strip():
-                        input_items = [{"role": "user", "content": last_plain_output}]
-                        input_source = "no-mem:last_output"
-                    else:
-                        input_items = list(initial_messages)
-                        input_source = "no-mem:fallback_initial"
-
-            prepared_items = sanitize_input_items(input_items)
+            # Input build using baton policy (LI parity)
+            prepared_items, baton_user_text, mem_id, mem_state, input_source = self._build_baton_input(
+                node_id=current_id,
+                g=g,
+                mem=mem,
+                initial_messages=initial_messages,
+                first_dispatch_done=first_dispatch_done,
+                last_plain_output=last_plain_output,
+                dbg=dbg,
+            )
 
             if dbg.log_inputs:
                 self.logger.debug(f"[input] source={input_source} items={len(prepared_items)} "
                                   f"preview={items_preview(prepared_items, dbg.preview_chars)}")
+                if mem_id:
+                    mem_info = f"{mem_id} (len={len(mem_state.items) if mem_state else 0})"
+                    self.logger.debug(f"[memory] attached={bool(mem_id)} mem_id={mem_info}")
 
             # Build agent with per-node runtime
+            # Restrict friendly_map only to allowed outgoing routes of current node
+            allowed_map = {rid: fs.agents[rid].name or rid for rid in (node.outputs or []) if rid in fs.agents}
             built = factory.build(
                 node=node,
                 node_runtime=node_rt,
                 preset=preset,
                 function_tools=function_tools,
                 force_router=False,  # auto on multi-output
-                friendly_map={aid: a.name or aid for aid, a in fs.agents.items()},
+                friendly_map=allowed_map,
                 handoffs_enabled=True,
                 context=agent_kwargs.get("context"),
             )
@@ -242,37 +341,43 @@ class FlowOrchestrator:
                 run_kwargs["trace_id"] = trace_id
 
             # Header for UI
-            title = f"\n\n**{built.name}**\n\n"
-            ctx.stream = title
+            ctx.set_agent_name(agent.name)
             bridge.on_step(ctx, begin)
             begin = False
             handler.begin = begin
-            if not use_partial_ctx:
-                handler.to_buffer(title)
 
             display_text = ""  # what we show to UI for this step
             next_id: Optional[str] = None
 
             # --- EXECUTION ---
             if stream and not multi_output:
-                # Full token streaming (single-output agent)
+                # Full token streaming (single-output agent) – collect full buffer for baton
                 result = Runner.run_streamed(agent, **run_kwargs)
                 handler.reset()
+                # Optional local accumulator; prefer handler.buffer after loop
+                last_chunk = ""
 
                 async for event in result.stream_events():
                     if bridge.stopped():
                         result.cancel()
                         bridge.on_stop(ctx)
                         break
-                    display_text, last_response_id = handler.handle(event, ctx)
-
-                # Prepare next inputs from result for memory update (if any)
-                input_items_next = result.to_input_list()
-                input_items_next = sanitize_input_items(input_items_next)
-
-                # Update memory only if attached
-                if mem_state:
-                    mem_state.update_from_result(input_items_next, last_response_id)
+                    chunk, last_response_id = handler.handle(event, ctx)
+                    if chunk:
+                        last_chunk = chunk
+
+                # Use full buffer if available (ensures baton sees complete output)
+                display_text = getattr(handler, "buffer", "") or last_chunk or ""
+
+                # Update memory strictly with baton + displayed text
+                self._update_memory_after_step(
+                    node_id=current_id,
+                    mem_state=mem_state,
+                    baton_user_text=baton_user_text,
+                    display_text=display_text,
+                    last_response_id=last_response_id,
+                    dbg=dbg,
+                )
 
                 # Route: first edge or END
                 outs = g.get_next(current_id)
@@ -307,13 +412,15 @@ class FlowOrchestrator:
                         decision = parse_route_output(raw_text, allowed_routes)
                         display_text = decision.content or ""
 
-                        # Prepare next inputs from streamed result, patch assistant content -> content
-                        input_items_next = result.to_input_list()
-                        input_items_next = patch_last_assistant_output(input_items_next, decision.content or "")
-
-                        # Update memory if attached
-                        if mem_state:
-                            mem_state.update_from_result(input_items_next, last_response_id)
+                        # Update memory with baton + displayed content
+                        self._update_memory_after_step(
+                            node_id=current_id,
+                            mem_state=mem_state,
+                            baton_user_text=baton_user_text,
+                            display_text=display_text,
+                            last_response_id=last_response_id,
+                            dbg=dbg,
+                        )
 
                         # Route decision
                         if decision.valid:
@@ -347,11 +454,15 @@ class FlowOrchestrator:
                             if not use_partial_ctx:
                                 handler.to_buffer(display_text)
 
-                        input_items_next = result.to_input_list()
-                        input_items_next = patch_last_assistant_output(input_items_next, display_text)
-
-                        if mem_state:
-                            mem_state.update_from_result(input_items_next, last_response_id)
+                        # Update memory with baton + displayed content
+                        self._update_memory_after_step(
+                            node_id=current_id,
+                            mem_state=mem_state,
+                            baton_user_text=baton_user_text,
+                            display_text=display_text,
+                            last_response_id=last_response_id,
+                            dbg=dbg,
+                        )
 
                         if decision.valid:
                             next_id = decision.route
@@ -372,11 +483,15 @@ class FlowOrchestrator:
                             if not use_partial_ctx:
                                 handler.to_buffer(display_text)
 
-                        input_items_next = result.to_input_list()
-                        input_items_next = patch_last_assistant_output(input_items_next, display_text)
-
-                        if mem_state:
-                            mem_state.update_from_result(input_items_next, last_response_id)
+                        # Update memory with baton + displayed content
+                        self._update_memory_after_step(
+                            node_id=current_id,
+                            mem_state=mem_state,
+                            baton_user_text=baton_user_text,
+                            display_text=display_text,
+                            last_response_id=last_response_id,
+                            dbg=dbg,
+                        )
 
                         if decision.valid:
                             next_id = decision.route
@@ -396,11 +511,15 @@ class FlowOrchestrator:
                         if not use_partial_ctx:
                             handler.to_buffer(display_text)
 
-                    input_items_next = result.to_input_list()
-                    input_items_next = sanitize_input_items(input_items_next)
-
-                    if mem_state:
-                        mem_state.update_from_result(input_items_next, last_response_id)
+                    # Update memory with baton + displayed text
+                    self._update_memory_after_step(
+                        node_id=current_id,
+                        mem_state=mem_state,
+                        baton_user_text=baton_user_text,
+                        display_text=display_text,
+                        last_response_id=last_response_id,
+                        dbg=dbg,
+                    )
 
                     outs = g.get_next(current_id)
                     next_id = outs[0] if outs else g.first_connected_end(current_id)
@@ -444,6 +563,10 @@ class FlowOrchestrator:
             else:
                 bridge.on_next(ctx)
 
+            # set next agent name if not at the end
+            if current_ids and current_ids[0] in fs.agents:
+                ctx.set_agent_name(fs.agents[current_ids[0]].name)
+
             # Step duration
             dur = perf_counter() - step_start
             self.logger.debug(f"[step {steps}] duration={dur:.3f}s")

pygpt_net/core/agents/runners/llama_workflow.py

@@ -1,3 +1,5 @@
+# core/agents/runners/llama_workflow.py
+
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 # ================================================== #
@@ -6,7 +8,7 @@
 # GitHub:  https://github.com/szczyglis-dev/py-gpt   #
 # MIT License                                        #
 # Created By  : Marcin Szczygliński                  #
-# Updated Date: 2025.08.24 02:00:00                  #
+# Updated Date: 2025.09.27 06:00:00                  #
 # ================================================== #
 
 import re
@@ -18,6 +20,7 @@ from llama_index.core.agent.workflow import (
     ToolCallResult,
     AgentStream,
     AgentOutput,
+    # AgentInput,  # not needed currently
 )
 from workflows.errors import WorkflowCancelledByUser
 
@@ -38,6 +41,7 @@ class LlamaWorkflow(BaseRunner):
         """
         super(LlamaWorkflow, self).__init__(window)
         self.window = window
+        self.last_response_id = None
 
     async def run(
             self,
@@ -177,12 +181,13 @@ class LlamaWorkflow(BaseRunner):
 
         prev_output = ctx.live_output
         if prev_output:
-            prev_output = self.filter_output(prev_output)  # remove all <execute>...</execute>
+            prev_output = self.filter_output(prev_output)  # remove all [!exec]...[/!exec]
 
         response_ctx.set_agent_final_response(ctx.agent_final_response)  # always set to further use
         response_ctx.set_output(prev_output)  # append from stream
         response_ctx.extra["agent_output"] = True  # mark as output response
         response_ctx.extra["agent_finish"] = True  # mark as finished
+        response_ctx.set_agent_name(ctx.get_agent_name()) # store last agent name
 
         if "agent_input" in response_ctx.extra:
             del response_ctx.extra["agent_input"]  # remove agent input from extra
@@ -254,8 +259,15 @@ class LlamaWorkflow(BaseRunner):
         item_ctx.output = ""  # empty to prevent render
         item_ctx.stream = ""  # for stream
 
+        # Keep last known agent name to avoid redundant ctx updates.
+        last_agent_name: Optional[str] = None
+
+        # Track whether current block has already produced user-visible tokens.
+        # This prevents creating empty DB items and preserves order.
+        content_written: bool = False
+        block_open: bool = False  # logical "block" opened after first StepEvent
+
         async for event in handler.stream_events():
-            print(event)
             if self.is_stopped():
                 # persist current output on stop
                 item_ctx.output = item_ctx.live_output
@@ -264,6 +276,7 @@ class LlamaWorkflow(BaseRunner):
                     self.end_stream(item_ctx, signals)
                 await handler.cancel_run()  # cancel, will raise WorkflowCancelledByUser
                 break
+
             if isinstance(event, ToolCallResult):
                 output = f"\n-----------\nExecution result:\n{event.tool_output}"
                 if verbose:
@@ -271,8 +284,11 @@ class LlamaWorkflow(BaseRunner):
                 formatted = "\n```output\n" + str(event.tool_output) + "\n```\n"
                 item_ctx.live_output += formatted
                 item_ctx.stream = formatted
+                content_written = True
                 if item_ctx.stream_agent_output and flush:
                     self.send_stream(item_ctx, signals, begin)
+                begin = False
+
             elif isinstance(event, ToolCall):
                 if "code" in event.tool_kwargs:
                     output = f"\n-----------\nTool call code:\n{event.tool_kwargs['code']}"
@@ -281,37 +297,97 @@ class LlamaWorkflow(BaseRunner):
                     formatted = "\n```python\n" + str(event.tool_kwargs['code']) + "\n```\n"
                     item_ctx.live_output += formatted
                     item_ctx.stream = formatted
+                    content_written = True
                     if item_ctx.stream_agent_output and flush:
                         self.send_stream(item_ctx, signals, begin)
+                    begin = False
+
             elif isinstance(event, StepEvent):
+                # UI splitting strategy aligned with OpenAI flow:
+                # - do NOT start a new DB item at the first StepEvent
+                # - only finalize the previous item if it already produced content
+                #   (prevents empty items and ordering glitches)
                 self.set_busy(signals)
                 if not use_partials:
+                    # We still want to propagate the name early if provided.
+                    try:
+                        meta = getattr(event, "meta", {}) or {}
+                        next_name = meta.get("agent_name")
+                        if next_name:
+                            last_agent_name = self._apply_agent_name_to_ctx(item_ctx, next_name, last_agent_name)
+                    except Exception:
+                        pass
+                    begin = True
                     continue
+
                 if verbose:
                     print("\n\n-----STEP-----\n\n")
                     print(f"[{event.name}] {event.index}/{event.total} meta={event.meta}")
-                if flush:
-                    item_ctx = self.on_next_ctx(
-                        item_ctx,
-                        signals=signals,
-                        begin=begin,
-                        stream=True,
-                    )
+
+                # If there was an open block with content -> finalize it to a new DB item.
+                if block_open and content_written:
+                    if flush:
+                        item_ctx = self.on_next_ctx(
+                            item_ctx,
+                            signals=signals,
+                            begin=begin,
+                            stream=True,
+                        )
+                    # Apply next agent name on the fresh ctx (so UI header is correct from token #1).
+                    try:
+                        meta = getattr(event, "meta", {}) or {}
+                        next_name = meta.get("agent_name")
+                        if next_name:
+                            last_agent_name = self._apply_agent_name_to_ctx(item_ctx, next_name, last_agent_name)
+                    except Exception:
+                        pass
+                else:
+                    # First step or previous step had no visible content: just propagate the name.
+                    try:
+                        meta = getattr(event, "meta", {}) or {}
+                        next_name = meta.get("agent_name")
+                        if next_name:
+                            last_agent_name = self._apply_agent_name_to_ctx(item_ctx, next_name, last_agent_name)
+                    except Exception:
+                        pass
+
+                # Prepare for the upcoming tokens (new block begins).
+                block_open = True
+                content_written = False
+                begin = True
+                continue
+
             elif isinstance(event, AgentStream):
+                # Update agent name from event if present; fallback to header parsing.
+                name = getattr(event, "current_agent_name", None)
+                if not name:
+                    name = self._guess_agent_name_from_text(getattr(event, "delta", "") or "")
+                if name:
+                    last_agent_name = self._apply_agent_name_to_ctx(item_ctx, name, last_agent_name)
+
                 if verbose:
                     print(f"{event.delta}", end="", flush=True)
                 if event.delta:
                     item_ctx.live_output += event.delta
                     item_ctx.stream = event.delta
+                    content_written = True
                     if item_ctx.stream_agent_output and flush:
                         self.send_stream(item_ctx, signals, begin)  # send stream to webview
                     begin = False
+
             elif isinstance(event, AgentOutput):
+                # Ensure final agent name is applied as well.
+                name = getattr(event, "current_agent_name", None)
+                if name:
+                    last_agent_name = self._apply_agent_name_to_ctx(item_ctx, name, last_agent_name)
                 thought, answer = self.extract_final_response(str(event))
                 if answer:
                     item_ctx.set_agent_final_response(answer)
                     if verbose:
                         print(f"\nFinal response: {answer}")
+                # Do not split the block here – we will either:
+                # - split on the next StepEvent, or
+                # - finalize once at the end (make_response), just like OpenAI flow does.
 
         return item_ctx
 
@@ -348,6 +424,40 @@ class LlamaWorkflow(BaseRunner):
         next_ctx.set_output("")
         next_ctx.partial = True
         next_ctx.extra["agent_output"] = True  # allow usage in history
-
+        next_ctx.set_agent_name(ctx.get_agent_name())  # propagate agent name
         self.send_response(next_ctx, signals, KernelEvent.APPEND_DATA)
-        return next_ctx
+
+        return next_ctx
+
+    # ===== helpers for agent name propagation =====
+
+    def _apply_agent_name_to_ctx(self, ctx: CtxItem, name: str, last_known: Optional[str]) -> str:
+        """
+        Apply agent name to your context, avoiding redundant updates.
+        Falls back to ctx.extra['agent_name'] if set_agent_name is unavailable.
+        """
+        if not name:
+            return last_known or ""
+        if last_known and last_known == name:
+            return last_known
+        try:
+            if hasattr(ctx, "set_agent_name") and callable(getattr(ctx, "set_agent_name")):
+                ctx.set_agent_name(name)
+            # Always mirror into extra for downstream consumers
+            ctx.extra["agent_name"] = name
+        except Exception:
+            ctx.extra["agent_name"] = name
+        return name
+
+    def _guess_agent_name_from_text(self, text: str) -> Optional[str]:
+        """
+        Try to infer agent name from header like '**Name**' which our workflow emits
+        before each agent block.
+        """
+        if not text:
+            return None
+        # Look for the first bold segment – keep it lenient
+        m = re.search(r"\*\*([^*]+?)\*\*", text)
+        if m:
+            return m.group(1).strip()
+        return None

pygpt_net/core/agents/runners/openai_workflow.py

@@ -6,7 +6,7 @@
 # GitHub:  https://github.com/szczyglis-dev/py-gpt   #
 # MIT License                                        #
 # Created By  : Marcin Szczygliński                  #
-# Updated Date: 2025.08.24 03:00:00                  #
+# Updated Date: 2025.09.26 17:00:00                  #
 # ================================================== #
 
 from typing import Dict, Any, List, Optional
@@ -237,6 +237,7 @@ class OpenAIWorkflow(BaseRunner):
         response_ctx.set_agent_final_response(output)  # always set to further use
         response_ctx.extra["agent_output"] = True  # mark as output response
         response_ctx.extra["agent_finish"] = True  # mark as finished
+        response_ctx.set_agent_name(ctx.get_agent_name()) # store last agent name
         response_ctx.msg_id = response_id  # set response id for OpenAI
 
         if ctx.agent_final_response:  # only if not empty

pygpt_net/core/node_editor/types.py

@@ -6,7 +6,7 @@
 # GitHub:  https://github.com/szczyglis-dev/py-gpt   #
 # MIT License                                        #
 # Created By  : Marcin Szczygliński                  #
-# Updated Date: 2025.09.25 15:00:00                  #
+# Updated Date: 2025.09.26 12:00:00                  #
 # ================================================== #
 
 from __future__ import annotations
@@ -34,11 +34,14 @@ class PropertySpec:
 class NodeTypeSpec:
     type_name: str
     title: Optional[str] = None
+    # UI-only human-readable label used for menus; never persisted nor used as an identifier
+    display_name: Optional[str] = None
     properties: List[PropertySpec] = field(default_factory=list)
     # Below are optional extensions for agent-flow needs:
     base_id: Optional[str] = None        # base prefix for friendly ids, e.g. "agent"
     export_kind: Optional[str] = None    # short kind for export, e.g. "agent", "start"
     bg_color: Optional[str] = None       # optional per-type background color (CSS/hex)
+    max_num: Optional[int] = None        # optional per-layout cap; None or <=0 means unlimited
 
 class NodeTypeRegistry:
     """Registry for node type specifications. Extend/override in subclasses."""
@@ -56,6 +59,15 @@ class NodeTypeRegistry:
     def get(self, type_name: str) -> Optional[NodeTypeSpec]:
         return self._types.get(type_name)
 
+    def display_name(self, type_name: str) -> str:
+        """Return UI label for a type: spec.display_name if non-empty, otherwise type_name."""
+        spec = self.get(type_name)
+        if spec:
+            dn = getattr(spec, "display_name", None)
+            if isinstance(dn, str) and dn.strip():
+                return dn
+        return type_name
+
     def _install_default_types(self):
         # Example/basic nodes kept intact
         self.register(NodeTypeSpec(