Upload 4 files

app/routers/auth_router.py

@@ -0,0 +1,44 @@
+from fastapi import APIRouter, Depends, HTTPException, status
+from sqlalchemy.orm import Session
+
+from ..database import get_db
+from .. import models
+from ..auth import verify_password, hash_password, create_access_token, get_current_user
+from ..schemas import UserCreate, UserOut, Token, LoginRequest
+
+router = APIRouter(prefix="/auth", tags=["Authentication"])
+
+
+@router.post("/signup", response_model=Token, status_code=status.HTTP_201_CREATED)
+def signup(payload: UserCreate, db: Session = Depends(get_db)):
+    if db.query(models.User).filter(models.User.username == payload.username).first():
+        raise HTTPException(status_code=400, detail="Username already taken")
+    if db.query(models.User).filter(models.User.email == payload.email).first():
+        raise HTTPException(status_code=400, detail="Email already registered")
+
+    user = models.User(
+        username=payload.username,
+        email=payload.email,
+        hashed_password=hash_password(payload.password),
+    )
+    db.add(user)
+    db.commit()
+    db.refresh(user)
+
+    token = create_access_token({"sub": user.username})
+    return Token(access_token=token, user=UserOut.model_validate(user))
+
+
+@router.post("/login", response_model=Token)
+def login(payload: LoginRequest, db: Session = Depends(get_db)):
+    user = db.query(models.User).filter(models.User.username == payload.username).first()
+    if not user or not verify_password(payload.password, user.hashed_password):
+        raise HTTPException(status_code=401, detail="Invalid username or password")
+
+    token = create_access_token({"sub": user.username})
+    return Token(access_token=token, user=UserOut.model_validate(user))
+
+
+@router.get("/me", response_model=UserOut)
+def me(current_user: models.User = Depends(get_current_user)):
+    return current_user

app/routers/proxy_config_router.py

@@ -0,0 +1,114 @@
+import json
+import secrets
+from fastapi import APIRouter, Depends, HTTPException, Request
+from sqlalchemy.orm import Session
+
+from ..database import get_db
+from .. import models
+from ..auth import get_current_user, encrypt_api_key
+from ..schemas import ProxyCreate, ProxyUpdate, ProxyOut
+
+router = APIRouter(prefix="/proxies", tags=["Proxy Management"])
+
+
+def _proxy_out_with_url(proxy: models.ProxyConfig, request: Request) -> dict:
+    base = str(request.base_url).rstrip("/")
+    return {
+        **ProxyOut.model_validate(proxy).model_dump(),
+        "proxy_url": f"{base}/{proxy.proxy_token}",
+    }
+
+
+@router.post("", status_code=201)
+def create_proxy(
+    payload: ProxyCreate,
+    request: Request,
+    db: Session = Depends(get_db),
+    current_user: models.User = Depends(get_current_user),
+):
+    token = secrets.token_urlsafe(24)
+    proxy = models.ProxyConfig(
+        user_id=current_user.id,
+        proxy_token=token,
+        name=payload.name,
+        openai_base_url=payload.openai_base_url.rstrip("/"),
+        encrypted_api_key=encrypt_api_key(payload.openai_api_key),
+        model_mapping=json.dumps(payload.model_mapping or {}),
+    )
+    db.add(proxy)
+    db.commit()
+    db.refresh(proxy)
+    return _proxy_out_with_url(proxy, request)
+
+
+@router.get("", response_model=list)
+def list_proxies(
+    request: Request,
+    db: Session = Depends(get_db),
+    current_user: models.User = Depends(get_current_user),
+):
+    proxies = db.query(models.ProxyConfig).filter(
+        models.ProxyConfig.user_id == current_user.id
+    ).all()
+    return [_proxy_out_with_url(p, request) for p in proxies]
+
+
+@router.get("/{proxy_id}")
+def get_proxy(
+    proxy_id: int,
+    request: Request,
+    db: Session = Depends(get_db),
+    current_user: models.User = Depends(get_current_user),
+):
+    proxy = db.query(models.ProxyConfig).filter(
+        models.ProxyConfig.id == proxy_id,
+        models.ProxyConfig.user_id == current_user.id,
+    ).first()
+    if not proxy:
+        raise HTTPException(status_code=404, detail="Proxy not found")
+    return _proxy_out_with_url(proxy, request)
+
+
+@router.patch("/{proxy_id}")
+def update_proxy(
+    proxy_id: int,
+    payload: ProxyUpdate,
+    request: Request,
+    db: Session = Depends(get_db),
+    current_user: models.User = Depends(get_current_user),
+):
+    proxy = db.query(models.ProxyConfig).filter(
+        models.ProxyConfig.id == proxy_id,
+        models.ProxyConfig.user_id == current_user.id,
+    ).first()
+    if not proxy:
+        raise HTTPException(status_code=404, detail="Proxy not found")
+
+    if payload.name is not None:
+        proxy.name = payload.name
+    if payload.openai_base_url is not None:
+        proxy.openai_base_url = payload.openai_base_url.rstrip("/")
+    if payload.openai_api_key is not None:
+        proxy.encrypted_api_key = encrypt_api_key(payload.openai_api_key)
+    if payload.model_mapping is not None:
+        proxy.model_mapping = json.dumps(payload.model_mapping)
+
+    db.commit()
+    db.refresh(proxy)
+    return _proxy_out_with_url(proxy, request)
+
+
+@router.delete("/{proxy_id}", status_code=204)
+def delete_proxy(
+    proxy_id: int,
+    db: Session = Depends(get_db),
+    current_user: models.User = Depends(get_current_user),
+):
+    proxy = db.query(models.ProxyConfig).filter(
+        models.ProxyConfig.id == proxy_id,
+        models.ProxyConfig.user_id == current_user.id,
+    ).first()
+    if not proxy:
+        raise HTTPException(status_code=404, detail="Proxy not found")
+    db.delete(proxy)
+    db.commit()

app/routers/proxy_endpoint_router.py

@@ -0,0 +1,50 @@
+from fastapi import APIRouter, Depends, HTTPException, Request
+from fastapi.responses import JSONResponse
+from sqlalchemy.orm import Session
+
+from ..database import get_db
+from .. import models
+from ..proxy_handler import handle_messages_request
+
+router = APIRouter(tags=["Proxy Endpoint"])
+
+
+def _get_proxy(token: str, db: Session) -> models.ProxyConfig:
+    proxy = db.query(models.ProxyConfig).filter(
+        models.ProxyConfig.proxy_token == token
+    ).first()
+    if not proxy:
+        raise HTTPException(status_code=404, detail="Invalid proxy token")
+    return proxy
+
+
+@router.post("/{token}/v1/messages")
+async def messages(
+    token: str,
+    request: Request,
+    db: Session = Depends(get_db),
+):
+    proxy = _get_proxy(token, db)
+    try:
+        body = await request.json()
+    except Exception:
+        raise HTTPException(status_code=400, detail="Invalid JSON body")
+
+    result = await handle_messages_request(body, proxy)
+    if isinstance(result, dict):
+        return JSONResponse(content=result)
+    return result
+
+
+@router.get("/{token}/v1/models")
+def models_list(token: str, db: Session = Depends(get_db)):
+    _get_proxy(token, db)
+    return JSONResponse(content={
+        "object": "list",
+        "data": [
+            {"id": "claude-3-opus-20240229",    "object": "model"},
+            {"id": "claude-3-sonnet-20240229",  "object": "model"},
+            {"id": "claude-3-haiku-20240307",   "object": "model"},
+            {"id": "claude-3-5-sonnet-20240620","object": "model"},
+        ],
+    })