mirror of
https://github.com/sortedcord/alemno-payments.git
synced 2026-07-22 04:02:49 +05:30
feat: Add google-genai & structured output
This commit is contained in:
@@ -1,3 +1,4 @@
|
||||
from typing import Optional
|
||||
from pydantic import Field, computed_field
|
||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
@@ -23,6 +24,9 @@ class Settings(BaseSettings):
|
||||
REDIS_PORT: int = Field(default=6379)
|
||||
REDIS_DB: int = Field(default=0)
|
||||
|
||||
# Gemini
|
||||
GEMINI_API_KEY: Optional[str] = Field(default=None)
|
||||
|
||||
@computed_field
|
||||
@property
|
||||
def DATABASE_URL(self) -> str:
|
||||
|
||||
314
app/worker.py
314
app/worker.py
@@ -1,7 +1,13 @@
|
||||
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
|
||||
@@ -13,6 +19,100 @@ 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",
|
||||
@@ -32,7 +132,7 @@ celery_app.conf.update(
|
||||
|
||||
def run_async(coro):
|
||||
"""Helper to run async functions synchronously in Celery worker context"""
|
||||
return asyncio.get_event_loop().run_until_complete(coro)
|
||||
return asyncio.run(coro)
|
||||
|
||||
|
||||
@celery_app.task(name="tasks.process_transaction_job")
|
||||
@@ -66,6 +166,148 @@ 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)
|
||||
|
||||
# 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"]
|
||||
@@ -78,7 +320,7 @@ async def _process_job_async(job_id: str, transactions: List[Dict[str, Any]]):
|
||||
merchant = t["merchant"].strip()
|
||||
amount = float(t["amount"])
|
||||
currency = t.get("currency", "INR").strip().upper()
|
||||
category = t.get("category", "General")
|
||||
category = t.get("category") or ""
|
||||
account_id = t.get("account_id")
|
||||
|
||||
# Spend calculation and currency conversion
|
||||
@@ -95,28 +337,62 @@ async def _process_job_async(job_id: str, transactions: List[Dict[str, Any]]):
|
||||
# Aggregate merchant spend (in USD for standardization)
|
||||
merchant_spends[merchant] = merchant_spends.get(merchant, 0.0) + amt_usd
|
||||
|
||||
# Heuristic Anomaly detection
|
||||
is_anomaly = False
|
||||
anomaly_reason = None
|
||||
# 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:
|
||||
is_anomaly = True
|
||||
anomaly_reason = "High value transaction (> $5,000 USD equivalent)"
|
||||
else:
|
||||
desc_lower = merchant.lower()
|
||||
if any(
|
||||
k in desc_lower
|
||||
for k in ["suspend", "hack", "error", "unknown", "fraud"]
|
||||
):
|
||||
is_anomaly = True
|
||||
anomaly_reason = "Suspicious merchant keyword"
|
||||
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
|
||||
|
||||
# Mock
|
||||
llm_failed = False
|
||||
llm_category = category.capitalize() if category else "General"
|
||||
llm_raw_response = f'{{"category": "{llm_category}", "confidence": 0.95, "status": "processed"}}'
|
||||
# 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,
|
||||
|
||||
Reference in New Issue
Block a user