import json import asyncio import shutil from datetime import datetime, timedelta from typing import Optional import httpx from fastapi import FastAPI, HTTPException, Depends, Request from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials from fastapi.staticfiles import StaticFiles from fastapi.responses import HTMLResponse, StreamingResponse from fastapi.templating import Jinja2Templates from jose import JWTError, jwt from pydantic import BaseModel from config import settings from database import ( init_db, create_user, get_user_by_username, verify_password, get_user_by_api_key, create_api_key, get_user_api_keys, revoke_api_key, generate_api_key ) app = FastAPI(title=settings.app_name) security = HTTPBearer() templates = Jinja2Templates(directory="templates") app.mount("/static", StaticFiles(directory="static"), name="static") init_db() OLLAMA_MODEL = "llama3.1:8b" OLLAMA_BASE = "http://localhost:11434" # --- Models --- class RegisterRequest(BaseModel): username: str email: str password: str class LoginRequest(BaseModel): username: str password: str class CreateApiKeyRequest(BaseModel): key_name: str class ChatRequest(BaseModel): message: str stream: bool = True # --- Auth Helpers --- def create_access_token(data: dict, expires_delta: Optional[timedelta] = None): to_encode = data.copy() expire = datetime.utcnow() + (expires_delta or timedelta(minutes=settings.access_token_expire_minutes)) to_encode.update({"exp": expire}) return jwt.encode(to_encode, settings.secret_key, algorithm=settings.algorithm) async def get_current_user(credentials: HTTPAuthorizationCredentials = Depends(security)): token = credentials.credentials if token.startswith("ollama_"): user = get_user_by_api_key(token) if not user: raise HTTPException(status_code=401, detail="Invalid API key") return user try: payload = jwt.decode(token, settings.secret_key, algorithms=[settings.algorithm]) username = payload.get("sub") if username is None: raise HTTPException(status_code=401, detail="Invalid token") except JWTError: raise HTTPException(status_code=401, detail="Invalid token") user = get_user_by_username(username) if not user: raise HTTPException(status_code=401, detail="User not found") return user async def get_user_or_none(request: Request): auth = request.headers.get("Authorization", "") if not auth.startswith("Bearer "): api_key = request.headers.get("X-API-Key", "") if api_key: return get_user_by_api_key(api_key) return None try: return await get_current_user(HTTPAuthorizationCredentials(scheme="Bearer", credentials=auth[7:])) except: return None # --- Web Pages --- @app.get("/", response_class=HTMLResponse) async def index(request: Request): return templates.TemplateResponse("index.html", {"request": request}) @app.get("/login", response_class=HTMLResponse) async def login_page(request: Request): return templates.TemplateResponse("login.html", {"request": request}) @app.get("/dashboard", response_class=HTMLResponse) async def dashboard_page(request: Request): return templates.TemplateResponse("dashboard.html", {"request": request}) @app.get("/chat", response_class=HTMLResponse) async def chat_page(request: Request): return templates.TemplateResponse("chat.html", {"request": request}) # --- Auth API --- @app.post("/api/register") async def register(req: RegisterRequest): user_id = create_user(req.username, req.email, req.password) if not user_id: raise HTTPException(status_code=400, detail="Username or email already taken") return {"message": "User created", "user_id": user_id} @app.post("/api/login") async def login(req: LoginRequest): user = get_user_by_username(req.username) if not user or not verify_password(req.password, user["password_hash"]): raise HTTPException(status_code=401, detail="Invalid credentials") token = create_access_token({"sub": user["username"]}) return {"access_token": token, "token_type": "bearer", "username": user["username"]} @app.get("/api/me") async def me(user: dict = Depends(get_current_user)): return { "username": user["username"], "email": user["email"], "api_key": user["api_key"], "created_at": user["created_at"] } # --- Model / Status --- @app.get("/api/status") async def status(): ollama_path = shutil.which("ollama") if not ollama_path: return {"ollama": False, "model": False, "message": "Ollama not installed"} try: proc = await asyncio.create_subprocess_exec(ollama_path, "list", stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE) stdout, _ = await asyncio.wait_for(proc.communicate(), timeout=5) installed = [line.split()[0] for line in stdout.decode().strip().split("\n")[1:] if line.strip()] model_ready = OLLAMA_MODEL in installed return { "ollama": True, "model": model_ready, "installed": installed, "message": "Ready" if model_ready else f"Model {OLLAMA_MODEL} not yet installed (pulling in background)" } except Exception as e: return {"ollama": False, "model": False, "message": str(e)} # --- API Key Management --- @app.post("/api/api-keys") async def create_new_api_key(req: CreateApiKeyRequest, user: dict = Depends(get_current_user)): key = create_api_key(user["id"], req.key_name) return {"api_key": key, "key_name": req.key_name} @app.get("/api/api-keys") async def list_api_keys(user: dict = Depends(get_current_user)): return {"api_keys": get_user_api_keys(user["id"])} @app.delete("/api/api-keys/{key_id}") async def delete_api_key(key_id: int, user: dict = Depends(get_current_user)): revoke_api_key(key_id, user["id"]) return {"message": "API key revoked"} # --- Chat API (streaming via SSE) --- @app.post("/api/chat") async def chat(req: ChatRequest, user: Optional[dict] = Depends(get_current_user)): if not user: raise HTTPException(status_code=401, detail="Authentication required") model_ready = await check_model_ready() if not model_ready: raise HTTPException(status_code=503, detail=f"{OLLAMA_MODEL} is still downloading. Please wait and try again.") async def generate(): async with httpx.AsyncClient(timeout=None) as client: async with client.stream( "POST", f"{OLLAMA_BASE}/api/chat", json={"model": OLLAMA_MODEL, "messages": [{"role": "user", "content": req.message}], "stream": True}, ) as resp: async for line in resp.aiter_lines(): if not line.strip(): continue try: chunk = json.loads(line) content = chunk.get("message", {}).get("content", "") if content: yield f"data: {json.dumps({'content': content})}\n\n" if chunk.get("done"): yield f"data: {json.dumps({'done': True})}\n\n" except json.JSONDecodeError: continue return StreamingResponse(generate(), media_type="text/event-stream") async def check_model_ready() -> bool: try: proc = await asyncio.create_subprocess_exec( shutil.which("ollama"), "list", stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE ) stdout, _ = await asyncio.wait_for(proc.communicate(), timeout=5) return any(OLLAMA_MODEL in line for line in stdout.decode().strip().split("\n")[1:]) except: return False # --- Startup --- if __name__ == "__main__": import uvicorn uvicorn.run(app, host=settings.host, port=settings.port)