tests: Add tests for llm summary

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(