853 lines
33 KiB
Python
853 lines
33 KiB
Python
"""
|
||
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 "<key>_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 "<key>_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)
|