RA.Aid/ra_aid/console/output.py

43 lines
1.6 KiB
Python

from typing import Any, Dict
from langchain_core.messages import AIMessage
from rich.markdown import Markdown
from rich.panel import Panel
# Import shared console instance
from .formatting import console
def print_agent_output(chunk: Dict[str, Any]) -> None:
"""Print only the agent's message content, not tool calls.
Args:
chunk: A dictionary containing agent or tool messages
"""
if "agent" in chunk and "messages" in chunk["agent"]:
messages = chunk["agent"]["messages"]
for msg in messages:
if isinstance(msg, AIMessage):
# Handle text content
if isinstance(msg.content, list):
for content in msg.content:
if content["type"] == "text" and content["text"].strip():
console.print(
Panel(Markdown(content["text"]), title="🤖 Assistant")
)
else:
if msg.content.strip():
console.print(
Panel(Markdown(msg.content.strip()), title="🤖 Assistant")
)
elif "tools" in chunk and "messages" in chunk["tools"]:
for msg in chunk["tools"]["messages"]:
if msg.status == "error" and msg.content:
console.print(
Panel(
Markdown(msg.content.strip()),
title="❌ Tool Error",
border_style="red bold",
)
)