Azure API Management Integration

Azure API Management (APIM) and Cognisafe solve different layers of the AI gateway problem. APIM owns cross-cutting enterprise concerns: subscription management, OAuth2/Azure AD validation, rate limiting by business tier, and unified API exposure across teams. Cognisafe owns AI-specific runtime security: per-request OWASP LLM scoring, jailbreak detection, PII scanning, cost attribution, and the full audit trail of every LLM exchange. Neither replaces the other. Together they form a complete enterprise AI gateway.

Architecture

Client application
    │
    │  Authorization: Bearer <AAD token>  (or Ocp-Apim-Subscription-Key)
    ▼
Azure API Management
    │  - Validates AAD token (validate-jwt policy)
    │  - Enforces rate limit (quota-by-key policy)
    │  - Injects X-Cognisafe-Project-ID header
    │  - Injects Authorization: Bearer <PROXY_API_KEY> (named value)
    │  - Routes to Cognisafe proxy backend
    ▼
Cognisafe Proxy  (Go, :8080 — ClusterIP inside AKS)
    │  - Authenticates with PROXY_API_KEY
    │  - Forwards request to upstream LLM
    │  - POSTs /internal/log (non-blocking) to Cognisafe API
    ▼
LLM Provider  (OpenAI / Azure OpenAI / Anthropic)
    │
    ◀──── response flows back through proxy → APIM → client

APIM never touches the LLM response payload — it passes through unchanged. Cognisafe intercepts and scores asynchronously, so neither APIM nor the client sees added latency from scoring.

Step 1: Create the APIM instance

Provision APIM

For internal enterprise use, Standard tier supports VNet integration. Developer tier is suitable for non-production.

az apim create \
  --resource-group <rg> \
  --name cognisafe-apim \
  --publisher-email platform@yourcompany.com \
  --publisher-name "YourCompany Platform" \
  --sku-name Standard \
  --location eastus

# Enable VNet integration (for private AKS backend)
az apim update \
  --resource-group <rg> \
  --name cognisafe-apim \
  --set virtualNetworkType=External \
  --virtual-network /subscriptions/<sub>/resourceGroups/<rg>/providers/Microsoft.Network/virtualNetworks/<vnet>

If your AKS cluster uses an internal load balancer (no public IP on the proxy Service), set virtualNetworkType=Internal and deploy APIM into the same VNet as AKS. The Cognisafe proxy backend URL will then be the internal ClusterIP or internal load balancer DNS name.

Import the OpenAI API spec

APIM can import the OpenAI API OpenAPI spec directly. This creates operations for /v1/chat/completions, /v1/completions, /v1/embeddings, etc.

# Download the OpenAI OpenAPI spec
curl -o openai-openapi.json \
  https://raw.githubusercontent.com/openai/openai-openapi/master/openapi.yaml

# Import into APIM
az apim api import \
  --resource-group <rg> \
  --service-name cognisafe-apim \
  --api-id cognisafe-llm \
  --display-name "Cognisafe LLM Gateway" \
  --path /llm \
  --specification-format OpenApi \
  --specification-path openai-openapi.json \
  --protocols https \
  --subscription-required true

After import, the effective endpoint clients call is: https://cognisafe-apim.azure-api.net/llm/v1/chat/completions

Set the backend URL

Point the APIM API backend at the Cognisafe proxy. If the proxy runs inside AKS with a ClusterIP service, you need either an internal load balancer IP or the AKS internal DNS name.

# If using AKS internal load balancer
PROXY_INTERNAL_IP=$(kubectl get service proxy \
  -n cognisafe-system \
  -o jsonpath='{.status.loadBalancer.ingress[0].ip}')

az apim api update \
  --resource-group <rg> \
  --service-name cognisafe-apim \
  --api-id cognisafe-llm \
  --set serviceUrl="http://${PROXY_INTERNAL_IP}:8080"

Alternatively, set the backend via the APIM portal: APIs → Cognisafe LLM Gateway → Settings → Web service URL.

Create Named Values for secrets

Named values store secrets in APIM without embedding them in policy XML. APIM can pull values from Azure Key Vault directly.

