Fix/mg/multimodal support (#42)

galuszkm · web-flow · commit bca982f88f63 · 2026-05-20T20:22:57.000+02:00
diff --git a/examples/01_minimal/main.py b/examples/01_minimal/main.py
@@ -29,10 +29,9 @@ def main() -> None:
             msg = input("You: ").strip()
             if not msg:
                 break
-            print("\nAgent is thinking...")
-            result = agent(msg)
-            print(f"\n\n{52 * '-'}")
-            print(f"Agent: {result}\n")
+            print()
+            agent(msg)
+            print("\n" + 52 * "-" + "\n")
     except KeyboardInterrupt:
         print("\nGoodbye!")
     finally:
diff --git a/examples/02_vars_and_anchors/main.py b/examples/02_vars_and_anchors/main.py
@@ -41,10 +41,9 @@ def main() -> None:
             msg = input("You: ").strip()
             if not msg:
                 break
-            print("\nAgent is thinking...")
-            result = agent(msg)
-            print(f"\n\n{52 * '-'}")
-            print(f"Agent: {result}\n")
+            print()
+            agent(msg)
+            print("\n" + 52 * "-" + "\n")
     except KeyboardInterrupt:
         print("\nGoodbye!")
     finally:
diff --git a/examples/03_tools/main.py b/examples/03_tools/main.py
@@ -30,10 +30,9 @@ def main() -> None:
             msg = input("You: ").strip()
             if not msg:
                 break
-            print("\nAgent is thinking...")
-            result = agent(msg)
-            print(f"\n\n{52 * '-'}")
-            print(f"Agent: {result}\n")
+            print()
+            agent(msg)
+            print("\n" + 52 * "-" + "\n")
     except KeyboardInterrupt:
         print("\nGoodbye!")
     finally:
diff --git a/examples/04_session/main.py b/examples/04_session/main.py
@@ -29,10 +29,9 @@ def main() -> None:
             msg = input("You: ").strip()
             if not msg:
                 break
-            print("\nAgent is thinking...")
-            result = agent(msg)
-            print(f"\n\n{52 * '-'}")
-            print(f"Agent: {result}\n")
+            print()
+            agent(msg)
+            print("\n" + 52 * "-" + "\n")
     except KeyboardInterrupt:
         print("\nGoodbye!")
     finally:
diff --git a/examples/05_hooks/main.py b/examples/05_hooks/main.py
@@ -32,10 +32,9 @@ def main() -> None:
             msg = input("You: ").strip()
             if not msg:
                 break
-            print("\nAgent is thinking...")
-            result = agent(msg)
-            print(f"\n\n{52 * '-'}")
-            print(f"Agent: {result}\n")
+            print()
+            agent(msg)
+            print("\n" + 52 * "-" + "\n")
     except KeyboardInterrupt:
         print("\nGoodbye!")
     finally:
diff --git a/examples/06_mcp/main.py b/examples/06_mcp/main.py
@@ -31,10 +31,9 @@ def main() -> None:
             msg = input("You: ").strip()
             if not msg:
                 break
-            print("\nAgent is thinking...")
-            result = agent(msg)
-            print(f"\n\n{52 * '-'}")
-            print(f"Agent: {result}\n")
+            print()
+            agent(msg)
+            print("\n" + 52 * "-" + "\n")
     except KeyboardInterrupt:
         print("\nGoodbye!")
     finally:
diff --git a/examples/07_delegate/main.py b/examples/07_delegate/main.py
@@ -29,10 +29,9 @@ def main() -> None:
             msg = input("You: ").strip()
             if not msg:
                 break
-            print("\nAgent is thinking...")
-            result = agent(msg)
-            print(f"\n\n{52 * '-'}")
-            print(f"Agent: {result}\n")
+            print()
+            agent(msg)
+            print("\n" + 52 * "-" + "\n")
     except KeyboardInterrupt:
         print("\nGoodbye!")
     finally:
diff --git a/examples/08_swarm/main.py b/examples/08_swarm/main.py
@@ -29,10 +29,9 @@ def main() -> None:
             msg = input("You: ").strip()
             if not msg:
                 break
-            print("\nSwarm is working...")
-            result = swarm(msg)
-            print(f"\n\n{52 * '-'}")
-            print(f"Agent: {result}\n")
+            print()
+            swarm(msg)
+            print("\n" + 52 * "-" + "\n")
     except KeyboardInterrupt:
         print("\nGoodbye!")
     finally:
diff --git a/examples/09_graph/main.py b/examples/09_graph/main.py
@@ -29,10 +29,9 @@ def main() -> None:
             msg = input("You: ").strip()
             if not msg:
                 break
-            print("\nPipeline is running...")
-            result = graph(msg)
-            print(f"\n\n{52 * '-'}")
-            print(f"Agent: {result}\n")
+            print()
+            graph(msg)
+            print("\n" + 52 * "-" + "\n")
     except KeyboardInterrupt:
         print("\nGoodbye!")
     finally:
diff --git a/examples/10_nested/main.py b/examples/10_nested/main.py
@@ -30,10 +30,9 @@ def main() -> None:
             msg = input("You: ").strip()
             if not msg:
                 break
-            print("\nAgent is thinking...")
-            result = agent(msg)
-            print(f"\n\n{52 * '-'}")
-            print(f"Agent: {result}\n")
+            print()
+            agent(msg)
+            print("\n" + 52 * "-" + "\n")
     except KeyboardInterrupt:
         print("\nGoodbye!")
     finally:
diff --git a/examples/11_multi_file_config/main.py b/examples/11_multi_file_config/main.py
@@ -32,10 +32,9 @@ def main() -> None:
             msg = input("You: ").strip()
             if not msg:
                 break
-            print("\nAgent is thinking...")
-            result = agent(msg)
-            print(f"\n\n{52 * '-'}")
-            print(f"Agent: {result}\n")
+            print()
+            agent(msg)
+            print("\n" + 52 * "-" + "\n")
     except KeyboardInterrupt:
         print("\nGoodbye!")
     finally:
diff --git a/examples/12_streaming/main.py b/examples/12_streaming/main.py
@@ -57,10 +57,9 @@ async def _main() -> None:
             msg = await asyncio.to_thread(lambda: input("You: ").strip())
             if not msg:
                 break
-            print("\nAgent is thinking...")
-            result = await _stream(msg, entry, queue)
-            print(f"\n\n{52 * '-'}")
-            print(f"Agent: {result}\n")
+            print()
+            await _stream(msg, entry, queue)
+            print("\n" + 52 * "-" + "\n")
     except KeyboardInterrupt:
         print("\nGoodbye!")
     finally:
diff --git a/examples/13_graph_conditions/main.py b/examples/13_graph_conditions/main.py
@@ -30,10 +30,9 @@ def main() -> None:
             msg = input("You: ").strip()
             if not msg:
                 break
-            print("\nPipeline is running...")
-            result = graph(msg)
-            print(f"\n\n{52 * '-'}")
-            print(f"Agent: {result}\n")
+            print()
+            graph(msg)
+            print("\n" + 52 * "-" + "\n")
     except KeyboardInterrupt:
         print("\nGoodbye!")
     finally:
diff --git a/examples/14_agent_factory/main.py b/examples/14_agent_factory/main.py
@@ -29,10 +29,9 @@ def main() -> None:
             msg = input("You: ").strip()
             if not msg:
                 break
-            print("\nAgent is thinking...")
-            result = agent(msg)
-            print(f"\n\n{52 * '-'}")
-            print(f"Agent: {result}\n")
+            print()
+            agent(msg)
+            print("\n" + 52 * "-" + "\n")
     except KeyboardInterrupt:
         print("\nGoodbye!")
     finally:
diff --git a/src/strands_compose/config/loaders/loaders.py b/src/strands_compose/config/loaders/loaders.py
@@ -39,7 +39,7 @@
     resolve_orchestrations,
     resolve_session_manager,
 )
