ani-sdhu/essp
0
1#!/usr/bin/env python32"""3ai_command.py4=============5 6A natural-language command bar for the energy/economic dashboard. You type7something like "compare henry hub price and US residential gas consumption" or8"add CPI" or "clear the chart", and it gets resolved to real series and9plotted.10 11How it stays reliable12---------------------13The language model does only the language understanding: it turns your sentence14into a small JSON plan (which items to plot, from which source, whether to15replace or add). It does NOT invent series IDs. The plan is then resolved by16the dashboard's existing deterministic search and fetch layer (FRED live17search and the EIA catalog index), so what actually gets plotted is always a18real series that the providers returned.19 20Provider21--------22Defaults to Groq, which exposes an OpenAI-compatible endpoint at23https://api.groq.com/openai/v1/chat/completions and a strict JSON mode via24response_format (Groq, 2026). Because it is OpenAI-compatible, you can point it25at any compatible endpoint by changing base_url and model.26 27Wiring into energy_data_dashboard.py28------------------------------------29In DashboardUI.__init__, after self.hub is set:30 31 from ai_command import LLMClient, CommandBar32 llm = LLMClient(33 api_key=read_key("GROQ_API_KEY", ".groq_key"),34 model=os.environ.get("GROQ_MODEL", "llama-3.3-70b-versatile"),35 )36 self.cmd = CommandBar(37 hub=self.hub,38 llm=llm,39 on_plot=self._ai_plot, # see adapter below40 get_context=self._ai_context, # see adapter below41 )42 43Add these two small adapters to DashboardUI:44 45 def _ai_plot(self, frames, replace):46 if replace:47 self.frames = list(frames)48 else:49 for df in frames:50 lbl = df["label"].iloc[0]51 self.frames = [f for f in self.frames52 if f["label"].iloc[0] != lbl] + [df]53 self._render()54 55 def _ai_context(self):56 return {57 "plotted": [f["label"].iloc[0] for f in self.frames],58 "recent_results": [f"{r['source']}:{r.get('id','')}"59 for r in getattr(self, "results", [])[:8]],60 }61 62Then drop self.cmd.panel() at the top of the main column in template().63 64Reference65---------66Groq. (2026). Text generation and structured outputs. GroqDocs.67 https://console.groq.com/docs/text-chat68"""69 70from __future__ import annotations71 72import json73import re74from typing import Any, Callable, Optional75 76import pandas as pd77import requests78import panel as pn79 80 81DEFAULT_BASE_URL = "https://api.groq.com/openai/v1"82DEFAULT_MODEL = "llama-3.3-70b-versatile"83 84PLAN_SCHEMA_HINT = """You translate a user's request about economic and energy85time series into a small JSON plan. Output ONLY a JSON object, no prose.86 87Sources:88- "FRED": macroeconomic and financial series (GDP, CPI, interest rates, broad89 price series such as the Henry Hub natural gas spot price, exchange rates).90- "EIA": United States energy statistics (natural gas, electricity, petroleum;91 often broken out by state and by sector such as residential or commercial).92- Use "any" only when you cannot tell.93 94JSON shape:95{96 "action": "plot" | "add" | "compare" | "clear",97 "items": [98 {99 "source": "FRED" | "EIA" | "any",100 "query": "short description used to search the source",101 "series_hint": "extra keyword to pick one EIA series, e.g. residential or Georgia",102 "exact_id": "a FRED series id only if you are confident, else empty"103 }104 ],105 "start": "optional start period, e.g. 1997-01 or 1997-01-01, else empty",106 "end": "optional end period, else empty"107}108 109Rules:110- "compare" and "plot" replace the chart with the listed items.111- "add" appends the listed items to whatever is already plotted.112- "clear" empties the chart and needs no items.113- Resolve references like "these two", "them", or "the second one" using the114 CONTEXT block (currently plotted series and recent search results). If the115 user says "compare these two" put the two referenced items in "items".116- For EIA energy series by sector or state, set source "EIA", put the dataset117 topic in "query" and the sector or state in "series_hint".118- Keep queries short and literal; do not invent identifiers.119 120Examples:121User: compare henry hub price and US residential natural gas consumption122{"action":"compare","items":[123 {"source":"FRED","query":"henry hub natural gas spot price","series_hint":"","exact_id":"DHHNGSP"},124 {"source":"EIA","query":"natural gas consumption","series_hint":"residential","exact_id":""}],125 "start":"","end":""}126 127User: add CPI since 2010128{"action":"add","items":[129 {"source":"FRED","query":"consumer price index all urban","series_hint":"","exact_id":"CPIAUCSL"}],130 "start":"2010-01-01","end":""}131 132User: clear it133{"action":"clear","items":[],"start":"","end":""}134"""135 136 137# ----------------------------------------------------------------------------138# LLM client (OpenAI-compatible; defaults to Groq). transport is injectable.139# ----------------------------------------------------------------------------140class LLMClient:141 def __init__(self, api_key: Optional[str], model: str = DEFAULT_MODEL,142 base_url: str = DEFAULT_BASE_URL,143 transport: Optional[Callable[[dict], dict]] = None):144 self.api_key = (api_key or "").strip()145 self.model = model146 self.base_url = base_url.rstrip("/")147 self._custom_transport = transport is not None148 self._transport = transport or self._http_post149 150 @property151 def available(self) -> bool:152 # Usable if we have a key for the real endpoint, or an injected transport.153 return bool(self.api_key) or self._custom_transport154 155 def _http_post(self, payload: dict) -> dict:156 r = requests.post(157 f"{self.base_url}/chat/completions",158 headers={"Authorization": f"Bearer {self.api_key}",159 "Content-Type": "application/json"},160 json=payload, timeout=40)161 if r.status_code != 200:162 raise RuntimeError(f"LLM HTTP {r.status_code}: {r.text[:200]}")163 return r.json()164 165 def plan(self, user_text: str, context: dict) -> dict:166 ctx = json.dumps(context or {}, ensure_ascii=False)167 payload = {168 "model": self.model,169 "temperature": 0.1,170 "response_format": {"type": "json_object"},171 "messages": [172 {"role": "system", "content": PLAN_SCHEMA_HINT},173 {"role": "user",174 "content": f"CONTEXT:\n{ctx}\n\nREQUEST:\n{user_text}"},175 ],176 }177 data = self._transport(payload)178 content = (data.get("choices", [{}])[0]179 .get("message", {}).get("content", "")) or ""180 return self._parse_json(content)181 182 @staticmethod183 def _parse_json(text: str) -> dict:184 text = text.strip()185 if text.startswith("```"):186 text = re.sub(r"^```(?:json)?|```$", "", text, flags=re.MULTILINE).strip()187 try:188 return json.loads(text)189 except Exception:190 m = re.search(r"\{.*\}", text, flags=re.DOTALL)191 if m:192 return json.loads(m.group(0))193 raise194 195 196# ----------------------------------------------------------------------------197# Command bar: resolve a plan against the existing DataHub and plot it198# ----------------------------------------------------------------------------199class CommandBar:200 def __init__(self, hub, llm: LLMClient,201 on_plot: Callable[[list, bool], None],202 get_context: Callable[[], dict] = lambda: {}):203 self.hub = hub204 self.llm = llm205 self.on_plot = on_plot206 self.get_context = get_context207 208 # -- resolution --------------------------------------------------------209 def _resolve_fred(self, item: dict, start: str, end: str) -> Optional[pd.DataFrame]:210 if self.hub.fred is None:211 return None212 exact = (item.get("exact_id") or "").strip()213 query = item.get("query", "")214 try:215 hits = self.hub.fred.search(query or exact, limit=8)216 except Exception:217 hits = []218 top = None219 if exact:220 top = next((h for h in hits if (h.get("id") or "").upper() == exact.upper()), None)221 if top is None:222 # Not in search results: fetch by id with id as the fallback name.223 df = self.hub.fetch_fred(exact, exact, item.get("units", ""), start, end)224 return df if (df is not None and not df.empty) else None225 if top is None:226 top = hits[0] if hits else None227 if top is None:228 return None229 df = self.hub.fetch_fred(top["id"], top.get("title", ""),230 top.get("units", ""), start, end)231 return df if not df.empty else None232 233 def _resolve_eia(self, item: dict, start: str, end: str) -> Optional[pd.DataFrame]:234 if self.hub.eia is None:235 return None236 datasets = self.hub.eia_dataset_search(item.get("query", ""), limit=20)237 if not datasets:238 return None239 # Prefer a dataset that exposes individual series when a hint is given.240 ds = next((d for d in datasets if d.get("series_facet")), datasets[0])241 route = ds["route"]242 freq = ds.get("default_frequency") or (243 ds.get("frequencies", [None]) or [None])[0] or ""244 facet = ds.get("series_facet")245 if facet:246 hint = (item.get("series_hint") or item.get("query") or "").strip()247 found = self.hub.eia_series(route, hint)248 results = found.get("results", [])249 if not results and hint:250 results = self.hub.eia_series(route, "").get("results", [])251 if not results:252 return None253 sid = results[0]["id"]254 sname = results[0].get("name", "")255 df = self.hub.fetch_eia(route, facet, sid, freq, start, end,256 series_name=sname)257 else:258 df = self.hub.fetch_eia(route, None, "", freq, start, end,259 series_name=ds.get("name", ""))260 return df if not df.empty else None261 262 def _resolve(self, item: dict, start: str, end: str):263 src = (item.get("source") or "any").upper()264 if src == "FRED":265 return self._resolve_fred(item, start, end), "FRED"266 if src == "EIA":267 return self._resolve_eia(item, start, end), "EIA"268 # "any": try FRED first, then EIA.269 df = self._resolve_fred(item, start, end)270 if df is not None:271 return df, "FRED"272 return self._resolve_eia(item, start, end), "EIA"273 274 # -- execute -----------------------------------------------------------275 def execute(self, user_text: str) -> str:276 user_text = (user_text or "").strip()277 if not user_text:278 return "Type a request first."279 try:280 plan = self.llm.plan(user_text, self.get_context())281 except Exception as exc:282 return f"Could not reach the language model: {exc}"283 284 action = (plan.get("action") or "plot").lower()285 start = (plan.get("start") or "").strip()286 end = (plan.get("end") or "").strip()287 items = plan.get("items") or []288 289 if action == "clear":290 self.on_plot([], True)291 return "Cleared the chart."292 293 if not items:294 return ("I could not tell which series you meant. Try naming them, "295 "for example: compare henry hub price and US residential "296 "gas consumption.")297 298 frames, resolved, missed = [], [], []299 for it in items:300 try:301 df, _src = self._resolve(it, start, end)302 except Exception as exc:303 df = None304 missed.append(f"{it.get('query','?')} ({exc})")305 if df is not None and not df.empty:306 frames.append(df)307 resolved.append(df["label"].iloc[0])308 else:309 missed.append(it.get("query", "?"))310 311 if not frames:312 return "Nothing matched. " + (313 "Tried: " + "; ".join(missed) if missed else "")314 315 replace = action != "add"316 self.on_plot(frames, replace)317 318 verb = "Plotted" if replace else "Added"319 msg = f"{verb} {', '.join(resolved)}."320 if missed:321 msg += f" Could not resolve: {', '.join(missed)}."322 return msg323 324 # -- panel widget ------------------------------------------------------325 def _find_groq_key(self) -> Optional[str]:326 import os327 from pathlib import Path328 v = os.environ.get("GROQ_API_KEY")329 if v and v.strip():330 return v.strip()331 names = [".groq_key", "_groq_key", "key.groq_key", "groq_key.txt", "groq.key"]332 dirs = [Path.cwd(), Path(__file__).resolve().parent, Path.home()]333 for d in dirs:334 for nm in names:335 p = d / nm336 try:337 if p.is_file():338 t = p.read_text(encoding="utf-8").strip()339 if t:340 return t341 except Exception:342 pass343 return None344 345 def panel(self) -> pn.Column:346 # Pull in a key from common locations if one was not passed in.347 if not self.llm.api_key:348 k = self._find_groq_key()349 if k:350 self.llm.api_key = k351 352 box = pn.widgets.TextInput(353 placeholder='Ask: "compare henry hub price and US residential gas '354 'consumption", "add CPI", "clear it"',355 sizing_mode="stretch_width")356 btn = pn.widgets.Button(name="Ask", button_type="primary", width=90)357 status = pn.pane.Markdown("", styles={"color": "#8FA3B8",358 "font-size": "12px"})359 360 def run(_=None):361 text = (box.value or "").strip()362 if not text:363 status.object = "Type a request first."364 return365 if not self.llm.api_key:366 k = self._find_groq_key()367 if k:368 self.llm.api_key = k369 if not self.llm.api_key:370 status.object = ("No Groq key found. Set GROQ_API_KEY, or save your "371 "key in a file named .groq_key or _groq_key next to "372 "the dashboard, then click Ask again.")373 return374 status.object = "Thinking..."375 status.object = self.execute(text)376 377 btn.on_click(run)378 box.param.watch(lambda e: run(), "value") # Enter submits379 380 # The box and button are always enabled, so the bar is always clickable.381 # position/z-index guard against a stacking-context overlay stealing clicks.382 return pn.Column(383 pn.Row(box, btn), status,384 sizing_mode="stretch_width",385 styles={"position": "relative", "z-index": "20",386 "pointer-events": "auto"})387