# Link APIM's managed identity to Key Vault
az apim update \
  --resource-group <rg> \
  --name cognisafe-apim \
  --enable-managed-identity true

APIM_IDENTITY=$(az apim show \
  --resource-group <rg> \
  --name cognisafe-apim \
  --query identity.principalId -o tsv)

az keyvault set-policy \
  --name <your-keyvault-name> \
  --object-id "$APIM_IDENTITY" \
  --secret-permissions get list

# Create Named Value backed by Key Vault
az apim nv create \
  --resource-group <rg> \
  --service-name cognisafe-apim \
  --named-value-id CognisafeProxyApiKey \
  --display-name "CognisafeProxyApiKey" \
  --secret true \
  --value-type keyVaultSecretReference \
  --key-vault-secret-identifier "https://<your-keyvault-name>.vault.azure.net/secrets/cognisafe-proxy-api-key"

Step 2: APIM inbound policy

This policy runs on every request before it reaches the Cognisafe proxy. It validates the Azure AD token, enforces rate limits, and rewrites headers.

<!-- Applied at API level: cognisafe-llm -->
<policies>
  <inbound>
    <base />

    <!-- 1. Validate Azure AD (Entra ID) Bearer token -->
    <validate-jwt header-name="Authorization" failed-validation-httpcode="401" failed-validation-error-message="Unauthorized — valid Azure AD token required">
      <openid-config url="https://login.microsoftonline.com/<tenant-id>/v2.0/.well-known/openid-configuration" />
      <audiences>
        <audience>api://<your-app-registration-client-id></audience>
      </audiences>
      <issuers>
        <issuer>https://sts.windows.net/<tenant-id>/</issuer>
        <issuer>https://login.microsoftonline.com/<tenant-id>/v2.0</issuer>
      </issuers>
      <required-claims>
        <claim name="roles" match="any">
          <value>LLMGateway.User</value>
          <value>LLMGateway.Admin</value>
        </claim>
      </required-claims>
    </validate-jwt>

    <!-- 2. Rate limiting: 1000 calls/hour per subscription key -->
    <rate-limit-by-key calls="1000" renewal-period="3600"
      counter-key="@(context.Subscription.Id)"
      increment-condition="@(context.Response.StatusCode >= 200 && context.Response.StatusCode < 300)"
      retry-after-header-name="Retry-After" />

    <!-- 3. Extract project ID from JWT claim and forward to Cognisafe -->
    <set-header name="X-Cognisafe-Project-ID" exists-action="override">
      <value>@{
        var jwt = context.Request.Headers.GetValueOrDefault("Authorization","").Split(' ').LastOrDefault();
        if (jwt != null) {
          var decoded = new System.IdentityModel.Tokens.Jwt.JwtSecurityToken(jwt);
          return decoded.Claims.FirstOrDefault(c => c.Type == "extension_ProjectID")?.Value ?? "default";
        }
        return "default";
      }</value>
    </set-header>

    <!-- 4. Replace client Authorization header with Cognisafe proxy API key -->
    <set-header name="Authorization" exists-action="override">
      <value>@("Bearer " + context.Variables["CognisafeProxyApiKey"])</value>
    </set-header>

    <!-- 5. Rewrite path: strip /llm prefix so proxy sees /v1/chat/completions -->
    <rewrite-uri template="@(context.Request.Url.Path.Replace("/llm", ""))" copy-unmatched-params="true" />

    <!-- 6. Forward client IP for audit logging -->
    <set-header name="X-Forwarded-For" exists-action="override">
      <value>@(context.Request.IpAddress)</value>
    </set-header>

    <!-- 7. Enforce HTTPS-only downstream -->
    <choose>
      <when condition="@(context.Request.Url.Scheme != "https")">
        <return-response>
          <set-status code="400" reason="HTTPS Required" />
          <set-body>{"error": "This endpoint requires HTTPS"}</set-body>
        </return-response>
      </when>
    </choose>
  </inbound>

  <backend>
    <base />
  </backend>

  <outbound>
    <!-- Pass through the proxy response unchanged. Do not modify LLM response bodies. -->
    <base />
    <!-- Expose rate limit headers to the client -->
    <set-header name="X-RateLimit-Remaining" exists-action="override">
      <value>@(context.Response.Headers.GetValueOrDefault("X-RateLimit-Remaining", ""))</value>
    </set-header>
  </outbound>

  <on-error>
    <base />
    <set-header name="Content-Type" exists-action="override">
      <value>application/json</value>
    </set-header>
    <set-body>@{
      return new JObject(
        new JProperty("error", new JObject(
          new JProperty("message", context.LastError.Message),
          new JProperty("source", context.LastError.Source),
          new JProperty("requestId", context.RequestId)
        ))
      ).ToString();
    }</set-body>
  </on-error>
