Update generator.py

generator.py

@@ -1,15 +1,17 @@
 import os
-from google.generativeai import Client
+import google.generativeai as genai
 
-API_KEY = os.getenv("GEMINI_API_KEY")
-client = Client(api_key=API_KEY)
+genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
+
+model = genai.GenerativeModel("gemini-pro")
 
 def generate_app(prompt: str):
-    system = "You are a full-stack developer."
-    response = client.chat(prompt=f"{system}\nGenerate complete React+Tailwind frontend and Node.js+Express+MongoDB backend for: {prompt}", model="gemini-pro", max_output_tokens=2000)
-    return {"full_code": response.last, "manifest": prompt}
-
-def enhance_app(bundle, prompt: str):
-    system = "You are the same full-stack developer, improving existing code."
-    response = client.chat(prompt=f"{system}\nHere is current code:\n{bundle['full_code']}\nEnhance by: {prompt}", model="gemini-pro", max_output_tokens=2000)
-    return {"full_code": response.last, "manifest": bundle["manifest"] + " | enhanced"}
+    system = "You are a full-stack developer. Build full React+Tailwind frontend and Express+MongoDB backend for:"
+    response = model.generate_content(f"{system}\n{prompt}")
+    return {"full_code": response.text, "manifest": prompt}
+
+def enhance_app(bundle, enhance_prompt: str):
+    system = "You are improving an existing app. Current code is below. Modify it as follows:"
+    prompt = f"{system}\n\n{bundle['full_code']}\n\nEnhance by: {enhance_prompt}"
+    response = model.generate_content(prompt)
+    return {"full_code": response.text, "manifest": bundle["manifest"] + " + " + enhance_prompt}