kairo-code 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

kairo/backend/services/agent/agent_crud_service.py

@@ -0,0 +1,150 @@
+"""
+Agent CRUD operations service.
+
+Handles registration, listing, retrieval, updating, and deletion of agents.
+"""
+
+import json
+import logging
+from datetime import datetime, timedelta, UTC
+
+from sqlalchemy import func, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from backend.models.agent import Agent, AgentMetrics1m
+from backend.schemas.agent import RegisterAgentRequest, UpdateAgentRequest
+
+logger = logging.getLogger(__name__)
+
+
+class AgentCrudService:
+    """Service for agent CRUD operations."""
+
+    def __init__(self, db: AsyncSession):
+        self.db = db
+
+    async def register(self, user_id: str, req: RegisterAgentRequest) -> Agent:
+        """Register a new agent."""
+        agent = Agent(
+            user_id=user_id,
+            name=req.name,
+            description=req.description,
+            system_prompt=req.system_prompt,
+            model_preference=req.model_preference,
+            agent_type=req.agent_type,
+            tools_config=json.dumps(req.tools) if req.tools else None,
+            state="created",
+        )
+        self.db.add(agent)
+        await self.db.commit()
+        await self.db.refresh(agent)
+        logger.info("Agent registered: user=%s name=%s id=%s", user_id, req.name, agent.id)
+        return agent
+
+    async def list_agents(self, user_id: str) -> list[Agent]:
+        """List all agents for a user."""
+        stmt = (
+            select(Agent)
+            .where(Agent.user_id == user_id)
+            .where(Agent.deleted_at.is_(None))
+            .order_by(Agent.created_at.desc())
+        )
+        result = await self.db.execute(stmt)
+        return list(result.scalars().all())
+
+    async def list_agents_with_metrics(self, user_id: str) -> list[dict]:
+        """List agents with 24h summary metrics."""
+        agents = await self.list_agents(user_id)
+        cutoff = datetime.now(UTC) - timedelta(hours=24)
+
+        result = []
+        for agent in agents:
+            metrics_stmt = (
+                select(
+                    func.sum(AgentMetrics1m.request_count).label("requests"),
+                    func.sum(AgentMetrics1m.error_count).label("errors"),
+                    func.sum(AgentMetrics1m.input_tokens + AgentMetrics1m.output_tokens).label("tokens"),
+                )
+                .where(AgentMetrics1m.agent_id == agent.id)
+                .where(AgentMetrics1m.bucket_time >= cutoff)
+            )
+            metrics_result = await self.db.execute(metrics_stmt)
+            row = metrics_result.one_or_none()
+
+            requests_24h = row.requests or 0 if row else 0
+            errors_24h = row.errors or 0 if row else 0
+            tokens_24h = row.tokens or 0 if row else 0
+            error_rate = (errors_24h / requests_24h * 100) if requests_24h > 0 else 0.0
+
+            result.append({
+                "id": agent.id,
+                "name": agent.name,
+                "description": agent.description,
+                "model_preference": agent.model_preference,
+                "agent_type": agent.agent_type,
+                "state": agent.state,
+                "sdk_version": agent.sdk_version,
+                "last_heartbeat_at": agent.last_heartbeat_at,
+                "created_at": agent.created_at,
+                "requests_24h": requests_24h,
+                "errors_24h": errors_24h,
+                "tokens_24h": tokens_24h,
+                "error_rate": round(error_rate, 2),
+            })
+
+        return result
+
+    async def get_agent(self, user_id: str, agent_id: str) -> Agent | None:
+        """Get agent by ID for a specific user."""
+        stmt = (
+            select(Agent)
+            .where(Agent.id == agent_id, Agent.user_id == user_id)
+            .where(Agent.deleted_at.is_(None))
+        )
+        result = await self.db.execute(stmt)
+        return result.scalar_one_or_none()
+
+    async def get_agent_by_id(self, agent_id: str) -> Agent | None:
+        """Get agent by ID only (for SDK auth)."""
+        stmt = select(Agent).where(Agent.id == agent_id).where(Agent.deleted_at.is_(None))
+        result = await self.db.execute(stmt)
+        return result.scalar_one_or_none()
+
+    async def update_agent(
+        self,
+        user_id: str,
+        agent_id: str,
+        req: UpdateAgentRequest
+    ) -> Agent | None:
+        """Update agent properties."""
+        agent = await self.get_agent(user_id, agent_id)
+        if not agent:
+            return None
+
+        if req.name is not None:
+            agent.name = req.name
+        if req.description is not None:
+            agent.description = req.description
+        if req.system_prompt is not None:
+            agent.system_prompt = req.system_prompt
+        if req.model_preference is not None:
+            agent.model_preference = req.model_preference
+        if req.tools is not None:
+            agent.tools_config = json.dumps(req.tools)
+
+        agent.updated_at = datetime.now(UTC)
+        await self.db.commit()
+        await self.db.refresh(agent)
+        return agent
+
+    async def delete_agent(self, user_id: str, agent_id: str) -> bool:
+        """Soft delete agent."""
+        agent = await self.get_agent(user_id, agent_id)
+        if not agent:
+            return False
+
+        agent.deleted_at = datetime.now(UTC)
+        agent.state = "disabled"
+        await self.db.commit()
+        logger.info("Agent deleted: id=%s user=%s", agent_id, user_id)
+        return True