</policies>

The validate-jwt policy caches the OIDC discovery document and JWKS keys. If your Azure AD app registration rotates signing keys (standard behaviour), APIM will auto-refresh within the configured TTL. Do not pin to specific key IDs — let the policy fetch from the OIDC jwks_uri.

Step 3: Per-IP rate limiting (secondary policy)

Apply this at the operation level on POST /v1/chat/completions if you want IP-level throttling in addition to per-subscription limits:

<inbound>
  <base />
  <rate-limit-by-key calls="60" renewal-period="60"
    counter-key="@(context.Request.IpAddress)"
    retry-after-header-name="Retry-After"
    retry-after-variable-name="retryAfterSeconds" />
</inbound>

Step 4: mTLS between APIM and Cognisafe proxy

For defence-in-depth, configure APIM to present a client certificate when calling the Cognisafe proxy backend. The proxy then validates that the certificate is from APIM before accepting the request.

Upload the client certificate to APIM

# Generate a self-signed cert for dev/staging; use your PKI for production
openssl req -x509 -newkey rsa:4096 -keyout apim-client.key -out apim-client.crt \
  -days 365 -nodes -subj "/CN=apim-cognisafe-client"

openssl pkcs12 -export -out apim-client.pfx \
  -inkey apim-client.key -in apim-client.crt \
  -passout pass:changeme

az apim certificate create \
  --resource-group <rg> \
  --service-name cognisafe-apim \
  --certificate-id apim-client-cert \
  --data @apim-client.pfx \
  --password changeme

Reference the certificate in the backend policy

<backend>
  <base />
  <forward-request timeout="120" follow-redirects="false">
    <proxy-certificate certificate-id="apim-client-cert" />
  </forward-request>
</backend>

Configure the Cognisafe proxy to require client certs

In your AKS Ingress (or directly in the Go proxy), set ssl_verify_client on (NGINX) or implement tls.RequireAndVerifyClientCert in the Go TLS config. Add the APIM client certificate’s CA to the proxy’s trusted CA pool.

Step 5: OAuth2 product and subscription setup

# Create a product for LLM access
az apim product create \
  --resource-group <rg> \
  --service-name cognisafe-apim \
  --product-id llm-gateway \
  --display-name "LLM Gateway" \
  --description "AI gateway via Cognisafe — OWASP-scored, audited" \
  --subscription-required true \
  --approval-required true \
  --subscriptions-limit 1 \
  --state published

# Link the Cognisafe LLM API to this product
az apim product api add \
  --resource-group <rg> \
  --service-name cognisafe-apim \
  --product-id llm-gateway \
  --api-id cognisafe-llm

Step 6: Send APIM access logs to Azure Monitor

# Create a diagnostic setting on APIM → Log Analytics
az monitor diagnostic-settings create \
  --resource /subscriptions/<sub>/resourceGroups/<rg>/providers/Microsoft.ApiManagement/service/cognisafe-apim \
  --name apim-diag \
  --workspace /subscriptions/<sub>/resourceGroups/<rg>/providers/microsoft.operationalinsights/workspaces/<workspace> \
  --logs '[{"category":"GatewayLogs","enabled":true,"retentionPolicy":{"days":90,"enabled":true}}]' \
  --metrics '[{"category":"AllMetrics","enabled":true}]'

Query APIM gateway logs in Log Analytics:

