Update app.py

app.py

@@ -12,7 +12,7 @@ def check_ollama():
     except:
         return False
 
-def chat_stream(message, history, temperature):
+def chat_stream(message, history, temperature, max_tokens):
     if not check_ollama():
         yield "⏳ Ollama starting... wait 30 seconds and try again."
         return
@@ -29,7 +29,7 @@ def chat_stream(message, history, temperature):
     try:
         response = requests.post(
             f"{OLLAMA_URL}/api/chat",
-            json={"model": MODEL, "messages": messages, "stream": True, "options": {"temperature": temperature}},
+            json={"model": MODEL, "messages": messages, "stream": True, "options": {"temperature": temperature, "num_predict": max_tokens}},
             stream=True, timeout=300
         )
         
@@ -46,7 +46,7 @@ def chat_stream(message, history, temperature):
     except Exception as e:
         yield f"Error: {e}"
 
-def generate_code(prompt, language):
+def generate_code(prompt, language, max_tokens):
     if not prompt.strip():
         return "Please describe what you want."
     if not check_ollama():
@@ -57,7 +57,7 @@ def generate_code(prompt, language):
     try:
         r = requests.post(
             f"{OLLAMA_URL}/api/generate",
-            json={"model": MODEL, "prompt": full_prompt, "stream": False, "options": {"temperature": 0.3}},
+            json={"model": MODEL, "prompt": full_prompt, "stream": False, "options": {"temperature": 0.3, "num_predict": max_tokens}},
             timeout=300
         )
         if r.status_code == 200:
@@ -74,7 +74,7 @@ def generate_code(prompt, language):
     except Exception as e:
         return f"Error: {e}"
 
-def explain_code(code):
+def explain_code(code, max_tokens):
     if not code.strip():
         return "Paste code to explain."
     if not check_ollama():
@@ -83,14 +83,14 @@ def explain_code(code):
     try:
         r = requests.post(
             f"{OLLAMA_URL}/api/generate",
-            json={"model": MODEL, "prompt": f"Explain this code:\n```\n{code}\n```", "stream": False},
+            json={"model": MODEL, "prompt": f"Explain this code:\n```\n{code}\n```", "stream": False, "options": {"num_predict": max_tokens}},
             timeout=300
         )
         return r.json().get("response", "") if r.status_code == 200 else f"Error: {r.text}"
     except Exception as e:
         return f"Error: {e}"
 
-def fix_code(code, error):
+def fix_code(code, error, max_tokens):
     if not code.strip():
         return "Paste code to fix."
     if not check_ollama():
@@ -101,7 +101,7 @@ def fix_code(code, error):
     try:
         r = requests.post(
             f"{OLLAMA_URL}/api/generate",
-            json={"model": MODEL, "prompt": prompt, "stream": False, "options": {"temperature": 0.3}},
+            json={"model": MODEL, "prompt": prompt, "stream": False, "options": {"temperature": 0.3, "num_predict": max_tokens}},
             timeout=300
         )
         return r.json().get("response", "") if r.status_code == 200 else f"Error: {r.text}"
@@ -112,7 +112,9 @@ with gr.Blocks(title="GOD Coding Machine", theme=gr.themes.Soft(primary_hue="pur
     
     gr.Markdown("# 🔥 GOD Coding Machine\n**Docker Edition** • Qwen2.5-Coder running locally • No rate limits!")
     
-    temperature = gr.Slider(0, 1, value=0.7, step=0.1, label="Temperature")
+    with gr.Row():
+        temperature = gr.Slider(0, 1, value=0.7, step=0.1, label="🌡️ Temperature")
+        max_tokens = gr.Slider(256, 4096, value=2048, step=256, label="📏 Max Tokens")
     
     with gr.Tabs():
         with gr.TabItem("💬 Chat"):
@@ -145,16 +147,16 @@ with gr.Blocks(title="GOD Coding Machine", theme=gr.themes.Soft(primary_hue="pur
                     fix_btn = gr.Button("Fix", variant="primary")
                 fix_output = gr.Markdown()
     
-    def respond(message, history, temp):
+    def respond(message, history, temp, tokens):
         history = history or []
-        for chunk in chat_stream(message, history, temp):
+        for chunk in chat_stream(message, history, temp, tokens):
             yield history + [[message, chunk]], ""
     
-    msg.submit(respond, [msg, chatbot, temperature], [chatbot, msg])
-    send.click(respond, [msg, chatbot, temperature], [chatbot, msg])
+    msg.submit(respond, [msg, chatbot, temperature, max_tokens], [chatbot, msg])
+    send.click(respond, [msg, chatbot, temperature, max_tokens], [chatbot, msg])
     clear.click(lambda: [], None, chatbot)
-    gen_btn.click(generate_code, [gen_prompt, gen_lang], gen_output)
-    explain_btn.click(explain_code, explain_input, explain_output)
-    fix_btn.click(fix_code, [fix_input, fix_error], fix_output)
+    gen_btn.click(generate_code, [gen_prompt, gen_lang, max_tokens], gen_output)
+    explain_btn.click(explain_code, [explain_input, max_tokens], explain_output)
+    fix_btn.click(fix_code, [fix_input, fix_error, max_tokens], fix_output)
 
 demo.launch(server_name="0.0.0.0", server_port=7860)