init
This commit is contained in:
@@ -0,0 +1,6 @@
|
||||
FROM python:3.12-slim
|
||||
WORKDIR /app
|
||||
COPY requirements.txt .
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
COPY app ./app
|
||||
CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000"]
|
||||
@@ -0,0 +1,95 @@
|
||||
from functools import lru_cache
|
||||
from typing import Any, Dict, Optional
|
||||
from urllib.parse import urljoin
|
||||
|
||||
import httpx
|
||||
from fastapi import HTTPException, Security
|
||||
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
|
||||
from jose import JWTError, jwk, jwt
|
||||
from pydantic import BaseModel, BaseSettings, Field, HttpUrl
|
||||
|
||||
security = HTTPBearer(auto_error=False)
|
||||
|
||||
class Settings(BaseSettings):
|
||||
oidc_issuer: HttpUrl = Field(..., env="OIDC_ISSUER")
|
||||
oidc_audience: str = Field(..., env="OIDC_AUDIENCE")
|
||||
oidc_client_id: str = Field(..., env="OIDC_CLIENT_ID")
|
||||
oidc_client_secret: Optional[str] = Field(None, env="OIDC_CLIENT_SECRET")
|
||||
oidc_verify_iss: bool = Field(True, env="OIDC_VERIFY_ISS")
|
||||
|
||||
class Config:
|
||||
env_file = ".env"
|
||||
|
||||
settings = Settings()
|
||||
|
||||
class TokenClaims(BaseModel):
|
||||
iss: str
|
||||
sub: str
|
||||
aud: Any
|
||||
exp: int
|
||||
iat: Optional[int]
|
||||
azp: Optional[str]
|
||||
email: Optional[str]
|
||||
preferred_username: Optional[str]
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def get_discovery() -> Dict[str, Any]:
|
||||
issuer = str(settings.oidc_issuer).rstrip("/")
|
||||
discovery_url = urljoin(issuer + "/", ".well-known/openid-configuration")
|
||||
with httpx.Client(timeout=10.0) as client:
|
||||
response = client.get(discovery_url)
|
||||
if response.status_code != 200:
|
||||
raise RuntimeError(f"Unable to load OIDC discovery document from {discovery_url}")
|
||||
return response.json()
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def get_jwks() -> Dict[str, Any]:
|
||||
discovery = get_discovery()
|
||||
jwks_uri = discovery.get("jwks_uri")
|
||||
if not jwks_uri:
|
||||
raise RuntimeError("OIDC discovery document does not contain jwks_uri")
|
||||
with httpx.Client(timeout=10.0) as client:
|
||||
response = client.get(jwks_uri)
|
||||
if response.status_code != 200:
|
||||
raise RuntimeError(f"Unable to load JWKS from {jwks_uri}")
|
||||
return response.json()
|
||||
|
||||
def get_signing_key(token: str) -> Dict[str, Any]:
|
||||
try:
|
||||
headers = jwt.get_unverified_header(token)
|
||||
except JWTError as exc:
|
||||
raise HTTPException(status_code=401, detail="Invalid JWT header") from exc
|
||||
|
||||
kid = headers.get("kid")
|
||||
if not kid:
|
||||
raise HTTPException(status_code=401, detail="JWT missing kid header")
|
||||
|
||||
jwks = get_jwks().get("keys", [])
|
||||
for key in jwks:
|
||||
if key.get("kid") == kid:
|
||||
return key
|
||||
|
||||
raise HTTPException(status_code=401, detail="Unable to find matching JWKS key")
|
||||
|
||||
async def validate_token(credentials: HTTPAuthorizationCredentials = Security(security)) -> TokenClaims:
|
||||
if not credentials or credentials.scheme.lower() != "bearer":
|
||||
raise HTTPException(status_code=401, detail="Missing Bearer authorization header")
|
||||
|
||||
token = credentials.credentials
|
||||
key = get_signing_key(token)
|
||||
public_key = jwk.construct(key)
|
||||
try:
|
||||
verified = jwt.decode(
|
||||
token,
|
||||
public_key,
|
||||
algorithms=[key.get("alg", "RS256")],
|
||||
audience=settings.oidc_audience,
|
||||
issuer=str(settings.oidc_issuer) if settings.oidc_verify_iss else None,
|
||||
)
|
||||
except JWTError as exc:
|
||||
raise HTTPException(status_code=401, detail="Invalid or expired token") from exc
|
||||
|
||||
try:
|
||||
return TokenClaims(**verified)
|
||||
except Exception as exc:
|
||||
raise HTTPException(status_code=401, detail="Unable to parse token claims") from exc
|
||||
@@ -0,0 +1,96 @@
|
||||
import os
|
||||
from typing import AsyncGenerator
|
||||
|
||||
import httpx
|
||||
from fastapi import Depends, FastAPI, HTTPException
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
from app.auth import get_discovery, settings, validate_token
|
||||
from app.schemas import AuthCallbackRequest, AuthConfigResponse, ChatRequest
|
||||
|
||||
OLLAMA_API_KEY = os.getenv("OLLAMA_API_KEY")
|
||||
OLLAMA_API_URL = os.getenv("OLLAMA_API_URL", "https://api.ollama.com/v1/chat/completions")
|
||||
if not OLLAMA_API_KEY:
|
||||
raise RuntimeError("Environment variable OLLAMA_API_KEY is required")
|
||||
|
||||
app = FastAPI(
|
||||
title="OIDC-Protected Ollama Proxy",
|
||||
description="FastAPI backend validating OIDC JWTs and streaming Ollama Cloud responses over SSE.",
|
||||
version="0.1.0",
|
||||
)
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
async def event_stream(request: ChatRequest) -> AsyncGenerator[str, None]:
|
||||
payload = {
|
||||
"model": request.model,
|
||||
"messages": [message.dict() for message in request.messages],
|
||||
"stream": True,
|
||||
}
|
||||
headers = {
|
||||
"Authorization": f"Bearer {OLLAMA_API_KEY}",
|
||||
"Content-Type": "application/json",
|
||||
"Accept": "text/event-stream",
|
||||
}
|
||||
|
||||
async with httpx.AsyncClient(timeout=None) as client:
|
||||
async with client.stream("POST", OLLAMA_API_URL, json=payload, headers=headers) as response:
|
||||
if response.status_code >= 400:
|
||||
text = await response.aread()
|
||||
raise HTTPException(status_code=502, detail=f"Ollama Cloud error: {text.decode('utf-8', 'ignore')}" )
|
||||
|
||||
async for chunk in response.aiter_text():
|
||||
if not chunk:
|
||||
continue
|
||||
for line in chunk.splitlines():
|
||||
if line.strip():
|
||||
yield f"data: {line}\n\n"
|
||||
|
||||
@app.get("/health")
|
||||
async def health() -> dict:
|
||||
return {"status": "ok"}
|
||||
|
||||
@app.get("/auth/config", response_model=AuthConfigResponse)
|
||||
async def auth_config() -> AuthConfigResponse:
|
||||
discovery = get_discovery()
|
||||
return AuthConfigResponse(
|
||||
issuer=discovery["issuer"],
|
||||
authorization_endpoint=discovery["authorization_endpoint"],
|
||||
token_endpoint=discovery["token_endpoint"],
|
||||
)
|
||||
|
||||
@app.post("/auth/callback")
|
||||
async def auth_callback(payload: AuthCallbackRequest):
|
||||
discovery = get_discovery()
|
||||
token_endpoint = discovery.get("token_endpoint")
|
||||
if not token_endpoint:
|
||||
raise HTTPException(status_code=502, detail="OIDC discovery document missing token_endpoint")
|
||||
|
||||
data = {
|
||||
"grant_type": "authorization_code",
|
||||
"client_id": settings.oidc_client_id,
|
||||
"code": payload.code,
|
||||
"code_verifier": payload.code_verifier,
|
||||
"redirect_uri": str(payload.redirect_uri),
|
||||
}
|
||||
if settings.oidc_client_secret:
|
||||
data["client_secret"] = settings.oidc_client_secret
|
||||
|
||||
async with httpx.AsyncClient(timeout=15.0) as client:
|
||||
response = await client.post(token_endpoint, data=data, headers={"Accept": "application/json"})
|
||||
|
||||
if response.status_code != 200:
|
||||
raise HTTPException(status_code=400, detail=f"OIDC token exchange failed: {response.text}")
|
||||
|
||||
return response.json()
|
||||
|
||||
@app.post("/chat", response_class=StreamingResponse)
|
||||
async def chat(request: ChatRequest, token=Depends(validate_token)):
|
||||
return StreamingResponse(event_stream(request), media_type="text/event-stream")
|
||||
@@ -0,0 +1,20 @@
|
||||
from pydantic import BaseModel, Field, HttpUrl
|
||||
from typing import Literal, List
|
||||
|
||||
class ChatMessage(BaseModel):
|
||||
role: Literal["user", "assistant", "system"]
|
||||
content: str
|
||||
|
||||
class ChatRequest(BaseModel):
|
||||
model: str = Field(default="llama2")
|
||||
messages: List[ChatMessage] = Field(default_factory=list)
|
||||
|
||||
class AuthCallbackRequest(BaseModel):
|
||||
code: str
|
||||
code_verifier: str
|
||||
redirect_uri: HttpUrl
|
||||
|
||||
class AuthConfigResponse(BaseModel):
|
||||
issuer: HttpUrl
|
||||
authorization_endpoint: HttpUrl
|
||||
token_endpoint: HttpUrl
|
||||
@@ -0,0 +1,5 @@
|
||||
fastapi==0.109.0
|
||||
uvicorn[standard]==0.23.2
|
||||
httpx==0.28.4
|
||||
python-jose[cryptography]==3.3.0
|
||||
pydantic==2.8.0
|
||||
Reference in New Issue
Block a user