extraction.py

import fitz
import docx
import tempfile
import os
import re

from typing import List, Dict, Any

from utils import (
    extract_latex_blocks,
    extract_equations_with_llm
)


# =====================================================
# PDF READER
# =====================================================

def read_pdf(file_path: str) -> str:

    text = ""

    try:

        doc = fitz.open(file_path)

        for page in doc:

            try:
                page_text = page.get_text()

                if page_text:
                    text += page_text + "\n"

            except Exception:
                continue

        doc.close()

    except Exception as e:

        print(f"ERREUR PDF: {e}")

    return text


# =====================================================
# DOCX READER
# =====================================================

def read_docx(uploaded_file) -> str:

    try:

        document = docx.Document(uploaded_file)

        paragraphs = []

        for p in document.paragraphs:

            if p.text.strip():
                paragraphs.append(p.text)

        return "\n".join(paragraphs)

    except Exception as e:

        print(f"ERREUR DOCX: {e}")

        return ""


# =====================================================
# TEXT CLEANER
# =====================================================

def clean_text(text: str) -> str:

    if not text:
        return ""

    text = text.replace("\x00", " ")

    text = re.sub(r"\s+", " ", text)

    return text.strip()


# =====================================================
# NORMALIZE EQUATIONS
# =====================================================

def normalize_equation(eq: Any) -> str:

    if isinstance(eq, dict):

        eq = eq.get("latex", "")

    if not isinstance(eq, str):
        eq = str(eq)

    eq = eq.strip()

    eq = eq.replace("\n", " ")

    return eq


# =====================================================
# BUILD RESULTS
# =====================================================

def build_equation_results(
    equations: List[Any],
    context: str
) -> List[Dict[str, Any]]:

    results = []

    seen = set()

    for i, eq in enumerate(equations):

        latex = normalize_equation(eq)

        if not latex:
            continue

        if latex in seen:
            continue

        seen.add(latex)

        results.append({
            "id": f"eq_{i}",
            "latex": latex,
            "context": context[:500]
        })

    return results


# =====================================================
# HEURISTIC EXTRACTION
# =====================================================

def heuristic_equation_detection(text: str):

    candidates = []

    lines = text.split("\n")

    patterns = [
        r".*=.*",
        r".+\+.+",
        r".+\-.+",
        r".+\\frac.+",
        r".+\\sum.+",
        r".+\\int.+",
        r".+\^.+",
    ]

    for line in lines:

        line = line.strip()

        if len(line) < 3:
            continue

        for pattern in patterns:

            if re.match(pattern, line):

                candidates.append(line)

                break

    return candidates


# =====================================================
# EXTRACT FROM TEXT
# =====================================================

def extract_from_text(
    text: str,
    use_llm: bool = True
) -> List[Dict[str, Any]]:

    text = clean_text(text)

    if not text:

        return []

    equations = []

    # -------------------------------------------------
    # LATEX EXTRACTION
    # -------------------------------------------------

    latex_equations = extract_latex_blocks(text)

    equations.extend(latex_equations)

    # -------------------------------------------------
    # HEURISTIC EXTRACTION
    # -------------------------------------------------

    if len(equations) == 0:

        heuristic_eqs = heuristic_equation_detection(text)

        equations.extend(heuristic_eqs)

    # -------------------------------------------------
    # LLM EXTRACTION
    # -------------------------------------------------

    if (
        use_llm
        and len(equations) == 0
        and len(text) > 30
    ):

        try:

            llm_eqs = extract_equations_with_llm(text)

            if isinstance(llm_eqs, list):

                equations.extend(llm_eqs)

        except Exception as e:

            print(f"ERREUR LLM EXTRACTION: {e}")

    # -------------------------------------------------
    # FALLBACK
    # -------------------------------------------------

    if len(equations) == 0:

        equations = [
            "E = mc^2"
        ]

    return build_equation_results(
        equations[:20],
        text
    )


# =====================================================
# EXTRACT FROM FILE
# =====================================================

def extract_from_file(
    uploaded_file
) -> List[Dict[str, Any]]:

    text = ""

    filename = uploaded_file.name.lower()

    # -------------------------------------------------
    # PDF
    # -------------------------------------------------

    if filename.endswith(".pdf"):

        with tempfile.NamedTemporaryFile(
            delete=False,
            suffix=".pdf"
        ) as tmp:

            tmp.write(uploaded_file.read())

            tmp_path = tmp.name

        try:

            text = read_pdf(tmp_path)

        finally:

            if os.path.exists(tmp_path):
                os.unlink(tmp_path)

    # -------------------------------------------------
    # DOCX
    # -------------------------------------------------

    elif filename.endswith(".docx"):

        text = read_docx(uploaded_file)

    # -------------------------------------------------
    # TXT
    # -------------------------------------------------

    else:

        try:

            text = uploaded_file.read().decode(
                "utf-8",
                errors="ignore"
            )

        except Exception:

            text = ""

    return extract_from_text(text)