implement-transaction-listing-with-smart-filtering-and-enable-auto-tool-exposure-via-decorator

R-Rudolf · R-Rudolf · commit 4978153a01b0 · 2025-08-14T08:11:54.000+02:00
diff --git a/spendee/spendee_firestore.py b/spendee/spendee_firestore.py
@@ -1,4 +1,5 @@
 from decimal import Decimal
+from typing import Callable, Dict
 from google.auth.credentials import Credentials
 from google.cloud import firestore
 from google.cloud.firestore_v1.base_query import FieldFilter
@@ -7,6 +8,7 @@
 import json
 import datetime
 import logging
+import re
 
 # Improvement ideas:
 # - Implement token expiration check in CustomFirebaseCredentials
@@ -17,8 +19,15 @@
 # Note: User document has field: firestoreDataExportDone, which is True.
 # Which is a trace of the migration from REST API to Firestore.
 
+
 logger = logging.getLogger(__name__)
 
+# MCP tool registry and decorator
+MCP_TOOLS: Dict[str, Callable] = {}
+def mcp_tool(func: Callable) -> Callable:
+    MCP_TOOLS[func.__name__] = func
+    return func
+
 class CustomFirebaseCredentials(Credentials):
     """Custom credentials that use existing Firebase access token"""
     
@@ -56,6 +65,8 @@ def __init__(self, email: str, password: str, base_url: str = 'https://api.spend
         self.email = access_token_data.get('email', None)
         self.user_name = access_token_data.get('name', None)
         self.wallet_name_map = { x['name']: x['id'] for x in self.list_wallets()}
+        self.category_name_map = { x['id']: x['name'] for x in self.list_categories()}
+        self.label_name_map = { x['id']: x['name'] for x in self.list_labels()}
         logger.info(f"SpendeeFirestore initialized for user_id={self.user_id}, email={self.email}")
 
     def _token_refresh(self):
@@ -107,6 +118,36 @@ def _json_serializer(self, obj):
     def as_json(self, obj):
         return json.dumps(obj, indent=2, default=self._json_serializer, ensure_ascii=False)
 
+    @staticmethod
+    def _matches_filters(value, filters):
+        for f in filters or []:
+            field = f.get("field")
+            op = f.get("op")
+            filter_value = f.get("value")
+            v = value.get(field)
+            if op == ">":
+                if not (v is not None and float(v) > float(filter_value)):
+                    return False
+            elif op == ">=":
+                if not (v is not None and float(v) >= float(filter_value)):
+                    return False
+            elif op == "<":
+                if not (v is not None and float(v) < float(filter_value)):
+                    return False
+            elif op == "<=":
+                if not (v is not None and float(v) <= float(filter_value)):
+                    return False
+            elif op == "=":
+                if not (str(v) == str(filter_value)):
+                    return False
+            elif op == "~=":
+                if not (v is not None and re.search(str(filter_value), str(v))):
+                    return False
+            else:
+                logger.warning(f"Unsupported filter op: {op}")
+                return False
+        return True
+
     # --- Spendee API methods ---
 
     def _get_raw_category(self, category_id: str, as_json: bool = False):
@@ -133,12 +174,16 @@ def _list_raw_categories(self, as_json: bool = False):
         return self.as_json(raw_data) if as_json else raw_data
     
 
+    # ...existing code...
     def list_categories(self, as_json: bool = False):
-        logger.info("Listing categories.")
         """
-        Returns a list of categories for the user.
+        Returns the list of categories of the user.
         If as_json is True, returns the data as a JSON string.
+
+        Each category has the fields: id, name, type
+        where type can be 'income' or 'expense'
         """
+        logger.info("Listing categories.")
         data = []
         for raw_data in self.client.collection(f'users/{self.user_id}/categories').get():
             data.append({
@@ -160,12 +205,15 @@ def _list_raw_labels(self, as_json: bool = False):
         return self.as_json(raw_data) if as_json else raw_data
 
 
+    @mcp_tool
     def list_labels(self, as_json: bool = False):
-        logger.info("Listing labels.")
         """
-        Returns a list of labels for the user.
+        Returns the list of labels used by the user.
         If as_json is True, returns the data as a JSON string.
+
+        Each label item has the fields: id, name
         """
+        logger.info("Listing labels.")
         data = []
         for raw_data in self.client.collection(f'users/{self.user_id}/labels').get():
             data.append({
@@ -177,11 +225,19 @@ def list_labels(self, as_json: bool = False):
         return self.as_json(data) if as_json else data
 
 
+    @mcp_tool
     def get_wallet_balance(self, wallet_id: str, start: str = None, end: str = None):
+        """Get the balance of a wallet for a specific timeframe.
+        The start and end parameters should be in ISO 8601 format. If not set,
+        no filtering is done.
+        Args:
+            wallet_id (str): Name of the wallet. (Should be equal to the results of list_wallets call)
+            start (str, optional): Start date in ISO 8601 format.
+            end (str, optional): End date in ISO 8601 format.
+        Returns:
+            int: The balance of the wallet.
         """
-        Returns the balance of the wallet with the given wallet_id for a specific timeframe.
-        The start and end parameters should be in ISO 8601 format. If not set, no filtering is done.
-        """
+
         logger.info(f"Calculating balance for wallet_id: {wallet_id}")
         query = self.client.collection(f'users/{self.user_id}/wallets/{wallet_id}/transactions')
 
@@ -213,7 +269,17 @@ def get_wallet_balance(self, wallet_id: str, start: str = None, end: str = None)
         return round(total + Decimal(str(starting_balance)))
 
 
+    @mcp_tool
     def list_wallets(self):
+        """List all wallets for the authenticated user. This is required before wallet related calls, to have exact string for names.
+        Each wallet is represented by an object with the following fields:
+        - id: Unique identifier of the wallet
+        - name: Name of the wallet
+        - type: Type of the wallet (e.g., cash, bank, etc.)
+        - currency: Currency of the wallet
+        - updatedAt: Last updated timestamp of the wallet
+        """
+
         logger.info("Listing wallets.")
         raw_data = [
             x.to_dict()
@@ -242,3 +308,101 @@ def _get_raw_transaction(self, wallet_id: str, transaction_id: str, as_json: boo
         obj = self.client.document(f"users/{self.user_id}/wallets/{wallet_id}/transactions/{transaction_id}").get().to_dict()
         logger.debug(f"Fetched raw transaction content: {obj}")
         return self.as_json(obj) if as_json else obj
+
+
+    @mcp_tool
+    def get_transaction(self, wallet_id: str, transaction_id: str, as_json: bool = False):
+        """Get a specific transaction by its ID from a wallet.
+        Args:
+            wallet_id (str): ID of the wallet.
+            transaction_id (str): ID of the transaction.
+            as_json (bool, optional): If True, returns the data as a JSON string.
+        Returns:
+            dict or str: The transaction data, either as a dictionary or JSON string.
+        """
+        value = self._get_raw_transaction(wallet_id, transaction_id)
+        # Convert category ID to category name using self.category_name_map
+        category_id = value.get("category", "")
+        category_name = self.category_name_map.get(category_id, None)
+
+        data = {
+            "note": value.get("note", ""),
+            "madeAt": value.get("madeAt", ""),
+            "category": category_name,
+            "type": value.get("type", ""),
+            "isPending": value.get("isPending", ""),
+            "amount": value.get("amount", ""),
+        }
+
+        logger.info(f"Getting transaction: wallet_id={wallet_id}, transaction_id={transaction_id}")
+        return data if not as_json else self.as_json(data)
+
+
+    @mcp_tool
+    def list_transactions(self, wallet_id: str, start: str, end: str = None, filters: list = None, limit: int = 20, fields: list = ["note", "madeAt", "category", "amount"], as_json: bool = False):
+        """
+        List transactions for a wallet, filtered by date range and dynamic filters.
+
+        Results are always sorted by 'madeAt' in descending order (most recent transactions first).
+        Each returned item has the same fields as get_transaction: id, note, madeAt, category (name), type, isPending, amount.
+        Optionally, the returned fields can be limited by the 'fields' parameter, default is ["note", "madeAt", "category", "amount"].
+
+        Args:
+            wallet_id (str): ID of the wallet.
+            start (str): Start date (ISO 8601, required).
+            end (str, optional): End date (ISO 8601).
+            filters (list, optional): List of filter dicts, e.g. [{"field": "amount", "op": ">=", "value": 100}].
+                The 'field' and 'op' values must be strings.
+                Supported operators: "=", "~=", ">", ">=", "<", "<=", where "~=" is regex match.
+            limit (int, optional): Max number of transactions to return (default 20).
+            fields (list, optional): List of field names to include in the result. Only supported fields are allowed.
+            as_json (bool, optional): Return as JSON string if True.
+        Returns:
+            list or str: List of transaction dicts or JSON string.
+        """
+        logger.info(f"Listing transactions for wallet_id={wallet_id}, start={start}, end={end}, filters={filters}, limit={limit}")
+        query = self.client.collection(f'users/{self.user_id}/wallets/{wallet_id}/transactions')
+
+        # Required start date
+        query = query.where(filter=FieldFilter("madeAt", ">=", datetime.datetime.fromisoformat(start)))
+        if end:
+            query = query.where(filter=FieldFilter("madeAt", "<=", datetime.datetime.fromisoformat(end)))
+
+        # Only order by 'madeAt' (descending), fetch all for post-filtering and limiting
+        query = query.order_by("madeAt", direction=firestore.Query.DESCENDING)
+        transactions = query.stream()
+
+        allowed_fields = {"id", "note", "madeAt", "category", "type", "isPending", "amount"}
+        if not isinstance(fields, list) or not all(isinstance(f, str) for f in fields):
+            raise ValueError("'fields' must be a list of strings.")
+        unsupported = set(fields) - allowed_fields
+        if unsupported:
+            raise ValueError(f"Unsupported fields requested: {unsupported}")
+
+        results = []
+        for transaction in transactions:
+            value = transaction.to_dict()
+            # Match get_transaction output fields
+            category_id = value.get("category", "")
+            category_name = self.category_name_map.get(category_id, None)
+            if category_name is None:
+                logger.warning(f"Category ID {category_id} not found in category_name_map, using None.")
+            data = {
+                "id": value.get("path", {}).get("transaction", ""),
+                "note": value.get("note", ""),
+                "madeAt": value.get("madeAt", ""),
+                "category": category_name,
+                "type": value.get("type", ""),
+                "isPending": value.get("isPending", ""),
+                "amount": value.get("amount", ""),
+            }
+            if not self._matches_filters(data, filters):
+                continue
+            if fields is not None:
+                data = {k: v for k, v in data.items() if k in fields}
+            results.append(data)
+            if len(results) >= limit:
+                break
+
+        logger.info(f"Found {len(results)} transactions.")
+        return self.as_json(results) if as_json else results
diff --git a/spendee/spendee_mcp.py b/spendee/spendee_mcp.py
@@ -18,7 +18,7 @@
 from starlette.responses import Response
 from starlette.types import Scope, Receive, Send
 
-from spendee.spendee_firestore import SpendeeFirestore
+from spendee.spendee_firestore import SpendeeFirestore, MCP_TOOLS
 
 # to start (after .venv setup):
 #   python spendee/spendee_mcp.py
@@ -50,8 +50,15 @@
 logger = logging.getLogger(__name__)
 
 mcp = FastMCP("spendee", host="0.0.0.0", port=PORT)
+
 spendee = SpendeeFirestore(EMAIL, PASSWORD)
 
+# Automatically register all MCP tools from the client
+for name, func in MCP_TOOLS.items():
+    # Bind the method to the spendee instance if it's a class method
+    bound_func = getattr(spendee, name)
+    mcp.tool()(bound_func)
+
 def main():
     #debug_secret(ACCEPTED_TOKEN, "MCP_TOKEN")
     #debug_secret(PASSWORD, "PASSWORD")
@@ -65,33 +72,6 @@ def main():
         logger.info("Using SSE transport")
         sse_server()
 
-
-@mcp.tool()
-def get_wallet_balance(wallet_id: str, start: str = None, end: str = None):
-    """Get the balance of a wallet for a specific timeframe.
-    The start and end parameters should be in ISO 8601 format. If not set,
-    no filtering is done.
-    Args:
-        wallet_id (str): Name of the wallet. (Should be equal to the results of list_wallets call)
-        start (str, optional): Start date in ISO 8601 format.
-        end (str, optional): End date in ISO 8601 format.
-    Returns:
-        int: The balance of the wallet.
-    """
-    return spendee.get_wallet_balance(wallet_id, start, end)
-
-@mcp.tool()
-def list_wallets():
-    """List all wallets for the authenticated user. This is required before wallet related calls, to have exact string for names.
-    Each wallet is represented by an object with the following fields:
-    - id: Unique identifier of the wallet
-    - name: Name of the wallet
-    - type: Type of the wallet (e.g., cash, bank, etc.)
-    - currency: Currency of the wallet
-    - updatedAt: Last updated timestamp of the wallet
-    """
-    return spendee.list_wallets()
-
 # Authentication middleware and server setup
 
 def debug_secret(secret, name):
@@ -191,3 +171,5 @@ async def sse_auth_middleware(scope: Scope, receive: Receive, send: Send):
 
 if __name__ == "__main__":
     main()
+else:
+    main()
