gdpr_logic.py

import json
import re
import time
from pathlib import Path
from typing import Any
from urllib.parse import urljoin

import requests
import torch
from bs4 import BeautifulSoup, Tag
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
from transformers import AutoModelForCausalLM, AutoTokenizer


GDPR_INDEX_URL = "https://gdpr.eu/tag/gdpr/"
CACHE_FILE = Path("gdpr_articles_cache.json")
MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
MAX_CONTEXT_ARTICLES = 4
REQUEST_DELAY_SECONDS = 0.15

QUICK_QUESTIONS = [
    "What is GDPR?",
    "What are the main GDPR principles?",
    "What rights do data subjects have?",
    "When is a Data Protection Officer required?",
    "What makes consent valid under the GDPR?",
    "How quickly must a personal data breach be reported?",
]


def clean_text(text: str) -> str:
    return re.sub(r"\s+", " ", text).strip()


def create_session() -> requests.Session:
    session = requests.Session()
    session.headers.update(
        {
            "User-Agent": (
                "Mozilla/5.0 (compatible; GDPR-Assistant/1.0)"
            ),
            "Accept-Language": "en-US,en;q=0.9",
        }
    )
    return session


def download_page(session: requests.Session, url: str) -> BeautifulSoup:
    response = session.get(url, timeout=40)
    response.raise_for_status()
    return BeautifulSoup(response.text, "lxml")


def get_article_links(
    session: requests.Session,
) -> dict[int, dict[str, Any]]:
    soup = download_page(session, GDPR_INDEX_URL)

    articles: dict[int, dict[str, Any]] = {}
    article_pattern = re.compile(
        r"\bArt\.\s*(\d+)\s+GDPR\b",
        flags=re.IGNORECASE,
    )

    for link in soup.find_all("a", href=True):
        link_text = clean_text(link.get_text(" ", strip=True))
        match = article_pattern.search(link_text)

        if not match:
            continue

        article_number = int(match.group(1))
        if not 1 <= article_number <= 99:
            continue

        articles[article_number] = {
            "number": article_number,
            "title": link_text,
            "url": urljoin(GDPR_INDEX_URL, link["href"]),
        }

    return dict(sorted(articles.items()))


def find_article_heading(
    soup: BeautifulSoup,
    article_number: int,
) -> Tag | None:
    heading_pattern = re.compile(
        rf"\bArt\.\s*{article_number}\s+GDPR\b",
        flags=re.IGNORECASE,
    )

    for heading_name in ("h1", "h2", "h3"):
        for heading in soup.find_all(heading_name):
            heading_text = clean_text(
                heading.get_text(" ", strip=True)
            )
            if heading_pattern.search(heading_text):
                return heading

    return None


def is_stop_heading(element: Tag) -> bool:
    if element.name not in {
        "h1", "h2", "h3", "h4", "h5", "h6"
    }:
        return False

    text = clean_text(
        element.get_text(" ", strip=True)
    ).lower()

    stop_phrases = (
        "suitable recitals",
        "related posts",
        "about gdpr.eu",
        "getting started",
        "templates",
        "technical review",
    )

    return any(phrase in text for phrase in stop_phrases)


def extract_article_content(
    heading: Tag,
) -> str:
    content_parts: list[str] = []
    seen: set[str] = set()

    for element in heading.find_all_next():
        if not isinstance(element, Tag):
            continue

        if is_stop_heading(element):
            break

        if element.name not in {"p", "li"}:
            continue

        text = clean_text(
            element.get_text(" ", strip=True)
        )

        if not text:
            continue

        if re.match(
            r"^(Chapter|Art\.\s*\d+\s+GDPR)",
            text,
            flags=re.IGNORECASE,
        ):
            continue

        normalized = text.casefold()
        if normalized in seen:
            continue

        seen.add(normalized)
        content_parts.append(text)

    return "\n".join(content_parts)


def extract_single_article(
    session: requests.Session,
    article_number: int,
    article_info: dict[str, Any],
) -> dict[str, Any] | None:
    soup = download_page(session, article_info["url"])
    heading = find_article_heading(soup, article_number)

    if heading is None:
        return None

    title = clean_text(heading.get_text(" ", strip=True))
    content = extract_article_content(heading)

    if not content:
        return None

    return {
        "number": article_number,
        "title": title,
        "url": article_info["url"],
        "content": content,
        "search_text": (
            f"Article {article_number}\n"
            f"{title}\n"
            f"{content}"
        ),
    }


def save_articles_to_cache(
    articles: list[dict[str, Any]],
) -> None:
    with CACHE_FILE.open("w", encoding="utf-8") as file:
        json.dump(
            articles,
            file,
            ensure_ascii=False,
            indent=2,
        )


def load_cache() -> list[dict[str, Any]]:
    if not CACHE_FILE.exists():
        return []

    try:
        with CACHE_FILE.open("r", encoding="utf-8") as file:
            data = json.load(file)
        return data if isinstance(data, list) else []
    except (OSError, json.JSONDecodeError):
        return []


def extract_all_articles(
    force_refresh: bool = False,
) -> list[dict[str, Any]]:
    if not force_refresh:
        cached = load_cache()
        if cached:
            return cached

    session = create_session()
    article_links = get_article_links(session)

    if not article_links:
        raise RuntimeError(
            "No GDPR article links were found."
        )

    extracted: list[dict[str, Any]] = []

    for number, info in article_links.items():
        try:
            article = extract_single_article(
                session,
                number,
                info,
            )
            if article:
                extracted.append(article)
        except requests.RequestException as error:
            print(
                f"Article {number} could not be downloaded: {error}"
            )

        time.sleep(REQUEST_DELAY_SECONDS)

    if not extracted:
        raise RuntimeError(
            "No GDPR article content could be extracted."
        )

    save_articles_to_cache(extracted)
    return extracted


