tests: Add tests for llm summary

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.

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(

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