Files
alemno-payments/app/worker.py

700 lines
28 KiB
Python

import asyncio
import json
import os
import statistics
from datetime import datetime, date
from typing import Any, Dict, List, Optional
from sqlalchemy import select
from celery import Celery
from pydantic import BaseModel, Field, AliasChoices
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
from app.core.observability import time_it, benchmark
exchange_rate_client = ExchangeRateClient()
def clean_json_response(text: str) -> str:
"""Extracts the first JSON object or array from the text, handling markdown fences and conversational prefix/suffix."""
text = text.strip()
start_idx = -1
for i, char in enumerate(text):
if char in ("{", "["):
start_idx = i
break
end_idx = -1
for i in range(len(text) - 1, -1, -1):
if text[i] in ("}", "]"):
end_idx = i
break
if start_idx != -1 and end_idx != -1 and end_idx > start_idx:
return text[start_idx : end_idx + 1]
return text
class TransactionClassification(BaseModel):
txn_id: str = Field(
description="The unique transaction ID (txn_id) from the input."
)
category: str = Field(
description="Assigned category matching one of the allowed categories list."
)
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]],
allowed_categories: List[str],
) -> 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 "
f"and assign one of the allowed categories: {', '.join(allowed_categories)} 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-2.5-flash-lite",
input=input_prompt,
response_format={
"type": "text",
"mime_type": "application/json",
"schema": BatchClassificationResponse.model_json_schema(),
},
),
)
clean_text = clean_json_response(interaction.output_text)
resp_obj = BatchClassificationResponse.model_validate_json(clean_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
class JobNarrativeSummary(BaseModel):
total_spend_by_currency: Dict[str, float] = Field(
validation_alias=AliasChoices(
"total_spend_by_currency", "spend_by_currency", "total_spend"
),
description="Total spend grouped by original currency",
)
top_3_merchants: Dict[str, float] = Field(
validation_alias=AliasChoices("top_3_merchants", "top_merchants", "top_3"),
description="Top 3 merchants by spend volume (value in USD equivalent)",
)
anomaly_count: int = Field(
validation_alias=AliasChoices(
"anomaly_count", "total_anomalies", "anomalies", "anomaly_transactions"
),
description="Total number of anomaly transactions flagged in the dataset",
)
spending_narrative: str = Field(
validation_alias=AliasChoices("spending_narrative", "narrative", "summary"),
description="A cohesive 2-3 sentence spending narrative",
)
risk_level: str = Field(
validation_alias=AliasChoices("risk_level", "overall_risk_level", "risk"),
description="Overall risk level: Low, Medium, or High",
)
@time_it("LLM Job Narrative Summary Request")
async def generate_job_summary_llm(
client: genai.Client,
db_transactions: List[Transaction],
) -> Optional[JobNarrativeSummary]:
"""
Calls the Gemini model to produce a structured JSON summary:
total spend by currency, top 3 merchants, anomaly count, a 2-3 sentence spending narrative, and a risk level.
"""
# Build a simplified list of transactions for context
tx_data = [
{
"merchant": t.merchant,
"amount": t.amount,
"currency": t.currency,
"category": t.category,
"is_anomaly": t.is_anomaly,
"anomaly_reason": t.anomaly_reason,
}
for t in db_transactions
]
input_prompt = (
"Analyze the following list of processed financial transactions and generate a structured JSON summary. "
"The returned JSON MUST contain exactly the following keys:\n"
' - "total_spend_by_currency": dictionary mapping currency symbol/code to total amount spent (e.g. {"USD": 150.0, "INR": 25000.0})\n'
' - "top_3_merchants": dictionary mapping the top 3 merchants by spend to their total spend value in USD equivalent (assuming 83 INR per USD for conversion)\n'
' - "anomaly_count": integer showing the total count of flagged anomaly transactions (is_anomaly=True)\n'
' - "spending_narrative": string containing a cohesive 2-3 sentence narrative summarizing key spending insights and largest merchants\n'
' - "risk_level": string showing the overall risk level of the job: "Low", "Medium", or "High" (pick one based on the nature of anomalies)\n\n'
f"{json.dumps(tx_data, indent=2)}"
)
for attempt in range(1, 4):
try:
logger.info(
f"Calling LLM to generate job narrative summary (Attempt {attempt})..."
)
loop = asyncio.get_running_loop()
interaction = await loop.run_in_executor(
None,
lambda: client.interactions.create(
model="gemini-2.5-flash-lite",
input=input_prompt,
response_format={
"type": "text",
"mime_type": "application/json",
"schema": JobNarrativeSummary.model_json_schema(),
},
),
)
clean_text = clean_json_response(interaction.output_text)
resp_obj = JobNarrativeSummary.model_validate_json(clean_text)
return resp_obj
except Exception as e:
logger.warning(f"Failed to generate LLM summary (Attempt {attempt}): {e}")
if attempt == 3:
break
await asyncio.sleep(2)
return None
# 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))
@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
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
# 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
merchant_spends: Dict[str, float] = {}
anomaly_count = 0
USD_TO_INR = await exchange_rate_client.get_usd_to_inr_rate()
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)
# 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:
# Load allowed categories from the database categories table
from app.db.models.category import Category
try:
res_cats = await db.execute(select(Category.name))
allowed_categories = [r[0] for r in res_cats.all()]
except Exception as db_err:
logger.warning(f"Could not load categories from database: {db_err}")
allowed_categories = []
if not allowed_categories:
allowed_categories = [
"Food",
"Shopping",
"Travel",
"Transport",
"Utilities",
"Cash Withdrawal",
"Entertainment",
"Other",
]
if api_key:
logger.info(
f"Found {len(uncategorized_txns)} uncategorized transactions. Initializing GenAI client..."
)
try:
client = genai.Client(api_key=api_key)
# Split into batches of 50 and fire all concurrently
batch_size = 50
batches = [
uncategorized_txns[i : i + batch_size]
for i in range(0, len(uncategorized_txns), batch_size)
]
logger.info(
f"Dispatching {len(batches)} classification batch(es) concurrently..."
)
with benchmark(
"LLM Concurrent Batch Classification (all batches)"
):
batch_results_list = await asyncio.gather(
*[
classify_transactions_batch(
client, batch, allowed_categories
)
for batch in batches
],
return_exceptions=True,
)
for batch_result in batch_results_list:
if isinstance(batch_result, Exception):
logger.error(
f"A classification batch failed: {batch_result}"
)
else:
llm_results.update(batch_result)
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 "Other" in allowed_categories
else (
allowed_categories[0] if allowed_categories else "Other"
)
)
if (
any(
kw in merchant_lower
for kw in [
"swiggy",
"zomato",
"restaurant",
"cafe",
"food",
]
)
and "Food" in allowed_categories
):
cat = "Food"
elif (
any(
kw in merchant_lower
for kw in ["ola", "uber", "taxi", "metro", "train"]
)
and "Transport" in allowed_categories
):
cat = "Transport"
elif (
any(
kw in merchant_lower
for kw in [
"amazon",
"flipkart",
"shopping",
"store",
"supermarket",
]
)
and "Shopping" in allowed_categories
):
cat = "Shopping"
elif (
any(
kw in merchant_lower
for kw in ["air", "flight", "hotel", "travel", "irctc"]
)
and "Travel" in allowed_categories
):
cat = "Travel"
elif (
any(
kw in merchant_lower
for kw in [
"electric",
"power",
"water",
"bill",
"utility",
"utilities",
]
)
and "Utilities" in allowed_categories
):
cat = "Utilities"
elif (
any(
kw in merchant_lower
for kw in ["atm", "cash", "withdrawal"]
)
and "Cash Withdrawal" in allowed_categories
):
cat = "Cash Withdrawal"
elif (
any(
kw in merchant_lower
for kw in [
"netflix",
"spotify",
"movie",
"show",
"entertainment",
]
)
and "Entertainment" in allowed_categories
):
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)
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
if api_key:
try:
client = genai.Client(api_key=api_key)
llm_summary = await generate_job_summary_llm(
client, db_transactions
)
except Exception as e:
logger.error(f"Error in LLM summary generation: {e}")
if llm_summary:
# Normalize risk level value to match capitalized Low/Medium/High in our DB
risk_level_str = llm_summary.risk_level.strip().lower()
if risk_level_str == "high":
risk_level = "High"
elif risk_level_str == "medium":
risk_level = "Medium"
else:
risk_level = "Low"
narrative = llm_summary.spending_narrative.strip()
top_merchants_json = {
merchant: round(spend, 2)
for merchant, spend in llm_summary.top_3_merchants.items()
}
anomaly_count = llm_summary.anomaly_count
else:
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}'."
)
sorted_merchants = sorted(
merchant_spends.items(), key=lambda x: x[1], reverse=True
)[:3]
top_merchants_json = {
merchant: round(spend, 2) for merchant, spend in sorted_merchants
}
# Create job 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()
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