A-Team-Security-Infra-Agent-Engineer/backend/services/embeddings.py

"""Embeddings service — chunking, metadata, ADB vector ingest."""
import os, json, uuid, time, re, hashlib
from pathlib import Path
from typing import Optional, List, Dict, Any

from config import DATA, REPORTS, WALLET_DIR, log, _chat_executor, _EMBED_STATUS_DIR
from database import db
from auth.jwt_auth import _config_log, _verify_config_access
from services.genai import (
    _get_adb_connection, _resolve_embed_config, _embed_text,
    _DIM_TO_MODEL, _get_table_embedding_dim, _get_active_adb_configs,
    _get_tables_for_config,
)

def _build_metadata_json(tenancy: str = "", compartments: str = "", section: str = "",
                         report_date: str = "", user_id: str = "", extra: dict = None) -> str:
    """Build a structured JSON metadata string for vector embeddings."""
    meta = {}
    if tenancy:
        meta["tenancy"] = tenancy
    if compartments:
        meta["compartments"] = compartments
    if section:
        meta["section"] = section
    if report_date:
        meta["report_date"] = report_date
    if user_id:
        meta["user_id"] = user_id
    if extra:
        meta.update(extra)
    return json.dumps(meta, ensure_ascii=False) if meta else ""


def _auto_register_table(adb_config_id: str, table_name: str, description: str = ""):
    """Auto-register a table in adb_vector_tables if not already present."""
    if not table_name:
        return
    with db() as c:
        exists = c.execute("SELECT 1 FROM adb_vector_tables WHERE adb_config_id=? AND table_name=? COLLATE NOCASE",
                           (adb_config_id, table_name)).fetchone()
        if not exists:
            c.execute("INSERT INTO adb_vector_tables (id, adb_config_id, table_name, description) VALUES (?,?,?,?)",
                      (str(uuid.uuid4()), adb_config_id, table_name, description))
            log.info(f"Auto-registered table '{table_name}' for ADB config {adb_config_id}")

def _ingest_documents_task(cfg: dict, genai_cfg: dict, documents: list, user_id: str, username: str,
                           table_name: str = None, tenancy: str = None, compartments: str = None,
                           report_date: str = None, task_id: str = None):
    """Background task: embed and insert documents into ADB via OCI GenAI.
    Tenancy and compartments are stored in METADATA as structured JSON for filtering."""
    import array
    emb_model = cfg.get("embedding_model_id", "cohere.embed-v4.0")
    table_name = table_name or cfg.get("table_name", "")
    # Auto-detect embedding dimension from existing table data and use matching model
    try:
        actual_dim = _get_table_embedding_dim(cfg, table_name)
        if actual_dim and actual_dim in _DIM_TO_MODEL:
            detected_model = _DIM_TO_MODEL[actual_dim]
            if detected_model != emb_model:
                log.info(f"Ingest: table '{table_name}' has {actual_dim} dims, switching model from {emb_model} to {detected_model}")
                emb_model = detected_model
    except Exception as e:
        log.warning(f"Ingest: failed to detect dimension for '{table_name}': {e}")
    total = len(documents)
    # Track status
    if task_id:
        _set_embed_status(task_id, {"status": "running", "table": table_name, "tenancy": tenancy or "",
                                       "inserted": 0, "total": total, "user_id": user_id, "message": "Iniciando embedding..."})
    # Auto-register table so it appears in multi-table RAG search
    _auto_register_table(cfg["id"], table_name)
    conn = _get_adb_connection(cfg)
    try:
        cur = conn.cursor()
        inserted = 0
        for i, doc in enumerate(documents):
            try:
                content = doc.get("content", "")
                if not content: continue
                embedding = _embed_text(content, genai_cfg, emb_model)
                vec = array.array('f', [float(x) for x in embedding])
                # Build structured metadata with tenancy isolation
                doc_tenancy = tenancy or doc.get("tenancy", "")
                doc_compartments = compartments or doc.get("compartments", "")
                metadata = _build_metadata_json(
                    tenancy=doc_tenancy,
                    compartments=doc_compartments,
                    section=doc.get("section", ""),
                    report_date=report_date or "",
                    user_id=user_id,
                    extra={"legacy_metadata": doc.get("metadata", "")} if doc.get("metadata") else None
                )
                cur.execute(f"""
                    INSERT INTO "{table_name}" (ID, TEXT, EMBEDDING, METADATA)
                    VALUES (HEXTORAW(:1), :2, :3, :4)
                """, [uuid.uuid4().hex.upper(), content, vec, metadata])
                inserted += 1
                if task_id:
                    _update_embed_status(task_id, {"inserted": inserted, "message": f"Embedding {inserted}/{total}..."})
            except Exception as e:
                log.error(f"Failed to ingest document: {e}")
        conn.commit()
        cur.close()
        msg = f"{inserted}/{total} documentos ingeridos em {table_name}" + (f" (tenancy: {tenancy})" if tenancy else "")
        log.info(f"Ingested {inserted}/{total} documents into {table_name}" +
                 (f" (tenancy={tenancy})" if tenancy else ""))
        _audit(user_id, username, "ingest_documents", cfg["id"], f"{inserted} documents")
        _config_log("adb", cfg["id"], cfg.get("config_name"), "success", "ingest", msg, user_id, username)
        if task_id:
            _update_embed_status(task_id, {"status": "done", "inserted": inserted, "message": msg})
    except Exception as e:
        log.error(f"Ingestion task failed: {e}")
        _config_log("adb", cfg["id"], cfg.get("config_name"), "error", "ingest", str(e)[:500], user_id, username)
        if task_id:
            _update_embed_status(task_id, {"status": "error", "message": str(e)[:300]})
    finally:
        conn.close()

