Update server.js

server.js

@@ -4,106 +4,166 @@ import { fileURLToPath } from "url";
 import dotenv from "dotenv";
 import bodyParser from "body-parser";
 
+import { PROVIDERS } from "./utils/providers.js";
+
+// Load environment variables from .env file
 dotenv.config();
 
 const app = express();
+
 const __filename = fileURLToPath(import.meta.url);
 const __dirname = path.dirname(__filename);
-const PORT = process.env.APP_PORT || 7860;
-const OPENROUTER_API_KEY = process.env.OPENROUTER_API_KEY;
+
+const PORT = process.env.APP_PORT || 3000;
 
 app.use(bodyParser.json());
 app.use(express.static(path.join(__dirname, "dist")));
 
 app.post("/api/ask-ai", async (req, res) => {
-  const { prompt, html, previousPrompt } = req.body;
+  const { prompt, html, previousPrompt, provider } = req.body;
   if (!prompt) {
-    return res.status(400).send({ ok: false, message: "Missing prompt" });
-  }
-
-  if (!OPENROUTER_API_KEY) {
-    return res.status(500).send({
+    return res.status(400).send({
       ok: false,
-      message: "Missing OpenRouter API Key",
+      message: "Missing required fields",
     });
   }
 
+  // Set up response headers for streaming
   res.setHeader("Content-Type", "text/plain");
   res.setHeader("Cache-Control", "no-cache");
   res.setHeader("Connection", "keep-alive");
 
-  const messages = [
-    {
-      role: "system",
-      content:
-        "Você é um especialista em landing pages. Gere um HTML bonito e persuasivo para a copy enviada. Responda apenas com código HTML válido e completo.",
-    },
-    ...(previousPrompt ? [{ role: "user", content: previousPrompt }] : []),
-    ...(html ? [{ role: "assistant", content: `The current code is: ${html}` }] : []),
-    { role: "user", content: prompt },
-  ];
-
-  try {
-    const response = await fetch("https://openrouter.ai/api/v1/chat/completions", {
-      method: "POST",
-      headers: {
-        Authorization: `Bearer ${OPENROUTER_API_KEY}`,
-        "Content-Type": "application/json",
-        "HTTP-Referer": "https://raypages.com", // personalize aqui
-        "X-Title": "RayPages AI",
-      },
-      body: JSON.stringify({
-        model: "deepseek-chat-v3.0",
-        stream: true,
-        messages,
-      }),
+  let systemPrompt = `ONLY USE HTML, CSS AND JAVASCRIPT. No explanations, ONLY CODE. If you want to use ICON make sure to import the library first. Try to create the best UI possible by using only HTML, CSS and JAVASCRIPT. Use as much as you can TailwindCSS for the CSS, if you can't do something with TailwindCSS, then use custom CSS (make sure to import <script src="https://cdn.jsdelivr.net/npm/@tailwindcss/browser@4"></script> in the head). Also, try to ellaborate as much as you can, to create something unique. ALWAYS GIVE THE RESPONSE INTO A SINGLE HTML FILE`;
+
+  let TOKENS_USED = prompt?.length;
+  if (previousPrompt) TOKENS_USED += previousPrompt.length;
+  if (html) TOKENS_USED += html.length;
+
+  const DEFAULT_PROVIDER = PROVIDERS.openrouter;
+  const selectedProvider =
+    provider === "auto"
+      ? DEFAULT_PROVIDER
+      : PROVIDERS[provider] ?? DEFAULT_PROVIDER;
+
+  if (provider !== "auto" && TOKENS_USED >= selectedProvider.max_tokens) {
+    return res.status(400).send({
+      ok: false,
+      openSelectProvider: true,
+      message: `Context is too long. ${selectedProvider.name} allow ${selectedProvider.max_tokens} max tokens.`,
     });
+  }
 
-    if (!response.ok || !response.body) {
-      return res.status(500).send({
-        ok: false,
-        message: "Failed to fetch OpenRouter response",
+  if (["local", "openrouter"].includes(selectedProvider.id)) {
+    try {
+      const { ApiKey, ApiUrl, Model } = req.body;
+      if (!ApiUrl || !Model) {
+        return res.status(400).send({
+          ok: false,
+          message: "Missing required fields for provider, set API KEY, BASE URL, and MODEL.",
+        });
+      }
+      const response = await fetch(`${ApiUrl}/chat/completions`, {
+        method: "POST",
+        headers: {
+          "Content-Type": "application/json",
+          Authorization: `Bearer ${ApiKey}`,
+        },
+        body: JSON.stringify({
+          model: Model,
+          messages: [
+            {
+              role: "system",
+              content: systemPrompt,
+            },
+            ...(previousPrompt
+              ? [
+                  {
+                    role: "user",
+                    content: previousPrompt,
+                  },
+                ]
+              : []),
+            ...(html
+              ? [
+                  {
+                    role: "assistant",
+                    content: `The current code is: ${html}.`,
+                  },
+                ]
+              : []),
+            {
+              role: "user",
+              content: prompt,
+            },
+          ],
+          stream: true
+        })
       });
-    }
+      if (!response.ok) {
+        const errorBody = await response.text();
+        throw new Error(`API Error: ${response.status} - ${errorBody}`);
+      }
 
-    const reader = response.body.getReader();
-    const decoder = new TextDecoder("utf-8");
-    let completeResponse = "";
+      const reader = response.body.getReader();
+      const decoder = new TextDecoder();
 
-    const stream = async () => {
       while (true) {
         const { done, value } = await reader.read();
         if (done) break;
 
-        const chunk = decoder.decode(value, { stream: true });
-        const matches = chunk.match(/"content":"([^"]*)"/g);
+        const chunk = decoder.decode(value);
+        const lines = chunk.split('\n').filter(line => line.trim());
 
-        if (matches) {
-          for (const match of matches) {
-            const text = match
-              .replace(/"content":"|"/g, "")
-              .replace(/\\n/g, "\n")
-              .replace(/\\t/g, "\t");
+        for (const line of lines) {
+          if (!line || !line.startsWith('data: ') || line.includes('[DONE]')) continue;
 
-            completeResponse += text;
-            res.write(text);
+          try {
+            if (line.includes('exceeded')) {
+              return res.status(402).send({
+                ok: false,
+                message: line,
+              });
+            }
+
+            const json = line.slice(6).trim();
 
-            if (completeResponse.includes("</html>")) {
-              res.end();
-              return;
+            if (!json.startsWith('{')) {
+              continue;
             }
+
+            const message = JSON.parse(json);
+            const content = message?.choices?.[0]?.delta?.content;
+
+            if (content) {
+              res.write(content);
+            }
+          } catch (e) {
+            console.error('Error on line stream:', e.message);
+            // continue ao invés de throw, pra não matar tudo por causa de uma linha podre
+            continue;
           }
         }
+
       }
       res.end();
-    };
-
-    stream();
-  } catch (err) {
-    res.status(500).send({
-      ok: false,
-      message: err.message || "AI request failed",
-    });
+    } catch (error) {
+      if (error.message.includes("exceeded")) {
+        return res.status(402).send({
+          ok: false,
+          message: error.message,
+        });
+      }
+      if (!res.headersSent) {
+        res.status(500).send({
+          ok: false,
+          message:
+            error.message || "An error occurred while processing your request.",
+        });
+      } else {
+        // Otherwise end the stream
+        res.end();
+      }
+    }
   }
 });
 
@@ -112,5 +172,5 @@ app.get("*", (_req, res) => {
 });
 
 app.listen(PORT, () => {
-  console.log(`✅ Server running on port ${PORT}`);
-});
+  console.log(`Server is running on port ${PORT}`);
+});