diff --git a/backend/app/alembic/versions/074_add_prompt_improvement_jobtype.py b/backend/app/alembic/versions/074_add_prompt_improvement_jobtype.py new file mode 100644 index 000000000..36381094d --- /dev/null +++ b/backend/app/alembic/versions/074_add_prompt_improvement_jobtype.py @@ -0,0 +1,27 @@ +"""add PROMPT_IMPROVEMENT jobtype enum value + +Revision ID: 074 +Revises: 073 +Create Date: 2026-07-09 00:00:00.000000 + +Prompt improvement moved from a synchronous request to a Celery job, so its runs +are tracked in the `job` table under this new jobtype. ALTER TYPE ... ADD VALUE +must run outside a transaction, hence the autocommit block. +""" + +from alembic import op + +revision = "074" +down_revision = "073" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + with op.get_context().autocommit_block(): + op.execute("ALTER TYPE jobtype ADD VALUE IF NOT EXISTS 'PROMPT_IMPROVEMENT'") + + +def downgrade() -> None: + # Postgres cannot drop a value from an enum type; leaving it is harmless. + pass diff --git a/backend/app/api/docs/evaluation/improve_prompt.md b/backend/app/api/docs/evaluation/improve_prompt.md index 399e223c9..811e58b86 100644 --- a/backend/app/api/docs/evaluation/improve_prompt.md +++ b/backend/app/api/docs/evaluation/improve_prompt.md @@ -1,13 +1,59 @@ -Generate a new, AI-improved prompt iteration for the configuration that was evaluated by this run. +Enqueue an AI prompt-improvement job for the configuration that was evaluated by this run. -The run must have `status = completed` and a `score_trace_url` pointing to the stored trace file. No request body is required — the service reads the trace file directly from S3, uploads it to the Anthropic Files API, and asks Claude to identify low-performing answers (by scores and divergence from ground truth) and rewrite the system prompt to address them. Only `completion.params.instructions` is changed — model, knowledge base, and all other settings are preserved. +The heavy LLM round-trip runs on a background worker, so this endpoint validates preconditions and returns immediately with `202 Accepted` and a `job_id`. The run must have `status = completed` and a `score_trace_url` pointing to the stored trace file. The worker reads the trace file, asks Claude to identify low-performing answers (by scores and divergence from ground truth) and rewrite the system prompt. Only `completion.params.instructions` is changed; model, knowledge base, and all other settings are preserved. -The new prompt is persisted as the next `config_version` for the evaluated configuration. AI provenance and the source evaluation run id are recorded in the version's `commit_message` field (prefixed with `[AI Generated]` for auditability). +On success the new prompt is persisted as the next `config_version` for the evaluated configuration, with AI provenance and the source evaluation run id recorded in the version's `commit_message` (prefixed with `[AI Generated]`). -**Error codes:** +**Request body (required):** + +```json +{ "callback_url": "https://your-service.example.com/webhooks/prompt-improvement" } +``` + +- `callback_url` — HTTPS webhook that receives the result. Must be `https://` and must not resolve to a private, loopback, or cloud-metadata address. + +**Callback delivery:** + +Once the worker finishes (success or failure), it POSTs a single JSON body to `callback_url`. Delivery is **best-effort, single-attempt** — there is no retry. The `config_version` is persisted regardless of callback outcome and remains recoverable from the configuration's version list, so a failed callback loses the notification, not the work. + +When a webhook signing secret is configured for the project, the request carries `X-Webhook-Signature` (HMAC-SHA256 of the body) and `X-Webhook-Timestamp` headers; verify them before trusting the payload. + +The body is a standard `APIResponse` envelope whose `data` is a `PromptImprovementJobPublic`: + +```json +{ + "success": true, + "data": { + "job_id": "…", + "status": "SUCCESS", + "config_version": { "…": "…" }, + "error_message": null + }, + "error": null, + "metadata": null +} +``` + +```json +{ + "success": false, + "data": { + "job_id": "…", + "status": "FAILED", + "config_version": null, + "error_message": "…" + }, + "error": "…", + "metadata": null +} +``` + +LLM and trace-download failures surface only via the failure callback, not on the `202` response. + +**Validation errors (returned synchronously before the job is created):** - `404 evaluation_not_found` — no run with this id in the caller's project. - `409 evaluation_not_completed` — run is not yet completed. - `409 source_config_unavailable` — the run's config or config_version has been deleted. - `422 traces_not_available` — the run has no `score_trace_url`; trace file is required. -- `502 trace_download_failed` — the trace file could not be retrieved from storage. -- `502 prompt_generation_failed` — the platform Anthropic key is not configured, or the LLM call failed. +- `422 invalid_callback_url` — `callback_url` is missing, not `https://`, or resolves to a blocked (private/loopback/metadata) address. +- `500 prompt_improvement_enqueue_failed` — the job could not be queued. diff --git a/backend/app/api/routes/evaluations/evaluation.py b/backend/app/api/routes/evaluations/evaluation.py index b51b1958a..ff3ca7f3b 100644 --- a/backend/app/api/routes/evaluations/evaluation.py +++ b/backend/app/api/routes/evaluations/evaluation.py @@ -12,23 +12,27 @@ Query, ) - from app.api.deps import AuthContextDep, SessionDep from app.api.permissions import Permission, require_permission from app.core.rate_monitor import monitor_rate from app.crud.evaluations import list_evaluation_runs as list_evaluation_runs_crud from app.crud.evaluations.core import group_traces_by_question_id -from app.models.config.version import ConfigVersionPublic -from app.models.evaluation import EvaluationRunPublic, RunModeEnum +from app.models.evaluation import ( + EvaluationRunPublic, + ImprovePromptRequest, + RunModeEnum, +) +from app.models.llm.response import LLMJobImmediatePublic from app.services.evaluations import ( get_evaluation_with_scores, - improve_prompt, + start_prompt_improvement_job, validate_and_start_batch_evaluation, validate_and_start_fast_evaluation, ) from app.utils import ( APIResponse, load_description, + validate_callback_url, ) logger = logging.getLogger(__name__) @@ -208,26 +212,43 @@ def get_evaluation_run_status( @router.post( "/{evaluation_id}/improve-prompt", description=load_description("evaluation/improve_prompt.md"), - response_model=APIResponse[ConfigVersionPublic], - status_code=201, + response_model=APIResponse[LLMJobImmediatePublic], + status_code=202, dependencies=[Depends(require_permission(Permission.REQUIRE_PROJECT))], ) def improve_evaluation_prompt( evaluation_id: int, session: SessionDep, auth_context: AuthContextDep, -) -> APIResponse[ConfigVersionPublic]: - """Generate an AI-improved prompt iteration from a completed evaluation run.""" + request: ImprovePromptRequest, +) -> APIResponse[LLMJobImmediatePublic]: + """Enqueue an AI prompt-improvement job for a completed evaluation run.""" logger.info( f"[improve_evaluation_prompt] Starting | evaluation_id={evaluation_id} " f"org_id={auth_context.organization_.id} project_id={auth_context.project_.id}" ) - new_version = improve_prompt( + # No global ValueError handler exists, so map the SSRF guard's ValueError to + # a 422 here rather than letting it surface as a 500. + try: + validate_callback_url(str(request.callback_url)) + except ValueError as exc: + raise HTTPException(status_code=422, detail=f"invalid_callback_url: {exc}") + + job = start_prompt_improvement_job( session=session, evaluation_id=evaluation_id, organization_id=auth_context.organization_.id, project_id=auth_context.project_.id, + callback_url=str(request.callback_url), ) - return APIResponse.success_response(data=new_version) + return APIResponse.success_response( + data=LLMJobImmediatePublic( + job_id=job.id, + status=job.status.value, + message="Prompt improvement is running; the result will be delivered to your callback_url.", + job_inserted_at=job.inserted_at, + job_updated_at=job.updated_at, + ) + ) diff --git a/backend/app/celery/tasks/job_execution.py b/backend/app/celery/tasks/job_execution.py index 0a6320612..000c6025a 100644 --- a/backend/app/celery/tasks/job_execution.py +++ b/backend/app/celery/tasks/job_execution.py @@ -246,6 +246,25 @@ def run_evaluation_batch_submission( ) +# Priority 6 (fast-eval tier): user-blocking interactive prompt iteration in the +# evaluation domain, above default batch work but below core LLM call/chain jobs. +@celery_app.task(bind=True, queue="default", priority=6) +@gevent_timeout(settings.CELERY_TASK_SOFT_TIME_LIMIT, "run_prompt_improvement") +def run_prompt_improvement(self, project_id: int, job_id: str, trace_id: str, **kwargs): + from app.services.evaluations.prompt_improvement import execute_prompt_improvement + + _set_trace(trace_id) + return _run_with_otel_parent( + self, + lambda: execute_prompt_improvement( + project_id=project_id, + job_id=job_id, + task_id=current_task.request.id, + **kwargs, + ), + ) + + @celery_app.task(bind=True, queue="default", priority=2) @gevent_timeout(settings.CELERY_TASK_SOFT_TIME_LIMIT, "run_stt_batch_submission") def run_stt_batch_submission( diff --git a/backend/app/celery/utils.py b/backend/app/celery/utils.py index 149b54b7e..7496620f0 100644 --- a/backend/app/celery/utils.py +++ b/backend/app/celery/utils.py @@ -180,6 +180,24 @@ def start_evaluation_batch_submission( return task_id +def start_prompt_improvement( + project_id: int, job_id: str, trace_id: str = "N/A", **kwargs +) -> str: + from app.celery.tasks.job_execution import run_prompt_improvement + + task_id = _enqueue_with_trace_context( + run_prompt_improvement, + project_id=project_id, + job_id=job_id, + trace_id=trace_id, + **kwargs, + ) + logger.info( + f"[start_prompt_improvement] Started job {job_id} with Celery task {task_id}" + ) + return task_id + + def start_stt_batch_submission( project_id: int, job_id: str, trace_id: str = "N/A", **kwargs ) -> str: diff --git a/backend/app/models/evaluation.py b/backend/app/models/evaluation.py index 95661dae9..92ef8770d 100644 --- a/backend/app/models/evaluation.py +++ b/backend/app/models/evaluation.py @@ -3,12 +3,13 @@ from typing import TYPE_CHECKING, Any, Optional from uuid import UUID -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, HttpUrl from sqlalchemy import Boolean, Column, Index, Text, UniqueConstraint, text from sqlalchemy.dialects.postgresql import ENUM, JSONB from sqlmodel import Field as SQLField, Relationship, SQLModel from app.core.util import now +from app.models.config.version import ConfigVersionPublic if TYPE_CHECKING: from .batch_job import BatchJob @@ -549,3 +550,20 @@ class EvaluationDatasetPublic(SQLModel): project_id: int inserted_at: datetime updated_at: datetime + + +class ImprovePromptRequest(SQLModel): + """Body for POST /evaluations/{id}/improve-prompt.""" + + callback_url: HttpUrl = Field( + description="HTTPS webhook that receives the result once the job finishes." + ) + + +class PromptImprovementJobPublic(SQLModel): + """Callback payload body: the new config_version once the job succeeds.""" + + job_id: UUID + status: str + config_version: ConfigVersionPublic | None = None + error_message: str | None = None diff --git a/backend/app/models/job.py b/backend/app/models/job.py index 32737e2d9..501407cef 100644 --- a/backend/app/models/job.py +++ b/backend/app/models/job.py @@ -22,6 +22,7 @@ class JobType(str, Enum): LLM_API = "LLM_API" LLM_CHAIN = "LLM_CHAIN" LLM_GUARDRAILS = "LLM_GUARDRAILS" + PROMPT_IMPROVEMENT = "PROMPT_IMPROVEMENT" class Job(SQLModel, table=True): @@ -67,7 +68,7 @@ class Job(SQLModel, table=True): job_type: JobType = Field( description="Type of job being executed (e.g., response, ingestion).", sa_column_kwargs={ - "comment": "Type of job being executed (e.g., RESPONSE, LLM_API, LLM_CHAIN, LLM_GUARDRAILS)" + "comment": "Type of job being executed (e.g., RESPONSE, LLM_API, LLM_CHAIN, LLM_GUARDRAILS, PROMPT_IMPROVEMENT)" }, ) meta: dict[str, Any] | None = Field( diff --git a/backend/app/services/evaluations/__init__.py b/backend/app/services/evaluations/__init__.py index 3c7ce196a..5a06cb7dc 100644 --- a/backend/app/services/evaluations/__init__.py +++ b/backend/app/services/evaluations/__init__.py @@ -12,7 +12,9 @@ validate_and_start_fast_evaluation, ) from app.services.evaluations.prompt_improvement import ( - improve_prompt, + execute_prompt_improvement, + start_prompt_improvement_job, + validate_improve_prompt, ) from app.services.evaluations.validators import ( ALLOWED_EXTENSIONS, diff --git a/backend/app/services/evaluations/prompt_improvement.py b/backend/app/services/evaluations/prompt_improvement.py index 068417469..eb7cfadef 100644 --- a/backend/app/services/evaluations/prompt_improvement.py +++ b/backend/app/services/evaluations/prompt_improvement.py @@ -1,5 +1,10 @@ """AI-assisted prompt improvement service. +Split into a fast request-side path (validate + enqueue a Celery job) and a +worker-side path that does the heavy Anthropic round-trip. The LLM call stays +synchronous — blocking in a Celery worker is fine, and it holds no FastAPI +threadpool thread. + Loads the evaluation run's stored score traces from object storage, asks Claude to rewrite the system prompt, and persists the result as a new config_version. """ @@ -7,22 +12,33 @@ import copy import json import logging +from uuid import UUID import anthropic +from asgi_correlation_id import correlation_id +from celery.exceptions import SoftTimeLimitExceeded from fastapi import HTTPException +from gevent import Timeout from sqlmodel import Session +from app.celery.utils import start_prompt_improvement from app.core.cloud.storage import get_cloud_storage from app.core.config import settings +from app.core.db import engine from app.core.storage_utils import load_json_from_object_store from app.crud.config.version import ConfigVersionCrud from app.crud.evaluations.core import get_evaluation_run_by_id +from app.crud.jobs import JobCrud from app.models.config.config import ConfigTag from app.models.config.version import ( + ConfigVersion, ConfigVersionPublic, ConfigVersionUpdate, ) +from app.models.evaluation import EvaluationRun, PromptImprovementJobPublic +from app.models.job import Job, JobStatus, JobType, JobUpdate from app.services.llm.providers.claude import ClaudeProvider +from app.utils import APIResponse, get_webhook_secret, send_callback logger = logging.getLogger(__name__) @@ -56,22 +72,44 @@ _COMPLETED_STATUS = "completed" -def improve_prompt( +def _resolve_source_version( + *, + session: Session, + run: EvaluationRun, + project_id: int, +) -> ConfigVersion | None: + """Resolve the exact config_version the run evaluated. + + read_one() raises 404 when the config itself is missing/soft-deleted and + returns None when only the version is gone — the caller treats both as the + run's source config no longer resolving. + """ + try: + return ConfigVersionCrud( + session=session, + config_id=run.config_id, + project_id=project_id, + tag=ConfigTag.DEFAULT, + ).read_one(version_number=run.config_version) + except HTTPException as exc: + if exc.status_code != 404: + raise + return None + + +def validate_improve_prompt( *, session: Session, evaluation_id: int, organization_id: int, project_id: int, -) -> ConfigVersionPublic: - """Run the full prompt-improvement flow synchronously and return the new version. +) -> EvaluationRun: + """Run the cheap DB precondition checks for prompt improvement. - Raises HTTPException for all domain errors so the route stays thin. + Raises HTTPException for every domain failure so the request-side caller + returns a real 4xx before any job is enqueued. No LLM call, no trace + download. Returns the run for reuse by the worker path. """ - logger.info( - f"[improve_prompt] Starting | evaluation_id={evaluation_id} " - f"project_id={project_id}" - ) - run = get_evaluation_run_by_id( session=session, evaluation_id=evaluation_id, @@ -97,7 +135,6 @@ def improve_prompt( "cannot improve prompt" ), ) - if run.config_id is None or run.config_version is None: # Both FKs are nullable; without them there's no version to improve. raise HTTPException( @@ -105,74 +142,347 @@ def improve_prompt( detail="source_config_unavailable: run has no config_id/config_version reference", ) - # read_one() raises 404 when the config itself is missing/soft-deleted and - # returns None when only the version is gone — both are 409 here, since the - # run referenced a config that no longer resolves. - try: - version = ConfigVersionCrud( - session=session, - config_id=run.config_id, - project_id=project_id, - tag=ConfigTag.DEFAULT, - ).read_one(version_number=run.config_version) - except HTTPException as exc: - if exc.status_code != 404: - raise - version = None - if version is None: + if _resolve_source_version(session=session, run=run, project_id=project_id) is None: raise HTTPException( status_code=409, detail="source_config_unavailable: the run's config or config_version is missing or soft-deleted", ) - blob = version.config_blob or {} - params = blob.get("completion", {}).get("params", {}) or {} - current_instructions = params.get("instructions") or "" + return run + - storage = get_cloud_storage(session=session, project_id=project_id) - traces = load_json_from_object_store(storage=storage, url=run.score_trace_url) - if traces is None: +def start_prompt_improvement_job( + *, + session: Session, + evaluation_id: int, + organization_id: int, + project_id: int, + callback_url: str, +) -> Job: + """Validate preconditions, create a job row, and enqueue the worker task. + + Returns the created Job immediately; the result is delivered to callback_url. + """ + validate_improve_prompt( + session=session, + evaluation_id=evaluation_id, + organization_id=organization_id, + project_id=project_id, + ) + + trace_id = correlation_id.get() or "N/A" + job = JobCrud(session=session).create( + job_type=JobType.PROMPT_IMPROVEMENT, + trace_id=trace_id, + project_id=project_id, + ) + + logger.info( + f"[start_prompt_improvement_job] Job created | job_id={job.id} " + f"evaluation_id={evaluation_id} project_id={project_id}" + ) + + try: + task_id = start_prompt_improvement( + project_id=project_id, + job_id=str(job.id), + trace_id=trace_id, + organization_id=organization_id, + evaluation_id=evaluation_id, + callback_url=callback_url, + ) + except Exception as exc: + logger.error( + f"[start_prompt_improvement_job] Failed to enqueue | job_id={job.id} " + f"evaluation_id={evaluation_id} | {exc}", + exc_info=True, + ) + JobCrud(session=session).update( + job.id, + JobUpdate( + status=JobStatus.FAILED, + error_message=f"Failed to queue prompt improvement: {exc}", + ), + ) raise HTTPException( - status_code=502, - detail="trace_download_failed: could not retrieve trace file from storage", + status_code=500, + detail="prompt_improvement_enqueue_failed: could not queue the job", ) - improved_instructions, rationale = _draft_improved_prompt( - current_instructions=current_instructions, - config_params=params, - traces=traces, + logger.info( + f"[start_prompt_improvement_job] Enqueued | job_id={job.id} task_id={task_id}" ) + return job - # Derive the new version from the *evaluated* version's blob (not the latest - # active one) so model, knowledge base, and other params stay apples-to-apples - # with what was actually scored — only the prompt text changes. - improved_blob = copy.deepcopy(blob) - improved_blob.setdefault("completion", {}).setdefault("params", {})[ - "instructions" - ] = improved_instructions - commit_message = ( - f"{AI_GENERATED_MARKER} improved from config version v{run.config_version} " - f"(Evaluation: {run.run_name}) {rationale}" - )[:COMMIT_MESSAGE_MAX_LENGTH] +def _build_improve_prompt_payload( + *, + job_id: UUID, + config_version: ConfigVersionPublic | None, + error_message: str | None, +) -> dict: + """Build the callback body: PromptImprovementJobPublic inside an APIResponse. + + Pre-dump the inner model to JSON so no UUID/datetime survives into the dict + send_callback serialises. + """ + status = JobStatus.FAILED if error_message else JobStatus.SUCCESS + job_public = PromptImprovementJobPublic( + job_id=job_id, + status=status.value, + config_version=config_version, + error_message=error_message, + ).model_dump(mode="json") + envelope = ( + APIResponse.failure_response(error=error_message, data=job_public) + if error_message + else APIResponse.success_response(data=job_public) + ) + return envelope.model_dump() + - new_version = ConfigVersionCrud( +def _resolve_success_config_version( + *, + session: Session, + evaluation_id: int, + organization_id: int, + project_id: int, + version_number: int | None, +) -> ConfigVersionPublic | None: + """Rebuild the minted config_version from a completed job's meta (redelivery).""" + if version_number is None: + return None + run = get_evaluation_run_by_id( session=session, - config_id=run.config_id, + evaluation_id=evaluation_id, + organization_id=organization_id, project_id=project_id, - ).create_or_raise( - ConfigVersionUpdate( - config_blob=improved_blob, - commit_message=commit_message, - ) ) + if run is None or run.config_id is None: + return None + version = ConfigVersionCrud( + session=session, + config_id=run.config_id, + project_id=project_id, + ).read_one(version_number=version_number) + return ConfigVersionPublic.model_validate(version) if version else None + + +def _send_improve_prompt_callback( + *, + callback_url: str | None, + payload: dict, + project_id: int, + organization_id: int, +) -> None: + """Best-effort single POST to callback_url; no-op when no URL was supplied.""" + if not callback_url: + return + webhook_secret = get_webhook_secret(project_id, organization_id) + send_callback(callback_url, payload, webhook_secret=webhook_secret) + + +def execute_prompt_improvement( + *, + project_id: int, + job_id: str, + organization_id: int, + evaluation_id: int, + **kwargs, +) -> dict: + """Worker entrypoint: run the full prompt-improvement flow and record it on the Job. + + Opens its own session (the request session is long gone). On success the new + config_version's id/version and the rationale land on Job.meta; on failure the + Job is marked FAILED with a clean message and the error is re-raised so Celery + records the failure. + """ + task_id = kwargs.get("task_id") + callback_url = kwargs.get("callback_url") + job_uuid = UUID(job_id) logger.info( - f"[improve_prompt] Done | evaluation_id={evaluation_id} " - f"new_version_id={new_version.id} version={new_version.version}" + f"[execute_prompt_improvement] Starting | job_id={job_id} " + f"evaluation_id={evaluation_id} project_id={project_id}" ) - return ConfigVersionPublic.model_validate(new_version) + with Session(engine) as session: + job_crud = JobCrud(session=session) + + # Idempotency (Celery redelivers): a SUCCESS job already minted a + # config_version and paid for the LLM call, so re-running would double-spend + # and create a duplicate version. A stuck PROCESSING job (worker died + # mid-run) is intentionally allowed to re-run. + existing = job_crud.get(job_id=job_uuid, project_id=project_id) + if existing and existing.status == JobStatus.SUCCESS: + logger.info( + f"[execute_prompt_improvement] Redelivery of completed job, skipping | " + f"job_id={job_id}" + ) + # At-least-once delivery: the first callback may have failed before the + # worker died, so re-send. The client may receive a duplicate callback. + _send_improve_prompt_callback( + callback_url=callback_url, + payload=_build_improve_prompt_payload( + job_id=job_uuid, + config_version=_resolve_success_config_version( + session=session, + evaluation_id=evaluation_id, + organization_id=organization_id, + project_id=project_id, + version_number=(existing.meta or {}).get("version"), + ), + error_message=None, + ), + project_id=project_id, + organization_id=organization_id, + ) + return {"success": True, **(existing.meta or {})} + + job_crud.update( + job_uuid, JobUpdate(status=JobStatus.PROCESSING, task_id=task_id) + ) + + try: + # Re-validate defensively: the run may have changed between enqueue and pickup. + run = validate_improve_prompt( + session=session, + evaluation_id=evaluation_id, + organization_id=organization_id, + project_id=project_id, + ) + version = _resolve_source_version( + session=session, run=run, project_id=project_id + ) + if version is None: + raise RuntimeError( + "source_config_unavailable: the run's config or config_version is missing or soft-deleted" + ) + + blob = version.config_blob or {} + params = blob.get("completion", {}).get("params", {}) or {} + current_instructions = params.get("instructions") or "" + + storage = get_cloud_storage(session=session, project_id=project_id) + traces = load_json_from_object_store( + storage=storage, url=run.score_trace_url + ) + if traces is None: + raise RuntimeError( + "trace_download_failed: could not retrieve trace file from storage" + ) + + improved_instructions, rationale = _draft_improved_prompt( + current_instructions=current_instructions, + config_params=params, + traces=traces, + ) + + # Derive the new version from the *evaluated* version's blob (not the + # latest active one) so model, knowledge base, and other params stay + # apples-to-apples with what was scored — only the prompt text changes. + improved_blob = copy.deepcopy(blob) + improved_blob.setdefault("completion", {}).setdefault("params", {})[ + "instructions" + ] = improved_instructions + + commit_message = ( + f"{AI_GENERATED_MARKER} improved from config version v{run.config_version} " + f"(Evaluation: {run.run_name}) {rationale}" + )[:COMMIT_MESSAGE_MAX_LENGTH] + + new_version = ConfigVersionCrud( + session=session, + config_id=run.config_id, + project_id=project_id, + ).create_or_raise( + ConfigVersionUpdate( + config_blob=improved_blob, + commit_message=commit_message, + ) + ) + + job_crud.update( + job_uuid, + JobUpdate( + status=JobStatus.SUCCESS, + meta={ + "version": new_version.version, + "rationale": rationale, + }, + ), + ) + + logger.info( + f"[execute_prompt_improvement] Done | job_id={job_id} " + f"evaluation_id={evaluation_id} new_version_id={new_version.id} " + f"version={new_version.version}" + ) + new_version_public = ConfigVersionPublic.model_validate(new_version) + result = {"success": True, "version": new_version.version} + + except (Timeout, SoftTimeLimitExceeded): + logger.warning( + f"[execute_prompt_improvement] Timeout | job_id={job_id} " + f"evaluation_id={evaluation_id}" + ) + error_message = "Task exceeded soft time limit" + job_crud.update( + job_uuid, + JobUpdate(status=JobStatus.FAILED, error_message=error_message), + ) + _send_improve_prompt_callback( + callback_url=callback_url, + payload=_build_improve_prompt_payload( + job_id=job_uuid, + config_version=None, + error_message=error_message, + ), + project_id=project_id, + organization_id=organization_id, + ) + raise + except Exception as exc: + # HTTPException (from re-validation) carries the clean detail; everything + # else falls back to str(exc). + error_message = getattr(exc, "detail", None) or str(exc) + logger.error( + f"[execute_prompt_improvement] Failed | job_id={job_id} " + f"evaluation_id={evaluation_id} | {error_message}", + exc_info=True, + ) + job_crud.update( + job_uuid, + JobUpdate(status=JobStatus.FAILED, error_message=error_message), + ) + _send_improve_prompt_callback( + callback_url=callback_url, + payload=_build_improve_prompt_payload( + job_id=job_uuid, + config_version=None, + error_message=error_message, + ), + project_id=project_id, + organization_id=organization_id, + ) + raise + + # Reached only on success (both except blocks re-raise). The callback + # lives outside the try so a delivery-side hiccup (e.g. get_webhook_secret + # hitting a transient DB error) can't flip the committed SUCCESS to FAILED + # and defeat the redelivery idempotency guard — that would double-spend the + # LLM call and mint a duplicate config_version. + _send_improve_prompt_callback( + callback_url=callback_url, + payload=_build_improve_prompt_payload( + job_id=job_uuid, + config_version=new_version_public, + error_message=None, + ), + project_id=project_id, + organization_id=organization_id, + ) + return result def _draft_improved_prompt( @@ -185,18 +495,15 @@ def _draft_improved_prompt( (improved_instructions, rationale). Uses structured outputs so the first text block is guaranteed-valid JSON. - Raises HTTPException on any LLM failure, mapping the cause to an HTTP status: - 500 (our misconfig/bug), 502 (bad upstream response), 503 (rate limited), - 504 (upstream timeout). + Runs inside a Celery worker, so failures raise plain exceptions (no client + sees an HTTP status here) — the worker turns them into Job FAILED. The + per-cause log lines are kept so operators can still triage the fault. """ if not settings.ANTHROPIC_API_KEY: # Missing platform key is a server-side misconfiguration, not an upstream fault. - raise HTTPException( - status_code=500, - detail=( - "prompt_generation_failed: the platform Anthropic key " - "(ANTHROPIC_API_KEY) is not configured" - ), + raise RuntimeError( + "prompt_generation_failed: the platform Anthropic key " + "(ANTHROPIC_API_KEY) is not configured" ) client = ClaudeProvider.create_client({"api_key": settings.ANTHROPIC_API_KEY}) @@ -245,12 +552,9 @@ def _draft_improved_prompt( "valid, not expired, and configured correctly.", exc_info=True, ) - raise HTTPException( - status_code=502, - detail=( - "prompt_generation_failed: Anthropic authentication failed — " - "verify the platform API key is valid and not expired" - ), + raise RuntimeError( + "prompt_generation_failed: Anthropic authentication failed — " + "verify the platform API key is valid and not expired" ) except anthropic.RateLimitError: @@ -259,12 +563,9 @@ def _draft_improved_prompt( "(code: 429): Hit Anthropic rate/quota — wait ≥1 min and retry.", exc_info=True, ) - raise HTTPException( - status_code=503, - detail=( - "prompt_generation_failed: Anthropic rate limit exceeded — " - "wait at least 1 minute and retry" - ), + raise RuntimeError( + "prompt_generation_failed: Anthropic rate limit exceeded — " + "wait at least 1 minute and retry" ) except anthropic.APITimeoutError: @@ -274,12 +575,9 @@ def _draft_improved_prompt( "(code: APITimeoutError): retry with a smaller payload.", exc_info=True, ) - raise HTTPException( - status_code=504, - detail=( - "prompt_generation_failed: Anthropic request timed out — " - "retry. If persistent, contact Kaapi" - ), + raise RuntimeError( + "prompt_generation_failed: Anthropic request timed out — " + "retry. If persistent, contact Kaapi" ) except anthropic.APIConnectionError: @@ -288,12 +586,9 @@ def _draft_improved_prompt( "(code: APIConnectionError): network or DNS issue reaching Anthropic.", exc_info=True, ) - raise HTTPException( - status_code=502, - detail=( - "prompt_generation_failed: network error reaching Anthropic — " - "check connectivity. If persistent, contact Kaapi" - ), + raise RuntimeError( + "prompt_generation_failed: network error reaching Anthropic — " + "check connectivity. If persistent, contact Kaapi" ) except anthropic.APIStatusError as exc: @@ -305,12 +600,9 @@ def _draft_improved_prompt( f"(code: {status}): {exc.message}.", exc_info=True, ) - raise HTTPException( - status_code=502, - detail=( - f"prompt_generation_failed: Anthropic returned HTTP {status} — " - "retry or contact Kaapi if persistent" - ), + raise RuntimeError( + f"prompt_generation_failed: Anthropic returned HTTP {status} — " + "retry or contact Kaapi if persistent" ) except Exception as exc: @@ -320,10 +612,7 @@ def _draft_improved_prompt( f"likely a Kaapi-side failure. Contact Kaapi if persistent.", exc_info=True, ) - raise HTTPException( - status_code=500, - detail=( - "prompt_generation_failed: unexpected error during prompt generation — " - "contact Kaapi if persistent" - ), + raise RuntimeError( + "prompt_generation_failed: unexpected error during prompt generation — " + "contact Kaapi if persistent" ) diff --git a/backend/app/tests/api/routes/test_improve_prompt.py b/backend/app/tests/api/routes/test_improve_prompt.py index 8b501d749..37fe469e0 100644 --- a/backend/app/tests/api/routes/test_improve_prompt.py +++ b/backend/app/tests/api/routes/test_improve_prompt.py @@ -1,62 +1,72 @@ -"""Tests for POST /evaluations/{evaluation_id}/improve-prompt. +"""Tests for the async prompt-improvement feature (callback-URL / webhook flow). -Covers the redesigned prompt-improvement behavior: no request body, traces -loaded from object storage via load_json_from_object_store, a single -client.messages.create call with structured outputs, and a ConfigVersionPublic -response. +The endpoint is: + - POST /evaluations/{id}/improve-prompt → 202 + a job handle + (LLMJobImmediatePublic). Requires a JSON body {"callback_url": "https://..."}. + The result is delivered to callback_url once the worker finishes — there is + no GET poll route anymore. + - execute_prompt_improvement → the Celery worker entrypoint that does the + Anthropic round-trip, mints the new config_version, and fires send_callback + on every exit (SUCCESS, timeout, generic failure, redelivery no-op). HTTP boundaries mocked (patched as bound in the service module): -- app.services.evaluations.prompt_improvement.ClaudeProvider.create_client - (returns a fake Anthropic client whose messages.create is stubbed) -- app.services.evaluations.prompt_improvement.get_cloud_storage -- app.services.evaluations.prompt_improvement.load_json_from_object_store - (returns already-parsed trace data, or None for the download-failed path) +- ClaudeProvider.create_client (fake Anthropic client) OR _draft_improved_prompt +- get_cloud_storage / load_json_from_object_store (traces) +- start_prompt_improvement (the Celery enqueue helper) — never touch a broker +- send_callback + get_webhook_secret — never make real outbound HTTP +- validate_callback_url at the route call site — real DNS is avoided for the + happy/domain paths; the SSRF-rejection tests use the real validator with + literal hosts that resolve without network. +- Session (the worker opens its own Session(engine); we redirect it at the + transactional db fixture, matching the doctransformer worker tests) DB is real (transactional db fixture; rolls back after each test). """ import contextlib import json +from contextlib import ExitStack +from dataclasses import dataclass from typing import Any, Iterator from unittest.mock import MagicMock, patch +from uuid import uuid4 -import anthropic import pytest -from fastapi.testclient import TestClient +from celery.exceptions import SoftTimeLimitExceeded +from fastapi import HTTPException +from gevent import Timeout from sqlmodel import Session, select from app.core.config import settings from app.crud.config.version import ConfigVersionCrud -from app.models import ConfigVersion, EvaluationDataset, EvaluationRun +from app.crud.jobs import JobCrud +from app.models import EvaluationDataset, EvaluationRun from app.models.config.config import ConfigTag +from app.models.job import Job, JobStatus, JobType from app.services.evaluations.prompt_improvement import ( AI_GENERATED_MARKER, COMMIT_MESSAGE_MAX_LENGTH, + execute_prompt_improvement, + start_prompt_improvement_job, + validate_improve_prompt, ) from app.tests.utils.auth import TestAuthContext from app.tests.utils.test_data import create_test_evaluation_dataset from app.tests.utils.utils import random_lower_string -IMPROVE_URL = "/api/v1/evaluations/{evaluation_id}/improve-prompt" +_SERVICE = "app.services.evaluations.prompt_improvement" +_ROUTE_VALIDATE = "app.api.routes.evaluations.evaluation.validate_callback_url" +POST_URL = "/api/v1/evaluations/{evaluation_id}/improve-prompt" +# The deleted GET poll route — any method here must now miss all routes. +OLD_POLL_URL = "/api/v1/evaluations/{evaluation_id}/improve-prompt/{job_id}" + +_CALLBACK_URL = "https://example.com/callback" _IMPROVED_INSTRUCTIONS = "You are an improved assistant. Answer precisely." _RATIONALE = "Tightened answer scoping to address weak categories." - -def _llm_json( - improved_instructions: str = _IMPROVED_INSTRUCTIONS, - rationale: str = _RATIONALE, -) -> str: - return json.dumps( - { - "improved_instructions": improved_instructions, - "rationale": rationale, - } - ) - - -# Already-parsed trace data — the new service path returns parsed JSON, not bytes. +# Already-parsed trace data — the service loads parsed JSON, not bytes. _TRACES: list[dict] = [ { "trace_id": "t1", @@ -68,12 +78,19 @@ def _llm_json( } ] -# Sentinel: "generate a default score_trace_url from the run id after commit" -_AUTO_URL = object() +_AUTO_URL = object() # "generate a valid score_trace_url from the run id after commit" + + +def _llm_json( + improved_instructions: str = _IMPROVED_INSTRUCTIONS, + rationale: str = _RATIONALE, +) -> str: + return json.dumps( + {"improved_instructions": improved_instructions, "rationale": rationale} + ) def _make_fake_claude_client(text_content: str | None = None) -> MagicMock: - """Return a fake Anthropic client whose messages.create yields a text block.""" if text_content is None: text_content = _llm_json() @@ -90,35 +107,62 @@ def _make_fake_claude_client(text_content: str | None = None) -> MagicMock: return client +@dataclass +class WorkerEnv: + """Handles yielded by ``_worker_env``: the LLM boundary and the callback sink.""" + + llm: MagicMock + callback: MagicMock + + @contextlib.contextmanager -def _patched_boundaries( +def _worker_env( + db: Session, *, + draft: MagicMock | None = None, claude_client: MagicMock | None = None, traces: Any = _TRACES, -) -> Iterator[MagicMock]: - """Patch the three HTTP boundaries and yield the fake Claude client. - - - ClaudeProvider.create_client → returns ``claude_client`` - - get_cloud_storage → returns a harmless MagicMock (load_json is patched, so - the storage object is never really exercised) - - load_json_from_object_store → returns ``traces`` (use None to exercise the - trace_download_failed path) +) -> Iterator[WorkerEnv]: + """Redirect the worker's Session(engine) at the test db and mock its boundaries. + + ``env.llm`` is the draft mock when ``draft`` is given (LLM step stubbed + wholesale), otherwise the fake Claude client (real _draft_improved_prompt + exercised). ``env.callback`` is the patched send_callback — no real HTTP ever. """ - if claude_client is None: - claude_client = _make_fake_claude_client() - - base = "app.services.evaluations.prompt_improvement" - with patch( - f"{base}.ClaudeProvider.create_client", - return_value=claude_client, - ), patch( - f"{base}.get_cloud_storage", - return_value=MagicMock(), - ), patch( - f"{base}.load_json_from_object_store", - return_value=traces, - ): - yield claude_client + session_cm = MagicMock() + session_cm.__enter__.return_value = db + session_cm.__exit__.return_value = None + + with ExitStack() as stack: + stack.enter_context(patch(f"{_SERVICE}.Session", return_value=session_cm)) + stack.enter_context( + patch(f"{_SERVICE}.get_cloud_storage", return_value=MagicMock()) + ) + stack.enter_context( + patch(f"{_SERVICE}.load_json_from_object_store", return_value=traces) + ) + stack.enter_context(patch(f"{_SERVICE}.get_webhook_secret", return_value=None)) + callback = stack.enter_context(patch(f"{_SERVICE}.send_callback")) + + if draft is not None: + stack.enter_context(patch(f"{_SERVICE}._draft_improved_prompt", draft)) + yield WorkerEnv(llm=draft, callback=callback) + else: + if claude_client is None: + claude_client = _make_fake_claude_client() + stack.enter_context( + patch( + f"{_SERVICE}.ClaudeProvider.create_client", + return_value=claude_client, + ) + ) + yield WorkerEnv(llm=claude_client, callback=callback) + + +def _callback_payload(callback: MagicMock) -> dict: + """Single send_callback invocation → the APIResponse envelope it POSTed.""" + assert callback.call_count == 1 + return callback.call_args.args[1] def _make_config_with_instructions( @@ -126,7 +170,6 @@ def _make_config_with_instructions( project_id: int, instructions: str = "You are a helpful assistant.", ) -> Any: - """Create a config whose config_blob has completion.params.instructions set.""" from app.crud.config import ConfigCrud from app.models.config.config import ConfigCreate from app.models.llm import KaapiCompletionConfig @@ -151,15 +194,16 @@ def _make_config_with_instructions( commit_message="Initial version", tag=ConfigTag.DEFAULT, ) - config_crud = ConfigCrud(session=db, project_id=project_id) - config, _ = config_crud.create_or_raise(config_create) + config, _ = ConfigCrud(session=db, project_id=project_id).create_or_raise( + config_create + ) return config def _make_completed_run( db: Session, config_id: Any, - config_version: int, + config_version: int | None, organization_id: int, project_id: int, dataset_id: int, @@ -167,14 +211,6 @@ def _make_completed_run( run_name: str | None = None, score_trace_url: Any = _AUTO_URL, ) -> EvaluationRun: - """Create and persist an EvaluationRun. - - score_trace_url behaviour: - - omitted (default _AUTO_URL): a valid S3 URL is generated after commit so the - run id is known. Use for all tests that mock the storage layer. - - None or "": stored verbatim (use to exercise the traces_not_available path). - - any str: stored verbatim. - """ if run_name is None: run_name = f"run-{random_lower_string()}" @@ -238,11 +274,8 @@ def config_with_instructions(db: Session, auth: TestAuthContext) -> Any: @pytest.fixture def anthropic_creds(monkeypatch: pytest.MonkeyPatch) -> None: - """Configure the platform-owned Anthropic key the service reads from settings.""" monkeypatch.setattr( - settings, - "ANTHROPIC_API_KEY", - "sk-ant-test-" + random_lower_string(), + settings, "ANTHROPIC_API_KEY", "sk-ant-test-" + random_lower_string() ) @@ -264,202 +297,28 @@ def completed_run( ) -class TestHappyPath: - """Completed run with traces → 201 with correct ConfigVersionPublic.""" - - def test_returns_201_with_new_config_version( - self, - client: TestClient, - headers: dict[str, str], - config_with_instructions: Any, - completed_run: EvaluationRun, - ) -> None: - """POST with no body → 201; response contains new config version at latest+1.""" - with _patched_boundaries(): - resp = client.post( - IMPROVE_URL.format(evaluation_id=completed_run.id), - headers=headers, - ) - - assert resp.status_code == 201, resp.text - body = resp.json()["data"] - assert body["version"] == 2 # initial version was 1 - assert body["config_id"] == str(config_with_instructions.id) - - def test_new_version_has_improved_instructions_in_db( - self, - client: TestClient, - headers: dict[str, str], - db: Session, - completed_run: EvaluationRun, - ) -> None: - """The new config version persisted in DB carries the improved instructions.""" - with _patched_boundaries(): - resp = client.post( - IMPROVE_URL.format(evaluation_id=completed_run.id), - headers=headers, - ) - - assert resp.status_code == 201, resp.text - new_version_id = resp.json()["data"]["id"] - stmt = select(ConfigVersion).where(ConfigVersion.id == new_version_id) - new_version = db.exec(stmt).one() - assert ( - new_version.config_blob["completion"]["params"]["instructions"] - == _IMPROVED_INSTRUCTIONS - ) - - def test_commit_message_marker_and_rationale( - self, - client: TestClient, - headers: dict[str, str], - completed_run: EvaluationRun, - ) -> None: - """commit_message starts with the AI marker and carries run name + rationale.""" - with _patched_boundaries(): - resp = client.post( - IMPROVE_URL.format(evaluation_id=completed_run.id), - headers=headers, - ) - - assert resp.status_code == 201, resp.text - commit_message = resp.json()["data"]["commit_message"] - assert commit_message.startswith(AI_GENERATED_MARKER) - assert f"(Evaluation: {completed_run.run_name})" in commit_message - assert _RATIONALE in commit_message - assert len(commit_message) <= COMMIT_MESSAGE_MAX_LENGTH - - def test_long_rationale_truncated_to_max_length( - self, - client: TestClient, - headers: dict[str, str], - completed_run: EvaluationRun, - ) -> None: - """A rationale longer than the cap forces real truncation at the boundary.""" - long_rationale = "x" * (COMMIT_MESSAGE_MAX_LENGTH + 200) - fake_client = _make_fake_claude_client(_llm_json(rationale=long_rationale)) - - with _patched_boundaries(claude_client=fake_client): - resp = client.post( - IMPROVE_URL.format(evaluation_id=completed_run.id), - headers=headers, - ) - - assert resp.status_code == 201, resp.text - commit_message = resp.json()["data"]["commit_message"] - assert len(commit_message) == COMMIT_MESSAGE_MAX_LENGTH - assert commit_message.startswith(AI_GENERATED_MARKER) - - def test_messages_create_uses_json_schema_output_config( - self, - client: TestClient, - headers: dict[str, str], - completed_run: EvaluationRun, - ) -> None: - """The single messages.create call requests structured json_schema output.""" - with _patched_boundaries() as fake_client: - resp = client.post( - IMPROVE_URL.format(evaluation_id=completed_run.id), - headers=headers, - ) - - assert resp.status_code == 201, resp.text +class TestValidateImprovePrompt: + """Fast DB precondition checks — no LLM, no trace download.""" - fake_client.messages.create.assert_called_once() - # The legacy Files API must be gone — no beta.* surface is touched. - fake_client.beta.files.upload.assert_not_called() - - output_config = fake_client.messages.create.call_args.kwargs["output_config"] - assert output_config["format"]["type"] == "json_schema" - assert "schema" in output_config["format"] - - def test_commit_message_has_no_metric_or_threshold( - self, - client: TestClient, - headers: dict[str, str], - completed_run: EvaluationRun, - ) -> None: - """The new design's provenance string must not carry metric=/threshold=.""" - with _patched_boundaries(): - resp = client.post( - IMPROVE_URL.format(evaluation_id=completed_run.id), - headers=headers, + def test_missing_run_raises_404(self, db: Session, auth: TestAuthContext) -> None: + with pytest.raises(HTTPException) as exc: + validate_improve_prompt( + session=db, + evaluation_id=9999999, + organization_id=auth.organization_id, + project_id=auth.project_id, ) - - assert resp.status_code == 201, resp.text - commit_message = resp.json()["data"]["commit_message"] - assert "metric=" not in commit_message - assert "threshold=" not in commit_message - - -class TestRunNotFound: - """404 when the evaluation run doesn't exist or is non-TEXT type.""" - - def test_nonexistent_run_returns_404( - self, - client: TestClient, - headers: dict[str, str], - anthropic_creds: None, - ) -> None: - resp = client.post( - IMPROVE_URL.format(evaluation_id=9999999), - headers=headers, - ) - - assert resp.status_code == 404, resp.text - assert "evaluation_not_found" in resp.json()["error"], resp.text - - def test_non_text_run_returns_404( - self, - client: TestClient, - headers: dict[str, str], - db: Session, - auth: TestAuthContext, - dataset: EvaluationDataset, - config_with_instructions: Any, - anthropic_creds: None, - ) -> None: - """get_evaluation_run_by_id filters type=='text'; a non-text run is invisible.""" - run = EvaluationRun( - run_name=f"run-stt-{random_lower_string()}", - dataset_name=f"ds-{random_lower_string()}", - dataset_id=dataset.id, - config_id=config_with_instructions.id, - config_version=1, - status="completed", - total_items=1, - type="stt", - organization_id=auth.organization_id, - project_id=auth.project_id, - score_trace_url="s3://test-bucket/traces.json", - ) - db.add(run) - db.commit() - db.refresh(run) - - resp = client.post( - IMPROVE_URL.format(evaluation_id=run.id), - headers=headers, - ) - - assert resp.status_code == 404, resp.text - assert "evaluation_not_found" in resp.json()["error"], resp.text - - -class TestNonCompletedStatus: - """409 when evaluation run is not in 'completed' status.""" + assert exc.value.status_code == 404 + assert "evaluation_not_found" in exc.value.detail @pytest.mark.parametrize("status", ["pending", "processing", "failed"]) - def test_non_completed_returns_409( + def test_non_completed_raises_409( self, status: str, - client: TestClient, - headers: dict[str, str], db: Session, auth: TestAuthContext, dataset: EvaluationDataset, config_with_instructions: Any, - anthropic_creds: None, ) -> None: run = _make_completed_run( db=db, @@ -470,30 +329,24 @@ def test_non_completed_returns_409( dataset_id=dataset.id, status=status, ) - - resp = client.post( - IMPROVE_URL.format(evaluation_id=run.id), - headers=headers, - ) - - assert resp.status_code == 409, resp.text - assert "evaluation_not_completed" in resp.json()["error"], resp.text - - -class TestTracesNotAvailable: - """422 when score_trace_url is empty/None (precondition unmet).""" + with pytest.raises(HTTPException) as exc: + validate_improve_prompt( + session=db, + evaluation_id=run.id, + organization_id=auth.organization_id, + project_id=auth.project_id, + ) + assert exc.value.status_code == 409 + assert "evaluation_not_completed" in exc.value.detail @pytest.mark.parametrize("trace_url", [None, ""]) - def test_missing_trace_url_returns_422( + def test_missing_trace_url_raises_422( self, trace_url: str | None, - client: TestClient, - headers: dict[str, str], db: Session, auth: TestAuthContext, dataset: EvaluationDataset, config_with_instructions: Any, - anthropic_creds: None, ) -> None: run = _make_completed_run( db=db, @@ -504,356 +357,485 @@ def test_missing_trace_url_returns_422( dataset_id=dataset.id, score_trace_url=trace_url, ) + with pytest.raises(HTTPException) as exc: + validate_improve_prompt( + session=db, + evaluation_id=run.id, + organization_id=auth.organization_id, + project_id=auth.project_id, + ) + assert exc.value.status_code == 422 + assert "traces_not_available" in exc.value.detail - resp = client.post( - IMPROVE_URL.format(evaluation_id=run.id), - headers=headers, - ) - - assert resp.status_code == 422, resp.text - assert "traces_not_available" in resp.json()["error"], resp.text - - -class TestSourceConfigUnavailable: - """409 when the source config or config_version is missing/soft-deleted.""" - - def test_soft_deleted_config_returns_409( + def test_missing_config_reference_raises_422( self, - client: TestClient, - headers: dict[str, str], db: Session, auth: TestAuthContext, dataset: EvaluationDataset, - anthropic_creds: None, ) -> None: - from app.core.util import now - from app.models.config.config import Config - - config = _make_config_with_instructions(db=db, project_id=auth.project_id) run = _make_completed_run( db=db, - config_id=config.id, - config_version=1, + config_id=None, + config_version=None, organization_id=auth.organization_id, project_id=auth.project_id, dataset_id=dataset.id, ) + with pytest.raises(HTTPException) as exc: + validate_improve_prompt( + session=db, + evaluation_id=run.id, + organization_id=auth.organization_id, + project_id=auth.project_id, + ) + assert exc.value.status_code == 422 + assert "source_config_unavailable" in exc.value.detail - stmt = select(Config).where(Config.id == config.id) - cfg = db.exec(stmt).one() - cfg.deleted_at = now() - db.add(cfg) - db.commit() - - resp = client.post( - IMPROVE_URL.format(evaluation_id=run.id), - headers=headers, - ) - - assert resp.status_code == 409, resp.text - assert "source_config_unavailable" in resp.json()["error"], resp.text - - def test_missing_config_version_returns_409( + def test_unresolvable_config_version_raises_409( self, - client: TestClient, - headers: dict[str, str], db: Session, auth: TestAuthContext, dataset: EvaluationDataset, - anthropic_creds: None, + config_with_instructions: Any, ) -> None: - config = _make_config_with_instructions(db=db, project_id=auth.project_id) run = _make_completed_run( db=db, - config_id=config.id, - config_version=99, # does not exist + config_id=config_with_instructions.id, + config_version=99, # version does not exist organization_id=auth.organization_id, project_id=auth.project_id, dataset_id=dataset.id, ) - - resp = client.post( - IMPROVE_URL.format(evaluation_id=run.id), - headers=headers, - ) - - assert resp.status_code == 409, resp.text - assert "source_config_unavailable" in resp.json()["error"], resp.text + with pytest.raises(HTTPException) as exc: + validate_improve_prompt( + session=db, + evaluation_id=run.id, + organization_id=auth.organization_id, + project_id=auth.project_id, + ) + assert exc.value.status_code == 409 + assert "source_config_unavailable" in exc.value.detail -class TestMissingAnthropicKey: - """500 when the platform ANTHROPIC_API_KEY is not configured (server misconfig).""" +class TestStartJobRoute: + """POST → 202 + job handle; enqueue behavior and fast-fail without enqueue.""" - def test_empty_api_key_returns_500( + def test_returns_202_with_job_handle_and_persists_job( self, - client: TestClient, + client: Any, headers: dict[str, str], + db: Session, + auth: TestAuthContext, completed_run: EvaluationRun, - monkeypatch: pytest.MonkeyPatch, ) -> None: - """Storage succeeds but the key check inside the LLM step fails with 500.""" - monkeypatch.setattr(settings, "ANTHROPIC_API_KEY", "") - - with _patched_boundaries(): + with patch(_ROUTE_VALIDATE), patch( + f"{_SERVICE}.start_prompt_improvement", return_value="task-1" + ) as enqueue: resp = client.post( - IMPROVE_URL.format(evaluation_id=completed_run.id), + POST_URL.format(evaluation_id=completed_run.id), headers=headers, + json={"callback_url": _CALLBACK_URL}, ) - assert resp.status_code == 500, resp.text - assert "prompt_generation_failed" in resp.json()["error"], resp.text - - -class TestTraceDownloadFailed: - """502 when the trace file cannot be loaded from object storage.""" - - def test_load_json_returns_none_returns_502( + assert resp.status_code == 202, resp.text + body = resp.json()["data"] + assert body["status"] == JobStatus.PENDING.value + job_id = body["job_id"] + + job = JobCrud(session=db).get(job_id=job_id, project_id=auth.project_id) + assert job is not None + assert job.job_type == JobType.PROMPT_IMPROVEMENT + + enqueue.assert_called_once() + kwargs = enqueue.call_args.kwargs + assert kwargs["project_id"] == auth.project_id + assert kwargs["job_id"] == job_id + assert kwargs["organization_id"] == auth.organization_id + assert kwargs["evaluation_id"] == completed_run.id + assert kwargs["callback_url"] == _CALLBACK_URL + + def test_missing_callback_url_returns_422( self, - client: TestClient, + client: Any, headers: dict[str, str], completed_run: EvaluationRun, ) -> None: - with _patched_boundaries(traces=None): + with patch(f"{_SERVICE}.start_prompt_improvement") as enqueue: resp = client.post( - IMPROVE_URL.format(evaluation_id=completed_run.id), + POST_URL.format(evaluation_id=completed_run.id), headers=headers, + json={}, ) - - assert resp.status_code == 502, resp.text - assert "trace_download_failed" in resp.json()["error"], resp.text - - -def _anthropic_exceptions() -> list[tuple[str, Exception, int]]: - """Build one instance of each SDK error plus the HTTP status the service maps it to. - - The mapping reflects who's at fault: rate limit -> 503 (upstream unavailable), - timeout -> 504 (upstream timeout), bad upstream response/network -> 502, and an - unexpected non-SDK error -> 500 (Kaapi-side bug). - """ - return [ - ( - "authentication", - anthropic.AuthenticationError( - message="auth failed", - response=MagicMock(status_code=401, headers={}), - body={}, - ), - 502, - ), - ( - "rate_limit", - anthropic.RateLimitError( - message="rate limited", - response=MagicMock(status_code=429, headers={}), - body={}, - ), - 503, - ), - ("timeout", anthropic.APITimeoutError(request=MagicMock()), 504), - ( - "connection", - anthropic.APIConnectionError(message="conn failed", request=MagicMock()), - 502, - ), - ( - "status", - anthropic.APIStatusError( - message="server error", - response=MagicMock(status_code=500, headers={}), - body={}, - ), - 502, - ), - ("generic", RuntimeError("something unexpected"), 500), - ] - - -class TestLLMFailureMapping: - """Each Anthropic/SDK failure from messages.create maps to its matching HTTP status.""" + assert resp.status_code == 422, resp.text + enqueue.assert_not_called() @pytest.mark.parametrize( - "name,exc,expected_status", - _anthropic_exceptions(), - ids=[name for name, _, _ in _anthropic_exceptions()], + "callback_url", + [ + "http://example.com/callback", # non-HTTPS scheme + "https://10.0.0.1/callback", # RFC 1918 private IP + "https://127.0.0.1/callback", # loopback + ], ) - def test_llm_error_maps_to_status( + def test_invalid_callback_url_rejected_by_ssrf_guard( self, - name: str, - exc: Exception, - expected_status: int, - client: TestClient, + callback_url: str, + client: Any, headers: dict[str, str], completed_run: EvaluationRun, ) -> None: - fake_client = _make_fake_claude_client() - fake_client.messages.create.side_effect = exc - - with _patched_boundaries(claude_client=fake_client): + # Real validate_callback_url runs; literal hosts resolve without DNS. + with patch(f"{_SERVICE}.start_prompt_improvement") as enqueue: resp = client.post( - IMPROVE_URL.format(evaluation_id=completed_run.id), + POST_URL.format(evaluation_id=completed_run.id), headers=headers, + json={"callback_url": callback_url}, ) + assert resp.status_code == 422, resp.text + assert "invalid_callback_url" in resp.json()["error"], resp.text + enqueue.assert_not_called() - assert resp.status_code == expected_status, resp.text - assert "prompt_generation_failed" in resp.json()["error"], resp.text + def test_enqueue_failure_marks_job_failed_and_raises_500( + self, + db: Session, + auth: TestAuthContext, + completed_run: EvaluationRun, + ) -> None: + with patch( + f"{_SERVICE}.start_prompt_improvement", + side_effect=RuntimeError("broker down"), + ): + with pytest.raises(HTTPException) as exc: + start_prompt_improvement_job( + session=db, + evaluation_id=completed_run.id, + organization_id=auth.organization_id, + project_id=auth.project_id, + callback_url=_CALLBACK_URL, + ) + + assert exc.value.status_code == 500 + + job = db.exec( + select(Job) + .where(Job.project_id == auth.project_id) + .where(Job.job_type == JobType.PROMPT_IMPROVEMENT) + ).one() + assert job.status == JobStatus.FAILED + assert job.error_message is not None - def test_malformed_json_response_returns_500( + @pytest.mark.parametrize( + "make_run_kwargs,expected_status,error_token", + [ + (None, 404, "evaluation_not_found"), + ({"status": "pending"}, 409, "evaluation_not_completed"), + ({"score_trace_url": None}, 422, "traces_not_available"), + ], + ) + def test_fast_validation_error_does_not_enqueue( self, - client: TestClient, + make_run_kwargs: dict | None, + expected_status: int, + error_token: str, + client: Any, headers: dict[str, str], - completed_run: EvaluationRun, + db: Session, + auth: TestAuthContext, + dataset: EvaluationDataset, + config_with_instructions: Any, + anthropic_creds: None, ) -> None: - """Non-JSON text from the model falls through to the generic 500 branch.""" - fake_client = _make_fake_claude_client("this is not json") + if make_run_kwargs is None: + evaluation_id = 9999999 + else: + run = _make_completed_run( + db=db, + config_id=config_with_instructions.id, + config_version=1, + organization_id=auth.organization_id, + project_id=auth.project_id, + dataset_id=dataset.id, + **make_run_kwargs, + ) + evaluation_id = run.id - with _patched_boundaries(claude_client=fake_client): + with patch(_ROUTE_VALIDATE), patch( + f"{_SERVICE}.start_prompt_improvement" + ) as enqueue: resp = client.post( - IMPROVE_URL.format(evaluation_id=completed_run.id), + POST_URL.format(evaluation_id=evaluation_id), headers=headers, + json={"callback_url": _CALLBACK_URL}, ) - assert resp.status_code == 500, resp.text - assert "prompt_generation_failed" in resp.json()["error"], resp.text + assert resp.status_code == expected_status, resp.text + assert error_token in resp.json()["error"], resp.text + enqueue.assert_not_called() + no_job = db.exec( + select(Job) + .where(Job.project_id == auth.project_id) + .where(Job.job_type == JobType.PROMPT_IMPROVEMENT) + ).first() + assert no_job is None -class TestTenantIsolation: - """404 when a run belongs to a different org/project.""" - def test_run_from_different_project_returns_404( +class TestExecuteWorker: + """Worker entrypoint: PROCESSING → SUCCESS/FAILED, callback delivery, redelivery.""" + + def _pending_job(self, db: Session, project_id: int) -> Job: + return JobCrud(session=db).create( + job_type=JobType.PROMPT_IMPROVEMENT, + trace_id="N/A", + project_id=project_id, + ) + + def test_happy_path_mints_new_version_and_records_meta( self, - client: TestClient, - headers: dict[str, str], db: Session, - superuser_api_key: TestAuthContext, - anthropic_creds: None, + auth: TestAuthContext, + completed_run: EvaluationRun, ) -> None: - su_dataset = create_test_evaluation_dataset( - db=db, - organization_id=superuser_api_key.organization_id, - project_id=superuser_api_key.project_id, - ) - su_config = _make_config_with_instructions( - db=db, - project_id=superuser_api_key.project_id, - ) - su_run = _make_completed_run( - db=db, - config_id=su_config.id, - config_version=1, - organization_id=superuser_api_key.organization_id, - project_id=superuser_api_key.project_id, - dataset_id=su_dataset.id, - ) + job = self._pending_job(db, auth.project_id) + + with _worker_env(db) as env: + result = execute_prompt_improvement( + project_id=auth.project_id, + job_id=str(job.id), + organization_id=auth.organization_id, + evaluation_id=completed_run.id, + task_id="celery-task-1", + callback_url=_CALLBACK_URL, + ) - resp = client.post( - IMPROVE_URL.format(evaluation_id=su_run.id), - headers=headers, # normal user headers, not superuser - ) + assert result["success"] is True - assert resp.status_code == 404, resp.text - assert "evaluation_not_found" in resp.json()["error"], resp.text + refreshed = db.get(Job, job.id) + assert refreshed.status == JobStatus.SUCCESS + assert refreshed.task_id == "celery-task-1" + assert refreshed.meta["version"] == 2 + assert refreshed.meta["rationale"] == _RATIONALE + crud = ConfigVersionCrud( + session=db, config_id=completed_run.config_id, project_id=auth.project_id + ) + v2 = crud.read_one(version_number=2) + assert v2 is not None + assert ( + v2.config_blob["completion"]["params"]["instructions"] + == _IMPROVED_INSTRUCTIONS + ) + assert v2.commit_message.startswith(AI_GENERATED_MARKER) + assert completed_run.run_name in v2.commit_message + assert _RATIONALE in v2.commit_message -class TestRepeatableIteration: - """Running improvement twice creates a version at latest+1 each time.""" + # structured-output request shape is what makes the JSON parse reliable + output_config = env.llm.messages.create.call_args.kwargs["output_config"] + assert output_config["format"]["type"] == "json_schema" - def test_second_improvement_creates_next_version( + def test_success_fires_callback_with_config_version( self, - client: TestClient, - headers: dict[str, str], db: Session, auth: TestAuthContext, - dataset: EvaluationDataset, - config_with_instructions: Any, - anthropic_creds: None, + completed_run: EvaluationRun, ) -> None: - config_id = config_with_instructions.id - - run1 = _make_completed_run( - db=db, - config_id=config_id, - config_version=1, - organization_id=auth.organization_id, - project_id=auth.project_id, - dataset_id=dataset.id, - ) - - with _patched_boundaries(): - resp1 = client.post( - IMPROVE_URL.format(evaluation_id=run1.id), - headers=headers, + job = self._pending_job(db, auth.project_id) + + with _worker_env(db) as env: + execute_prompt_improvement( + project_id=auth.project_id, + job_id=str(job.id), + organization_id=auth.organization_id, + evaluation_id=completed_run.id, + callback_url=_CALLBACK_URL, ) - assert resp1.status_code == 201, resp1.text - assert resp1.json()["data"]["version"] == 2 - run2 = _make_completed_run( - db=db, - config_id=config_id, - config_version=2, - organization_id=auth.organization_id, - project_id=auth.project_id, - dataset_id=dataset.id, - ) + assert env.callback.call_args.args[0] == _CALLBACK_URL + payload = _callback_payload(env.callback) + assert payload["success"] is True + assert payload["data"]["status"] == JobStatus.SUCCESS.value + assert payload["data"]["error_message"] is None + assert payload["data"]["config_version"] is not None + assert payload["data"]["config_version"]["version"] == 2 - with _patched_boundaries(): - resp2 = client.post( - IMPROVE_URL.format(evaluation_id=run2.id), - headers=headers, + def test_long_rationale_truncated_in_commit_message( + self, + db: Session, + auth: TestAuthContext, + completed_run: EvaluationRun, + ) -> None: + job = self._pending_job(db, auth.project_id) + long_rationale = "x" * (COMMIT_MESSAGE_MAX_LENGTH + 200) + draft = MagicMock(return_value=(_IMPROVED_INSTRUCTIONS, long_rationale)) + + with _worker_env(db, draft=draft): + execute_prompt_improvement( + project_id=auth.project_id, + job_id=str(job.id), + organization_id=auth.organization_id, + evaluation_id=completed_run.id, + callback_url=_CALLBACK_URL, ) - assert resp2.status_code == 201, resp2.text - assert resp2.json()["data"]["version"] == 3 + v2 = ConfigVersionCrud( + session=db, config_id=completed_run.config_id, project_id=auth.project_id + ).read_one(version_number=2) + assert len(v2.commit_message) == COMMIT_MESSAGE_MAX_LENGTH + assert v2.commit_message.startswith(AI_GENERATED_MARKER) -class TestPriorVersionsPreserved: - """Pre-existing config_version rows are unchanged after a new improvement.""" + def test_llm_runtime_error_marks_failed_and_fires_failure_callback( + self, + db: Session, + auth: TestAuthContext, + completed_run: EvaluationRun, + ) -> None: + job = self._pending_job(db, auth.project_id) + draft = MagicMock(side_effect=RuntimeError("prompt_generation_failed: boom")) + + with _worker_env(db, draft=draft) as env: + with pytest.raises(RuntimeError): + execute_prompt_improvement( + project_id=auth.project_id, + job_id=str(job.id), + organization_id=auth.organization_id, + evaluation_id=completed_run.id, + callback_url=_CALLBACK_URL, + ) + + refreshed = db.get(Job, job.id) + assert refreshed.status == JobStatus.FAILED + assert "prompt_generation_failed" in refreshed.error_message + + payload = _callback_payload(env.callback) + assert payload["success"] is False + assert payload["data"]["status"] == JobStatus.FAILED.value + assert payload["data"]["config_version"] is None + assert payload["data"]["error_message"] + assert "prompt_generation_failed" in payload["data"]["error_message"] + + # no version minted on failure + v2 = ConfigVersionCrud( + session=db, config_id=completed_run.config_id, project_id=auth.project_id + ).read_one(version_number=2) + assert v2 is None + + def test_trace_download_failure_fires_failure_callback( + self, + db: Session, + auth: TestAuthContext, + completed_run: EvaluationRun, + ) -> None: + job = self._pending_job(db, auth.project_id) + + with _worker_env(db, traces=None) as env: + with pytest.raises(RuntimeError): + execute_prompt_improvement( + project_id=auth.project_id, + job_id=str(job.id), + organization_id=auth.organization_id, + evaluation_id=completed_run.id, + callback_url=_CALLBACK_URL, + ) + + refreshed = db.get(Job, job.id) + assert refreshed.status == JobStatus.FAILED + + payload = _callback_payload(env.callback) + assert payload["success"] is False + assert payload["data"]["status"] == JobStatus.FAILED.value + assert "trace_download_failed" in payload["data"]["error_message"] - def test_prior_versions_still_retrievable( + @pytest.mark.parametrize( + "exc", [Timeout(1), SoftTimeLimitExceeded()], ids=["gevent", "celery"] + ) + def test_timeout_marks_failed_and_fires_failure_callback( self, - client: TestClient, - headers: dict[str, str], + exc: Exception, db: Session, auth: TestAuthContext, - dataset: EvaluationDataset, - config_with_instructions: Any, - anthropic_creds: None, + completed_run: EvaluationRun, ) -> None: - config_id = config_with_instructions.id + job = self._pending_job(db, auth.project_id) + draft = MagicMock(side_effect=exc) + + with _worker_env(db, draft=draft) as env: + with pytest.raises(type(exc)): + execute_prompt_improvement( + project_id=auth.project_id, + job_id=str(job.id), + organization_id=auth.organization_id, + evaluation_id=completed_run.id, + callback_url=_CALLBACK_URL, + ) + + refreshed = db.get(Job, job.id) + assert refreshed.status == JobStatus.FAILED + assert refreshed.error_message == "Task exceeded soft time limit" + + payload = _callback_payload(env.callback) + assert payload["success"] is False + assert payload["data"]["status"] == JobStatus.FAILED.value + assert payload["data"]["error_message"] == "Task exceeded soft time limit" + + def test_redelivery_of_success_job_is_noop_and_resends_callback( + self, + db: Session, + auth: TestAuthContext, + completed_run: EvaluationRun, + ) -> None: + job = self._pending_job(db, auth.project_id) + + # First run mints v2 and lands SUCCESS. + with _worker_env(db) as env: + execute_prompt_improvement( + project_id=auth.project_id, + job_id=str(job.id), + organization_id=auth.organization_id, + evaluation_id=completed_run.id, + callback_url=_CALLBACK_URL, + ) + assert env.llm.messages.create.call_count == 1 + crud = ConfigVersionCrud( - session=db, config_id=config_id, project_id=auth.project_id + session=db, config_id=completed_run.config_id, project_id=auth.project_id ) + assert crud.read_one(version_number=3) is None + first_meta = db.get(Job, job.id).meta + + # Redelivery: SUCCESS job must not re-call the LLM or mint a duplicate, + # but at-least-once delivery re-sends the success callback. + draft = MagicMock(side_effect=AssertionError("LLM must not be re-called")) + with _worker_env(db, draft=draft) as env: + result = execute_prompt_improvement( + project_id=auth.project_id, + job_id=str(job.id), + organization_id=auth.organization_id, + evaluation_id=completed_run.id, + callback_url=_CALLBACK_URL, + ) - original_version = crud.read_one(version_number=1) - assert original_version is not None - original_instructions = original_version.config_blob["completion"]["params"][ - "instructions" - ] + draft.assert_not_called() + assert result["success"] is True + assert result["version"] == first_meta["version"] + assert crud.read_one(version_number=3) is None - run = _make_completed_run( - db=db, - config_id=config_id, - config_version=1, - organization_id=auth.organization_id, - project_id=auth.project_id, - dataset_id=dataset.id, - ) + payload = _callback_payload(env.callback) + assert payload["success"] is True + assert payload["data"]["status"] == JobStatus.SUCCESS.value + assert payload["data"]["config_version"]["version"] == first_meta["version"] - with _patched_boundaries(): - resp = client.post( - IMPROVE_URL.format(evaluation_id=run.id), - headers=headers, - ) - assert resp.status_code == 201, resp.text +class TestPollRouteRemoved: + """The GET poll endpoint was removed with the switch to callbacks.""" - db.expire_all() - still_v1 = crud.read_one(version_number=1) - assert still_v1 is not None - assert ( - still_v1.config_blob["completion"]["params"]["instructions"] - == original_instructions + def test_old_poll_path_no_longer_routed( + self, + client: Any, + headers: dict[str, str], + completed_run: EvaluationRun, + ) -> None: + resp = client.get( + OLD_POLL_URL.format(evaluation_id=completed_run.id, job_id=uuid4()), + headers=headers, ) - - v2 = crud.read_one(version_number=2) - assert v2 is not None - assert v2.commit_message is not None - assert v2.commit_message.startswith(AI_GENERATED_MARKER) + assert resp.status_code == 404, resp.text diff --git a/docs/wiki/domain-map.md b/docs/wiki/domain-map.md index 240cb3646..e9cacdaed 100644 --- a/docs/wiki/domain-map.md +++ b/docs/wiki/domain-map.md @@ -27,7 +27,7 @@ APIKey → Organization, Project, User # programmatic access | ConfigVersion | config/version.py | Config | resolved by `LLMCallConfig` saved references (logical) | | LlmCall | llm/request.py | Job, LlmChain, Org, Project | Langfuse traces (logical); analytics | | LlmChain | llm/request.py | Org, Project | LlmCall | -| Job | job.py | Project | LlmCall; Celery job execution (logical) | +| Job | job.py | Project | LlmCall; Celery job execution (logical); evaluation prompt improvement (`JobType.PROMPT_IMPROVEMENT`, logical) | | BatchJob | batch_job.py | Org, Project | EvaluationRun, Assessment; batch polling cron (logical) | | EvaluationDataset | evaluation.py | Org, Project, Language | EvaluationRun, STTSample (via stt_evaluation), Assessment | | EvaluationRun | evaluation.py | Dataset, Config, BatchJob, Org, Project, Language | STTResult, TTSResult; Langfuse scores (logical); console UI (logical) | diff --git a/docs/wiki/modules/evaluations.md b/docs/wiki/modules/evaluations.md index ed81b6099..7f296c6e9 100644 --- a/docs/wiki/modules/evaluations.md +++ b/docs/wiki/modules/evaluations.md @@ -28,6 +28,7 @@ Key `EvaluationRun` JSONB fields: `score` (per-trace `scores` + `summary_scores` ## Async - Provider batches polled by cron (`crud/evaluations/cron.py`); no long-lived Celery task per run. +- Prompt improvement is job-based with callback delivery: `POST /evaluations/{id}/improve-prompt` validates preconditions, enqueues a `Job` (`JobType.PROMPT_IMPROVEMENT`, `models/job.py`) run by Celery task `run_prompt_improvement` (`celery/tasks/job_execution.py`), and returns `202` with an `LLMJobImmediatePublic` handle. On finish the worker POSTs a single best-effort callback to the caller-supplied `callback_url` (SSRF-guarded via `validate_callback_url`): an `APIResponse[PromptImprovementJobPublic]` (`models/evaluation.py`) carrying the new `ConfigVersion` on success or `error_message` on failure. The `ConfigVersion` is persisted regardless of callback outcome. Celery redelivery of a `SUCCESS` job re-sends the callback without re-running the LLM. ## External - OpenAI Batch + embeddings, Gemini Batch, Langfuse (per-trace scores + summary), object storage (datasets/audio), kaapi-frontend console (results + annotation).