feat: add dedicated Embeddings tab with report chunking and file upload
Add a new Embeddings menu in the sidebar for managing vector embeddings
separately from ADB config. Supports embedding CIS reports (chunked by
section) and uploading .txt files (chunked by paragraphs). Includes
listing and deleting individual embeddings from the ADB vector store.
- Add _chunk_report_by_section() for CIS report sectional chunking
- Add _chunk_text_file() for paragraph-based text chunking
- Add POST /api/embeddings/report/{rid} endpoint
- Add POST /api/embeddings/upload endpoint (multipart)
- Add GET /api/embeddings/{vid}/list with pagination
- Add DELETE /api/embeddings/{vid}/{doc_id}
- Add Embeddings tab in frontend with 3 cards (list, report, upload)
- Remove ingestion section from ADB Vector tab
This commit is contained in:
134
backend/app.py
134
backend/app.py
@@ -953,6 +953,140 @@ def _ingest_documents_task(cfg: dict, genai_cfg: dict, documents: list, user_id:
|
||||
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)
|
||||
findings = report_data.get("findings", {})
|
||||
tenancy = report_data.get("tenancy", "unknown")
|
||||
generated_at = report_data.get("generated_at", "")
|
||||
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 = [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}",
|
||||
"metadata": f"section: {section_name}, total: {len(checks)}, passed: {passed}, failed: {failed}, review: {review}"
|
||||
})
|
||||
return documents
|
||||
|
||||
def _chunk_text_file(text: str, filename: str, chunk_size: int = 1000) -> list:
|
||||
"""Split text into chunks by paragraphs or fixed size."""
|
||||
paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]
|
||||
documents = []
|
||||
current_chunk = ""
|
||||
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
|
||||
current_chunk = 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):
|
||||
"""Load ADB config and its linked GenAI config. Returns (adb_cfg, genai_cfg) or raises."""
|
||||
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)
|
||||
if not cfg.get("genai_config_id"):
|
||||
raise HTTPException(400, "GenAI config not linked to this ADB connection")
|
||||
with db() as c:
|
||||
gc = c.execute("SELECT * FROM genai_configs WHERE id=?", (cfg["genai_config_id"],)).fetchone()
|
||||
if not gc: raise HTTPException(400, "Linked GenAI config not found")
|
||||
return cfg, dict(gc)
|
||||
|
||||
@app.post("/api/embeddings/report/{rid}")
|
||||
async def embed_report(rid: str, req: dict, bg: BackgroundTasks, u=Depends(require("admin","user"))):
|
||||
vid = req.get("adb_config_id")
|
||||
if not vid: raise HTTPException(400, "adb_config_id is required")
|
||||
with db() as c:
|
||||
r = c.execute("SELECT report_data,tenancy_name FROM reports WHERE id=? AND status='completed'", (rid,)).fetchone()
|
||||
if not r: raise HTTPException(404, "Report not found or not completed")
|
||||
report_data = r["report_data"]
|
||||
if isinstance(report_data, str):
|
||||
try: report_data = json.loads(report_data)
|
||||
except: raise HTTPException(400, "Invalid report data")
|
||||
documents = _chunk_report_by_section(report_data)
|
||||
if not documents: raise HTTPException(400, "No sections found in report")
|
||||
cfg, gc = _get_adb_and_genai(vid)
|
||||
bg.add_task(_ingest_documents_task, cfg, gc, documents, u["id"], u["username"])
|
||||
_audit(u["id"], u["username"], "embed_report", rid, f"{len(documents)} sections")
|
||||
return {"ok": True, "message": f"Embedding de {len(documents)} seções iniciado", "sections": len(documents)}
|
||||
|
||||
@app.post("/api/embeddings/upload")
|
||||
async def embed_upload(adb_config_id: str = Form(...), file: UploadFile = File(...), bg: BackgroundTasks = None, u=Depends(require("admin","user"))):
|
||||
if not file.filename.lower().endswith(".txt"):
|
||||
raise HTTPException(400, "Only .txt files are supported")
|
||||
content = (await file.read()).decode("utf-8", errors="replace")
|
||||
if not content.strip(): raise HTTPException(400, "File is empty")
|
||||
documents = _chunk_text_file(content, file.filename)
|
||||
if not documents: raise HTTPException(400, "No content chunks found")
|
||||
cfg, gc = _get_adb_and_genai(adb_config_id)
|
||||
bg.add_task(_ingest_documents_task, cfg, gc, documents, u["id"], u["username"])
|
||||
_audit(u["id"], u["username"], "embed_upload", file.filename, f"{len(documents)} chunks")
|
||||
return {"ok": True, "message": f"Embedding de {len(documents)} chunks iniciado", "chunks": len(documents), "filename": file.filename}
|
||||
|
||||
@app.get("/api/embeddings/{vid}/list")
|
||||
async def list_embeddings(vid: str, limit: int = Query(50), offset: int = Query(0), u=Depends(current_user)):
|
||||
with db() as c:
|
||||
cfg = c.execute("SELECT * FROM adb_vector_configs WHERE id=?", (vid,)).fetchone()
|
||||
if not cfg: raise HTTPException(404)
|
||||
try:
|
||||
conn = _get_adb_connection(dict(cfg))
|
||||
cur = conn.cursor()
|
||||
table_name = cfg["table_name"] or "CIS_EMBEDDINGS"
|
||||
cur.execute(f"SELECT COUNT(*) FROM {table_name}")
|
||||
total = cur.fetchone()[0]
|
||||
cur.execute(f"""
|
||||
SELECT ID, SOURCE, METADATA, CREATED_AT FROM {table_name}
|
||||
ORDER BY CREATED_AT DESC
|
||||
OFFSET :1 ROWS FETCH NEXT :2 ROWS ONLY
|
||||
""", [offset, limit])
|
||||
rows = []
|
||||
for row in cur:
|
||||
rows.append({"id": row[0], "source": row[1], "metadata": row[2],
|
||||
"created_at": str(row[3]) if row[3] else None})
|
||||
cur.close(); conn.close()
|
||||
return {"total": total, "offset": offset, "limit": limit, "documents": rows}
|
||||
except Exception as e:
|
||||
raise HTTPException(500, f"Erro ao listar embeddings: {str(e)[:500]}")
|
||||
|
||||
@app.delete("/api/embeddings/{vid}/{doc_id}")
|
||||
async def delete_embedding(vid: str, doc_id: str, u=Depends(require("admin","user"))):
|
||||
with db() as c:
|
||||
cfg = c.execute("SELECT * FROM adb_vector_configs WHERE id=?", (vid,)).fetchone()
|
||||
if not cfg: raise HTTPException(404)
|
||||
try:
|
||||
conn = _get_adb_connection(dict(cfg))
|
||||
cur = conn.cursor()
|
||||
table_name = cfg["table_name"] or "CIS_EMBEDDINGS"
|
||||
cur.execute(f"DELETE FROM {table_name} WHERE ID = :1", [doc_id])
|
||||
conn.commit()
|
||||
cur.close(); conn.close()
|
||||
return {"ok": True}
|
||||
except Exception as e:
|
||||
raise HTTPException(500, f"Erro ao deletar: {str(e)[:500]}")
|
||||
|
||||
# ── Reports ───────────────────────────────────────────────────────────────────
|
||||
@app.post("/api/reports/run")
|
||||
async def run_report(req: RunReportReq, bg: BackgroundTasks, u=Depends(require("admin","user"))):
|
||||
|
||||
Reference in New Issue
Block a user