Skip to content

feat(energy): compose the starter agent's toolbelt inline in the notebook #164

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ethancjackson merged 1 commit into from
Jul 6, 2026
Merged

feat(energy): compose the starter agent's toolbelt inline in the notebook #164

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
67 changes: 22 additions & 45 deletions implementations/energy_oil_forecasting/99_starter_agent.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,19 +4,7 @@
"cell_type": "markdown",
"id": "cell-00",
"metadata": {},
"source": [
"# WTI Crude Oil — Your Starter Agent\n",
"\n",
"**If you're not sure what to do next, continue from here.**\n",
"\n",
"This notebook is a fresh, hackable agent for the WTI crude-oil use case — deliberately *not* wired into the numbered curriculum. It gives you our common building blocks behind simple toggles, so you can start building something of your own:\n",
"\n",
"- **optional news search** — bounded, cutoff-aware Google Search (proxy-only)\n",
"- **optional code execution** — an E2B Python sandbox\n",
"- **two lightweight skills** — *tool-usage playbooks* in `starter_agent/skills/`\n",
"\n",
"It does two things: lets you **talk to the agent** (open-ended, Track 2) and **score one real forecast** (Track 1). The live cells are gated by `RUN_AGENT` so a fresh `Run All` is safe and free; flip it to `True` to actually call the model.\n"
]
"source": "# WTI Crude Oil — Your Starter Agent\n\n**If you're not sure what to do next, continue from here.**\n\nThis notebook is a fresh, hackable agent for the WTI crude-oil use case — deliberately *not* wired into the numbered curriculum. An agent is a **persona** plus a **toolbelt**, and you assemble that toolbelt right here in the notebook from a menu of one-line tool factories:\n\n- **`news_search()`** — bounded, cutoff-aware Google Search (proxy-only)\n- **`arima_forecast()`** — an AutoARIMA statistical anchor the agent can call directly (no code-gen)\n- **`code_sandbox()`** — an E2B Python sandbox for the agent to compute its own diagnostics\n- each tool pulls in its own *playbook* skill from `starter_agent/skills/`\n\nThe factories live in `starter_agent/tools.py` — open it to see how a tool is built, or add your own. It does two things: lets you **talk to the agent** (open-ended, Track 2) and **score one real forecast** (Track 1). The live cells are gated by `RUN_AGENT` so a fresh `Run All` is safe and free; flip it to `True` to actually call the model.\n"
},
{
"cell_type": "code",
Expand Down Expand Up @@ -53,6 +41,7 @@
"from energy_oil_forecasting.starter_agent import (\n",
" build_starter_agent_config,\n",
" build_starter_agent_predictor,\n",
" tools,\n",
")\n",
"\n",
"\n",
Expand All @@ -63,12 +52,7 @@
"cell_type": "markdown",
"id": "cell-02",
"metadata": {},
"source": [
"---\n",
"## 1. Meet your agent\n",
"\n",
"`build_starter_agent_config` returns an `AgentConfig` with two toggles. The default turns **news search on** (proxy-only, no extra key) and **code execution off** (it needs `E2B_API_KEY` and is slower). Flip them and re-run — the loaded skills follow the enabled tools.\n"
]
"source": "---\n## 1. Build your agent's toolbelt\n\nThis is where you compose the agent. `build_starter_agent_config` takes a `tools=[...]` list — the toolbelt — and folds each tool onto the agent (its config, its skill, its instructions). **Comment a line to drop a tool; uncomment to add one**, then re-run. That's the whole model: an agent is a persona plus the tools you hand it.\n"
},
{
"cell_type": "code",
Expand All @@ -77,17 +61,24 @@
"metadata": {},
"outputs": [],
"source": [
"config = build_starter_agent_config(\n",
" model=AGENT_MODEL,\n",
" enable_search=True, # ← cutoff-aware Google Search (proxy-only)\n",
" enable_code_exec=False, # ← E2B Python sandbox (needs E2B_API_KEY); try True!\n",
")\n",
"\n",
"print(\"Agent:\", config.name)\n",
"print(\"Search enabled: \", config.context_retrieval.enabled)\n",
"print(\"Code-exec enabled: \", config.code_execution.enabled)\n",
"print(\"Skills loaded: \", [p.name for p in config.skills_dirs])\n",
"print(\"\\n── System instruction (edit this in starter_agent/agent.py) ──\\n\")\n",
"# ── Your agent's toolbelt ──────────────────────────────\n",
"# Each factory returns one tool. Comment a line to drop it, uncomment to add it.\n",
"# See starter_agent/tools.py for how each is built — and to write your own.\n",
"toolbelt = [\n",
" tools.news_search(), # cutoff-aware Google Search (proxy-only, no extra key)\n",
" tools.arima_forecast(), # AutoARIMA anchor — the agent calls a forecast directly, no code-gen\n",
" # tools.code_sandbox(), # E2B Python sandbox (needs E2B_API_KEY, slower) — uncomment to add\n",
"]\n",
"\n",
"config = build_starter_agent_config(model=AGENT_MODEL, tools=toolbelt)\n",
"\n",
"print(\"Agent: \", config.name)\n",
"print(\"Toolbelt:\", [t.label for t in toolbelt])\n",
"print(\" search enabled: \", config.context_retrieval.enabled)\n",
"print(\" forecast tool: \", bool(config.function_tools))\n",
"print(\" code-exec enabled:\", config.code_execution.enabled)\n",
"print(\"Skills loaded: \", [p.name for p in config.skills_dirs])\n",
"print(\"\\n── System instruction (edit the persona in starter_agent/agent.py) ──\\n\")\n",
"print(config.instruction[:1200], \"...\")"
]
},
Expand Down Expand Up @@ -191,21 +182,7 @@
"cell_type": "markdown",
"id": "cell-08",
"metadata": {},
"source": [
"---\n",
"## 4. Make it yours\n",
"\n",
"This agent is a starting point. Here are concrete next steps, easiest first — each is a small edit, then re-run the cells above.\n",
"\n",
"1. **Flip code execution on.** Set `enable_code_exec=True` in §1 (needs `E2B_API_KEY`). The agent loads the `code-analysis-playbook` skill and can compute its own diagnostics before forecasting. Compare the rationale.\n",
"2. **Edit the agent's personality.** Open `starter_agent/agent.py` and change `_build_starter_instruction()` — make it more cautious, more contrarian, focused on one driver. Re-run §1 to see the new instruction.\n",
"3. **Sharpen the skills.** The two files in `starter_agent/skills/` are short on purpose. Add your best queries to `research-playbook`, or a new diagnostic to `code-analysis-playbook`. The agent picks them up automatically.\n",
"4. **Change the question and the origin.** Try a different `QUESTION` in §2 and a different origin in §3.\n",
"5. **Add a tool.** Give the agent a conventional forecast tool as a statistical anchor — see `analyst_agent.build_wti_tool_config` for the `ForecastTool` pattern.\n",
"6. **Score it properly.** Run it across several origins with `backtest()` (see `04_systematic_backtest_eval.ipynb`) and compare CRPS against the baselines.\n",
"\n",
"Bigger ideas — an agent that *learns* a strategy (notebooks 05–06), news vs. no-news lift, live prospective forecasting — are in the use-case `README.md` and `planning-docs/roadmap.md`.\n"
]
"source": "---\n## 4. Make it yours\n\nThis agent is a starting point. Here are concrete next steps, easiest first — each is a small edit, then re-run the cells above.\n\n1. **Change the toolbelt.** In §1, uncomment `tools.code_sandbox()` (needs `E2B_API_KEY`) to let the agent compute its own diagnostics, or drop `arima_forecast()` and compare the rationale with and without a statistical anchor. Adding a tool automatically loads its playbook skill and its instructions.\n2. **Edit the agent's personality.** Open `starter_agent/agent.py` and change `_build_starter_instruction()` — make it more cautious, more contrarian, focused on one driver. Re-run §1 to see the new instruction.\n3. **Sharpen the skills.** The files in `starter_agent/skills/` are short on purpose. Add your best queries to `research-playbook`, or a new diagnostic to `code-analysis-playbook`. The agent picks them up automatically.\n4. **Change the question and the origin.** Try a different `QUESTION` in §2 and a different origin in §3.\n5. **Write your own tool.** Open `starter_agent/tools.py` and add a factory that returns a `ToolSpec` — point `arima_forecast()` at a different series, swap AutoARIMA for another predictor, or wrap a brand-new function tool. Then add it to the toolbelt in §1.\n6. **Score it properly.** Run it across several origins with `backtest()` (see `04_systematic_backtest_eval.ipynb`) and compare CRPS against the baselines.\n\nBigger ideas — an agent that *learns* a strategy (notebooks 05–06), news vs. no-news lift, live prospective forecasting — are in the use-case `README.md` and `planning-docs/roadmap.md`.\n"
}
],
"metadata": {
Expand Down
7 changes: 5 additions & 2 deletions implementations/energy_oil_forecasting/starter_agent/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,11 @@
"""WTI starter agent — a fresh, hackable template for your own exploration.

Exports the toggle-driven :class:`AgentConfig` factory and the predictor
convenience factory. See ``99_starter_agent.ipynb`` and ``agent.py``.
Exports the toolbelt-driven :class:`AgentConfig` factory, the predictor
convenience factory, and the :mod:`tools` module of per-tool factories you
compose in the notebook. See ``99_starter_agent.ipynb`` and ``agent.py``.
"""

from energy_oil_forecasting.starter_agent import tools
from energy_oil_forecasting.starter_agent.agent import (
build_starter_agent_config,
build_starter_agent_predictor,
Expand All @@ -13,4 +15,5 @@
__all__ = [
"build_starter_agent_config",
"build_starter_agent_predictor",
"tools",
]
117 changes: 56 additions & 61 deletions implementations/energy_oil_forecasting/starter_agent/agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@

import json
from pathlib import Path
from typing import Any, Callable
from typing import Any, Callable, Sequence

from aieng.forecasting.data.context import ForecastContext
from aieng.forecasting.evaluation.task import ForecastingTask
Expand All @@ -46,13 +46,13 @@
# Reuse the existing WTI prompt builder + history compression — these serialise
# the task/context into the agent's JSON payload and are not worth duplicating.
from energy_oil_forecasting.analyst_agent import WtiPriceForecastPromptBuilder
from energy_oil_forecasting.starter_agent.tools import ToolSpec, news_search


# Skills live next to this module.
# Skills live next to this module. The forecasting contract is always loaded;
# each tool loads its own playbook via its ToolSpec (see tools.py).
_SKILLS_ROOT = Path(__file__).parent / "skills"
_FORECASTING_SKILL = _SKILLS_ROOT / "forecasting"
_RESEARCH_SKILL = _SKILLS_ROOT / "research-playbook"
_CODE_ANALYSIS_SKILL = _SKILLS_ROOT / "code-analysis-playbook"


# ---------------------------------------------------------------------------
Expand Down Expand Up @@ -88,88 +88,82 @@ def _build_starter_instruction() -> str:
_STARTER_INSTRUCTION = _build_starter_instruction()


_CONTEXT_RETRIEVAL_INSTRUCTION = """\
You are an oil-market intelligence specialist with web search.

Return a concise structured markdown summary (3-5 paragraphs) covering, as the
query warrants: WTI/Brent price level and trend; OPEC+ supply decisions;
geopolitical risk in the Persian Gulf and key shipping lanes; US SPR / energy
policy; notable supply-disruption signals; and published analyst price targets.

Ground every claim in the search results you actually retrieve. When a cutoff
date is specified, never report or speculate about events after it.

Before finalizing your summary, reason step by step: (1) for each candidate \
fact, judge its actual recency from the substance of the result itself, \
never from a source's claimed publish date or byline timestamp — those are \
frequently stale or updated after original publication; (2) discard \
anything you cannot confidently place before the cutoff date; (3) only then \
write your summary. Do not supplement the search results with your own \
background/training knowledge — if the results are insufficient, say so \
explicitly rather than filling gaps from memory.\
"""


# ---------------------------------------------------------------------------
# Config factory
# ---------------------------------------------------------------------------


def build_starter_agent_config(
model: str = LITE_MODEL,
search_model: str = LITE_MODEL,
*,
enable_search: bool = True,
enable_code_exec: bool = False,
tools: Sequence[ToolSpec] = (),
) -> AgentConfig:
"""Build the WTI starter :class:`AgentConfig`.
"""Build the WTI starter :class:`AgentConfig` from a toolbelt.

An agent is a persona plus a list of tools. The persona is fixed here (edit
``_build_starter_instruction``); the *toolbelt* is what you compose in the
notebook — a list of :class:`~energy_oil_forecasting.starter_agent.tools.ToolSpec`
from the factories in :mod:`energy_oil_forecasting.starter_agent.tools`::

from energy_oil_forecasting.starter_agent import tools, build_starter_agent_config

config = build_starter_agent_config(
model=AGENT_MODEL,
tools=[tools.news_search(), tools.arima_forecast()],
)

Each spec lands in a different ``AgentConfig`` field; this function folds the
list, routing every fragment to the right place — search sub-agent, code
sandbox, function tools — and loading each tool's playbook skill and prompt
supplement. Adding or removing a tool is one line in the notebook.

Parameters
----------
model : str
Model for the analyst agent (default: lite). Pass the advanced model
(``"gemini-3.5-flash"``) for higher-quality runs.
search_model : str
Model for the bounded web-search sub-tool.
enable_search : bool, default=True
Wire a cutoff-aware ``search_web`` tool and load the
``research-playbook`` skill. Proxy-only — no extra API key.
enable_code_exec : bool, default=False
Wire an E2B Python sandbox and load the ``code-analysis-playbook``
skill. Needs ``E2B_API_KEY`` and is slower, so it is off by default —
flip it on to let the agent compute its own diagnostics.
tools : Sequence[ToolSpec], default=()
The agent's toolbelt. Build entries with the factories in ``tools.py``
(``news_search()``, ``code_sandbox()``, ``arima_forecast()``), or write
your own factory that returns a ``ToolSpec``.

Returns
-------
AgentConfig
"""
# Every attached skill is loaded on demand: ADK injects each skill's name +
# description into the system prompt, and the agent reads the full SKILL.md
# only when relevant — so toggling a tool just adds its skill, no persona edits.
# The forecasting contract is always loaded. Each tool's own playbook is
# loaded on demand: ADK injects each skill's name + description into the
# system prompt, and the agent reads the full SKILL.md only when relevant —
# so adding a tool just adds its skill, no persona edits.
skills_dirs: list[Path] = [_FORECASTING_SKILL]
if enable_search:
skills_dirs.append(_RESEARCH_SKILL)
if enable_code_exec:
skills_dirs.append(_CODE_ANALYSIS_SKILL)

context_retrieval = (
ContextRetrievalConfig(
enabled=True,
instruction=_CONTEXT_RETRIEVAL_INSTRUCTION,
search_model=search_model,
)
if enable_search
else ContextRetrievalConfig()
)
instruction = _STARTER_INSTRUCTION
context_retrieval = ContextRetrievalConfig()
code_execution = CodeExecutionConfig()
function_tools: list[Any] = []
max_output_tokens: int | None = None

for spec in tools:
if spec.skill_dir is not None:
skills_dirs.append(spec.skill_dir)
if spec.instruction_supplement:
instruction += spec.instruction_supplement
if spec.context_retrieval is not None:
context_retrieval = spec.context_retrieval
if spec.code_execution is not None:
code_execution = spec.code_execution
if spec.function_tool is not None:
function_tools.append(spec.function_tool)
if spec.max_output_tokens is not None:
max_output_tokens = max(max_output_tokens or 0, spec.max_output_tokens)

return AgentConfig(
name="wti_starter_agent",
model=model,
instruction=_STARTER_INSTRUCTION,
# 16k headroom: enough for a complete run_code script + structured output.
max_output_tokens=16_384 if enable_code_exec else None,
instruction=instruction,
max_output_tokens=max_output_tokens,
context_retrieval=context_retrieval,
code_execution=CodeExecutionConfig(enabled=enable_code_exec),
code_execution=code_execution,
function_tools=function_tools,
skills_dirs=skills_dirs,
)

Expand Down Expand Up @@ -232,5 +226,6 @@ def build_starter_agent_predictor(config: AgentConfig) -> AgentPredictor:
def __getattr__(name: str) -> Any:
"""Expose ``root_agent`` lazily for schema-free interactive use via ``adk web``."""
if name == "root_agent":
return build_adk_agent(build_starter_agent_config())
# Interactive default: news search on (proxy-only, no extra key).
return build_adk_agent(build_starter_agent_config(tools=[news_search()]))
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
Loading
Loading