hindsight-api 0.1.4__py3-none-any.whl → 0.1.6__py3-none-any.whl

hindsight_api/engine/query_analyzer.py

@@ -4,11 +4,12 @@ Query analysis abstraction for the memory system.
 Provides an interface for analyzing natural language queries to extract
 structured information like temporal constraints.
 """
-from abc import ABC, abstractmethod
-from typing import Optional
-from datetime import datetime, timedelta
+
 import logging
 import re
+from abc import ABC, abstractmethod
+from datetime import datetime, timedelta
+
 from pydantic import BaseModel, Field
 
 logger = logging.getLogger(__name__)
@@ -20,6 +21,7 @@ class TemporalConstraint(BaseModel):
 
     Represents a time range with start and end dates.
     """
+
     start_date: datetime = Field(description="Start of the time range (inclusive)")
     end_date: datetime = Field(description="End of the time range (inclusive)")
 
@@ -33,9 +35,9 @@ class QueryAnalysis(BaseModel):
 
     Contains extracted structured information like temporal constraints.
     """
-    temporal_constraint: Optional[TemporalConstraint] = Field(
-        default=None,
-        description="Extracted temporal constraint, if any"
+
+    temporal_constraint: TemporalConstraint | None = Field(
+        default=None, description="Extracted temporal constraint, if any"
     )
 
 
@@ -58,9 +60,7 @@ class QueryAnalyzer(ABC):
         pass
 
     @abstractmethod
-    def analyze(
-        self, query: str, reference_date: Optional[datetime] = None
-    ) -> QueryAnalysis:
+    def analyze(self, query: str, reference_date: datetime | None = None) -> QueryAnalysis:
         """
         Analyze a natural language query.
 
@@ -95,11 +95,10 @@ class DateparserQueryAnalyzer(QueryAnalyzer):
         """Load dateparser (lazy import)."""
         if self._search_dates is None:
             from dateparser.search import search_dates
+
             self._search_dates = search_dates
 
