MuXinCG
222 lines
7.0 KiB
Python
222 lines
7.0 KiB
Python
"""Utility functions for EvoScientist.
|
|
|
|
This module primarily contains helpers for displaying messages and prompts in
|
|
notebooks, and lightweight configuration loaders used by the agent runtime.
|
|
"""
|
|
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import yaml
|
|
from rich.console import Console
|
|
from rich.panel import Panel
|
|
from rich.text import Text
|
|
|
|
console = Console()
|
|
|
|
|
|
def format_message_content(message):
|
|
"""Convert message content to displayable string.
|
|
|
|
Args:
|
|
message: A LangChain message object with content attribute.
|
|
|
|
Returns:
|
|
Formatted string representation of the message content.
|
|
"""
|
|
parts = []
|
|
tool_calls_processed = False
|
|
|
|
# Handle main content
|
|
if isinstance(message.content, str):
|
|
parts.append(message.content)
|
|
elif isinstance(message.content, list):
|
|
# Handle complex content like tool calls (Anthropic format)
|
|
for item in message.content:
|
|
if item.get("type") == "text":
|
|
parts.append(item["text"])
|
|
elif item.get("type") == "tool_use":
|
|
parts.append(f"\n🔧 Tool Call: {item['name']}")
|
|
parts.append(f" Args: {json.dumps(item['input'], indent=2)}")
|
|
parts.append(f" ID: {item.get('id', 'N/A')}")
|
|
tool_calls_processed = True
|
|
else:
|
|
parts.append(str(message.content))
|
|
|
|
# Handle tool calls attached to the message (OpenAI format) - only if not already processed
|
|
if (
|
|
not tool_calls_processed
|
|
and hasattr(message, "tool_calls")
|
|
and message.tool_calls
|
|
):
|
|
for tool_call in message.tool_calls:
|
|
parts.append(f"\n🔧 Tool Call: {tool_call['name']}")
|
|
parts.append(f" Args: {json.dumps(tool_call['args'], indent=2)}")
|
|
parts.append(f" ID: {tool_call['id']}")
|
|
|
|
return "\n".join(parts)
|
|
|
|
|
|
def format_messages(messages):
|
|
"""Format and display a list of messages with Rich formatting.
|
|
|
|
Args:
|
|
messages: List of LangChain message objects to display.
|
|
"""
|
|
for m in messages:
|
|
msg_type = m.__class__.__name__.replace("Message", "")
|
|
content = format_message_content(m)
|
|
|
|
if msg_type == "Human":
|
|
console.print(Panel(content, title="🧑 Human", border_style="blue"))
|
|
elif msg_type == "Ai":
|
|
console.print(Panel(content, title="🤖 Assistant", border_style="green"))
|
|
elif msg_type == "Tool":
|
|
console.print(Panel(content, title="🔧 Tool Output", border_style="yellow"))
|
|
else:
|
|
console.print(Panel(content, title=f"📝 {msg_type}", border_style="white"))
|
|
|
|
|
|
def format_message(messages):
|
|
"""Alias for format_messages for backward compatibility."""
|
|
return format_messages(messages)
|
|
|
|
|
|
def show_prompt(prompt_text: str, title: str = "Prompt", border_style: str = "blue"):
|
|
"""Display a prompt with rich formatting and XML tag highlighting.
|
|
|
|
Args:
|
|
prompt_text: The prompt string to display
|
|
title: Title for the panel (default: "Prompt")
|
|
border_style: Border color style (default: "blue")
|
|
"""
|
|
# Create a formatted display of the prompt
|
|
formatted_text = Text(prompt_text)
|
|
formatted_text.highlight_regex(r"<[^>]+>", style="bold blue") # Highlight XML tags
|
|
formatted_text.highlight_regex(
|
|
r"##[^#\n]+", style="bold magenta"
|
|
) # Highlight headers
|
|
formatted_text.highlight_regex(
|
|
r"###[^#\n]+", style="bold cyan"
|
|
) # Highlight sub-headers
|
|
|
|
# Display in a panel for better presentation
|
|
console.print(
|
|
Panel(
|
|
formatted_text,
|
|
title=f"[bold green]{title}[/bold green]",
|
|
border_style=border_style,
|
|
padding=(1, 2),
|
|
)
|
|
)
|
|
|
|
|
|
def load_subagents(
|
|
config_path: Path,
|
|
*,
|
|
tool_registry: dict[str, Any],
|
|
prompt_refs: dict[str, str] | None = None,
|
|
) -> list[dict[str, Any]]:
|
|
"""Load subagent definitions from YAML and wire up tools.
|
|
|
|
NOTE: This is a custom utility. deepagents does not natively load subagents
|
|
from files - they're normally defined inline in the create_deep_agent() call.
|
|
We externalize to YAML here to keep configuration separate from code.
|
|
|
|
Supported YAML schemas:
|
|
|
|
1) Mapping style (recommended):
|
|
planner-agent:
|
|
description: "..."
|
|
tools: [think_tool]
|
|
system_prompt: |
|
|
...
|
|
research-agent:
|
|
description: "..."
|
|
tools: [tavily_search, think_tool]
|
|
system_prompt_ref: RESEARCHER_INSTRUCTIONS
|
|
|
|
2) List style (legacy):
|
|
subagents:
|
|
- name: planner-agent
|
|
description: "..."
|
|
tools: [think_tool]
|
|
system_prompt: |
|
|
...
|
|
"""
|
|
prompt_refs = prompt_refs or {}
|
|
|
|
with config_path.open(encoding="utf-8") as f:
|
|
config = yaml.safe_load(f) or {}
|
|
|
|
if not isinstance(config, dict) or not config:
|
|
raise ValueError("subagent.yaml must be a mapping or contain 'subagents:'")
|
|
|
|
subagents: list[dict[str, Any]] = []
|
|
|
|
def _build_one(name: str, spec: dict[str, Any]) -> dict[str, Any]:
|
|
subagent: dict[str, Any] = {
|
|
"name": name,
|
|
"description": spec.get("description", ""),
|
|
}
|
|
|
|
if "system_prompt_ref" in spec:
|
|
ref = spec["system_prompt_ref"]
|
|
if ref not in prompt_refs:
|
|
raise ValueError(f"Unknown system_prompt_ref '{ref}' for subagent '{name}'")
|
|
subagent["system_prompt"] = prompt_refs[ref]
|
|
else:
|
|
subagent["system_prompt"] = spec.get("system_prompt", "")
|
|
|
|
if "model" in spec:
|
|
subagent["model"] = spec["model"]
|
|
|
|
if "skills" in spec:
|
|
subagent["skills"] = spec["skills"]
|
|
|
|
if "tools" in spec:
|
|
subagent["tools"] = [tool_registry[t] for t in spec["tools"]]
|
|
|
|
return subagent
|
|
|
|
# Legacy list style
|
|
if "subagents" in config:
|
|
items = config.get("subagents")
|
|
if not isinstance(items, list) or not items:
|
|
raise ValueError("subagent.yaml must contain a non-empty 'subagents:' list")
|
|
for item in items:
|
|
if not isinstance(item, dict):
|
|
continue
|
|
name = item.get("name")
|
|
if not name:
|
|
raise ValueError("Each subagent entry must have a 'name'")
|
|
subagents.append(_build_one(name, item))
|
|
return subagents
|
|
|
|
# Mapping style: {<name>: <spec>}
|
|
for name, spec in config.items():
|
|
if not isinstance(spec, dict):
|
|
continue
|
|
subagents.append(_build_one(name, spec))
|
|
|
|
return subagents
|
|
|
|
|
|
def load_subagent(
|
|
config_path: Path,
|
|
name: str,
|
|
*,
|
|
tool_registry: dict[str, Any],
|
|
prompt_refs: dict[str, str] | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Load a single sub-agent by name from YAML."""
|
|
for agent in load_subagents(
|
|
config_path,
|
|
tool_registry=tool_registry,
|
|
prompt_refs=prompt_refs,
|
|
):
|
|
if agent.get("name") == name:
|
|
return agent
|
|
raise KeyError(f"Sub-agent not found: {name}")
|