diff --git a/backend/access_token.py b/backend/access_token.py new file mode 100644 index 0000000..6a9bbd4 --- /dev/null +++ b/backend/access_token.py @@ -0,0 +1,19 @@ +import json +import httpx + + +async def validate_access_token(access_code: str): + async with httpx.AsyncClient() as client: + url = "http://localhost:8000/screenshot_to_code/validate_access_token" + data = json.dumps({"access_code": access_code}) + 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 117c58c..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) @@ -82,13 +82,13 @@ async def stream_code(websocket: WebSocket): openai_api_key = None if "accessCode" in params and params["accessCode"]: print("Access code - using platform API key") - if params["accessCode"] == "hi": + 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. Please try again.", + "value": "Invalid access code or you're out of credits. Please try again.", } ) return