"""Date arithmetic solver using python-dateutil. Handles: "N days/weeks/months after/before ", "what day of the week is ", and "how many days between and ". Only fires when the prompt matches one of these shapes exactly and all date tokens parse unambiguously; anything looser falls through to Tier 1. """ from __future__ import annotations import re from datetime import date, datetime from dateutil import parser as dateutil_parser from dateutil.relativedelta import relativedelta from routing_agent.classifier import TaskType from routing_agent.models import Task from routing_agent.solvers import SolverResult _UNIT_TO_KWARG = { "day": "days", "days": "days", "week": "weeks", "weeks": "weeks", "month": "months", "months": "months", "year": "years", "years": "years", } _OFFSET_RE = re.compile( r"(?P\d+)\s+(?Pdays?|weeks?|months?|years?)\s+" r"(?Pafter|before)\s+(?P.+?)\s*\??$", re.IGNORECASE, ) _WEEKDAY_RE = re.compile( r"(?:what\s+day\s+of\s+the\s+week|which\s+weekday|weekday)\s+" r"(?:is|was|falls\s+on)?\s*(?P.+?)\s*\??$", re.IGNORECASE, ) _BETWEEN_RE = re.compile( r"how\s+many\s+days\s+(?:are\s+)?between\s+(?P.+?)\s+and\s+(?P.+?)\s*" r"(?:\(.*\))?\s*\??$", re.IGNORECASE, ) _FROM_TO_RE = re.compile( r"how\s+many\s+days\s+(?:are\s+there\s+)?from\s+(?P.+?)\s+to\s+(?P.+?)\s*\??$", re.IGNORECASE, ) _WEEKDAY_NAMES = ( "Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday", "Sunday", ) def _parse_date(text: str) -> date | None: cleaned = text.strip().strip("?.!,") if not cleaned: return None try: parsed = dateutil_parser.parse(cleaned, fuzzy=False, default=datetime(2000, 1, 1)) except (ValueError, OverflowError): return None return parsed.date() def try_solve(task: Task, task_type: TaskType) -> SolverResult: """Attempt to solve a date-math task deterministically. Only fires for the three supported shapes (offset, weekday-of, days- between); any prompt that doesn't cleanly match one of these regexes, or whose date tokens fail to parse, returns confident=False. """ if task_type != TaskType.DATE_MATH: return SolverResult(answer=None, confident=False) prompt = task.prompt.strip() offset_match = _OFFSET_RE.search(prompt) if offset_match: base = _parse_date(offset_match.group("date")) if base is None: return SolverResult(answer=None, confident=False) amount = int(offset_match.group("num")) unit_kwarg = _UNIT_TO_KWARG[offset_match.group("unit").lower()] if offset_match.group("direction").lower() == "before": amount = -amount result = base + relativedelta(**{unit_kwarg: amount}) return SolverResult(answer=result.isoformat(), confident=True) weekday_match = _WEEKDAY_RE.search(prompt) if weekday_match: target = _parse_date(weekday_match.group("date")) if target is None: return SolverResult(answer=None, confident=False) return SolverResult(answer=_WEEKDAY_NAMES[target.weekday()], confident=True) between_match = _BETWEEN_RE.search(prompt) or _FROM_TO_RE.search(prompt) if between_match: first = _parse_date(between_match.group("date1")) second = _parse_date(between_match.group("date2")) if first is None or second is None: return SolverResult(answer=None, confident=False) delta_days = abs((second - first).days) return SolverResult(answer=str(delta_days), confident=True) return SolverResult(answer=None, confident=False)