Update app/services/autocategorizer.py

app/services/autocategorizer.py

@@ -7,6 +7,7 @@ import re
 import time
 from typing import Callable, Dict, List, Optional
 
+from bson import ObjectId
 from fastapi import HTTPException
 from motor.motor_asyncio import AsyncIOMotorCollection
 from openai import AsyncOpenAI
@@ -18,14 +19,137 @@ from app.schemas.categories import CategoryPrediction
 class AutoCategoryService:
     """Classifies transaction notes into the closest Mongo-backed category."""
 
+    # Curated categories requested by the client. When enabled via settings.use_static_categories,
+    # we bypass Mongo reads to avoid noisy data and long scans.
+    _STATIC_CATEGORIES: List[Dict[str, object]] = [
+        {
+            "title": "Food & Drinks",
+            "subcategories": ["Groceries", "Restaurant, Fast - Food", "Bar, Cafe", "Food & Drink"],
+        },
+        {
+            "title": "Investments",
+            "subcategories": [
+                "Investments",
+                "Realty",
+                "Vehicles, Chattels",
+                "Finacial investments",
+                "Savings",
+                "Collections",
+            ],
+        },
+        {
+            "title": "Communication,PC",
+            "subcategories": ["Communication,PC", "Phone", "Internet", "Software, app, games", "Postal services"],
+        },
+        {
+            "title": "Financial Expenses",
+            "subcategories": [
+                "Financial expenses",
+                "Taxes",
+                "Insurances",
+                "Loan, interests",
+                "Fines",
+                "Advisory",
+                "Charges, Fees",
+                "Child Support",
+            ],
+        },
+        {
+            "title": "Life & Entertainment",
+            "subcategories": [
+                "Life & Entertainment",
+                "Health, Care, Doctor",
+                "Wellness, Beauty",
+                "Active sport, Fitness",
+                "Culture, sport events",
+                "Life events",
+                "Hobbies",
+                "Education, Development",
+                "Books, Audio, subscription",
+                "TV, Streaming",
+                "Holiday, Trip, Hotels",
+                "Charity, Gifts",
+                "Alcohol, Tobacco",
+                "Lottery, Gamblings",
+            ],
+        },
+        {
+            "title": "Vehicle",
+            "subcategories": [
+                "Vehicle",
+                "Fuel",
+                "Parking",
+                "Vehicle maintenance",
+                "Rentals",
+                "Vehicle insurance",
+                "Leasing",
+            ],
+        },
+        {
+            "title": "Transportation",
+            "subcategories": ["Transportation", "Public transport", "Taxi", "Long distance", "Business trips"],
+        },
+        {
+            "title": "Housing",
+            "subcategories": [
+                "Housing",
+                "Rent",
+                "Mortgage",
+                "Energy, utilities",
+                "Services",
+                "Maintenance, repairs",
+                "Property insurance",
+            ],
+        },
+        {
+            "title": "Shopping",
+            "subcategories": [
+                "Shopping",
+                "Clothes & shoes",
+                "Jewels & Accessories",
+                "Health & Beauty",
+                "Kids",
+                "Home & Garden",
+                "Pets & Animals",
+                "Electronics",
+                "Gift",
+                "Stationary",
+                "Free time",
+                "Chemist",
+            ],
+        },
+        {
+            "title": "Income",
+            "subcategories": [
+                "Income",
+                "Wage, Invoices",
+                "Sale",
+                "Rental income",
+                "Dues & grants",
+                "Lending, renting",
+                "Checks, coupons",
+                "Lottery, gambling",
+                "Refunds",
+                "Child support",
+                "Gifts",
+                "Account Manage",
+            ],
+        },
+    ]
+
+    _categories_timeout_seconds = 15.0
+    _model_timeout_seconds = 20.0
+
     def __init__(
         self,
         collection_getter: Callable[[], AsyncIOMotorCollection],
+        subcategory_collection_getter: Callable[[], AsyncIOMotorCollection],
         openai_client: AsyncOpenAI,
         model: str,
         cache_ttl_seconds: int,
     ) -> None:
         self._collection_getter = collection_getter
+        self._subcategory_collection_getter = subcategory_collection_getter
         self._openai_client = openai_client
         self._model = model
         self._cache_ttl_seconds = cache_ttl_seconds
@@ -36,8 +160,19 @@ class AutoCategoryService:
     def _collection(self) -> AsyncIOMotorCollection:
         return self._collection_getter()
 
