diff --git a/README.md b/README.md index eb4a3c4..f81f5b8 100644 --- a/README.md +++ b/README.md @@ -13,7 +13,7 @@ An asynchronous transactional data cleaning, validation, and LLM powered categor - Reads multi format dates and handles missing transaction IDs by auto generating unique fallbacks, and parses messy amounts with currency signs and commas cleanly. - Offloads validation, spend conversion, and analysis to a Celery background worker backed by Redis. -- Integrates with the Google GenAI SDK (`gemini-2.5-flash-lite` 75% cheaper than `gemini-3.5-flash-lite with no tangible decrease in performance) using structured output schemas to categorize transaction records into database-defined custom categories. +- Integrates with the Google GenAI SDK (`gemini-2.5-flash-lite` 75% cheaper than `gemini-3.5-flash-lite with no tangible decrease in performance) using structured output schemas to categorize transaction records into database- defined custom categories. - Uses Redis caching with a 24-hour TTL to store live exchange rates, falling back to a hardcoded rate of `93.0` if offline. - Statistical Anomaly Detection: - Flags transactions exceeding 3 times the account's median spend as statistical outliers. diff --git a/app/worker.py b/app/worker.py index fd12491..6859bd1 100644 --- a/app/worker.py +++ b/app/worker.py @@ -3,10 +3,10 @@ import json import os import statistics from datetime import datetime, date -from typing import Any, Dict, List +from typing import Any, Dict, List, Optional from sqlalchemy import select from celery import Celery -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, AliasChoices from google import genai from app.core.config import settings from app.core.logging import logger @@ -19,6 +19,27 @@ from app.clients.exchange_rate_client import ExchangeRateClient 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." @@ -80,9 +101,8 @@ async def classify_transactions_batch( ), ) - resp_obj = BatchClassificationResponse.model_validate_json( - interaction.output_text - ) + clean_text = clean_json_response(interaction.output_text) + resp_obj = BatchClassificationResponse.model_validate_json(clean_text) results = {} for item in resp_obj.classifications: @@ -113,6 +133,94 @@ async def classify_transactions_batch( 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", + ) + + +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", @@ -473,31 +581,58 @@ async def _process_job_async(job_id: str, transactions: List[Dict[str, Any]]): 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" + # 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}") - 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}'." - ) + 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" - # 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 - } + 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 summary = JobSummary( diff --git a/tests/services/test_llm_summary.py b/tests/services/test_llm_summary.py new file mode 100644 index 0000000..0bb3b91 --- /dev/null +++ b/tests/services/test_llm_summary.py @@ -0,0 +1,77 @@ +import pytest +from app.worker import generate_job_summary_llm, clean_json_response +from app.db.models.transaction import Transaction +from google import genai +from unittest.mock import MagicMock, patch + + +def test_clean_json_response_markdown(): + # Test stripping markdown code fences + dirty_json = '```json\n{"key": "value"}\n```' + assert clean_json_response(dirty_json) == '{"key": "value"}' + + # Test stripping markdown code fences with conversational prefix + dirty_conversational = ( + 'Here is the data:\n```\n{"key": "value"}\n```\nHope it helps.' + ) + assert clean_json_response(dirty_conversational) == '{"key": "value"}' + + +@pytest.mark.asyncio +async def test_generate_job_summary_llm_success(): + # 1. Setup mock transactions + t1 = Transaction( + merchant="Zomato", + amount=100.0, + currency="USD", + category="Food", + is_anomaly=False, + ) + t2 = Transaction( + merchant="Ola", + amount=2500.0, + currency="INR", + category="Transport", + is_anomaly=True, + anomaly_reason="Outlier", + ) + tx_list = [t1, t2] + + # 2. Mock client and return json + mock_client = MagicMock(spec=genai.Client) + mock_interaction = MagicMock() + mock_interaction.output_text = """ + ```json + { + "total_spend_by_currency": {"USD": 100.0, "INR": 2500.0}, + "top_3_merchants": {"Ola": 30.12, "Zomato": 100.0}, + "anomaly_count": 1, + "spending_narrative": "Cohesive two sentence summary of spendings. Mostly food and transit.", + "risk_level": "Medium" + } + ``` + """ + mock_client.interactions.create.return_value = mock_interaction + + # 3. Call and assert + result = await generate_job_summary_llm(mock_client, tx_list) + assert result is not None + assert result.anomaly_count == 1 + assert result.risk_level == "Medium" + assert result.total_spend_by_currency["USD"] == 100.0 + assert result.total_spend_by_currency["INR"] == 2500.0 + assert "Ola" in result.top_3_merchants + + +@pytest.mark.asyncio +async def test_generate_job_summary_llm_retry_failure(): + mock_client = MagicMock(spec=genai.Client) + mock_client.interactions.create.side_effect = Exception("API Key Expired") + + tx_list = [Transaction(merchant="Test", amount=1.0, currency="USD")] + + # Patch sleep to execute instantly + with patch("asyncio.sleep", return_value=None): + result = await generate_job_summary_llm(mock_client, tx_list) + + assert result is None