Spaces:
Running
Running
| """ | |
| Base Agent Runner - Abstract class for all AI agents | |
| Provides common functionality for insight detection and notification dispatch. | |
| """ | |
| from abc import ABC, abstractmethod | |
| from typing import List, Dict, Any, Optional | |
| from datetime import datetime | |
| import uuid | |
| import logging | |
| import json | |
| from sqlalchemy.ext.asyncio import AsyncSession | |
| from sqlalchemy.future import select | |
| from database.db import AsyncSessionLocal | |
| from database.orm import UserProfile, WorkspaceMember | |
| from database.workspace_ops import get_workspace_members | |
| from services.notification_service import ( | |
| enqueue_notification, | |
| NotificationJob, | |
| should_notify, | |
| in_dnd_window, | |
| check_rate_limit | |
| ) | |
| logger = logging.getLogger(__name__) | |
| class Insight: | |
| """Data class for AI-generated insights""" | |
| def __init__( | |
| self, | |
| title: str, | |
| body: str, | |
| severity: str, # 'low' | 'medium' | 'high' | |
| score: float = None, | |
| metadata: Dict[str, Any] = None, | |
| chart_payload: Dict[str, Any] = None | |
| ): | |
| self.title = title | |
| self.body = body | |
| self.severity = severity | |
| self.score = score | |
| self.metadata = metadata or {} | |
| self.chart_payload = chart_payload | |
| class AgentRunner(ABC): | |
| """Base class for all AI agents""" | |
| def __init__(self, agent_name: str): | |
| self.agent_name = agent_name | |
| self.logger = logging.getLogger(f"agents.{agent_name}") | |
| async def detect_insights(self, workspace_id: str) -> List[Insight]: | |
| """ | |
| Detect insights for a workspace | |
| Must be implemented by each agent | |
| """ | |
| pass | |
| async def run(self, workspace_id: str): | |
| """ | |
| Main execution method | |
| 1. Detect insights | |
| 2. Persist to DB | |
| 3. Log execution | |
| 4. Dispatch notifications | |
| """ | |
| self.logger.info(f"Running {self.agent_name} for workspace {workspace_id}") | |
| try: | |
| # Log agent start | |
| await self._log_execution(workspace_id, 'started', f"{self.agent_name} execution started") | |
| # Detect insights | |
| insights = await self.detect_insights(workspace_id) | |
| self.logger.info(f"Detected {len(insights)} insights for workspace {workspace_id}") | |
| for insight in insights: | |
| await self._process_insight(workspace_id, insight) | |
| # Log agent completion | |
| await self._log_execution( | |
| workspace_id, | |
| 'completed', | |
| f"{self.agent_name} completed: {len(insights)} insights generated" | |
| ) | |
| except Exception as e: | |
| self.logger.error(f"Agent execution failed: {e}", exc_info=True) | |
| # Log error | |
| await self._log_execution( | |
| workspace_id, | |
| 'error', | |
| f"{self.agent_name} failed: {str(e)}", | |
| {'error': str(e)} | |
| ) | |
| async def _process_insight(self, workspace_id: str, insight: Insight): | |
| """Process a single insight: persist, log, and notify""" | |
| try: | |
| # 1. Persist insight to database | |
| saved_insight = await self._persist_insight(workspace_id, insight) | |
| if not saved_insight: | |
| self.logger.error("Failed to persist insight") | |
| return | |
| insight_id = saved_insight['id'] | |
| # 2. Log insight creation | |
| await self._log_execution( | |
| workspace_id, | |
| 'insight_created', | |
| f"Insight created: {insight.title}", | |
| {'insight_id': insight_id, 'severity': insight.severity} | |
| ) | |
| # 3. Dispatch notifications to workspace members | |
| await self._dispatch_notifications(workspace_id, insight_id, insight) | |
| except Exception as e: | |
| self.logger.error(f"Failed to process insight: {e}", exc_info=True) | |
| async def _persist_insight(self, workspace_id: str, insight: Insight) -> Optional[Dict[str, Any]]: | |
| """Save insight to database via direct SQL (ai_insights table)""" | |
| try: | |
| async with AsyncSessionLocal() as db: | |
| # Use raw SQL insert since we don't have a dedicated AIInsight ORM model yet | |
| from sqlalchemy import text | |
| insight_id = str(uuid.uuid4()) | |
| stmt = text(""" | |
| INSERT INTO ai_insights (id, workspace_id, title, body, severity, score, metadata, chart_payload, created_by_agent, created_at) | |
| VALUES (:id, :workspace_id, :title, :body, :severity, :score, :metadata, :chart_payload, :agent, :created_at) | |
| ON CONFLICT DO NOTHING | |
| RETURNING id | |
| """) | |
| result = await db.execute(stmt, { | |
| 'id': insight_id, | |
| 'workspace_id': workspace_id, | |
| 'title': insight.title, | |
| 'body': insight.body, | |
| 'severity': insight.severity, | |
| 'score': insight.score, | |
| 'metadata': json.dumps(insight.metadata) if insight.metadata else '{}', | |
| 'chart_payload': json.dumps(insight.chart_payload) if insight.chart_payload else None, | |
| 'agent': self.agent_name, | |
| 'created_at': datetime.utcnow() | |
| }) | |
| await db.commit() | |
| return {'id': insight_id} | |
| except Exception as e: | |
| self.logger.error(f"Failed to persist insight: {e}") | |
| # If the ai_insights table doesn't exist yet, return a generated ID | |
| # so notification dispatch can still proceed | |
| return {'id': str(uuid.uuid4())} | |
| async def _dispatch_notifications(self, workspace_id: str, insight_id: str, insight: Insight): | |
| """Send notifications to all workspace members based on their preferences""" | |
| try: | |
| # Get workspace members from PostgreSQL | |
| async with AsyncSessionLocal() as db: | |
| members = await get_workspace_members(db, workspace_id) | |
| self.logger.info(f"Dispatching notifications to {len(members)} members") | |
| for member in members: | |
| user_id = member['user_id'] | |
| # Use default notification settings | |
| # Default: email enabled, push disabled, medium threshold | |
| default_settings = { | |
| 'ai_insights': True, | |
| 'email_notifications': True, | |
| 'push_notifications': False, | |
| 'severity_threshold': 'medium', | |
| 'dnd_start': None, | |
| 'dnd_end': None | |
| } | |
| # Check if AI insights are enabled | |
| if not default_settings.get('ai_insights', True): | |
| continue | |
| # Check severity threshold | |
| if not should_notify(insight.severity, default_settings.get('severity_threshold', 'medium')): | |
| self.logger.info(f"Insight severity {insight.severity} below threshold for user {user_id}") | |
| continue | |
| # Check Do Not Disturb window | |
| if in_dnd_window(default_settings.get('dnd_start'), default_settings.get('dnd_end')): | |
| self.logger.info(f"User {user_id} is in DND window") | |
| continue | |
| # Check rate limit | |
| if await check_rate_limit(workspace_id, user_id): | |
| self.logger.warning(f"Rate limit exceeded for user {user_id}") | |
| continue | |
| # Enqueue notification | |
| job = NotificationJob( | |
| insight_id=insight_id, | |
| workspace_id=workspace_id, | |
| user_id=user_id, | |
| channels={ | |
| 'email': default_settings.get('email_notifications', True), | |
| 'push': default_settings.get('push_notifications', False) | |
| }, | |
| payload={ | |
| 'title': insight.title, | |
| 'body': insight.body, | |
| 'chart_payload': insight.chart_payload | |
| } | |
| ) | |
| await enqueue_notification(job) | |
| self.logger.info(f"Notification enqueued for user {user_id}") | |
| except Exception as e: | |
| self.logger.error(f"Failed to dispatch notifications: {e}", exc_info=True) | |
| async def _log_execution( | |
| self, | |
| workspace_id: str, | |
| status: str, | |
| message: str, | |
| metadata: Dict[str, Any] = None | |
| ): | |
| """Log agent execution to database""" | |
| try: | |
| async with AsyncSessionLocal() as db: | |
| from sqlalchemy import text | |
| stmt = text(""" | |
| INSERT INTO agent_logs (id, agent_name, workspace_id, status, message, metadata, created_at) | |
| VALUES (:id, :agent_name, :workspace_id, :status, :message, :metadata, :created_at) | |
| ON CONFLICT DO NOTHING | |
| """) | |
| await db.execute(stmt, { | |
| 'id': str(uuid.uuid4()), | |
| 'agent_name': self.agent_name, | |
| 'workspace_id': workspace_id, | |
| 'status': status, | |
| 'message': message, | |
| 'metadata': json.dumps(metadata or {}), | |
| 'created_at': datetime.utcnow() | |
| }) | |
| await db.commit() | |
| except Exception as e: | |
| # Don't let logging failures crash the agent | |
| self.logger.error(f"Failed to log execution: {e}") | |