# ── Embeddings ────────────────────────────────────────────────────────────────
def _chunk_report_by_section(report_data: dict) -> list:
    """Chunk a CIS report into documents grouped by section."""
    if isinstance(report_data, str):
        report_data = json.loads(report_data)
    if isinstance(report_data, list):
        report_data = {"findings": {str(i): item for i, item in enumerate(report_data)}, "tenancy": "unknown"}
    findings = report_data.get("findings", {})
    tenancy = report_data.get("tenancy", "unknown")
    generated_at = report_data.get("generated_at", "")
    regions = report_data.get("regions", [])
    compartments = report_data.get("compartments", [])
    # Build context header for each chunk
    ctx_parts = [f"Tenancy: {tenancy}"]
    if regions:
        ctx_parts.append(f"Regions: {', '.join(regions)}")
    if compartments:
        ctx_parts.append(f"Compartments: {', '.join(compartments[:50])}")
    ctx_header = "\n".join(ctx_parts)
    sections = {}
    for cid, check in findings.items():
        sec = check.get("section", "Other")
        sections.setdefault(sec, [])
        sections[sec].append(check)
    documents = []
    for section_name, checks in sections.items():
        passed = sum(1 for c in checks if c.get("status") == "PASS")
        failed = sum(1 for c in checks if c.get("status") == "FAIL")
        review = sum(1 for c in checks if c.get("status") == "REVIEW")
        lines = [ctx_header, "", f"Section: {section_name}", f"Total checks: {len(checks)}, Passed: {passed}, Failed: {failed}, Review: {review}", ""]
        for c in checks:
            status = c.get("status", "REVIEW")
            lines.append(f"- [{c.get('id', '')}] {c.get('title', '')} — Status: {status}")
            if c.get("findings"):
                for f in c["findings"]:
                    lines.append(f"  Finding: {f}")
        documents.append({
            "content": "\n".join(lines),
            "source": f"CIS Report - {tenancy} - {generated_at}",
            "section": section_name,
            "tenancy": tenancy,
            "compartments": ", ".join(compartments[:50]),
            "metadata": f"tenancy: {tenancy}, section: {section_name}, total: {len(checks)}, passed: {passed}, failed: {failed}, review: {review}"
        })
    return documents

