feat: parallel subagent invocations (#69)

Author: ottokruse

.gitignore

@@ -4,6 +4,7 @@
 __pycache__/
 *.egg-info
 dist
+build
 .env
 
 # Directory for local CDK context files

README.md

@@ -18,7 +18,7 @@
     "\n",
     "We can implement the multi-agent architecture by simply registering subordinate agents as tools with a higher-level agent. To be able to use an agent as tool, the agent must have a `name` and a `description`. Providing an `input_schema` is optional, but can be used to signal to the supervisor agent that it should send certain fields to the subordinate agent.\n",
     "\n",
-    "In the example here we just have a 2-level hierarchy: one supervisor, and 2 subordinate agents. You could add more levels if you want; one of the subordinate agents could itself be a supervisor to its own set of subordinate agents."
+    "In the example here we just have a 2-level hierarchy: one supervisor, and 2 subordinate agents. You could add more levels if you want; one of the subordinate agents could itself be a supervisor to its own set of subordinate agents.\n"
    ]
   },
   {
@@ -39,7 +39,7 @@
    "id": "00d4232c",
    "metadata": {},
    "source": [
-    "## Weather agent"
+    "## Weather agent\n"
    ]
   },
   {
@@ -75,7 +75,7 @@
    "id": "b8a8e867",
    "metadata": {},
    "source": [
-    "## Events agent"
+    "## Events agent\n"
    ]
   },
   {
@@ -130,7 +130,7 @@
    "id": "415af9d9",
    "metadata": {},
    "source": [
-    "## Supervisor agent"
+    "## Supervisor agent\n"
    ]
   },
   {
@@ -178,7 +178,7 @@
    "source": [
     "## Chat with the supervisor agent\n",
     "\n",
-    "The supervisor agent will use the subordinate agents to help the user:"
+    "The supervisor agent will use the subordinate agents to help the user:\n"
    ]
   },
   {
@@ -202,7 +202,9 @@
    "source": [
     "## Tracing\n",
     "\n",
-    "Each agent has its own set of traces that you can inspect"
+    "### All traces from supervisor (includes subagent traces)\n",
+    "\n",
+    "The supervisor's traces include all traces from subagent invocations:\n"
    ]
   },
   {
@@ -216,32 +218,127 @@
     "supervisor_ui.launch()"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "subcontext_explanation",
+   "metadata": {},
+   "source": [
+    "### Accessing individual subagent traces\n",
+    "\n",
+    "When agents are used as tools, each invocation gets its own **subcontext ID** to maintain separate conversational memory. This enables:\n",
+    "\n",
+    "- **Parallel invocations**: The supervisor could call both agents simultaneously, each with separate memory\n",
+    "- **Sequential invocations without memory carryover**: Calling the weather agent twice won't leak context between invocations\n",
+    "- **Follow-up messages**: You can send additional messages to a specific subagent invocation by reusing its subcontext ID\n",
+    "\n",
+    "To view traces for a specific subagent invocation, we need to:\n",
+    "\n",
+    "1. Extract the subcontext ID from the supervisor's traces\n",
+    "2. Set the correct conversation ID and subcontext ID on the subagent\n",
+    "3. Then access the subagent's traces\n",
+    "\n",
+    "First, let's add a helper function to extract subcontext IDs:\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "helper_function",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from collections import defaultdict\n",
+    "\n",
+    "\n",
+    "def get_subagent_subcontext_ids(traces):\n",
+    "    \"\"\"\n",
+    "    Extract subcontext IDs used for each subagent invocation from traces.\n",
+    "\n",
+    "    Returns a dict mapping subagent name -> set of subcontext IDs used\n",
+    "    \"\"\"\n",
+    "    result = defaultdict(set)\n",
+    "\n",
+    "    for trace in traces:\n",
+    "        if (\n",
+    "            trace.attributes.get(\"ai.trace.type\") == \"tool-invocation\"\n",
+    "            and \"ai.tool.subagent.subcontext.id\" in trace.attributes\n",
+    "            and \"ai.tool.name\" in trace.attributes\n",
+    "        ):\n",
+    "            subagent_name = trace.attributes[\"ai.tool.name\"]\n",
+    "            subcontext_id = trace.attributes[\"ai.tool.subagent.subcontext.id\"]\n",
+    "            result[subagent_name].add(subcontext_id)\n",
+    "\n",
+    "    return result"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "weather_agent_traces_explanation",
+   "metadata": {},
+   "source": [
+    "Now let's access the weather agent's traces for this specific invocation:\n"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
    "id": "995050bc",
    "metadata": {},
    "outputs": [],
    "source": [
+    "# Store the conversation ID\n",
+    "conversation_id = supervisor.conversation_id\n",
+    "\n",
+    "# Extract subcontext IDs from the supervisor's traces\n",
+    "subcontext_ids = get_subagent_subcontext_ids(supervisor.traces)\n",
+    "\n",
+    "# Get the specific subcontext_id used for weather_agent\n",
+    "weather_subcontext_id = subcontext_ids[\"transfer_to_weather_agent\"].pop()\n",
+    "\n",
+    "# Set the weather agent to the correct conversation and subcontext\n",
+    "weather_agent.set_conversation_id(conversation_id, subcontext_id=weather_subcontext_id)\n",
+    "\n",
+    "# Now weather_agent.traces returns the correct traces for this invocation\n",
+    "print(f\"Weather agent subcontext ID: {weather_subcontext_id}\")\n",
+    "print(f\"Number of traces: {len(weather_agent.traces)}\")\n",
+    "\n",
     "weather_agent_ui = traces_ui(weather_agent.traces)\n",
     "weather_agent_ui.launch()"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "events_agent_traces_explanation",
+   "metadata": {},
+   "source": [
+    "And the events agent's traces:\n"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
    "id": "bc2a53db",
    "metadata": {},
    "outputs": [],
    "source": [
+    "# Get the specific subcontext_id used for events_agent\n",
+    "events_subcontext_id = subcontext_ids[\"transfer_to_events_agent\"].pop()\n",
+    "\n",
+    "# Set the events agent to the correct conversation and subcontext\n",
+    "events_agent.set_conversation_id(conversation_id, subcontext_id=events_subcontext_id)\n",
+    "\n",
+    "# Now events_agent.traces returns the correct traces for this invocation\n",
+    "print(f\"Events agent subcontext ID: {events_subcontext_id}\")\n",
+    "print(f\"Number of traces: {len(events_agent.traces)}\")\n",
+    "\n",
     "events_agent_ui = traces_ui(events_agent.traces)\n",
     "events_agent_ui.launch()"
    ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": ".venv",
+   "display_name": "generative-ai-toolkit",
    "language": "python",
    "name": "python3"
   },

examples/multi_agent.ipynb

@@ -152,7 +152,8 @@
     "\n",
     "for trace in in_memory_tracer.get_traces(\n",
     "    attribute_filter={\n",
-    "        \"ai.conversation.id\": conversation_id  # filter traces by conversation id\n",
+    "        \"ai.conversation.id\": conversation_id,  # filter traces by conversation id\n",
+    "        \"ai.subcontext.id\": None\n",
     "    }\n",
     "):\n",
     "    print(trace.as_human_readable())\n",
@@ -361,6 +362,7 @@
     "    attribute_filter={\n",
     "        \"ai.conversation.id\": conversation_id,\n",
     "        \"ai.auth.context\": auth_context,\n",
+    "        \"ai.subcontext.id\": None\n",
     "    }\n",
     "):\n",
     "    print(trace.as_human_readable())\n",
@@ -681,6 +683,7 @@
     "    attribute_filter={\n",
     "        \"ai.conversation.id\": conversation_id,\n",
     "        \"ai.auth.context\": auth_context,\n",
+    "        \"ai.subcontext.id\": None\n",
     "    }\n",
     "):\n",
     "    print(trace.as_human_readable())\n",

examples/tracing101.ipynb

@@ -37,39 +37,26 @@ run-agent = [
     "opentelemetry-proto>=1.31,<2.0",
     "pydantic>=2.10,<3.0",
 ]
-evaluate = ["nltk>=3.9,<4.0", "scikit-learn>=1.5,<2.0"]
+evaluate = [
+    "nltk>=3.9,<4.0",
+    "scikit-learn>=1.5,<2.0",
+]
 all = [
+    "generative-ai-toolkit[run-agent,evaluate]",
     "boto3-stubs[bedrock-runtime,dynamodb]>=1.37,<2.0",
-    "flask>=3.1,<4.0",
     "gradio>=5.23,<6.0",
-    "gunicorn>=23.0,<24.0",
     "ipython>=8.30,<9.0",
-    "nltk>=3.9,<4.0",
     "mcp>=1.8,<2.0",
-    "opentelemetry-proto>=1.31,<2.0",
     "pandas>=2.2,<3.0",
-    "pydantic>=2.10,<3.0",
-    "scikit-learn>=1.5,<2.0",
     "tabulate>=0.9,<1.0",
 ]
 dev = [
-    "boto3-stubs[bedrock-runtime,dynamodb]>=1.37,<2.0",
-    "flask>=3.1,<4.0",
-    "gradio>=5.23,<6.0",
-    "gunicorn>=23.0,<24.0",
-    "ipython>=8.30,<9.0",
-    "mcp>=1.8,<2.0",
-    "nltk>=3.9,<4.0",
-    "opentelemetry-proto>=1.31,<2.0",
-    "pandas>=2.2,<3.0",
+    "generative-ai-toolkit[all]",
     "playwright>=1.52,<2.0",
-    "pydantic>=2.10,<3.0",
+    "pytest>=8.4,<9.0",
     "pytest-cov>=6.2,<7.0",
     "pytest-playwright>=0.7,<0.8",
-    "pytest>=8.4,<9.0",
     "ruff>=0.12,<0.13",
-    "scikit-learn>=1.5,<2.0",
-    "tabulate>=0.9,<1.0",
 ]
 
 [tool.setuptools.packages.find]

pyproject.toml

@@ -12,7 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from generative_ai_toolkit.agent.agent import Agent, BedrockConverseAgent
+from generative_ai_toolkit.agent.agent import Agent
+from generative_ai_toolkit.agent.bedrock_converse_agent import BedrockConverseAgent
 from generative_ai_toolkit.agent.tool import BedrockConverseTool, Tool
 
 __all__ = ["BedrockConverseAgent", "BedrockConverseTool", "Agent", "Tool"]