| 1 | from typing import Any |
| 2 | from typing import Dict |
| 3 | from typing import Optional |
| 4 | |
| 5 | from fastapi import HTTPException |
| 6 | from loguru import logger |
| 7 | from sqlalchemy.ext.asyncio import AsyncSession |
| 8 | from sqlalchemy.future import select |
| 9 | from sqlalchemy.orm import selectinload |
| 10 | |
| 11 | from app.connectors.talon.schema.talon import TalonInvestigateRequest |
| 12 | from app.connectors.talon.schema.talon import TalonInvestigateResponse |
| 13 | from app.connectors.talon.schema.talon import TalonJobResponse |
| 14 | from app.connectors.talon.schema.talon import TalonMessageRequest |
| 15 | from app.connectors.talon.schema.talon import TalonMessageResponse |
| 16 | from app.connectors.talon.schema.talon import TalonStatusResponse |
| 17 | from app.connectors.talon.schema.talon import TalonTemplatesResponse |
| 18 | from app.connectors.talon.utils.universal import send_get_request |
| 19 | from app.connectors.talon.utils.universal import send_post_request |
| 20 | from app.connectors.talon.utils.universal import send_post_request_sse |
| 21 | from app.db.universal_models import AiAnalystJob |
| 22 | |
| 23 | |
| 24 | async def send_talon_message(request: TalonMessageRequest) -> TalonMessageResponse: |
| 25 | """ |
| 26 | Send a message to Talon for ad-hoc analyst prompts. |
| 27 | |
| 28 | Args: |
| 29 | request: The message request containing the message and sender. |
| 30 | |
| 31 | Returns: |
| 32 | TalonMessageResponse with the Talon response. |
| 33 | """ |
| 34 | logger.info(f"Sending message to Talon: {request.message}") |
| 35 | response = await send_post_request( |
| 36 | endpoint="/message", |
| 37 | data=request.model_dump(), |
| 38 | timeout=600, |
| 39 | ) |
| 40 | if not response.get("success"): |
| 41 | raise HTTPException( |
| 42 | status_code=500, |
| 43 | detail=response.get("message", "Failed to send message to Talon"), |
| 44 | ) |
| 45 | return TalonMessageResponse( |
| 46 | success=True, |
| 47 | message="Message sent to Talon successfully", |
| 48 | data=response.get("data"), |
| 49 | ) |
| 50 | |
| 51 | |
| 52 | async def stream_talon_message(request: TalonMessageRequest): |
| 53 | """ |
| 54 | Stream a message response from Talon via SSE. |
| 55 | |
| 56 | Args: |
| 57 | request: The message request containing the message and sender. |
| 58 | |
| 59 | Yields: |
| 60 | str: SSE lines from Talon. |
| 61 | """ |
| 62 | logger.info(f"Streaming message to Talon: {request.message}") |
| 63 | async for chunk in send_post_request_sse( |
| 64 | endpoint="/message", |
| 65 | data=request.model_dump(), |
| 66 | ): |
| 67 | yield chunk |
| 68 | |
| 69 | |
| 70 | async def investigate_alert(request: TalonInvestigateRequest) -> TalonInvestigateResponse: |
| 71 | """ |
| 72 | Trigger an investigation for a specific alert. |
| 73 | |
| 74 | Args: |
| 75 | request: The investigation request containing the alert ID, customer code, and sender. |
| 76 | |
| 77 | Returns: |
| 78 | TalonInvestigateResponse with the investigation result. |
| 79 | """ |
| 80 | logger.info(f"Triggering Talon investigation for alert ID: {request.alert_id}") |
| 81 | response = await send_post_request( |
| 82 | endpoint="/investigate", |
| 83 | data={"alert_id": request.alert_id, "customer_code": request.customer_code, "sender": request.sender}, |
| 84 | ) |
| 85 | if not response.get("success"): |
| 86 | raise HTTPException( |
| 87 | status_code=500, |
| 88 | detail=response.get("message", "Failed to trigger Talon investigation"), |
| 89 | ) |
| 90 | return TalonInvestigateResponse( |
| 91 | success=True, |
| 92 | message="Investigation triggered successfully", |
| 93 | data=response.get("data"), |
| 94 | ) |
| 95 | |
| 96 | |
| 97 | async def get_talon_status() -> TalonStatusResponse: |
| 98 | """ |
| 99 | Get the current status of the Talon service including queue and job overview. |
| 100 | |
| 101 | Returns: |
| 102 | TalonStatusResponse with the status data. |
| 103 | """ |
| 104 | logger.info("Fetching Talon status") |
| 105 | response = await send_get_request(endpoint="/status") |
| 106 | if not response.get("success"): |
| 107 | raise HTTPException( |
| 108 | status_code=500, |
| 109 | detail=response.get("message", "Failed to get Talon status"), |
| 110 | ) |
| 111 | return TalonStatusResponse( |
| 112 | success=True, |
| 113 | message="Talon status retrieved successfully", |
| 114 | data=response.get("data"), |
| 115 | ) |
| 116 | |
| 117 | |
| 118 | async def get_talon_job(alert_id: int, session: AsyncSession) -> TalonJobResponse: |
| 119 | """ |
| 120 | Get the job status and report for a specific alert from the database. |
| 121 | |
| 122 | Args: |
| 123 | alert_id: The alert ID to look up. |
| 124 | session: The database session. |
| 125 | |
| 126 | Returns: |
| 127 | TalonJobResponse with the job data. |
| 128 | """ |
| 129 | logger.info(f"Fetching Talon job for alert ID: {alert_id}") |
| 130 | result = await session.execute( |
| 131 | select(AiAnalystJob) |
| 132 | .where(AiAnalystJob.alert_id == alert_id) |
| 133 | .options(selectinload(AiAnalystJob.reports)) |
| 134 | .order_by(AiAnalystJob.created_at.desc()), |
| 135 | ) |
| 136 | jobs = result.scalars().all() |
| 137 | if not jobs: |
| 138 | raise HTTPException( |
| 139 | status_code=404, |
| 140 | detail=f"No job found for alert {alert_id}", |
| 141 | ) |
| 142 | # Use the most recent job for top-level metadata |
| 143 | job = jobs[0] |
| 144 | # Aggregate reports from all jobs |
| 145 | all_reports = [] |
| 146 | for j in jobs: |
| 147 | for r in j.reports or []: |
| 148 | all_reports.append( |
| 149 | { |
| 150 | "id": r.id, |
| 151 | "job_id": j.id, |
| 152 | "alert_id": j.alert_id, |
| 153 | "customer_code": j.customer_code, |
| 154 | "severity_assessment": r.severity_assessment, |
| 155 | "summary": r.summary, |
| 156 | "report_markdown": r.report_markdown, |
| 157 | "recommended_actions": r.recommended_actions, |
| 158 | "created_at": r.created_at.isoformat(), |
| 159 | }, |
| 160 | ) |
| 161 | return TalonJobResponse( |
| 162 | success=True, |
| 163 | message="Talon job retrieved successfully", |
| 164 | data={ |
| 165 | "id": job.id, |
| 166 | "alert_id": job.alert_id, |
| 167 | "customer_code": job.customer_code, |
| 168 | "status": job.status, |
| 169 | "alert_type": job.alert_type, |
| 170 | "triggered_by": job.triggered_by, |
| 171 | "template_used": job.template_used, |
| 172 | "created_at": job.created_at.isoformat(), |
| 173 | "started_at": job.started_at.isoformat() if job.started_at else None, |
| 174 | "completed_at": job.completed_at.isoformat() if job.completed_at else None, |
| 175 | "error_message": job.error_message, |
| 176 | "reports": all_reports, |
| 177 | }, |
| 178 | ) |
| 179 | |
| 180 | |
| 181 | async def replay_investigation( |
| 182 | alert_id: int, |
| 183 | customer_code: str, |
| 184 | template_override: str, |
| 185 | sender: str = "copilot-replay", |
| 186 | ) -> Dict[str, Any]: |
| 187 | """ |
| 188 | Trigger an investigation replay with a forced template via Talon's |
| 189 | POST /investigate endpoint. |
| 190 | |
| 191 | Args: |
| 192 | alert_id: CoPilot alert ID to re-investigate. |
| 193 | customer_code: Customer code for the alert. |
| 194 | template_override: Template filename to force (validated upstream). |
| 195 | sender: Audit identifier for the replay. |
| 196 | |
| 197 | Returns: |
| 198 | Raw Talon response envelope (success, message, data). |
| 199 | """ |
| 200 | logger.info( |
| 201 | f"Replaying Talon investigation for alert {alert_id} " f"with template_override={template_override}", |
| 202 | ) |
| 203 | response = await send_post_request( |
| 204 | endpoint="/investigate", |
| 205 | data={ |
| 206 | "alert_id": alert_id, |
| 207 | "customer_code": customer_code, |
| 208 | "template_override": template_override, |
| 209 | "sender": sender, |
| 210 | }, |
| 211 | ) |
| 212 | if not response.get("success"): |
| 213 | raise HTTPException( |
| 214 | status_code=500, |
| 215 | detail=response.get("message", "Failed to replay Talon investigation"), |
| 216 | ) |
| 217 | return response |
| 218 | |
| 219 | |
| 220 | async def list_talon_templates() -> TalonTemplatesResponse: |
| 221 | """ |
| 222 | List the prompt templates available in NanoClaw's CoPilot group. |
| 223 | Powers the "Re-run with different template" picker in the review UI. |
| 224 | |
| 225 | NanoClaw returns {templates: [{filename, size_bytes, modified_at, first_line}]}. |
| 226 | We surface that envelope directly — template bodies stay server-side. |
| 227 | """ |
| 228 | logger.info("Fetching Talon templates list") |
| 229 | response = await send_get_request(endpoint="/templates") |
| 230 | if not response.get("success"): |
| 231 | raise HTTPException( |
| 232 | status_code=500, |
| 233 | detail=response.get("message", "Failed to list Talon templates"), |
| 234 | ) |
| 235 | data = response.get("data") or {} |
| 236 | raw_templates = data.get("templates") if isinstance(data, dict) else None |
| 237 | if raw_templates is None: |
| 238 | raw_templates = [] |
| 239 | return TalonTemplatesResponse( |
| 240 | success=True, |
| 241 | message=f"{len(raw_templates)} templates retrieved", |
| 242 | templates=raw_templates, |
| 243 | ) |
| 244 | |
| 245 | |
| 246 | async def search_palace_lessons( |
| 247 | customer_code: str, |
| 248 | query: str, |
| 249 | room: Optional[str] = None, |
| 250 | limit: int = 5, |
| 251 | ) -> Dict[str, Any]: |
| 252 | """ |
| 253 | Preview similar MemPalace lessons via Talon's GET /palace/search endpoint. |
| 254 | Read-only — never mutates the palace. |
| 255 | |
| 256 | Args: |
| 257 | customer_code: Customer whose wing to search. |
| 258 | query: Semantic search query. |
| 259 | room: Optional room filter (environment, false_positives, assets, threat_intel, alerts). |
| 260 | limit: Max hits to return (clamped 1–25 upstream). |
| 261 | |
| 262 | Returns: |
| 263 | Raw Talon response envelope (success, message, data). |
| 264 | """ |
| 265 | logger.info( |
| 266 | f"Searching MemPalace for customer={customer_code} room={room} query={query!r} limit={limit}", |
| 267 | ) |
| 268 | params: Dict[str, Any] = { |
| 269 | "customer_code": customer_code, |
| 270 | "query": query, |
| 271 | "limit": limit, |
| 272 | } |
| 273 | if room: |
| 274 | params["room"] = room |
| 275 | response = await send_get_request(endpoint="/palace/search", params=params) |
| 276 | if not response.get("success"): |
| 277 | raise HTTPException( |
| 278 | status_code=500, |
| 279 | detail=response.get("message", "Failed to search MemPalace"), |
| 280 | ) |
| 281 | return response |