233 lines
7.7 KiB
Python
233 lines
7.7 KiB
Python
"""Agent routes.
|
|
|
|
Backend responsibilities are intentionally minimal:
|
|
GET /agents/catalog — static catalog for UI display
|
|
POST /agents/can-create — billing eligibility check
|
|
POST /agents/trigger — trigger a local agent run
|
|
|
|
Agent configuration is owned by the Electron app and is not persisted
|
|
in backend agent-config tables.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import uuid
|
|
from datetime import datetime, timezone
|
|
|
|
from fastapi import APIRouter, Depends, HTTPException, status
|
|
from sqlalchemy import func, select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.api.deps import get_current_user
|
|
from app.billing.tier_manager import FEATURES
|
|
from app.core.agent_runner import is_agent_running, run_local_agent
|
|
from app.core.device_manager import device_manager
|
|
from app.db import get_session
|
|
from app.models import AgentRunLog, LocalAgentConfig
|
|
from app.schemas import (
|
|
AgentCatalogItem,
|
|
AgentCreationCheckRequest,
|
|
AgentCreationCheckResponse,
|
|
AgentRunLogResponse,
|
|
AgentTriggerRequest,
|
|
UserProfile,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
router = APIRouter(prefix="/agents", tags=["agents"])
|
|
|
|
|
|
# ── Datetime helpers ──────────────────────────────────────────────────
|
|
|
|
def _dt_ms(dt: datetime) -> int:
|
|
return int(dt.timestamp() * 1000)
|
|
|
|
|
|
def _dt_ms_opt(dt: datetime | None) -> int | None:
|
|
return int(dt.timestamp() * 1000) if dt else None
|
|
|
|
|
|
def _to_data_types(values: list[str]) -> list[str]:
|
|
normalize = {
|
|
"task": "tasks", "tasks": "tasks",
|
|
"note": "notes", "notes": "notes",
|
|
"timeline": "timelines", "timelines": "timelines", "timelineEvents": "timelines",
|
|
"project": "projects", "projects": "projects",
|
|
}
|
|
seen: set[str] = set()
|
|
result: list[str] = []
|
|
for v in values:
|
|
mapped = normalize.get(v)
|
|
if mapped and mapped not in seen:
|
|
seen.add(mapped)
|
|
result.append(mapped)
|
|
return result
|
|
|
|
|
|
def _to_run_log_response(log: AgentRunLog) -> AgentRunLogResponse:
|
|
return AgentRunLogResponse(
|
|
id=log.id,
|
|
agent_id=log.agent_id,
|
|
agent_type=log.agent_type, # type: ignore[arg-type]
|
|
status=log.status, # type: ignore[arg-type]
|
|
items_processed=log.items_processed,
|
|
items_created=log.items_created,
|
|
errors=log.errors or [],
|
|
started_at=_dt_ms(log.started_at),
|
|
completed_at=_dt_ms_opt(log.completed_at),
|
|
)
|
|
|
|
|
|
def _enforce_agent_limit(tier: str, current_count: int) -> int:
|
|
limit: int = FEATURES.get(tier, FEATURES["free"])["batch_active"]
|
|
if limit != -1 and current_count >= limit:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_403_FORBIDDEN,
|
|
detail=f"Agent limit ({limit}) reached for your tier. Upgrade to create more.",
|
|
)
|
|
return limit
|
|
|
|
|
|
async def _enforce_run_frequency(
|
|
tier: str,
|
|
user_id: str,
|
|
db: AsyncSession,
|
|
) -> None:
|
|
"""Raise HTTP 402 if the user has exceeded their daily batch run limit."""
|
|
limit: int = FEATURES.get(tier, FEATURES["free"])["batch_runs_per_day"]
|
|
if limit == -1:
|
|
return # unlimited
|
|
|
|
today_start = datetime.now(timezone.utc).replace(
|
|
hour=0, minute=0, second=0, microsecond=0
|
|
)
|
|
result = await db.execute(
|
|
select(func.count(AgentRunLog.id)).where(
|
|
AgentRunLog.user_id == user_id,
|
|
AgentRunLog.started_at >= today_start,
|
|
)
|
|
)
|
|
runs_today: int = result.scalar_one()
|
|
|
|
if runs_today >= limit:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_402_PAYMENT_REQUIRED,
|
|
detail=f"Daily batch run limit ({limit}) reached for your tier. Upgrade for more runs.",
|
|
)
|
|
|
|
|
|
# ── Catalog ───────────────────────────────────────────────────────────
|
|
|
|
@router.get("/catalog", response_model=list[AgentCatalogItem])
|
|
async def get_agent_catalog(
|
|
current_user: UserProfile = Depends(get_current_user),
|
|
) -> list[AgentCatalogItem]:
|
|
"""Return the static list of available agent types and their descriptions."""
|
|
return [
|
|
AgentCatalogItem(
|
|
type="local_directory",
|
|
name="Local Directory Monitor",
|
|
description="Watches local directories, extracts data from files using AI",
|
|
),
|
|
AgentCatalogItem(
|
|
type="gmail",
|
|
name="Gmail Connector",
|
|
description="Scans Gmail inbox, extracts tasks/notes from emails",
|
|
),
|
|
AgentCatalogItem(
|
|
type="teams",
|
|
name="Microsoft Teams Connector",
|
|
description="Monitors Teams messages, extracts action items",
|
|
),
|
|
AgentCatalogItem(
|
|
type="outlook",
|
|
name="Outlook Connector",
|
|
description="Scans Outlook inbox, extracts tasks/notes",
|
|
),
|
|
]
|
|
|
|
|
|
@router.post("/can-create", response_model=AgentCreationCheckResponse)
|
|
async def can_create_agent(
|
|
body: AgentCreationCheckRequest,
|
|
current_user: UserProfile = Depends(get_current_user),
|
|
) -> AgentCreationCheckResponse:
|
|
"""Check if the user can create one more agent based on billing tier.
|
|
|
|
Since configuration is client-owned, the Electron app sends its current
|
|
active agent count and the backend applies tier limits.
|
|
"""
|
|
limit: int = FEATURES.get(current_user.tier, FEATURES["free"])["batch_active"]
|
|
allowed = limit == -1 or body.active_agents < limit
|
|
return AgentCreationCheckResponse(
|
|
allowed=allowed,
|
|
tier=current_user.tier,
|
|
active_agents=body.active_agents,
|
|
limit=limit,
|
|
)
|
|
|
|
|
|
@router.post("/trigger", response_model=AgentRunLogResponse, status_code=status.HTTP_202_ACCEPTED)
|
|
async def trigger_agent_run(
|
|
body: AgentTriggerRequest,
|
|
current_user: UserProfile = Depends(get_current_user),
|
|
db: AsyncSession = Depends(get_session),
|
|
) -> AgentRunLogResponse:
|
|
"""Trigger a local agent run using client-provided configuration."""
|
|
_enforce_agent_limit(current_user.tier, body.active_agents)
|
|
await _enforce_run_frequency(current_user.tier, current_user.id, db)
|
|
|
|
last_run_dt = (
|
|
datetime.fromtimestamp(body.last_run_at / 1000, tz=timezone.utc)
|
|
if body.last_run_at
|
|
else None
|
|
)
|
|
config = LocalAgentConfig(
|
|
id=str(uuid.uuid4()),
|
|
user_id=current_user.id,
|
|
device_id=body.device_id,
|
|
name="Local Directory Monitor",
|
|
directory_paths=[body.directory],
|
|
data_types=_to_data_types(body.what_to_extract),
|
|
prompt_template=body.custom_agent_prompt or "",
|
|
agent_config=body.agent_config,
|
|
file_extensions=[],
|
|
schedule_cron=body.batch_interval,
|
|
enabled=True,
|
|
last_run_at=last_run_dt,
|
|
)
|
|
|
|
# Use the FE's stable agent_id if provided, fall back to the ephemeral config id.
|
|
stable_agent_id = body.agent_id or config.id
|
|
|
|
if is_agent_running(stable_agent_id):
|
|
raise HTTPException(
|
|
status_code=status.HTTP_409_CONFLICT,
|
|
detail="Agent is already running. Only one run per agent is allowed at a time.",
|
|
)
|
|
|
|
run_log = AgentRunLog(
|
|
agent_id=stable_agent_id,
|
|
agent_type="local",
|
|
user_id=current_user.id,
|
|
status="running",
|
|
)
|
|
db.add(run_log)
|
|
await db.commit()
|
|
await db.refresh(run_log)
|
|
|
|
run_context = {
|
|
"type": "agent_batch",
|
|
"run_id": run_log.id,
|
|
"agent_id": stable_agent_id,
|
|
}
|
|
|
|
asyncio.create_task(
|
|
run_local_agent(current_user.id, config, run_log, device_manager, run_context)
|
|
)
|
|
|
|
return _to_run_log_response(run_log)
|