from __future__ import annotations
import asyncio
import inspect
import json
import queue
import threading
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from pathlib import Path
from typing import Any
from backend import config
from backend.logging_audit import audit_event
from backend.run_report import write_upload_reports
class JobStatus(str, Enum):
PENDING = "pending"
RUNNING = "running"
COMPLETED = "completed"
FAILED = "failed"
@dataclass
class PipelineJob:
job_id: str
user_n: str
ip: str
status: JobStatus = JobStatus.PENDING
stage: str = "pending"
created_at: str = field(
default_factory=lambda: datetime.now(timezone.utc)
.isoformat()
.replace("+00:00", "Z")
)
updated_at: str = field(
default_factory=lambda: datetime.now(timezone.utc)
.isoformat()
.replace("+00:00", "Z")
)
error: str | None = None
uploaded_files: list[str] = field(default_factory=list)
artifacts: dict[str, str] = field(default_factory=dict)
stage_counts: dict[str, Any] = field(default_factory=dict)
_thread: threading.Thread | None = None
_queue: queue.Queue[str | None] = field(default_factory=queue.Queue)
@property
def root_dir(self) -> Path:
return config.job_dir(self.job_id)
@property
def uploads_dir(self) -> Path:
return config.job_uploads_dir(self.job_id)
@property
def ingestion_dir(self) -> Path:
return config.job_ingestion_dir(self.job_id)
@property
def chunking_dir(self) -> Path:
return config.job_chunking_dir(self.job_id)
@property
def embedding_dir(self) -> Path:
return config.job_embedding_dir(self.job_id)
_JOBS: dict[str, PipelineJob] = {}
_JOBS_LOCK = threading.Lock()
def _utc_now() -> str:
return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
def _sse(event_type: str, payload: dict[str, Any]) -> str:
return f"event: {event_type}\ndata: {json.dumps(payload, ensure_ascii=True)}\n\n"
def _emit(job: PipelineJob, event_type: str, payload: dict[str, Any]) -> None:
body = {
"job_id": job.job_id,
"stage": job.stage,
"status": job.status.value,
**payload,
}
job._queue.put(_sse(event_type, body))
def _ensure_stage_supports_output_dir(callable_obj: Any, stage_name: str) -> None:
params = inspect.signature(callable_obj).parameters
if "output_dir" not in params:
raise RuntimeError(
f"{stage_name} stage cannot be constrained to backend output path using existing API."
)
def create_job(user_n: str, ip: str) -> PipelineJob:
config.bootstrap()
job_id = uuid.uuid4().hex[:12]
job = PipelineJob(job_id=job_id, user_n=user_n, ip=ip or "unknown")
for path in (
job.root_dir,
job.uploads_dir,
job.ingestion_dir,
job.chunking_dir,
job.embedding_dir,
):
path.mkdir(parents=True, exist_ok=True)
with _JOBS_LOCK:
_JOBS[job_id] = job
_emit(job, "job_created", {"message": "job accepted"})
return job
def get_job(job_id: str) -> PipelineJob:
with _JOBS_LOCK:
job = _JOBS.get(job_id)
if job is None:
raise KeyError(f"Job not found: {job_id}")
return job
def record_uploads(job_id: str, filenames: list[str]) -> None:
job = get_job(job_id)
job.uploaded_files = filenames
job.updated_at = _utc_now()
_emit(job, "upload_received", {"files": filenames, "file_count": len(filenames)})
audit_event(
"kb_upload",
{"files": filenames, "file_count": len(filenames)},
user_n=job.user_n,
ip=job.ip,
job_id=job.job_id,
)
def start_job(job_id: str) -> None:
job = get_job(job_id)
if job._thread and job._thread.is_alive():
return
thread = threading.Thread(target=_run_job, args=(job,), daemon=True)
job._thread = thread
thread.start()
def _stage_start(job: PipelineJob, stage: str) -> None:
job.stage = stage
job.status = JobStatus.RUNNING
job.updated_at = _utc_now()
_emit(job, "stage_started", {"stage_name": stage})
audit_event(
"stage_started",
{"stage_name": stage},
user_n=job.user_n,
ip=job.ip,
job_id=job.job_id,
)
def _stage_done(job: PipelineJob, stage: str, payload: dict[str, Any]) -> None:
job.updated_at = _utc_now()
_emit(job, "stage_completed", {"stage_name": stage, **payload})
audit_event(
"stage_completed",
{"stage_name": stage, **payload},
user_n=job.user_n,
ip=job.ip,
job_id=job.job_id,
)
def _run_job(job: PipelineJob) -> None:
try:
modules = config.get_stage_modules()
ingestion_engine = modules["ingestion_engine"]
ingestion_events = modules["ingestion_events"]
chunking = modules["chunking"]
embedding_chunks = modules["embedding_chunks"]
vectordb_upsert = modules["vectordb_upsert"]
if not any(path.is_file() for path in job.uploads_dir.rglob("*")):
raise RuntimeError("No uploaded files were provided for ingestion.")
# Ingestion
_stage_start(job, "ingestion")
callback_emitter = ingestion_events.CallbackEmitter(
lambda event: _on_ingestion_event(job, ingestion_events, event)
)
ingestion_cfg = ingestion_engine.IngestionConfig(
input_dir=job.uploads_dir,
output_dir=job.ingestion_dir,
workers=config.env_int("INGEST_WORKERS", 4),
)
engine = ingestion_engine.IngestionEngine(ingestion_cfg, emitter=callback_emitter)
ingestion_stats = engine.run()
cleaned_path = job.ingestion_dir / "cleaned_documents.jsonl"
if not cleaned_path.exists():
raise RuntimeError("Ingestion did not produce cleaned_documents.jsonl.")
job.artifacts["cleaned_documents_jsonl"] = str(cleaned_path)
job.stage_counts["ingestion"] = ingestion_stats.get("counts", {})
_stage_done(job, "ingestion", {"counts": ingestion_stats.get("counts", {})})
# Chunking
_stage_start(job, "chunking")
_ensure_stage_supports_output_dir(chunking.chunk_file, "chunking")
chunk_stats = chunking.chunk_file(
input_path=cleaned_path,
output_dir=job.chunking_dir,
chunk_size=config.env_int("CHUNK_SIZE", 2000),
chunk_overlap=config.env_int("CHUNK_OVERLAP", 200),
)
chunked_path = job.chunking_dir / "chunked_documents.jsonl"
if not chunked_path.exists():
raise RuntimeError("Chunking did not produce chunked_documents.jsonl.")
job.artifacts["chunked_documents_jsonl"] = str(chunked_path)
job.stage_counts["chunking"] = chunk_stats.get("counts", {})
_stage_done(job, "chunking", {"counts": chunk_stats.get("counts", {})})
# Embedding
_stage_start(job, "embedding")
_ensure_stage_supports_output_dir(embedding_chunks.embed_chunks, "embedding")
embed_stats = embedding_chunks.embed_chunks(
input_path=str(chunked_path),
output_dir=str(job.embedding_dir),
model_name=config.env_value(
"EMBEDDING_MODEL", "nvidia/llama-3.2-nv-embedqa-1b-v2"
),
api_key=config.env_value("NVIDIA_API_KEY", ""),
truncate="NONE",
batch_size=config.env_int("EMBED_BATCH_SIZE", 32),
normalize=False,
)
embedding_path = job.embedding_dir / "embeddings_mtsamples.jsonl"
if not embedding_path.exists():
raise RuntimeError("Embedding stage did not produce embeddings_mtsamples.jsonl.")
job.artifacts["embeddings_jsonl"] = str(embedding_path)
job.stage_counts["embedding"] = embed_stats.get("counts", {})
_stage_done(job, "embedding", {"counts": embed_stats.get("counts", {})})
# Upload reports (per file stage artifacts copied under backend/run_report/upload)
_stage_start(job, "upload_report")
report_summary = write_upload_reports(
job_id=job.job_id,
uploaded_files=job.uploaded_files,
ingestion_jsonl=cleaned_path,
chunking_jsonl=chunked_path,
embedding_jsonl=embedding_path,
)
job.artifacts["run_report_upload"] = str(config.RUN_REPORT_UPLOAD_DIR)
_stage_done(job, "upload_report", {"summary": report_summary})
# Dense upsert
_stage_start(job, "upsert_dense")
dense_result = vectordb_upsert.upsert_dense(
embedding_path,
namespace=config.env_value("DOC_NAMESPACE", ""),
)
job.stage_counts["upsert_dense"] = {
"stored_count": dense_result.get("stored_count", 0)
}
_stage_done(job, "upsert_dense", {"result": dense_result})
# Sparse upsert
_stage_start(job, "upsert_sparse")
sparse_result = vectordb_upsert.upsert_sparse(
embedding_path,
namespace=config.env_value("DOC_NAMESPACE", ""),
)
job.stage_counts["upsert_sparse"] = {
"stored_count": sparse_result.get("stored_count", 0)
}
_stage_done(job, "upsert_sparse", {"result": sparse_result})
job.status = JobStatus.COMPLETED
job.stage = "completed"
job.updated_at = _utc_now()
_emit(job, "pipeline_completed", {"artifacts": job.artifacts, "counts": job.stage_counts})
audit_event(
"pipeline_completed",
{"artifacts": job.artifacts, "counts": job.stage_counts},
user_n=job.user_n,
ip=job.ip,
job_id=job.job_id,
)
except Exception as exc:
job.status = JobStatus.FAILED
job.error = str(exc)
job.updated_at = _utc_now()
_emit(job, "pipeline_failed", {"error": str(exc)})
audit_event(
"pipeline_failed",
{"error": str(exc)},
user_n=job.user_n,
ip=job.ip,
job_id=job.job_id,
)
finally:
job._queue.put(None)
def _on_ingestion_event(job: PipelineJob, ingestion_events_module: Any, event: Any) -> None:
payload = ingestion_events_module.event_to_dict(event)
event_type = payload.get("event_type", "ingestion_event")
_emit(job, event_type, payload)
async def stream_events(job_id: str):
job = get_job(job_id)
while True:
try:
value = job._queue.get_nowait()
except queue.Empty:
if job.status in {JobStatus.COMPLETED, JobStatus.FAILED}:
return
await asyncio.sleep(0.05)
continue
if value is None:
return
yield value
def job_status(job_id: str) -> dict[str, Any]:
job = get_job(job_id)
return {
"job_id": job.job_id,
"status": job.status.value,
"stage": job.stage,
"created_at": job.created_at,
"updated_at": job.updated_at,
"uploaded_files": job.uploaded_files,
"artifacts": job.artifacts,
"counts": job.stage_counts,
"error": job.error,
}