Update app.py

app.py

@@ -9,7 +9,27 @@ from IPython.display import display, Markdown
 api_key = os.getenv("GEMINI_API_KEY")
 genai.configure(api_key = api_key)
 
-# model = genai.GenerativeModel(model_name='gemini-2.0-flash-exp')
+chat_model = genai.GenerativeModel(model_name='gemini-2.0-flash-exp')
+generationConfig = {temperature: 0.7, topK: 1, topP: 1, maxOutputTokens: 1024,};
+safetySettings = [
+    {
+      category: HarmCategory.HARM_CATEGORY_HARASSMENT,
+      threshold: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+    },
+    {
+      category: HarmCategory.HARM_CATEGORY_HATE_SPEECH,
+      threshold: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+    },
+    {
+      category: HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
+      threshold: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+    },
+    {
+      category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
+      threshold: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+    },
+]
+
 # query = "Tell me about yourself"
 # response = model.generate_content(f"""You are an friendly assistant.Answer the following with high precision.{query}""")display(Markdown(response.text))
 
@@ -24,16 +44,22 @@ def ask_openai(prompt, secret_word):
     system_message = os.getenv("SYSTM_MSG")
     
     try:
-        response = client.chat.completions.create(
-            model="gpt-4o",
-            messages=[
-                {"role": "system", "content": system_message},
-                {"role": "user", "content": prompt}
-            ],
-            temperature=0.7,
-            max_tokens=100,
-        )
-        return response.choices[0].message.content
+        chat = model.start_chat(generationConfig, safetySettings,
+                                history=[{"role": "system", "content": system_message}, {"role": "user", "content": prompt},])
+        
+        print(repsonse)
+        return response.text
+            
+        # response = client.chat.completions.create(
+        #     model="gpt-4o",
+        #     messages=[
+        #         {"role": "system", "content": system_message},
+        #         {"role": "user", "content": prompt}
+        #     ],
+        #     temperature=0.7,
+        #     max_tokens=100,
+        # )
+        # return response.choices[0].message.content
     except Exception as e:
         return f"Error: {str(e)}"
 
@@ -61,13 +87,15 @@ def jailbreak_game(user_prompt, current_level, attempts):
         else:
             NY_msg = "✨ May 2025 bring you endless joy, success, and fireworks of happiness! Happy New Year! ✨"
             try:
-                response = client.chat.completions.create(
-                    model="gpt-4o",
-                    messages=[{"role": "user", "content": "Generate a unique Happy New year wish. Your response should be only the New Year wish."}],
-                    temperature=0.1,
-                    max_tokens=100,
-                )
-                NY_msg = response.choices[0].message.content
+                chat = model.start_chat(generationConfig, safetySettings,
+                                history=[{"role": "user", "content": "Generate a unique Happy New year wish. Your response should be only the New Year wish."},])
+                # response = client.chat.completions.create(
+                #     model="gpt-4o",
+                #     messages=[{"role": "user", "content": "Generate a unique Happy New year wish. Your response should be only the New Year wish."}],
+                #     temperature=0.1,
+                #     max_tokens=100,
+                # )
+                NY_msg = response.text #response.choices[0].message.content
             except Exception as e:
                 NY_msg = f"✨ May 2025 bring you endless joy, success, and fireworks of happiness! Happy New Year! ✨"