Senior AI Software Engineer - Interview Revision

from datetime import datetime, timezone
from typing import Literal
from uuid import UUID

from pydantic import BaseModel, ConfigDict, IPvAnyAddress, field_validator

class AlertContext(BaseModel):
    tenant_id: str
    source: str

class ThreatEvent(BaseModel):
    model_config = ConfigDict(frozen=True, extra="forbid")

    event_id: UUID
    severity: Literal["low", "medium", "high", "critical"]
    timestamp: datetime
    source_ip: IPvAnyAddress
    context: AlertContext

    @field_validator("timestamp")
    @classmethod
    def must_be_timezone_aware(cls, value: datetime) -> datetime:
        if value.tzinfo is None or value.tzinfo.utcoffset(value) is None:
            raise ValueError("timestamp must include timezone")
        return value.astimezone(timezone.utc)

from datetime import datetime
from enum import Enum

from pydantic import BaseModel, ConfigDict, Field

class OSType(str, Enum):
    linux = "linux"
    windows = "windows"
    macos = "macos"

class EndpointAgent(BaseModel):
    model_config = ConfigDict(frozen=True)
    agent_id: str
    hostname: str
    os_type: OSType
    threat_score: float = Field(ge=0, le=100)
    tags: list[str] = []
    last_seen: datetime

class AgentHeartbeat(BaseModel):
    model_config = ConfigDict(frozen=True)
    heartbeat_id: str
    status: str
    agent: EndpointAgent
    received_at: datetime

from typing import Literal

from pydantic import BaseModel
from pydantic_ai import Agent

class ThreatSummary(BaseModel):
    severity: Literal["low", "medium", "high", "critical"]
    affected_hosts: list[str]
    recommended_action: Literal[
        "monitor", "isolate-host", "rotate-credentials", "open-investigation"
    ]

triage_agent = Agent(
    "openai:gpt-4.1",
    result_type=ThreatSummary,
    system_prompt="Summarise threat telemetry for SOC analysts.",
)

result = await triage_agent.run("Summarise this endpoint incident")
summary = result.output

from contextlib import asynccontextmanager

import httpx
from fastapi import FastAPI

@asynccontextmanager
async def lifespan(app: FastAPI):
    async with httpx.AsyncClient(timeout=10.0) as client:
        app.state.http_client = client
        app.state.llm_client = build_llm_client(client)
        yield

app = FastAPI(lifespan=lifespan)

from fastapi.responses import StreamingResponse

async def stream_llm_tokens(prompt: str):
    async for chunk in llm_client.stream(prompt):
        yield f"data: {chunk.text}\n\n"

@router.get("/analysis/stream")
async def stream_analysis(prompt: str):
    return StreamingResponse(
        stream_llm_tokens(prompt),
        media_type="text/event-stream",
    )

async def fetch_signal(signal_id: str) -> dict:
    return await api_client.get(signal_id)

async def fetch_all(ids: list[str]) -> list[dict]:
    first = await fetch_signal(ids[0])
    rest = await asyncio.gather(*(fetch_signal(i) for i in ids[1:]))
    asyncio.create_task(write_audit_log(ids))
    return [first, *rest]

semaphore = asyncio.Semaphore(10)

async def score_event(event: ThreatEvent) -> ThreatSummary:
    async with semaphore:
        return await llm_gateway.score(event)

results = await asyncio.gather(*(score_event(event) for event in events))

from concurrent.futures import ThreadPoolExecutor

executor = ThreadPoolExecutor(max_workers=4)

async def scan_binary(binary: bytes) -> ScanResult:
    loop = asyncio.get_running_loop()
    return await loop.run_in_executor(
        executor,
        cpu_heavy_fn,
        binary,
    )

from fastapi import BackgroundTasks

@router.post("/alerts/{event_id}/notify")
async def notify(event_id: str, background_tasks: BackgroundTasks):
    event = await service.get_event(event_id)
    background_tasks.add_task(send_alert_webhook, event)
    return {"status": "queued", "event_id": event_id}

@asynccontextmanager
async def lifespan(app: FastAPI):
    async with httpx.AsyncClient() as http_client:
        redis = await redis.from_url(settings.redis_url)
        llm = build_llm_gateway(http_client)
        app.state.http_client = http_client
        app.state.redis = redis
        app.state.llm = llm
        yield
        await redis.close()

def get_llm_client(request: Request) -> LLMGateway:
    return request.app.state.llm

@router.post("/triage")
async def triage(
    payload: ThreatEventRequest,
    db: AsyncSession = Depends(get_db_session),
    user: User = Depends(get_current_user),
    llm: LLMGateway = Depends(get_llm_client),
):
    return await service.triage(payload, db, user, llm)

class ErrorResponse(BaseModel):
    code: str
    message: str
    trace_id: str

@app.exception_handler(ServiceUnavailableError)
async def service_unavailable_handler(request: Request, exc: ServiceUnavailableError):
    payload = ErrorResponse(
        code="upstream_unavailable",
        message=str(exc),
        trace_id=request.state.trace_id,
    )
    return JSONResponse(status_code=503, content=payload.model_dump())

raise HTTPException(status_code=404, detail={"code": "not_found", "entity": "event"})

trace_id_var: ContextVar[str] = ContextVar("trace_id")

@app.middleware("http")
async def inject_trace_id(request: Request, call_next):
    trace_id = request.headers.get("x-trace-id", str(uuid4()))
    request.state.trace_id = trace_id
    token = trace_id_var.set(trace_id)
    try:
        response = await call_next(request)
        response.headers["x-trace-id"] = trace_id
        return response
    finally:
        trace_id_var.reset(token)

@agent.tool
async def lookup_cve(cve_id: str) -> CVERecord:
    return await cve_client.fetch(cve_id)

@agent.tool
async def search_endpoint(hostname: str) -> EndpointRecord:
    return await graph_client.lookup_host(hostname)

result = await agent.run("Investigate suspicious PowerShell activity")

@tool
async def list_compromised_hosts(campaign_id: str) -> list[str]:
    return await singularity_client.list_hosts(campaign_id)

@tool
async def fetch_indicator(indicator: str) -> IndicatorRecord:
    return await singularity_client.get_indicator(indicator)