Skip to content

refactor: prune ZIP for prompt optimization job #1031

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
leonardmq merged 7 commits into from
Feb 14, 2026
Merged

refactor: prune ZIP for prompt optimization job #1031

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
8175df0
refactor: filter out unneeded artefacts from the project ZIP for gepa
leonardmq Feb 14, 2026
2e3d690
Merge branch 'leonard/kil-348-gepa-billing-payment-credit-2' of githu…
leonardmq Feb 14, 2026
bc2c663
refactor: remove token budget from client side (now remote server con…
leonardmq Feb 14, 2026
a8d45a9
fix: cr review, non-interactive variant
leonardmq Feb 14, 2026
6a57b52
fix: cr review (switch to temp dir instead of temp file, for Windows …
leonardmq Feb 14, 2026
3151293
Merge branch 'leonard/kil-403-gepa-exclude-docs-from-zip' of github.c…
leonardmq Feb 14, 2026
c04ad39
Merge pull request #1032 from Kiln-AI/leonard/kil-402-gepa-remove-lig…
leonardmq Feb 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 0 additions & 4 deletions app/desktop/studio_server/api_client/kiln_ai_server_client/models/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,6 @@
from .answer_option_with_selection import AnswerOptionWithSelection
from .api_key_verification_result import ApiKeyVerificationResult
from .body_start_gepa_job_v1_jobs_gepa_job_start_post import BodyStartGepaJobV1JobsGepaJobStartPost
from .body_start_gepa_job_v1_jobs_gepa_job_start_post_token_budget import (
BodyStartGepaJobV1JobsGepaJobStartPostTokenBudget,
)
from .body_start_sample_job_v1_jobs_sample_job_start_post import BodyStartSampleJobV1JobsSampleJobStartPost
from .check_entitlements_v1_check_entitlements_get_response_check_entitlements_v1_check_entitlements_get import (
CheckEntitlementsV1CheckEntitlementsGetResponseCheckEntitlementsV1CheckEntitlementsGet,
Expand Down Expand Up @@ -62,7 +59,6 @@
"AnswerOptionWithSelection",
"ApiKeyVerificationResult",
"BodyStartGepaJobV1JobsGepaJobStartPost",
"BodyStartGepaJobV1JobsGepaJobStartPostTokenBudget",
"BodyStartSampleJobV1JobsSampleJobStartPost",
"CheckEntitlementsV1CheckEntitlementsGetResponseCheckEntitlementsV1CheckEntitlementsGet",
"CheckModelSupportedResponse",
Expand Down
13 changes: 0 additions & 13 deletions ...pi_client/kiln_ai_server_client/models/body_start_gepa_job_v1_jobs_gepa_job_start_post.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,9 +8,6 @@
from attrs import field as _attrs_field

from .. import types
from ..models.body_start_gepa_job_v1_jobs_gepa_job_start_post_token_budget import (
BodyStartGepaJobV1JobsGepaJobStartPostTokenBudget,
)
from ..types import File

T = TypeVar("T", bound="BodyStartGepaJobV1JobsGepaJobStartPost")
Expand All @@ -20,23 +17,19 @@
class BodyStartGepaJobV1JobsGepaJobStartPost:
"""
Attributes:
token_budget (BodyStartGepaJobV1JobsGepaJobStartPostTokenBudget): The token budget to use
task_id (str): The task ID
target_run_config_id (str): The target run config ID
eval_ids (list[str]): The list of eval IDs to use for optimization
project_zip (File): The project zip file
"""

token_budget: BodyStartGepaJobV1JobsGepaJobStartPostTokenBudget
task_id: str
target_run_config_id: str
eval_ids: list[str]
project_zip: File
additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)

def to_dict(self) -> dict[str, Any]:
token_budget = self.token_budget.value

task_id = self.task_id

target_run_config_id = self.target_run_config_id
Expand All @@ -49,7 +42,6 @@ def to_dict(self) -> dict[str, Any]:
field_dict.update(self.additional_properties)
field_dict.update(
{
"token_budget": token_budget,
"task_id": task_id,
"target_run_config_id": target_run_config_id,
"eval_ids": eval_ids,
Expand All @@ -62,8 +54,6 @@ def to_dict(self) -> dict[str, Any]:
def to_multipart(self) -> types.RequestFiles:
files: types.RequestFiles = []

files.append(("token_budget", (None, str(self.token_budget.value).encode(), "text/plain")))

files.append(("task_id", (None, str(self.task_id).encode(), "text/plain")))

files.append(("target_run_config_id", (None, str(self.target_run_config_id).encode(), "text/plain")))
Expand All @@ -81,8 +71,6 @@ def to_multipart(self) -> types.RequestFiles:
@classmethod
def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
d = dict(src_dict)
token_budget = BodyStartGepaJobV1JobsGepaJobStartPostTokenBudget(d.pop("token_budget"))

task_id = d.pop("task_id")

target_run_config_id = d.pop("target_run_config_id")
Expand All @@ -92,7 +80,6 @@ def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
project_zip = File(payload=BytesIO(d.pop("project_zip")))

body_start_gepa_job_v1_jobs_gepa_job_start_post = cls(
token_budget=token_budget,
task_id=task_id,
target_run_config_id=target_run_config_id,
eval_ids=eval_ids,
Expand Down
10 changes: 0 additions & 10 deletions ...n_ai_server_client/models/body_start_gepa_job_v1_jobs_gepa_job_start_post_token_budget.py
View file
Open in desktop

This file was deleted.

121 changes: 43 additions & 78 deletions app/desktop/studio_server/gepa_job_api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,8 @@
import asyncio
import io
import logging
import zipfile
import tempfile
from pathlib import Path
from typing import Literal, cast

from app.desktop.studio_server.api_client.kiln_ai_server_client.api.jobs import (
check_model_supported_v1_jobs_gepa_job_check_model_supported_get,
Expand All @@ -17,9 +16,6 @@
from app.desktop.studio_server.api_client.kiln_ai_server_client.models.body_start_gepa_job_v1_jobs_gepa_job_start_post import (
BodyStartGepaJobV1JobsGepaJobStartPost,
)
from app.desktop.studio_server.api_client.kiln_ai_server_client.models.body_start_gepa_job_v1_jobs_gepa_job_start_post_token_budget import (
BodyStartGepaJobV1JobsGepaJobStartPostTokenBudget,
)
from app.desktop.studio_server.api_client.kiln_ai_server_client.models.http_validation_error import (
HTTPValidationError,
)
Expand All @@ -38,8 +34,13 @@
eval_from_id,
task_run_config_from_id,
)
from app.desktop.studio_server.utils.copilot_utils import check_response_error
from fastapi import FastAPI, HTTPException
from kiln_ai.datamodel import GepaJob, Project, Prompt
from kiln_ai.cli.commands.package_project import (
PackageForTrainingConfig,
package_project_for_training,
)
from kiln_ai.datamodel import GepaJob, Prompt
from kiln_ai.datamodel.task import TaskRunConfig
from kiln_ai.utils.config import Config
from kiln_ai.utils.lock import shared_async_lock_manager
Expand Down Expand Up @@ -101,7 +102,6 @@ def _get_api_key() -> str:


class StartGepaJobRequest(BaseModel):
token_budget: Literal["light", "medium", "heavy"]
target_run_config_id: str
eval_ids: list[str]

Expand Down Expand Up @@ -312,52 +312,6 @@ async def update_gepa_job_and_create_artifacts(
return gepa_job


def zip_project(project: Project) -> bytes:
"""
Create a ZIP file of the entire project directory.
Returns the ZIP file as bytes.
"""
if not project.path:
raise ValueError("Project path is not set")
project_path = Path(project.path).parent

# Skip common directories that shouldn't be included
skip_patterns = {
".git",
"__pycache__",
".pytest_cache",
"node_modules",
".venv",
"venv",
".DS_Store",
".vscode",
".idea",
}

buffer = io.BytesIO()
file_count = 0
with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
for file_path in project_path.rglob("*"):
# Skip if any parent directory matches skip patterns
if any(skip_dir in file_path.parts for skip_dir in skip_patterns):
continue

if file_path.is_file():
arcname = file_path.relative_to(project_path)
try:
zip_file.write(file_path, arcname=arcname)
file_count += 1
except Exception as e:
logger.warning(f"Skipping file {file_path}: {e}")

buffer.seek(0)
zip_bytes = buffer.getvalue()
logger.info(
f"Created project ZIP with {file_count} files, total size: {len(zip_bytes)} bytes"
)
return zip_bytes


def connect_gepa_job_api(app: FastAPI):
@app.get("/api/projects/{project_id}/tasks/{task_id}/gepa_jobs/check_run_config")
async def check_run_config(
Expand Down Expand Up @@ -515,7 +469,8 @@ async def start_gepa_job(
Creates and saves a GepaJob datamodel to track the job.
"""
task = task_from_id(project_id, task_id)
if not task.parent:
project = task.parent_project()
if not project:
raise HTTPException(status_code=404, detail="Project not found")

try:
Expand All @@ -538,49 +493,59 @@ async def start_gepa_job(
status_code=500, detail="Server client not authenticated"
)

# Create ZIP file of the project
project_zip_bytes = zip_project(cast(Project, task.parent))
with tempfile.TemporaryDirectory(prefix="kiln_gepa_") as tmpdir:
tmp_file = Path(tmpdir) / "kiln_gepa_project.zip"
package_project_for_training(
project=project,
task_ids=[task_id],
run_config_id=request.target_run_config_id,
eval_ids=request.eval_ids,
output=tmp_file,
config=PackageForTrainingConfig(
include_documents=False,
exclude_task_runs=False,
exclude_eval_config_runs=True,
),
)
zip_bytes = tmp_file.read_bytes()
logger.info(
f"Created project ZIP, total size: {len(zip_bytes)} bytes and file name: {tmp_file.name}"
)

# Create the File object for the SDK
project_zip_file = File(
payload=io.BytesIO(project_zip_bytes),
file_name="project.zip",
mime_type="application/zip",
)
project_zip_file = File(
payload=io.BytesIO(zip_bytes),
file_name="project.zip",
mime_type="application/zip",
)

# Create the request body
body = BodyStartGepaJobV1JobsGepaJobStartPost(
token_budget=BodyStartGepaJobV1JobsGepaJobStartPostTokenBudget(
request.token_budget
),
task_id=task_id,
target_run_config_id=request.target_run_config_id,
project_zip=project_zip_file,
eval_ids=request.eval_ids,
)

response = await start_gepa_job_v1_jobs_gepa_job_start_post.asyncio(
client=server_client, body=body
)

if isinstance(response, HTTPValidationError):
error_detail = (
str(response.detail)
if hasattr(response, "detail")
else "Validation error"
detailed_response = (
await start_gepa_job_v1_jobs_gepa_job_start_post.asyncio_detailed(
client=server_client, body=body
)
raise HTTPException(status_code=422, detail=error_detail)
)
check_response_error(
detailed_response,
default_detail="Failed to start GEPA job: unexpected error from server",
)

if response is None:
response = detailed_response.parsed
if response is None or isinstance(response, HTTPValidationError):
raise HTTPException(
status_code=500,
detail="Failed to start GEPA job: No response from server",
detail="Failed to start GEPA job: unexpected response from server",
Comment on lines +529 to +543
Copy link
Contributor

@coderabbitai coderabbitai bot Feb 14, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Find check_response_error function definition
rg -n -A 30 'def check_response_error' --type py

Repository: Kiln-AI/Kiln

Length of output: 1720

HTTPValidationError isinstance check is unreachable for 422 responses

check_response_error raises HTTPException immediately on any non-200 status code (including 422). The isinstance(response, HTTPValidationError) check at line 540 cannot be reached when a 422 status is returned, since the exception is raised first. Remove the unreachable check or restructure the error handling to distinguish between different failure modes.

🤖 Prompt for AI Agents 
In `@app/desktop/studio_server/gepa_job_api.py` around lines 529 - 543, The
current flow calls check_response_error(detailed_response, ...) which raises an
HTTPException for any non-200 status (including 422), so the subsequent
isinstance(response, HTTPValidationError) branch in the
start_gepa_job_v1_jobs_gepa_job_start_post handling is unreachable; fix by
handling validation errors before calling check_response_error: inspect
detailed_response.status_code (or detailed_response.parsed) immediately after
awaiting start_gepa_job_v1_jobs_gepa_job_start_post.asyncio_detailed and if
status_code == 422 and isinstance(detailed_response.parsed, HTTPValidationError)
raise a 422 HTTPException with the parsed validation error, otherwise call
check_response_error(detailed_response, ...) and proceed to use response as
before.

)

gepa_job = GepaJob(
name=generate_memorable_name(),
job_id=response.job_id,
token_budget=request.token_budget,
target_run_config_id=request.target_run_config_id,
latest_status=JobStatus.PENDING,
eval_ids=request.eval_ids,
Expand Down
Loading
Loading