From bb743dfd55f26aca5d4ece6b7a5d31c59393284b Mon Sep 17 00:00:00 2001 From: scosman Date: Mon, 7 Oct 2024 10:30:43 -0400 Subject: [PATCH] Add task APIs --- libs/studio/kiln_studio/server.py | 6 +- libs/studio/kiln_studio/task_management.py | 35 ++++++ .../kiln_studio/test_task_management.py | 114 ++++++++++++++++++ 3 files changed, 154 insertions(+), 1 deletion(-) create mode 100644 libs/studio/kiln_studio/task_management.py create mode 100644 libs/studio/kiln_studio/test_task_management.py diff --git a/libs/studio/kiln_studio/server.py b/libs/studio/kiln_studio/server.py index 7e90264..7f452c0 100644 --- a/libs/studio/kiln_studio/server.py +++ b/libs/studio/kiln_studio/server.py @@ -8,6 +8,7 @@ from .project_management import connect_project_management from .provider_management import connect_provider_management from .settings import connect_settings +from .task_management import connect_task_management from .webhost import connect_webhost @@ -29,10 +30,13 @@ def ping(): connect_project_management(app) connect_provider_management(app) + connect_task_management(app) connect_settings(app) - connect_webhost(app) connect_custom_errors(app) + # Important: webhost must be last, it handles all other URLs + connect_webhost(app) + return app diff --git a/libs/studio/kiln_studio/task_management.py b/libs/studio/kiln_studio/task_management.py new file mode 100644 index 0000000..be9795f --- /dev/null +++ b/libs/studio/kiln_studio/task_management.py @@ -0,0 +1,35 @@ +import os +from pathlib import Path + +from fastapi import FastAPI +from fastapi.responses import JSONResponse + +from libs.core.kiln_ai.datamodel import Project, Task + + +def connect_task_management(app: FastAPI): + @app.post("/api/task") + async def create_task(task: Task, project_path: str): + if not os.path.exists(project_path): + return JSONResponse( + status_code=400, + content={ + "message": "Parent project not found. Can't create task.", + }, + ) + + try: + parent_project = Project.load_from_file(Path(project_path)) + task.parent = parent_project + except Exception as e: + return JSONResponse( + status_code=500, + content={ + "message": f"Failed to load parent project: {e}", + }, + ) + + task.save_to_file() + returnTask = task.model_dump() + returnTask["path"] = task.path + return returnTask diff --git a/libs/studio/kiln_studio/test_task_management.py b/libs/studio/kiln_studio/test_task_management.py new file mode 100644 index 0000000..31d2627 --- /dev/null +++ b/libs/studio/kiln_studio/test_task_management.py @@ -0,0 +1,114 @@ +import json +from pathlib import Path +from unittest.mock import patch + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from libs.core.kiln_ai.datamodel import Project, Task +from libs.studio.kiln_studio.task_management import connect_task_management + + +@pytest.fixture +def app(): + app = FastAPI() + connect_task_management(app) + return app + + +@pytest.fixture +def client(app): + return TestClient(app) + + +def test_create_task_success(client, tmp_path): + project_path = tmp_path / "test_project" + project_path.mkdir() + + task_data = { + "name": "Test Task", + "description": "This is a test task", + } + + with patch( + "libs.core.kiln_ai.datamodel.Project.load_from_file" + ) as mock_load, patch("libs.core.kiln_ai.datamodel.Task.save_to_file") as mock_save: + mock_load.return_value = Project(name="Test Project") + mock_save.return_value = None + + response = client.post(f"/api/task?project_path={project_path}", json=task_data) + + assert response.status_code == 200 + res = response.json() + assert res["name"] == "Test Task" + assert res["description"] == "This is a test task" + assert "path" in res + assert res["id"] is not None + assert res["priority"] == 2 + + +def test_create_task_project_not_found(client, tmp_path): + non_existent_path = tmp_path / "non_existent" + + task_data = { + "name": "Test Task", + "description": "This is a test task", + } + + response = client.post( + f"/api/task?project_path={non_existent_path}", json=task_data + ) + + assert response.status_code == 400 + assert response.json()["message"] == "Parent project not found. Can't create task." + + +def test_create_task_project_load_error(client, tmp_path): + project_path = tmp_path / "test_project" + project_path.mkdir() + + task_data = { + "name": "Test Task", + "description": "This is a test task", + } + + with patch("libs.core.kiln_ai.datamodel.Project.load_from_file") as mock_load: + mock_load.side_effect = Exception("Failed to load project") + + response = client.post(f"/api/task?project_path={project_path}", json=task_data) + + assert response.status_code == 500 + assert "Failed to load parent project" in response.json()["message"] + + +def test_create_task_real_project(client, tmp_path): + project_path = tmp_path / "real_project" / Project.base_filename() + project_path.parent.mkdir() + + # Create a real Project + project = Project(name="Real Project", path=str(project_path)) + project.save_to_file() + + task_data = { + "name": "Real Task", + "description": "This is a real task", + } + + response = client.post(f"/api/task?project_path={project.path}", json=task_data) + + assert response.status_code == 200 + res = response.json() + assert res["name"] == "Real Task" + assert res["description"] == "This is a real task" + assert "path" in res + assert res["id"] is not None + assert res["priority"] == 2 + + # Verify the task file on disk + task_from_disk = Task.load_from_file(Path(res["path"])) + + assert task_from_disk.name == "Real Task" + assert task_from_disk.description == "This is a real task" + assert task_from_disk.id == res["id"] + assert task_from_disk.priority == 2