alemno-payments/app/worker.py

import asyncio
import json
import os
import statistics
from datetime import datetime, date
from typing import Any, Dict, List
from sqlalchemy import select
from celery import Celery
from pydantic import BaseModel, Field
from google import genai
from app.core.config import settings
from app.core.logging import logger
from app.db.session import AsyncSessionLocal
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

exchange_rate_client = ExchangeRateClient()


class TransactionClassification(BaseModel):
    txn_id: str = Field(
        description="The unique transaction ID (txn_id) from the input."
    )
    category: str = Field(
        description="Assigned category, must be one of: Food, Shopping, Travel, Transport, Utilities, Cash Withdrawal, Entertainment, or Other."
    )


class BatchClassificationResponse(BaseModel):
    classifications: List[TransactionClassification]


async def classify_transactions_batch(
    client: genai.Client,
    batch_txns: List[Dict[str, Any]],
) -> Dict[str, Dict[str, Any]]:
    """
    Calls the Gemini model to classify a batch of transactions.
    Implements up to 3 retries (4 attempts total) with exponential backoff.
    Returns a dictionary mapping txn_id to classification results.
    """
    # Format the input for the LLM
    prompt_data = []
    for t in batch_txns:
        prompt_data.append(
            {
                "txn_id": t["txn_id"],
                "merchant": t["merchant"],
                "amount": t["amount"],
                "currency": t.get("currency", "INR"),
            }
        )

    input_prompt = (
        "Classify the following financial transactions. For each transaction, match its txn_id "
        "and assign one of the allowed categories: Food, Shopping, Travel, Transport, Utilities, "
        "Cash Withdrawal, Entertainment, or Other based on the merchant name and details:\n\n"
        f"{json.dumps(prompt_data, indent=2)}"
    )

    for attempt in range(1, 5):
        try:
            logger.info(
                f"Sending batch of {len(batch_txns)} transactions to LLM (Attempt {attempt})..."
            )
            loop = asyncio.get_running_loop()
            # Wrap synchronous genai SDK call in executor
            interaction = await loop.run_in_executor(
                None,
                lambda: client.interactions.create(
                    model="gemini-3.5-flash",
                    input=input_prompt,
                    response_format={
                        "type": "text",
                        "mime_type": "application/json",
                        "schema": BatchClassificationResponse.model_json_schema(),
                    },
                ),
            )

            resp_obj = BatchClassificationResponse.model_validate_json(
                interaction.output_text
            )

            results = {}
            for item in resp_obj.classifications:
                item_json = json.dumps(item.model_dump())
                results[item.txn_id] = {
                    "category": item.category,
                    "raw_response": item_json,
                    "failed": False,
                }
            return results
        except Exception as e:
            logger.warning(f"LLM Classification attempt {attempt} failed: {e}")
            if attempt == 4:
                break
            sleep_seconds = 2**attempt
            await asyncio.sleep(sleep_seconds)

    logger.error(
        "LLM Classification completely failed after all retries. Falling back to default."
    )
    results = {}
    for t in batch_txns:
        results[t["txn_id"]] = {
            "category": "Other",
            "raw_response": None,
            "failed": True,
        }
    return results


# Define Celery app
celery_app = Celery(
    "alemno_worker",
    broker=settings.CELERY_BROKER_URL,
    backend=settings.CELERY_RESULT_BACKEND,
)

# Optional configurations
celery_app.conf.update(
    task_serializer="json",
    accept_content=["json"],
    result_serializer="json",
    timezone="UTC",
    enable_utc=True,
)


def run_async(coro):
    """Helper to run async functions synchronously in Celery worker context"""
    return asyncio.run(coro)


@celery_app.task(name="tasks.process_transaction_job")
def process_transaction_job(job_id: str, transactions: List[Dict[str, Any]]):
    """
    Background job to clean transactions, detect anomalies, calculate spend breakdown,
    generate narrative summary, and persist results to Transaction and JobSummary tables.
    """
    logger.info(f"Starting Celery task to process job_id: {job_id}")
    return run_async(_process_job_async(job_id, transactions))


