feat: Implemented observability layer and benchmarking

app/worker.py

@@ -15,6 +15,7 @@ from app.repositories.job_repository import JobRepository
 from app.db.models.transaction import Transaction
 from app.db.models.job_summary import JobSummary
 from app.clients.exchange_rate_client import ExchangeRateClient
+from app.core.observability import time_it, benchmark
 
 exchange_rate_client = ExchangeRateClient()
 
@@ -53,6 +54,7 @@ class BatchClassificationResponse(BaseModel):
     classifications: List[TransactionClassification]
 
 
+@time_it("LLM Batch Classification Request")
 async def classify_transactions_batch(
     client: genai.Client,
     batch_txns: List[Dict[str, Any]],
@@ -160,6 +162,7 @@ class JobNarrativeSummary(BaseModel):
     )
 
 
+@time_it("LLM Job Narrative Summary Request")
 async def generate_job_summary_llm(
     client: genai.Client,
     db_transactions: List[Transaction],
@@ -253,22 +256,26 @@ def process_transaction_job(job_id: str, transactions: List[Dict[str, Any]]):
     return run_async(_process_job_async(job_id, transactions))
 
 
+@time_it("Process Job Background Task")
 async def _process_job_async(job_id: str, transactions: List[Dict[str, Any]]):
     from app.db.session import engine
 
-    await engine.dispose()
-    async with AsyncSessionLocal() as db:
-        repo = JobRepository(db)
-        job = await repo.get_by_id(job_id)
-        if not job:
-            logger.error(f"Job with ID {job_id} not found in database.")
-            return
+    with benchmark("Dispose Engine"):
+        await engine.dispose()
+
+    async with AsyncSessionLocal() as db:
+        with benchmark("Initial Job Setup"):
+            repo = JobRepository(db)
+            job = await repo.get_by_id(job_id)
+            if not job:
+                logger.error(f"Job with ID {job_id} not found in database.")
+                return
 
-        try:
             # Update status to processing
             await repo.update(job, status="processing")
             await db.commit()
 
+        try:
             db_transactions = []
             total_spend_inr = 0.0
             total_spend_usd = 0.0
@@ -277,43 +284,44 @@ async def _process_job_async(job_id: str, transactions: List[Dict[str, Any]]):
 
             USD_TO_INR = await exchange_rate_client.get_usd_to_inr_rate()
 
-            # 1. Fetch existing transaction amounts and currencies for all accounts in this job to calculate median
-            account_ids = {
-                t.get("account_id") for t in transactions if t.get("account_id")
-            }
-            existing_amounts_inr = {}
-            if account_ids:
-                stmt = select(
-                    Transaction.account_id, Transaction.amount, Transaction.currency
-                ).where(Transaction.account_id.in_(account_ids))
-                res = await db.execute(stmt)
-                for acc_id, amt, curr in res.all():
-                    if acc_id not in existing_amounts_inr:
-                        existing_amounts_inr[acc_id] = []
-                    # Standardize to INR for median calculation
-                    amt_inr = amt * USD_TO_INR if curr.upper() == "USD" else amt
-                    existing_amounts_inr[acc_id].append(amt_inr)
+            with benchmark("Compute Account History Medians"):
+                # 1. Fetch existing transaction amounts and currencies for all accounts in this job to calculate median
+                account_ids = {
+                    t.get("account_id") for t in transactions if t.get("account_id")
+                }
+                existing_amounts_inr = {}
+                if account_ids:
+                    stmt = select(
+                        Transaction.account_id, Transaction.amount, Transaction.currency
+                    ).where(Transaction.account_id.in_(account_ids))
+                    res = await db.execute(stmt)
+                    for acc_id, amt, curr in res.all():
+                        if acc_id not in existing_amounts_inr:
+                            existing_amounts_inr[acc_id] = []
+                        # Standardize to INR for median calculation
+                        amt_inr = amt * USD_TO_INR if curr.upper() == "USD" else amt
+                        existing_amounts_inr[acc_id].append(amt_inr)
 
-            # Group incoming transaction amounts standardized to INR
-            incoming_amounts_inr = {}
-            for t in transactions:
-                acc_id = t.get("account_id")
-                if acc_id:
-                    if acc_id not in incoming_amounts_inr:
-                        incoming_amounts_inr[acc_id] = []
-                    amt = float(t["amount"])
-                    curr = t.get("currency", "INR").strip().upper()
-                    amt_inr = amt * USD_TO_INR if curr == "USD" else amt
-                    incoming_amounts_inr[acc_id].append(amt_inr)
+                # Group incoming transaction amounts standardized to INR
+                incoming_amounts_inr = {}
+                for t in transactions:
+                    acc_id = t.get("account_id")
+                    if acc_id:
+                        if acc_id not in incoming_amounts_inr:
+                            incoming_amounts_inr[acc_id] = []
+                        amt = float(t["amount"])
+                        curr = t.get("currency", "INR").strip().upper()
+                        amt_inr = amt * USD_TO_INR if curr == "USD" else amt
+                        incoming_amounts_inr[acc_id].append(amt_inr)
 
-            # Precalculate median of (existing + incoming) transactions in INR
-            account_medians_inr = {}
-            for acc_id in account_ids:
-                all_amts = existing_amounts_inr.get(
-                    acc_id, []
-                ) + incoming_amounts_inr.get(acc_id, [])
-                if all_amts:
-                    account_medians_inr[acc_id] = statistics.median(all_amts)
+                # Precalculate median of (existing + incoming) transactions in INR
+                account_medians_inr = {}
+                for acc_id in account_ids:
+                    all_amts = existing_amounts_inr.get(
+                        acc_id, []
+                    ) + incoming_amounts_inr.get(acc_id, [])
+                    if all_amts:
+                        account_medians_inr[acc_id] = statistics.median(all_amts)
 
             # 2. Call LLM to classify transactions without a category
             # Find all uncategorized incoming transactions
@@ -579,7 +587,8 @@ async def _process_job_async(job_id: str, transactions: List[Dict[str, Any]]):
                 )
                 db_transactions.append(db_txn)
 
