src/sanitizer.py

from __future__ import annotations

import hashlib
import re
import uuid

from src.schemas import SanitizedTask, UserMemory


PRIVATE_PRONOUN_PATTERNS = [
    r"\bmy name is\s+[^,.!?]+",
    r"\bI am called\s+[^,.!?]+",
    r"\bI live in\s+[^,.!?]+",
    r"\bmy email is\s+[^,.!?]+",
]


def _keywords(text: str) -> list[str]:
    words = re.findall(r"[A-Za-z][A-Za-z0-9_-]{3,}", text.lower())
    stop = {"this", "that", "with", "from", "have", "what", "when", "where", "these", "those", "about"}
    out: list[str] = []
    for word in words:
        if word not in stop and word not in out:
            out.append(word)
    return out[:12]


def sanitize_task(user_message: str, memory: UserMemory, approved_context: list[str] | None = None) -> SanitizedTask:
    sanitized = user_message
    removed: list[str] = []
    approved_context = approved_context or []

    for pattern in PRIVATE_PRONOUN_PATTERNS:
        sanitized, count = re.subn(pattern, "[private detail hidden]", sanitized, flags=re.IGNORECASE)
        if count:
            removed.append(pattern)

    for fact in memory.accepted():
        fact_text = fact.text.strip()
        if fact_text and fact_text.lower() in sanitized.lower() and fact_text not in approved_context:
            sanitized = re.sub(re.escape(fact_text), "[private memory hidden]", sanitized, flags=re.IGNORECASE)
            removed.append(fact_text)
        if fact.kind == "profile" and fact_text.lower().startswith("user name:"):
            name = fact_text.split(":", 1)[1].strip()
            if name and name.lower() in sanitized.lower() and fact_text not in approved_context:
                sanitized = re.sub(re.escape(name), "[private name hidden]", sanitized, flags=re.IGNORECASE)
                removed.append(name)

    return SanitizedTask(
        run_id=f"run_{uuid.uuid4().hex[:12]}",
        original_question_hash=hashlib.sha256(user_message.encode("utf-8")).hexdigest()[:16],
        sanitized_query=sanitized.strip(),
        neutral_keywords=_keywords(sanitized),
        removed_or_hidden_terms=removed,
        user_approved_context=approved_context,
    )