diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 22edb61..644e5a8 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -4,7 +4,7 @@ on: push: branches: [main,development] pull_request: - branches: [main,development] + branches: [main, development, 'development-*'] jobs: lint: @@ -23,7 +23,7 @@ jobs: uses: astral-sh/setup-uv@v6 - name: Install linter - run: uv pip install --system ruff + run: uv pip install --system ruff==0.16.1 - name: Run linter run: | diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 281b6b1..af033e0 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -9,7 +9,7 @@ on: - 'alembic.ini' - '.github/workflows/tests.yml' pull_request: - branches: [main,development] + branches: [main, development, 'development-*'] paths: - '**.py' - 'requirements.txt' diff --git a/app/api/router.py b/app/api/router.py index cc59847..f182c94 100644 --- a/app/api/router.py +++ b/app/api/router.py @@ -1,6 +1,6 @@ from fastapi import APIRouter -from app.api.routes import forms, input, jobs, system, templates, weather, zipcode +from app.api.routes import extraction, forms, input, jobs, system, templates, weather, zipcode from app.core.config import API_PREFIX api_router = APIRouter() @@ -10,4 +10,5 @@ api_router.include_router(jobs.router, prefix=API_PREFIX) api_router.include_router(weather.router, prefix=API_PREFIX) api_router.include_router(zipcode.router, prefix=API_PREFIX) -api_router.include_router(input.router, prefix=API_PREFIX) \ No newline at end of file +api_router.include_router(input.router, prefix=API_PREFIX) +api_router.include_router(extraction.router, prefix=API_PREFIX) \ No newline at end of file diff --git a/app/api/routes/extraction.py b/app/api/routes/extraction.py new file mode 100644 index 0000000..76258b0 --- /dev/null +++ b/app/api/routes/extraction.py @@ -0,0 +1,131 @@ +from uuid import UUID + +from fastapi import APIRouter, Depends +from sqlmodel import Session + +from app.api.deps import get_db +from app.api.schemas.enums import ExtractionStatus, InputStatus +from app.api.schemas.extraction import ( + ExtractionCompleted, + ExtractionJobResponse, + ExtractionProcessing, + ExtractionRequest, +) +from app.api.schemas.incident_contract import IncidentContract +from app.core.config import ( + ESTIMATED_EXTRACTION_SECONDS, + EXTRACTION_POLL_INTERVAL_SECONDS, +) +from app.core.errors.base import AppError +from app.db.repositories import ( + get_extraction, + get_extraction_by_input, + get_incident_by_extract, + get_input, +) +from app.services.extraction import ExtractionService +from app.services.llm import check_ollama_available + +router = APIRouter(prefix="/extract", tags=["extraction"]) + + +@router.post("/{input_id}", response_model=ExtractionJobResponse, status_code=202) +def create_extraction( + input_id: UUID, + body: ExtractionRequest | None = None, + db: Session = Depends(get_db), +): + record = get_input(db, input_id) + if record is None: + raise AppError( + f"Input with ID {input_id} not found", + status_code=404, + error_code="INPUT_NOT_FOUND", + ) + + if record.status != InputStatus.ready: + raise AppError( + f"Input is in '{record.status}' state. Wait until status is 'ready'.", + status_code=409, + error_code="INPUT_NOT_READY", + detail={"current_status": record.status}, + ) + + existing = get_extraction_by_input(db, input_id) + if existing is not None: + raise AppError( + "An extraction already exists for this input", + status_code=409, + error_code="EXTRACTION_EXISTS", + detail={"existing_extract_id": str(existing.extract_id)}, + ) + + if not check_ollama_available(): + raise AppError( + "Ollama LLM service is not available", + status_code=503, + error_code="LLM_UNAVAILABLE", + ) + + model_override = body.model_override if body else None + svc = ExtractionService() + extraction, job = svc.start_extraction(db, input_id, model_override=model_override) + + return ExtractionJobResponse( + extract_id=extraction.extract_id, + input_id=input_id, + job_id=job.job_id, + status=extraction.status, + queued_at=extraction.created_at, + estimated_seconds=ESTIMATED_EXTRACTION_SECONDS, + poll_url=f"/api/v1/extract/{extraction.extract_id}", + ) + + +@router.get("/{extract_id}", response_model=ExtractionCompleted | ExtractionProcessing) +def get_extraction_result(extract_id: UUID, db: Session = Depends(get_db)): + extraction = get_extraction(db, extract_id) + if extraction is None: + raise AppError( + f"Extraction with ID {extract_id} not found", + status_code=404, + error_code="EXTRACT_NOT_FOUND", + ) + + if extraction.status == ExtractionStatus.completed: + incident = get_incident_by_extract(db, extract_id) + contract = IncidentContract.model_validate( + (incident.incident_contract if incident else None) or {} + ) + return ExtractionCompleted( + extract_id=extraction.extract_id, + input_id=extraction.input_id, + incident_id=incident.incident_id if incident else None, + status="completed", + incident_contract=contract, + completed_at=extraction.completed_at, + model_used=extraction.model_used, + processing_time_seconds=extraction.processing_time_seconds, + corrections=extraction.corrections, + ) + + retry_after = ( + EXTRACTION_POLL_INTERVAL_SECONDS + if extraction.status == ExtractionStatus.processing + else None + ) + partial = ( + IncidentContract.model_validate(extraction.partial_result) + if extraction.partial_result + else None + ) + return ExtractionProcessing( + extract_id=extraction.extract_id, + input_id=extraction.input_id, + status=extraction.status, + started_at=extraction.started_at, + retry_after_seconds=retry_after, + error_type=extraction.error_type, + error_detail=extraction.error_detail, + partial_result=partial, + ) diff --git a/app/api/schemas/extraction.py b/app/api/schemas/extraction.py index 962c5d6..621d8e7 100644 --- a/app/api/schemas/extraction.py +++ b/app/api/schemas/extraction.py @@ -4,7 +4,7 @@ from typing import Any, Literal from uuid import UUID -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, ConfigDict from app.api.schemas.enums import FieldSource from app.api.schemas.incident_contract import IncidentContract @@ -47,6 +47,20 @@ class ExtractionRequest(BaseModel): # Responses # --------------------------------------------------------------------------- +class ExtractionJobResponse(BaseModel): + """202 body for POST /extract/{input_id}. Carries the ids the client needs + to poll the extraction plus the underlying async job.""" + + extract_id: UUID + input_id: UUID + job_id: str + job_type: str = "extraction" + status: str + queued_at: datetime | None = None + estimated_seconds: int | None = None + poll_url: str + + class Correction(BaseModel): """One manual correction applied to the contract via PATCH.""" diff --git a/app/core/celery.py b/app/core/celery.py index a91146f..0111166 100644 --- a/app/core/celery.py +++ b/app/core/celery.py @@ -16,7 +16,7 @@ result_expires=86400, ) -celery_app.conf.include = ["app.tasks.fill", "app.tasks.purge", "app.tasks.transcribe"] +celery_app.conf.include = ["app.tasks.fill", "app.tasks.purge", "app.tasks.transcribe", "app.tasks.extract"] # Optional Celery Beat schedule — runs purge_old_submissions once a day. # Enable by running: celery -A app.core.celery beat diff --git a/app/core/config.py b/app/core/config.py index 4a7ca89..f654dc9 100644 --- a/app/core/config.py +++ b/app/core/config.py @@ -58,6 +58,13 @@ # or transcribing. Value matches the contract example (contracts/path/input.yaml). INPUT_POLL_INTERVAL_SECONDS = 5 +# Polling hint returned by GET /extract/{id} while an extraction is still +# processing. Matches the contract example (contracts/path/extraction.yaml). +EXTRACTION_POLL_INTERVAL_SECONDS = 5 + +# Advisory estimate returned in the 202 body of POST /extract/{input_id}. +ESTIMATED_EXTRACTION_SECONDS = int(os.getenv("ESTIMATED_EXTRACTION_SECONDS", "60")) + # --- API Versioning ------------------------------------------------------- API_PREFIX = "/api/v1" diff --git a/app/core/errors/handlers.py b/app/core/errors/handlers.py index ce8988a..6d923f2 100644 --- a/app/core/errors/handlers.py +++ b/app/core/errors/handlers.py @@ -1,3 +1,5 @@ +import json + from fastapi import Request from fastapi.exceptions import RequestValidationError from fastapi.responses import JSONResponse @@ -6,6 +8,23 @@ from app.core.errors.base import AppError +def _jsonable(value): + """Return value untouched if JSON-serializable, else a string form. + + FastAPI puts the offending input on each validation error. When a JSON-body + endpoint is called with the wrong content-type, that input is the raw + request body as bytes, which JSONResponse cannot encode. Coerce anything + non-serializable so the 422 renders instead of turning into a 500. + """ + try: + json.dumps(value) + return value + except (TypeError, ValueError): + if isinstance(value, bytes): + return value.decode("utf-8", "replace") + return str(value) + + def register_exception_handlers(app): @app.exception_handler(AppError) async def app_error_handler(request: Request, exc: AppError): @@ -25,7 +44,7 @@ async def validation_error_handler(request: Request, exc: RequestValidationError validation_errors.append({ "field": field or None, "issue": error.get("msg"), - "value": error.get("input"), + "value": _jsonable(error.get("input")), }) return JSONResponse( status_code=422, diff --git a/app/db/repositories.py b/app/db/repositories.py index fbce020..e2f565b 100644 --- a/app/db/repositories.py +++ b/app/db/repositories.py @@ -102,6 +102,11 @@ def get_extraction(session: Session, extract_id: UUID) -> Extraction | None: return session.get(Extraction, extract_id) +def get_extraction_by_input(session: Session, input_id: UUID) -> Extraction | None: + statement = select(Extraction).where(Extraction.input_id == input_id) + return session.exec(statement).first() + + def update_extraction(session: Session, extraction: Extraction) -> Extraction: session.add(extraction) session.commit() diff --git a/app/services/extraction.py b/app/services/extraction.py new file mode 100644 index 0000000..c3aa251 --- /dev/null +++ b/app/services/extraction.py @@ -0,0 +1,53 @@ +"""Extraction service. + +Owns the write path that turns a ready input into a queued extraction: it +creates the extraction row, creates the async job, and dispatches the worker. +The route stays a thin HTTP handler and calls straight into here. + +The worker itself is stubbed for now (app/tasks/extract.py); the real chunked +extraction is #630. +""" + +from datetime import datetime, timezone + +from sqlmodel import Session + +from app.api.schemas.enums import ExtractionStatus +from app.db.repositories import create_extraction, create_job, update_job +from app.models import Extraction, Job +from app.tasks.extract import extract_task + + +class ExtractionService: + def start_extraction( + self, + session: Session, + input_id, + model_override: str | None = None, + ) -> tuple[Extraction, Job]: + """Create the extraction row and job, then dispatch the worker. + + The extraction starts in ``processing`` so a poll right after the 202 + sees the in-flight shape. Mirrors the transcription flow: the job row is + created first with a known job_id, dispatched, then its celery_task_id + is backfilled once the broker returns a task id. + """ + now = datetime.now(timezone.utc) + extraction = Extraction( + input_id=input_id, + status=ExtractionStatus.processing, + started_at=now, + model_used=model_override, + created_at=now, + updated_at=now, + ) + extraction = create_extraction(session, extraction) + + job = Job(celery_task_id="", job_type="extraction", status="queued", model=model_override) + job = create_job(session, job) + + result = extract_task.delay(str(extraction.extract_id), job.job_id) + job.celery_task_id = result.id + job = update_job(session, job) + + return extraction, job diff --git a/app/services/llm.py b/app/services/llm.py index e2d1639..f73e51e 100644 --- a/app/services/llm.py +++ b/app/services/llm.py @@ -9,6 +9,14 @@ logger = get_logger(__name__) +def check_ollama_available() -> bool: + """Return True if the Ollama server responds with a successful status.""" + try: + return requests.get(f"{OLLAMA_HOST}/api/tags", timeout=3).ok + except requests.exceptions.RequestException: + return False + + class LLM: def __init__(self, transcript_text: str=None, target_fields: list=None, json_dict: dict=None, model: str=None): self._transcript_text = transcript_text diff --git a/app/tasks/extract.py b/app/tasks/extract.py new file mode 100644 index 0000000..d072743 --- /dev/null +++ b/app/tasks/extract.py @@ -0,0 +1,47 @@ +import logging +from datetime import datetime, timezone +from uuid import UUID + +from app.api.schemas.enums import ExtractionStatus +from app.core.celery import celery_app +from app.db.database import get_session +from app.db.repositories import get_extraction, get_job_by_uuid, update_extraction, update_job + +logger = logging.getLogger(__name__) + + +@celery_app.task(name="extract_incident") +def extract_task(extract_id_str: str, job_id_str: str) -> dict: + """Stub extraction worker. + + #629 ships the queue slice: the extraction row and job exist and the task + runs, but the real chunked extraction (split the narrative into field + groups, call the model, validate, stitch a contract, create the draft + incident) lands in #630. For now the task only marks the job in-flight and + leaves the extraction in its ``processing`` state, which is what + GET /extract/{id} reports back to pollers. + """ + session = next(get_session()) + extract_id = UUID(extract_id_str) + try: + extraction = get_extraction(session, extract_id) + job = get_job_by_uuid(session, job_id_str) + + now = datetime.now(timezone.utc) + if extraction: + extraction.status = ExtractionStatus.processing + extraction.started_at = extraction.started_at or now + extraction.updated_at = now + update_extraction(session, extraction) + if job: + job.status = "processing" + job.updated_at = now + update_job(session, job) + + logger.info( + "extract_task stub ran for extraction %s; real worker is #630", + extract_id_str, + ) + return {"extract_id": extract_id_str, "job_id": job_id_str} + finally: + session.close() diff --git a/requirements-dev.txt b/requirements-dev.txt index 71a9a7b..092631f 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -2,4 +2,4 @@ # datamodel-code-generator turns the incident contract into Pydantic models # (see scripts/generate_contract_models.py, run via `make generate-contract-models`). datamodel-code-generator==0.25.9 -ruff==0.6.9 +ruff==0.16.1 diff --git a/ruff.toml b/ruff.toml new file mode 100644 index 0000000..c04f533 --- /dev/null +++ b/ruff.toml @@ -0,0 +1,7 @@ +target-version = "py311" +line-length = 100 + +[lint] +# Pin the rule set explicitly. Ruff's defaults changed in 0.16, so relying on +# them means a version bump silently changes what CI enforces. +select = ["E4", "E7", "E9", "F"] diff --git a/tests/test_v1_extraction.py b/tests/test_v1_extraction.py new file mode 100644 index 0000000..e92f7f7 --- /dev/null +++ b/tests/test_v1_extraction.py @@ -0,0 +1,239 @@ +"""Tests for POST /api/v1/extract/{input_id} and GET /api/v1/extract/{extract_id}. + +Endpoint tests only — dispatch is mocked (no broker) and Ollama availability is +patched. The real chunked worker lands in #630, so there is no task unit here; +the stub task is exercised indirectly by the POST dispatch assertions. +""" + +from datetime import datetime, timezone +from unittest.mock import MagicMock, patch +from uuid import UUID, uuid4 + +from app.api.schemas.enums import ExtractionStatus, InputStatus, InputType, ReportStatus +from app.db.repositories import ( + create_extraction, + create_incident, + create_input, + get_extraction, + get_job_by_uuid, +) +from app.models import Extraction, Incident, Input + +POST_URL = "/api/v1/extract" +GET_URL = "/api/v1/extract" + +_CONTRACT = { + "schema_version": "1.1.0", + "schema_name": "fireform_incident_contract", + "incident": {"name": "Bear Creek Wildfire"}, + "location": {"city": "Reno", "state": "NV"}, +} + + +# --------------------------------------------------------------------------- +# Seed helpers +# --------------------------------------------------------------------------- + +def _ready_input(db, status=InputStatus.ready) -> Input: + now = datetime.now(timezone.utc) + record = Input( + input_type=InputType.text, + status=status, + transcript="Structure fire at 42 Oak St, two engines on scene, one civilian injury.", + character_count=70, + word_count=13, + created_at=now, + updated_at=now, + ) + return create_input(db, record) + + +def _completed_extraction_with_incident(db, input_id) -> tuple[Extraction, Incident]: + now = datetime.now(timezone.utc) + extraction = create_extraction( + db, + Extraction( + input_id=input_id, + status=ExtractionStatus.completed, + started_at=now, + completed_at=now, + model_used="qwen2.5:1.5b", + processing_time_seconds=42.0, + ), + ) + incident = create_incident( + db, + Incident( + extract_id=extraction.extract_id, + status=ReportStatus.draft, + incident_contract=_CONTRACT, + ), + ) + return extraction, incident + + +# --------------------------------------------------------------------------- +# POST /api/v1/extract/{input_id} +# --------------------------------------------------------------------------- + +class TestCreateExtraction: + + def _post(self, client, input_id, body=None, ollama_up=True, celery_id="celery-extract-001"): + mock_result = MagicMock() + mock_result.id = celery_id + with patch("app.api.routes.extraction.check_ollama_available", return_value=ollama_up), \ + patch("app.services.extraction.extract_task") as mock_task: + mock_task.delay.return_value = mock_result + resp = client.post(f"{POST_URL}/{input_id}", json=body) + return resp, mock_task + + def test_202_returns_required_fields(self, client, db): + inp = _ready_input(db) + resp, _ = self._post(client, inp.input_id) + assert resp.status_code == 202 + body = resp.json() + assert body["status"] == "processing" + assert body["job_type"] == "extraction" + assert body["input_id"] == str(inp.input_id) + assert "extract_id" in body + assert "job_id" in body + assert body["poll_url"] == f"/api/v1/extract/{body['extract_id']}" + assert body["estimated_seconds"] == 60 + + def test_202_creates_processing_extraction_row(self, client, db): + inp = _ready_input(db) + resp, _ = self._post(client, inp.input_id) + extraction = get_extraction(db, UUID(resp.json()["extract_id"])) + assert extraction is not None + assert extraction.status == ExtractionStatus.processing + assert extraction.input_id == inp.input_id + assert extraction.started_at is not None + + def test_202_creates_extraction_job_row(self, client, db): + inp = _ready_input(db) + resp, _ = self._post(client, inp.input_id) + job = get_job_by_uuid(db, resp.json()["job_id"]) + assert job is not None + assert job.job_type == "extraction" + assert job.celery_task_id == "celery-extract-001" + + def test_202_dispatch_called_with_extract_id_and_job_id(self, client, db): + inp = _ready_input(db) + resp, mock_task = self._post(client, inp.input_id) + body = resp.json() + mock_task.delay.assert_called_once() + args = mock_task.delay.call_args[0] + assert args[0] == body["extract_id"] + assert args[1] == body["job_id"] + + def test_202_model_override_stored_on_job(self, client, db): + inp = _ready_input(db) + resp, _ = self._post(client, inp.input_id, body={"model_override": "llama3:8b"}) + job = get_job_by_uuid(db, resp.json()["job_id"]) + assert job.model == "llama3:8b" + + def test_404_input_not_found(self, client, db): + resp, _ = self._post(client, uuid4()) + assert resp.status_code == 404 + assert resp.json()["error_code"] == "INPUT_NOT_FOUND" + + def test_409_input_not_ready(self, client, db): + inp = _ready_input(db, status=InputStatus.transcribing) + resp, _ = self._post(client, inp.input_id) + assert resp.status_code == 409 + body = resp.json() + assert body["error_code"] == "INPUT_NOT_READY" + assert body["detail"]["current_status"] == "transcribing" + + def test_409_extraction_already_exists(self, client, db): + inp = _ready_input(db) + existing = create_extraction(db, Extraction(input_id=inp.input_id)) + resp, _ = self._post(client, inp.input_id) + assert resp.status_code == 409 + body = resp.json() + assert body["error_code"] == "EXTRACTION_EXISTS" + assert body["detail"]["existing_extract_id"] == str(existing.extract_id) + + def test_503_ollama_unavailable(self, client, db): + inp = _ready_input(db) + resp, _ = self._post(client, inp.input_id, ollama_up=False) + assert resp.status_code == 503 + assert resp.json()["error_code"] == "LLM_UNAVAILABLE" + + def test_422_non_json_body_does_not_500(self, client, db): + # A wrong content-type puts the raw bytes body on the validation error; + # the error handler must still render a 422 rather than blow up on + # serializing bytes. + inp = _ready_input(db) + with patch("app.api.routes.extraction.check_ollama_available", return_value=True): + resp = client.post( + f"{POST_URL}/{inp.input_id}", + content=b"{}", + headers={"Content-Type": "application/x-www-form-urlencoded"}, + ) + assert resp.status_code == 422 + assert resp.json()["error_code"] == "VALIDATION_ERROR" + + +# --------------------------------------------------------------------------- +# GET /api/v1/extract/{extract_id} +# --------------------------------------------------------------------------- + +class TestGetExtraction: + + def test_200_processing_shape(self, client, db): + inp = _ready_input(db) + extraction = create_extraction( + db, + Extraction( + input_id=inp.input_id, + status=ExtractionStatus.processing, + started_at=datetime.now(timezone.utc), + ), + ) + resp = client.get(f"{GET_URL}/{extraction.extract_id}") + assert resp.status_code == 200 + body = resp.json() + assert body["status"] == "processing" + assert body["extract_id"] == str(extraction.extract_id) + assert body["input_id"] == str(inp.input_id) + assert body["retry_after_seconds"] == 5 + assert "incident_contract" not in body + + def test_200_completed_shape_embeds_contract(self, client, db): + inp = _ready_input(db) + extraction, incident = _completed_extraction_with_incident(db, inp.input_id) + resp = client.get(f"{GET_URL}/{extraction.extract_id}") + assert resp.status_code == 200 + body = resp.json() + assert body["status"] == "completed" + assert body["incident_id"] == str(incident.incident_id) + assert body["model_used"] == "qwen2.5:1.5b" + assert body["processing_time_seconds"] == 42.0 + assert body["incident_contract"]["incident"]["name"] == "Bear Creek Wildfire" + assert body["incident_contract"]["location"]["city"] == "Reno" + + def test_200_failed_shape(self, client, db): + inp = _ready_input(db) + extraction = create_extraction( + db, + Extraction( + input_id=inp.input_id, + status=ExtractionStatus.failed, + started_at=datetime.now(timezone.utc), + error_type="LLM_UNAVAILABLE", + error_detail="Ollama connection refused", + ), + ) + resp = client.get(f"{GET_URL}/{extraction.extract_id}") + assert resp.status_code == 200 + body = resp.json() + assert body["status"] == "failed" + assert body["error_type"] == "LLM_UNAVAILABLE" + assert body["error_detail"] == "Ollama connection refused" + assert body["retry_after_seconds"] is None + + def test_404_extraction_not_found(self, client, db): + resp = client.get(f"{GET_URL}/{uuid4()}") + assert resp.status_code == 404 + assert resp.json()["error_code"] == "EXTRACT_NOT_FOUND"