Feat/budget limits

src/lampyrid/clients/firefly.py

@@ -15,6 +15,9 @@
     AvailableBudgetArray,
     BudgetArray,
     BudgetLimitArray,
+    BudgetLimitSingle,
+    BudgetLimitStore,
+    BudgetLimitUpdate,
     BudgetSingle,
     BudgetStore,
     InsightGroup,
@@ -282,6 +285,33 @@ async def get_budget_limits(
         r.raise_for_status()
         return BudgetLimitArray.model_validate(r.json())
 
+    async def create_budget_limit(
+        self, budget_id: str, budget_limit_store: BudgetLimitStore
+    ) -> BudgetLimitSingle:
+        """Create a budget limit for a budget."""
+        payload = self._serialize_model(budget_limit_store)
+        r = await self._client.post(f'/api/v1/budgets/{budget_id}/limits', json=payload)
+        self._handle_api_error(r, payload)
+        r.raise_for_status()
+        return BudgetLimitSingle.model_validate(r.json())
+
+    async def update_budget_limit(
+        self, budget_id: str, limit_id: str, budget_limit_update: BudgetLimitUpdate
+    ) -> BudgetLimitSingle:
+        """Update an existing budget limit."""
+        payload = self._serialize_model(budget_limit_update, exclude_unset=True)
+        r = await self._client.put(f'/api/v1/budgets/{budget_id}/limits/{limit_id}', json=payload)
+        self._handle_api_error(r, payload)
+        r.raise_for_status()
+        return BudgetLimitSingle.model_validate(r.json())
+
+    async def delete_budget_limit(self, budget_id: str, limit_id: str) -> bool:
+        """Delete a budget limit by ID."""
+        r = await self._client.delete(f'/api/v1/budgets/{budget_id}/limits/{limit_id}')
+        self._handle_api_error(r)
+        r.raise_for_status()
+        return r.status_code == 204
+
     async def create_budget(self, budget_store: BudgetStore) -> BudgetSingle:
         """Create a new budget."""
         payload = self._serialize_model(budget_store)

src/lampyrid/models/lampyrid_models.py

@@ -8,6 +8,7 @@
 from .firefly_models import (
     AccountRead,
     AccountTypeFilter,
+    BudgetLimitRead,
     BudgetRead,
     ShortAccountTypeProperty,
     TransactionArray,
@@ -750,6 +751,126 @@ def validate_auto_budget(self) -> 'CreateBudgetRequest':
         return self
 
 
+class BudgetLimit(BaseModel):
+    """A budget limit: the spendable amount set for a budget over a specific period."""
+
+    id: str = Field(..., description='Unique identifier for the budget limit', examples=['5'])
+    budget_id: str = Field(
+        ..., description='ID of the budget this limit belongs to', examples=['3']
+    )
+    amount: float = Field(
+        ..., description='Amount allocated to the budget for this period', examples=[500.0]
+    )
+    start_date: date = Field(..., description='First day this limit applies to (inclusive)')
+    end_date: date = Field(..., description='Last day this limit applies to (inclusive)')
+    currency_code: Optional[str] = Field(
+        None, description='Currency code (ISO 4217) for the amount', examples=['USD']
+    )
+    spent: Optional[float] = Field(
+        None,
+        description='Amount already spent against this limit in the period (positive number)',
+        examples=[120.0],
+    )
+    notes: Optional[str] = Field(None, description='Optional notes attached to this budget limit')
+
+    @classmethod
+    def from_budget_limit_read(cls, budget_limit_read: 'BudgetLimitRead') -> 'BudgetLimit':
+        """Create a BudgetLimit instance from a Firefly BudgetLimitRead object."""
+        attrs = budget_limit_read.attributes
+
+        spent: Optional[float] = None
+        if attrs.spent:
+            spent = 0.0
+            for spent_entry in attrs.spent:
+                if spent_entry.sum:
+                    spent += abs(float(spent_entry.sum))
+
+        return cls(
+            id=budget_limit_read.id,
+            budget_id=attrs.budget_id,
+            amount=float(attrs.amount),
+            start_date=attrs.start.date(),
+            end_date=attrs.end.date(),
+            currency_code=attrs.currency_code,
+            spent=spent,
+            notes=attrs.notes,
+        )
+
+
+class _BudgetLimitRequestBase(BaseModel):
+    """Shared base for budget-limit requests: identify a budget and an optional period."""
+
+    model_config = ConfigDict(extra='forbid')
+
+    budget_id: Optional[str] = Field(
+        None,
+        description='ID of the budget (from list_budgets). Provide budget_id OR budget_name.',
+    )
+    budget_name: Optional[str] = Field(
+        None,
+        description=(
+            'Name of the budget if the ID is unknown (e.g., "Groceries"). '
+            'Provide budget_id OR budget_name. If both are given, budget_id wins.'
+        ),
+    )
+    start_date: Optional[date] = Field(
+        None,
+        description=(
+            'Start date of the period (YYYY-MM-DD), inclusive. '
+            'Defaults to the first day of the current month. '
+            'Must be provided together with end_date.'
+        ),
+    )
+    end_date: Optional[date] = Field(
+        None,
+        description=(
+            'End date of the period (YYYY-MM-DD), inclusive. '
+            'Defaults to the last day of the current month. '
+            'Must be provided together with start_date.'
+        ),
+    )
+
+    @model_validator(mode='after')
+    def _validate_budget_ref_and_period(self):
+        """Require a budget reference and ensure dates are provided together."""
+        if self.budget_id is None and self.budget_name is None:
+            raise ValueError('Provide either budget_id or budget_name.')
+        if (self.start_date is None) != (self.end_date is None):
+            raise ValueError('Provide both start_date and end_date, or neither.')
+        if (
+            self.start_date is not None
+            and self.end_date is not None
+            and self.end_date < self.start_date
+        ):
+            raise ValueError('end_date must not be before start_date.')
+        return self
+
+
+class SetBudgetLimitRequest(_BudgetLimitRequestBase):
+    """Request to set (create or update) a budget limit for a period."""
+
+    amount: float = Field(
+        ...,
+        description='Spendable amount to allocate to the budget for the period (e.g., 500.0)',
+        gt=0,
+    )
+    currency_code: Optional[str] = Field(
+        None,
+        description=(
+            "Currency code (ISO 4217, e.g., 'USD'). Defaults to the user's primary currency."
+        ),
+    )
+    notes: Optional[str] = Field(None, description='Optional notes to attach to this budget limit')
+
+
+class ListBudgetLimitsRequest(_BudgetLimitRequestBase):
+    """Request to list budget limits for a budget, optionally within a date range."""
+
+
+class DeleteBudgetLimitRequest(_BudgetLimitRequestBase):
+    """Request to delete the budget limit for a budget and period."""
+
+
 class CreateBulkTransactionsRequest(BaseModel):
     """Create multiple transactions in one operation."""
 

src/lampyrid/services/budgets.py

@@ -4,28 +4,35 @@
 operations between the MCP tools and the Firefly III client.
 """
 
-from datetime import date
-from typing import List
+from datetime import date, timedelta
+from typing import List, Optional, Tuple
 
 from ..clients.firefly import FireflyClient
 from ..models.firefly_models import (
     AutoBudgetPeriod,
     AutoBudgetPeriodEnum,
     AutoBudgetType,
     AutoBudgetTypeEnum,
+    BudgetLimitRead,
+    BudgetLimitStore,
+    BudgetLimitUpdate,
     BudgetStore,
 )
 from ..models.lampyrid_models import (
     AvailableBudget,
     Budget,
+    BudgetLimit,
     BudgetSpending,
     BudgetSummary,
     CreateBudgetRequest,
+    DeleteBudgetLimitRequest,
     GetAvailableBudgetRequest,
     GetBudgetRequest,
     GetBudgetSpendingRequest,
     GetBudgetSummaryRequest,
+    ListBudgetLimitsRequest,
     ListBudgetsRequest,
+    SetBudgetLimitRequest,
 )
 
 
@@ -235,3 +242,147 @@ async def create_budget(self, req: CreateBudgetRequest) -> Budget:
 
         budget_single = await self._client.create_budget(budget_store)
         return Budget.from_budget_read(budget_single.data)
+
+    async def _resolve_budget_id(self, budget_id: Optional[str], budget_name: Optional[str]) -> str:
+        """Resolve a budget reference (ID or name) to a budget ID.
+
+        If budget_id is provided it is returned as-is. Otherwise the budget is looked up
+        by name (case-insensitive exact match) from the list of all budgets.
+
+        Raises:
+                ValueError: If no budget reference is given, the name is not found, or the
+                        name matches more than one budget.
+
+        """
+        if budget_id is not None:
+            return budget_id
+
+        if budget_name is None:
+            raise ValueError('Provide either budget_id or budget_name.')
+
+        budgets_array = await self._client.get_budgets()
+        matches = [
+            budget
+            for budget in budgets_array.data
+            if budget.attributes.name.lower() == budget_name.lower()
+        ]
+
+        if not matches:
+            raise ValueError(f'No budget found with name {budget_name!r}.')
+        if len(matches) > 1:
+            raise ValueError(
+                f'Multiple budgets match name {budget_name!r}; use budget_id to disambiguate.'
+            )
+        return matches[0].id
+
+    @staticmethod
+    def _resolve_period(start_date: Optional[date], end_date: Optional[date]) -> Tuple[date, date]:
+        """Resolve a period, defaulting to the current calendar month when omitted.
+
+        Both dates must be provided together or neither. The request models already
+        enforce this, but the helper validates it too in case it is called directly.
+        """
+        if (start_date is None) != (end_date is None):
+            raise ValueError('Provide both start_date and end_date, or neither.')
+
+        if start_date is not None and end_date is not None:
+            return start_date, end_date
+
+        today = date.today()
+        first = today.replace(day=1)
+        last = (first.replace(day=28) + timedelta(days=4)).replace(day=1) - timedelta(days=1)
+        return first, last
+
+    async def _find_limit_for_period(
+        self, budget_id: str, start: date, end: date
+    ) -> Optional[BudgetLimitRead]:
+        """Find an existing budget limit that exactly matches the given period."""
+        limits_array = await self._client.get_budget_limits(budget_id, start, end)
+        for limit in limits_array.data:
+            attrs = limit.attributes
+            if (
+                attrs.start is not None
+                and attrs.end is not None
+                and attrs.start.date() == start
+                and attrs.end.date() == end
+            ):
+                return limit
+        return None
+
+    async def set_budget_limit(self, req: SetBudgetLimitRequest) -> BudgetLimit:
+        """Set (create or update) a budget limit for a budget and period.
+
+        If a limit already exists for the exact period it is updated; otherwise a new
+        limit is created. This makes the operation an idempotent upsert from the
+        caller's perspective.
+
+        Args:
+                req: Request containing the budget reference, amount, and optional period.
+
+        Returns:
+                The created or updated budget limit.
+
+        """
+        budget_id = await self._resolve_budget_id(req.budget_id, req.budget_name)
+        start, end = self._resolve_period(req.start_date, req.end_date)
+
+        existing = await self._find_limit_for_period(budget_id, start, end)
+
+        if existing is not None:
+            limit_update = BudgetLimitUpdate(amount=str(req.amount))
+            if req.notes is not None:
+                limit_update.notes = req.notes
+            limit_single = await self._client.update_budget_limit(
+                budget_id, existing.id, limit_update
+            )
+        else:
+            limit_store = BudgetLimitStore(
+                budget_id=budget_id,
+                start=start,
+                end=end,
+                amount=str(req.amount),
+                currency_code=req.currency_code,
+                notes=req.notes,
+            )
+            limit_single = await self._client.create_budget_limit(budget_id, limit_store)
+
+        return BudgetLimit.from_budget_limit_read(limit_single.data)
+
+    async def list_budget_limits(self, req: ListBudgetLimitsRequest) -> List[BudgetLimit]:
+        """List budget limits for a budget, optionally filtered by date range.
+
+        Args:
+                req: Request containing the budget reference and optional date range.
+
+        Returns:
+                List of budget limits set for the budget.
+
+        """
+        budget_id = await self._resolve_budget_id(req.budget_id, req.budget_name)
+        limits_array = await self._client.get_budget_limits(budget_id, req.start_date, req.end_date)
+        return [BudgetLimit.from_budget_limit_read(limit) for limit in limits_array.data]
+
+    async def delete_budget_limit(self, req: DeleteBudgetLimitRequest) -> bool:
+        """Delete the budget limit for a budget and period.
+
+        Args:
+                req: Request containing the budget reference and optional period.
+
+        Returns:
+                True if the limit was deleted.
+
+        Raises:
+                ValueError: If no limit exists for the resolved period.
+
+        """
+        budget_id = await self._resolve_budget_id(req.budget_id, req.budget_name)
+        start, end = self._resolve_period(req.start_date, req.end_date)
+
+        existing = await self._find_limit_for_period(budget_id, start, end)
+        if existing is None:
+            raise ValueError(
+                f'No budget limit set for budget {budget_id} for period '
+                f'{start.isoformat()} to {end.isoformat()}.'
+            )
+
+        return await self._client.delete_budget_limit(budget_id, existing.id)