ApiManagementGatewayLogs
| where TimeGenerated > ago(1h)
| where ApiId == "cognisafe-llm"
| project TimeGenerated, OperationName, ResponseCode, DurationMs,
          ClientProtocol, BackendUrl, ClientIp,
          RequestHeaders["X-Cognisafe-Project-ID"]
| order by TimeGenerated desc

Step 7: End-to-end test

Obtain an Azure AD token

TOKEN=$(az account get-access-token \
  --resource api://<your-app-registration-client-id> \
  --query accessToken -o tsv)

Call the APIM endpoint

curl -s -X POST \
  "https://cognisafe-apim.azure-api.net/llm/v1/chat/completions" \
  -H "Authorization: Bearer $TOKEN" \
  -H "Ocp-Apim-Subscription-Key: <your-subscription-key>" \
  -H "Content-Type: application/json" \
  -d '{
    "model": "gpt-4o-mini",
    "messages": [{"role": "user", "content": "Hello"}],
    "max_tokens": 50
  }' | jq .

Verify Cognisafe received the request

# Check the API logs
kubectl logs -n cognisafe-system -l app=api --tail=20

# Or query the database directly
kubectl exec -n cognisafe-system \
  $(kubectl get pod -n cognisafe-system -l app=api -o name | head -1) \
  -- python -c "
import asyncio, asyncpg, os
async def run():
    conn = await asyncpg.connect(os.environ['POSTGRES_URL'])
    rows = await conn.fetch('SELECT id, model, created_at FROM llm_requests ORDER BY created_at DESC LIMIT 5')
    for r in rows: print(dict(r))
    await conn.close()
asyncio.run(run())
"

Troubleshooting

Symptom	Likely cause	Fix
401 from APIM	JWT validation failure	Check token audience matches `api://<client-id>`; verify `roles` claim includes `LLMGateway.User`
401 from Cognisafe proxy	Wrong `PROXY_API_KEY` in Named Value	Re-sync Named Value from Key Vault; check `Authorization` header in APIM trace
429 from APIM	Rate limit hit	Review `rate-limit-by-key` policy; check `Retry-After` header
502 Bad Gateway	APIM cannot reach proxy backend	Verify VNet peering / internal LB IP; check NSG allows APIM subnet → AKS node subnet on port 8080
Requests logged in APIM but not in Cognisafe	Path rewrite mismatch	Enable APIM tracing (`Ocp-Apim-Trace: true`) and inspect the rewritten URL sent to backend

Getting Started

SDKs

LLM Providers

Self-hosted Models

Self-hosting

Integrations

Reference Architectures

Deployment Blueprints

Safety & Scoring

Azure API Management Integration

Architecture

Step 1: Create the APIM instance

Step 2: APIM inbound policy

Step 3: Per-IP rate limiting (secondary policy)

Step 4: mTLS between APIM and Cognisafe proxy

Upload the client certificate to APIM

Reference the certificate in the backend policy

Configure the Cognisafe proxy to require client certs

Step 5: OAuth2 product and subscription setup

Step 6: Send APIM access logs to Azure Monitor

Step 7: End-to-end test

Troubleshooting

​Architecture

​Step 1: Create the APIM instance

​Step 2: APIM inbound policy

​Step 3: Per-IP rate limiting (secondary policy)

​Step 4: mTLS between APIM and Cognisafe proxy

​Upload the client certificate to APIM

​Reference the certificate in the backend policy

​Configure the Cognisafe proxy to require client certs

​Step 5: OAuth2 product and subscription setup

​Step 6: Send APIM access logs to Azure Monitor

​Step 7: End-to-end test

​Troubleshooting

Architecture

Step 1: Create the APIM instance

Step 2: APIM inbound policy

Step 3: Per-IP rate limiting (secondary policy)

Step 4: mTLS between APIM and Cognisafe proxy

Upload the client certificate to APIM

Reference the certificate in the backend policy

Configure the Cognisafe proxy to require client certs

Step 5: OAuth2 product and subscription setup

Step 6: Send APIM access logs to Azure Monitor

Step 7: End-to-end test

Troubleshooting