chaim
4ad569d54c
Old behavior: clicking "עיבוד מחדש" on a source preserved the existing
metadata (title, identifier, dates, summary, labels) and only re-chunked
+ re-embedded the file. Sources whose original AI classification produced
weak metadata (e.g. title="החלטה" with no summary) had no path to recover
short of deleting and re-uploading.
New behavior: re-process now goes through the same two-phase pipeline as
a fresh upload — classify (AI extracts metadata) → awaiting_review (user
sees suggestions in the batch review screen and edits) → commit (kb_source
metadata is overwritten + chunks are replaced + labels are replaced).
Implementation:
• start_reingest now returns {job_id, batch_id} and pre-fills metadata_json
with the existing source's labels + summary so the form has a fallback
if the classifier times out.
• The reingest route schedules ingest_jobs.process_classify_stage instead
of admin_sources.process_reingest_job.
• commit_job branches on job.source_id: NULL → fresh upload (creates new
kb_source); set → re-ingest (UPDATE kb_source + replace chunks).
• New process_reingest_embed_stage replaces the old process_reingest_job;
it bypasses _EDITABLE_FIELDS so the user can also change kind on
re-process (e.g. caselaw uploaded as circular).
The browser opens the same batch review screen, so the UX matches uploads
exactly. Re-ingests are single-file batches.
Frontend changes (KnowledgeBase extension v0.9.0) ship in a separate repo.
584 lines
22 KiB
Python
584 lines
22 KiB
Python
"""Source management for the KB admin panel (Phase 3 / Task #14).
|
||
|
||
Today the only way to clean up or correct a source is `DELETE FROM
|
||
kb_source WHERE id=...` on the shared PG plus an mc rm in MinIO. This
|
||
module backs the new sources table in the EspoCRM "ניהול" tab:
|
||
|
||
list_sources — table view with chunk count and last-ingest status
|
||
update_source — patch the user-editable metadata
|
||
(title, identifier, dates, source_url)
|
||
delete_source — hard delete (cascades chunks; drops S3 object)
|
||
start_reingest — kick off the two-phase re-ingest pipeline:
|
||
classify → awaiting_review → commit. The
|
||
classify stage reuses
|
||
ingest_jobs.process_classify_stage so the
|
||
user gets fresh AI metadata suggestions in
|
||
the same review UI as fresh uploads.
|
||
process_reingest_embed_stage — final stage after the user commits:
|
||
UPDATE kb_source metadata + replace labels +
|
||
re-chunk and re-embed (the original
|
||
process_reingest_job logic, minus the
|
||
"preserve metadata" rule).
|
||
|
||
The reingest path differs from ingest_source: ingest_source creates a
|
||
NEW kb_source row and supersedes the previous version. Reingest keeps
|
||
the same row id (so cached _lastSearch chunk_index references in the
|
||
browser still point at a valid source), wipes its chunks, and inserts
|
||
fresh ones from the same file. As of v0.9.0, hand-edited metadata is
|
||
NO LONGER preserved silently — the user reviews AI suggestions and
|
||
edits them on the batch review screen before committing, so any
|
||
overwrite is explicit. Sources whose AI-extracted metadata was thin
|
||
(e.g. title="החלטה" only) get fresh suggestions on every re-process.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import datetime as _dt
|
||
import json
|
||
import logging
|
||
import re
|
||
from typing import Optional
|
||
|
||
from api.services.kb import ingest as kb_ingest
|
||
from api.services.kb import s3 as kb_s3
|
||
from api.services.kb import voyage
|
||
from api.services.kb.chunker import chunk as chunk_text
|
||
from api.services.kb.db import get_pool
|
||
|
||
logger = logging.getLogger("shira.kb.admin_sources")
|
||
|
||
# Whitelist of fields the user can edit through PUT /admin/kb/sources/{id}.
|
||
# Anything else (kind, topic_id, checksum, source_id, ...) is structural
|
||
# and must not be flipped from a metadata edit.
|
||
_EDITABLE_FIELDS = ("title", "identifier", "source_url", "published_at",
|
||
"effective_at", "description")
|
||
|
||
|
||
# ── Listing ─────────────────────────────────────────────────────────────────
|
||
|
||
async def list_sources(
|
||
*,
|
||
topic_id: Optional[int] = None,
|
||
kind: Optional[str] = None,
|
||
label_id: Optional[int] = None,
|
||
limit: int = 200,
|
||
) -> list[dict]:
|
||
"""Source rows enriched with chunk_count, last_ingest_status, and labels.
|
||
|
||
The last-ingest column is a LATERAL join against kb_ingest_job — picks
|
||
the most recent job row for the source (regardless of status), so the
|
||
UI can flag sources whose last re-ingest failed.
|
||
|
||
Labels are aggregated as a JSON array via a correlated subquery —
|
||
keeps the row count clean (no N×labels Cartesian product).
|
||
"""
|
||
limit = max(1, min(500, limit))
|
||
where: list[str] = ["s.superseded_by IS NULL"]
|
||
args: list = []
|
||
if topic_id is not None:
|
||
args.append(topic_id)
|
||
where.append(f"s.topic_id = ${len(args)}")
|
||
if kind:
|
||
args.append(kind)
|
||
where.append(f"s.kind = ${len(args)}")
|
||
if label_id is not None:
|
||
args.append(label_id)
|
||
where.append(f"EXISTS (SELECT 1 FROM kb_source_label sl "
|
||
f"WHERE sl.source_id = s.id AND sl.label_id = ${len(args)})")
|
||
args.append(limit)
|
||
sql = f"""
|
||
SELECT
|
||
s.id,
|
||
s.kind,
|
||
s.title,
|
||
s.identifier,
|
||
s.published_at,
|
||
s.effective_at,
|
||
s.source_url,
|
||
s.original_path,
|
||
s.checksum,
|
||
s.topic_id,
|
||
s.description,
|
||
s.ai_classified_at,
|
||
s.created_at,
|
||
s.updated_at,
|
||
(SELECT COUNT(*) FROM kb_chunk c WHERE c.source_id = s.id) AS chunk_count,
|
||
(
|
||
SELECT json_build_object(
|
||
'job_id', j.id,
|
||
'status', j.status,
|
||
'created_at', j.created_at,
|
||
'completed_at', j.completed_at,
|
||
'error_message', j.error_message
|
||
)
|
||
FROM kb_ingest_job j
|
||
WHERE j.source_id = s.id
|
||
ORDER BY j.created_at DESC
|
||
LIMIT 1
|
||
) AS last_ingest,
|
||
COALESCE(
|
||
(
|
||
SELECT json_agg(json_build_object(
|
||
'id', l.id, 'slug', l.slug, 'name', l.name,
|
||
'topic_id', l.topic_id
|
||
) ORDER BY l.name)
|
||
FROM kb_label l
|
||
JOIN kb_source_label sl ON sl.label_id = l.id
|
||
WHERE sl.source_id = s.id
|
||
),
|
||
'[]'::json
|
||
) AS labels
|
||
FROM kb_source s
|
||
WHERE {' AND '.join(where)}
|
||
ORDER BY s.kind, s.title
|
||
LIMIT ${len(args)}
|
||
"""
|
||
pool = await get_pool()
|
||
async with pool.acquire() as conn:
|
||
rows = await conn.fetch(sql, *args)
|
||
return [_row_to_dict(r) for r in rows]
|
||
|
||
|
||
async def get_source(source_id: int) -> Optional[dict]:
|
||
pool = await get_pool()
|
||
async with pool.acquire() as conn:
|
||
row = await conn.fetchrow(
|
||
"SELECT * FROM kb_source WHERE id = $1",
|
||
source_id,
|
||
)
|
||
return _row_to_dict(row) if row else None
|
||
|
||
|
||
# ── Editing ─────────────────────────────────────────────────────────────────
|
||
|
||
async def update_source(
|
||
source_id: int,
|
||
fields: dict,
|
||
*,
|
||
label_ids: Optional[list[int]] = None,
|
||
) -> dict:
|
||
"""UPDATE only whitelisted columns. Returns the post-update row.
|
||
|
||
If `label_ids` is provided (not None), the source's label set is
|
||
REPLACED to exactly those ids — usage_count adjusted both ways.
|
||
Pass an empty list to clear all labels; pass None to leave labels
|
||
untouched.
|
||
"""
|
||
from api.services.kb import labels as kb_labels # local import: cycle
|
||
|
||
clean: dict = {}
|
||
for k in _EDITABLE_FIELDS:
|
||
if k not in fields:
|
||
continue
|
||
v = fields[k]
|
||
if k in ("published_at", "effective_at"):
|
||
clean[k] = kb_ingest._coerce_date(v)
|
||
elif v is None or v == "":
|
||
clean[k] = None
|
||
else:
|
||
clean[k] = str(v).strip()
|
||
if "title" in clean and not clean["title"]:
|
||
raise ValueError("title cannot be empty")
|
||
|
||
pool = await get_pool()
|
||
|
||
if clean:
|
||
set_pieces = []
|
||
args: list = [source_id]
|
||
for k, v in clean.items():
|
||
args.append(v)
|
||
set_pieces.append(f"{k} = ${len(args)}")
|
||
sql = f"""
|
||
UPDATE kb_source
|
||
SET {', '.join(set_pieces)}
|
||
WHERE id = $1
|
||
RETURNING id
|
||
"""
|
||
async with pool.acquire() as conn:
|
||
updated = await conn.fetchrow(sql, *args)
|
||
if not updated:
|
||
raise LookupError(f"source {source_id} not found")
|
||
else:
|
||
# Caller passed nothing editable; verify the source still exists
|
||
# so we can return a sensible 404 vs blindly returning empty.
|
||
async with pool.acquire() as conn:
|
||
existing = await conn.fetchrow(
|
||
"SELECT id FROM kb_source WHERE id = $1", source_id
|
||
)
|
||
if not existing:
|
||
raise LookupError(f"source {source_id} not found")
|
||
|
||
if label_ids is not None:
|
||
await kb_labels.replace_for_source(source_id, label_ids)
|
||
|
||
# Re-read with labels via list_sources([id])-style query — reuse
|
||
# get_source() for simplicity, then fetch labels separately.
|
||
row = await get_source(source_id)
|
||
if row is not None:
|
||
row["labels"] = await kb_labels.get_labels_for_source(source_id)
|
||
return row
|
||
|
||
|
||
# ── Deletion ────────────────────────────────────────────────────────────────
|
||
|
||
async def delete_source(source_id: int) -> dict:
|
||
"""Hard-delete a source. kb_chunk rows cascade via FK; the underlying
|
||
S3 object is deleted on a best-effort basis (a stray file is harmless,
|
||
but cluttered).
|
||
"""
|
||
src = await get_source(source_id)
|
||
if not src:
|
||
raise LookupError(f"source {source_id} not found")
|
||
|
||
pool = await get_pool()
|
||
async with pool.acquire() as conn:
|
||
result = await conn.execute(
|
||
"DELETE FROM kb_source WHERE id = $1", source_id
|
||
)
|
||
deleted = result.endswith(" 1")
|
||
|
||
# Best-effort S3 cleanup — never fail the whole delete on a stray
|
||
# object that might already be gone.
|
||
s3_deleted = False
|
||
bucket, key = _parse_s3_uri(src.get("original_path") or "")
|
||
if key:
|
||
try:
|
||
kb_s3._client().delete_object(Bucket=bucket or kb_s3._bucket(), Key=key)
|
||
s3_deleted = True
|
||
except Exception as exc:
|
||
logger.warning("[admin.delete] S3 delete failed for %s: %s", key, exc)
|
||
|
||
return {"deleted": deleted, "source_id": source_id, "s3_deleted": s3_deleted}
|
||
|
||
|
||
# ── Re-ingest ───────────────────────────────────────────────────────────────
|
||
|
||
async def start_reingest(
|
||
source_id: int, requested_by_user: Optional[str]
|
||
) -> dict:
|
||
"""Create a kb_ingest_job tied to the existing source_id.
|
||
|
||
The job enters the same two-phase pipeline as a fresh upload —
|
||
classify (AI metadata extraction) → awaiting_review → embed
|
||
(re-chunk + UPDATE kb_source) — so the user can review suggestions
|
||
on the batch review screen before anything overwrites kb_source.
|
||
|
||
Returns {"job_id", "batch_id"}; the route schedules
|
||
`ingest_jobs.process_classify_stage(job_id)` after this, and the
|
||
browser opens `/admin/kb/batch/<batch_id>` to poll the review UI.
|
||
"""
|
||
import uuid as _uuid
|
||
|
||
src = await get_source(source_id)
|
||
if not src:
|
||
raise LookupError(f"source {source_id} not found")
|
||
if not src.get("original_path"):
|
||
raise ValueError(
|
||
f"source {source_id} has no original_path — re-ingest needs the original file"
|
||
)
|
||
|
||
# Pre-flight: confirm the file actually exists in S3 before queuing
|
||
# the job. Saves the user from waiting 30s on a doomed job.
|
||
bucket, key = _parse_s3_uri(src["original_path"])
|
||
if not key:
|
||
raise ValueError(f"unparseable original_path: {src['original_path']}")
|
||
try:
|
||
kb_s3._client().head_object(Bucket=bucket or kb_s3._bucket(), Key=key)
|
||
except Exception as exc:
|
||
raise ValueError(f"original file unavailable: {exc}")
|
||
|
||
# Reuse the kb_ingest_job table; mark this row as a re-ingest by
|
||
# pre-filling source_id (uploads start with source_id NULL). The
|
||
# batch_id is auto-generated so the existing review UI works
|
||
# without changes — re-ingests are single-file batches. Existing
|
||
# labels go into metadata_json so the review form can pre-fill
|
||
# them as a fallback when the AI classifier hasn't run yet (or
|
||
# returns empty).
|
||
from api.services.kb import labels as kb_labels # local import: cycle
|
||
existing_labels = await kb_labels.get_labels_for_source(source_id)
|
||
metadata = {
|
||
"title": src.get("title"),
|
||
"identifier": src.get("identifier"),
|
||
"published_at": _date_to_iso(src.get("published_at")),
|
||
"effective_at": _date_to_iso(src.get("effective_at")),
|
||
"source_url": src.get("source_url"),
|
||
"summary": src.get("description"),
|
||
"labels": [lbl.get("name") for lbl in existing_labels if lbl.get("name")],
|
||
"reingest": True,
|
||
}
|
||
batch_id = _uuid.uuid4().hex
|
||
pool = await get_pool()
|
||
async with pool.acquire() as conn:
|
||
row = await conn.fetchrow(
|
||
"""
|
||
INSERT INTO kb_ingest_job
|
||
(source_id, original_filename, s3_key, kind, topic_id,
|
||
metadata_json, status, requested_by_user, batch_id)
|
||
VALUES ($1, $2, $3, $4, $5, $6::jsonb, 'queued', $7, $8)
|
||
RETURNING id
|
||
""",
|
||
source_id,
|
||
src.get("title") or _filename_from_key(key),
|
||
key,
|
||
src["kind"],
|
||
src["topic_id"],
|
||
json.dumps(metadata, ensure_ascii=False),
|
||
requested_by_user,
|
||
batch_id,
|
||
)
|
||
job_id = row["id"]
|
||
logger.info(
|
||
"[admin.reingest] queued job_id=%s source_id=%s batch_id=%s",
|
||
job_id, source_id, batch_id,
|
||
)
|
||
return {"job_id": job_id, "batch_id": batch_id}
|
||
|
||
|
||
async def process_reingest_embed_stage(
|
||
job_id: int, fields: dict, label_ids: list[int],
|
||
) -> None:
|
||
"""Final stage of the two-phase re-ingest pipeline.
|
||
|
||
Called from `ingest_jobs.commit_job` after the user reviewed the AI
|
||
suggestions. The job is already in `status='processing',
|
||
processing_stage='embedding'` at this point — commit_job set that
|
||
transition before scheduling us.
|
||
|
||
Steps:
|
||
1. UPDATE kb_source with the user-confirmed metadata (kind,
|
||
title, identifier, source_url, dates, description=summary).
|
||
2. Replace the source's labels with `label_ids`.
|
||
3. Re-parse the file → chunk → embed → atomic DELETE+INSERT
|
||
of kb_chunk rows + checksum update.
|
||
4. Mark the job done.
|
||
|
||
Failures here are real failures (the user already confirmed). They
|
||
surface as `status='failed'` so the polling UI lets the user retry.
|
||
"""
|
||
from api.services.kb import labels as kb_labels # local import: cycle
|
||
|
||
pool = await get_pool()
|
||
async with pool.acquire() as conn:
|
||
job_row = await conn.fetchrow(
|
||
"SELECT * FROM kb_ingest_job WHERE id = $1", job_id,
|
||
)
|
||
if not job_row:
|
||
logger.warning("[reingest-embed] job %s vanished", job_id)
|
||
return
|
||
|
||
job = _row_to_dict(job_row)
|
||
source_id = job["source_id"]
|
||
if source_id is None:
|
||
await _fail_job(pool, job_id, "re-ingest job has no source_id")
|
||
return
|
||
|
||
try:
|
||
# Step 1: UPDATE kb_source with the user-confirmed metadata.
|
||
# commit_job's `fields` dict mirrors JobCommit: kind, title,
|
||
# identifier, source_url, published_at, effective_at, summary.
|
||
# kb_source.description is the schema name for "summary".
|
||
kind = (fields.get("kind") or job["kind"]).strip()
|
||
title = (fields.get("title") or "").strip()
|
||
if not title:
|
||
raise ValueError("title is empty after commit — aborting re-ingest")
|
||
|
||
await _update_source_metadata_for_reingest(
|
||
pool=pool,
|
||
source_id=source_id,
|
||
kind=kind,
|
||
title=title,
|
||
identifier=fields.get("identifier"),
|
||
source_url=fields.get("source_url"),
|
||
published_at=fields.get("published_at"),
|
||
effective_at=fields.get("effective_at"),
|
||
description=fields.get("summary"),
|
||
)
|
||
|
||
# Step 2: replace labels (empty list clears all).
|
||
await kb_labels.replace_for_source(source_id, label_ids or [])
|
||
|
||
# Step 3: re-chunk + re-embed.
|
||
bucket, key = _parse_s3_uri(job["s3_key"])
|
||
if not key:
|
||
key = job["s3_key"] # bare key, no s3:// prefix
|
||
data = kb_s3._client().get_object(
|
||
Bucket=bucket or kb_s3._bucket(),
|
||
Key=key,
|
||
)["Body"].read()
|
||
if not data:
|
||
raise kb_ingest.IngestError(f"empty file at {key}")
|
||
|
||
text = kb_ingest.parse(data, job.get("original_filename") or key)
|
||
if not text:
|
||
raise kb_ingest.IngestError("no text extracted")
|
||
|
||
chunks = chunk_text(text, kind)
|
||
if not chunks:
|
||
raise kb_ingest.IngestError("chunker produced zero chunks")
|
||
|
||
# Embed first (slow, network-bound) so the actual DB swap is fast.
|
||
texts = [c["content"] for c in chunks]
|
||
vectors = await voyage.embed(texts, input_type="document")
|
||
|
||
import hashlib
|
||
new_checksum = hashlib.sha256(text.encode("utf-8")).hexdigest()
|
||
|
||
async with pool.acquire() as conn:
|
||
async with conn.transaction():
|
||
# Drop existing chunks and reinsert in one tx so we never
|
||
# leave the source partially-chunked.
|
||
await conn.execute(
|
||
"DELETE FROM kb_chunk WHERE source_id = $1", source_id
|
||
)
|
||
rows = [
|
||
(
|
||
source_id,
|
||
c["chunk_index"],
|
||
c["heading_path"],
|
||
c["section_ref"],
|
||
c["content"],
|
||
vectors[i],
|
||
c["token_count"],
|
||
c.get("page_number"),
|
||
)
|
||
for i, c in enumerate(chunks)
|
||
]
|
||
await conn.executemany(
|
||
"""
|
||
INSERT INTO kb_chunk
|
||
(source_id, chunk_index, heading_path, section_ref,
|
||
content, embedding, token_count, page_number)
|
||
VALUES ($1,$2,$3,$4,$5,$6,$7,$8)
|
||
""",
|
||
rows,
|
||
)
|
||
await conn.execute(
|
||
"UPDATE kb_source SET checksum = $2 WHERE id = $1",
|
||
source_id, new_checksum,
|
||
)
|
||
await conn.execute(
|
||
"""
|
||
UPDATE kb_ingest_job
|
||
SET status = 'done',
|
||
processing_stage = NULL,
|
||
chunks_created = $2,
|
||
completed_at = now()
|
||
WHERE id = $1
|
||
""",
|
||
job_id, len(chunks),
|
||
)
|
||
logger.info(
|
||
"[reingest-embed] done job_id=%s source_id=%s chunks=%d labels=%d",
|
||
job_id, source_id, len(chunks), len(label_ids or []),
|
||
)
|
||
except Exception as exc:
|
||
logger.exception("[reingest-embed] failed job_id=%s", job_id)
|
||
await _fail_job(pool, job_id, str(exc))
|
||
|
||
|
||
async def _update_source_metadata_for_reingest(
|
||
*, pool, source_id: int, kind: str, title: str,
|
||
identifier: Optional[str], source_url: Optional[str],
|
||
published_at: Optional[str], effective_at: Optional[str],
|
||
description: Optional[str],
|
||
) -> None:
|
||
"""UPDATE kb_source after a re-ingest commit. Bypasses _EDITABLE_FIELDS
|
||
so we can also update `kind` (the user explicitly chose it on the
|
||
review screen — e.g. caselaw uploaded as circular by mistake).
|
||
"""
|
||
pub = kb_ingest._coerce_date(published_at) if published_at else None
|
||
eff = kb_ingest._coerce_date(effective_at) if effective_at else None
|
||
async with pool.acquire() as conn:
|
||
await conn.execute(
|
||
"""
|
||
UPDATE kb_source
|
||
SET kind = $2,
|
||
title = $3,
|
||
identifier = $4,
|
||
source_url = $5,
|
||
published_at = $6,
|
||
effective_at = $7,
|
||
description = $8
|
||
WHERE id = $1
|
||
""",
|
||
source_id,
|
||
kind,
|
||
title,
|
||
(identifier or None),
|
||
(source_url or None),
|
||
pub,
|
||
eff,
|
||
(description or None),
|
||
)
|
||
|
||
|
||
# ── helpers ─────────────────────────────────────────────────────────────────
|
||
|
||
async def _fail_job(pool, job_id: int, msg: str) -> None:
|
||
async with pool.acquire() as conn:
|
||
await conn.execute(
|
||
"""
|
||
UPDATE kb_ingest_job
|
||
SET status = 'failed',
|
||
error_message = $2,
|
||
completed_at = now()
|
||
WHERE id = $1
|
||
""",
|
||
job_id, msg[:4000],
|
||
)
|
||
|
||
|
||
def _row_to_dict(row) -> Optional[dict]:
|
||
if row is None:
|
||
return None
|
||
out: dict = {}
|
||
for k, v in dict(row).items():
|
||
if isinstance(v, _dt.datetime):
|
||
out[k] = v.isoformat()
|
||
elif isinstance(v, _dt.date):
|
||
out[k] = v.isoformat()
|
||
elif k == "metadata_json" and isinstance(v, str):
|
||
try:
|
||
out[k] = json.loads(v)
|
||
except Exception:
|
||
out[k] = {}
|
||
elif k == "last_ingest" and isinstance(v, str):
|
||
try:
|
||
out[k] = json.loads(v)
|
||
except Exception:
|
||
out[k] = None
|
||
elif k == "labels" and isinstance(v, str):
|
||
try:
|
||
out[k] = json.loads(v)
|
||
except Exception:
|
||
out[k] = []
|
||
else:
|
||
out[k] = v
|
||
return out
|
||
|
||
|
||
_S3_URI = re.compile(r"^s3://([^/]+)/(.+)$")
|
||
|
||
def _parse_s3_uri(uri: str) -> tuple[str, str]:
|
||
"""('bucket', 'key/path') or ('', '') if not an s3:// URI."""
|
||
if not uri:
|
||
return ("", "")
|
||
m = _S3_URI.match(uri)
|
||
if not m:
|
||
return ("", "")
|
||
return (m.group(1), m.group(2))
|
||
|
||
|
||
def _filename_from_key(key: str) -> str:
|
||
return key.rsplit("/", 1)[-1] if key else "source"
|
||
|
||
|
||
def _date_to_iso(value) -> Optional[str]:
|
||
if value is None:
|
||
return None
|
||
if isinstance(value, str):
|
||
return value[:10] if value else None
|
||
if isinstance(value, _dt.date):
|
||
return value.isoformat()
|
||
return None
|