Friday, October 30, 2026

LiteLLM Gateway – Custom LLM Integration with JWT Authentication

 1. Objective

Configure a LiteLLM Gateway to access a custom corporate LLM that requires two sequential operations:

  1. Authentication: Call an authentication endpoint using credentials to retrieve a JWT access token.

  2. Inference: Use the retrieved token as a Bearer token when calling an OpenAI-compatible LLM endpoint.

The solution should support token caching, automatic refresh, concurrent requests, error handling and multiple models.

The calling application should continue using the standard LangChain ChatOpenAI class without implementing the JWT authentication logic.

2. Target architecture

Recommended implementation: Deploy a lightweight JWT-aware adapter behind LiteLLM. LiteLLM handles public model aliases and routing; the adapter handles corporate authentication and forwards OpenAI-compatible requests.

This avoids modifying LiteLLM internals and allows the adapter to be tested independently.

3. Project structure

Use the following files:

corporate-llm-gateway/
├── docker-compose.yml
├── .env
├── litellm/
│   └── config.yaml
└── adapter/
    ├── Dockerfile
    ├── requirements.txt
    └── app.py

4. Implement the JWT-aware adapter

The adapter exposes an OpenAI-compatible /v1/chat/completions endpoint. It authenticates with the corporate identity service, caches the access token and forwards inference requests.

This example assumes OAuth 2.0 Client Credentials and a corporate API that already accepts OpenAI-compatible request and response bodies.

adapter/requirements.txt

fastapi
uvicorn[standard]
httpx

adapter/app.py

import asyncio        return token@app.get("/healthz")async def healthz():    return {"status": "ok"}@app.post("/v1/chat/completions")async def chat_completions(request: Request):    if request.headers.get("authorization") != f"Bearer {ADAPTER_KEY}":        raise HTTPException(status_code=401, detail="Unauthorized")    payload = await request.body()    client = request.app.state.http    try:        access_token = await get_token(client)        headers = {            "Authorization": f"Bearer {access_token}",            "Content-Type": "application/json",        }        # A streaming HTTP request also works for normal JSON responses.        upstream_request = client.build_request(            "POST",            CORPORATE_LLM_URL,            content=payload,            headers=headers,        )        upstream = await client.send(            upstream_request,            stream=True,        )        if upstream.status_code == 401:            await upstream.aclose()            access_token = await get_token(                client, force_refresh=True            )            headers["Authorization"] = f"Bearer {access_token}"            upstream_request = client.build_request(                "POST",                CORPORATE_LLM_URL,                content=payload,                headers=headers,            )            upstream = await client.send(                upstream_request,                stream=True,            )        if upstream.status_code >= 400:            status = upstream.status_code            await upstream.aclose()            # Avoid exposing corporate error bodies containing            # sensitive information.            return Response(                content='{"error":"Upstream LLM request failed"}',                status_code=status,                media_type="application/json",            )        content_type = upstream.headers.get(            "content-type", "application/json"        )

This is a functional reference implementation for local integration testing, not a fully hardened production adapter.

It supports normal chat responses and streaming response bodies. It also refreshes the token before expiry and retries once when the upstream rejects an expired token.

For production, add per-request correlation IDs, safe structured logging, metrics, controlled retry policies and explicit cancellation handling. If the authentication endpoint returns a JWT without expires_in, adapt expiry handling to the identity provider's documented contract rather than assuming the default lifetime.

5. Configure LiteLLM

LiteLLM can route a model to an OpenAI-compatible endpoint using the openai/ provider prefix and api_base.

liteLLM
+1

litellm/config.yaml

model_list:
  - model_name: corporate-llm
    litellm_params:
      model: openai/corporate-model
      api_base: http://jwt-adapter:8000/v1
      api_key: os.environ/ADAPTER_KEY

general_settings:
  master_key: os.environ/LITELLM_MASTER_KEY

There are two distinct credentials:

  • LITELLM_MASTER_KEY: Used by your LangChain application to authenticate with LiteLLM.

  • ADAPTER_KEY: Used by LiteLLM to authenticate with the internal adapter.

The adapter obtains the corporate JWT using CLIENT_ID and CLIENT_SECRET. Neither the JWT nor the corporate credentials need to be exposed to your LangChain application.

6. Docker Compose deployment

adapter/Dockerfile

FROM python:3.12-slim

WORKDIR /app

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY app.py .

CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]

docker-compose.yml

services:
  jwt-adapter:
    build: ./adapter
    environment:
      TOKEN_URL: ${TOKEN_URL}
      CLIENT_ID: ${CLIENT_ID}
      CLIENT_SECRET: ${CLIENT_SECRET}
      OAUTH_SCOPE: ${OAUTH_SCOPE:-}
      CORPORATE_LLM_URL: ${CORPORATE_LLM_URL}
      ADAPTER_KEY: ${ADAPTER_KEY}
    expose:
      - "8000"
    restart: unless-stopped

  litellm:
    image: docker.litellm.ai/berriai/litellm:latest
    ports:
      - "127.0.0.1:4000:4000"
    volumes:
      - ./litellm/config.yaml:/app/config.yaml:ro
    environment:
      LITELLM_MASTER_KEY: ${LITELLM_MASTER_KEY}
      ADAPTER_KEY: ${ADAPTER_KEY}
    command:
      - "--config"
      - "/app/config.yaml"
      - "--port"
      - "4000"
    depends_on:
      - jwt-adapter
    restart: unless-stopped

