mirror of
https://github.com/sortedcord/alemno-payments.git
synced 2026-07-22 04:02:49 +05:30
feat: Add API endpoint for managing categories
This commit is contained in:
147
app/worker.py
147
app/worker.py
@@ -24,7 +24,7 @@ class TransactionClassification(BaseModel):
|
||||
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."
|
||||
description="Assigned category matching one of the allowed categories list."
|
||||
)
|
||||
|
||||
|
||||
@@ -35,6 +35,7 @@ class BatchClassificationResponse(BaseModel):
|
||||
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.
|
||||
@@ -55,8 +56,7 @@ async def classify_transactions_batch(
|
||||
|
||||
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"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)}"
|
||||
)
|
||||
|
||||
@@ -146,6 +146,9 @@ def process_transaction_job(job_id: str, transactions: List[Dict[str, Any]]):
|
||||
|
||||
|
||||
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)
|
||||
@@ -215,6 +218,28 @@ async def _process_job_async(job_id: str, transactions: List[Dict[str, Any]]):
|
||||
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..."
|
||||
@@ -226,7 +251,7 @@ async def _process_job_async(job_id: str, transactions: List[Dict[str, Any]]):
|
||||
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
|
||||
client, batch, allowed_categories
|
||||
)
|
||||
llm_results.update(batch_results)
|
||||
except Exception as e:
|
||||
@@ -247,58 +272,92 @@ async def _process_job_async(job_id: str, transactions: List[Dict[str, Any]]):
|
||||
# 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 = (
|
||||
"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"]
|
||||
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",
|
||||
]
|
||||
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"]
|
||||
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",
|
||||
]
|
||||
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"]
|
||||
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",
|
||||
]
|
||||
elif (
|
||||
any(
|
||||
kw in merchant_lower
|
||||
for kw in [
|
||||
"netflix",
|
||||
"spotify",
|
||||
"movie",
|
||||
"show",
|
||||
"entertainment",
|
||||
]
|
||||
)
|
||||
and "Entertainment" in allowed_categories
|
||||
):
|
||||
cat = "Entertainment"
|
||||
|
||||
|
||||
Reference in New Issue
Block a user