class GDPRSearchIndex:
    def __init__(self, articles: list[dict[str, Any]]):
        if not articles:
            raise ValueError(
                "At least one article is required."
            )

        self.articles = articles
        self.vectorizer = TfidfVectorizer(
            stop_words="english",
            ngram_range=(1, 2),
            max_features=30000,
        )

        documents = [
            article["search_text"]
            for article in articles
        ]

        self.article_matrix = (
            self.vectorizer.fit_transform(documents)
        )

    def search(
        self,
        query: str,
        limit: int = MAX_CONTEXT_ARTICLES,
    ) -> list[dict[str, Any]]:
        query_vector = self.vectorizer.transform([query])

        scores = cosine_similarity(
            query_vector,
            self.article_matrix,
        ).flatten()

        ranked_indexes = scores.argsort()[::-1]
        results: list[dict[str, Any]] = []

        for index in ranked_indexes:
            score = float(scores[index])
            if score <= 0:
                continue

            article = dict(self.articles[index])
            article["score"] = score
            results.append(article)

            if len(results) >= limit:
                break

        return results


print(f"Loading model: {MODEL_NAME}")

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_NAME,
    trust_remote_code=True,
)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype="auto",
    device_map="auto",
    trust_remote_code=True,
)

print("Model loaded.")


def build_context(
    articles: list[dict[str, Any]],
) -> str:
    return "\n\n---\n\n".join(
        (
            f"ARTICLE {article['number']}\n"
            f"TITLE: {article['title']}\n"
            f"SOURCE: {article['url']}\n\n"
            f"{article['content']}"
        )
        for article in articles
    )


def generate_answer(
    question: str,
    articles: list[dict[str, Any]],
) -> str:
    context = build_context(articles)

    messages = [
        {
            "role": "system",
            "content": (
                "You are a GDPR document assistant. "
                "Answer only from the supplied GDPR article text. "
                "If the answer is not in the supplied articles, say so. "
                "Explain the answer clearly, mention relevant article "
                "numbers, and do not present the response as formal "
                "legal advice."
            ),
        },
        {
            "role": "user",
            "content": (
                f"GDPR CONTEXT:\n\n{context}\n\n"
                f"QUESTION:\n{question}"
            ),
        },
    ]

    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )

    model_inputs = tokenizer(
        prompt,
        return_tensors="pt",
        truncation=True,
        max_length=7000,
    ).to(model.device)

    with torch.inference_mode():
        output_ids = model.generate(
            **model_inputs,
            max_new_tokens=350,
            do_sample=True,
            temperature=0.2,
            top_p=0.9,
            repetition_penalty=1.08,
            pad_token_id=tokenizer.eos_token_id,
        )

    generated_ids = output_ids[
        :,
        model_inputs["input_ids"].shape[1]:
    ]

    return tokenizer.batch_decode(
        generated_ids,
        skip_special_tokens=True,
    )[0].strip()


def format_sources(
    articles: list[dict[str, Any]],
) -> str:
    return "\n".join(
        (
            f"- [Article {article['number']}: "
            f"{article['title']}]({article['url']})"
        )
        for article in articles
    )


gdpr_articles: list[dict[str, Any]] = []
search_index: GDPRSearchIndex | None = None


def initialize_data() -> str:
    global gdpr_articles, search_index

    try:
        gdpr_articles = extract_all_articles(
            force_refresh=False
        )
        search_index = GDPRSearchIndex(gdpr_articles)

        return (
            f"✅ GDPR knowledge base ready: "
            f"**{len(gdpr_articles)} articles loaded**."
        )
    except Exception as error:
        gdpr_articles = []
        search_index = None

        return (
            "❌ Could not initialize the GDPR knowledge base. "
            f"`{error}`"
        )


def refresh_data() -> str:
    global gdpr_articles, search_index

    try:
        gdpr_articles = extract_all_articles(
            force_refresh=True
        )
        search_index = GDPRSearchIndex(gdpr_articles)

        return (
            f"✅ GDPR knowledge base refreshed: "
            f"**{len(gdpr_articles)} articles loaded**."
        )
    except Exception as error:
        return (
            "❌ Could not refresh the GDPR website. "
            f"`{error}`"
        )


def ask_question(
    question: str,
    history: list[dict[str, str]] | None,
):
    history = history or []

    if not question or not question.strip():
        return history, ""

    question = question.strip()
    history.append(
        {
            "role": "user",
            "content": question,
        }
    )

    if search_index is None:
        history.append(
            {
                "role": "assistant",
                "content": (
                    "The GDPR knowledge base is unavailable. "
                    "Please refresh the website and try again."
                ),
            }
        )
        return history, ""

    relevant_articles = search_index.search(question)

    if not relevant_articles:
        history.append(
            {
                "role": "assistant",
                "content": (
                    "I could not find relevant information in "
                    "the loaded GDPR articles."
                ),
            }
        )
        return history, ""

    try:
        answer = generate_answer(
            question,
            relevant_articles,
        )

        sources = format_sources(relevant_articles)

        complete_answer = (
            f"{answer}\n\n"
            f"### Sources used\n{sources}\n\n"
            "*This response is based on the loaded GDPR articles "
            "and is not formal legal advice.*"
        )
    except Exception as error:
        complete_answer = (
            "Relevant GDPR articles were found, but the model "
            f"could not generate an answer. `{error}`"
        )

    history.append(
        {
            "role": "assistant",
            "content": complete_answer,
        }
    )

    return history, ""


def clear_chat():
    return [], ""


INITIAL_STATUS = initialize_data()