Spaces:
Sleeping
Sleeping
| """ | |
| Trace Environment Inference Script. | |
| Runs a simple hardcoded policy against the Trace environment server | |
| to demonstrate the full episode lifecycle. | |
| Usage: | |
| # Make sure server is running first: | |
| # cd trace && uvicorn server.app:app --host 127.0.0.1 --port 8000 | |
| python inference.py --url http://127.0.0.1:8000 | |
| """ | |
| import logging | |
| import sys | |
| import os | |
| # Ensure project root is importable | |
| sys.path.insert(0, os.path.dirname(__file__)) | |
| from client import TraceEnvClient | |
| from models import TraceAction | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger(__name__) | |
| def run_inference(base_url: str = "http://localhost:8000"): | |
| """ | |
| Run a full inference episode against the Trace environment. | |
| """ | |
| logger.info("Connecting to Trace environment...") | |
| with TraceEnvClient(base_url=base_url).sync() as client: | |
| # 1. Reset | |
| logger.info("Step 0: Resetting environment...") | |
| result = client.reset() | |
| obs = result.observation | |
| logger.info(f" Instruction: {obs.instruction}") | |
| logger.info(f" Sources: {obs.available_sources}") | |
| # 2. Plan | |
| logger.info("Step 1: Planning...") | |
| result = client.step(TraceAction( | |
| action_type="PLAN", | |
| content=( | |
| "Step 1: Fast text-first retrieval for transactional emails. " | |
| "Step 2: Deep retrieval for attachment-heavy mails (PDF/image/doc). " | |
| "Step 3: Retrieve historical transactions from Google Sheets ledger. " | |
| "Step 4: Export merged data into a local DOCX report." | |
| ) | |
| )) | |
| logger.info(f" Context: {result.observation.context}") | |
| # 3. Quick Retrieval (Text-first) | |
| logger.info("Step 2: Quick retrieval (Text-first)...") | |
| result = client.step(TraceAction( | |
| action_type="RETRIEVE", | |
| content="newer_than:180d (receipt OR invoice OR payment OR transaction OR order OR bill OR booking OR ride OR food)", | |
| source="gmail", | |
| metadata={"pass": "quick_text", "max_results": 50} | |
| )) | |
| logger.info(f" Context: {result.observation.context}") | |
| # 4. Deep Retrieval (Attachments) | |
| logger.info("Step 3: Deep retrieval (Attachments)...") | |
| result = client.step(TraceAction( | |
| action_type="RETRIEVE", | |
| content="newer_than:180d (receipt OR invoice OR payment OR transaction OR order OR bill OR booking OR ride OR food)", | |
| source="gmail", | |
| metadata={"pass": "deep_attachments", "max_results": 50, "analyse_images": True} | |
| )) | |
| logger.info(f" Context: {result.observation.context}") | |
| # Track Gmail-only summary for provenance | |
| gmail_summary = result.observation.metadata.get("transactions_summary", {}) | |
| gmail_count = gmail_summary.get("count", 0) | |
| gmail_spend = gmail_summary.get("total_spend", 0.0) | |
| gmail_by_cat = gmail_summary.get("by_category", {}) | |
| logger.info(f" Gmail data: {gmail_count} transactions, ₹{gmail_spend:,.2f}") | |
| logger.info(f" Gmail categories: {gmail_by_cat}") | |
| # 5. RETRIEVE from Google Sheets (gets historical data) | |
| # The env merges Gmail retrieved_data with Sheets data automatically | |
| logger.info("Step 4: Retrieving historical data from Google Sheets...") | |
| result = client.step(TraceAction( | |
| action_type="RETRIEVE", | |
| content="Read all existing transactions from the ledger.", | |
| source="sheets" | |
| )) | |
| logger.info(f" Context: {result.observation.context}") | |
| # Extract MERGED summary (Gmail + Sheets combined by the env) | |
| summary = result.observation.metadata.get("transactions_summary", {}) | |
| 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) | |
| sheets_new = result.observation.metadata.get("new_from_sheets", 0) | |
| logger.info(f" Merged: {tx_count} total ({gmail_count} Gmail + {sheets_new} Sheets-only)") | |
| logger.info(f" Combined categories: {by_category}") | |
| # 6. EXPORT all transactions to local DOCX | |
| logger.info("Step 5: Exporting transactions to DOCX report...") | |
| result = client.step(TraceAction( | |
| action_type="EXPORT", | |
| content="Export the combined financial transactions into a DOCX report." | |
| )) | |
| logger.info(f" Context: {result.observation.context}") | |
| # 7. Memorize findings | |
| logger.info("Step 6: Memorizing...") | |
| result = client.step(TraceAction( | |
| action_type="MEMORIZE", | |
| content=( | |
| f"Consolidated {tx_count} transactions (Gmail: {gmail_count}, " | |
| f"Sheets-only: {sheets_new}). Total: ₹{total_spend:,.2f}. " | |
| f"Categories: {', '.join(f'{k}: ₹{v:,.2f}' for k, v in by_category.items())}. " | |
| f"Report generated locally." | |
| ) | |
| )) | |
| # 8. Submit answer — includes both data sources | |
| logger.info("Step 7: Submitting answer...") | |
| cat_breakdown = ", ".join(f"{k}: ₹{v:,.2f}" for k, v in list(by_category.items())[:5]) | |
| result = client.step(TraceAction( | |
| action_type="ANSWER", | |
| content=( | |
| f"Total spend: ₹{total_spend:,.2f} across {tx_count} transactions " | |
| f"(Gmail: {gmail_count}, Sheets historical: {sheets_new}). " | |
| f"Top category: {top_category} (₹{top_category_spend:,.2f}). " | |
| f"Breakdown: {cat_breakdown}. " | |
| "Summary: Financial audit complete — Gmail + Sheets data merged and local DOCX report generated." | |
| ) | |
| )) | |
| logger.info(f" Final Reward: {result.reward}") | |
| logger.info(f" Done: {result.done}") | |
| return result | |
| if __name__ == "__main__": | |
| import argparse | |
| parser = argparse.ArgumentParser(description="Trace Environment Inference") | |
| parser.add_argument("--url", type=str, default="http://localhost:8000") | |
| args = parser.parse_args() | |
| run_inference(base_url=args.url) | |