.env

TOKEN_URL=https://auth.company.com/oauth2/token
CLIENT_ID=your-client-id
CLIENT_SECRET=your-client-secret
OAUTH_SCOPE=llm.invoke

CORPORATE_LLM_URL=https://llm.company.com/v1/chat/completions

ADAPTER_KEY=replace-with-generated-internal-secret
LITELLM_MASTER_KEY=sk-replace-with-generated-gateway-key

Use independently generated high-entropy secrets for both keys. Keep .env out of version control and restrict its file permissions. In production, pin dependency versions and the LiteLLM container image.

Start the services:

docker compose up -d --build

The gateway is bound to VM loopback for local testing. If the application runs on another VM, expose LiteLLM through a secured HTTPS reverse proxy rather than publishing port 4000 openly.

7. Test using LangChain ChatOpenAI

Your application does not need to know about the JWT authentication process.

import osfrom langchain_openai import ChatOpenAIllm = ChatOpenAI(    model="corporate-llm",    base_url="http://localhost:4000/v1",    api_key=os.environ["LITELLM_MASTER_KEY"],    temperature=0.2,    timeout=120,)response = llm.invoke(    "Explain why a BGP session can remain in Idle state.")print(response.content)

For streaming:

for chunk in llm.stream(    "Explain Kubernetes pod networking."):    print(chunk.content, end="", flush=True)

Request processing sequence

  1. LangChain sends a request to LiteLLM using its gateway key.

  2. LiteLLM resolves the corporate-llm model alias and forwards the request to the JWT adapter.

  3. The adapter checks its token cache. If necessary, it calls the corporate authentication endpoint.

  4. The adapter sends the OpenAI-compatible request with Authorization: Bearer <JWT>.

  5. The corporate LLM returns its response through the adapter and LiteLLM to LangChain.

Subsequent requests reuse the cached token until refresh is required.

8. Production considerations

Area

Recommended implementation

Token management

Cache and refresh before expiry

Concurrent requests

Prevent simultaneous refreshes

Authentication failures

Refresh once on upstream HTTP 401

Credentials

Kubernetes Secrets or external secret manager

Network security

HTTPS, internal-only adapter access

High availability

Token cache per replica or shared token manager

Streaming

Preserve SSE chunks and cancellation

Observability

Log latency, errors and token refresh events without exposing credentials

Model support

Verify tool calling, structured output and streaming

Model routing

Stable LiteLLM aliases mapped to corporate model IDs

Implementation detail: The example forwards the JSON request unchanged. Therefore, the corporate API must accept the model ID corporate-model and the standard OpenAI chat-completion payload. If the corporate API uses a different model ID or expects additional headers, implement that translation in the adapter.

Also, the example covers chat completions only. Embeddings, the Responses API and other endpoints require separate adapter routes.

9. Instructions for your coding agent

Task: Implement JWT-authenticated Corporate LLM Integration with LiteLLM

Implement a reusable integration that enables LiteLLM Gateway to route inference requests to a corporate OpenAI-compatible LLM endpoint requiring JWT authentication.

Architecture

LangChain / LangGraph → LiteLLM Proxy → JWT Authentication Adapter → Corporate LLM API.

Functional requirements

  1. Configure LiteLLM with a custom model alias and OpenAI-compatible upstream endpoint.

  2. Implement a Python FastAPI adapter that obtains JWT access tokens from a configurable authentication endpoint.

  3. Support OAuth 2.0 Client Credentials authentication, with configurable token URL, credentials and scope.

  4. Cache access tokens and refresh them before expiry.

  5. Prevent duplicate token refresh requests under concurrent workloads.

  6. Inject the access token into upstream requests using the Bearer authorization header.

  7. Support standard OpenAI chat-completion requests, including streaming.

  8. Retry authentication once when an upstream request receives HTTP 401.

  9. Propagate appropriate HTTP errors without exposing credentials or sensitive upstream responses.

  10. Support multiple model aliases and configurable upstream model mappings.

  11. Keep all URLs, model IDs, credentials and timeouts environment-driven.

  12. Provide Docker Compose deployment for local development and Helm/Kubernetes deployment for production.

  13. Add health checks, structured logging, timeout controls and metrics.

  14. Implement automated tests covering token caching, expiry, refresh, concurrent requests, upstream errors, normal inference and SSE streaming.

  15. Document the integration using LangChain ChatOpenAI, including synchronous and asynchronous calls.

Security requirements

  • Never expose JWTs, client secrets or gateway credentials in application logs.

  • Use TLS verification for corporate authentication and inference endpoints.

  • Restrict direct access to the internal adapter.

  • Store credentials using Kubernetes Secrets or a dedicated secrets manager.

  • Avoid forwarding incoming gateway authorization headers to the corporate endpoint.

Acceptance criteria

The existing LangChain/LangGraph application must be able to invoke the corporate LLM using only model, base_url and api_key, without requiring any application-level JWT handling.

Use version-pinned dependencies and provide an executable local integration test with mocked authentication and LLM endpoints.

Recommendation: For your architecture, this adapter-based approach provides a clean separation between model routing and corporate authentication. It also allows the same LiteLLM instance to serve Ollama, standard cloud models and JWT-protected corporate models.

Reference: LiteLLM official documentation.

No comments:

Post a Comment