From e9140c331b286c93b988d80c9fb9d6a227dbffd1 Mon Sep 17 00:00:00 2001 From: Abi Raja Date: Tue, 30 Jan 2024 13:51:35 -0500 Subject: [PATCH] for a subscription, regardless of the key setting, it should just work --- backend/routes/generate_code.py | 109 +++++++++++++++++--------------- 1 file changed, 57 insertions(+), 52 deletions(-) diff --git a/backend/routes/generate_code.py b/backend/routes/generate_code.py index 349afbc..a37a4e7 100644 --- a/backend/routes/generate_code.py +++ b/backend/routes/generate_code.py @@ -56,8 +56,6 @@ async def stream_code(websocket: WebSocket): # TODO: Are the values always strings? params: Dict[str, str] = await websocket.receive_json() - print("Received params") - # Read the code config settings from the request. Fall back to default if not provided. generated_code_config = "" if "generatedCodeConfig" in params and params["generatedCodeConfig"]: @@ -66,64 +64,71 @@ async def stream_code(websocket: WebSocket): # Track how this generation is being paid for payment_method: PaymentMethod = PaymentMethod.UNKNOWN - - # Get the OpenAI API key from the request. Fall back to environment variable if not provided. - # If neither is provided, we throw an error. + # Track the OpenAI API key to use openai_api_key = None - if "accessCode" in params and params["accessCode"]: - print("Access code - using platform API key") - res = await validate_access_token(params["accessCode"]) - if res["success"]: - payment_method = PaymentMethod.ACCESS_CODE + + auth_token = params.get("authToken") + if not auth_token: + await throw_error("You need to be logged in to use screenshot to code") + raise Exception("No auth token") + + # Get the OpenAI key by waterfalling through the different payment methods + # 1. Subcription + # 2. Access code + # 3. User's API key from client-side settings dialog + # 4. User's API key from environment variable + + # If the user is a subscriber, use the platform API key + # TODO: Rename does_user_have_subscription_credits + res = await does_user_have_subscription_credits(auth_token) + if res.status != "not_subscriber": + if res.status == "subscriber_has_credits": + payment_method = PaymentMethod.SUBSCRIPTION openai_api_key = os.environ.get("PLATFORM_OPENAI_API_KEY") - else: - await websocket.send_json( - { - "type": "error", - "value": res["failure_reason"], - } + print("Subscription - using platform API key") + elif res.status == "subscriber_has_no_credits": + await throw_error( + "Your subscription has run out of monthly credits. Contact support and we can add more credits to your account for free." ) - return - else: - auth_token = params.get("authToken") - if auth_token: - # TODO: Rename does_user_have_subscription_credits - res = await does_user_have_subscription_credits(auth_token) - if res.status == "not_subscriber": - # Keep going for non-subscriber users - pass - elif res.status == "subscriber_has_credits": - payment_method = PaymentMethod.SUBSCRIPTION - openai_api_key = os.environ.get("PLATFORM_OPENAI_API_KEY") - elif res.status == "subscriber_has_no_credits": - return await throw_error( - "Your subscription has run out of monthly credits. Contact support and we can add more credits to your account for free." - ) - else: - return await throw_error("Unknown error occurred. Contact support.") - + raise Exception("User has no credits") else: - # Log but keep going for users - print("Missing auth token") + await throw_error("Unknown error occurred. Contact support.") + raise Exception("Unknown error occurred when checking subscription credits") - if params["openAiApiKey"]: - openai_api_key = params["openAiApiKey"] - payment_method = PaymentMethod.OPENAI_API_KEY - print("Using OpenAI API key from client-side settings dialog") - elif not openai_api_key: - openai_api_key = os.environ.get("OPENAI_API_KEY") - if openai_api_key: - print("Using OpenAI API key from environment variable") + # For non-subscribers, if they have an access code, validate it + # and use the platform API key + if not openai_api_key: + accessCode = params.get("accessCode", None) + if accessCode: + print("Access code - using platform API key") + res = await validate_access_token(accessCode) + if res["success"]: + payment_method = PaymentMethod.ACCESS_CODE + openai_api_key = os.environ.get("PLATFORM_OPENAI_API_KEY") + else: + await throw_error(res["failure_reason"]) + raise Exception("Invalid access code: " + accessCode) + # If we still don't have an API key, use the user's API key from client-side settings dialog + if not openai_api_key: + openai_api_key = params.get("openAiApiKey", None) + payment_method = PaymentMethod.OPENAI_API_KEY + print("Using OpenAI API key from client-side settings dialog") + + # If we still don't have an API key, use the user's API key from environment variable + if not openai_api_key: + openai_api_key = os.environ.get("OPENAI_API_KEY") + payment_method = PaymentMethod.OPENAI_API_KEY + if openai_api_key: + print("Using OpenAI API key from environment variable") + + # If we still don't have an API key, throw an error if not openai_api_key: print("OpenAI API key not found") - await websocket.send_json( - { - "type": "error", - "value": "No OpenAI API key found. Please add your API key in the settings dialog or add it to backend/.env file.", - } + await throw_error( + "No OpenAI API key found. Please add your API key in the settings dialog or add it to backend/.env file." ) - return + raise Exception("No OpenAI API key found") # Validate the generated code config if not generated_code_config in get_args(Stack): @@ -218,7 +223,7 @@ async def stream_code(websocket: WebSocket): image_cache = create_alt_url_mapping(params["history"][-2]) - pprint_prompt(prompt_messages) + # pprint_prompt(prompt_messages) if SHOULD_MOCK_AI_RESPONSE: completion = await mock_completion(process_chunk)