59a59f8869
Major architectural improvement to replace in-memory task storage with database-backed persistence while maintaining backward compatibility. Changes: - Created Task SQLAlchemy model matching database schema - Added Workflow and Execution SQLAlchemy models - Created TaskService for database CRUD operations - Updated UnifiedCoordinator to use database persistence - Modified task APIs to leverage database storage - Added task loading from database on coordinator initialization - Implemented status change persistence during task execution - Enhanced task cleanup with database support - Added comprehensive task statistics from database Benefits: - Tasks persist across application restarts - Better scalability and reliability - Historical task data retention - Comprehensive task filtering and querying - Maintains in-memory cache for performance 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
220 lines
9.6 KiB
Python
220 lines
9.6 KiB
Python
"""
|
|
Task service for database operations
|
|
Handles CRUD operations for tasks and integrates with the UnifiedCoordinator
|
|
"""
|
|
|
|
from typing import List, Optional, Dict, Any
|
|
from sqlalchemy.orm import Session
|
|
from sqlalchemy import desc, func
|
|
from datetime import datetime, timedelta
|
|
import uuid
|
|
|
|
from ..models.task import Task as ORMTask
|
|
from ..core.unified_coordinator import Task as CoordinatorTask, TaskStatus, AgentType
|
|
from ..core.database import SessionLocal
|
|
|
|
|
|
class TaskService:
|
|
"""Service for managing task persistence and database operations"""
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
def create_task(self, coordinator_task: CoordinatorTask) -> ORMTask:
|
|
"""Create a task in the database from a coordinator task"""
|
|
with SessionLocal() as db:
|
|
try:
|
|
# Convert coordinator task to database task
|
|
db_task = ORMTask(
|
|
id=uuid.UUID(coordinator_task.id) if isinstance(coordinator_task.id, str) else coordinator_task.id,
|
|
title=coordinator_task.context.get('title', f"Task {coordinator_task.type.value}"),
|
|
description=coordinator_task.context.get('description', ''),
|
|
priority=coordinator_task.priority,
|
|
status=coordinator_task.status.value,
|
|
assigned_agent_id=coordinator_task.assigned_agent,
|
|
workflow_id=uuid.UUID(coordinator_task.workflow_id) if coordinator_task.workflow_id else None,
|
|
metadata={
|
|
'type': coordinator_task.type.value,
|
|
'context': coordinator_task.context,
|
|
'payload': coordinator_task.payload,
|
|
'dependencies': coordinator_task.dependencies,
|
|
'created_at': coordinator_task.created_at,
|
|
'completed_at': coordinator_task.completed_at,
|
|
'result': coordinator_task.result
|
|
}
|
|
)
|
|
|
|
if coordinator_task.status == TaskStatus.IN_PROGRESS and coordinator_task.created_at:
|
|
db_task.started_at = datetime.fromtimestamp(coordinator_task.created_at)
|
|
|
|
if coordinator_task.status == TaskStatus.COMPLETED and coordinator_task.completed_at:
|
|
db_task.completed_at = datetime.fromtimestamp(coordinator_task.completed_at)
|
|
|
|
db.add(db_task)
|
|
db.commit()
|
|
db.refresh(db_task)
|
|
|
|
return db_task
|
|
|
|
except Exception as e:
|
|
db.rollback()
|
|
raise e
|
|
|
|
def update_task(self, task_id: str, coordinator_task: CoordinatorTask) -> Optional[ORMTask]:
|
|
"""Update a task in the database"""
|
|
with SessionLocal() as db:
|
|
try:
|
|
# Convert string ID to UUID if needed
|
|
uuid_id = uuid.UUID(task_id) if isinstance(task_id, str) else task_id
|
|
|
|
db_task = db.query(ORMTask).filter(ORMTask.id == uuid_id).first()
|
|
if not db_task:
|
|
return None
|
|
|
|
# Update fields from coordinator task
|
|
db_task.title = coordinator_task.context.get('title', db_task.title)
|
|
db_task.description = coordinator_task.context.get('description', db_task.description)
|
|
db_task.priority = coordinator_task.priority
|
|
db_task.status = coordinator_task.status.value
|
|
db_task.assigned_agent_id = coordinator_task.assigned_agent
|
|
|
|
# Update metadata
|
|
db_task.metadata = {
|
|
'type': coordinator_task.type.value,
|
|
'context': coordinator_task.context,
|
|
'payload': coordinator_task.payload,
|
|
'dependencies': coordinator_task.dependencies,
|
|
'created_at': coordinator_task.created_at,
|
|
'completed_at': coordinator_task.completed_at,
|
|
'result': coordinator_task.result
|
|
}
|
|
|
|
# Update timestamps based on status
|
|
if coordinator_task.status == TaskStatus.IN_PROGRESS and not db_task.started_at:
|
|
db_task.started_at = datetime.utcnow()
|
|
|
|
if coordinator_task.status == TaskStatus.COMPLETED and not db_task.completed_at:
|
|
db_task.completed_at = datetime.utcnow()
|
|
|
|
db.commit()
|
|
db.refresh(db_task)
|
|
|
|
return db_task
|
|
|
|
except Exception as e:
|
|
db.rollback()
|
|
raise e
|
|
|
|
def get_task(self, task_id: str) -> Optional[ORMTask]:
|
|
"""Get a task by ID"""
|
|
with SessionLocal() as db:
|
|
uuid_id = uuid.UUID(task_id) if isinstance(task_id, str) else task_id
|
|
return db.query(ORMTask).filter(ORMTask.id == uuid_id).first()
|
|
|
|
def get_tasks(self, status: Optional[str] = None, agent_id: Optional[str] = None,
|
|
workflow_id: Optional[str] = None, limit: int = 100) -> List[ORMTask]:
|
|
"""Get tasks with optional filtering"""
|
|
with SessionLocal() as db:
|
|
query = db.query(ORMTask)
|
|
|
|
if status:
|
|
query = query.filter(ORMTask.status == status)
|
|
if agent_id:
|
|
query = query.filter(ORMTask.assigned_agent_id == agent_id)
|
|
if workflow_id:
|
|
uuid_workflow_id = uuid.UUID(workflow_id) if isinstance(workflow_id, str) else workflow_id
|
|
query = query.filter(ORMTask.workflow_id == uuid_workflow_id)
|
|
|
|
return query.order_by(desc(ORMTask.created_at)).limit(limit).all()
|
|
|
|
def get_pending_tasks(self, limit: int = 50) -> List[ORMTask]:
|
|
"""Get pending tasks ordered by priority"""
|
|
with SessionLocal() as db:
|
|
return db.query(ORMTask).filter(
|
|
ORMTask.status == 'pending'
|
|
).order_by(
|
|
ORMTask.priority.asc(), # Lower number = higher priority
|
|
ORMTask.created_at.asc()
|
|
).limit(limit).all()
|
|
|
|
def delete_task(self, task_id: str) -> bool:
|
|
"""Delete a task"""
|
|
with SessionLocal() as db:
|
|
try:
|
|
uuid_id = uuid.UUID(task_id) if isinstance(task_id, str) else task_id
|
|
task = db.query(ORMTask).filter(ORMTask.id == uuid_id).first()
|
|
if task:
|
|
db.delete(task)
|
|
db.commit()
|
|
return True
|
|
return False
|
|
except Exception as e:
|
|
db.rollback()
|
|
raise e
|
|
|
|
def cleanup_completed_tasks(self, max_age_hours: int = 24) -> int:
|
|
"""Clean up old completed tasks"""
|
|
with SessionLocal() as db:
|
|
try:
|
|
cutoff_time = datetime.utcnow() - timedelta(hours=max_age_hours)
|
|
|
|
deleted_count = db.query(ORMTask).filter(
|
|
ORMTask.status.in_(['completed', 'failed']),
|
|
ORMTask.completed_at < cutoff_time
|
|
).delete(synchronize_session=False)
|
|
|
|
db.commit()
|
|
return deleted_count
|
|
|
|
except Exception as e:
|
|
db.rollback()
|
|
raise e
|
|
|
|
def coordinator_task_from_orm(self, orm_task: ORMTask) -> CoordinatorTask:
|
|
"""Convert ORM task back to coordinator task"""
|
|
metadata = orm_task.metadata or {}
|
|
|
|
# Extract fields from metadata
|
|
task_type = AgentType(metadata.get('type', 'general_ai'))
|
|
context = metadata.get('context', {})
|
|
payload = metadata.get('payload', {})
|
|
dependencies = metadata.get('dependencies', [])
|
|
result = metadata.get('result')
|
|
created_at = metadata.get('created_at', orm_task.created_at.timestamp() if orm_task.created_at else None)
|
|
completed_at = metadata.get('completed_at')
|
|
|
|
# Convert status
|
|
status = TaskStatus(orm_task.status) if orm_task.status in [s.value for s in TaskStatus] else TaskStatus.PENDING
|
|
|
|
return CoordinatorTask(
|
|
id=str(orm_task.id),
|
|
type=task_type,
|
|
priority=orm_task.priority,
|
|
status=status,
|
|
context=context,
|
|
payload=payload,
|
|
assigned_agent=orm_task.assigned_agent_id,
|
|
result=result,
|
|
created_at=created_at,
|
|
completed_at=completed_at,
|
|
workflow_id=str(orm_task.workflow_id) if orm_task.workflow_id else None,
|
|
dependencies=dependencies
|
|
)
|
|
|
|
def get_task_statistics(self) -> Dict[str, Any]:
|
|
"""Get task statistics"""
|
|
with SessionLocal() as db:
|
|
total_tasks = db.query(ORMTask).count()
|
|
pending_tasks = db.query(ORMTask).filter(ORMTask.status == 'pending').count()
|
|
in_progress_tasks = db.query(ORMTask).filter(ORMTask.status == 'in_progress').count()
|
|
completed_tasks = db.query(ORMTask).filter(ORMTask.status == 'completed').count()
|
|
failed_tasks = db.query(ORMTask).filter(ORMTask.status == 'failed').count()
|
|
|
|
return {
|
|
'total_tasks': total_tasks,
|
|
'pending_tasks': pending_tasks,
|
|
'in_progress_tasks': in_progress_tasks,
|
|
'completed_tasks': completed_tasks,
|
|
'failed_tasks': failed_tasks,
|
|
'success_rate': completed_tasks / total_tasks if total_tasks > 0 else 0
|
|
} |