anusuya-gurusamy/forensic-audit-openenv
0
1from __future__ import annotations2 3import json4import logging5import os6import uuid7from pathlib import Path8from typing import Any, Dict, Set, Tuple9 10from .ledger_query import LedgerQueryEngine11from .models import (12 Action,13 AuditObservation,14 CompanyData,15 IssueVerdictAction,16 LinkEvidenceAction,17 NavigateAction,18 QueryLedgerAction,19 VerdictPayload,20)21from .tasks import get_grader22 23MAX_STEPS = 1524# Support both local dev (data/ next to src/) and container (data/ in WORKDIR)25_default_data = Path(__file__).parent.parent / "data"26if not _default_data.exists():27 _default_data = Path("data")28DATA_DIR = Path(os.getenv("DATA_DIR", str(_default_data)))29 30COMPANY_MAP = {31 "TechCorp": "company_1",32 "RetailCo": "company_2",33 "ManufacturingInc": "company_3",34 "FinanceHub": "company_4",35 "HealthcarePlus": "company_5",36 "company_1": "company_1",37 "company_2": "company_2",38 "company_3": "company_3",39 "company_4": "company_4",40 "company_5": "company_5",41}42 43ANOMALY_ALERTS: Dict[str, list[str]] = {44 "company_1": ["AR spike detected: Accounts Receivable increased 40% YoY"],45 "company_2": ["12 duplicate invoice payments detected in Accounts Payable"],46 "company_3": ["Narrative claims strong cash generation but OCF is negative"],47 "company_4": [],48 "company_5": [],49}50 51AVAILABLE_ACTIONS = [52 "navigate_to(income_statement)",53 "navigate_to(balance_sheet)",54 "navigate_to(cashflow)",55 "navigate_to(general_ledger)",56 "navigate_to(sub_ledger)",57 "navigate_to(narrative)",58 "query_ledger(filters={})",59 "link_evidence(source_id, target_id, relationship)",60 "issue_verdict(conclusion, rationale, evidence_chain)",61]62 63 64# ---------------------------------------------------------------------------65# RewardComputer66# ---------------------------------------------------------------------------67 68class RewardComputer:69 CLAMP_MIN: float = -2.070 CLAMP_MAX: float = 3.071 72 def delta_navigate(self, view: str, visited: Set[str]) -> float:73 return 0.1 if view not in visited else 0.074 75 def delta_query(self, row_count: int) -> float:76 return 0.3 if row_count > 0 else -0.0577 78 def delta_link_evidence(self, link: Dict[str, Any], scratchpad: list[Dict]) -> float:79 if self._is_contradictory(link, scratchpad):80 return -0.1581 return 0.282 83 def delta_loop(self, action_repr: str, history: list[str]) -> float:84 return -0.1 if action_repr in history else 0.085 86 def delta_verdict(self, task_score: float, evidence_complete: bool) -> float:87 if task_score > 0.5:88 return 1.0 if evidence_complete else 0.589 return 0.090 91 # ------------------------------------------------------------------92 # Upgrade 2: Six Sigma Redundancy Reward93 # Reward +0.3 when agent verifies the same financial fact across94 # multiple independent data sources (GL + narrative, or CF + BS).95 # This trains agents toward High-Reliability Organization (HRO) standards.96 # ------------------------------------------------------------------97 98 def delta_redundancy(self, source_id: str, scratchpad: list[Dict]) -> float:99 """100 +0.3 if this link corroborates a fact already in the scratchpad101 from a different data source (cross-system verification).102 """103 CROSS_SOURCE_PAIRS = [104 ("income_statement", "narrative"),105 ("cashflow", "balance_sheet"),106 ("general_ledger", "sub_ledger"),107 ("general_ledger", "narrative"),108 ("cashflow", "narrative"),109 ]110 src_lower = source_id.lower()111 for existing in scratchpad:112 existing_src = str(existing.get("source_id", "")).lower()113 for a, b in CROSS_SOURCE_PAIRS:114 if (a in src_lower and b in existing_src) or (b in src_lower and a in existing_src):115 return 0.3116 return 0.0117 118 def clamp(self, reward: float) -> float:119 return max(self.CLAMP_MIN, min(self.CLAMP_MAX, reward))120 121 def _is_contradictory(self, link: Dict[str, Any], scratchpad: list[Dict]) -> bool:122 src = link.get("source_id", "")123 tgt = link.get("target_id", "")124 rel = link.get("relationship", "")125 for existing in scratchpad:126 if existing.get("source_id") == src and existing.get("target_id") == tgt:127 existing_rel = existing.get("relationship", "")128 if existing_rel and existing_rel != rel:129 return True130 return False131 132 133# ---------------------------------------------------------------------------134# EpisodeLogger135# ---------------------------------------------------------------------------136 137class EpisodeLogger:138 def __init__(self) -> None:139 self._logger = logging.getLogger("forensic_audit_env")140 if not self._logger.handlers:141 handler = logging.StreamHandler()142 handler.setFormatter(logging.Formatter("%(message)s"))143 self._logger.addHandler(handler)144 self._logger.setLevel(logging.INFO)145 146 def log_start(self, episode_id: str, company: str, task_id: int) -> None:147 self._logger.info(json.dumps({148 "tag": "[START]", "episode_id": episode_id,149 "company": company, "task_id": task_id,150 }))151 152 def log_step(self, step: int, action_type: str, reward_delta: float) -> None:153 self._logger.info(json.dumps({154 "tag": "[STEP]", "step": step,155 "action_type": action_type, "reward_delta": round(reward_delta, 4),156 }))157 158 def log_end(self, episode_id: str, total_reward: float, task_score: float, reason: str) -> None:159 self._logger.info(json.dumps({160 "tag": "[END]", "episode_id": episode_id,161 "total_reward": round(total_reward, 4),162 "task_score": round(task_score, 4), "reason": reason,163 }))164 165 166# ---------------------------------------------------------------------------167# ForensicAuditEnv168# ---------------------------------------------------------------------------169 170class ForensicAuditEnv:171 def __init__(self, data_dir: Path = DATA_DIR) -> None:172 self._data_dir = data_dir173 self._reward = RewardComputer()174 self._logger = EpisodeLogger()175 self._reset_state()176 177 def _reset_state(self) -> None:178 self._episode_id: str = ""179 self._step_count: int = 0180 self._cumulative_reward: float = 0.0181 self._visited_views: Set[str] = set()182 self._action_history: list[str] = []183 self._scratchpad: list[Dict[str, Any]] = []184 self._current_view: str = "income_statement"185 self._active: bool = False186 self._done: bool = False187 self._company_data: CompanyData | None = None188 self._task_id: int = 1189 self._task_score: float = 0.0190 self._ledger_engine: LedgerQueryEngine | None = None191 self._sub_ledger_engine: LedgerQueryEngine | None = None192 193 def reset(self, company_id: str, task_id: int, force: bool = False) -> AuditObservation:194 if self._active and not force:195 raise EpisodeActiveError("Episode in progress. Use force=true to reset.")196 197 # Close existing DB connections198 if self._ledger_engine:199 self._ledger_engine.close()200 if self._sub_ledger_engine:201 self._sub_ledger_engine.close()202 203 self._reset_state()204 205 cid = COMPANY_MAP.get(company_id)206 if cid is None:207 raise CompanyNotFoundError(f"Company '{company_id}' not found")208 if task_id not in (1, 2, 3):209 raise ValueError("task_id must be 1, 2, or 3")210 211 company_dir = self._data_dir / cid212 self._company_data = self._load_company(cid, company_dir)213 self._task_id = task_id214 self._episode_id = str(uuid.uuid4())215 self._active = True216 self._current_view = "income_statement"217 218 self._ledger_engine = LedgerQueryEngine(self._company_data.ledger_db_path)219 self._sub_ledger_engine = LedgerQueryEngine(self._company_data.sub_ledger_db_path)220 221 self._logger.log_start(self._episode_id, cid, task_id)222 return self._build_observation()223 224 def step(self, action: Action) -> Tuple[AuditObservation, float, bool, Dict[str, Any]]:225 if not self._active:226 raise NoActiveEpisodeError("No active episode. Call /reset first.")227 if self._done:228 raise EpisodeDoneError("Episode is complete. Call /reset to start a new episode.")229 230 action_repr = action.model_dump_json()231 loop_delta = self._reward.delta_loop(action_repr, self._action_history)232 self._action_history.append(action_repr)233 self._cumulative_reward += loop_delta234 235 delta = 0.0236 done = False237 info: Dict[str, Any] = {}238 239 if isinstance(action, NavigateAction):240 delta = self._handle_navigate(action)241 elif isinstance(action, QueryLedgerAction):242 delta = self._handle_query(action)243 elif isinstance(action, LinkEvidenceAction):244 delta = self._handle_link(action)245 elif isinstance(action, IssueVerdictAction):246 delta, done, info = self._handle_verdict(action)247 248 self._cumulative_reward += delta249 self._step_count += 1250 251 if self._step_count >= MAX_STEPS and not done:252 done = True253 self._logger.log_end(254 self._episode_id,255 self._reward.clamp(self._cumulative_reward),256 self._task_score,257 "step_limit",258 )259 260 if done:261 self._done = True262 self._active = False263 264 total_delta = loop_delta + delta265 self._logger.log_step(self._step_count, action.action_type, total_delta)266 267 clamped = self._reward.clamp(self._cumulative_reward)268 obs = self._build_observation()269 return obs, clamped, done, info270 271 def get_state(self) -> AuditObservation:272 return self._build_observation()273 274 # ------------------------------------------------------------------275 # Action handlers276 # ------------------------------------------------------------------277 278 def _handle_navigate(self, action: NavigateAction) -> float:279 delta = self._reward.delta_navigate(action.view, self._visited_views)280 self._visited_views.add(action.view)281 self._current_view = action.view282 return delta283 284 def _handle_query(self, action: QueryLedgerAction) -> float:285 engine = (286 self._sub_ledger_engine287 if self._current_view == "sub_ledger"288 else self._ledger_engine289 )290 table = "sub_ledger" if self._current_view == "sub_ledger" else "transactions"291 if engine is None:292 return -0.05293 rows = engine.query(action.filters, table=table)294 self._visible_query_results = rows295 return self._reward.delta_query(len(rows))296 297 def _handle_link(self, action: LinkEvidenceAction) -> float:298 link = {299 "source_id": action.source_id,300 "target_id": action.target_id,301 "relationship": action.relationship,302 }303 delta = self._reward.delta_link_evidence(link, self._scratchpad)304 # Upgrade 2: Six Sigma redundancy reward for cross-source verification305 delta += self._reward.delta_redundancy(action.source_id, self._scratchpad)306 self._scratchpad.append(link)307 return delta308 309 def _handle_verdict(310 self, action: IssueVerdictAction311 ) -> Tuple[float, bool, Dict[str, Any]]:312 payload = VerdictPayload(313 conclusion=action.conclusion,314 rationale=action.rationale,315 evidence_chain=action.evidence_chain,316 scratchpad=self._scratchpad,317 task_id=self._task_id,318 flagged_tx_ids=self._extract_tx_ids(action),319 )320 grader = get_grader(self._task_id)321 score = grader.grade(payload)322 self._task_score = score323 324 evidence_complete = len(self._scratchpad) >= 3325 verdict_delta = self._reward.delta_verdict(score, evidence_complete)326 327 self._logger.log_end(328 self._episode_id,329 self._reward.clamp(self._cumulative_reward + verdict_delta),330 score,331 "verdict",332 )333 return verdict_delta, True, {"task_score": score, "verdict": action.conclusion}334 335 def _extract_tx_ids(self, action: IssueVerdictAction) -> list[str]:336 ids: list[str] = []337 for entry in action.evidence_chain:338 tx = entry.get("tx_id") or entry.get("transaction_id")339 if tx:340 ids.append(str(tx))341 for key in ("tx_ids", "duplicate_ids", "flagged_ids"):342 val = entry.get(key)343 if isinstance(val, list):344 ids.extend(str(v) for v in val)345 return ids346 347 # ------------------------------------------------------------------348 # Observation builder349 # ------------------------------------------------------------------350 351 def _build_observation(self) -> AuditObservation:352 visible = self._get_visible_data()353 cid = self._company_data.company_id if self._company_data else "company_1"354 alerts = ANOMALY_ALERTS.get(cid, [])355 progress = min(1.0, self._step_count / MAX_STEPS)356 357 return AuditObservation(358 current_view=self._current_view,359 visible_data=visible,360 scratchpad=list(self._scratchpad),361 available_actions=AVAILABLE_ACTIONS,362 anomaly_alerts=alerts,363 investigation_progress=progress,364 )365 366 def _get_visible_data(self) -> Dict[str, Any]:367 if self._company_data is None:368 return {}369 cd = self._company_data370 view = self._current_view371 372 if view == "income_statement":373 data = cd.income_statement.model_dump()374 # Upgrade 3: enrich with narrative description for unstructured reasoning375 data["_narrative_summary"] = (376 f"Revenue of ${data['revenue']/1e6:.1f}M with operating margin "377 f"{data['operating_margin']*100:.1f}%. Net income ${data['net_income']/1e6:.1f}M."378 )379 return data380 elif view == "balance_sheet":381 data = cd.balance_sheet.model_dump()382 data["_narrative_summary"] = (383 f"Total assets ${data['total_assets']/1e6:.1f}M. "384 f"Accounts receivable ${data['accounts_receivable']/1e6:.1f}M "385 f"({data['accounts_receivable']/data['total_assets']*100:.1f}% of assets). "386 f"Equity ${data['equity']/1e6:.1f}M."387 )388 return data389 elif view == "cashflow":390 data = cd.cashflow.model_dump()391 ocf = data["operating_cash_flow"]392 data["_narrative_summary"] = (393 f"Operating cash flow ${ocf/1e6:.1f}M "394 f"({'positive — healthy cash generation' if ocf > 0 else 'NEGATIVE — cash burn despite reported profits'}). "395 f"Investing ${data['investing_cash_flow']/1e6:.1f}M, "396 f"Financing ${data['financing_cash_flow']/1e6:.1f}M."397 )398 return data399 elif view == "narrative":400 data = cd.narrative.model_dump()401 # Upgrade 3: add a flat text summary of all chunks for easier agent parsing402 data["_full_text"] = " | ".join(403 f"[{c['chunk_id']}] {c['text']}" for c in data.get("chunks", [])404 )405 return data406 elif view in ("general_ledger", "sub_ledger"):407 results = getattr(self, "_visible_query_results", [])408 return {409 "rows": results,410 "row_count": len(results),411 "hint": "Use query_ledger action to filter transactions. Filters: account_code, date_from, date_to, amount_min, amount_max, vendor_id",412 }413 return {}414 415 # ------------------------------------------------------------------416 # Data loader417 # ------------------------------------------------------------------418 419 def _load_company(self, company_id: str, company_dir: Path) -> CompanyData:420 from .models import BalanceSheet, CashFlow, IncomeStatement, Narrative421 422 def load(fname: str) -> dict:423 return json.loads((company_dir / fname).read_text())424 425 name_map = {426 "company_1": "TechCorp",427 "company_2": "RetailCo",428 "company_3": "ManufacturingInc",429 "company_4": "FinanceHub",430 "company_5": "HealthcarePlus",431 }432 433 return CompanyData(434 company_id=company_id,435 company_name=name_map.get(company_id, company_id),436 income_statement=IncomeStatement(**load("income_statement.json")),437 balance_sheet=BalanceSheet(**load("balance_sheet.json")),438 cashflow=CashFlow(**load("cashflow.json")),439 narrative=Narrative(**load("narrative.json")),440 ledger_db_path=company_dir / "ledger.db",441 sub_ledger_db_path=company_dir / "sub_ledger.db",442 )443 444 445# ---------------------------------------------------------------------------446# Custom exceptions (used by FastAPI layer)447# ---------------------------------------------------------------------------448 449class EpisodeActiveError(Exception):450 pass451 452class NoActiveEpisodeError(Exception):453 pass454 455class EpisodeDoneError(Exception):456 pass457 458class CompanyNotFoundError(Exception):459 pass460