|
13 | 13 | # See the License for the specific language governing permissions and |
14 | 14 | # limitations under the License. |
15 | 15 |
|
| 16 | +import functools |
16 | 17 | import json |
17 | 18 | import logging |
| 19 | +import os |
18 | 20 | import time |
19 | | -from typing import Any, Dict, List |
| 21 | +from datetime import datetime |
| 22 | +from pathlib import Path |
| 23 | +from typing import Any, Dict, List, Optional |
20 | 24 |
|
21 | | -from langchain_core.callbacks import AsyncCallbackHandler |
| 25 | +from langchain_core.callbacks import AsyncCallbackHandler, BaseCallbackHandler |
22 | 26 | from rai.communication.ros2 import ROS2Connector, ROS2HRIMessage, ROS2Message |
23 | 27 |
|
24 | 28 |
|
@@ -220,3 +224,196 @@ def send_heartbeat(self): |
220 | 224 | msg_type="std_msgs/msg/Header", |
221 | 225 | target=self.heartbeat_topic, |
222 | 226 | ) |
| 227 | + |
| 228 | + |
| 229 | +def _stringify_content(content: Any) -> str: |
| 230 | + """Flatten a message's content to text. Multimodal blocks (VLM image parts) |
| 231 | + are summarized as ``[image]`` so the log stays readable and small.""" |
| 232 | + if content is None: |
| 233 | + return "" |
| 234 | + if isinstance(content, str): |
| 235 | + return content |
| 236 | + if isinstance(content, list): |
| 237 | + parts = [] |
| 238 | + for block in content: |
| 239 | + if isinstance(block, dict): |
| 240 | + btype = block.get("type") |
| 241 | + if btype == "text": |
| 242 | + parts.append(block.get("text", "")) |
| 243 | + elif btype in ("image_url", "image"): |
| 244 | + parts.append("[image]") |
| 245 | + else: |
| 246 | + parts.append(f"[{btype or 'block'}]") |
| 247 | + else: |
| 248 | + parts.append(str(block)) |
| 249 | + return "\n".join(p for p in parts if p) |
| 250 | + return str(content) |
| 251 | + |
| 252 | + |
| 253 | +def _guard(fn): |
| 254 | + """Never let a tracing error break the agent run.""" |
| 255 | + |
| 256 | + @functools.wraps(fn) |
| 257 | + def wrapper(self, *args, **kwargs): |
| 258 | + try: |
| 259 | + return fn(self, *args, **kwargs) |
| 260 | + except Exception as exc: # noqa: BLE001 - tracing must not crash the agent |
| 261 | + self.logger.warning(f"trace callback {fn.__name__} failed: {exc}") |
| 262 | + |
| 263 | + return wrapper |
| 264 | + |
| 265 | + |
| 266 | +class ConversationFileCallback(BaseCallbackHandler): |
| 267 | + """Dump the full LLM conversation (orchestrator + subagents) to a run directory. |
| 268 | +
|
| 269 | + Register it in the orchestrator's ``langchain_callbacks`` list. That list is |
| 270 | + passed to ``astream(config={"callbacks": ...}, subgraphs=True)``, and LangGraph |
| 271 | + propagates config callbacks into every subgraph, so both the megamind |
| 272 | + orchestrator and its executor subagents are captured in one place. |
| 273 | +
|
| 274 | + Writes two files per run: |
| 275 | + - ``log.txt`` human-readable transcript (messages, tool calls, results) |
| 276 | + - ``trace.jsonl`` one JSON record per event (machine-readable) |
| 277 | +
|
| 278 | + Enabled by default to ``runs/<timestamp>/``. Override the directory with |
| 279 | + ``AMM_TRACE_DIR``; disable entirely with ``AMM_TRACE=0``. |
| 280 | + """ |
| 281 | + |
| 282 | + # Run inline on the event loop so concurrent subgraph events don't interleave |
| 283 | + # partial lines in the files. |
| 284 | + run_inline = True |
| 285 | + |
| 286 | + def __init__(self, out_dir: str | Path): |
| 287 | + self.out_dir = Path(out_dir) |
| 288 | + self.out_dir.mkdir(parents=True, exist_ok=True) |
| 289 | + self.text_path = self.out_dir / "log.txt" |
| 290 | + self.jsonl_path = self.out_dir / "trace.jsonl" |
| 291 | + self.logger = logging.getLogger(__name__) |
| 292 | + self.logger.info(f"Conversation trace -> {self.out_dir}") |
| 293 | + |
| 294 | + @classmethod |
| 295 | + def from_env(cls) -> Optional["ConversationFileCallback"]: |
| 296 | + """Build from env, or return None if disabled (``AMM_TRACE=0``).""" |
| 297 | + if os.getenv("AMM_TRACE", "1") == "0": |
| 298 | + return None |
| 299 | + out_dir = os.getenv("AMM_TRACE_DIR") or os.path.join( |
| 300 | + "runs", datetime.now().strftime("%Y-%m-%d_%H-%M-%S") |
| 301 | + ) |
| 302 | + return cls(out_dir) |
| 303 | + |
| 304 | + # ── writers ────────────────────────────────────────────────────────────── |
| 305 | + @staticmethod |
| 306 | + def _ts() -> str: |
| 307 | + return datetime.now().strftime("%H:%M:%S") |
| 308 | + |
| 309 | + def _write_text(self, block: str) -> None: |
| 310 | + with open(self.text_path, "a", encoding="utf-8") as f: |
| 311 | + f.write(block.rstrip() + "\n") |
| 312 | + |
| 313 | + def _write_json(self, record: Dict[str, Any]) -> None: |
| 314 | + record = {"ts": datetime.now().isoformat(timespec="seconds"), **record} |
| 315 | + with open(self.jsonl_path, "a", encoding="utf-8") as f: |
| 316 | + f.write(json.dumps(record, default=str) + "\n") |
| 317 | + |
| 318 | + @staticmethod |
| 319 | + def _node(metadata: Any, tags: Any) -> str: |
| 320 | + """Which graph node / subagent produced this event.""" |
| 321 | + if isinstance(metadata, dict): |
| 322 | + node = metadata.get("langgraph_node") |
| 323 | + if node: |
| 324 | + return str(node) |
| 325 | + for tag in tags or []: |
| 326 | + if isinstance(tag, str) and ":" in tag: |
| 327 | + return tag |
| 328 | + return "?" |
| 329 | + |
| 330 | + # ── callbacks ──────────────────────────────────────────────────────────── |
| 331 | + @_guard |
| 332 | + def on_chat_model_start( |
| 333 | + self, serialized, messages, *, tags=None, metadata=None, **kwargs |
| 334 | + ): |
| 335 | + node = self._node(metadata, tags) |
| 336 | + rendered = [] |
| 337 | + for msg_list in messages: |
| 338 | + for m in msg_list: |
| 339 | + rendered.append( |
| 340 | + { |
| 341 | + "role": getattr(m, "type", m.__class__.__name__), |
| 342 | + "content": _stringify_content(getattr(m, "content", "")), |
| 343 | + } |
| 344 | + ) |
| 345 | + lines = [f"[{self._ts()}] [{node}] >> LLM input ({len(rendered)} msgs)"] |
| 346 | + for r in rendered: |
| 347 | + lines.append(f" {r['role']}: {r['content']}") |
| 348 | + self._write_text("\n".join(lines)) |
| 349 | + self._write_json({"event": "llm_start", "node": node, "messages": rendered}) |
| 350 | + |
| 351 | + @_guard |
| 352 | + def on_llm_end(self, response, *, tags=None, metadata=None, **kwargs): |
| 353 | + node = self._node(metadata, tags) |
| 354 | + try: |
| 355 | + gen = response.generations[0][0] |
| 356 | + except (AttributeError, IndexError): |
| 357 | + return |
| 358 | + message = getattr(gen, "message", None) |
| 359 | + content = _stringify_content( |
| 360 | + getattr(message, "content", getattr(gen, "text", "")) |
| 361 | + ) |
| 362 | + tool_calls = getattr(message, "tool_calls", None) or [] |
| 363 | + lines = [f"[{self._ts()}] [{node}] << LLM output"] |
| 364 | + if content: |
| 365 | + lines.append(f" {content}") |
| 366 | + for tc in tool_calls: |
| 367 | + args = json.dumps(tc.get("args", {}), default=str) |
| 368 | + lines.append(f" tool_call: {tc.get('name')}({args})") |
| 369 | + self._write_text("\n".join(lines)) |
| 370 | + self._write_json( |
| 371 | + { |
| 372 | + "event": "llm_end", |
| 373 | + "node": node, |
| 374 | + "content": content, |
| 375 | + "tool_calls": tool_calls, |
| 376 | + } |
| 377 | + ) |
| 378 | + |
| 379 | + @_guard |
| 380 | + def on_tool_start( |
| 381 | + self, serialized, input_str, *, tags=None, metadata=None, **kwargs |
| 382 | + ): |
| 383 | + name = (serialized or {}).get("name", "tool") |
| 384 | + node = self._node(metadata, tags) |
| 385 | + self._write_text(f"[{self._ts()}] [{node}] -> TOOL {name}({input_str})") |
| 386 | + self._write_json( |
| 387 | + {"event": "tool_start", "node": node, "tool": name, "input": input_str} |
| 388 | + ) |
| 389 | + |
| 390 | + @_guard |
| 391 | + def on_tool_end(self, output, *, tags=None, metadata=None, **kwargs): |
| 392 | + node = self._node(metadata, tags) |
| 393 | + text = _stringify_content(getattr(output, "content", output)) |
| 394 | + self._write_text(f"[{self._ts()}] [{node}] <- TOOL result: {text}") |
| 395 | + self._write_json({"event": "tool_end", "node": node, "output": text}) |
| 396 | + |
| 397 | + @_guard |
| 398 | + def on_chain_start(self, serialized, inputs, *, parent_run_id=None, **kwargs): |
| 399 | + # Only the root run (no parent) marks a task boundary; inner nodes/subgraphs |
| 400 | + # start constantly and would be noise. |
| 401 | + if parent_run_id is None: |
| 402 | + self._write_text(f"[{self._ts()}] ===== TASK START =====") |
| 403 | + self._write_json({"event": "task_start"}) |
| 404 | + |
| 405 | + @_guard |
| 406 | + def on_chain_end(self, outputs, *, parent_run_id=None, **kwargs): |
| 407 | + if parent_run_id is None: |
| 408 | + self._write_text(f"[{self._ts()}] ===== TASK COMPLETE =====") |
| 409 | + self._write_json({"event": "task_end"}) |
| 410 | + |
| 411 | + @_guard |
| 412 | + def on_llm_error(self, error, **kwargs): |
| 413 | + self._write_text(f"[{self._ts()}] !! LLM ERROR: {error}") |
| 414 | + self._write_json({"event": "llm_error", "error": str(error)}) |
| 415 | + |
| 416 | + @_guard |
| 417 | + def on_tool_error(self, error, **kwargs): |
| 418 | + self._write_text(f"[{self._ts()}] !! TOOL ERROR: {error}") |
| 419 | + self._write_json({"event": "tool_error", "error": str(error)}) |
0 commit comments