|
7 | 7 |
|
8 | 8 | import os |
9 | 9 | from pathlib import Path |
10 | | -from typing import Literal |
| 10 | +from typing import Annotated, Literal, TypedDict |
11 | 11 |
|
12 | 12 | from langchain_core.messages import SystemMessage, HumanMessage |
13 | 13 | from langchain_core.tools import tool |
14 | 14 | from langchain_openai import ChatOpenAI |
15 | 15 | from langgraph.graph import StateGraph, MessagesState, END |
| 16 | +from langgraph.graph.message import add_messages |
16 | 17 | from langgraph.prebuilt import ToolNode |
17 | 18 | from langgraph_sdk import get_client |
18 | 19 |
|
19 | | -from src.aviation_tools import ( |
20 | | - get_airport_info, |
21 | | - find_routes, |
22 | | - lookup_flight, |
23 | | -) |
24 | | - |
25 | 20 | PROMPTS_DIR = Path(__file__).parent.parent / "prompts" |
26 | 21 |
|
27 | 22 | # ── generate_title node (inline; matches Pattern D from spec |
@@ -75,87 +70,110 @@ async def generate_title(state: MessagesState, config) -> dict: |
75 | 70 | return {} |
76 | 71 |
|
77 | 72 |
|
78 | | -_RESEARCH_PROMPT = """You are a Research Agent for trip planning. Your job is to gather |
79 | | -destination intel about airports the traveler is considering. Use the |
80 | | -get_airport_info tool to look up airport details (city, weather, terminals, |
81 | | -runways) for any airport codes mentioned in the task description. |
| 73 | +_RESEARCH_PROMPT = """You are a Research Agent for trip planning. Given a task |
| 74 | +describing one or more airports, return destination intel about them: city, |
| 75 | +typical weather, major terminals, and notable travel considerations. Draw on |
| 76 | +general knowledge of the airports named in the task. |
| 77 | +
|
| 78 | +Return a concise 2-4 sentence summary. If an airport isn't recognizable, say so.""" |
| 79 | + |
| 80 | +_BOOKING_PROMPT = """You are a Booking Agent for trip planning. Given an origin and |
| 81 | +destination in the task description, describe realistic flight options between |
| 82 | +them: which major carriers fly the route nonstop, typical durations, and rough |
| 83 | +fare expectations. |
| 84 | +
|
| 85 | +Return a concise summary listing 2-3 plausible options with airline, an example |
| 86 | +flight number, times, and price-or-aircraft info.""" |
| 87 | + |
| 88 | +_ITINERARY_PROMPT = """You are an Itinerary Agent for trip planning. Synthesize a |
| 89 | +final trip plan from the research + booking outputs you receive in the task |
| 90 | +description. |
| 91 | +
|
| 92 | +Return a clean 3-5 sentence itinerary summarizing the recommended flight choice, |
| 93 | +what to expect on arrival (weather), and any practical tips (e.g., terminal info, |
| 94 | +buffer time). Be helpful and concise.""" |
82 | 95 |
|
83 | | -Return a concise 2-4 sentence summary of what you found. If a code isn't |
84 | | -recognized, say so.""" |
85 | 96 |
|
86 | | -_BOOKING_PROMPT = """You are a Booking Agent for trip planning. Your job is to find |
87 | | -flight options between the origin and destination airports in the task |
88 | | -description. Use find_routes to list available flights, and lookup_flight |
89 | | -if the user mentioned a specific flight number. |
| 97 | +# subagent_type → system prompt. Keyed by the same Literal the `task` tool |
| 98 | +# exposes, so one parameterized subgraph serves all three specialists. |
| 99 | +_SUBAGENT_PROMPTS: dict[str, str] = { |
| 100 | + "research": _RESEARCH_PROMPT, |
| 101 | + "booking": _BOOKING_PROMPT, |
| 102 | + "itinerary": _ITINERARY_PROMPT, |
| 103 | +} |
90 | 104 |
|
91 | | -Return a concise summary listing 2-3 best flight options with airline, |
92 | | -flight number, times, and price-or-aircraft info. If no flights are found, |
93 | | -say so and suggest alternatives.""" |
94 | 105 |
|
95 | | -_ITINERARY_PROMPT = """You are an Itinerary Agent for trip planning. Your job is to |
96 | | -synthesize a final trip plan from research + booking outputs you receive in |
97 | | -the task description. |
| 106 | +class SubagentState(TypedDict): |
| 107 | + """Child-graph state. `subagent_type` selects the system prompt.""" |
| 108 | + messages: Annotated[list, add_messages] |
| 109 | + subagent_type: str |
| 110 | + task_description: str |
98 | 111 |
|
99 | | -Return a clean 3-5 sentence itinerary summarizing the recommended flight |
100 | | -choice, what to expect on arrival (weather), and any practical tips |
101 | | -(e.g., delays, terminal info). Be helpful and concise.""" |
102 | 112 |
|
| 113 | +async def _subagent_node(state: SubagentState) -> dict: |
| 114 | + """Focused subagent: a single role-prompted LLM call. Kept to ONE LLM call |
| 115 | + (no within-subagent tool loop) so each subagent's request has a unique, |
| 116 | + stable discriminator (its role-specific task_description) — this lets the |
| 117 | + aimock e2e replay match it deterministically. The within-subagent tool |
| 118 | + calling is exercised by the dedicated tool-calls cap; here the focus is |
| 119 | + subagent orchestration + the inline subagent card. The returned message |
| 120 | + streams under this subgraph's `tools:<call_id>` namespace, which the |
| 121 | + @threadplane/langgraph SubagentTracker matches to surface the card.""" |
| 122 | + subagent_type = state["subagent_type"] |
| 123 | + task_description = state["task_description"] |
| 124 | + system_prompt = _SUBAGENT_PROMPTS.get(subagent_type, _ITINERARY_PROMPT) |
103 | 125 |
|
104 | | -async def _run_subagent(role: str, task_description: str, system_prompt: str, tools: list): |
105 | | - """Run a single subagent: LLM bound with role-specific tools, single tool loop.""" |
106 | 126 | llm = ChatOpenAI(model="gpt-5-mini", streaming=True) |
107 | | - if tools: |
108 | | - llm = llm.bind_tools(tools) |
109 | | - messages = [ |
| 127 | + response = await llm.ainvoke([ |
110 | 128 | SystemMessage(content=system_prompt), |
111 | 129 | HumanMessage(content=task_description), |
112 | | - ] |
113 | | - # Allow up to 3 tool-loop iterations |
114 | | - for _ in range(3): |
115 | | - response = await llm.ainvoke(messages) |
116 | | - messages.append(response) |
117 | | - tool_calls = getattr(response, "tool_calls", None) |
118 | | - if not tool_calls: |
119 | | - return response.content |
120 | | - # Execute tool calls inline |
121 | | - for tc in tool_calls: |
122 | | - tool_name = tc["name"] |
123 | | - tool_args = tc["args"] |
124 | | - target = next((t for t in tools if t.name == tool_name), None) |
125 | | - if target is None: |
126 | | - tool_result = f"Tool {tool_name} not available" |
127 | | - else: |
128 | | - tool_result = await target.ainvoke(tool_args) |
129 | | - from langchain_core.messages import ToolMessage |
130 | | - messages.append(ToolMessage(content=str(tool_result), tool_call_id=tc["id"])) |
131 | | - return response.content |
| 130 | + ]) |
| 131 | + return {"messages": [response]} |
132 | 132 |
|
133 | 133 |
|
134 | | -@tool |
135 | | -async def task(role: Literal["research", "booking", "itinerary"], task_description: str) -> str: |
136 | | - """Delegate a subtask to a specialized subagent. |
| 134 | +# Compiled child graph. Invoking it from inside the `task` tool makes LangGraph |
| 135 | +# nest its run under a `tools:<call_id>` namespace, which the @threadplane/langgraph |
| 136 | +# SubagentTracker matches to the registered `task` dispatch to surface a card. |
| 137 | +_subagent_builder = StateGraph(SubagentState) |
| 138 | +_subagent_builder.add_node("subagent", _subagent_node) |
| 139 | +_subagent_builder.set_entry_point("subagent") |
| 140 | +_subagent_builder.add_edge("subagent", END) |
| 141 | +subagent_subgraph = _subagent_builder.compile() |
| 142 | + |
137 | 143 |
|
138 | | - Roles: |
139 | | - - research: gathers destination intel (airports, weather, conditions) |
140 | | - - booking: finds flight options between origin and destination |
141 | | - - itinerary: synthesizes a final trip plan combining research + bookings |
| 144 | +def _final_text(messages: list) -> str: |
| 145 | + """Last non-empty string content from the child graph's messages.""" |
| 146 | + for msg in reversed(messages or []): |
| 147 | + content = getattr(msg, "content", None) |
| 148 | + if isinstance(content, str) and content.strip(): |
| 149 | + return content |
| 150 | + if isinstance(content, list): |
| 151 | + parts = [b.get("text", "") for b in content if isinstance(b, dict) and b.get("type") == "text"] |
| 152 | + if any(p.strip() for p in parts): |
| 153 | + return "\n".join(parts) |
| 154 | + return "(no subagent output)" |
| 155 | + |
| 156 | + |
| 157 | +@tool |
| 158 | +async def task(subagent_type: Literal["research", "booking", "itinerary"], task_description: str) -> str: |
| 159 | + """Delegate a subtask to a specialized subagent subgraph. |
142 | 160 |
|
143 | 161 | Args: |
144 | | - role: One of "research", "booking", "itinerary". |
145 | | - task_description: Plain-English description of what the subagent |
146 | | - should do (e.g., "Gather info on LAX and JFK airports", or |
147 | | - "Find morning flights from LAX to JFK"). |
| 162 | + subagent_type: Which specialist to dispatch — "research" (airport / |
| 163 | + destination intel), "booking" (flight options between origin and |
| 164 | + destination), or "itinerary" (final trip plan synthesizing research |
| 165 | + + bookings). This label also identifies the subagent in the UI. |
| 166 | + task_description: Plain-English description of what the subagent should |
| 167 | + do (e.g., "Gather info on LAX and JFK airports"). |
148 | 168 |
|
149 | 169 | Returns: |
150 | 170 | The subagent's final answer as a string. |
151 | 171 | """ |
152 | | - if role == "research": |
153 | | - return await _run_subagent(role, task_description, _RESEARCH_PROMPT, [get_airport_info]) |
154 | | - if role == "booking": |
155 | | - return await _run_subagent(role, task_description, _BOOKING_PROMPT, [find_routes, lookup_flight]) |
156 | | - if role == "itinerary": |
157 | | - return await _run_subagent(role, task_description, _ITINERARY_PROMPT, []) |
158 | | - return f"Unknown role: {role}" |
| 172 | + result = await subagent_subgraph.ainvoke( |
| 173 | + {"subagent_type": subagent_type, "task_description": task_description, "messages": []} |
| 174 | + ) |
| 175 | + messages = result.get("messages") if isinstance(result, dict) else None |
| 176 | + return _final_text(messages) |
159 | 177 |
|
160 | 178 |
|
161 | 179 | def build_subagents_graph(): |
|
0 commit comments