def _chunk_cis_pdf(text: str, filename: str, target_chars: int = 7000, overlap_chars: int = 500) -> list:
    """Chunk a CIS Foundations Benchmark PDF by recommendation number.
    Each recommendation (1.1, 1.2, etc.) becomes one or more chunks with overlap.
    Port of the JavaScript embedding pipeline."""
    import re as _re

    def normalize(t):
        t = t.replace('\r', '\n')
        t = _re.sub(r'[ \t]+\n', '\n', t)
        t = _re.sub(r'\n{3,}', '\n\n', t)
        return t.strip()

    def strip_page_headers(t):
        # Remove "Page XX" both standalone and at start of lines
        t = _re.sub(r'^\s*Page\s+\d+\s*$', '', t, flags=_re.MULTILINE | _re.IGNORECASE)
        t = _re.sub(r'^Page\s+\d+\s+', '', t, flags=_re.MULTILINE | _re.IGNORECASE)
        return t

    def remove_toc(t):
        # Remove everything from "Table of Contents" up to the actual recommendations section
        # The real content starts with "Recommendations\n1 Identity" or "Profile Applicability"
        toc_start = _re.search(r'\bTable of Contents\b', t, _re.IGNORECASE)
        if not toc_start:
            return t
        # Find where actual recommendation content begins
        content_start = _re.search(r'\bRecommendations\s*\n\s*1\s+Identity', t[toc_start.start():], _re.IGNORECASE)
        if not content_start:
            content_start = _re.search(r'\bProfile Applicability\b', t[toc_start.start():], _re.IGNORECASE)
        if not content_start:
            content_start = _re.search(r'\bOverview\b', t[toc_start.start():], _re.IGNORECASE)
        if not content_start:
            return t
        end_pos = toc_start.start() + content_start.start()
        if end_pos <= toc_start.start():
            return t
        return normalize(t[:toc_start.start()] + '\n\n' + t[end_pos:])

    def is_chapter_header(line):
        l = line.strip()
        return bool(_re.match(r'^\d+\s+[A-Za-z].+', l)) and not _re.match(r'^\d+\.\d+', l)

    def is_rec_header_start(line):
        l = line.strip()
        # Must be "1.1 Word..." but NOT a TOC line (with dots/page numbers)
        if not _re.match(r'^\d+\.\d+(\.\d+)?\s+[A-Z]', l):
            return False
        # Skip TOC lines: contain "...." or end with a page number
        if '....' in l or _re.search(r'\.\s*\d+\s*$', l):
            return False
        return True

    def header_looks_complete(h):
        # Complete if has (Manual)/(Automated) or ends with a closing paren
        if _re.search(r'\(\s*(Manual|Automated)\s*\)', h, _re.IGNORECASE):
            return True
        # Also stop if next line starts a known section like "Profile Applicability"
        return False

    def chunk_text(t):
        if not t:
            return []
        paragraphs = [p.strip() for p in t.split('\n\n') if p.strip()]
        chunks = []
        buf = ""
        def push():
            nonlocal buf
            b = buf.strip()
            if b:
                chunks.append(b)
            buf = ""
        for p in paragraphs:
            if len(p) > target_chars:
                push()
                i = 0
                while i < len(p):
                    chunks.append(p[i:i + target_chars].strip())
                    i += max(1, target_chars - overlap_chars)
                continue
            candidate = f"{buf}\n\n{p}" if buf else p
            if len(candidate) <= target_chars:
                buf = candidate
            else:
                push()
                if chunks and overlap_chars > 0:
                    prev = chunks[-1]
                    overlap = prev[max(0, len(prev) - overlap_chars):]
                    buf = f"{overlap}\n\n{p}".strip()
                else:
                    buf = p
        push()
        return chunks

    def remove_appendix(t):
        """Remove appendix sections (Assessment Status, Change History, etc.) that pollute embeddings."""
        for marker in [r'\bAppendix\b', r'\bAssessment Status\b', r'\bChange History\b',
                       r'\bCIS Controls v\d', r'\bDate\s+Version\s+Changes']:
            m = _re.search(marker, t, _re.IGNORECASE)
            if m and m.start() > len(t) * 0.7:  # only cut if in last 30% of doc
                t = t[:m.start()].rstrip()
                break
        return t

    # Pipeline
    text = normalize(text)
    text = strip_page_headers(text)
    text = remove_toc(text)
    text = remove_appendix(text)
    lines = text.split('\n')

    # Segment by recommendation
    segments = []
    current = None
    current_chapter = ""

    i = 0
    while i < len(lines):
        line = lines[i]
        if is_chapter_header(line):
            current_chapter = line.strip()
        if is_rec_header_start(line):
            if current:
                segments.append(current)
            header = line.strip()
            j = i + 1
            while j < len(lines) and not header_looks_complete(header):
                next_line = lines[j].strip()
                if is_rec_header_start(next_line):
                    break
                # Stop consuming if we hit a known section start
                if next_line.startswith('Profile Applicability') or next_line.startswith('Description:'):
                    break
                if next_line:
                    header = _re.sub(r'\s+', ' ', f"{header} {next_line}").strip()
                j += 1
            i = j - 1
            current = {"header": header, "chapter": current_chapter, "body_lines": []}
            i += 1
            continue
        if current:
            current["body_lines"].append(line)
        i += 1

    if current:
        segments.append(current)

    # Generate chunks
    documents = []
    for seg in segments:
        body = normalize('\n'.join(seg["body_lines"]))
        if not body:
            continue
        rec_match = _re.match(r'^(\d+(\.\d+)+)', seg["header"])
        rec_number = rec_match.group(1) if rec_match else "unknown"
        canonical = normalize('\n'.join(filter(None, [
            f"Recommendation: {seg['header']}",
            f"Chapter: {seg['chapter']}" if seg['chapter'] else "",
            "",
            body,
        ])))
        chunks = chunk_text(canonical)
        for idx, chunk in enumerate(chunks):
            documents.append({
                "content": chunk,
                "source": filename,
                "metadata": json.dumps({
                    "filename": filename,
                    "recommendationNumber": rec_number,
                    "chapter": seg["chapter"],
                    "source": "CIS-OCI-PDF",
                    "chunkIndex": idx + 1,
                    "chunkCount": len(chunks),
                }),
            })

    log.info(f"CIS PDF chunking: {len(segments)} recommendations → {len(documents)} chunks from {filename}")
    return documents


def _chunk_text_file(text: str, filename: str, chunk_size: int = 1000, overlap: int = 200) -> list:
    """Split text into chunks by paragraphs with overlap to avoid losing context at boundaries."""
    paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]
    documents = []
    current_chunk = ""
    prev_tail = ""  # last N chars of previous chunk for overlap
    chunk_num = 1
    for para in paragraphs:
        if len(current_chunk) + len(para) + 2 > chunk_size and current_chunk:
            documents.append({"content": current_chunk, "source": filename, "metadata": f"chunk: {chunk_num}"})
            chunk_num += 1
            # Keep overlap from end of current chunk
            prev_tail = current_chunk[-overlap:] if len(current_chunk) > overlap else current_chunk
            current_chunk = prev_tail + "\n\n" + para
        else:
            current_chunk = current_chunk + "\n\n" + para if current_chunk else para
    if current_chunk:
        documents.append({"content": current_chunk, "source": filename, "metadata": f"chunk: {chunk_num}"})
    return documents

def _get_adb_and_genai(vid: str, oci_config_id: str = None, user_id: str = None):
    """Load ADB config and resolve embed config (scoped to user_id).
    Priority: ADB.genai_config_id → genai by oci_config_id → oci_config directly → user's default."""
    with db() as c:
        cfg = c.execute("SELECT * FROM adb_vector_configs WHERE id=?", (vid,)).fetchone()
    if not cfg: raise HTTPException(404, "ADB config not found")
    cfg = dict(cfg)
    genai_cfg = None
    if cfg.get("genai_config_id"):
        with db() as c:
            row = c.execute("SELECT * FROM genai_configs WHERE id=?", (cfg["genai_config_id"],)).fetchone()
            if row: genai_cfg = dict(row)
    gc = _resolve_embed_config(oci_config_id=oci_config_id, genai_cfg=genai_cfg, user_id=user_id or cfg.get("user_id"))
    return cfg, gc



def _chunk_summary_csv(csv_path: str, tenancy: str, extract_date: str) -> list:
    """Chunk the cis_summary_report.csv into documents with tenancy/date metadata.
    Each row becomes a document with structured content for vector search."""
    import csv as csvmod
    p = Path(csv_path)
    if not p.exists():
        return []
    documents = []
    with open(p, "r", encoding="utf-8") as f:
        rows = list(csvmod.DictReader(f))
    if not rows:
        return []
    # Group rows by section for richer context
    sections: dict = {}
    for row in rows:
        sec = row.get("Section", "Unknown")
        sections.setdefault(sec, []).append(row)
    for sec_name, sec_rows in sections.items():
        lines = []
        for r in sec_rows:
            rec = r.get("Recommendation #", "")
            title = r.get("Title", "")
            compliant = r.get("Compliant", "")
            pct = r.get("Compliance Percentage Per Recommendation", "")
            findings = r.get("Findings", "0")
            total = r.get("Total", "0")
            lines.append(f"Recommendation {rec}: {title} | Status: {compliant} | Compliance: {pct}% | Findings: {findings}/{total}")
        content = (
            f"Tenancy: {tenancy}\n"
            f"Extract Date: {extract_date}\n"
            f"Section: {sec_name}\n"
            f"Total Recommendations: {len(sec_rows)}\n\n"
            + "\n".join(lines)
        )
        documents.append({
            "content": content,
            "section": sec_name,
            "tenancy": tenancy,
            "metadata": json.dumps({"tenancy": tenancy, "extract_date": extract_date, "section": sec_name})
        })
    return documents


# ── CIS Report CSV → ADB Table Mapping ──
_CIS_TABLE_MAP = {
    "Identity_and_Access_Management": "identityandaccess",
    "Networking": "networking",
    "Compute": "computeinstances",
    "Logging_and_Monitoring": "loggingandmonitoring",
    "Storage_Object_Storage": "objectstorage",
    "Storage_Block_Volumes": "storageblockvolume",
    "Storage_File_Storage_Service": "filestorageservice",
    "Asset_Management": "assetmanagement",
}

def _purge_table_by_tenancy(cfg: dict, table_name: str, tenancy: str, extract_date: str = "") -> int:
    """Delete existing embeddings from a table for a specific tenancy (and optionally extract_date).
    Returns number of rows deleted."""
    try:
        conn = _get_adb_connection(cfg)
        cur = conn.cursor()
        if extract_date:
            cur.execute(f"""DELETE FROM "{table_name}" WHERE
                JSON_VALUE(METADATA, '$.tenancy') = :1 AND
                JSON_VALUE(METADATA, '$.extract_date') = :2""", [tenancy, extract_date])
        else:
            cur.execute(f"""DELETE FROM "{table_name}" WHERE
                JSON_VALUE(METADATA, '$.tenancy') = :1""", [tenancy])
        deleted = cur.rowcount
        conn.commit()
        cur.close()
        conn.close()
        if deleted:
            log.info(f"Purged {deleted} rows from {table_name} (tenancy={tenancy}, date={extract_date or 'all'})")
        return deleted
    except Exception as e:
        log.warning(f"Purge failed for {table_name}: {e}")
        return 0