async def _process_job_async(job_id: str, transactions: List[Dict[str, Any]]):
    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

        try:
            # Update status to processing
            await repo.update(job, status="processing")
            await db.commit()

            db_transactions = []
            total_spend_inr = 0.0
            total_spend_usd = 0.0
            merchant_spends: Dict[str, float] = {}
            anomaly_count = 0

            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)

            # 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)

            # 2. Call LLM to classify transactions without a category
            # Find all uncategorized incoming transactions
            uncategorized_txns = []
            for t in transactions:
                cat = t.get("category")
                if not cat or cat.strip().lower() in {"", "general", "other", "none"}:
                    uncategorized_txns.append(t)

            llm_results = {}
            api_key = settings.GEMINI_API_KEY or os.environ.get("GEMINI_API_KEY")
            if uncategorized_txns:
                if api_key:
                    logger.info(
                        f"Found {len(uncategorized_txns)} uncategorized transactions. Initializing GenAI client..."
                    )
                    try:
                        client = genai.Client(api_key=api_key)
                        # Batch transactions in chunks of 50
                        batch_size = 50
                        for i in range(0, len(uncategorized_txns), batch_size):
                            batch = uncategorized_txns[i : i + batch_size]
                            batch_results = await classify_transactions_batch(
                                client, batch
                            )
                            llm_results.update(batch_results)
                    except Exception as e:
                        logger.error(
                            f"Failed to initialize GenAI Client or classify: {e}"
                        )
                        # Fallback for all
                        for t in uncategorized_txns:
                            llm_results[t["txn_id"]] = {
                                "category": "Other",
                                "raw_response": None,
                                "failed": True,
                            }
                else:
                    logger.warning(
                        "GEMINI_API_KEY is not set. Defaulting to local mock categorization."
                    )
                    # Mock categorization mapping based on merchant name keywords
                    for t in uncategorized_txns:
                        merchant_lower = t["merchant"].lower()
                        cat = "Other"
                        if any(
                            kw in merchant_lower
                            for kw in ["swiggy", "zomato", "restaurant", "cafe", "food"]
                        ):
                            cat = "Food"
                        elif any(
                            kw in merchant_lower
                            for kw in ["ola", "uber", "taxi", "metro", "train"]
                        ):
                            cat = "Transport"
                        elif any(
                            kw in merchant_lower
                            for kw in [
                                "amazon",
                                "flipkart",
                                "shopping",
                                "store",
                                "supermarket",
                            ]
                        ):
                            cat = "Shopping"
                        elif any(
                            kw in merchant_lower
                            for kw in ["air", "flight", "hotel", "travel", "irctc"]
                        ):
                            cat = "Travel"
                        elif any(
                            kw in merchant_lower
                            for kw in [
                                "electric",
                                "power",
                                "water",
                                "bill",
                                "utility",
                                "utilities",
                            ]
                        ):
                            cat = "Utilities"
                        elif any(
                            kw in merchant_lower for kw in ["atm", "cash", "withdrawal"]
                        ):
                            cat = "Cash Withdrawal"
                        elif any(
                            kw in merchant_lower
                            for kw in [
                                "netflix",
                                "spotify",
                                "movie",
                                "show",
                                "entertainment",
                            ]
                        ):
                            cat = "Entertainment"

                        llm_results[t["txn_id"]] = {
                            "category": cat,
                            "raw_response": f'{{"category": "{cat}", "confidence": 0.8, "status": "mocked_offline"}}',
                            "failed": False,
                        }

            for t in transactions:
                txn_id = t["txn_id"]
                raw_date = t["date"]
                # Parse date string to datetime.date object
                parsed_date = (
                    date.fromisoformat(raw_date)
                    if isinstance(raw_date, str)
                    else raw_date
                )
                merchant = t["merchant"].strip()
                amount = float(t["amount"])
                currency = t.get("currency", "INR").strip().upper()
                category = t.get("category") or ""
                account_id = t.get("account_id")

                # Spend calculation and currency conversion
                if currency == "USD":
                    amt_usd = amount
                    amt_inr = amount * USD_TO_INR
                else:
                    amt_inr = amount
                    amt_usd = amount / USD_TO_INR

                total_spend_inr += amt_inr
                total_spend_usd += amt_usd

                # Aggregate merchant spend (in USD for standardization)
                merchant_spends[merchant] = merchant_spends.get(merchant, 0.0) + amt_usd

                # Anomaly detection checks
                reasons = []

                # A. 3x Median Outlier Check
                if account_id and account_id in account_medians_inr:
                    median_inr = account_medians_inr[account_id]
                    if amt_inr > 3.0 * median_inr:
                        reasons.append(
                            "Statistical outlier (amount exceeds 3x account median)"
                        )

                # B. Domestic Brand transacted in USD check
                domestic_brands = {"swiggy", "ola", "irctc"}
                merchant_lower = merchant.lower()
                if currency == "USD" and any(
                    brand in merchant_lower for brand in domestic_brands
                ):
                    reasons.append("Domestic brand transacted in USD")

                # C. High value transaction check
                if amt_usd > 5000.0:
                    reasons.append("High value transaction (> $5,000 USD equivalent)")

                # D. Suspicious keyword check
                desc_lower = merchant.lower()
                if any(
                    k in desc_lower
                    for k in ["suspend", "hack", "error", "unknown", "fraud"]
                ):
                    reasons.append("Suspicious merchant keyword")

                is_anomaly = len(reasons) > 0
                anomaly_reason = "; ".join(reasons) if is_anomaly else None

                if is_anomaly:
                    anomaly_count += 1

                # Check if it was uncategorized initially and assign LLM category
                cat_lower = category.strip().lower()
                is_uncategorized = not cat_lower or cat_lower in {
                    "",
                    "general",
                    "other",
                    "none",
                }

                if is_uncategorized and txn_id in llm_results:
                    classification = llm_results[txn_id]
                    llm_category = classification["category"].strip()
                    llm_raw_response = classification["raw_response"]
                    llm_failed = classification["failed"]
                    category = llm_category
                else:
                    llm_category = category.strip() if category else "Other"
                    llm_raw_response = None
                    llm_failed = False

                db_txn = Transaction(
                    job_id=job_id,
                    txn_id=txn_id,
                    date=parsed_date,
                    merchant=merchant,
                    amount=amount,
                    currency=currency,
                    status="cleaned",
                    category=category,
                    account_id=account_id,
                    is_anomaly=is_anomaly,
                    anomaly_reason=anomaly_reason,
                    llm_category=llm_category,
                    llm_raw_response=llm_raw_response,
                    llm_failed=llm_failed,
                )
                db_transactions.append(db_txn)

            await repo.add_transactions(db_transactions)

            top_merchant = (
                max(merchant_spends, key=merchant_spends.get)
                if merchant_spends
                else "None"
            )
            risk_level = "Low"
            if anomaly_count > 2:
                risk_level = "High"
            elif anomaly_count > 0:
                risk_level = "Medium"

            narrative = (
                f"Successfully parsed and clean-processed {len(transactions)} transaction records. "
                f"Total spend is INR {total_spend_inr:,.2f} (${total_spend_usd:,.2f} USD). "
                f"Top merchant by spend volume is '{top_merchant}' with total of ${merchant_spends.get(top_merchant, 0):,.2f} USD. "
                f"Scan flagged {anomaly_count} potential anomalies, resulting in a overall risk level of '{risk_level}'."
            )

            # Build top merchants dict sorted by spend limit (keep top 5)
            sorted_merchants = sorted(
                merchant_spends.items(), key=lambda x: x[1], reverse=True
            )[:5]
            top_merchants_json = {
                merchant: round(spend, 2) for merchant, spend in sorted_merchants
            }

            # 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)

            # 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:
            logger.exception(f"Error processing job {job_id}")
            await repo.update(
                job,
                status="failed",
                error_message=str(e),
                completed_at=datetime.utcnow(),
            )
            await db.commit()
            raise e