feat: add setup notebook and update env example for Athena

This commit is contained in:
2026-06-10 07:02:34 -04:00
parent a2420ed692
commit faa7d20b3e
27 changed files with 2483 additions and 151 deletions

View File

@@ -53,6 +53,8 @@ load_dotenv()
logger = logging.getLogger(__name__)
API_PREFIX = "/api/v1/tei"
ORBIT_PREFIX = "/api/v1/orbit"
ENGAGEMENT_PREFIX = "/api/v1/engagement"
class AthenaAPIError(Exception):
@@ -242,6 +244,34 @@ class TEIClient:
"""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,
@@ -286,10 +316,18 @@ class TEIClient:
)
def reorder_fields(self, report_public_id: str, field_ids: list[int]) -> dict:
"""Bulk-reorder fields. Spec: PATCH /reports/{id}/fields/reorder/."""
"""
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_ids": field_ids},
data={
"field_order": [
{"id": fid, "sort_order": i + 1}
for i, fid in enumerate(field_ids)
]
},
)
# ─────────────────────────────────────────────
@@ -345,96 +383,338 @@ class TEIClient:
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 _normalize_value(value: dict) -> dict:
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]:
"""
Normalize a value-row dict into the shape the API expects.
Expand one friendly value dict into wire-format rows for the bulk
``PUT /values/`` endpoint (one row per field/year, per Athena_TEI.md).
Accepts any of the following input forms and produces a uniform
wire-format dict::
Accepted input forms::
# annual fields
{"field_key": "A1", "year_1": 100, "year_2": 200, "year_3": 300, ...}
{"field_key": "A1", "year_values": {"1": 100, "2": 200, "3": 300}, ...}
# 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, ...}
Returns a dict like::
# 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": {...}}
{"field_key": "A1",
"year_values": {"1": 100.0, "2": 200.0, "3": 300.0},
"risk_adjustment": 0.15,
"notes": ""}
Output rows look like::
{"field_key": "A1", "year": 1, "value": "100", ...}
{"field_key": "rate", "year": None, "value": "0.10", ...}
"""
out: dict[str, Any] = {}
if "field_key" in value:
out["field_key"] = value["field_key"]
elif "field" in value:
out["field_key"] = value["field"]
field_key = value.get("field_key") or value.get("field")
if not field_key:
return []
# Collect annual year_N keys into year_values
year_values: dict[str, float] = {}
if "year_values" in value and isinstance(value["year_values"], dict):
year_values: dict[int, float] = {}
if isinstance(value.get("year_values"), dict):
for k, v in value["year_values"].items():
year_values[str(k)] = float(v) if v is not None else 0.0
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[str(n)] = float(raw) if raw is not None else 0.0
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:
out["year_values"] = year_values
if "value" in value and value["value"] is not None and not year_values:
out["value"] = value["value"]
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:
out["initial"] = float(value["initial"])
if value.get("risk_adjustment") is not None:
out["risk_adjustment"] = float(value["risk_adjustment"])
if value.get("notes"):
out["notes"] = str(value["notes"])
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."""
"""
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):
# Could be {"values": [...]} envelope, the TEITool wrapper, or a page
if "values" in result and isinstance(result["values"], list):
return result["values"]
if "results" in result and isinstance(result["results"], list):
return result["results"]
return []
if isinstance(result, list):
return result
return []
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. See ``_normalize_value`` for accepted shapes.
Bulk-update field values. Accepts friendly dicts (``year_values`` /
``year_N`` / ``value`` / ``initial``); see ``_rows_from_value``.
"""
payload = {"values": [self._normalize_value(v) for v in values]}
return self._put(f"{API_PREFIX}/tools/{public_id}/values/", data=payload)
def patch_value(self, public_id: str, field_key: str, **changes) -> dict:
"""
Patch a single field value by its ``field_key``.
Accepts the same shorthand as ``update_values`` (``year_1=…``, etc).
"""
body = self._normalize_value({"field_key": field_key, **changes})
body.pop("field_key", None) # carried in URL
return self._patch(
f"{API_PREFIX}/tools/{public_id}/values/{field_key}/", data=body
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
# ─────────────────────────────────────────────
@@ -467,11 +747,18 @@ class TEIClient:
return result["versions"]
return []
def save_version(self, public_id: str, note: str = "") -> dict:
"""Snapshot current values + summary as a new version."""
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={"note": note},
data={"date": date or datetime.now().strftime("%Y-%m-%d"), "note": note},
)
def get_version(self, public_id: str, version_number: int) -> dict:
@@ -555,9 +842,11 @@ class TEIClient:
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)
print(f" Net Present Value: ${_f(s.get('npv')):>16,.0f}")
print(f" ROI: {_f(s.get('roi')):>15,.0f}%")
payback = s.get("payback_months")
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)

View File

@@ -108,8 +108,8 @@ class TEIValue:
"""
A field value for a specific TEI tool instance.
The exact wire format is not fully pinned in the OpenAPI spec; we use a
convention that the client `_normalize_value` helper builds:
The exact wire format is not fully pinned in the OpenAPI spec; we use
the client's friendly value convention (see ``TEIClient.get_values``):
- annual fields: {field_key, year_values: {"1": ..., "2": ...},
risk_adjustment, notes}