1. Objective
Configure a LiteLLM Gateway to access a custom corporate LLM that requires two sequential operations:
Authentication: Call an authentication endpoint using credentials to retrieve a JWT access token.
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.py4. 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]
httpxadapter/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/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_KEYThere 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-keyUse 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 --buildThe 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
LangChain sends a request to LiteLLM using its gateway key.
LiteLLM resolves the
corporate-llmmodel alias and forwards the request to the JWT adapter.The adapter checks its token cache. If necessary, it calls the corporate authentication endpoint.
The adapter sends the OpenAI-compatible request with
Authorization: Bearer <JWT>.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
Configure LiteLLM with a custom model alias and OpenAI-compatible upstream endpoint.
Implement a Python FastAPI adapter that obtains JWT access tokens from a configurable authentication endpoint.
Support OAuth 2.0 Client Credentials authentication, with configurable token URL, credentials and scope.
Cache access tokens and refresh them before expiry.
Prevent duplicate token refresh requests under concurrent workloads.
Inject the access token into upstream requests using the Bearer authorization header.
Support standard OpenAI chat-completion requests, including streaming.
Retry authentication once when an upstream request receives HTTP 401.
Propagate appropriate HTTP errors without exposing credentials or sensitive upstream responses.
Support multiple model aliases and configurable upstream model mappings.
Keep all URLs, model IDs, credentials and timeouts environment-driven.
Provide Docker Compose deployment for local development and Helm/Kubernetes deployment for production.
Add health checks, structured logging, timeout controls and metrics.
Implement automated tests covering token caching, expiry, refresh, concurrent requests, upstream errors, normal inference and SSE streaming.
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