← All Go packages

github.com/anulum/director-ai/gateway/internal/risk

package risk // import "github.com/anulum/director-ai/gateway/internal/risk"

Package risk mirrors the Python PromptRiskScorer heuristic so the gateway can
refuse obvious attacks before any Python RPC is issued. The implementation
is deliberately identical to “director_ai/core/routing/scorer.py“'s length /
structural / marker heuristic so traffic gets the same verdict regardless of
which tier handles it.

The gateway never sees the sanitiser or injection signals — those require the
Python model to be loaded — so this Go path is heuristic-only. Callers that want
the full blend must reach back into the Python “PromptRiskScorer“ over gRPC;
this package covers the cheap first line of defence.

CONSTANTS

const (
	DefaultRulesThreshold  = 0.2
	DefaultEmbedThreshold  = 0.55
	DefaultRejectThreshold = 0.92
)
    Defaults mirror “director_ai.core.routing.RiskRouter“.


TYPES

type Budget struct {
	// Has unexported fields.
}
    Budget is a sliding-window risk budget. Safe for concurrent use.

func NewBudget(allowance, windowSeconds float64, clock func() time.Time) (*Budget, error)
    NewBudget builds a budget with “allowance“ per “window“. Callers pass “nil“
    for “clock“ to use “time.Now“.

func (b *Budget) AllowanceFor(tenantID string) float64
    AllowanceFor returns the effective allowance for “tenantID“.

func (b *Budget) Reserve(tenantID string, risk float64) BudgetEntry
    Reserve attempts to charge “risk“ (clamped to [0, 1]) against “tenantID“'s
    ledger. Pass zero to read the state without charging.

func (b *Budget) Reset(tenantID string)
    Reset clears one tenant's ledger. Passing the empty string clears every
    tenant.

func (b *Budget) SetAllowance(tenantID string, allowance float64) error
    SetAllowance overrides the allowance for a single tenant.

func (b *Budget) Snapshot(tenantID string) BudgetEntry
    Snapshot reads the current ledger without charging.

type BudgetEntry struct {
	TenantID      string
	WindowSeconds float64
	Allowance     float64
	Consumed      float64
	Remaining     float64
	Events        int
	Accepted      bool
}
    BudgetEntry mirrors the Python “BudgetEntry“ dataclass. “Accepted“ tells
    the caller whether the last reservation succeeded — without it, a non-zero
    “Remaining“ could reflect either "reservation applied" or "reservation
    refused, ledger untouched".

func (e BudgetEntry) Exhausted() bool
    Exhausted reports whether the caller should reject the request.

type Components struct {
	Heuristic float64
	Sanitiser float64
	Injection float64
	Combined  float64
}
    Components mirrors the Python “RiskComponents“ dataclass. Only the heuristic
    channel is populated here; “Sanitiser“ and “Injection“ stay at zero.

type Middleware struct {
	Scorer            *Scorer
	Budget            *Budget
	RulesThreshold    float64
	EmbedThreshold    float64
	RejectThreshold   float64
	TenantFromRequest func(*http.Request) string
}
    Middleware refuses obvious attack prompts and throttles tenants whose
    sliding-window risk budget is exhausted. The decision is stamped onto
    response headers so downstream auditors see why a request was accepted or
    rejected:

      - X-Risk-Score — combined risk in [0, 1]
      - X-Risk-Backend — chosen scorer backend (rules/embed/nli)
      - X-Risk-Action — allow / reject
      - X-Risk-Reason — human-readable reason
      - X-Risk-Remaining — remaining budget for the tenant

    The middleware extracts the prompt from the request body when the payload
    is a JSON object with a “messages“ array (OpenAI chat) or a “prompt“ field
    (legacy completion). If the body is not JSON or the prompt cannot be found,
    the middleware passes the request through untouched — failing-open preserves
    compatibility with clients that use an unexpected wire shape.

func NewMiddleware(scorer *Scorer, budget *Budget) *Middleware
    NewMiddleware returns a wired middleware with sane defaults. The tenant
    resolver falls back to “auth.FingerprintFromContext“ so a gateway with
    API-key auth binds risk to the key's fingerprint; unauthenticated
    deployments pass empty tenant IDs, which is still acceptable because the
    budget is per-key.

func (m *Middleware) Enabled() bool
    Handler wraps next with risk scoring. Enabled reports whether the middleware
    will do anything; callers that did not configure a scorer should skip the
    wrap entirely.

func (m *Middleware) Handler(next http.Handler) http.Handler
    Handler returns the net/http-compatible handler.

type Scorer struct {
	// Has unexported fields.
}
    Scorer produces a “[0, 1]“ risk score from a prompt. Thread safe — the regex
    slices are never mutated after construction.

func NewScorer() *Scorer
    NewScorer builds a Scorer with the defaults that match the Python side:
    length saturates at 8000 characters, weights sum to 1.

func NewScorerWithMaxLength(maxSafeLength int) (*Scorer, error)
    NewScorerWithMaxLength lets callers tune the length-saturation threshold;
    the weights stay at the Python defaults.

func (s *Scorer) Score(prompt string) Components
    Score returns a “Components“ record for “prompt“. Mirrors
    :meth:`PromptRiskScorer.score` — empty or whitespace-only prompts return all
    zeros.