diff --git a/backend/access_token.py b/backend/access_token.py new file mode 100644 index 0000000..e61ef12 --- /dev/null +++ b/backend/access_token.py @@ -0,0 +1,27 @@ +import json +import os +import httpx + + +async def validate_access_token(access_code: str): + async with httpx.AsyncClient() as client: + url = ( + "https://backend.buildpicoapps.com/screenshot_to_code/validate_access_token" + ) + data = json.dumps( + { + "access_code": access_code, + "secret": os.environ.get("PICO_BACKEND_SECRET"), + } + ) + headers = {"Content-Type": "application/json"} + + response = await client.post(url, content=data, headers=headers) + response_data = response.json() + + if response_data["success"]: + print("Access token is valid.") + return True + else: + print(f"Access token validation failed: {response_data['failure_reason']}") + return False diff --git a/backend/main.py b/backend/main.py index 4b75f71..4c5823b 100644 --- a/backend/main.py +++ b/backend/main.py @@ -1,6 +1,5 @@ # Load environment variables first from dotenv import load_dotenv -from pydantic import BaseModel load_dotenv() @@ -16,6 +15,7 @@ from mock import mock_completion from image_generation import create_alt_url_mapping, generate_images from prompts import assemble_prompt from routes import screenshot +from access_token import validate_access_token app = FastAPI(openapi_url=None, docs_url=None, redoc_url=None) @@ -79,13 +79,27 @@ async def stream_code(websocket: WebSocket): # Get the OpenAI API key from the request. Fall back to environment variable if not provided. # If neither is provided, we throw an error. - if params["openAiApiKey"]: - openai_api_key = params["openAiApiKey"] - print("Using OpenAI API key from client-side settings dialog") + openai_api_key = None + if "accessCode" in params and params["accessCode"]: + print("Access code - using platform API key") + if await validate_access_token(params["accessCode"]): + openai_api_key = os.environ.get("PLATFORM_OPENAI_API_KEY") + else: + await websocket.send_json( + { + "type": "error", + "value": "Invalid access code or you're out of credits. Please try again.", + } + ) + return else: - openai_api_key = os.environ.get("OPENAI_API_KEY") - if openai_api_key: - print("Using OpenAI API key from environment variable") + if params["openAiApiKey"]: + openai_api_key = params["openAiApiKey"] + print("Using OpenAI API key from client-side settings dialog") + else: + openai_api_key = os.environ.get("OPENAI_API_KEY") + if openai_api_key: + print("Using OpenAI API key from environment variable") if not openai_api_key: print("OpenAI API key not found") diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx index e9ba1ce..9dab79e 100644 --- a/frontend/src/App.tsx +++ b/frontend/src/App.tsx @@ -51,6 +51,7 @@ function App() { isImageGenerationEnabled: true, editorTheme: EditorTheme.COBALT, isTermOfServiceAccepted: false, + accessCode: null, }, "setting" ); diff --git a/frontend/src/components/SettingsDialog.tsx b/frontend/src/components/SettingsDialog.tsx index 215b0ea..8a8be03 100644 --- a/frontend/src/components/SettingsDialog.tsx +++ b/frontend/src/components/SettingsDialog.tsx @@ -15,6 +15,7 @@ import { Label } from "./ui/label"; import { Input } from "./ui/input"; import { Select, SelectContent, SelectItem, SelectTrigger } from "./ui/select"; import { capitalize } from "../lib/utils"; +import { IS_RUNNING_ON_CLOUD } from "../config"; interface Props { settings: Settings; @@ -38,6 +39,31 @@ function SettingsDialog({ settings, setSettings }: Props) { Settings + + {/* Access code */} + {IS_RUNNING_ON_CLOUD && ( +
+ + + + setSettings((s) => ({ + ...s, + accessCode: e.target.value, + })) + } + /> +
+ )} +