-    def analyze(
-        self, query: str, reference_date: Optional[datetime] = None
-    ) -> QueryAnalysis:
+    def analyze(self, query: str, reference_date: datetime | None = None) -> QueryAnalysis:
         """
         Analyze query using dateparser.
 
@@ -126,9 +125,9 @@ class DateparserQueryAnalyzer(QueryAnalyzer):
 
         # Use dateparser's search_dates to find temporal expressions
         settings = {
-            'RELATIVE_BASE': reference_date,
-            'PREFER_DATES_FROM': 'past',
-            'RETURN_AS_TIMEZONE_AWARE': False,
+            "RELATIVE_BASE": reference_date,
+            "PREFER_DATES_FROM": "past",
+            "RETURN_AS_TIMEZONE_AWARE": False,
         }
 
         results = self._search_dates(query, settings=settings)
@@ -137,11 +136,8 @@ class DateparserQueryAnalyzer(QueryAnalyzer):
             return QueryAnalysis(temporal_constraint=None)
 
         # Filter out false positives (common words parsed as dates)
-        false_positives = {'do', 'may', 'march', 'will', 'can', 'sat', 'sun', 'mon', 'tue', 'wed', 'thu', 'fri'}
-        valid_results = [
-            (text, date) for text, date in results
-            if text.lower() not in false_positives or len(text) > 3
-        ]
+        false_positives = {"do", "may", "march", "will", "can", "sat", "sun", "mon", "tue", "wed", "thu", "fri"}
+        valid_results = [(text, date) for text, date in results if text.lower() not in false_positives or len(text) > 3]
 
         if not valid_results:
             return QueryAnalysis(temporal_constraint=None)
@@ -153,84 +149,94 @@ class DateparserQueryAnalyzer(QueryAnalyzer):
         start_date = parsed_date.replace(hour=0, minute=0, second=0, microsecond=0)
         end_date = parsed_date.replace(hour=23, minute=59, second=59, microsecond=999999)
 
-        return QueryAnalysis(
-            temporal_constraint=TemporalConstraint(
-                start_date=start_date,
-                end_date=end_date
-            )
-        )
+        return QueryAnalysis(temporal_constraint=TemporalConstraint(start_date=start_date, end_date=end_date))
 
-    def _extract_period(
-        self, query: str, reference_date: datetime
-    ) -> Optional[TemporalConstraint]:
+    def _extract_period(self, query: str, reference_date: datetime) -> TemporalConstraint | None:
         """
         Extract period-based temporal expressions (week, month, year, weekend).
 
         These need special handling as they represent date ranges, not single dates.
         Supports multiple languages.
         """
+
         def constraint(start: datetime, end: datetime) -> TemporalConstraint:
             return TemporalConstraint(
                 start_date=start.replace(hour=0, minute=0, second=0, microsecond=0),
-                end_date=end.replace(hour=23, minute=59, second=59, microsecond=999999)
+                end_date=end.replace(hour=23, minute=59, second=59, microsecond=999999),
             )
 
         # Yesterday patterns (English, Spanish, Italian, French, German)
-        if re.search(r'\b(yesterday|ayer|ieri|hier|gestern)\b', query, re.IGNORECASE):
+        if re.search(r"\b(yesterday|ayer|ieri|hier|gestern)\b", query, re.IGNORECASE):
             d = reference_date - timedelta(days=1)
             return constraint(d, d)
 
         # Today patterns
-        if re.search(r'\b(today|hoy|oggi|aujourd\'?hui|heute)\b', query, re.IGNORECASE):
+        if re.search(r"\b(today|hoy|oggi|aujourd\'?hui|heute)\b", query, re.IGNORECASE):
             return constraint(reference_date, reference_date)
 
         # "a couple of days ago" / "a few days ago" patterns
         # These are imprecise so we create a range
-        if re.search(r'\b(a\s+)?couple\s+(of\s+)?days?\s+ago\b', query, re.IGNORECASE):
+        if re.search(r"\b(a\s+)?couple\s+(of\s+)?days?\s+ago\b", query, re.IGNORECASE):
             # "a couple of days" = approximately 2 days, give range of 1-3 days
             return constraint(reference_date - timedelta(days=3), reference_date - timedelta(days=1))
 
-        if re.search(r'\b(a\s+)?few\s+days?\s+ago\b', query, re.IGNORECASE):
+        if re.search(r"\b(a\s+)?few\s+days?\s+ago\b", query, re.IGNORECASE):
             # "a few days" = approximately 3-4 days, give range of 2-5 days
             return constraint(reference_date - timedelta(days=5), reference_date - timedelta(days=2))
 
         # "a couple of weeks ago" / "a few weeks ago" patterns
-        if re.search(r'\b(a\s+)?couple\s+(of\s+)?weeks?\s+ago\b', query, re.IGNORECASE):
+        if re.search(r"\b(a\s+)?couple\s+(of\s+)?weeks?\s+ago\b", query, re.IGNORECASE):
             # "a couple of weeks" = approximately 2 weeks, give range of 1-3 weeks
             return constraint(reference_date - timedelta(weeks=3), reference_date - timedelta(weeks=1))
 
-        if re.search(r'\b(a\s+)?few\s+weeks?\s+ago\b', query, re.IGNORECASE):
+        if re.search(r"\b(a\s+)?few\s+weeks?\s+ago\b", query, re.IGNORECASE):
             # "a few weeks" = approximately 3-4 weeks, give range of 2-5 weeks
             return constraint(reference_date - timedelta(weeks=5), reference_date - timedelta(weeks=2))
 
         # "a couple of months ago" / "a few months ago" patterns
-        if re.search(r'\b(a\s+)?couple\s+(of\s+)?months?\s+ago\b', query, re.IGNORECASE):
+        if re.search(r"\b(a\s+)?couple\s+(of\s+)?months?\s+ago\b", query, re.IGNORECASE):
             # "a couple of months" = approximately 2 months, give range of 1-3 months
             return constraint(reference_date - timedelta(days=90), reference_date - timedelta(days=30))
 
-        if re.search(r'\b(a\s+)?few\s+months?\s+ago\b', query, re.IGNORECASE):
+        if re.search(r"\b(a\s+)?few\s+months?\s+ago\b", query, re.IGNORECASE):
             # "a few months" = approximately 3-4 months, give range of 2-5 months
             return constraint(reference_date - timedelta(days=150), reference_date - timedelta(days=60))
 
         # Last week patterns (English, Spanish, Italian, French, German)
-        if re.search(r'\b(last\s+week|la\s+semana\s+pasada|la\s+settimana\s+scorsa|la\s+semaine\s+derni[eè]re|letzte\s+woche)\b', query, re.IGNORECASE):
+        if re.search(
+            r"\b(last\s+week|la\s+semana\s+pasada|la\s+settimana\s+scorsa|la\s+semaine\s+derni[eè]re|letzte\s+woche)\b",
+            query,
+            re.IGNORECASE,
+        ):
             start = reference_date - timedelta(days=reference_date.weekday() + 7)
             return constraint(start, start + timedelta(days=6))
 
         # Last month patterns
-        if re.search(r'\b(last\s+month|el\s+mes\s+pasado|il\s+mese\s+scorso|le\s+mois\s+dernier|letzten?\s+monat)\b', query, re.IGNORECASE):
+        if re.search(
+            r"\b(last\s+month|el\s+mes\s+pasado|il\s+mese\s+scorso|le\s+mois\s+dernier|letzten?\s+monat)\b",
+            query,
+            re.IGNORECASE,
+        ):
             first = reference_date.replace(day=1)
             end = first - timedelta(days=1)
             start = end.replace(day=1)
             return constraint(start, end)
 
         # Last year patterns
-        if re.search(r'\b(last\s+year|el\s+a[ñn]o\s+pasado|l\'anno\s+scorso|l\'ann[ée]e\s+derni[eè]re|letztes?\s+jahr)\b', query, re.IGNORECASE):
+        if re.search(
+            r"\b(last\s+year|el\s+a[ñn]o\s+pasado|l\'anno\s+scorso|l\'ann[ée]e\s+derni[eè]re|letztes?\s+jahr)\b",
+            query,
+            re.IGNORECASE,
+        ):
             year = reference_date.year - 1
             return constraint(datetime(year, 1, 1), datetime(year, 12, 31))
 
         # Last weekend patterns
-        if re.search(r'\b(last\s+weekend|el\s+fin\s+de\s+semana\s+pasado|lo\s+scorso\s+fine\s+settimana|le\s+week-?end\s+dernier|letztes?\s+wochenende)\b', query, re.IGNORECASE):
+        if re.search(
+            r"\b(last\s+weekend|el\s+fin\s+de\s+semana\s+pasado|lo\s+scorso\s+fine\s+settimana|le\s+week-?end\s+dernier|letztes?\s+wochenende)\b",
+            query,
+            re.IGNORECASE,
+        ):
             days_since_sat = (reference_date.weekday() + 2) % 7
             if days_since_sat == 0:
                 days_since_sat = 7
@@ -239,22 +245,22 @@ class DateparserQueryAnalyzer(QueryAnalyzer):
 
         # Month + Year patterns (e.g., "June 2024", "junio 2024", "giugno 2024")
         month_patterns = {
-            'january|enero|gennaio|janvier|januar': 1,
-            'february|febrero|febbraio|f[ée]vrier|februar': 2,
-            'march|marzo|mars|m[äa]rz': 3,
-            'april|abril|aprile|avril': 4,
-            'may|mayo|maggio|mai': 5,
-            'june|junio|giugno|juin|juni': 6,
-            'july|julio|luglio|juillet|juli': 7,
-            'august|agosto|ao[uû]t': 8,
-            'september|septiembre|settembre|septembre': 9,
-            'october|octubre|ottobre|octobre|oktober': 10,
-            'november|noviembre|novembre': 11,
-            'december|diciembre|dicembre|d[ée]cembre|dezember': 12,
+            "january|enero|gennaio|janvier|januar": 1,
+            "february|febrero|febbraio|f[ée]vrier|februar": 2,
+            "march|marzo|mars|m[äa]rz": 3,
+            "april|abril|aprile|avril": 4,
+            "may|mayo|maggio|mai": 5,
+            "june|junio|giugno|juin|juni": 6,
+            "july|julio|luglio|juillet|juli": 7,
+            "august|agosto|ao[uû]t": 8,
+            "september|septiembre|settembre|septembre": 9,
+            "october|octubre|ottobre|octobre|oktober": 10,
+            "november|noviembre|novembre": 11,
+            "december|diciembre|dicembre|d[ée]cembre|dezember": 12,
         }
 
         for pattern, month_num in month_patterns.items():
-            match = re.search(rf'\b({pattern})\s+(\d{{4}})\b', query, re.IGNORECASE)
+            match = re.search(rf"\b({pattern})\s+(\d{{4}})\b", query, re.IGNORECASE)
             if match:
                 year = int(match.group(2))
                 start = datetime(year, month_num, 1)
@@ -279,11 +285,7 @@ class TransformerQueryAnalyzer(QueryAnalyzer):
     - Model size: ~80M params (~300MB download)
     """
 