-            await repo.add_transactions(db_transactions)
+            with benchmark("Save Cleaned Transactions to DB"):
+                await repo.add_transactions(db_transactions)
 
             # Generate narrative and risk summary using LLM if api_key is available
             llm_summary = None
@@ -635,26 +644,27 @@ async def _process_job_async(job_id: str, transactions: List[Dict[str, Any]]):
                 }
 
             # Create job summary
-            summary = JobSummary(
-                job_id=job_id,
-                total_spend_inr=round(total_spend_inr, 2),
-                total_spend_usd=round(total_spend_usd, 2),
-                top_merchants=top_merchants_json,
-                anomaly_count=anomaly_count,
-                narrative=narrative,
-                risk_level=risk_level,
-            )
-            await repo.add_summary(summary)
+            with benchmark("Save Summary and Complete Job"):
+                summary = JobSummary(
+                    job_id=job_id,
+                    total_spend_inr=round(total_spend_inr, 2),
+                    total_spend_usd=round(total_spend_usd, 2),
+                    top_merchants=top_merchants_json,
+                    anomaly_count=anomaly_count,
+                    narrative=narrative,
+                    risk_level=risk_level,
+                )
+                await repo.add_summary(summary)
 
-            # Update job state
-            await repo.update(
-                job,
-                status="completed",
-                row_count_raw=len(transactions),
-                row_count_clean=len(db_transactions),
-                completed_at=datetime.utcnow(),
-            )
-            await db.commit()
+                # Update job state
+                await repo.update(
+                    job,
+                    status="completed",
+                    row_count_raw=len(transactions),
+                    row_count_clean=len(db_transactions),
+                    completed_at=datetime.utcnow(),
+                )
+                await db.commit()
             logger.info(f"Job {job_id} processing completed successfully.")
 
         except Exception as e: