GustavoNogueira/A-Team-Security-Infra-Agent-Engineer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: embedding reliability - chunking, truncation, progress bar, shared status

Browse Source 
- Text chunking: max 8000 chars per document, large CSV rows split with context header
- Text truncation: _embed_text truncates at 8000 chars before sending to model (prevents 400 errors)
- Purge before re-embed: deletes old docs by tenancy+date before inserting new ones
- Shared embedding status: file-based status (/data/.embed_status/) instead of in-memory dict (fixes multi-worker visibility)
- Progress bar: shows green (OK) + red (failures) segments, percentage, processed/total count
- Better error logging: extracts readable error message from OCI API errors
This commit is contained in:
nogueiraguh
2026-03-21 10:48:21 -03:00
parent 73c60212f7
commit ab2d0a3e9b
2 changed files with 197 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

202
backend/app.py
View File

@@ -41,7 +41,28 @@ for d in [DATA, OCI_DIR, REPORTS, MCP_DIR, WALLET_DIR]:
_running_reports: dict[str, asyncio.subprocess.Process] = {} # rid → subprocess
_running_terraform: dict[str, asyncio.subprocess.Process] = {} # wid → subprocess
_embedding_status: dict[str, dict] = {} # task_id → {status, message, table, tenancy, inserted, total}
_EMBED_STATUS_DIR = DATA / ".embed_status"
_EMBED_STATUS_DIR.mkdir(parents=True, exist_ok=True)
def _set_embed_status(task_id: str, data: dict):
"""Write embedding status to shared file (visible across all workers)."""
(_EMBED_STATUS_DIR / f"{task_id}.json").write_text(json.dumps(data), encoding="utf-8")
def _get_embed_status(task_id: str) -> dict | None:
"""Read embedding status from shared file."""
p = _EMBED_STATUS_DIR / f"{task_id}.json"
if p.exists():
try:
return json.loads(p.read_text(encoding="utf-8"))
except Exception:
return None
return None
def _update_embed_status(task_id: str, updates: dict):
"""Update embedding status (read-modify-write)."""
current = _get_embed_status(task_id) or {}
current.update(updates)
_set_embed_status(task_id, current)
TERRAFORM_DIR = DATA / "terraform"
TERRAFORM_DIR.mkdir(parents=True, exist_ok=True)
_chat_executor = concurrent.futures.ThreadPoolExecutor(max_workers=10, thread_name_prefix="chat")
@@ -2811,9 +2832,12 @@ def _resolve_embed_config(oci_config_id: str = None, genai_cfg: dict = None) ->
}
raise HTTPException(400, "Nenhuma credencial OCI configurada para gerar embeddings.")
def _embed_text(text: str, genai_cfg: dict, embedding_model_id: str) -> list:
def _embed_text(text: str, genai_cfg: dict, embedding_model_id: str, max_chars: int = 8000) -> list:
"""Generate embedding using OCI GenAI embed endpoint."""
import oci
# Truncate text to fit model token limit (~8192 tokens)
if len(text) > max_chars:
text = text[:max_chars]
config_path = str(OCI_DIR / genai_cfg["oci_config_id"] / "config")
config = oci.config.from_file(config_path, "DEFAULT")
endpoint = genai_cfg.get("endpoint") or f"https://inference.generativeai.{genai_cfg.get('genai_region','us-ashburn-1')}.oci.oraclecloud.com"
@@ -2828,7 +2852,7 @@ def _embed_text(text: str, genai_cfg: dict, embedding_model_id: str) -> list:
emb_ref = emb_info.get("ocids", {}).get(region) or embedding_model_id
embed_detail.serving_mode = oci.generative_ai_inference.models.OnDemandServingMode(model_id=emb_ref)
embed_detail.compartment_id = genai_cfg.get("compartment_id", "")
embed_detail.truncate = "NONE"
embed_detail.truncate = "END"
embed_detail.input_type = "SEARCH_QUERY"
response = client.embed_text(embed_detail)
return response.data.embeddings[0]
@@ -3489,8 +3513,8 @@ def _ingest_documents_task(cfg: dict, genai_cfg: dict, documents: list, user_id:
total = len(documents)
# Track status
if task_id:
_embedding_status[task_id] = {"status": "running", "table": table_name, "tenancy": tenancy or "",
"inserted": 0, "total": total, "message": "Iniciando embedding..."}
_set_embed_status(task_id, {"status": "running", "table": table_name, "tenancy": tenancy or "",
"inserted": 0, "total": total, "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)
@@ -3519,7 +3543,7 @@ def _ingest_documents_task(cfg: dict, genai_cfg: dict, documents: list, user_id:
""", [uuid.uuid4().hex.upper(), content, vec, metadata])
inserted += 1
if task_id:
_embedding_status[task_id].update({"inserted": inserted, "message": f"Embedding {inserted}/{total}..."})
_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()
@@ -3530,12 +3554,12 @@ def _ingest_documents_task(cfg: dict, genai_cfg: dict, documents: list, user_id:
_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:
_embedding_status[task_id].update({"status": "done", "inserted": inserted, "message": msg})
_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:
_embedding_status[task_id].update({"status": "error", "message": str(e)[:300]})
_update_embed_status(task_id, {"status": "error", "message": str(e)[:300]})
finally:
conn.close()
@@ -3693,6 +3717,31 @@ _CIS_TABLE_MAP = {
"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":
@@ -3703,9 +3752,11 @@ def _resolve_table_for_csv(filename: str) -> str | None:
return None
def _chunk_findings_csv(csv_path: str, tenancy: str, extract_date: str) -> list:
"""Chunk a CIS findings CSV into documents. Each row becomes a document with structured content."""
import csv as csvmod
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 []
@@ -3716,26 +3767,58 @@ def _chunk_findings_csv(csv_path: str, tenancy: str, extract_date: str) -> list:
return []
skip_cols = {"extract_date", "deep_link", "domain_deeplink", "defined_tags",
"freeform_tags", "system_tags", "external_identifier"}
meta = json.dumps({"tenancy": tenancy, "extract_date": extract_date, "source": p.name})
for row in rows:
parts = []
parts.append(f"Tenancy: {tenancy}")
parts.append(f"Extract Date: {extract_date}")
# Build context header (always repeated in each chunk)
header_parts = [f"Tenancy: {tenancy}", f"Extract Date: {extract_date}"]
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"):
import re
m = re.search(r',\s*"([^"]+)"', val)
m = _re.search(r',\s*"([^"]+)"', val)
val = m.group(1) if m else val
parts.append(f"{col}: {val}")
content = "\n".join(parts)
if len(content) > 50:
documents.append({
"content": content,
"tenancy": tenancy,
"metadata": json.dumps({"tenancy": tenancy, "extract_date": extract_date, "source": p.name})
})
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
@@ -3796,11 +3879,11 @@ async def embed_report(rid: str, req: dict, bg: BackgroundTasks, u=Depends(requi
# Calculate totals for status tracking
total_docs = sum(len(d) for d in table_docs.values())
tables_used = list(table_docs.keys())
_embedding_status[task_id] = {
_set_embed_status(task_id, {
"status": "running", "table": ", ".join(tables_used), "tenancy": tenancy,
"inserted": 0, "total": total_docs,
"message": f"Embedding {total_docs} documentos em {len(tables_used)} tabelas..."
}
})
def _bg_embed_all():
"""Background: embed documents into their respective tables sequentially."""
@@ -3810,6 +3893,10 @@ async def embed_report(rid: str, req: dict, bg: BackgroundTasks, u=Depends(requi
errors = []
for tbl, docs in table_docs.items():
_auto_register_table(cfg["id"], tbl)
# Purge old data for this tenancy/date before inserting
purged = _purge_table_by_tenancy(cfg, tbl, tenancy, extract_date)
if purged:
_update_embed_status(task_id, {"message": f"Purged {purged} old docs from {tbl}..."})
# Auto-detect embedding model based on table dimension
emb_model = default_model
try:
@@ -3823,9 +3910,12 @@ async def embed_report(rid: str, req: dict, bg: BackgroundTasks, u=Depends(requi
conn = _get_adb_connection(cfg)
cur = conn.cursor()
for doc in docs:
processed = 0
try:
content = doc.get("content", "")
if not content: continue
if not content:
processed += 1
continue
embedding = _embed_text(content, gc, emb_model)
vec = array.array('f', [float(x) for x in embedding])
metadata = _build_metadata_json(
@@ -3836,12 +3926,13 @@ async def embed_report(rid: str, req: dict, bg: BackgroundTasks, u=Depends(requi
cur.execute(f'INSERT INTO "{tbl}" (ID, TEXT, EMBEDDING, METADATA) VALUES (HEXTORAW(:1), :2, :3, :4)',
[uuid.uuid4().hex.upper(), content, vec, metadata])
inserted_total += 1
_embedding_status[task_id].update({
"inserted": inserted_total,
"message": f"Embedding {inserted_total}/{total_docs} — tabela: {tbl}"
})
except Exception as e:
log.error(f"Failed to embed doc in {tbl}: {e}")
log.warning(f"Embed skip in {tbl}: {str(e)[:120]}")
processed += 1
_update_embed_status(task_id, {
"inserted": inserted_total,
"message": f"Embedding {inserted_total}/{processed} OK — {tbl}"
})
conn.commit()
cur.close()
conn.close()
@@ -3853,7 +3944,7 @@ async def embed_report(rid: str, req: dict, bg: BackgroundTasks, u=Depends(requi
msg = f"{inserted_total}/{total_docs} documentos em {len(tables_used)} tabelas ({', '.join(tables_used)})"
if tenancy: msg += f" — tenancy: {tenancy}"
if errors: msg += f" | Erros: {'; '.join(errors)}"
_embedding_status[task_id].update({
_update_embed_status(task_id, {
"status": "done" if not errors else "done",
"inserted": inserted_total, "message": msg
})
@@ -3874,7 +3965,7 @@ async def embed_report(rid: str, req: dict, bg: BackgroundTasks, u=Depends(requi
@app.get("/api/embeddings/status/{task_id}")
async def embedding_status(task_id: str, u=Depends(current_user)):
"""Check embedding task progress."""
st = _embedding_status.get(task_id)
st = _get_embed_status(task_id)
if not st:
return {"status": "unknown", "message": "Task not found"}
return st
@@ -3943,46 +4034,65 @@ async def embed_report_section(rid: str, req: dict, bg: BackgroundTasks, u=Depen
total_docs = sum(len(d) for d in table_docs.values())
tables_used = list(table_docs.keys())
task_id = str(uuid.uuid4())
_embedding_status[task_id] = {
_set_embed_status(task_id, {
"status": "running", "table": ", ".join(tables_used), "tenancy": tenancy,
"inserted": 0, "total": total_docs, "message": f"Embedding {total_docs} docs em {', '.join(tables_used)}..."
}
})
def _bg():
default_model = cfg.get("embedding_model_id", "cohere.embed-v4.0")
inserted = 0
skipped = 0
processed = 0
for tbl, docs in table_docs.items():
purged = _purge_table_by_tenancy(cfg, tbl, tenancy, extract_date)
if purged:
_update_embed_status(task_id, {"message": f"Purged {purged} old docs from {tbl}..."})
_auto_register_table(cfg["id"], tbl)
# Auto-detect embedding model based on table dimension
emb_model = default_model
try:
actual_dim = _get_table_embedding_dim(cfg, tbl)
if actual_dim and actual_dim in _DIM_TO_MODEL:
emb_model = _DIM_TO_MODEL[actual_dim]
log.info(f"Table {tbl}: dim={actual_dim}, using model {emb_model}")
except Exception as e:
log.warning(f"Could not detect dim for {tbl}: {e}")
except Exception:
pass
try:
conn = _get_adb_connection(cfg)
cur = conn.cursor()
for doc in docs:
try:
content = doc.get("content", "")
if not content: continue
if not content:
skipped += 1
processed += 1
continue
embedding = _embed_text(content, gc, emb_model)
vec = array.array('f', [float(x) for x in embedding])
metadata = _build_metadata_json(tenancy=doc.get("tenancy", tenancy), section=doc.get("section", ""), report_date=extract_date)
cur.execute(f'INSERT INTO "{tbl}" (ID, TEXT, EMBEDDING, METADATA) VALUES (HEXTORAW(:1), :2, :3, :4)',
[uuid.uuid4().hex.upper(), content, vec, metadata])
inserted += 1
_embedding_status[task_id].update({"inserted": inserted, "message": f"Embedding {inserted}/{total_docs} — {tbl}"})
except Exception as e:
log.error(f"Section embed error in {tbl}: {e}")
skipped += 1
err_str = str(e)
if "message" in err_str:
import re as _re
m = _re.search(r'"message"\s*:\s*"(.*?)"', err_str)
log.warning(f"Embed skip in {tbl}: {m.group(1)[:200] if m else err_str[:200]}")
else:
log.warning(f"Embed skip in {tbl}: {err_str[:200]}")
processed += 1
_update_embed_status(task_id, {
"inserted": inserted, "skipped": skipped, "processed": processed, "total": total_docs,
"message": f"{tbl}: {inserted} OK, {skipped} falhas ({processed}/{total_docs})"
})
conn.commit(); cur.close(); conn.close()
except Exception as e:
log.error(f"Section embed connection error {tbl}: {e}")
msg = f"{inserted}/{total_docs} docs em {', '.join(tables_used)} (tenancy: {tenancy})"
_embedding_status[task_id].update({"status": "done", "inserted": inserted, "message": msg})
log.error(f"Embed connection error {tbl}: {e}")
msg = f"{inserted}/{total_docs} embeddings OK"
if skipped: msg += f", {skipped} falhas"
msg += f" em {', '.join(tables_used)} (tenancy: {tenancy})"
_update_embed_status(task_id, {"status": "done", "inserted": inserted, "skipped": skipped, "processed": processed, "message": msg})
_audit(u["id"], u["username"], "embed_section", rid, msg)
_chat_executor.submit(_bg)