+    def _subcategory_collection(self) -> AsyncIOMotorCollection:
+        return self._subcategory_collection_getter()
+
     async def categorize(self, notes: str) -> CategoryPrediction:
-        categories = await self._get_categories()
+        try:
+            categories = await asyncio.wait_for(
+                self._get_categories(), timeout=self._categories_timeout_seconds
+            )
+        except asyncio.TimeoutError as exc:
+            raise HTTPException(status_code=504, detail="Timed out loading categories from database.") from exc
+        except Exception as exc:
+            raise HTTPException(status_code=502, detail="Failed to load categories from database.") from exc
+
         if not categories:
             raise HTTPException(status_code=500, detail="No categories configured.")
 
@@ -66,15 +201,22 @@ class AutoCategoryService:
         )
 
         try:
-            response = await self._openai_client.responses.create(
-                response_format={"type": "json_object"},
-                **request_payload,
+            response = await asyncio.wait_for(
+                self._create_model_response(request_payload),
+                timeout=self._model_timeout_seconds,
             )
         except TypeError as exc:
             # Older openai-python clients (pre 1.3x) do not yet support response_format.
             if "response_format" not in str(exc):
                 raise
-            response = await self._openai_client.responses.create(**request_payload)
+            response = await asyncio.wait_for(
+                self._openai_client.responses.create(**request_payload),
+                timeout=self._model_timeout_seconds,
+            )
+        except asyncio.TimeoutError as exc:
+            raise HTTPException(status_code=504, detail="Timed out waiting for model response.") from exc
+        except Exception as exc:
+            raise HTTPException(status_code=502, detail="Failed to call the model API.") from exc
 
         try:
             payload = self._parse_response_payload(response)
@@ -192,12 +334,50 @@ class AutoCategoryService:
             if self._cached_categories and (now - self._last_loaded) < self._cache_ttl_seconds:
                 return self._cached_categories
 
+            if settings.use_static_categories:
+                self._cached_categories = self._STATIC_CATEGORIES
+                self._last_loaded = now
+                return self._cached_categories
+
+            # Use headcategories + categories to avoid scanning millions of raw transaction titles.
+            head_collection = self._collection()
+            subcategory_collection = self._subcategory_collection()
+
+            pipeline = [
+                {"$match": {"type": "EXPENSE", "categories": {"$type": "array", "$ne": []}}},
+                {"$group": {"_id": "$title", "category_ids": {"$first": "$categories"}}},
+            ]
+            head_docs = await head_collection.aggregate(pipeline).to_list(length=1000)
+
+            all_ids: set[ObjectId] = set()
+            for doc in head_docs:
+                for cid in doc.get("category_ids") or []:
+                    if isinstance(cid, ObjectId):
+                        all_ids.add(cid)
+
+            subcategory_titles: Dict[ObjectId, str] = {}
+            if all_ids:
+                cursor = subcategory_collection.find({"_id": {"$in": list(all_ids)}}, {"title": 1})
+                async for subdoc in cursor:
+                    title = subdoc.get("title")
+                    if isinstance(title, str) and title.strip():
+                        subcategory_titles[subdoc["_id"]] = title.strip()
+
             categories: List[Dict[str, object]] = []
-            async for doc in self._collection().find({}, {"_id": 0}):
+            for doc in head_docs:
+                raw_title = doc.get("_id")
+                if not isinstance(raw_title, str) or not raw_title.strip():
+                    continue
+
+                ids = [cid for cid in (doc.get("category_ids") or []) if isinstance(cid, ObjectId)]
+                subcategories = sorted({subcategory_titles[cid] for cid in ids if cid in subcategory_titles})
+                if not subcategories:
+                    continue
+
                 categories.append(
                     {
-                        "title": doc.get("title", ""),
-                        "subcategories": doc.get("subcategories", []),
+                        "title": raw_title.strip(),
+                        "subcategories": subcategories,
                     }
                 )
 
@@ -205,6 +385,12 @@ class AutoCategoryService:
             self._last_loaded = now
             return categories
 
+    async def _create_model_response(self, request_payload: Dict[str, object]):
+        return await self._openai_client.responses.create(
+            response_format={"type": "json_object"},
+            **request_payload,
+        )
+
     @staticmethod
     def _format_categories(categories: List[Dict[str, object]]) -> str:
         lines = []