144 lines
4.8 KiB
Python
144 lines
4.8 KiB
Python
"""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:
|
|
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()
|