from dotenv import load_dotenv
from utils.agent_visualizer import (
display_agent_response,
print_activity,
reset_activity_context,
visualize_conversation,
)
from haijun_agent_sdk import HaijunAgentOptions, HaijunSDKClient
load_dotenv()
Define the model to use throughout this notebook
Using Opus 4.6 for its superior planning and reasoning capabilities
MODEL = "haijun-opus-4-6"
print(f"š Notebook configured to use: {MODEL}")
ī
š Notebook configured to use: haijun-opus-4-6 01 - The Chief of Staff Agent Introduction In notebook 00, we built a simple research agent. In this notebook, we'll incrementally introduce key Haijun Code SDK features for building comprehensive agents. For each introduced feature, we'll explain:
multiple sources
Provide executive summaries
with actionable recommendations
Basic Features
Understanding Agent Data Source Preferences
What Just Happened: By adding to our prompt, we guided the agent to rely on the HAIJUN.md context rather than seeking more granular data from CSV files.
financial_forecast.py
: Models ARR growth scenarios (base/optimistic/pessimistic) given the current
2.4
M
A
R
R
g
r
o
w
i
n
g
a
t
15
2.4M ARR growing at 15% MoM.Critical for financial-analyst to project Series B readiness and validate the
2.4
M
A
R
R
g
r
o
w
in
g
a
t
15
30M fundraising target.
decision_matrix.py
: Creates weighted decision matrices for strategic choices like the SmartDev acquisition or office expansion. Helps chief of staff systematically evaluate complex decisions with multiple stakeholders and criteria.
messages = []
async with HaijunSDKClient(
options=HaijunAgentOptions(
model=MODEL,
allowed_tools=["Bash", "Read"],
cwd="chief_of_staff_agent", # Points to subdirectory where our agent is defined
)
) as agent:
await agent.query(
"Use your simple calculation script with a total runway of 2904829 and a monthly burn of 121938."
)
async for msg in agent.receive_response():
print_activity(msg)
messages.append(msg)
Display the response with HTML rendering
display_agent_response(messages)
ī
š¤ Using: Glob() š¤ Using: Glob() š¤ Using: Glob() ā Tool completed ā Tool completed ā Tool completed š¤ Thinking... š¤ Using: Read() ā Tool completed š¤ Thinking... š¤ Using: Bash() ā Tool completed š¤ Thinking...
p" style="padding-left:8ch;text-indent:-8ch"> output_dir: Directory to save plan files
title: Title for the plan document
Returns:
Path to the saved plan file
"""
plans_dir = Path(output_dir)
plans_dir.mkdir(exist_ok=True)
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
plan_file = plans_dir / f"plan_{timestamp}.md"
with open(plan_file, "w") as f:
f.write(f"# {title}\n\n")
f.write(f"Created: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
f.write(f"Prompt: {prompt_summary}\n")
f.write(f"Model: {model_name}\n")
f.write(f"Plan Source: {plan_source}\n\n")
f.write("---\n\n")
f.write(plan_content)
f.write("\n\n---\n\n")
f.write("This plan was generated in plan mode and has not been executed.\n")
return plan_file
def capture_message_content(
msg: Any,
plan_content: list[str],
write_tool_content: list[str],
write_tool_paths: list[str],
) -> None:
"""
Process a streaming message and capture relevant plan content.
This function extracts content from three potential sources:
- Text blocks in message content
- Write tool call parameters
- Final result attribute
Args:
msg: The message object from the agent stream
plan_content: List to append text content to
write_tool_content: List to append Write tool content to
write_tool_paths: List to append Write tool file paths to
"""
Source 1: Text blocks from message content
if hasattr(msg, "content"):
for block in msg.content:
if hasattr(block, "text"):
plan_content.append(block.text)
Source 2: Write tool calls
if hasattr(block, "type") and block.type == "tool_use":
if hasattr(block, "name") and block.name == "Write":
if hasattr(block, "input") and isinstance(block.input, dict):
if "content" in block.input:
write_tool_content.append(block.input["content"])
if "file_path" in block.input:
write_tool_paths.append(block.input["file_path"])
Source 3: Final result
if hasattr(msg, "result") and msg.result:
plan_content.append(msg.result)
print("ā Plan Mode helper functions loaded")
ī
ā
Plan Mode helper functions loaded # ============================================================================= # Plan Mode Configuration # ============================================================================= # Note: MODEL is defined in cell-0 as "haijun-opus-4-6" # Opus excels at complex planning tasks # The prompt is carefully crafted to: # 1. Provide explicit context (since Opus prefers explicit information) # 2. Request XML-tagged output for reliable extraction # 3. Prevent file-writing so we can capture the plan programmatically PLAN_PROMPT = """Restructure our engineering team for AI focus. CONTEXT (from HAIJUN.md): You are the Chief of Staff for TechStart Inc, a 50-person B2B SaaS startup that raised $10M Series A. - Current engineering team: 25 people (Backend: 12, Frontend: 8, DevOps: 5) - Monthly burn rate: ~$500K, Runway: 20 months - Senior Engineer compensation: $180K-$220K + equity CRITICAL OUTPUT INSTRUCTIONS: 1. DO NOT use the Write tool - Output your plan directly in your response text 2. DO NOT save to any files - I will handle saving the plan myself 3. Wrap your ENTIRE plan inside