-    def __init__(
-        self,
-        model_name: str = "google/flan-t5-small",
-        device: str = "cpu"
-    ):
+    def __init__(self, model_name: str = "google/flan-t5-small", device: str = "cpu"):
         """
         Initialize T5 query analyzer.
 
@@ -304,11 +306,10 @@ class TransformerQueryAnalyzer(QueryAnalyzer):
             return
 
         try:
-            from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+            from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
         except ImportError:
             raise ImportError(
-                "transformers is required for TransformerQueryAnalyzer. "
-                "Install it with: pip install transformers"
+                "transformers is required for TransformerQueryAnalyzer. Install it with: pip install transformers"
             )
 
         logger.info(f"Loading query analyzer model: {self.model_name}...")
@@ -322,9 +323,7 @@ class TransformerQueryAnalyzer(QueryAnalyzer):
         """Lazy load the T5 model for temporal extraction (calls load())."""
         self.load()
 
-    def _extract_with_rules(
-        self, query: str, reference_date: datetime
-    ) -> Optional[TemporalConstraint]:
+    def _extract_with_rules(self, query: str, reference_date: datetime) -> TemporalConstraint | None:
         """
         Extract temporal expressions using rule-based patterns.
 
@@ -332,6 +331,7 @@ class TransformerQueryAnalyzer(QueryAnalyzer):
         patterns that need model-based extraction.
         """
         import re
+
         query_lower = query.lower()
 
         def get_last_weekday(weekday: int) -> datetime:
@@ -343,50 +343,60 @@ class TransformerQueryAnalyzer(QueryAnalyzer):
         def constraint(start: datetime, end: datetime) -> TemporalConstraint:
             return TemporalConstraint(
                 start_date=start.replace(hour=0, minute=0, second=0, microsecond=0),
-                end_date=end.replace(hour=23, minute=59, second=59, microsecond=999999)
+                end_date=end.replace(hour=23, minute=59, second=59, microsecond=999999),
             )
 
         # Yesterday
-        if re.search(r'\byesterday\b', query_lower):
+        if re.search(r"\byesterday\b", query_lower):
             d = reference_date - timedelta(days=1)
             return constraint(d, d)
 
         # Last week
-        if re.search(r'\blast\s+week\b', query_lower):
+        if re.search(r"\blast\s+week\b", query_lower):
             start = reference_date - timedelta(days=reference_date.weekday() + 7)
             return constraint(start, start + timedelta(days=6))
 
         # Last month
-        if re.search(r'\blast\s+month\b', query_lower):
+        if re.search(r"\blast\s+month\b", query_lower):
             first = reference_date.replace(day=1)
             end = first - timedelta(days=1)
             start = end.replace(day=1)
             return constraint(start, end)
 
         # Last year
-        if re.search(r'\blast\s+year\b', query_lower):
+        if re.search(r"\blast\s+year\b", query_lower):
             y = reference_date.year - 1
             return constraint(datetime(y, 1, 1), datetime(y, 12, 31))
 
         # Last weekend
-        if re.search(r'\blast\s+weekend\b', query_lower):
+        if re.search(r"\blast\s+weekend\b", query_lower):
             sat = get_last_weekday(5)
             return constraint(sat, sat + timedelta(days=1))
 
         # Last <weekday>
-        weekdays = {'monday': 0, 'tuesday': 1, 'wednesday': 2, 'thursday': 3,
-                    'friday': 4, 'saturday': 5, 'sunday': 6}
+        weekdays = {"monday": 0, "tuesday": 1, "wednesday": 2, "thursday": 3, "friday": 4, "saturday": 5, "sunday": 6}
         for name, num in weekdays.items():
-            if re.search(rf'\blast\s+{name}\b', query_lower):
+            if re.search(rf"\blast\s+{name}\b", query_lower):
                 d = get_last_weekday(num)
                 return constraint(d, d)
 
         # Month + Year: "June 2024", "in March 2023"
-        months = {'january': 1, 'february': 2, 'march': 3, 'april': 4, 'may': 5,
-                  'june': 6, 'july': 7, 'august': 8, 'september': 9, 'october': 10,
-                  'november': 11, 'december': 12}
+        months = {
+            "january": 1,
+            "february": 2,
+            "march": 3,
+            "april": 4,
+            "may": 5,
+            "june": 6,
+            "july": 7,
+            "august": 8,
+            "september": 9,
+            "october": 10,
+            "november": 11,
+            "december": 12,
+        }
         for name, num in months.items():
-            match = re.search(rf'\b{name}\s+(\d{{4}})\b', query_lower)
+            match = re.search(rf"\b{name}\s+(\d{{4}})\b", query_lower)
             if match:
                 year = int(match.group(1))
                 if num == 12:
@@ -397,9 +407,7 @@ class TransformerQueryAnalyzer(QueryAnalyzer):
 
         return None
 
-    def analyze(
-        self, query: str, reference_date: Optional[datetime] = None
-    ) -> QueryAnalysis:
+    def analyze(self, query: str, reference_date: datetime | None = None) -> QueryAnalysis:
         """
         Analyze query for temporal expressions.
 
