> ## Documentation Index
> Fetch the complete documentation index at: https://docs.cognisafe.uk/llms.txt
> Use this file to discover all available pages before exploring further.

# Azure API Management Integration

> Insert Cognisafe behind Azure APIM for enterprise AI gateway patterns

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

<Steps>
  <Step title="Provision APIM">
    For internal enterprise use, **Standard** tier supports VNet integration. **Developer** tier is suitable for non-production.

    ```bash theme={null}
    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>
    ```

    <Note>
      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.
    </Note>
  </Step>

  <Step title="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.

    ```bash theme={null}
    # 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`
  </Step>

  <Step title="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.

    ```bash theme={null}
    # 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**.
  </Step>

  <Step title="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.

    ```bash theme={null}
    # 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>
</Steps>

## 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.

```xml theme={null}
<!-- 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>
```

<Warning>
  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`.
</Warning>

## 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:

```xml theme={null}
<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

```bash theme={null}
# 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

```xml theme={null}
<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

```bash theme={null}
# 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

```bash theme={null}
# 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:

```kusto theme={null}
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

<Steps>
  <Step title="Obtain an Azure AD token">
    ```bash theme={null}
    TOKEN=$(az account get-access-token \
      --resource api://<your-app-registration-client-id> \
      --query accessToken -o tsv)
    ```
  </Step>

  <Step title="Call the APIM endpoint">
    ```bash theme={null}
    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 .
    ```
  </Step>

  <Step title="Verify Cognisafe received the request">
    ```bash theme={null}
    # 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())
    "
    ```
  </Step>
</Steps>

## 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        |