# """ from __future__ import annotations import time import uuid from dataclasses import dataclass, field from typing import Any, Optional from .schemas import TraceAction, TraceObservation, EpisodeState from .world_model import SemanticWorldModel from ..agents.planner import PlannerAgent from ..agents.retriever import RetrieverAgent from ..agents.memory import MemoryAgent from ..agents.verifier import VerifierAgent from ..rewards.reward_fn import compute_reward from ..rewards.anti_hack import AntiHackGuard class TraceEnv: """ OpenEnv-compatible environment for the Trace project. Episode lifecycle: reset(task) -> observation step(action) -> observation, reward, done, info state() -> EpisodeState (for logging/debugging) The environment simulates a user with a federated digital footprint. The agent is given a long-horizon instruction (e.g. "Audit all receipts from 2022-2024") and must plan sub-tasks, retrieve data from virtual sources, and synthesize a verified result. """ DEFAULT_MAX_STEPS = 20 # hard limit per episode DEFAULT_TIMEOUT_SECONDS = 300 # wall-clock timeout (per-step) def __init__(self, config: dict): self.config = config self.max_steps = config.get("max_steps", self.DEFAULT_MAX_STEPS) self.timeout_seconds = config.get("timeout_seconds", self.DEFAULT_TIMEOUT_SECONDS) self.world_model = SemanticWorldModel(config) self.planner = PlannerAgent(config) self.retriever = RetrieverAgent(config) self.memory = MemoryAgent(config) self.verifier = VerifierAgent(config) self.anti_hack = AntiHackGuard() self._episode_id: Optional[str] = None self._steps: int = 0 self._step_start_time: float = 0.0 self._state: Optional[EpisodeState] = None # ------------------------------------------------------------------ # OpenEnv interface # ------------------------------------------------------------------ def reset(self, task: dict) -> TraceObservation: """ Start a fresh episode with a new task. Args: task: { "instruction": str, # natural-language goal "difficulty": "easy"|"medium"|"hard", "available_sources": list[str], # e.g. ["gmail", "drive"] "ground_truth": dict, # for reward computation } Returns: TraceObservation: initial observation for the agent. """ self._episode_id = str(uuid.uuid4()) self._steps = 0 self._step_start_time = time.time() self.world_model.initialize(task) self.memory.reset() self.anti_hack.reset() self._state = EpisodeState( episode_id=self._episode_id, task=task, plan=[], retrieved_data=[], verified=False, steps=0, done=False, ) obs = TraceObservation( episode_id=self._episode_id, step=0, instruction=task["instruction"], available_sources=task["available_sources"], context="", memory_summary=self.memory.summarize(), world_state=self.world_model.snapshot(), ) return obs def step(self, action: TraceAction) -> tuple[TraceObservation, float, bool, dict]: """ Execute one agent action and return the next state. Action types: - PLAN: decompose instruction into sub-tasks - RETRIEVE: fetch data from a virtual source - MEMORIZE: store a finding into episodic memory - VERIFY: verify the current plan/answer against world model - ANSWER: submit the final synthesized answer Returns: (observation, reward, done, info) """ assert self._state is not None, "Call reset() before step()" self._steps += 1 self._state.steps = self._steps # ── Timeout / step-limit guards ───────────────────────────────── # Per-step timeout: only the current step's processing time matters, # not idle time between API calls. self._step_start_time = time.time() if self._steps > self.max_steps: return self._terminate(reason="max_steps") # ── Anti-hack validation ───────────────────────────────────────── hack_flag = self.anti_hack.check(action) if hack_flag: reward = compute_reward( action, self._state, hack_penalty=True ) info_dict = {"hack": hack_flag} obs = self._build_obs(f"[ANTI-HACK] {hack_flag}", metadata=info_dict) return obs, reward, False, info_dict # ── Dispatch action ────────────────────────────────────────────── result_context = "" action_type = action.action_type.strip().upper() if action_type == "PLAN": plan = self.planner.decompose( action.content, self._state.task ) self._state.plan = plan result_context = f"Plan created: {plan}" elif action_type == "RETRIEVE": data = self.retriever.fetch( source=action.source, query=action.content, world_model=self.world_model, metadata=action.metadata, ) if not isinstance(data, list): data = [data] # Inject real data into world model so visible_items updates self.world_model.inject_real_data(action.source, data) self._state.retrieved_data.extend(data) result_context = f"Retrieved {len(data)} items from {action.source}" info = {} # ── Gmail processing: merge and summarize all retrieved transactions ── if action.source == "gmail" and data: try: from ..tools.transaction_parser import parse_transactions_bulk from ..tools.dashboard_renderer import render_dashboard # We parse everything we've retrieved so far to ensure deduplication # and a cumulative summary. parsed_all = parse_transactions_bulk(self._state.retrieved_data) summary = parsed_all.get("summary", {}) transactions = parsed_all.get("transactions", []) total_spend = summary.get("total_spend", 0.0) tx_count = summary.get("count", 0) by_category = summary.get("by_category", {}) top_category = next(iter(by_category.keys()), "unknown") top_category_spend = by_category.get(top_category, 0.0) dashboard_html = render_dashboard(parsed_all) result_context = ( f"Step summary: Retrieved {len(data)} new items. " f"Cumulative Audit: {tx_count} total transactions | " f"Total Spend: ₹{total_spend:,.2f} | " f"Top Category: {top_category} (₹{top_category_spend:,.2f})" ) info = { "gmail_query": action.content, "transactions_summary": summary, "transactions": transactions, "dashboard_html": dashboard_html, "dashboard_generated": True, "cumulative_count": tx_count, "cumulative_spend": total_spend, } except Exception as e: info = { "dashboard_generated": False, "dashboard_error": str(e), } elif action.source == "sheets": try: from ..tools.sheets_tool import fetch_and_summarize from ..tools.transaction_parser import parse_transactions_bulk summary = fetch_and_summarize() sheet_txs = summary.get("transactions", []) # Deduplicate before extending: only add Sheets rows # whose IDs are not already present from Gmail retrieval existing_ids = { item.get("id") for item in self._state.retrieved_data if item.get("id") } # Keep all sheet transactions; parse_transactions_bulk will handle merging self._state.retrieved_data.extend(sheet_txs) # Build cumulative summary from ALL retrieved data (Gmail + Sheets) parsed_all = parse_transactions_bulk(self._state.retrieved_data) summary = parsed_all.get("summary", {}) all_txs = parsed_all.get("transactions", []) # Calculate overlapping items for logging gmail_ids = { item.get("id") for item in self._state.retrieved_data if item.get("_source") != "sheets" and item.get("id") } overlapping = sum(1 for tx in sheet_txs if tx.get("id") in gmail_ids) result_context = ( f"Retrieved {len(sheet_txs)} items from Google Sheets " f"({len(sheet_txs) - overlapping} new, {overlapping} already in Gmail). " f"Cumulative Audit: {summary.get('count', 0)} total transactions | " f"Total Spend: ₹{summary.get('total_spend', 0.0):,.2f}" ) info = { "source": "sheets", "sheets_count": len(sheet_txs), "new_from_sheets": len(sheet_txs) - overlapping, "transactions": all_txs, # merged Gmail + Sheets "transactions_summary": summary, } except Exception as e: result_context = f"Error retrieving from Sheets: {e}" info = {"error": str(e)} else: info = {} elif action_type == "MEMORIZE": self.memory.store(action.content, action.metadata) result_context = "Stored to episodic memory." elif action_type == "VERIFY": verification = self.verifier.verify( claim=action.content, world_model=self.world_model, memory=self.memory, ) self._state.verified = verification["passed"] result_context = f"Verification: {verification}" elif action_type == "SYNC": # Sync retrieved transactions to Google Sheets try: from ..tools.sheets_tool import append_transactions, fetch_and_summarize # Get transactions from the current state (parsed if available in info) # In a real scenario, we'd pull from world model or previous info. # For the demo, we'll use the last retrieved transactions if they exist. # However, retriever already puts data into world_model. # We need to get the parsed transactions. # Let's assume we want to sync whatever we found in the last Gmail pass. # But a cleaner way is to sync all transactions in the world model. all_tx = [] # In a real implementation, we'd query the world model for all transactions. # For now, let's use the retrieved_data directly if it looks like Gmail data. # Or better, the env maintains a list of parsed transactions. from ..tools.transaction_parser import parse_transactions_bulk parsed = parse_transactions_bulk(self._state.retrieved_data) transactions = parsed.get("transactions", []) sheet_url = append_transactions(transactions) if sheet_url: # After sync, fetch the full summary to verify ledger_summary = fetch_and_summarize() total_ledger = ledger_summary.get("total_spend", 0.0) result_context = ( f"Synced {len(transactions)} transactions to Google Sheets: {sheet_url}. " f"Current Ledger Total: ₹{total_ledger:,.2f}" ) info = { "sheet_url": sheet_url, "ledger_summary": ledger_summary, "sync_count": len(transactions) } else: result_context = "Failed to sync to Google Sheets. Check credentials." info = {"error": "Sync failed"} except Exception as e: result_context = f"Error during SYNC: {e}" info = {"error": str(e)} elif action_type == "EXPORT": # Export retrieved transactions to a DOCX report try: from ..tools.transaction_parser import parse_transactions_bulk from ..tools.report_tool import export_transactions_to_docx # Use all currently retrieved data (Gmail + Sheets) parsed = parse_transactions_bulk(self._state.retrieved_data) transactions = parsed.get("transactions", []) report_path = export_transactions_to_docx(transactions) result_context = ( f"Exported {len(transactions)} transactions to DOCX report at: {report_path}." ) info = { "report_path": report_path, "export_count": len(transactions) } except Exception as e: result_context = f"Error during EXPORT: {e}" info = {"error": str(e)} elif action_type == "ANSWER": self._state.final_answer = action.content reward = compute_reward(action, self._state) info_dict = {"final_answer": action.content} obs = self._build_obs("Episode complete.", metadata=info_dict) self._state.done = True return obs, reward, True, info_dict else: result_context = f"Unknown action type: {action.action_type}" # ── Intermediate reward & next observation ─────────────────────── info_dict = info if "info" in locals() else {} reward = compute_reward(action, self._state) obs = self._build_obs(result_context, metadata=info_dict) return obs, reward, False, info_dict def state(self) -> EpisodeState: """Return full episode state (for logging/debugging).""" return self._state # ------------------------------------------------------------------ # Internal helpers # ------------------------------------------------------------------ def _build_obs(self, context: str, metadata: Optional[dict] = None) -> TraceObservation: return TraceObservation( episode_id=self._episode_id, step=self._steps, instruction=self._state.task["instruction"], available_sources=self._state.task["available_sources"], context=context, memory_summary=self.memory.summarize(), world_state=self.world_model.snapshot(), metadata=metadata or {}, ) def _terminate(self, reason: str) -> tuple[TraceObservation, float, bool, dict]: info_dict = {"termination_reason": reason} obs = self._build_obs(f"Episode terminated: {reason}", metadata=info_dict) reward = compute_reward(None, self._state, terminal_penalty=True) self._state.done = True return obs, reward, True, info_dict