""" TEI Client — Athena API wrapper for Palladium. Endpoints (per Athena API.yaml, all under ``/api/v1/tei/``): Reports (templates) GET /reports/ list_reports GET /reports/{public_id}/ get_report GET /reports/{public_id}/fields/ list_fields PATCH /reports/{public_id}/fields/reorder/ reorder_fields Tools (instances) GET /tools/ list_tools POST /tools/ create_tool GET /tools/{public_id}/ get_tool PATCH /tools/{public_id}/ update_tool DELETE /tools/{public_id}/ delete_tool Values (data entry) GET /tools/{public_id}/values/ get_values PUT /tools/{public_id}/values/ update_values (bulk) PATCH /tools/{public_id}/values/{field_key}/ patch_value (single) Calculation & summary POST /tools/{public_id}/calculate/ calculate GET /tools/{public_id}/summary/ get_summary GET /summary/ aggregate_summary Versions GET /tools/{public_id}/versions/ list_versions POST /tools/{public_id}/versions/ save_version GET /tools/{public_id}/versions/{n}/ get_version Export GET /tools/{public_id}/export/ export Authentication uses the ``Authorization: Api-Key {key}`` header. """ from __future__ import annotations import json import logging import os from datetime import datetime from typing import Any import requests from dotenv import load_dotenv load_dotenv() logger = logging.getLogger(__name__) API_PREFIX = "/api/v1/tei" ORBIT_PREFIX = "/api/v1/orbit" ENGAGEMENT_PREFIX = "/api/v1/engagement" class AthenaAPIError(Exception): """Raised when Athena returns a non-success response.""" def __init__(self, status_code: int, detail: str, url: str): self.status_code = status_code self.detail = detail self.url = url super().__init__(f"Athena API {status_code} at {url}: {detail}") class TEIClient: """ Client for Athena's TEI Calculator API. Wraps every TEI endpoint and provides a few convenience helpers used by the Palladium notebooks and Streamlit app. Environment variables (read via python-dotenv): ATHENA_BASE_URL e.g. https://athena.nttdata.com ATHENA_API_KEY Api-Key value (admin-issued) Example:: from core.tei_client import TEIClient client = TEIClient() client.test_connection() for r in client.list_reports(): print(r["name"]) """ def __init__( self, base_url: str | None = None, api_key: str | None = None, timeout: int = 30, ): self.base_url = (base_url or os.getenv("ATHENA_BASE_URL", "")).rstrip("/") self.api_key = api_key or os.getenv("ATHENA_API_KEY", "") self.timeout = timeout if not self.base_url: raise ValueError( "ATHENA_BASE_URL is required. Set it in .env or pass base_url." ) if not self.api_key: raise ValueError( "ATHENA_API_KEY is required. Set it in .env or pass api_key." ) self.session = requests.Session() self.session.headers.update( { "Authorization": f"Api-Key {self.api_key}", "Content-Type": "application/json", "Accept": "application/json", } ) logger.info("TEIClient initialised for %s", self.base_url) # ───────────────────────────────────────────── # Internal HTTP helpers # ───────────────────────────────────────────── def _url(self, path: str) -> str: if not path.startswith("/"): path = f"/{path}" return f"{self.base_url}{path}" def _request( self, method: str, path: str, params: dict | None = None, json_data: Any | None = None, ) -> Any: url = self._url(path) logger.debug("%s %s", method.upper(), url) try: response = self.session.request( method=method, url=url, params=params, json=json_data, timeout=self.timeout, ) except requests.ConnectionError as e: raise AthenaAPIError(0, f"Connection failed: {e}", url) from e except requests.Timeout as e: raise AthenaAPIError(408, "Request timed out", url) from e if response.status_code >= 400: try: payload = response.json() detail = payload.get("detail") or json.dumps(payload) except (json.JSONDecodeError, ValueError, AttributeError): detail = response.text raise AthenaAPIError(response.status_code, detail, url) if response.status_code == 204 or not response.content: return {} return response.json() def _get(self, path: str, params: dict | None = None) -> Any: return self._request("GET", path, params=params) def _post(self, path: str, data: Any | None = None) -> Any: return self._request("POST", path, json_data=data) def _put(self, path: str, data: Any | None = None) -> Any: return self._request("PUT", path, json_data=data) def _patch(self, path: str, data: Any | None = None) -> Any: return self._request("PATCH", path, json_data=data) def _delete(self, path: str) -> Any: return self._request("DELETE", path) def _paginated(self, path: str, params: dict | None = None) -> list[dict]: """ Fetch all pages of a paginated list endpoint. Athena uses the standard DRF page/results envelope:: {"count": N, "next": url|None, "previous": ..., "results": [...]} """ out: list[dict] = [] result = self._get(path, params=params) while True: if isinstance(result, list): out.extend(result) return out if not isinstance(result, dict): return out out.extend(result.get("results", []) or []) next_url = result.get("next") if not next_url: return out # Follow absolute next URL try: result = self.session.get(next_url, timeout=self.timeout).json() except Exception: # pragma: no cover – defensive return out # ───────────────────────────────────────────── # Connection test # ───────────────────────────────────────────── def test_connection(self) -> dict: """Verify API connectivity and authentication.""" try: result = self._get(f"{API_PREFIX}/reports/") count = ( result.get("count", len(result.get("results", []))) if isinstance(result, dict) else len(result) ) return { "status": "ok", "base_url": self.base_url, "authenticated": True, "reports_found": count, "timestamp": datetime.now().isoformat(), } except AthenaAPIError as e: return { "status": "error", "base_url": self.base_url, "authenticated": e.status_code != 401, "error_code": e.status_code, "detail": e.detail, "timestamp": datetime.now().isoformat(), } # ───────────────────────────────────────────── # Reports (templates) # ───────────────────────────────────────────── def list_reports(self) -> list[dict]: """List all TEI report templates (auto-paginated).""" return self._paginated(f"{API_PREFIX}/reports/") def get_report(self, public_id: str) -> dict: """Get a TEI report template by its public_id.""" return self._get(f"{API_PREFIX}/reports/{public_id}/") def create_report( self, name: str, vendor: str, version: str = "1.0", description: str = "", analysis_period_years: int = 3, discount_rate: float | str = "0.10", status: str = "draft", ) -> dict: """Create a new TEI report template (admin only).""" return self._post( f"{API_PREFIX}/reports/", data={ "name": name, "vendor": vendor, "version": version, "description": description, "analysis_period_years": analysis_period_years, "discount_rate": str(discount_rate), "status": status, }, ) def update_report(self, public_id: str, **changes) -> dict: """Patch report template metadata (e.g. ``status='active'``).""" return self._patch(f"{API_PREFIX}/reports/{public_id}/", data=changes) def list_fields( self, report_public_id: str, table: str | None = None, ) -> list[dict]: """ Get field definitions for a report. Args: report_public_id: The report template's public_id (12-char short UUID). table: Optional filter — ``'benefits'`` or ``'costs'``. Returns a list of field-definition dicts. See ``TEIField.from_dict`` for the expected shape. """ params = {"table": table} if table else None rows = self._paginated( f"{API_PREFIX}/reports/{report_public_id}/fields/", params=params ) # Defensive — server-side filter may not be implemented; filter locally. if table: rows = [r for r in rows if r.get("table") == table] rows.sort(key=lambda r: (r.get("table", ""), r.get("sort_order") or 0)) return rows def create_field(self, report_public_id: str, field: dict) -> dict: """Create a new field definition under a report (admin only).""" return self._post( f"{API_PREFIX}/reports/{report_public_id}/fields/", data=field ) def update_field(self, report_public_id: str, field_id: int, **changes) -> dict: """Patch one field definition by its integer id.""" return self._patch( f"{API_PREFIX}/reports/{report_public_id}/fields/{field_id}/", data=changes, ) def delete_field(self, report_public_id: str, field_id: int) -> dict: return self._delete( f"{API_PREFIX}/reports/{report_public_id}/fields/{field_id}/" ) def reorder_fields(self, report_public_id: str, field_ids: list[int]) -> dict: """ Bulk-reorder fields. Spec body: ``{"field_order": [{"id", "sort_order"}]}``. ``field_ids`` is the desired order; sort_order is assigned 1..N. """ return self._patch( f"{API_PREFIX}/reports/{report_public_id}/fields/reorder/", data={ "field_order": [ {"id": fid, "sort_order": i + 1} for i, fid in enumerate(field_ids) ] }, ) # ───────────────────────────────────────────── # Tools (instances) # ───────────────────────────────────────────── def list_tools(self) -> list[dict]: """List TEI tool instances owned by the current API key.""" return self._paginated(f"{API_PREFIX}/tools/") def get_tool(self, public_id: str) -> dict: """Get a TEI tool instance by public_id.""" return self._get(f"{API_PREFIX}/tools/{public_id}/") def create_tool( self, report_public_id: str, proposal: int | None = None, engagement: int | None = None, name: str | None = None, status: str = "draft", ) -> dict: """ Create a new TEI tool instance from a report template. Athena scopes a TEI tool to a *Proposal* (which itself belongs to an Opportunity) and/or an *Engagement*. Pass the integer PK of either or both to link the tool. """ data: dict[str, Any] = {"report": report_public_id, "status": status} if proposal is not None: data["proposal"] = proposal if engagement is not None: data["engagement"] = engagement if name: data["name"] = name return self._post(f"{API_PREFIX}/tools/", data=data) def update_tool( self, public_id: str, name: str | None = None, status: str | None = None, ) -> dict: """Update tool metadata. Only ``name`` and ``status`` are mutable.""" data: dict[str, Any] = {} if name is not None: data["name"] = name if status is not None: data["status"] = status return self._patch(f"{API_PREFIX}/tools/{public_id}/", data=data) def delete_tool(self, public_id: str) -> dict: return self._delete(f"{API_PREFIX}/tools/{public_id}/") # ───────────────────────────────────────────── # CRM context — clients, opportunities, proposals, engagements # # A TEI tool must attach to a Proposal OR an Engagement. These methods # let Palladium browse the CRM (Athena's "Orbit" module) so the user # selects an existing record — and client data (industry, agent counts, # revenue) flows into the study without re-entry. # ───────────────────────────────────────────── def list_clients(self, search: str | None = None) -> list[dict]: """List CRM clients, optionally filtered by name/legal name/overview.""" params = {"search": search} if search else None return self._paginated(f"{ORBIT_PREFIX}/clients/", params=params) def get_client(self, client_id: int) -> dict: """Full client record — vertical, employee_count, revenue, contact_center_agent_count, supervisor_count, location_count, …""" return self._get(f"{ORBIT_PREFIX}/clients/{int(client_id)}/") def client_profile(self, client_id: int) -> dict: """The TEI-relevant subset of a client record (for assumptions).""" c = self.get_client(client_id) keys = ( "id", "name", "abbreviated_name", "vertical", "client_type", "employee_count", "revenue", "contact_center_agent_count", "service_desk_agent_count", "supervisor_count", "location_count", ) return {k: c.get(k) for k in keys} def list_opportunities(self, search: str | None = None) -> list[dict]: """List opportunities (each embeds its read-only ``client``).""" params = {"search": search} if search else None return self._paginated(f"{ORBIT_PREFIX}/opportunities/", params=params) def create_opportunity(self, name: str, client_id: int, **extra) -> dict: """Create an opportunity for a client (sandbox/demo convenience).""" return self._post( f"{ORBIT_PREFIX}/opportunities/", data={"name": name, "client_id": int(client_id), **extra}, ) def list_proposals( self, opportunity_id: int | None = None, status: str | None = None, search: str | None = None, ) -> list[dict]: """List proposals (each embeds ``opportunity`` → ``client``).""" params: dict[str, Any] = {} if opportunity_id is not None: params["opportunity_id"] = int(opportunity_id) if status: params["status"] = status if search: params["search"] = search return self._paginated(f"{ORBIT_PREFIX}/proposals/", params=params or None) def create_proposal( self, name: str, opportunity_id: int, status: str = "Draft", **extra ) -> dict: """Create a proposal under an opportunity.""" return self._post( f"{ORBIT_PREFIX}/proposals/", data={ "name": name, "opportunity_id": int(opportunity_id), "status": status, **extra, }, ) def list_engagements(self, search: str | None = None) -> list[dict]: """List engagements (summary rows include ``client_name``).""" params = {"search": search} if search else None return self._paginated(f"{ENGAGEMENT_PREFIX}/engagements/", params=params) def proposals_for_client(self, client_id: int) -> list[dict]: """All proposals whose opportunity belongs to ``client_id``.""" out = [] for p in self.list_proposals(): opp = p.get("opportunity") or {} cli = opp.get("client") or {} if cli.get("id") == int(client_id): out.append(p) return out def engagements_for_client(self, client_name: str) -> list[dict]: """All engagements matching a client's name (summary list filter).""" return [ e for e in self.list_engagements(search=client_name) if (e.get("client_name") or "").lower() == client_name.lower() ] # ───────────────────────────────────────────── # Values (data entry) # ───────────────────────────────────────────── #: Suffix used for companion non-annual fields that hold a cost's #: Year-0 "Initial" amount (the TEI API has no native year-0 concept). INITIAL_SUFFIX = "_initial" @staticmethod def _coerce_float(raw: Any, default: float = 0.0) -> float: try: return float(raw) if raw is not None and raw != "" else default except (TypeError, ValueError): return default @classmethod def _rows_from_value(cls, value: dict) -> list[dict]: """ Expand one friendly value dict into wire-format rows for the bulk ``PUT /values/`` endpoint (one row per field/year, per Athena_TEI.md). Accepted input forms:: # annual fields (either shorthand) {"field_key": "A1", "year_1": 100, "year_2": 200, ...} {"field_key": "A1", "year_values": {"1": 100, "2": 200}, ...} # non-annual scalars {"field_key": "rate", "value": 0.10, ...} # costs with a Year-0 component → companion "_initial" # non-annual field (must exist on the report; see provisioning) {"field_key": "impl", "initial": 1_000_000, "year_values": {...}} Output rows look like:: {"field_key": "A1", "year": 1, "value": "100", ...} {"field_key": "rate", "year": None, "value": "0.10", ...} """ field_key = value.get("field_key") or value.get("field") if not field_key: return [] year_values: dict[int, float] = {} if isinstance(value.get("year_values"), dict): for k, v in value["year_values"].items(): year_values[int(k)] = cls._coerce_float(v) for key, raw in value.items(): if key.startswith("year_"): try: n = int(key.split("_", 1)[1]) except ValueError: continue year_values[n] = cls._coerce_float(raw) risk_adjustment = ( float(value["risk_adjustment"]) if value.get("risk_adjustment") is not None else None ) notes = str(value["notes"]) if value.get("notes") else None rows: list[dict] = [] if year_values: for i, year in enumerate(sorted(year_values)): row: dict[str, Any] = { "field_key": field_key, "year": year, "value": str(year_values[year]), "risk_adjustment": ( str(risk_adjustment) if risk_adjustment is not None else None ), # Attach narrative notes to the first year row only. "notes": notes if i == 0 else None, } rows.append(row) elif value.get("value") is not None: rows.append( { "field_key": field_key, "year": None, "value": str(value["value"]), "risk_adjustment": ( str(risk_adjustment) if risk_adjustment is not None else None ), "notes": notes, } ) # Year-0 "Initial" amount → companion non-annual field. if value.get("initial") is not None: rows.append( { "field_key": f"{field_key}{cls.INITIAL_SUFFIX}", "year": None, "value": str(cls._coerce_float(value["initial"])), "risk_adjustment": None, "notes": None, } ) return rows @classmethod def _friendly_value_row(cls, raw: dict) -> dict: """ Convert one GET ``/values/`` row (documented shape) into the friendly internal shape used by the notebooks and the Streamlit app:: {"field_key", "label", "table", "category", "field_type", "is_annual", "risk_adjustment", "year_values": {"1": 100.0, ...}, "value": 0.10, "notes": "…"} """ out = { k: raw.get(k) for k in ( "id", "field_key", "label", "table", "category", "field_type", "is_annual", "notes", ) } out["risk_adjustment"] = ( cls._coerce_float(raw.get("risk_adjustment"), 0.0) if raw.get("risk_adjustment") is not None else None ) years = raw.get("years") if isinstance(years, dict): # documented annual shape year_values: dict[str, float] = {} notes_parts: list[str] = [] for y in sorted(years, key=int): cell = years[y] or {} year_values[str(int(y))] = cls._coerce_float(cell.get("value")) if cell.get("risk_adjustment") is not None: out["risk_adjustment"] = cls._coerce_float( cell["risk_adjustment"], out["risk_adjustment"] or 0.0 ) if cell.get("notes"): notes_parts.append(str(cell["notes"])) out["year_values"] = year_values if notes_parts and not out.get("notes"): out["notes"] = " | ".join(notes_parts) elif isinstance(raw.get("year_values"), dict): # already friendly out["year_values"] = { str(int(k)): cls._coerce_float(v) for k, v in raw["year_values"].items() } else: # non-annual scalar out["value"] = cls._coerce_float(raw.get("value")) return out def get_values(self, public_id: str) -> list[dict]: """ Get all current field values for a TEI tool instance. Returns friendly rows (see ``_friendly_value_row``). Companion ``*_initial`` fields are folded into their parent cost row as an ``initial`` key rather than returned as standalone rows. """ result = self._get(f"{API_PREFIX}/tools/{public_id}/values/") if isinstance(result, dict): raw_rows = ( result.get("values") or result.get("results") or [] ) elif isinstance(result, list): raw_rows = result else: raw_rows = [] rows = [self._friendly_value_row(r) for r in raw_rows if isinstance(r, dict)] # Fold "_initial" companions into their parent row. by_key = {r["field_key"]: r for r in rows if r.get("field_key")} folded: list[dict] = [] for row in rows: fk = row.get("field_key") or "" if fk.endswith(self.INITIAL_SUFFIX): parent = by_key.get(fk[: -len(self.INITIAL_SUFFIX)]) if parent is not None: parent["initial"] = row.get("value", 0.0) continue # folded — don't emit standalone folded.append(row) return folded def update_values(self, public_id: str, values: list[dict]) -> dict: """ Bulk-update field values. Accepts friendly dicts (``year_values`` / ``year_N`` / ``value`` / ``initial``); see ``_rows_from_value``. """ rows: list[dict] = [] for v in values: rows.extend(self._rows_from_value(v)) return self._put( f"{API_PREFIX}/tools/{public_id}/values/", data={"values": rows} ) def patch_value( self, public_id: str, field_key: str, year: int | None = None, value: Any | None = None, risk_adjustment: float | None = None, notes: str | None = None, ) -> dict: """ Patch a single field value. ``year`` is required for annual fields (passed as a query param per the API spec); omit it for non-annual fields. """ body: dict[str, Any] = {} if value is not None: body["value"] = str(value) if risk_adjustment is not None: body["risk_adjustment"] = str(risk_adjustment) if notes is not None: body["notes"] = notes path = f"{API_PREFIX}/tools/{public_id}/values/{field_key}/" if year is not None: path = f"{path}?year={int(year)}" return self._patch(path, data=body) # ───────────────────────────────────────────── # Calculation & summary # ───────────────────────────────────────────── def calculate(self, public_id: str) -> dict: """Trigger server-side calculation; returns the updated summary.""" return self._post(f"{API_PREFIX}/tools/{public_id}/calculate/") def get_summary(self, public_id: str) -> dict: """Return the most-recent summary (404 if never calculated).""" return self._get(f"{API_PREFIX}/tools/{public_id}/summary/") def aggregate_summary(self) -> dict: """Aggregate NPV across all tools owned by the current API key.""" return self._get(f"{API_PREFIX}/summary/") # ───────────────────────────────────────────── # Versions # ───────────────────────────────────────────── def list_versions(self, public_id: str) -> list[dict]: """List all saved version snapshots for a TEI tool.""" result = self._get(f"{API_PREFIX}/tools/{public_id}/versions/") if isinstance(result, list): return result if isinstance(result, dict): if "results" in result and isinstance(result["results"], list): return result["results"] if "versions" in result and isinstance(result["versions"], list): return result["versions"] return [] def save_version( self, public_id: str, note: str = "", date: str | None = None ) -> dict: """ Snapshot current values + summary as a new version. ``date`` defaults to today (ISO ``YYYY-MM-DD``); the API requires it. Saving a version auto-triggers ``/calculate/`` server-side. """ return self._post( f"{API_PREFIX}/tools/{public_id}/versions/", data={"date": date or datetime.now().strftime("%Y-%m-%d"), "note": note}, ) def get_version(self, public_id: str, version_number: int) -> dict: """Get a single version's full snapshot.""" return self._get( f"{API_PREFIX}/tools/{public_id}/versions/{int(version_number)}/" ) # ───────────────────────────────────────────── # Export # ───────────────────────────────────────────── def export(self, public_id: str) -> dict: """ Return the LLM-ready export payload for the report pipeline. The shape is determined by Athena and consumed by Peitho / html2docx; Palladium's ``core.export.report_data`` builds on this. """ return self._get(f"{API_PREFIX}/tools/{public_id}/export/") # ───────────────────────────────────────────── # Convenience # ───────────────────────────────────────────── def get_benefits(self, public_id: str) -> list[dict]: """Return only benefit-table values (table='benefits').""" return [v for v in self.get_values(public_id) if v.get("table") == "benefits"] def get_costs(self, public_id: str) -> list[dict]: """Return only cost-table values (table='costs').""" return [v for v in self.get_values(public_id) if v.get("table") == "costs"] def get_tool_with_data(self, public_id: str) -> dict: """ Bundle a tool, its field definitions, current values, and summary. Convenience for notebook initialisation. The summary is allowed to 404 (returned as ``None``) when the tool has never been calculated. """ tool = self.get_tool(public_id) report_pid = tool.get("report") if isinstance(report_pid, dict): report_pid = report_pid.get("id") or report_pid.get("public_id") fields = self.list_fields(report_pid) if report_pid else [] values = self.get_values(public_id) try: summary = self.get_summary(public_id) except AthenaAPIError as e: if e.status_code == 404: summary = None else: raise return { "tool": tool, "fields": fields, "values": values, "summary": summary, } # ───────────────────────────────────────────── # Display # ───────────────────────────────────────────── def __repr__(self) -> str: # pragma: no cover return f"TEIClient(base_url='{self.base_url}')" def print_summary(self, public_id: str) -> None: """Pretty-print a financial summary block for notebooks/REPL.""" s = self.get_summary(public_id) def _f(v: Any, default: float = 0.0) -> float: try: return float(v) if v is not None else default except (TypeError, ValueError): return default print("═" * 56) print(" TEI Financial Summary") print("═" * 56) print(f" Total Benefits (PV): ${_f(s.get('total_benefits_pv')):>16,.0f}") print(f" Total Costs (PV): ${_f(s.get('total_costs_pv')):>16,.0f}") print("─" * 56) npv = s.get("net_present_value", s.get("npv")) roi = s.get("roi_percentage", s.get("roi")) print(f" Net Present Value: ${_f(npv):>16,.0f}") print(f" ROI: {_f(roi):>15,.0f}%") payback = s.get("payback_period_months", s.get("payback_months")) payback_str = f"{_f(payback):.1f} months" if payback is not None else "N/A" print(f" Payback: {payback_str:>17}") print("═" * 56)