"""Task execution engine: runs tasks via AI Workers. MVP supports two execution modes: 1. single_call - one LLM call 2. agent_loop - multi-step reasoning (simplified: multiple calls with self-reflection) """ import json from datetime import datetime from typing import Optional from sqlalchemy.orm import Session from app.models.task import Task from app.models.worker import AIWorker from app.models.artifact import Artifact from app.models.review import Review from app.services.llm_service import LLMGateway from app.services.worker_service import WorkerService from app.services.task_service import TaskService from app.services.cost_service import CostService from app.core.exceptions import ( TaskNotExecutableError, WorkerOfflineError, BudgetExceededError, ) from app.schemas.task import TaskExecutionResult class ExecutionService: """Orchestrates task execution through AI Workers.""" @staticmethod def execute( db: Session, tenant_id: int, task_id: int, override_input: Optional[str] = None, ) -> TaskExecutionResult: task = TaskService.get_task(db, tenant_id, task_id) if not task: raise TaskNotExecutableError(f"Task #{task_id} not found") if task.status in ("done", "cancelled"): raise TaskNotExecutableError(f"Task is already {task.status}") # Get worker if not task.worker_id: raise TaskNotExecutableError("No AI Worker assigned to this task") worker = WorkerService.get_worker(db, tenant_id, task.worker_id) if not worker or not worker.is_active: raise WorkerOfflineError(task.worker_id) if worker.current_task_count >= worker.max_concurrent_tasks: raise WorkerOfflineError(f"{worker.name} is at capacity") # Budget check CostService.check_budget(db, tenant_id, task, worker) # Mark task in progress task.status = "in_progress" task.started_at = datetime.utcnow() worker.current_task_count += 1 db.commit() # Parse input input_str = override_input or task.input_data or "" try: input_data = json.loads(input_str) if input_str else {} except json.JSONDecodeError: input_data = {"prompt": input_str} prompt = input_data.get("prompt", task.description or task.title) context = input_data.get("context", "") # Build messages messages = [{"role": "system", "content": worker.system_prompt}] if context: messages.append({"role": "user", "content": f"Context:\n{context}"}) messages.append({"role": "user", "content": prompt}) try: # Branch by task type if task.task_type == "agent_loop": return ExecutionService._execute_agent_loop( db, tenant_id, task, worker, messages ) # Single-call mode output, cost_log = LLMGateway.call( worker=worker, messages=messages, db=db, tenant_id=tenant_id, project_id=task.project_id, task_id=task.id, ) # Save output task.output_data = json.dumps( {"content": output, "cost_log_id": cost_log.id}, ensure_ascii=False, ) # Create artifact artifact = Artifact( tenant_id=tenant_id, task_id=task.id, project_id=task.project_id, title=task.title, content=output, artifact_type="text", version=1, created_by_worker_id=worker.id, ) db.add(artifact) db.flush() # Create review gate if needed if task.requires_review: review = Review( tenant_id=tenant_id, task_id=task.id, project_id=task.project_id, review_type="approval", status="pending", review_content=output, submitted_at=datetime.utcnow(), ) db.add(review) task.status = "review" task.review_status = "pending" result_status = "needs_review" else: task.status = "done" task.completed_at = datetime.utcnow() result_status = "success" db.commit() return TaskExecutionResult( task_id=task.id, status=result_status, output=output, token_usage=cost_log.total_tokens, cost_cents=cost_log.cost_cents, duration_ms=cost_log.duration_ms, artifact_id=artifact.id, ) except Exception as e: task.status = "pending" # reset to allow retry db.commit() raise finally: worker.current_task_count = max(0, worker.current_task_count - 1) db.commit() @staticmethod def _execute_agent_loop(db, tenant_id, task, worker, messages): """Agent loop: multi-step reasoning with up to 5 iterations.""" max_iterations = 5 all_messages = list(messages) final_output = "" total_cost_cents = 0 total_tokens = 0 total_duration = 0 for i in range(max_iterations): output, cost_log = LLMGateway.call( worker=worker, messages=all_messages, db=db, tenant_id=tenant_id, project_id=task.project_id, task_id=task.id, ) total_cost_cents += cost_log.cost_cents total_tokens += cost_log.total_tokens total_duration += cost_log.duration_ms # Check if agent wants to continue if output.strip().startswith("NEEDS_TOOL:"): tool_request = output.replace("NEEDS_TOOL:", "").strip() tool_result = f"[Tool simulation] Processed: {tool_request}" all_messages.append({"role": "assistant", "content": output}) all_messages.append({"role": "user", "content": f"Tool result: {tool_result}\nPlease continue."}) final_output = output continue else: final_output = output break # Save output task.output_data = json.dumps( {"content": final_output, "iterations": i + 1, "total_cost": total_cost_cents}, ensure_ascii=False, ) artifact = Artifact( tenant_id=tenant_id, task_id=task.id, project_id=task.project_id, title=task.title, content=final_output, artifact_type="text", version=1, created_by_worker_id=worker.id, ) db.add(artifact) db.flush() if task.requires_review: review = Review( tenant_id=tenant_id, task_id=task.id, project_id=task.project_id, review_type="approval", status="pending", review_content=final_output, submitted_at=datetime.utcnow(), ) db.add(review) task.status = "review" task.review_status = "pending" result_status = "needs_review" else: task.status = "done" task.completed_at = datetime.utcnow() result_status = "success" db.commit() return TaskExecutionResult( task_id=task.id, status=result_status, output=final_output, token_usage=total_tokens, cost_cents=total_cost_cents, duration_ms=total_duration, artifact_id=artifact.id, )