ani-sdhu/essp

1#!/usr/bin/env python32"""Offline test for ai_command: a mocked LLM produces plans, and the real3resolution pipeline runs against mocked EIA + FRED. No network or keys."""4import json5 6import eia_search as es7import energy_data_dashboard as dash8import ai_command as ai9 10# --- Mock EIA (root -> natural-gas -> cons -> sum, with series facet) --------11EIA_NODES = {12    "": {"id": "", "name": "EIA", "routes": [{"id": "natural-gas", "name": "Natural Gas"}]},13    "natural-gas": {"id": "natural-gas", "name": "Natural Gas",14                    "routes": [{"id": "cons", "name": "Consumption"}]},15    "natural-gas/cons": {"id": "cons", "name": "Consumption",16                         "routes": [{"id": "sum", "name": "Consumption Summary"}]},17    "natural-gas/cons/sum": {18        "id": "sum", "name": "Natural Gas Consumption Summary",19        "description": "Monthly natural gas consumed by sector.",20        "frequency": [{"id": "monthly", "description": "Monthly"}],21        "facets": [{"id": "series", "description": "Series"}],22        "data": {"value": {"units": "MMcf"}},23        "startPeriod": "1973-01", "endPeriod": "2024-12",24        "defaultFrequency": "monthly"},25}26EIA_SERIES = {"id": "series", "facets": [27    {"id": "N3010US3", "name": "Residential Consumers, U.S."},28    {"id": "N3020US3", "name": "Commercial Consumers, U.S."},29]}30 31 32def eia_transport(url, params):33    pd_ = {}34    for k, v in params:35        pd_.setdefault(k, []).append(v)36    if url.endswith("/facet/series"):37        return {"response": EIA_SERIES}38    if url.endswith("/data"):39        sid = pd_.get("facets[series][]", ["N3010US3"])[0]40        desc = {"N3010US3": "Natural Gas Deliveries to Residential Consumers, U.S.",41                "N3020US3": "Natural Gas Deliveries to Commercial Consumers, U.S."}42        rows = [{"period": f"2024-{m:02d}", "series": sid,43                 "series-description": desc.get(sid, sid),44                 "value": str(100 + m), "value-units": "MMcf"} for m in range(1, 13)]45        return {"response": {"total": "12", "data": rows}}46    return {"response": EIA_NODES[url.replace(es.API_BASE, "").strip("/")]}47 48 49def fred_transport(endpoint, params):50    if endpoint == "series/search":51        return {"seriess": [52            {"id": "DHHNGSP", "title": "Henry Hub Natural Gas Spot Price",53             "frequency": "Daily", "units": "Dollars per Million BTU",54             "observation_start": "1997-01-07", "observation_end": "2025-12-01",55             "popularity": 60}]}56    if endpoint == "series/observations":57        return {"observations": [58            {"date": "2000-01-01", "value": "2.4"},59            {"date": "2000-02-01", "value": "2.6"}]}60    raise AssertionError(endpoint)61 62 63# --- Mock LLM transport: returns a plan based on the request text ------------64SEEN_CONTEXT = {}65 66 67def llm_transport(payload):68    user_msg = payload["messages"][-1]["content"]69    SEEN_CONTEXT["last"] = user_msg70    low = user_msg.lower()71    if "clear" in low:72        plan = {"action": "clear", "items": [], "start": "", "end": ""}73    elif "add" in low:74        plan = {"action": "add", "items": [75            {"source": "FRED", "query": "henry hub", "series_hint": "",76             "exact_id": "DHHNGSP"}], "start": "", "end": ""}77    elif "compare" in low:78        plan = {"action": "compare", "items": [79            {"source": "FRED", "query": "henry hub natural gas spot price",80             "series_hint": "", "exact_id": "DHHNGSP"},81            {"source": "EIA", "query": "natural gas consumption",82             "series_hint": "residential", "exact_id": ""}],83            "start": "", "end": ""}84    else:85        plan = {"action": "plot", "items": [86            {"source": "EIA", "query": "natural gas consumption",87             "series_hint": "residential", "exact_id": ""}],88            "start": "", "end": ""}89    return {"choices": [{"message": {"content": json.dumps(plan)}}]}90 91 92def run():93    fails = []94    def check(name, cond):95        print(("PASS" if cond else "FAIL"), "-", name)96        if not cond:97            fails.append(name)98 99    eia_client = es.EIAClient("K", transport=eia_transport)100    eia_index = es.build_index(eia_client, log=lambda *_: None)101    fred = dash.FREDClient("K", transport=fred_transport)102    hub = dash.DataHub(eia_client, fred, eia_index=eia_index)103 104    llm = ai.LLMClient(api_key="dummy", transport=llm_transport)105    check("llm reports available", llm.available is True)106 107    # Plan parsing (incl. fenced JSON robustness)108    p = llm.plan("compare a and b", {"plotted": []})109    check("plan parses to dict", isinstance(p, dict) and p["action"] == "compare")110    check("fenced json parses",111          ai.LLMClient._parse_json('```json\n{"action":"clear","items":[]}\n```')["action"] == "clear")112 113    # Capture plot calls114    calls = []115    def on_plot(frames, replace):116        calls.append((list(frames), replace))117    ctx = {"plotted": ["EIA:N3010US3"], "recent_results": ["FRED:DHHNGSP"]}118    bar = ai.CommandBar(hub, llm, on_plot=on_plot, get_context=lambda: ctx)119 120    # 1. Compare -> replace with two resolved frames (one FRED, one EIA)121    msg = bar.execute("compare henry hub price and US residential gas consumption")122    check("compare produced a plot call", len(calls) == 1)123    frames, replace = calls[-1]124    check("compare replaces chart", replace is True)125    check("compare resolved two series", len(frames) == 2)126    labels = {f["label"].iloc[0] for f in frames}127    ids = {f["series"].iloc[0] for f in frames}128    check("FRED series id resolved", "DHHNGSP" in ids)129    check("EIA residential id resolved", "N3010US3" in ids)130    check("FRED label is plain language",131          "Henry Hub Natural Gas Spot Price" in labels)132    check("EIA label is plain language",133          any("Residential Consumers" in l for l in labels))134    check("status uses plain language",135          "Plotted" in msg and "Henry Hub" in msg)136 137    # Context actually reached the model138    check("context passed to LLM", "N3010US3" in SEEN_CONTEXT.get("last", ""))139 140    # 2. Add -> append (replace False), one frame141    bar.execute("add henry hub")142    frames2, replace2 = calls[-1]143    check("add appends (replace False)", replace2 is False)144    check("add resolved one series", len(frames2) == 1)145 146    # 3. Clear -> replace with empty147    cmsg = bar.execute("clear it")148    frames3, replace3 = calls[-1]149    check("clear sends empty replace", replace3 is True and frames3 == [])150    check("clear status", "Cleared" in cmsg)151 152    # 4. Plain plot -> one EIA series, replace153    bar.execute("show me residential natural gas use")154    frames4, replace4 = calls[-1]155    check("plain plot replaces with one frame",156          replace4 is True and len(frames4) == 1)157 158    # 5. Disabled client (no key, default transport) yields inactive bar159    bad = ai.LLMClient(api_key="")160    check("no-key client unavailable", bad.available is False)161    bar2 = ai.CommandBar(hub, bad, on_plot=on_plot)162    panel = bar2.panel()163    check("command bar builds a Panel object", panel is not None)164 165    print()166    if fails:167        print(f"{len(fails)} FAILED: {fails}"); return 1168    print("ALL TESTS PASSED"); return 0169 170 171if __name__ == "__main__":172    raise SystemExit(run())173