company-ai-platform/app/modules/observability/service.py

from datetime import timedelta
from typing import Any

from sqlalchemy import select, text
from sqlalchemy.orm import Session

from app.core.config import get_settings
from app.core.constants import ActorValue
from app.core.utils.time import utc_now
from app.modules.ai_memory.service import AIMemoryService
from app.modules.audit.constants import (
    AuditAction,
    AuditRiskLevel,
    AuditSource,
    AuditTargetType,
)
from app.modules.audit.schemas import AuditLogCreate
from app.modules.audit.service import AuditService
from app.modules.events.constants import EventStatus
from app.modules.events.service import EventService
from app.modules.observability.constants import (
    HeartbeatStatus,
    ObservabilityKey,
    ObservabilityMetricKey,
    ObservabilityStatus,
)
from app.modules.observability.models import SystemHeartbeat
from app.modules.workflows.constants import WorkflowStatus
from app.modules.workflows.service import WorkflowService


class ObservabilityService:
    """Build health, readiness, and JSON metrics for V3 operations."""

    def __init__(self, db: Session):
        self.db = db

    def live(self) -> dict[str, str]:
        return {ObservabilityKey.STATUS: ObservabilityStatus.OK}

    def ready(self) -> dict[str, Any]:
        checks = {
            ObservabilityKey.DATABASE: self._database_check(),
            ObservabilityKey.REDIS: self._redis_check(),
            ObservabilityKey.EVENTS: self._events_check(),
            ObservabilityKey.WORKFLOWS: self._workflows_check(),
            ObservabilityKey.HEARTBEATS: self._heartbeats_check(),
        }
        degraded = any(
            item[ObservabilityKey.STATUS]
            in {ObservabilityStatus.DEGRADED, ObservabilityStatus.ERROR}
            for item in checks.values()
        )
        return {
            ObservabilityKey.STATUS: (
                ObservabilityStatus.DEGRADED if degraded else ObservabilityStatus.OK
            ),
            ObservabilityKey.CHECKS: checks,
        }

    def metrics(self) -> dict[str, Any]:
        return {
            ObservabilityKey.METRICS: {
                ObservabilityKey.EVENTS: EventService(self.db).count_by_status(),
                ObservabilityKey.WORKFLOWS: WorkflowService(self.db).count_by_status(),
                ObservabilityKey.AI_MEMORY: AIMemoryService(self.db).count_by_status(),
                ObservabilityKey.HEARTBEATS: self.heartbeat_summary(),
            }
        }

    def record_heartbeat(
        self,
        component: str,
        instance_id: str,
        status_value: str = HeartbeatStatus.OK,
        actor: str = ActorValue.SYSTEM,
    ) -> dict[str, Any]:
        now = utc_now()
        record = self.db.execute(
            select(SystemHeartbeat).where(
                SystemHeartbeat.component == component,
                SystemHeartbeat.instance_id == instance_id,
            )
        ).scalar_one_or_none()
        if record is None:
            record = SystemHeartbeat(
                component=component,
                instance_id=instance_id,
                status=status_value,
                last_seen_at=now,
            )
            self.db.add(record)
        else:
            record.status = status_value
            record.last_seen_at = now
            record.updated_at = now
        self.db.commit()
        self.db.refresh(record)
        AuditService(self.db).log(
            AuditLogCreate(
                actor=actor,
                source=AuditSource.OBSERVABILITY,
                action=AuditAction.HEARTBEAT,
                target_type=AuditTargetType.HEARTBEAT,
                target_id=f"{component}:{instance_id}",
                risk_level=AuditRiskLevel.LOW,
                response_payload={
                    ObservabilityMetricKey.COMPONENT: component,
                    ObservabilityMetricKey.INSTANCE_ID: instance_id,
                    ObservabilityKey.STATUS: status_value,
                    ObservabilityMetricKey.LAST_SEEN_AT: record.last_seen_at.isoformat(),
                },
            )
        )
        return self._serialize_heartbeat(record)

    def heartbeat_summary(self) -> dict[str, Any]:
        records = list(self.db.execute(select(SystemHeartbeat)).scalars())
        threshold = self._heartbeat_stale_threshold()
        stale = [item for item in records if item.last_seen_at < threshold]
        active = len(records) - len(stale)
        last_seen_at = max((item.last_seen_at for item in records), default=None)
        return {
            ObservabilityMetricKey.TOTAL: len(records),
            ObservabilityMetricKey.ACTIVE: active,
            ObservabilityMetricKey.STALE: len(stale),
            ObservabilityMetricKey.LAST_SEEN_AT: (
                last_seen_at.isoformat() if last_seen_at else None
            ),
            ObservabilityMetricKey.ITEMS: [
                self._serialize_heartbeat(item) for item in records
            ],
        }

    def _database_check(self) -> dict[str, Any]:
        try:
            self.db.execute(text("select 1")).scalar()
        except Exception as exc:
            return {
                ObservabilityKey.STATUS: ObservabilityStatus.ERROR,
                ObservabilityMetricKey.ERROR: str(exc),
            }
        return {ObservabilityKey.STATUS: ObservabilityStatus.OK}

    def _redis_check(self) -> dict[str, Any]:
        settings = get_settings()
        if not settings.task_queue_enabled:
            return {ObservabilityKey.STATUS: ObservabilityStatus.SKIPPED}
        try:
            from redis import Redis

            Redis.from_url(settings.redis_url, socket_connect_timeout=1).ping()
        except Exception as exc:
            return {
                ObservabilityKey.STATUS: ObservabilityStatus.DEGRADED,
                ObservabilityMetricKey.ERROR: str(exc),
            }
        return {ObservabilityKey.STATUS: ObservabilityStatus.OK}

    def _events_check(self) -> dict[str, Any]:
        counts = EventService(self.db).count_by_status()
        failed = counts.get(EventStatus.FAILED, 0)
        return {
            ObservabilityKey.STATUS: (
                ObservabilityStatus.DEGRADED if failed else ObservabilityStatus.OK
            ),
            ObservabilityMetricKey.PENDING: counts.get(EventStatus.PENDING, 0),
            ObservabilityMetricKey.FAILED: failed,
        }

    def _workflows_check(self) -> dict[str, Any]:
        counts = WorkflowService(self.db).count_by_status()
        failed = counts.get(WorkflowStatus.FAILED, 0)
        return {
            ObservabilityKey.STATUS: (
                ObservabilityStatus.DEGRADED if failed else ObservabilityStatus.OK
            ),
            ObservabilityMetricKey.RUNNING: counts.get(WorkflowStatus.RUNNING, 0),
            ObservabilityMetricKey.FAILED: failed,
        }

    def _heartbeats_check(self) -> dict[str, Any]:
        summary = self.heartbeat_summary()
        total = summary[ObservabilityMetricKey.TOTAL]
        stale = summary[ObservabilityMetricKey.STALE]
        if total == 0:
            return {ObservabilityKey.STATUS: ObservabilityStatus.SKIPPED}
        return {
            ObservabilityKey.STATUS: (
                ObservabilityStatus.DEGRADED if stale else ObservabilityStatus.OK
            ),
            ObservabilityMetricKey.TOTAL: total,
            ObservabilityMetricKey.STALE: stale,
            ObservabilityMetricKey.LAST_SEEN_AT: summary[
                ObservabilityMetricKey.LAST_SEEN_AT
            ],
        }

    @staticmethod
    def _serialize_heartbeat(record: SystemHeartbeat) -> dict[str, Any]:
        return {
            ObservabilityMetricKey.COMPONENT: record.component,
            ObservabilityMetricKey.INSTANCE_ID: record.instance_id,
            ObservabilityKey.STATUS: record.status,
            ObservabilityMetricKey.LAST_SEEN_AT: record.last_seen_at.isoformat(),
        }

    @staticmethod
    def _heartbeat_stale_threshold() -> Any:
        settings = get_settings()
        return utc_now() - timedelta(seconds=settings.heartbeat_interval_seconds * 3)