def _resolve_table_for_csv(filename: str) -> str | None:
    """Map a CIS report CSV filename to its ADB vector table."""
    if filename == "cis_summary_report.csv":
        return "summaryreportcsvvector"
    for pattern, table in _CIS_TABLE_MAP.items():
        if pattern in filename:
            return table
    return None


def _chunk_findings_csv(csv_path: str, tenancy: str, extract_date: str, max_chars: int = 8000) -> list:
    """Chunk a CIS findings CSV into documents. Each row becomes one or more documents.
    If a row exceeds max_chars (~6000 tokens), it's split into smaller chunks with
    a context header (tenancy, resource name, ID) repeated in each part."""
    import csv as csvmod, re as _re
    p = Path(csv_path)
    if not p.exists():
        return []
    documents = []
    with open(p, "r", encoding="utf-8") as f:
        rows = list(csvmod.DictReader(f))
    if not rows:
        return []
    skip_cols = {"extract_date", "deep_link", "domain_deeplink", "defined_tags",
                 "freeform_tags", "system_tags", "external_identifier"}
    # Extract CIS recommendation number from filename (e.g., cis_Identity_and_Access_Management_1-1.csv → 1.1)
    rec_match = _re.search(r'_(\d+)-(\d+(?:\.\d+)?)\.csv$', p.name)
    cis_rec = f"{rec_match.group(1)}.{rec_match.group(2)}" if rec_match else ""
    # Extract section name from filename
    sec_match = _re.search(r'^cis_(.+?)_\d+-', p.name)
    cis_section = sec_match.group(1).replace("_", " ") if sec_match else ""

    meta = json.dumps({"tenancy": tenancy, "extract_date": extract_date, "source": p.name, "cis_recommendation": cis_rec})

    for row in rows:
        # Build context header (always repeated in each chunk)
        header_parts = [f"Tenancy: {tenancy}", f"Extract Date: {extract_date}"]
        if cis_rec:
            header_parts.append(f"CIS Recommendation: {cis_rec}")
        if cis_section:
            header_parts.append(f"Section: {cis_section}")
        header_parts.append(f"Status: Non-Compliant")
        body_parts = []
        # Identify key fields for the header
        name = row.get("name") or row.get("display_name") or row.get("username") or ""
        rid = row.get("id", "")
        if name:
            header_parts.append(f"Resource: {name}")
        if rid:
            header_parts.append(f"ID: {rid}")

        for col, val in row.items():
            if col.lower() in skip_cols or not val or not val.strip():
                continue
            if col.lower() in ("name", "display_name", "username", "id"):
                continue  # already in header
            # Clean HYPERLINK formulas
            if val.startswith("=HYPERLINK"):
                m = _re.search(r',\s*"([^"]+)"', val)
                val = m.group(1) if m else val
            body_parts.append(f"{col}: {val}")

        header = "\n".join(header_parts)
        body = "\n".join(body_parts)
        full_content = header + "\n" + body

        if len(full_content) <= max_chars:
            if len(full_content) > 50:
                documents.append({"content": full_content, "tenancy": tenancy, "metadata": meta})
        else:
            # Split body into chunks, each prefixed with context header
            chunk_size = max_chars - len(header) - 50  # reserve space for header + part label
            chunks = []
            current = ""
            for line in body_parts:
                if len(current) + len(line) + 2 > chunk_size and current:
                    chunks.append(current)
                    current = line
                else:
                    current = current + "\n" + line if current else line
            if current:
                chunks.append(current)

            for i, chunk in enumerate(chunks):
                part_label = f"(part {i + 1}/{len(chunks)})" if len(chunks) > 1 else ""
                content = f"{header}\n{part_label}\n{chunk}".strip()
                if len(content) > 50:
                    documents.append({"content": content, "tenancy": tenancy, "metadata": meta})

    return documents