@@ -435,11 +443,11 @@ class TransformerQueryAnalyzer(QueryAnalyzer):
         last_saturday = get_last_weekday(5)
 
         # Build prompt for T5
-        prompt = f"""Today is {reference_date.strftime('%Y-%m-%d')}. Extract date range or "none".
+        prompt = f"""Today is {reference_date.strftime("%Y-%m-%d")}. Extract date range or "none".
 
 June 2024 = 2024-06-01 to 2024-06-30
-yesterday = {yesterday.strftime('%Y-%m-%d')} to {yesterday.strftime('%Y-%m-%d')}
-last Saturday = {last_saturday.strftime('%Y-%m-%d')} to {last_saturday.strftime('%Y-%m-%d')}
+yesterday = {yesterday.strftime("%Y-%m-%d")} to {yesterday.strftime("%Y-%m-%d")}
+last Saturday = {last_saturday.strftime("%Y-%m-%d")} to {last_saturday.strftime("%Y-%m-%d")}
 what is the weather = none
 {query} ="""
 
@@ -448,13 +456,7 @@ what is the weather = none
         inputs = {k: v.to(self.device) for k, v in inputs.items()}
 
         with self._no_grad():
-            outputs = self._model.generate(
-                **inputs,
-                max_new_tokens=30,
-                num_beams=3,
-                do_sample=False,
-                temperature=1.0
-            )
+            outputs = self._model.generate(**inputs, max_new_tokens=30, num_beams=3, do_sample=False, temperature=1.0)
 
         result = self._tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
 
@@ -466,14 +468,14 @@ what is the weather = none
         """Get torch.no_grad context manager."""
         try:
             import torch
+
             return torch.no_grad()
         except ImportError:
             from contextlib import nullcontext
+
             return nullcontext()
 
-    def _parse_generated_output(
-        self, result: str, reference_date: datetime
-    ) -> Optional[TemporalConstraint]:
+    def _parse_generated_output(self, result: str, reference_date: datetime) -> TemporalConstraint | None:
         """
         Parse T5 generated output into TemporalConstraint.
 
@@ -492,7 +494,8 @@ what is the weather = none
         try:
             # Parse "YYYY-MM-DD to YYYY-MM-DD"
             import re
-            pattern = r'(\d{4}-\d{2}-\d{2})\s+to\s+(\d{4}-\d{2}-\d{2})'
+
+            pattern = r"(\d{4}-\d{2}-\d{2})\s+to\s+(\d{4}-\d{2}-\d{2})"
             match = re.search(pattern, result, re.IGNORECASE)
 
             if match:
@@ -513,7 +516,7 @@ what is the weather = none
 
                 return TemporalConstraint(start_date=start_date, end_date=end_date)
 
-        except (ValueError, AttributeError) as e:
+        except (ValueError, AttributeError):
             return None
 
         return None