kairo/backend/services/agent/agent_events_service.py

@@ -0,0 +1,103 @@
+"""
+Agent events service.
+
+Handles event logging and querying for agent activity tracking.
+"""
+
+import logging
+from datetime import datetime, timedelta, UTC
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from backend.models.agent import Agent, AgentEvent
+from backend.services.agent.constants import DEFAULT_EVENTS_RETENTION_DAYS
+
+logger = logging.getLogger(__name__)
+
+
+class AgentEventsService:
+    """Service for agent event management."""
+
+    def __init__(self, db: AsyncSession):
+        self.db = db
+
+    async def log_event(
+        self,
+        agent_id: str,
+        event_type: str,
+        event_data: dict | None = None,
+        error_type: str | None = None,
+        error_message: str | None = None,
+        client_ip: str | None = None,
+    ) -> AgentEvent:
+        """Log an agent event."""
+        event = AgentEvent(
+            agent_id=agent_id,
+            event_type=event_type,
+            event_data=event_data,
+            error_type=error_type,
+            error_message=error_message,
+            client_ip=client_ip,
+        )
+        self.db.add(event)
+        return event
+
+    async def get_events(
+        self,
+        agent_id: str,
+        event_type: str | None = None,
+        limit: int = 50,
+    ) -> list[AgentEvent]:
+        """Get agent events with optional type filter."""
+        stmt = (
+            select(AgentEvent)
+            .where(AgentEvent.agent_id == agent_id)
+            .order_by(AgentEvent.created_at.desc())
+            .limit(limit)
+        )
+        if event_type:
+            stmt = stmt.where(AgentEvent.event_type == event_type)
+
+        result = await self.db.execute(stmt)
+        return list(result.scalars().all())
+
+    async def get_events_for_user(
+        self,
+        user_id: str,
+        agent_id: str,
+        event_type: str | None = None,
+        limit: int = 50,
+    ) -> list[AgentEvent] | None:
+        """Get events for an agent owned by a user. Returns None if agent not found."""
+        # Verify agent ownership
+        agent_stmt = (
+            select(Agent)
+            .where(Agent.id == agent_id, Agent.user_id == user_id)
+            .where(Agent.deleted_at.is_(None))
+        )
+        agent_result = await self.db.execute(agent_stmt)
+        if not agent_result.scalar_one_or_none():
+            return None
+
+        return await self.get_events(agent_id, event_type, limit)
+
+    async def cleanup_old_events(
+        self,
+        retention_days: int = DEFAULT_EVENTS_RETENTION_DAYS
+    ) -> int:
+        """Clean up old events based on retention policy."""
+        cutoff = datetime.now(UTC) - timedelta(days=retention_days)
+        stmt = select(AgentEvent).where(AgentEvent.created_at < cutoff)
+        result = await self.db.execute(stmt)
+        events = list(result.scalars().all())
+        count = len(events)
+
+        for event in events:
+            await self.db.delete(event)
+
+        if count > 0:
+            await self.db.commit()
+            logger.info("Cleaned up %d old events", count)
+
+        return count

kairo/backend/services/agent/agent_heartbeat_service.py

