74 lines
2.4 KiB
Python
74 lines
2.4 KiB
Python
|
"""AI Worker service: CRUD + status management."""
|
||
|
|
from sqlalchemy.orm import Session
|
||
|
|
from typing import List, Optional
|
||
|
|
from app.models.worker import AIWorker
|
||
|
|
from app.schemas.worker import WorkerCreate, WorkerUpdate
|
||
|
|
|
||
|
|
|
||
|
|
class WorkerService:
|
||
|
|
@staticmethod
|
||
|
|
def list_workers(
|
||
|
|
db: Session, tenant_id: int, skip: int = 0, limit: int = 100
|
||
|
|
) -> List[AIWorker]:
|
||
|
|
return (
|
||
|
|
db.query(AIWorker)
|
||
|
|
.filter(AIWorker.tenant_id == tenant_id)
|
||
|
|
.order_by(AIWorker.created_at.desc())
|
||
|
|
.offset(skip)
|
||
|
|
.limit(limit)
|
||
|
|
.all()
|
||
|
|
)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def get_worker(db: Session, tenant_id: int, worker_id: int) -> Optional[AIWorker]:
|
||
|
|
return (
|
||
|
|
db.query(AIWorker)
|
||
|
|
.filter(AIWorker.id == worker_id, AIWorker.tenant_id == tenant_id)
|
||
|
|
.first()
|
||
|
|
)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def create_worker(db: Session, tenant_id: int, req: WorkerCreate) -> AIWorker:
|
||
|
|
worker = AIWorker(
|
||
|
|
tenant_id=tenant_id,
|
||
|
|
name=req.name,
|
||
|
|
description=req.description,
|
||
|
|
provider=req.provider,
|
||
|
|
model_name=req.model_name,
|
||
|
|
api_key=req.api_key,
|
||
|
|
base_url=req.base_url,
|
||
|
|
temperature=req.temperature,
|
||
|
|
system_prompt=req.system_prompt,
|
||
|
|
tools=req.tools,
|
||
|
|
max_cost_per_task_cents=req.max_cost_per_task_cents,
|
||
|
|
max_cost_per_month_cents=req.max_cost_per_month_cents,
|
||
|
|
available_hours=req.available_hours,
|
||
|
|
max_concurrent_tasks=req.max_concurrent_tasks,
|
||
|
|
status="active",
|
||
|
|
)
|
||
|
|
db.add(worker)
|
||
|
|
db.commit()
|
||
|
|
db.refresh(worker)
|
||
|
|
return worker
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def update_worker(
|
||
|
|
db: Session, tenant_id: int, worker_id: int, req: WorkerUpdate
|
||
|
|
) -> AIWorker:
|
||
|
|
worker = WorkerService.get_worker(db, tenant_id, worker_id)
|
||
|
|
if not worker:
|
||
|
|
raise ValueError(f"Worker #{worker_id} not found")
|
||
|
|
for field, value in req.model_dump(exclude_unset=True).items():
|
||
|
|
setattr(worker, field, value)
|
||
|
|
db.commit()
|
||
|
|
db.refresh(worker)
|
||
|
|
return worker
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def delete_worker(db: Session, tenant_id: int, worker_id: int) -> bool:
|
||
|
|
worker = WorkerService.get_worker(db, tenant_id, worker_id)
|
||
|
|
if not worker:
|
||
|
|
return False
|
||
|
|
db.delete(worker)
|
||
|
|
db.commit()
|
||
|
|
return True
|