support screenshotting a URL with ScreenshotOne

backend/main.py

@@ -1,5 +1,6 @@
 # Load environment variables first
 from dotenv import load_dotenv
+from pydantic import BaseModel
 
 load_dotenv()
 
@@ -9,19 +10,35 @@ import os
 import traceback
 from datetime import datetime
 from fastapi import FastAPI, WebSocket
-
+from fastapi.middleware.cors import CORSMiddleware
 from llm import stream_openai_response
 from mock import mock_completion
 from image_generation import create_alt_url_mapping, generate_images
 from prompts import assemble_prompt
+from routes import screenshot
+
+app = FastAPI(openapi_url=None, docs_url=None, redoc_url=None)
+
+# Configure CORS
+
+# Configure CORS settings
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
-app = FastAPI()
 
 # Useful for debugging purposes when you don't want to waste GPT4-Vision credits
 # Setting to True will stream a mock response instead of calling the OpenAI API
 SHOULD_MOCK_AI_RESPONSE = False
 
 
+app.include_router(screenshot.router)
+
+
 def write_logs(prompt_messages, completion):
     # Get the logs path from environment, default to the current working directory
     logs_path = os.environ.get("LOGS_PATH", os.getcwd())

backend/poetry.lock

@@ -471,4 +471,4 @@ files = [
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.10"
-content-hash = "5e4aa03dda279f66a9b3d30f7327109bcfd395795470d95f8c563897ce1bff84"
+content-hash = "b8d248a44a5eea9638a7726096de77d7a9aa8c00673da806534da2c228ffabb4"

backend/pyproject.toml

@@ -13,6 +13,7 @@ websockets = "^12.0"
 openai = "^1.2.4"
 python-dotenv = "^1.0.0"
 beautifulsoup4 = "^4.12.2"
+httpx = "^0.25.1"
 
 [build-system]
 requires = ["poetry-core"]

backend/routes/screenshot.py

@@ -0,0 +1,64 @@
+import base64
+from fastapi import APIRouter
+from pydantic import BaseModel
+import httpx
+
+router = APIRouter()
+
+
+def bytes_to_data_url(image_bytes: bytes, mime_type: str) -> str:
+    base64_image = base64.b64encode(image_bytes).decode("utf-8")
+    return f"data:{mime_type};base64,{base64_image}"
+
+
+async def capture_screenshot(target_url, api_key, device="desktop") -> bytes:
+    api_base_url = "https://api.screenshotone.com/take"
+
+    params = {
+        "access_key": api_key,
+        "url": target_url,
+        "full_page": "true",
+        "device_scale_factor": "1",
+        "format": "png",
+        "block_ads": "true",
+        "block_cookie_banners": "true",
+        "block_trackers": "true",
+        "cache": "false",
+        "viewport_width": "342",
+        "viewport_height": "684",
+    }
+
+    if device == "desktop":
+        params["viewport_width"] = "1280"
+        params["viewport_height"] = "832"
+
+    async with httpx.AsyncClient(timeout=60) as client:
+        response = await client.get(api_base_url, params=params)
+        if response.status_code == 200 and response.content:
+            return response.content
+        else:
+            raise Exception("Error taking screenshot")
+
+
+class ScreenshotRequest(BaseModel):
+    url: str
+    apiKey: str
+
+
+class ScreenshotResponse(BaseModel):
+    url: str
+
+
+@router.post("/api/screenshot")
+async def app_screenshot(request: ScreenshotRequest):
+    # Extract the URL from the request body
+    url = request.url
+    api_key = request.apiKey
+
+    # TODO: Add error handling
+    image_bytes = await capture_screenshot(url, api_key=api_key)
+
+    # Convert the image bytes to a data url
+    data_url = bytes_to_data_url(image_bytes, "image/png")
+
+    return ScreenshotResponse(url=data_url)

frontend/src/App.tsx

@@ -21,6 +21,7 @@ import { Settings } from "./types";
 import { IS_RUNNING_ON_CLOUD } from "./config";
 import { PicoBadge } from "./components/PicoBadge";
 import { OnboardingNote } from "./components/OnboardingNote";
+import { UrlInputSection } from "./components/UrlInputSection";
 
 function App() {
   const [appState, setAppState] = useState<"INITIAL" | "CODING" | "CODE_READY">(
@@ -33,8 +34,9 @@ function App() {
   const [history, setHistory] = useState<string[]>([]);
   const [settings, setSettings] = useState<Settings>({
     openAiApiKey: null,
+    screenshotOneApiKey: null,
     isImageGenerationEnabled: true,
-    editorTheme: "cobalt"
+    editorTheme: "cobalt",
   });
 
   const downloadCode = () => {
@@ -202,9 +204,13 @@ function App() {
 
       <main className="py-2 lg:pl-96">
         {appState === "INITIAL" && (
-          <>
+          <div className="flex flex-col justify-center items-center gap-y-10">
             <ImageUpload setReferenceImages={doCreate} />
-          </>
+            <UrlInputSection
+              doCreate={doCreate}
+              screenshotOneApiKey={settings.screenshotOneApiKey}
+            />
+          </div>
         )}
 
         {(appState === "CODING" || appState === "CODE_READY") && (
@@ -231,7 +237,10 @@ function App() {
                 <Preview code={generatedCode} device="mobile" />
               </TabsContent>
               <TabsContent value="code">
-                <CodeMirror code={generatedCode} editorTheme={settings.editorTheme} />
+                <CodeMirror
+                  code={generatedCode}
+                  editorTheme={settings.editorTheme}
+                />
               </TabsContent>
             </Tabs>
           </div>

frontend/src/components/SettingsDialog.tsx

@@ -74,6 +74,26 @@ function SettingsDialog({ settings, setSettings }: Props) {
               }))
             }
           />
+
+          <Label htmlFor="screenshot-one-api-key">
+            <div>ScreenshotOne API key</div>
+            <div className="font-light mt-2">
+              Never stored. Get 500 screenshots for free by month by signing up.
+            </div>
+          </Label>
+
+          <Input
+            id="screenshot-one-api-key"
+            placeholder="ScreenshotOne API key"
+            value={settings.screenshotOneApiKey || ""}
+            onChange={(e) =>
+              setSettings((s) => ({
+                ...s,
+                screenshotOneApiKey: e.target.value,
+              }))
+            }
+          />
+
           <Label htmlFor="editor-theme">
             <div>Editor Theme</div>
           </Label>

frontend/src/components/UrlInputSection.tsx

@@ -0,0 +1,83 @@
+import { useState } from "react";
+import { HTTP_BACKEND_URL } from "../config";
+import { Button } from "./ui/button";
+import { Input } from "./ui/input";
+import { toast } from "react-hot-toast";
+
+interface Props {
+  screenshotOneApiKey: string | null;
+  doCreate: (urls: string[]) => void;
+}
+
+export function UrlInputSection({ doCreate, screenshotOneApiKey }: Props) {
+  const [isLoading, setIsLoading] = useState(false);
+  const [referenceUrl, setReferenceUrl] = useState("");
+  const isDisabled = !screenshotOneApiKey;
+
+  async function takeScreenshot() {
+    if (!screenshotOneApiKey) {
+      toast.error("Please add a Screenshot API key in the settings dialog");
+      return;
+    }
+
+    if (!referenceUrl) {
+      toast.error("Please enter a URL");
+      return;
+    }
+
+    if (referenceUrl) {
+      try {
+        setIsLoading(true);
+        const response = await fetch(`${HTTP_BACKEND_URL}/api/screenshot`, {
+          method: "POST",
+          body: JSON.stringify({
+            url: referenceUrl,
+            apiKey: screenshotOneApiKey,
+          }),
+          headers: {
+            "Content-Type": "application/json",
+          },
+        });
+
+        if (!response.ok) {
+          throw new Error("Failed to capture screenshot");
+        }
+
+        const res = await response.json();
+        doCreate([res.url]);
+      } catch (error) {
+        console.error(error);
+        toast.error(
+          "Failed to capture screenshot. Look at the console and your backend logs for more details."
+        );
+      } finally {
+        setIsLoading(false);
+      }
+    }
+  }
+
+  return (
+    <div className="w-[400px] gap-y-2 flex flex-col">
+      <div className="text-gray-500 text-sm">Or screenshot a URL...</div>
+      <Input
+        placeholder="Enter URL"
+        onChange={(e) => setReferenceUrl(e.target.value)}
+        value={referenceUrl}
+      />
+      <Button onClick={takeScreenshot} disabled={isDisabled || isLoading}>
+        {isLoading ? "Capturing..." : "Capture"}
+      </Button>
+      {isDisabled && (
+        <div className="flex space-y-4 bg-slate-200 p-2 rounded text-stone-800 text-sm">
+          <span>
+            To screenshot a URL, add a{" "}
+            <a href="https://screenshotone.com/" className="underline">
+              ScreenshotOne API key
+            </a>{" "}
+            in the settings dialog.
+          </span>
+        </div>
+      )}
+    </div>
+  );
+}

frontend/src/config.ts

@@ -1,3 +1,6 @@
 // Default to false if set to anything other than "true" or unset
 export const IS_RUNNING_ON_CLOUD =
   import.meta.env.VITE_IS_DEPLOYED === "true" || false;
+
+export const HTTP_BACKEND_URL =
+  import.meta.env.VITE_HTTP_BACKEND_URL || "http://127.0.0.1:7001";

frontend/src/types.ts

@@ -1,5 +1,6 @@
 export interface Settings {
   openAiApiKey: string | null;
+  screenshotOneApiKey: string | null;
   isImageGenerationEnabled: boolean;
   editorTheme: string;
 }