@@ -0,0 +1,207 @@
+"""
+Agent heartbeat service.
+
+Handles heartbeat processing, agent state management, and background state checks.
+"""
+
+import logging
+from datetime import datetime, timedelta, UTC
+
+from sqlalchemy import select, or_
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from backend.models.agent import Agent, AgentMetrics1m
+from backend.schemas.agent import AgentHeartbeatRequest
+from backend.services.agent.agent_events_service import AgentEventsService
+from backend.services.agent.agent_commands_service import AgentCommandsService
+from backend.services.agent.constants import (
+    DEFAULT_HEARTBEAT_TIMEOUT_SECONDS,
+    DEFAULT_STALE_AGENT_DAYS,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class AgentHeartbeatService:
+    """Service for agent heartbeat processing and state management."""
+
+    def __init__(self, db: AsyncSession):
+        self.db = db
+        self._events_service = AgentEventsService(db)
+        self._commands_service = AgentCommandsService(db)
+
+    async def heartbeat(
+        self,
+        agent: Agent,
+        req: AgentHeartbeatRequest,
+        client_ip: str | None = None
+    ) -> tuple[Agent, list[dict]]:
+        """
+        Process heartbeat and return agent + pending commands.
+
+        Args:
+            agent: The agent to process heartbeat for
+            req: The heartbeat request data
+            client_ip: Optional client IP address
+
+        Returns:
+            Tuple of (updated agent, list of pending commands)
+        """
+        now = datetime.now(UTC)
+        is_first_connection = agent.first_connected_at is None
+
+        # Update agent status
+        agent.status = req.status
+        agent.last_heartbeat_at = now
+
+        # Update state based on status
+        self._update_agent_state(agent, req, now, is_first_connection)
+
+        # Update SDK info
+        if req.sdk_version:
+            agent.sdk_version = req.sdk_version
+        if req.host_info:
+            agent.host_info = req.host_info.model_dump()
+
+        # Record metrics if provided
+        if req.metrics_since_last_heartbeat:
+            await self._record_metrics(agent.id, req.metrics_since_last_heartbeat)
+
+        # Log event
+        await self._events_service.log_event(
+            agent_id=agent.id,
+            event_type="heartbeat",
+            event_data={
+                "status": req.status,
+                "sdk_version": req.sdk_version,
+                "queue_depth": req.queue_depth,
+                "active_request": req.active_request,
+            },
+            client_ip=client_ip,
+        )
+
+        agent.updated_at = now
+        await self.db.commit()
+        await self.db.refresh(agent)
+
+        # Get pending commands
+        commands = await self._commands_service.get_pending_commands(agent.id)
+
+        return agent, commands
+
+    def _update_agent_state(
+        self,
+        agent: Agent,
+        req: AgentHeartbeatRequest,
+        now: datetime,
+        is_first_connection: bool
+    ) -> None:
+        """Update agent state based on heartbeat status."""
+        if req.status in ("online", "busy", "idle"):
+            agent.state = "online" if req.status == "online" else req.status
+            agent.last_online_at = now
+            if is_first_connection:
+                agent.first_connected_at = now
+        elif req.status == "error":
+            agent.state = "error"
+            if req.last_error:
+                agent.last_error_at = now
+                agent.last_error_message = req.last_error.get("message", "Unknown error")
+
+    async def _record_metrics(self, agent_id: str, metrics) -> None:
+        """Record metrics in 1-minute bucket."""
+        now = datetime.now(UTC)
+        bucket_time = now.replace(second=0, microsecond=0)
+
+        stmt = select(AgentMetrics1m).where(
+            AgentMetrics1m.agent_id == agent_id,
+            AgentMetrics1m.bucket_time == bucket_time
+        )
+        result = await self.db.execute(stmt)
+        existing = result.scalar_one_or_none()
+
+        if existing:
+            existing.request_count += metrics.requests_completed + metrics.requests_failed
+            existing.error_count += metrics.requests_failed
+            existing.input_tokens += metrics.input_tokens
+            existing.output_tokens += metrics.output_tokens
+            existing.total_latency_ms += metrics.total_latency_ms
+            existing.tool_calls += metrics.tool_calls
+            existing.tool_errors += metrics.tool_errors
+        else:
+            new_metrics = AgentMetrics1m(
+                agent_id=agent_id,
+                bucket_time=bucket_time,
+                request_count=metrics.requests_completed + metrics.requests_failed,
+                error_count=metrics.requests_failed,
+                input_tokens=metrics.input_tokens,
+                output_tokens=metrics.output_tokens,
+                total_latency_ms=metrics.total_latency_ms,
+                tool_calls=metrics.tool_calls,
+                tool_errors=metrics.tool_errors,
+            )
+            self.db.add(new_metrics)
+
+    async def mark_stale_agents_offline(
+        self,
+        threshold_seconds: int = DEFAULT_HEARTBEAT_TIMEOUT_SECONDS
+    ) -> int:
+        """Mark agents as offline if no heartbeat within threshold."""
+        cutoff = datetime.now(UTC) - timedelta(seconds=threshold_seconds)
+        stmt = (
+            select(Agent)
+            .where(Agent.state.in_(["online", "busy", "idle"]))
+            .where(Agent.deleted_at.is_(None))
+            .where(
+                or_(
+                    Agent.last_heartbeat_at < cutoff,
+                    Agent.last_heartbeat_at.is_(None)
+                )
+            )
+        )
+        result = await self.db.execute(stmt)
+        agents = list(result.scalars().all())
+
+        for agent in agents:
+            old_state = agent.state
+            agent.state = "offline"
+            agent.status = "offline"
+            agent.updated_at = datetime.now(UTC)
+
+            await self._events_service.log_event(
+                agent_id=agent.id,
+                event_type="state_change",
+                event_data={
+                    "from": old_state,
+                    "to": "offline",
+                    "reason": "heartbeat_timeout"
+                },
+            )
+
+        if agents:
+            await self.db.commit()
+            logger.info("Marked %d agents as offline", len(agents))
+
+        return len(agents)
+
+    async def mark_stale_agents(self, days: int = DEFAULT_STALE_AGENT_DAYS) -> int:
+        """Mark offline agents as stale after days."""
+        cutoff = datetime.now(UTC) - timedelta(days=days)
+        stmt = (
+            select(Agent)
+            .where(Agent.state == "offline")
+            .where(Agent.deleted_at.is_(None))
+            .where(Agent.last_online_at < cutoff)
+        )
+        result = await self.db.execute(stmt)
+        agents = list(result.scalars().all())
+
+        for agent in agents:
+            agent.state = "stale"
+            agent.updated_at = datetime.now(UTC)
+
+        if agents:
+            await self.db.commit()
+            logger.info("Marked %d agents as stale", len(agents))
+
+        return len(agents)