topoteretes--cognee/cognee/modules/users/api_key/create_api_key.py

import secrets
from pydantic_settings import BaseSettings, SettingsConfigDict

from cognee.modules.users.models import User
from cognee.shared.logging_utils import get_logger
from cognee.infrastructure.databases.relational import get_relational_engine
from cognee.modules.users.models.UserApiKey import UserApiKey
from .exceptions import ApiKeyCreationError
from .get_api_keys import get_api_keys
from .hash_api_key import prepare_api_key


logger = get_logger(__name__)


class ApiKeySettings(BaseSettings):
    max_user_api_keys: int = 10

    model_config = SettingsConfigDict(env_file=".env", extra="allow")


apiKeySettings = ApiKeySettings()


async def create_api_key(user: User, name: str = None):
    existing_api_keys = await get_api_keys(user)

    if len(existing_api_keys) >= apiKeySettings.max_user_api_keys:
        raise ApiKeyCreationError("You have reached the maximum number of API keys.")

    relational_engine = get_relational_engine()

    api_key = generate_api_key()
    prepared_api_key = prepare_api_key(api_key)
    label = api_key[:8] + "****"

    async with relational_engine.get_async_session() as session:
        user_api_key = UserApiKey(
            user_id=user.id,
            api_key=prepared_api_key,
            label=label,
            name=name,
        )
        session.add(user_api_key)

        try:
            await session.commit()
            user_api_key.api_key = api_key  # return raw key so caller can show it once
            return user_api_key
        except Exception as error:
            logger.error(f"Failed to create API key for user {user.id}: {str(error)}")
            await session.rollback()
            raise ApiKeyCreationError("Failed to create API key, please try again.")


def generate_api_key():
    return secrets.token_hex(32)