feat: Implemented observability layer and benchmarking

app/clients/exchange_rate_client.py

@@ -4,6 +4,9 @@ from app.core.config import settings
 from app.core.logging import logger
 
 
+from app.core.observability import time_it
+
+
 class ExchangeRateClient:
     def __init__(self):
         self.default_rate = 93.0
@@ -20,6 +23,7 @@ class ExchangeRateClient:
             logger.warning(f"Could not connect to Redis for exchange rate caching: {e}")
             self.redis_client = None
 
+    @time_it("Get USD to INR Exchange Rate")
     async def get_usd_to_inr_rate(self) -> float:
         """
         Retrieves the USD to INR conversion rate.

app/core/observability.py

@@ -0,0 +1,59 @@
+import functools
+import inspect
+import time
+from contextlib import contextmanager
+from typing import Optional
+from app.core.logging import logger
+
+
+@contextmanager
+def benchmark(name: str):
+    """
+    Context manager to benchmark a block of code.
+    Example:
+        with benchmark("Database save"):
+            await repo.save(data)
+    """
+    start_time = time.perf_counter()
+    try:
+        yield
+    finally:
+        elapsed = time.perf_counter() - start_time
+        logger.info(f"[BENCHMARK] {name} completed in {elapsed:.4f} seconds")
+
+
+def time_it(name: Optional[str] = None):
+    """
+    Decorator to benchmark synchronous or asynchronous functions.
+    Example:
+        @time_it("Retrieve exchange rate")
+        async def get_rate():
+            ...
+    """
+
+    def decorator(func):
+        label = name or func.__name__
+
+        @functools.wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            start = time.perf_counter()
+            try:
+                return await func(*args, **kwargs)
+            finally:
+                elapsed = time.perf_counter() - start
+                logger.info(f"[BENCHMARK] {label} completed in {elapsed:.4f} seconds")
+
+        @functools.wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            start = time.perf_counter()
+            try:
+                return func(*args, **kwargs)
+            finally:
+                elapsed = time.perf_counter() - start
+                logger.info(f"[BENCHMARK] {label} completed in {elapsed:.4f} seconds")
+
+        if inspect.iscoroutinefunction(func):
+            return async_wrapper
+        return sync_wrapper
+
+    return decorator

app/services/job_service.py

@@ -6,10 +6,14 @@ from app.utils.csv_parser import parse_and_validate_csv
 from app.worker import process_transaction_job
 
 
+from app.core.observability import time_it
+
+
 class JobService:
     def __init__(self, repo: JobRepository):
         self.repo = repo
 
+    @time_it("Upload CSV Service")
     async def upload_csv(self, filename: str, content: bytes) -> Job:
         parsed_transactions = parse_and_validate_csv(content)
 
@@ -19,17 +23,20 @@ class JobService:
 
         return job
 
+    @time_it("Get Job Status Service")
     async def get_job_status(self, job_id: str) -> Job:
         job = await self.repo.get_by_id(job_id)
         if not job:
             raise JobNotFoundException(job_id)
         return job
 
+    @time_it("Get Job Results Service")
     async def get_job_results(self, job_id: str) -> Job:
         job = await self.repo.get_by_id(job_id)
         if not job:
             raise JobNotFoundException(job_id)
         return job
 
+    @time_it("List Jobs Service")
     async def list_jobs(self, status: Optional[str] = None) -> List[Job]:
         return await self.repo.list_all(status=status)

app/utils/csv_parser.py

@@ -4,9 +4,12 @@ from datetime import datetime
 from typing import Any, Dict, List
 from app.core.exceptions import InvalidCSVException
 
+from app.core.observability import time_it
+
 REQUIRED_COLUMNS = {"txn_id", "date", "merchant", "amount"}
 
 
+@time_it("Parse and Validate CSV")
 def parse_and_validate_csv(content: bytes) -> List[Dict[str, Any]]:
     """
     Parses CSV content from bytes, validates headers, and checks field types.

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: