100 lines
3.5 KiB
Python
100 lines
3.5 KiB
Python
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",
|
|
"code": payload.code,
|
|
"code_verifier": payload.code_verifier,
|
|
"redirect_uri": str(payload.redirect_uri),
|
|
"client_id": settings.oidc_client_id,
|
|
}
|
|
|
|
# Some OIDC providers expect client credentials via HTTP Basic auth
|
|
auth = None
|
|
if settings.oidc_client_secret:
|
|
auth = (settings.oidc_client_id, 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"}, auth=auth)
|
|
|
|
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")
|