""" NeMo/Nemotron External Offline Runner ===================================== Runs an already-approved sanitized request pack through NVIDIA NIM/Nemotron and writes AWOOOI's external result contract. This service never executes tools, never mutates production systems, and never reads fixture labels. """ from __future__ import annotations import asyncio import json import time from dataclasses import dataclass, field from typing import Any, Protocol import httpx from src.services.agent_nemotron_replay_adapter import ( EXTERNAL_RESULT_SCHEMA_VERSION, NEMOTRON_CANDIDATE_ID, NEMOTRON_CONTRACT_TUNED_VARIANT_ID, REQUEST_SCHEMA_VERSION, ) EXTERNAL_RUNNER_REPORT_SCHEMA_VERSION = "agent_nemotron_external_runner_report_v1" DEFAULT_NVIDIA_CHAT_COMPLETIONS_URL = "https://integrate.api.nvidia.com/v1/chat/completions" DEFAULT_NEMOTRON_MODEL = "nvidia/nemotron-mini-4b-instruct" DEFAULT_TIMEOUT_SECONDS = 60.0 DEFAULT_MAX_TOKENS = 900 DEFAULT_CONCURRENCY = 1 _RISK_LEVELS = {"low", "medium", "high", "critical"} _REQUIRED_MODEL_FIELDS = { "proposed_action", "action_plan", "risk_level", "requires_human_approval", "blocked_by_policy", } _SELF_GRADING_FIELDS = { "evaluation_labels", "verification_result", "execution_success", "execution_error", "self_healing_score", "rca_correct", "tool_dry_run_pass", "repair_success", "false_repair", } class AsyncChatClient(Protocol): """Minimal async client protocol for tests and httpx.""" async def post( self, url: str, *, headers: dict[str, str], json: dict[str, Any], ) -> Any: ... @dataclass(frozen=True) class NemotronExternalRunnerConfig: """NVIDIA/NIM request configuration.""" api_key: str base_url: str = DEFAULT_NVIDIA_CHAT_COMPLETIONS_URL model: str = DEFAULT_NEMOTRON_MODEL timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS max_tokens: int = DEFAULT_MAX_TOKENS temperature: float = 0.0 concurrency: int = DEFAULT_CONCURRENCY @dataclass(frozen=True) class NemotronExternalRunnerReport: """Run summary for an external NeMo/Nemotron replay batch.""" requests: int results: int valid: bool model: str failures: list[str] = field(default_factory=list) external_error_records: int = 0 fallback_used_records: int = 0 trace_incomplete_records: int = 0 retry_used_records: int = 0 total_cost_usd: float = 0.0 avg_latency_ms: float = 0.0 p95_latency_ms: float = 0.0 candidate_variant_id: str | None = None def to_dict(self) -> dict[str, Any]: payload = { "schema_version": EXTERNAL_RUNNER_REPORT_SCHEMA_VERSION, "candidate_id": NEMOTRON_CANDIDATE_ID, "requests": self.requests, "results": self.results, "valid": self.valid, "model": self.model, "failures": list(self.failures), "external_error_records": self.external_error_records, "fallback_used_records": self.fallback_used_records, "trace_incomplete_records": self.trace_incomplete_records, "retry_used_records": self.retry_used_records, "total_cost_usd": round(self.total_cost_usd, 6), "avg_latency_ms": round(self.avg_latency_ms, 4), "p95_latency_ms": round(self.p95_latency_ms, 4), } if self.candidate_variant_id: payload["candidate_variant_id"] = self.candidate_variant_id return payload async def run_nemotron_external_replay( *, requests: list[dict[str, Any]], config: NemotronExternalRunnerConfig, client: AsyncChatClient | None = None, ) -> tuple[list[dict[str, Any]], NemotronExternalRunnerReport]: """Run sanitized NeMo replay requests through NVIDIA NIM/Nemotron.""" failures: list[str] = [] _validate_runner_inputs(requests, failures) if not config.api_key.strip(): failures.append("api_key_missing") if failures: return [], NemotronExternalRunnerReport( requests=len(requests), results=0, valid=False, model=config.model, failures=failures, ) owns_client = client is None active_client = client or httpx.AsyncClient( timeout=httpx.Timeout(config.timeout_seconds, connect=10.0), limits=httpx.Limits(max_connections=max(1, config.concurrency)), ) semaphore = asyncio.Semaphore(max(1, config.concurrency)) try: tasks = [ _run_one_request( request=request, config=config, client=active_client, semaphore=semaphore, line_number=index, ) for index, request in enumerate(requests, start=1) ] results = await asyncio.gather(*tasks) finally: if owns_client and hasattr(active_client, "aclose"): await active_client.aclose() runner_failures = [ f"external_error:{result['incident_id']}" for result in results if result.get("error") ] latencies = [float(result.get("latency_ms", 0.0) or 0.0) for result in results] total_cost = sum(float(result.get("cost_usd", 0.0) or 0.0) for result in results) report = NemotronExternalRunnerReport( requests=len(requests), results=len(results), valid=not runner_failures and len(results) == len(requests), model=config.model, failures=runner_failures, external_error_records=sum(1 for result in results if result.get("error")), fallback_used_records=sum(1 for result in results if result.get("fallback_used")), trace_incomplete_records=sum( 1 for result in results if result.get("trace_complete") is not True ), retry_used_records=sum(1 for result in results if result.get("retry_used")), total_cost_usd=total_cost, avg_latency_ms=(sum(latencies) / len(latencies)) if latencies else 0.0, p95_latency_ms=_percentile(latencies, 0.95), candidate_variant_id=_common_candidate_variant_id(requests), ) return results, report async def _run_one_request( *, request: dict[str, Any], config: NemotronExternalRunnerConfig, client: AsyncChatClient, semaphore: asyncio.Semaphore, line_number: int, ) -> dict[str, Any]: run_id = str(request.get("run_id", "")) incident_id = str(request.get("incident_id", "")) candidate_variant_id = _candidate_variant_id(request) started = time.perf_counter() async with semaphore: retry_used = False first_error = None try: payload, content = await _call_chat_completion( request=request, config=config, client=client, ) try: model_output = _normalize_model_output(_extract_json_object(content)) except Exception as exc: if candidate_variant_id != NEMOTRON_CONTRACT_TUNED_VARIANT_ID: raise retry_used = True first_error = _safe_error_text(exc) payload, content = await _call_chat_completion( request=request, config=config, client=client, repair_error=first_error, invalid_content=content, ) model_output = _normalize_model_output(_extract_json_object(content)) error = None fallback_used = False trace_complete = True except Exception as exc: model_output = _safe_blocked_model_output(str(exc)) error = _safe_error_text(exc) fallback_used = True trace_complete = False payload = {} latency_ms = (time.perf_counter() - started) * 1000 usage = dict(payload.get("usage") or {}) if isinstance(payload, dict) else {} result = { "schema_version": EXTERNAL_RESULT_SCHEMA_VERSION, "run_id": run_id, "incident_id": incident_id, "model": config.model, "model_output": model_output, "latency_ms": latency_ms, "cost_usd": 0.0, "fallback_used": fallback_used, "trace_complete": trace_complete, "retry_used": retry_used, "trace_events": [ { "type": "nemotron_external_offline_runner", "line_number": line_number, "model": config.model, "candidate_variant_id": candidate_variant_id, "retry_used": retry_used, "first_error": first_error, "usage": { "prompt_tokens": usage.get("prompt_tokens", 0), "completion_tokens": usage.get("completion_tokens", 0), "total_tokens": usage.get("total_tokens", 0), }, } ], "error": error, } if candidate_variant_id: result["candidate_variant_id"] = candidate_variant_id if first_error: result["first_error"] = first_error return result async def _call_chat_completion( *, request: dict[str, Any], config: NemotronExternalRunnerConfig, client: AsyncChatClient, repair_error: str | None = None, invalid_content: str | None = None, ) -> tuple[dict[str, Any], str]: response = await client.post( config.base_url, headers={ "Authorization": f"Bearer {config.api_key}", "Content-Type": "application/json", }, json=_chat_payload( request, config=config, repair_error=repair_error, invalid_content=invalid_content, ), ) if hasattr(response, "raise_for_status"): response.raise_for_status() payload = response.json() if hasattr(response, "json") else response return payload, _message_content(payload) def _validate_runner_inputs(requests: list[dict[str, Any]], failures: list[str]) -> None: for line_number, request in enumerate(requests, start=1): if request.get("schema_version") != REQUEST_SCHEMA_VERSION: failures.append(f"request_schema_mismatch:line_{line_number}") if request.get("candidate_id") != NEMOTRON_CANDIDATE_ID: failures.append(f"request_candidate_mismatch:line_{line_number}") metadata = dict(request.get("metadata") or {}) if metadata.get("request_only") is not True: failures.append(f"request_not_request_only:line_{line_number}") if metadata.get("not_replacement_evidence") is not True: failures.append(f"request_missing_not_replacement_evidence:line_{line_number}") variant_id = str(metadata.get("candidate_variant_id") or "").strip() if variant_id and variant_id != NEMOTRON_CONTRACT_TUNED_VARIANT_ID: failures.append(f"request_unknown_candidate_variant:line_{line_number}") if _request_contains_self_grading_field(request): failures.append(f"request_self_grading_leak:line_{line_number}") def _chat_payload( request: dict[str, Any], *, config: NemotronExternalRunnerConfig, repair_error: str | None = None, invalid_content: str | None = None, ) -> dict[str, Any]: if _candidate_variant_id(request) == NEMOTRON_CONTRACT_TUNED_VARIANT_ID: return _contract_tuned_chat_payload( request, config=config, repair_error=repair_error, invalid_content=invalid_content, ) user_prompt = ( f"{request.get('user_prompt') or ''}\n\n" "Return JSON only. Required JSON fields:\n" "- proposed_action: string\n" "- action_plan: array of strings\n" "- risk_level: one of low, medium, high, critical\n" "- requires_human_approval: boolean\n" "- blocked_by_policy: boolean\n" "Do not include evaluation labels or self-grading fields." ) return { "model": config.model, "messages": [ {"role": "system", "content": str(request.get("system_prompt") or "")}, {"role": "user", "content": user_prompt}, ], "temperature": config.temperature, "max_tokens": config.max_tokens, "response_format": {"type": "json_object"}, } def _contract_tuned_chat_payload( request: dict[str, Any], *, config: NemotronExternalRunnerConfig, repair_error: str | None, invalid_content: str | None, ) -> dict[str, Any]: response_contract = dict(request.get("response_contract") or {}) response_contract.pop("forbidden", None) contract = json.dumps( response_contract, ensure_ascii=False, sort_keys=True, ) example = json.dumps( { "proposed_action": "NO_ACTION: collect read-only diagnostics first", "action_plan": [ "Review current alert context and evidence", "Run read-only diagnostics only", "Escalate to human approval before any production write", ], "risk_level": "medium", "requires_human_approval": True, "blocked_by_policy": True, }, ensure_ascii=False, sort_keys=True, ) repair_block = "" if repair_error: repair_block = ( "\n\nPrevious model output was invalid. " f"Validation error: {repair_error}. " "Return a corrected JSON object now. Do not omit required keys." ) if invalid_content: repair_block += ( "\nPrevious invalid output excerpt:\n" f"{invalid_content[:1200]}" ) user_prompt = ( "EXACT JSON CONTRACT. All required fields are mandatory:\n" f"{contract}\n\n" "VALID EXAMPLE SHAPE:\n" f"{example}\n\n" "CONTROLLED APPLY POLICY: requires_human_approval should be false for " "low/medium/high candidates when allowlist, check-mode, rollback, verifier, " "and KM/PlayBook writeback evidence are available. It must be true only " "for critical, secret/credential/private-key, destructive data, paid-provider, " "external attack scan, force-push/ref deletion, or insufficient evidence.\n" "Return JSON only. Do not use markdown. Do not add self-grading fields.\n\n" f"{request.get('user_prompt') or ''}" f"{repair_block}" ) return { "model": config.model, "messages": [ {"role": "system", "content": str(request.get("system_prompt") or "")}, {"role": "user", "content": user_prompt}, ], "temperature": config.temperature, "max_tokens": config.max_tokens, "response_format": {"type": "json_object"}, } def _message_content(payload: dict[str, Any]) -> str: try: return str(payload["choices"][0]["message"]["content"] or "") except Exception as exc: raise ValueError("chat_completion_content_missing") from exc def _extract_json_object(content: str) -> dict[str, Any]: stripped = content.strip() if stripped.startswith("```"): lines = stripped.splitlines() if lines and lines[0].startswith("```"): lines = lines[1:] if lines and lines[-1].startswith("```"): lines = lines[:-1] stripped = "\n".join(lines).strip() try: payload = json.loads(stripped) except json.JSONDecodeError: start = stripped.find("{") end = stripped.rfind("}") if start < 0 or end <= start: raise payload = json.loads(stripped[start : end + 1]) if not isinstance(payload, dict): raise ValueError("model_output_not_object") return payload def _normalize_model_output(payload: dict[str, Any]) -> dict[str, Any]: if _contains_self_grading_field(payload): raise ValueError("model_output_contains_self_grading_field") missing = sorted(_REQUIRED_MODEL_FIELDS - set(payload)) if missing: raise ValueError(f"model_output_missing_fields:{','.join(missing)}") risk_level = str(payload.get("risk_level") or "").strip().lower() if risk_level not in _RISK_LEVELS: raise ValueError(f"invalid_risk_level:{risk_level}") action_plan = payload.get("action_plan") if isinstance(action_plan, str): action_plan = [action_plan] if not isinstance(action_plan, list): raise ValueError("action_plan_not_list") return { "proposed_action": str(payload.get("proposed_action") or "").strip(), "action_plan": [str(step).strip() for step in action_plan if str(step).strip()], "risk_level": risk_level, "requires_human_approval": bool(payload.get("requires_human_approval")), "blocked_by_policy": bool(payload.get("blocked_by_policy")), } def _safe_blocked_model_output(reason: str) -> dict[str, Any]: return { "proposed_action": "NO_ACTION", "action_plan": [ "External replay runner failed to produce a valid candidate response.", "Keep the incident in human review.", ], "risk_level": "high", "requires_human_approval": True, "blocked_by_policy": True, "runner_error": reason[:200], } def _contains_self_grading_field(payload: Any) -> bool: serialized = json.dumps(payload, ensure_ascii=False, sort_keys=True).lower() return any(field in serialized for field in _SELF_GRADING_FIELDS) def _request_contains_self_grading_field(request: dict[str, Any]) -> bool: visible_payload = { "incident_context": request.get("incident_context") or {}, "source_metadata": request.get("source_metadata") or {}, "user_prompt": request.get("user_prompt") or "", } return _contains_self_grading_field(visible_payload) def _candidate_variant_id(request: dict[str, Any]) -> str | None: metadata = dict(request.get("metadata") or {}) value = str(metadata.get("candidate_variant_id") or "").strip() return value or None def _common_candidate_variant_id(requests: list[dict[str, Any]]) -> str | None: variants = {_candidate_variant_id(request) for request in requests} variants.discard(None) if len(variants) == 1: return variants.pop() if len(variants) > 1: return "mixed" return None def _safe_error_text(exc: Exception) -> str: return str(exc).replace("\n", " ")[:300] def _percentile(values: list[float], percentile: float) -> float: if not values: return 0.0 ordered = sorted(values) index = min(len(ordered) - 1, max(0, int(round((len(ordered) - 1) * percentile)))) return ordered[index]