🔧 fix(api_key.py): import missing dependencies and fix import paths for auth_utils and get_current_active_user

✨ feat(api_key.py): add new route to save store API key for a user
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-10-18 17:03:53 -03:00
commit 596fc8ed41

View file

@ -1,7 +1,8 @@
from uuid import UUID from uuid import UUID
from fastapi import APIRouter, HTTPException, Depends from fastapi import APIRouter, HTTPException, Depends, Body
from httpx import get
from langflow.api.v1.schemas import ApiKeysResponse from langflow.api.v1.schemas import ApiKeysResponse
from langflow.services.auth.utils import get_current_active_user from langflow.services.auth import utils as auth_utils
from langflow.services.database.models.api_key.api_key import ( from langflow.services.database.models.api_key.api_key import (
ApiKeyCreate, ApiKeyCreate,
UnmaskedApiKeyRead, UnmaskedApiKeyRead,
@ -14,9 +15,19 @@ from langflow.services.database.models.api_key.crud import (
delete_api_key, delete_api_key,
) )
from langflow.services.database.models.user.user import User from langflow.services.database.models.user.user import User
from langflow.services.deps import get_session from langflow.services.deps import (
get_store_service,
get_session,
get_settings_service,
)
from typing import TYPE_CHECKING
from sqlmodel import Session from sqlmodel import Session
if TYPE_CHECKING:
from langflow.services.store.service import StoreService
from langflow.services.settings.service import SettingsService
router = APIRouter(tags=["APIKey"], prefix="/api_key") router = APIRouter(tags=["APIKey"], prefix="/api_key")
@ -24,7 +35,7 @@ router = APIRouter(tags=["APIKey"], prefix="/api_key")
@router.get("/", response_model=ApiKeysResponse) @router.get("/", response_model=ApiKeysResponse)
def get_api_keys_route( def get_api_keys_route(
db: Session = Depends(get_session), db: Session = Depends(get_session),
current_user: User = Depends(get_current_active_user), current_user: User = Depends(auth_utils.auth_utils.get_current_active_user),
): ):
try: try:
user_id = current_user.id user_id = current_user.id
@ -38,7 +49,7 @@ def get_api_keys_route(
@router.post("/", response_model=UnmaskedApiKeyRead) @router.post("/", response_model=UnmaskedApiKeyRead)
def create_api_key_route( def create_api_key_route(
req: ApiKeyCreate, req: ApiKeyCreate,
current_user: User = Depends(get_current_active_user), current_user: User = Depends(auth_utils.get_current_active_user),
db: Session = Depends(get_session), db: Session = Depends(get_session),
): ):
try: try:
@ -51,7 +62,7 @@ def create_api_key_route(
@router.delete("/{api_key_id}") @router.delete("/{api_key_id}")
def delete_api_key_route( def delete_api_key_route(
api_key_id: UUID, api_key_id: UUID,
current_user=Depends(get_current_active_user), current_user=Depends(auth_utils.get_current_active_user),
db: Session = Depends(get_session), db: Session = Depends(get_session),
): ):
try: try:
@ -59,3 +70,22 @@ def delete_api_key_route(
return {"detail": "API Key deleted"} return {"detail": "API Key deleted"}
except Exception as e: except Exception as e:
raise HTTPException(status_code=400, detail=str(e)) from e raise HTTPException(status_code=400, detail=str(e)) from e
@router.post("/store")
def save_store_api_key(
api_key: Body(str, embed=True),
current_user: User = Depends(auth_utils.get_current_active_user),
db: Session = Depends(get_session),
settings_service=Depends(get_settings_service),
):
try:
# Encrypt the API key
encrypted = auth_utils.encrypt_api_key(
api_key, fernet=auth_utils.get_fernet(settings_service)
)
current_user.store_api_key = encrypted
db.commit()
return {"detail": "API Key saved"}
except Exception as e:
raise HTTPException(status_code=400, detail=str(e)) from e