fix: Restore all files lost during destructive rebase
A previous `git pull --rebase origin main` dropped 177 local commits,
losing 3400+ files across admin-v2, backend, studio-v2, website,
klausur-service, and many other services. The partial restore attempt
(660295e2) only recovered some files.
This commit restores all missing files from pre-rebase ref 98933f5e
while preserving post-rebase additions (night-scheduler, night-mode UI,
NightModeWidget dashboard integration).
Restored features include:
- AI Module Sidebar (FAB), OCR Labeling, OCR Compare
- GPU Dashboard, RAG Pipeline, Magic Help
- Klausur-Korrektur (8 files), Abitur-Archiv (5+ files)
- Companion, Zeugnisse-Crawler, Screen Flow
- Full backend, studio-v2, website, klausur-service
- All compliance SDKs, agent-core, voice-service
- CI/CD configs, documentation, scripts
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
625
klausur-service/backend/training_api.py
Normal file
625
klausur-service/backend/training_api.py
Normal file
@@ -0,0 +1,625 @@
|
||||
"""
|
||||
Training API - Endpoints for managing AI training jobs
|
||||
|
||||
Provides endpoints for:
|
||||
- Starting/stopping training jobs
|
||||
- Monitoring training progress
|
||||
- Managing model versions
|
||||
- Configuring training parameters
|
||||
- SSE streaming for real-time metrics
|
||||
|
||||
Phase 2.2: Server-Sent Events for live progress
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import uuid
|
||||
import asyncio
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Optional, List, Dict, Any
|
||||
from enum import Enum
|
||||
from dataclasses import dataclass, field, asdict
|
||||
from fastapi import APIRouter, HTTPException, BackgroundTasks, Request
|
||||
from fastapi.responses import StreamingResponse
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# ENUMS & MODELS
|
||||
# ============================================================================
|
||||
|
||||
class TrainingStatus(str, Enum):
|
||||
QUEUED = "queued"
|
||||
PREPARING = "preparing"
|
||||
TRAINING = "training"
|
||||
VALIDATING = "validating"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
PAUSED = "paused"
|
||||
CANCELLED = "cancelled"
|
||||
|
||||
|
||||
class ModelType(str, Enum):
|
||||
ZEUGNIS = "zeugnis"
|
||||
KLAUSUR = "klausur"
|
||||
GENERAL = "general"
|
||||
|
||||
|
||||
# Request/Response Models
|
||||
class TrainingConfig(BaseModel):
|
||||
"""Configuration for a training job."""
|
||||
name: str = Field(..., description="Name for the training job")
|
||||
model_type: ModelType = Field(ModelType.ZEUGNIS, description="Type of model to train")
|
||||
bundeslaender: List[str] = Field(..., description="List of Bundesland codes to include")
|
||||
batch_size: int = Field(16, ge=1, le=128)
|
||||
learning_rate: float = Field(0.00005, ge=0.000001, le=0.1)
|
||||
epochs: int = Field(10, ge=1, le=100)
|
||||
warmup_steps: int = Field(500, ge=0, le=10000)
|
||||
weight_decay: float = Field(0.01, ge=0, le=1)
|
||||
gradient_accumulation: int = Field(4, ge=1, le=32)
|
||||
mixed_precision: bool = Field(True, description="Use FP16 mixed precision training")
|
||||
|
||||
|
||||
class TrainingMetrics(BaseModel):
|
||||
"""Metrics from a training job."""
|
||||
precision: float = 0.0
|
||||
recall: float = 0.0
|
||||
f1_score: float = 0.0
|
||||
accuracy: float = 0.0
|
||||
loss_history: List[float] = []
|
||||
val_loss_history: List[float] = []
|
||||
|
||||
|
||||
class TrainingJob(BaseModel):
|
||||
"""A training job with full details."""
|
||||
id: str
|
||||
name: str
|
||||
model_type: ModelType
|
||||
status: TrainingStatus
|
||||
progress: float
|
||||
current_epoch: int
|
||||
total_epochs: int
|
||||
loss: float
|
||||
val_loss: float
|
||||
learning_rate: float
|
||||
documents_processed: int
|
||||
total_documents: int
|
||||
started_at: Optional[datetime]
|
||||
estimated_completion: Optional[datetime]
|
||||
completed_at: Optional[datetime]
|
||||
error_message: Optional[str]
|
||||
metrics: TrainingMetrics
|
||||
config: TrainingConfig
|
||||
|
||||
|
||||
class ModelVersion(BaseModel):
|
||||
"""A trained model version."""
|
||||
id: str
|
||||
job_id: str
|
||||
version: str
|
||||
model_type: ModelType
|
||||
created_at: datetime
|
||||
metrics: TrainingMetrics
|
||||
is_active: bool
|
||||
size_mb: float
|
||||
bundeslaender: List[str]
|
||||
|
||||
|
||||
class DatasetStats(BaseModel):
|
||||
"""Statistics about the training dataset."""
|
||||
total_documents: int
|
||||
total_chunks: int
|
||||
training_allowed: int
|
||||
by_bundesland: Dict[str, int]
|
||||
by_doc_type: Dict[str, int]
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# IN-MEMORY STATE (Replace with database in production)
|
||||
# ============================================================================
|
||||
|
||||
@dataclass
|
||||
class TrainingState:
|
||||
"""Global training state."""
|
||||
jobs: Dict[str, dict] = field(default_factory=dict)
|
||||
model_versions: Dict[str, dict] = field(default_factory=dict)
|
||||
active_job_id: Optional[str] = None
|
||||
|
||||
|
||||
_state = TrainingState()
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# HELPER FUNCTIONS
|
||||
# ============================================================================
|
||||
|
||||
async def simulate_training_progress(job_id: str):
|
||||
"""Simulate training progress (replace with actual training logic)."""
|
||||
global _state
|
||||
|
||||
if job_id not in _state.jobs:
|
||||
return
|
||||
|
||||
job = _state.jobs[job_id]
|
||||
job["status"] = TrainingStatus.TRAINING.value
|
||||
job["started_at"] = datetime.now().isoformat()
|
||||
|
||||
total_steps = job["total_epochs"] * 100 # Simulate 100 steps per epoch
|
||||
current_step = 0
|
||||
|
||||
while current_step < total_steps and job["status"] == TrainingStatus.TRAINING.value:
|
||||
# Update progress
|
||||
progress = (current_step / total_steps) * 100
|
||||
current_epoch = current_step // 100 + 1
|
||||
|
||||
# Simulate decreasing loss
|
||||
base_loss = 0.8 * (1 - progress / 100) + 0.1
|
||||
loss = base_loss + (0.05 * (0.5 - (current_step % 100) / 100))
|
||||
val_loss = loss * 1.1
|
||||
|
||||
# Update job state
|
||||
job["progress"] = progress
|
||||
job["current_epoch"] = min(current_epoch, job["total_epochs"])
|
||||
job["loss"] = round(loss, 4)
|
||||
job["val_loss"] = round(val_loss, 4)
|
||||
job["documents_processed"] = int((progress / 100) * job["total_documents"])
|
||||
|
||||
# Update metrics
|
||||
job["metrics"]["loss_history"].append(round(loss, 4))
|
||||
job["metrics"]["val_loss_history"].append(round(val_loss, 4))
|
||||
job["metrics"]["precision"] = round(0.5 + (progress / 200), 3)
|
||||
job["metrics"]["recall"] = round(0.45 + (progress / 200), 3)
|
||||
job["metrics"]["f1_score"] = round(0.47 + (progress / 200), 3)
|
||||
job["metrics"]["accuracy"] = round(0.6 + (progress / 250), 3)
|
||||
|
||||
# Keep only last 50 history points
|
||||
if len(job["metrics"]["loss_history"]) > 50:
|
||||
job["metrics"]["loss_history"] = job["metrics"]["loss_history"][-50:]
|
||||
job["metrics"]["val_loss_history"] = job["metrics"]["val_loss_history"][-50:]
|
||||
|
||||
# Estimate completion
|
||||
if progress > 0:
|
||||
elapsed = (datetime.now() - datetime.fromisoformat(job["started_at"])).total_seconds()
|
||||
remaining = (elapsed / progress) * (100 - progress)
|
||||
job["estimated_completion"] = (datetime.now() + timedelta(seconds=remaining)).isoformat()
|
||||
|
||||
current_step += 1
|
||||
await asyncio.sleep(0.5) # Simulate work
|
||||
|
||||
# Mark as completed
|
||||
if job["status"] == TrainingStatus.TRAINING.value:
|
||||
job["status"] = TrainingStatus.COMPLETED.value
|
||||
job["progress"] = 100
|
||||
job["completed_at"] = datetime.now().isoformat()
|
||||
|
||||
# Create model version
|
||||
version_id = str(uuid.uuid4())
|
||||
_state.model_versions[version_id] = {
|
||||
"id": version_id,
|
||||
"job_id": job_id,
|
||||
"version": f"v{len(_state.model_versions) + 1}.0",
|
||||
"model_type": job["model_type"],
|
||||
"created_at": datetime.now().isoformat(),
|
||||
"metrics": job["metrics"],
|
||||
"is_active": True,
|
||||
"size_mb": 245.7,
|
||||
"bundeslaender": job["config"]["bundeslaender"],
|
||||
}
|
||||
|
||||
_state.active_job_id = None
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# ROUTER
|
||||
# ============================================================================
|
||||
|
||||
router = APIRouter(prefix="/api/v1/admin/training", tags=["Training"])
|
||||
|
||||
|
||||
@router.get("/jobs", response_model=List[dict])
|
||||
async def list_training_jobs():
|
||||
"""Get all training jobs."""
|
||||
return list(_state.jobs.values())
|
||||
|
||||
|
||||
@router.get("/jobs/{job_id}", response_model=dict)
|
||||
async def get_training_job(job_id: str):
|
||||
"""Get details for a specific training job."""
|
||||
if job_id not in _state.jobs:
|
||||
raise HTTPException(status_code=404, detail="Job not found")
|
||||
return _state.jobs[job_id]
|
||||
|
||||
|
||||
@router.post("/jobs", response_model=dict)
|
||||
async def create_training_job(config: TrainingConfig, background_tasks: BackgroundTasks):
|
||||
"""Create and start a new training job."""
|
||||
global _state
|
||||
|
||||
# Check if there's already an active job
|
||||
if _state.active_job_id:
|
||||
active_job = _state.jobs.get(_state.active_job_id)
|
||||
if active_job and active_job["status"] in [
|
||||
TrainingStatus.TRAINING.value,
|
||||
TrainingStatus.PREPARING.value,
|
||||
]:
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail="Another training job is already running"
|
||||
)
|
||||
|
||||
# Create job
|
||||
job_id = str(uuid.uuid4())
|
||||
job = {
|
||||
"id": job_id,
|
||||
"name": config.name,
|
||||
"model_type": config.model_type.value,
|
||||
"status": TrainingStatus.QUEUED.value,
|
||||
"progress": 0,
|
||||
"current_epoch": 0,
|
||||
"total_epochs": config.epochs,
|
||||
"loss": 1.0,
|
||||
"val_loss": 1.0,
|
||||
"learning_rate": config.learning_rate,
|
||||
"documents_processed": 0,
|
||||
"total_documents": len(config.bundeslaender) * 50, # Estimate
|
||||
"started_at": None,
|
||||
"estimated_completion": None,
|
||||
"completed_at": None,
|
||||
"error_message": None,
|
||||
"metrics": {
|
||||
"precision": 0.0,
|
||||
"recall": 0.0,
|
||||
"f1_score": 0.0,
|
||||
"accuracy": 0.0,
|
||||
"loss_history": [],
|
||||
"val_loss_history": [],
|
||||
},
|
||||
"config": config.dict(),
|
||||
}
|
||||
|
||||
_state.jobs[job_id] = job
|
||||
_state.active_job_id = job_id
|
||||
|
||||
# Start training in background
|
||||
background_tasks.add_task(simulate_training_progress, job_id)
|
||||
|
||||
return {"id": job_id, "status": "queued", "message": "Training job created"}
|
||||
|
||||
|
||||
@router.post("/jobs/{job_id}/pause", response_model=dict)
|
||||
async def pause_training_job(job_id: str):
|
||||
"""Pause a running training job."""
|
||||
if job_id not in _state.jobs:
|
||||
raise HTTPException(status_code=404, detail="Job not found")
|
||||
|
||||
job = _state.jobs[job_id]
|
||||
if job["status"] != TrainingStatus.TRAINING.value:
|
||||
raise HTTPException(status_code=400, detail="Job is not running")
|
||||
|
||||
job["status"] = TrainingStatus.PAUSED.value
|
||||
return {"success": True, "message": "Training paused"}
|
||||
|
||||
|
||||
@router.post("/jobs/{job_id}/resume", response_model=dict)
|
||||
async def resume_training_job(job_id: str, background_tasks: BackgroundTasks):
|
||||
"""Resume a paused training job."""
|
||||
if job_id not in _state.jobs:
|
||||
raise HTTPException(status_code=404, detail="Job not found")
|
||||
|
||||
job = _state.jobs[job_id]
|
||||
if job["status"] != TrainingStatus.PAUSED.value:
|
||||
raise HTTPException(status_code=400, detail="Job is not paused")
|
||||
|
||||
job["status"] = TrainingStatus.TRAINING.value
|
||||
_state.active_job_id = job_id
|
||||
background_tasks.add_task(simulate_training_progress, job_id)
|
||||
|
||||
return {"success": True, "message": "Training resumed"}
|
||||
|
||||
|
||||
@router.post("/jobs/{job_id}/cancel", response_model=dict)
|
||||
async def cancel_training_job(job_id: str):
|
||||
"""Cancel a training job."""
|
||||
if job_id not in _state.jobs:
|
||||
raise HTTPException(status_code=404, detail="Job not found")
|
||||
|
||||
job = _state.jobs[job_id]
|
||||
job["status"] = TrainingStatus.CANCELLED.value
|
||||
job["completed_at"] = datetime.now().isoformat()
|
||||
|
||||
if _state.active_job_id == job_id:
|
||||
_state.active_job_id = None
|
||||
|
||||
return {"success": True, "message": "Training cancelled"}
|
||||
|
||||
|
||||
@router.delete("/jobs/{job_id}", response_model=dict)
|
||||
async def delete_training_job(job_id: str):
|
||||
"""Delete a training job."""
|
||||
if job_id not in _state.jobs:
|
||||
raise HTTPException(status_code=404, detail="Job not found")
|
||||
|
||||
job = _state.jobs[job_id]
|
||||
if job["status"] == TrainingStatus.TRAINING.value:
|
||||
raise HTTPException(status_code=400, detail="Cannot delete running job")
|
||||
|
||||
del _state.jobs[job_id]
|
||||
return {"success": True, "message": "Job deleted"}
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# MODEL VERSIONS
|
||||
# ============================================================================
|
||||
|
||||
@router.get("/models", response_model=List[dict])
|
||||
async def list_model_versions():
|
||||
"""Get all trained model versions."""
|
||||
return list(_state.model_versions.values())
|
||||
|
||||
|
||||
@router.get("/models/{version_id}", response_model=dict)
|
||||
async def get_model_version(version_id: str):
|
||||
"""Get details for a specific model version."""
|
||||
if version_id not in _state.model_versions:
|
||||
raise HTTPException(status_code=404, detail="Model version not found")
|
||||
return _state.model_versions[version_id]
|
||||
|
||||
|
||||
@router.post("/models/{version_id}/activate", response_model=dict)
|
||||
async def activate_model_version(version_id: str):
|
||||
"""Set a model version as active."""
|
||||
if version_id not in _state.model_versions:
|
||||
raise HTTPException(status_code=404, detail="Model version not found")
|
||||
|
||||
# Deactivate all other versions of same type
|
||||
model = _state.model_versions[version_id]
|
||||
for v in _state.model_versions.values():
|
||||
if v["model_type"] == model["model_type"]:
|
||||
v["is_active"] = False
|
||||
|
||||
model["is_active"] = True
|
||||
return {"success": True, "message": "Model activated"}
|
||||
|
||||
|
||||
@router.delete("/models/{version_id}", response_model=dict)
|
||||
async def delete_model_version(version_id: str):
|
||||
"""Delete a model version."""
|
||||
if version_id not in _state.model_versions:
|
||||
raise HTTPException(status_code=404, detail="Model version not found")
|
||||
|
||||
model = _state.model_versions[version_id]
|
||||
if model["is_active"]:
|
||||
raise HTTPException(status_code=400, detail="Cannot delete active model")
|
||||
|
||||
del _state.model_versions[version_id]
|
||||
return {"success": True, "message": "Model deleted"}
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# DATASET STATS
|
||||
# ============================================================================
|
||||
|
||||
@router.get("/dataset/stats", response_model=dict)
|
||||
async def get_dataset_stats():
|
||||
"""Get statistics about the training dataset."""
|
||||
# Get stats from zeugnis sources
|
||||
from metrics_db import get_zeugnis_stats
|
||||
|
||||
zeugnis_stats = await get_zeugnis_stats()
|
||||
|
||||
return {
|
||||
"total_documents": zeugnis_stats.get("total_documents", 0),
|
||||
"total_chunks": zeugnis_stats.get("total_documents", 0) * 12, # Estimate ~12 chunks per doc
|
||||
"training_allowed": zeugnis_stats.get("training_allowed_documents", 0),
|
||||
"by_bundesland": {
|
||||
bl["bundesland"]: bl.get("doc_count", 0)
|
||||
for bl in zeugnis_stats.get("per_bundesland", [])
|
||||
},
|
||||
"by_doc_type": {
|
||||
"verordnung": 150,
|
||||
"schulordnung": 80,
|
||||
"handreichung": 45,
|
||||
"erlass": 30,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# TRAINING STATUS
|
||||
# ============================================================================
|
||||
|
||||
@router.get("/status", response_model=dict)
|
||||
async def get_training_status():
|
||||
"""Get overall training system status."""
|
||||
active_job = None
|
||||
if _state.active_job_id and _state.active_job_id in _state.jobs:
|
||||
active_job = _state.jobs[_state.active_job_id]
|
||||
|
||||
return {
|
||||
"is_training": _state.active_job_id is not None and active_job is not None and
|
||||
active_job["status"] == TrainingStatus.TRAINING.value,
|
||||
"active_job_id": _state.active_job_id,
|
||||
"total_jobs": len(_state.jobs),
|
||||
"completed_jobs": sum(
|
||||
1 for j in _state.jobs.values()
|
||||
if j["status"] == TrainingStatus.COMPLETED.value
|
||||
),
|
||||
"failed_jobs": sum(
|
||||
1 for j in _state.jobs.values()
|
||||
if j["status"] == TrainingStatus.FAILED.value
|
||||
),
|
||||
"model_versions": len(_state.model_versions),
|
||||
"active_models": sum(1 for m in _state.model_versions.values() if m["is_active"]),
|
||||
}
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# SERVER-SENT EVENTS (SSE) ENDPOINTS
|
||||
# ============================================================================
|
||||
|
||||
async def training_metrics_generator(job_id: str, request: Request):
|
||||
"""
|
||||
SSE generator for streaming training metrics.
|
||||
|
||||
Yields JSON-encoded training status updates every 500ms.
|
||||
"""
|
||||
while True:
|
||||
# Check if client disconnected
|
||||
if await request.is_disconnected():
|
||||
break
|
||||
|
||||
# Get job status
|
||||
if job_id not in _state.jobs:
|
||||
yield f"data: {json.dumps({'error': 'Job not found'})}\n\n"
|
||||
break
|
||||
|
||||
job = _state.jobs[job_id]
|
||||
|
||||
# Build metrics response
|
||||
metrics_data = {
|
||||
"job_id": job["id"],
|
||||
"status": job["status"],
|
||||
"progress": job["progress"],
|
||||
"current_epoch": job["current_epoch"],
|
||||
"total_epochs": job["total_epochs"],
|
||||
"current_step": int(job["progress"] * job["total_epochs"]),
|
||||
"total_steps": job["total_epochs"] * 100,
|
||||
"elapsed_time_ms": 0,
|
||||
"estimated_remaining_ms": 0,
|
||||
"metrics": {
|
||||
"loss": job["loss"],
|
||||
"val_loss": job["val_loss"],
|
||||
"accuracy": job["metrics"]["accuracy"],
|
||||
"learning_rate": job["learning_rate"]
|
||||
},
|
||||
"history": [
|
||||
{
|
||||
"epoch": i + 1,
|
||||
"step": (i + 1) * 10,
|
||||
"loss": loss,
|
||||
"val_loss": job["metrics"]["val_loss_history"][i] if i < len(job["metrics"]["val_loss_history"]) else None,
|
||||
"learning_rate": job["learning_rate"],
|
||||
"timestamp": 0
|
||||
}
|
||||
for i, loss in enumerate(job["metrics"]["loss_history"][-50:])
|
||||
]
|
||||
}
|
||||
|
||||
# Calculate elapsed time
|
||||
if job["started_at"]:
|
||||
started = datetime.fromisoformat(job["started_at"])
|
||||
metrics_data["elapsed_time_ms"] = int((datetime.now() - started).total_seconds() * 1000)
|
||||
|
||||
# Calculate remaining time
|
||||
if job["estimated_completion"]:
|
||||
estimated = datetime.fromisoformat(job["estimated_completion"])
|
||||
metrics_data["estimated_remaining_ms"] = max(0, int((estimated - datetime.now()).total_seconds() * 1000))
|
||||
|
||||
# Send SSE event
|
||||
yield f"data: {json.dumps(metrics_data)}\n\n"
|
||||
|
||||
# Check if job completed
|
||||
if job["status"] in [TrainingStatus.COMPLETED.value, TrainingStatus.FAILED.value, TrainingStatus.CANCELLED.value]:
|
||||
break
|
||||
|
||||
# Wait before next update
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
|
||||
@router.get("/metrics/stream")
|
||||
async def stream_training_metrics(job_id: str, request: Request):
|
||||
"""
|
||||
SSE endpoint for streaming training metrics.
|
||||
|
||||
Streams real-time training progress for a specific job.
|
||||
|
||||
Usage:
|
||||
const eventSource = new EventSource('/api/v1/admin/training/metrics/stream?job_id=xxx')
|
||||
eventSource.onmessage = (event) => {
|
||||
const data = JSON.parse(event.data)
|
||||
console.log(data.progress, data.metrics.loss)
|
||||
}
|
||||
"""
|
||||
if job_id not in _state.jobs:
|
||||
raise HTTPException(status_code=404, detail="Job not found")
|
||||
|
||||
return StreamingResponse(
|
||||
training_metrics_generator(job_id, request),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"Connection": "keep-alive",
|
||||
"X-Accel-Buffering": "no" # Disable nginx buffering
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def batch_ocr_progress_generator(images_count: int, request: Request):
|
||||
"""
|
||||
SSE generator for batch OCR progress simulation.
|
||||
|
||||
In production, this would integrate with actual OCR processing.
|
||||
"""
|
||||
import random
|
||||
|
||||
for i in range(images_count):
|
||||
# Check if client disconnected
|
||||
if await request.is_disconnected():
|
||||
break
|
||||
|
||||
# Simulate processing time
|
||||
await asyncio.sleep(random.uniform(0.3, 0.8))
|
||||
|
||||
progress_data = {
|
||||
"type": "progress",
|
||||
"current": i + 1,
|
||||
"total": images_count,
|
||||
"progress_percent": ((i + 1) / images_count) * 100,
|
||||
"elapsed_ms": (i + 1) * 500,
|
||||
"estimated_remaining_ms": (images_count - i - 1) * 500,
|
||||
"result": {
|
||||
"text": f"Sample recognized text for image {i + 1}",
|
||||
"confidence": round(random.uniform(0.7, 0.98), 2),
|
||||
"processing_time_ms": random.randint(200, 600),
|
||||
"from_cache": random.random() < 0.2
|
||||
}
|
||||
}
|
||||
|
||||
yield f"data: {json.dumps(progress_data)}\n\n"
|
||||
|
||||
# Send completion event
|
||||
yield f"data: {json.dumps({'type': 'complete', 'total_time_ms': images_count * 500, 'processed_count': images_count})}\n\n"
|
||||
|
||||
|
||||
@router.get("/ocr/stream")
|
||||
async def stream_batch_ocr(images_count: int, request: Request):
|
||||
"""
|
||||
SSE endpoint for streaming batch OCR progress.
|
||||
|
||||
Simulates batch OCR processing with progress updates.
|
||||
In production, integrate with actual TrOCR batch processing.
|
||||
|
||||
Args:
|
||||
images_count: Number of images to process
|
||||
|
||||
Usage:
|
||||
const eventSource = new EventSource('/api/v1/admin/training/ocr/stream?images_count=10')
|
||||
eventSource.onmessage = (event) => {
|
||||
const data = JSON.parse(event.data)
|
||||
if (data.type === 'progress') {
|
||||
console.log(`${data.current}/${data.total}`)
|
||||
}
|
||||
}
|
||||
"""
|
||||
if images_count < 1 or images_count > 100:
|
||||
raise HTTPException(status_code=400, detail="images_count must be between 1 and 100")
|
||||
|
||||
return StreamingResponse(
|
||||
batch_ocr_progress_generator(images_count, request),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"Connection": "keep-alive",
|
||||
"X-Accel-Buffering": "no"
|
||||
}
|
||||
)
|
||||
Reference in New Issue
Block a user