-from ..schema import AppConfig, SwarmOrchestrationDef
+from ..schema import AppConfig, GraphOrchestrationDef, SwarmOrchestrationDef
 from .helpers import merge_raw_configs, parse_single_source, sanitize_collection_keys
 from .validators import validate_references
 
@@ -236,19 +236,24 @@ def load_session(
             session_id_override=session_id,
         )
 
-    # Agents used in Swarm orchestrations cannot have session_manager set.
-    # Temporary until strands-agents supports swarm agents with session persistence.
-    swarm_agent_names: set[str] = set()
+    # Agents used in Swarm or Graph orchestrations cannot have session_manager set.
+    # Temporary until strands-agents supports session persistence for orchestration node agents.
+    orchestration_agent_names: set[str] = set()
     for orch in config.orchestrations.values():
         if isinstance(orch, SwarmOrchestrationDef):
-            swarm_agent_names.update(orch.agents)
+            orchestration_agent_names.update(orch.agents)
+        elif isinstance(orch, GraphOrchestrationDef):
+            orchestration_agent_names.add(orch.entry_name)
+            for edge in orch.edges:
+                orchestration_agent_names.add(edge.from_agent)
+                orchestration_agent_names.add(edge.to_agent)
 
     agents = resolve_agents(
         agent_defs=config.agents,
         models=infra.models,
         mcp_clients=infra.clients,
         session_manager=session_manager,
-        swarm_agent_names=swarm_agent_names,
+        orchestration_agent_names=orchestration_agent_names,
     )
     orchestrators = resolve_orchestrations(
         config,
diff --git a/src/strands_compose/config/resolvers/agents.py b/src/strands_compose/config/resolvers/agents.py
@@ -36,7 +36,7 @@ def build_agent_from_def(
     extra_tools: list[Any] | None = None,
     extra_hooks: list[Any] | None = None,
     session_manager_override: SessionManager | None = None,
-    swarm_agent_names: set[str] | None = None,
+    orchestration_agent_names: set[str] | None = None,
 ) -> Agent:
     """Build a single Agent from an AgentDef blueprint.
 
@@ -54,13 +54,13 @@ def build_agent_from_def(
         extra_tools: Additional tools to append (e.g. delegate tools).
         extra_hooks: Additional hooks to append (e.g. orchestration-level hooks).
         session_manager_override: If set, overrides session manager resolution entirely.
-        swarm_agent_names: Agent names in swarm orchestrations (fail-fast on SM).
+        orchestration_agent_names: Agent names in swarm/graph orchestrations (fail-fast on SM).
 
     Returns:
         A freshly constructed strands Agent.
 
     Raises:
-        ConfigurationError: If a swarm agent has a session manager.
+        ConfigurationError: If a swarm or graph agent has a session manager.
         TypeError: If a custom factory doesn't return an Agent.
     """
     # 1. Resolve model — inline ModelDef or global name ref
@@ -94,17 +94,21 @@ def build_agent_from_def(
         else:
             agent_session = session_manager
 
-    # Fail fast: swarm node agents cannot carry a session manager.
-    if swarm_agent_names and name in swarm_agent_names and agent_session is not None:
+    # Fail fast: swarm/graph node agents cannot carry a session manager.
+    if (
+        orchestration_agent_names
+        and name in orchestration_agent_names
+        and agent_session is not None
+    ):
         source = (
             "per-agent 'session_manager:' field"
             if agent_def.session_manager is not None
             else "global 'session_manager:' in config"
         )
         raise ConfigurationError(
-            f"Agent '{name}' is in swarm orchestration and cannot have a session manager "
+            f"Agent '{name}' is in a swarm or graph orchestration and cannot have a session manager "
             f"(source: {source}).\n"
-            "Strands does not yet support session persistence for Swarm node agents.\n"
+            "Strands does not yet support session persistence for Swarm or Graph node agents.\n"
             f"Fix: Add 'session_manager: ~' to agent '{name}' to opt out of the global default."
         )
 
@@ -159,7 +163,7 @@ def resolve_agents(
     models: dict[str, Model],
     mcp_clients: dict[str, StrandsMCPClient],
     session_manager: SessionManager | None,
-    swarm_agent_names: set[str] | None = None,
+    orchestration_agent_names: set[str] | None = None,
 ) -> dict[str, Agent]:
     """Resolve all agents -- flat, no dependencies between them.
 
@@ -176,14 +180,14 @@ def resolve_agents(
         models: Resolved model instances keyed by name.
         mcp_clients: Resolved MCP clients keyed by name.
         session_manager: Optional shared session manager for all agents.
-        swarm_agent_names: Names of agents that participate in swarm orchestrations
+        orchestration_agent_names: Names of agents in swarm or graph orchestrations
             (these agents must not have a session manager).
 
     Returns:
         Resolved agents dict keyed by name.
 
     Raises:
-        ConfigurationError if an agent used in a Swarm orchestration has a session manager assigned.
+        ConfigurationError if an agent used in a Swarm or Graph orchestration has a session manager assigned.
         TypeError if agent factory does not return a strands Agent instance.
     """
     resolved: dict[str, Agent] = {}
@@ -195,7 +199,7 @@ def resolve_agents(
             models=models,
             mcp_clients=mcp_clients,
             session_manager=session_manager,
-            swarm_agent_names=swarm_agent_names,
+            orchestration_agent_names=orchestration_agent_names,
         )
         resolved[name] = agent
         logger.info("agent=<%s>, type=<%s> | resolved agent", name, agent_def.type or "Agent")
diff --git a/src/strands_compose/tools/extractors.py b/src/strands_compose/tools/extractors.py
diff --git a/src/strands_compose/tools/wrappers.py b/src/strands_compose/tools/wrappers.py
diff --git a/tests/unit/config/loaders/test_helpers_extended.py b/tests/unit/config/loaders/test_helpers_extended.py
diff --git a/tests/unit/config/loaders/test_load_session.py b/tests/unit/config/loaders/test_load_session.py
diff --git a/tests/unit/config/resolvers/orchestrations/test_tools.py b/tests/unit/config/resolvers/orchestrations/test_tools.py
diff --git a/tests/unit/config/resolvers/test_agents.py b/tests/unit/config/resolvers/test_agents.py
diff --git a/tests/unit/mcp/test_server.py b/tests/unit/mcp/test_server.py
diff --git a/uv.lock b/uv.lock
