From 058978e82fb339f71f06e3ba6b64446940b7de53 Mon Sep 17 00:00:00 2001 From: Ethan Jackson Date: Mon, 6 Jul 2026 09:55:12 -0400 Subject: [PATCH] feat(energy): compose the starter agent's toolbelt inline in the notebook MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the starter agent's boolean toggles (enable_search / enable_code_exec) with an explicit tools=[...] toolbelt, so the notebook shows how an agent is assembled — a persona plus a list of tools — instead of hiding it behind flags. - New starter_agent/tools.py: a ToolSpec descriptor plus one factory per tool (news_search, code_sandbox, arima_forecast). Each spec carries the config fragment it fills, its playbook skill, prompt supplement, and token floor. - build_starter_agent_config now folds a Sequence[ToolSpec] onto AgentConfig, routing each fragment to the right field. The plumbing stays in the module; the notebook composition is comment-a-line-to-toggle. - Wires the previously demo-only ForecastTool (AutoARIMA run_forecast) into the starter agent via arima_forecast() — the agent can now call a statistical forecast directly, no code-gen. arima_forecast() accepts a data_service and series_id/frequency so it generalizes beyond WTI. - 99_starter_agent.ipynb §1 and §4 updated to the toolbelt model; add-a-tool step now points at tools.py instead of the analyst_agent reference. - test_starter_agent.py pins the fold: each factory routes to exactly one AgentConfig field, and composition/order are order-independent. Co-Authored-By: Claude Opus 4.8 --- .../99_starter_agent.ipynb | 67 ++--- .../starter_agent/__init__.py | 7 +- .../starter_agent/agent.py | 117 +++++---- .../starter_agent/tools.py | 231 ++++++++++++++++++ .../test_starter_agent.py | 111 +++++++++ 5 files changed, 425 insertions(+), 108 deletions(-) create mode 100644 implementations/energy_oil_forecasting/starter_agent/tools.py create mode 100644 implementations/tests/energy_oil_forecasting/test_starter_agent.py diff --git a/implementations/energy_oil_forecasting/99_starter_agent.ipynb b/implementations/energy_oil_forecasting/99_starter_agent.ipynb index 20b4c2b7..de7e1bdc 100644 --- a/implementations/energy_oil_forecasting/99_starter_agent.ipynb +++ b/implementations/energy_oil_forecasting/99_starter_agent.ipynb @@ -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", @@ -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", @@ -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", @@ -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], \"...\")" ] }, @@ -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": { diff --git a/implementations/energy_oil_forecasting/starter_agent/__init__.py b/implementations/energy_oil_forecasting/starter_agent/__init__.py index f70dae73..1e6cb962 100644 --- a/implementations/energy_oil_forecasting/starter_agent/__init__.py +++ b/implementations/energy_oil_forecasting/starter_agent/__init__.py @@ -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, @@ -13,4 +15,5 @@ __all__ = [ "build_starter_agent_config", "build_starter_agent_predictor", + "tools", ] diff --git a/implementations/energy_oil_forecasting/starter_agent/agent.py b/implementations/energy_oil_forecasting/starter_agent/agent.py index b7d1027e..2f554617 100644 --- a/implementations/energy_oil_forecasting/starter_agent/agent.py +++ b/implementations/energy_oil_forecasting/starter_agent/agent.py @@ -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 @@ -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" # --------------------------------------------------------------------------- @@ -88,28 +88,6 @@ 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 # --------------------------------------------------------------------------- @@ -117,59 +95,75 @@ def _build_starter_instruction() -> str: 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, ) @@ -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}") diff --git a/implementations/energy_oil_forecasting/starter_agent/tools.py b/implementations/energy_oil_forecasting/starter_agent/tools.py new file mode 100644 index 00000000..de93b77c --- /dev/null +++ b/implementations/energy_oil_forecasting/starter_agent/tools.py @@ -0,0 +1,231 @@ +"""The starter agent's toolbelt — one factory per tool, composed in the notebook. + +An agent is a *persona* plus a *list of tools*. This module makes that list the +thing you edit: each function here returns a :class:`ToolSpec` describing one +capability, and you assemble an agent by handing a list of them to +:func:`~energy_oil_forecasting.starter_agent.agent.build_starter_agent_config`:: + + from energy_oil_forecasting.starter_agent import tools, build_starter_agent_config + + config = build_starter_agent_config( + model=AGENT_MODEL, + tools=[ + tools.news_search(), # cutoff-aware Google Search (proxy-only) + tools.arima_forecast(), # AutoARIMA statistical anchor — no code-gen + # tools.code_sandbox(), # E2B Python sandbox (needs E2B_API_KEY) + ], + ) + +Each tool lands in a *different* field of the underlying ``AgentConfig`` (search +is a sub-agent, code execution is a sandbox capability, the forecast is a +plain function tool). A :class:`ToolSpec` carries everything a tool needs — the +config fragment it fills, its playbook skill, and any prompt supplement — so the +config factory can route it without the notebook ever touching ADK plumbing. + +To add your own tool, write a factory that returns a ``ToolSpec``. Point it at a +different series, swap AutoARIMA for another predictor, or wrap a brand-new +function tool — the notebook composition and the config fold both keep working. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from aieng.forecasting.data import DataService +from aieng.forecasting.methods.agentic import ForecastTool +from aieng.forecasting.methods.agentic.agent_factory import ( + CodeExecutionConfig, + ContextRetrievalConfig, +) +from aieng.forecasting.methods.numerical.darts_arima import DartsAutoARIMAPredictor +from aieng.forecasting.models import LITE_MODEL +from energy_oil_forecasting.data import WTI_SERIES_ID, build_wti_service + + +# Skills live next to this module; each tool loads its own playbook. +_SKILLS_ROOT = Path(__file__).parent / "skills" +_RESEARCH_SKILL = _SKILLS_ROOT / "research-playbook" +_CODE_ANALYSIS_SKILL = _SKILLS_ROOT / "code-analysis-playbook" + + +# --------------------------------------------------------------------------- +# ToolSpec — the seam between the notebook's toolbelt and AgentConfig +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class ToolSpec: + """One item on the agent's toolbelt. + + A tool is more than a callable: it may need a config fragment, a skill that + teaches the agent how to use it, and a line of instruction. This descriptor + bundles all of that so + :func:`~energy_oil_forecasting.starter_agent.agent.build_starter_agent_config` + can fold a list of specs onto a single ``AgentConfig`` — routing each field + to the right place — without the caller knowing the internals. + + Attributes + ---------- + label : str + Short human-readable name, shown when the config is printed. + context_retrieval : ContextRetrievalConfig or None + Web-search sub-agent config, if this tool provides search. + code_execution : CodeExecutionConfig or None + E2B sandbox config, if this tool provides code execution. + function_tool : Any or None + A ready-to-register ADK function tool (e.g. from + ``ForecastTool.as_function_tool()``). + skill_dir : Path or None + A playbook skill directory to load alongside the tool. + instruction_supplement : str + Text appended to the agent's system instruction when this tool is on. + max_output_tokens : int or None + A per-tool floor on the response budget (e.g. code execution needs + headroom for a full script). The config takes the max across all tools. + """ + + label: str + context_retrieval: ContextRetrievalConfig | None = None + code_execution: CodeExecutionConfig | None = None + function_tool: Any | None = None + skill_dir: Path | None = None + instruction_supplement: str = "" + max_output_tokens: int | None = None + + +# --------------------------------------------------------------------------- +# Tool-specific prompt text +# --------------------------------------------------------------------------- + + +_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.\ +""" + + +def _forecast_tool_supplement(series_id: str, frequency: str) -> str: + """Instruction appended when the statistical forecast tool is attached. + + Kept in the factory (not hard-coded) so a tool built for a different series + or frequency describes itself correctly to the agent. + """ + return f""" + +## Statistical forecast tool + +You have access to `run_forecast`, a conventional statistical baseline +(AutoARIMA) you can call directly. Unlike open-ended code, this tool has a fixed, +auditable interface and returns a structured forecast you can reason from. + +Call it ONCE before producing your forecast, with: +- `series_id`: "{series_id}" +- `cutoff_date`: the `as_of` date from the payload (YYYY-MM-DD). This is the + information cutoff — the model uses only data on or before it. +- `horizons`: the `horizons` list from the payload. +- `frequency`: "{frequency}" (the business calendar the series trades on). + +The tool returns JSON with point forecasts and 80%/90% prediction intervals per +horizon. Treat it as a disciplined statistical anchor: combine it with any +market context you have. You may adjust away from the baseline when fundamentals +or geopolitical risk justify it — document your reasoning in the `rationale` +fields.\ +""" + + +# --------------------------------------------------------------------------- +# Tool factories — each returns one ToolSpec +# --------------------------------------------------------------------------- + + +def news_search(*, search_model: str = LITE_MODEL) -> ToolSpec: + """Build a cutoff-aware Google Search tool, run by a bounded sub-agent (proxy-only). + + Wires a ``search_web`` tool and loads the ``research-playbook`` skill. No + extra API key — everything routes through the Vector proxy. + + Parameters + ---------- + search_model : str + Model for the web-search sub-agent. Defaults to the lite model. + """ + return ToolSpec( + label="news_search", + context_retrieval=ContextRetrievalConfig( + enabled=True, + instruction=_CONTEXT_RETRIEVAL_INSTRUCTION, + search_model=search_model, + ), + skill_dir=_RESEARCH_SKILL, + ) + + +def code_sandbox() -> ToolSpec: + """Build an E2B Python sandbox for the agent to compute its own diagnostics. + + Wires the code-execution capability and loads the ``code-analysis-playbook`` + skill. Needs ``E2B_API_KEY`` and is slower than the other tools, so it is + off by default — add it to the toolbelt to turn it on. + """ + return ToolSpec( + label="code_sandbox", + code_execution=CodeExecutionConfig(enabled=True), + skill_dir=_CODE_ANALYSIS_SKILL, + # 16k headroom: enough for a complete run_code script + structured output. + max_output_tokens=16_384, + ) + + +def arima_forecast( + *, + series_id: str = WTI_SERIES_ID, + frequency: str = "B", + num_samples: int = 200, + data_service: DataService | None = None, +) -> ToolSpec: + """Build a conventional statistical anchor: AutoARIMA behind a `run_forecast` tool. + + Lets the agent invoke a statistical forecast *directly* — a rigid, auditable + interface — instead of writing forecasting code. In contrast to + :func:`code_sandbox` (open-ended), this trades flexibility for control and + reproducibility. The tool reads series data server-side; it never enters the + LLM context. + + Parameters + ---------- + series_id : str + Series the agent should forecast. Defaults to the WTI target. + frequency : str + Business calendar passed to the predictor (``"B"`` for WTI). + num_samples : int + Monte Carlo sample count for AutoARIMA. Kept modest to bound latency. + data_service : DataService or None + Pre-populated data service. When ``None``, a cache-backed WTI service is + built. Pass one to point the tool at your own series (or to avoid a data + fetch in tests). + """ + service = data_service if data_service is not None else build_wti_service() + tool = ForecastTool(service, predictor=DartsAutoARIMAPredictor(num_samples=num_samples)) + return ToolSpec( + label="arima_forecast", + function_tool=tool.as_function_tool(), + instruction_supplement=_forecast_tool_supplement(series_id, frequency), + ) diff --git a/implementations/tests/energy_oil_forecasting/test_starter_agent.py b/implementations/tests/energy_oil_forecasting/test_starter_agent.py new file mode 100644 index 00000000..6963eb94 --- /dev/null +++ b/implementations/tests/energy_oil_forecasting/test_starter_agent.py @@ -0,0 +1,111 @@ +"""Tests for the starter-agent toolbelt fold. + +Each tool factory returns a :class:`ToolSpec` that lands in a *different* +``AgentConfig`` field (search sub-agent, code sandbox, function tool). These +tests pin that routing: composing a toolbelt must populate exactly the right +fields, load the right skills, and append the right instruction supplements — +so a refactor can't silently drop a tool onto the wrong field. + +Construction is offline: ``arima_forecast`` is handed an empty ``DataService`` +so no series data is fetched (data is only read when the agent calls the tool). +""" + +from __future__ import annotations + +from aieng.forecasting.data import DataService +from energy_oil_forecasting.starter_agent import build_starter_agent_config, tools + + +def _empty_arima() -> object: + """Build an ``arima_forecast`` spec whose tool reads from an empty service (no fetch).""" + return tools.arima_forecast(data_service=DataService()) + + +# --------------------------------------------------------------------------- +# Empty toolbelt — the bare persona +# --------------------------------------------------------------------------- + + +def test_empty_toolbelt_is_bare_persona() -> None: + """No tools: only the forecasting skill, no capabilities, unmodified instruction.""" + config = build_starter_agent_config(tools=[]) + + assert config.name == "wti_starter_agent" + assert not config.context_retrieval.enabled + assert not config.code_execution.enabled + assert list(config.function_tools) == [] + assert config.max_output_tokens is None + assert [p.name for p in config.skills_dirs] == ["forecasting"] + + +# --------------------------------------------------------------------------- +# Each factory routes to exactly one AgentConfig field +# --------------------------------------------------------------------------- + + +def test_news_search_wires_context_retrieval_and_skill() -> None: + config = build_starter_agent_config(tools=[tools.news_search()]) + + assert config.context_retrieval.enabled + assert config.context_retrieval.instruction.strip() + assert "research-playbook" in [p.name for p in config.skills_dirs] + # News search touches nothing else. + assert not config.code_execution.enabled + assert list(config.function_tools) == [] + + +def test_arima_forecast_wires_function_tool_and_supplement() -> None: + base = build_starter_agent_config(tools=[]) + config = build_starter_agent_config(tools=[_empty_arima()]) + + assert len(config.function_tools) == 1 + # Appends its instruction supplement to the persona; adds no skill of its own. + assert len(config.instruction) > len(base.instruction) + assert "run_forecast" in config.instruction + assert [p.name for p in config.skills_dirs] == ["forecasting"] + assert not config.context_retrieval.enabled + assert not config.code_execution.enabled + + +def test_code_sandbox_wires_code_execution_skill_and_token_budget() -> None: + config = build_starter_agent_config(tools=[tools.code_sandbox()]) + + assert config.code_execution.enabled + assert "code-analysis-playbook" in [p.name for p in config.skills_dirs] + # Code execution needs response headroom for a full script. + assert config.max_output_tokens == 16_384 + assert not config.context_retrieval.enabled + assert list(config.function_tools) == [] + + +# --------------------------------------------------------------------------- +# Composition — folding several tools onto one config +# --------------------------------------------------------------------------- + + +def test_full_toolbelt_composes_all_fields() -> None: + """Search + forecast + sandbox each land in their own field simultaneously.""" + config = build_starter_agent_config( + tools=[tools.news_search(), _empty_arima(), tools.code_sandbox()], + ) + + assert config.context_retrieval.enabled + assert config.code_execution.enabled + assert len(config.function_tools) == 1 + assert "run_forecast" in config.instruction + assert config.max_output_tokens == 16_384 + assert [p.name for p in config.skills_dirs] == [ + "forecasting", + "research-playbook", + "code-analysis-playbook", + ] + + +def test_toolbelt_order_is_independent_of_field_routing() -> None: + """Reordering the toolbelt doesn't change which field each tool fills.""" + forward = build_starter_agent_config(tools=[tools.news_search(), _empty_arima()]) + reverse = build_starter_agent_config(tools=[_empty_arima(), tools.news_search()]) + + assert forward.context_retrieval.enabled == reverse.context_retrieval.enabled + assert len(forward.function_tools) == len(reverse.function_tools) == 1 + assert {p.name for p in forward.skills_dirs} == {p.name for p in reverse.skills_dirs}