Update app.py

app.py

@@ -1,3 +1,4 @@
+import os
 import re
 from http import HTTPStatus
 from typing import Dict, List, Optional, Tuple
@@ -15,12 +16,10 @@ from bs4 import BeautifulSoup
 import html2text
 import json
 import time
-import os
 
 import gradio as gr
 from huggingface_hub import InferenceClient
 from tavily import TavilyClient
-import openai
 
 # Gradio supported languages for syntax highlighting
 GRADIO_SUPPORTED_LANGUAGES = [
@@ -115,41 +114,45 @@ Removing the paragraph...
 
 # Available models
 AVAILABLE_MODELS = [
+    {
+        "name": "Moonshot Kimi-K2",
+        "id": "moonshotai/Kimi-K2-Instruct",
+        "description": "Moonshot AI Kimi-K2-Instruct model for code generation and general tasks"
+    },
     {
         "name": "DeepSeek V3",
         "id": "deepseek-ai/DeepSeek-V3-0324",
-        "description": "DeepSeek V3 model for code generation", "provider": "deepseek"
+        "description": "DeepSeek V3 model for code generation"
     },
     {
-        "name": "DeepSeek R1",
+        "name": "DeepSeek R1", 
         "id": "deepseek-ai/DeepSeek-R1-0528",
-        "description": "DeepSeek R1 model for code generation", "provider": "deepseek"
+        "description": "DeepSeek R1 model for code generation"
     },
     {
-    #    "name": "GPT-4o-mini",
-    #    "id": "gpt-4o-mini",
-    #    "description": "OpenAI GPT-4o-mini model for code generation and general tasks",
-    #    "provider": "openai"
-    #},  # Commented out to avoid requiring extra setup for all users
-    #{
-        "name": "Gemini 1.5 Pro",
-        "id": "gemini-1.5-pro-latest",
-        "description": "Google Gemini 1.5 Pro model for code generation and general tasks", "provider": "gemini"
+        "name": "ERNIE-4.5-VL",
+        "id": "baidu/ERNIE-4.5-VL-424B-A47B-Base-PT",
+        "description": "ERNIE-4.5-VL model for multimodal code generation with image support"
     },
     {
-        "name": "GPT-4o",
-        "id": "gpt-4o",
-        "description": "OpenAI GPT-4o model for code generation and general tasks", "provider": "openai"
+        "name": "MiniMax M1",
+        "id": "MiniMaxAI/MiniMax-M1-80k",
+        "description": "MiniMax M1 model for code generation and general tasks"
     },
     {
         "name": "Qwen3-235B-A22B",
         "id": "Qwen/Qwen3-235B-A22B",
-        "description": "Qwen3-235B-A22B model for code generation and general tasks", "provider": "huggingface"
+        "description": "Qwen3-235B-A22B model for code generation and general tasks"
     },
     {
-       "name": "GLM-4.1V-9B-Thinking",
-       "id": "THUDM/GLM-4.1V-9B-Thinking",
-       "description": "GLM-4.1V-9B-Thinking model for multimodal code generation with image support",
+        "name": "SmolLM3-3B",
+        "id": "HuggingFaceTB/SmolLM3-3B",
+        "description": "SmolLM3-3B model for code generation and general tasks"
+    },
+    {
+        "name": "GLM-4.1V-9B-Thinking",
+        "id": "THUDM/GLM-4.1V-9B-Thinking",
+        "description": "GLM-4.1V-9B-Thinking model for multimodal code generation with image support"
     }
 ]
 
@@ -213,51 +216,16 @@ DEMO_LIST = [
 ]
 
 # HF Inference Client
-HF_TOKEN = os.getenv("HF_TOKEN")
-OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
-GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
-DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")
+HF_TOKEN = os.getenv('HF_TOKEN')
 
 def get_inference_client(model_id):
-   """Return an appropriate client based on model_id and configured provider."""
-   model_info = next((m for m in AVAILABLE_MODELS if m["id"] == model_id), None)
-   if not model_info:
-       raise ValueError(f"Model with id '{model_id}' not found.")
-
-   provider = model_info.get("provider")
-
-   if provider == "deepseek":
-       if not DEEPSEEK_API_KEY:
-           raise ValueError("DEEPSEEK_API_KEY environment variable not set.")
-       return "deepseek_client" # Return a string placeholder, logic is handled in generation_code
-
-   if provider == "groq":
-       return InferenceClient(
-           provider="groq",
-           api_key=HF_TOKEN,  # Groq uses HF token as well
-           bill_to="huggingface"  # Assuming billing through HF for Groq
-       )
-   elif provider == "huggingface":
-       if not HF_TOKEN:
-           raise ValueError("HF_TOKEN environment variable not set for Hugging Face models.")
-       return InferenceClient(
-           provider="auto",  # Use HF's auto provider for other HF models
-           api_key=HF_TOKEN,
-           bill_to="huggingface"
-       )
-   elif provider == "openai":
-       if not OPENAI_API_KEY:
-           raise ValueError("OPENAI_API_KEY environment variable not set.")
-       openai.api_key = OPENAI_API_KEY
-       return openai
-   elif provider == "gemini":
-       if not GEMINI_API_KEY:
-           raise ValueError("GEMINI_API_KEY environment variable not set.")
-       import google.generativeai as genai
-       genai.configure(api_key=GEMINI_API_KEY)
-       return genai
-   else:
-       raise ValueError(f"Unsupported provider: {provider}")
+    """Return an InferenceClient with provider based on model_id."""
+    provider = "groq" if model_id == "moonshotai/Kimi-K2-Instruct" else "auto"
+    return InferenceClient(
+        provider=provider,
+        api_key=HF_TOKEN,
+        bill_to="huggingface"
+    )
 
 # Type definitions
 History = List[Tuple[str, str]]
@@ -283,7 +251,7 @@ def history_to_messages(history: History, system: str) -> Messages:
             text_content = ""
             for item in user_content:
                 if isinstance(item, dict) and item.get("type") == "text":
-                    text_content += str(item.get("text", ""))
+                    text_content += item.get("text", "")
             user_content = text_content if text_content else str(user_content)
         
         messages.append({'role': 'user', 'content': user_content})
@@ -300,7 +268,7 @@ def messages_to_history(messages: Messages) -> Tuple[str, History]:
             text_content = ""
             for item in user_content:
                 if isinstance(item, dict) and item.get("type") == "text":
-                    text_content += str(item.get("text", ""))
+                    text_content += item.get("text", "")
             user_content = text_content if text_content else str(user_content)
         
         history.append([user_content, r['content']])
@@ -315,7 +283,7 @@ def history_to_chatbot_messages(history: History) -> List[Dict[str, str]]:
             text_content = ""
             for item in user_msg:
                 if isinstance(item, dict) and item.get("type") == "text":
-                    text_content += str(item.get("text", ""))
+                    text_content += item.get("text", "")
             user_msg = text_content if text_content else str(user_msg)
         
         messages.append({"role": "user", "content": user_msg})
@@ -1011,139 +979,41 @@ This will help me create a better design for you."""
     else:
         messages.append({'role': 'user', 'content': enhanced_query})
     try:
-        if _current_model["provider"] == "openai":
-            # Use OpenAI client
-            stream = client.chat.completions.create(
-                model=_current_model["id"],
-                messages=[{"role": m["role"], "content": m["content"]} for m in messages],
-                stream=True,
-                max_tokens=5000,
-            )
-
-            content = ""
-            for chunk in stream:
-                if chunk.choices[0].delta.content:
-                    content += chunk.choices[0].delta.content
-                    clean_code = remove_code_block(content)
-                    if has_existing_html:
-                        # ... (rest of the modification logic) ...
-                        if clean_code.strip().startswith("<!DOCTYPE html>") or clean_code.strip().startswith("<html"):
-                            yield {
-                                code_output: gr.update(value=clean_code, language=get_gradio_language(language)),
-                                history_output: history_to_chatbot_messages(_history),
-                                sandbox: send_to_sandbox(clean_code) if language == "html" else "<div style='padding:1em;color:#888;text-align:center;'>Preview is only available for HTML. Please download your code using the download button above.</div>",
-                            }
-                        else:
-                            last_html = _history[-1][1] if _history else ""
-                            modified_html = apply_search_replace_changes(last_html, clean_code)
-                            clean_html = remove_code_block(modified_html)
-                            yield {
-                                code_output: gr.update(value=clean_html, language=get_gradio_language(language)),
-                                history_output: history_to_chatbot_messages(_history),
-                                sandbox: send_to_sandbox(clean_html) if language == "html" else "<div style='padding:1em;color:#888;text-align:center;'>Preview is only available for HTML. Please download your code using the download button above.</div>",
-                            }
-                    else:
+        completion = client.chat.completions.create(
+            model=_current_model["id"],
+            messages=messages,
+            stream=True,
+            max_tokens=5000
+        )
+        content = ""
+        for chunk in completion:
+            if chunk.choices[0].delta.content:
+                content += chunk.choices[0].delta.content
+                clean_code = remove_code_block(content)
+                search_status = " (with web search)" if enable_search and tavily_client else ""
+                if has_existing_html:
+                    # Fallback: If the model returns a full HTML file, use it directly
+                    if clean_code.strip().startswith("<!DOCTYPE html>") or clean_code.strip().startswith("<html"):
                         yield {
                             code_output: gr.update(value=clean_code, language=get_gradio_language(language)),
                             history_output: history_to_chatbot_messages(_history),
                             sandbox: send_to_sandbox(clean_code) if language == "html" else "<div style='padding:1em;color:#888;text-align:center;'>Preview is only available for HTML. Please download your code using the download button above.</div>",
                         }
-        elif _current_model["provider"] == "gemini":
-            # Use Gemini client (assumes you've adapted the message format)
-            chat = client.GenerativeModel(_current_model["id"]).start_chat(history=[])  # Assuming it uses history
-            for message in messages[1:]:
-                if message["role"] == "user":
-                    chat.send_message(message["content"])
-                elif message["role"] == "assistant":
-                    chat.send_message(message["content"])
-
-            response = chat.send_message(messages[-1]["content"], stream=True)
-            content = ""
-            for chunk in response:
-                content += chunk.text
-                clean_code = remove_code_block(content)
-                if has_existing_html:
-                   # ... (rest of the modification logic) ...
-                   if clean_code.strip().startswith("<!DOCTYPE html>") or clean_code.strip().startswith("<html"):
-                       yield {
-                           code_output: gr.update(value=clean_code, language=get_gradio_language(language)),
-                           history_output: history_to_chatbot_messages(_history),
-                           sandbox: send_to_sandbox(clean_code) if language == "html" else "<div style='padding:1em;color:#888;text-align:center;'>Preview is only available for HTML. Please download your code using the download button above.</div>",
-                       }
-                   else:
-                       last_html = _history[-1][1] if _history else ""
-                       modified_html = apply_search_replace_changes(last_html, clean_code)
-                       clean_html = remove_code_block(modified_html)
-                       yield {
-                           code_output: gr.update(value=clean_html, language=get_gradio_language(language)),
-                           history_output: history_to_chatbot_messages(_history),
-                           sandbox: send_to_sandbox(clean_html) if language == "html" else "<div style='padding:1em;color:#888;text-align:center;'>Preview is only available for HTML. Please download your code using the download button above.</div>",
-                       }
+                    else:
+                        last_html = _history[-1][1] if _history else ""
+                        modified_html = apply_search_replace_changes(last_html, clean_code)
+                        clean_html = remove_code_block(modified_html)
+                        yield {
+                            code_output: gr.update(value=clean_html, language=get_gradio_language(language)),
+                            history_output: history_to_chatbot_messages(_history),
+                            sandbox: send_to_sandbox(clean_html) if language == "html" else "<div style='padding:1em;color:#888;text-align:center;'>Preview is only available for HTML. Please download your code using the download button above.</div>",
+                        }
                 else:
-                   yield {
-                       code_output: gr.update(value=clean_code, language=get_gradio_language(language)),
-                       history_output: history_to_chatbot_messages(_history),
-                       sandbox: send_to_sandbox(clean_code) if language == "html" else "<div style='padding:1em;color:#888;text-align:center;'>Preview is only available for HTML. Please download your code using the download button above.</div>",
-                   }
-        elif _current_model["provider"] == "deepseek":
-            # Use requests to call Deepseek API directly
-            url = "https://api.deepseek.com/v1/chat/completions"
-            headers = {
-                "Authorization": f"Bearer {DEEPSEEK_API_KEY}",
-                "Content-Type": "application/json"
-            }
-            payload = {
-                "model": _current_model["id"],  # Use the full model ID from AVAILABLE_MODELS
-                "messages": [{"role": m["role"], "content": m["content"]} for m in messages],
-                "stream": True,
-                "max_tokens": 5000
-            }
-            
-            response = requests.post(url, headers=headers, json=payload, stream=True)
-            response.raise_for_status()
-            
-            content = ""
-            for line in response.iter_lines():
-                if line:
-                    line_str = line.decode('utf-8')
-                    if line_str.startswith('data: '):
-                        json_data = line_str[6:]
-                        if json_data.strip() != '[DONE]':
-                            chunk = json.loads(json_data)
-                            if chunk['choices'][0]['delta'].get('content'):
-                                content += chunk['choices'][0]['delta']['content']
-                                clean_code = remove_code_block(content)
-                                yield {
-                                    code_output: gr.update(value=clean_code, language=get_gradio_language(language)),
-                                    history_output: history_to_chatbot_messages(_history),
-                                    sandbox: send_to_sandbox(clean_code) if language == "html" else "<div style='padding:1em;color:#888;text-align:center;'>Preview is only available for HTML. Please download your code using the download button above.</div>",
-                                }
-
-        else:
-            # Use Hugging Face Inference Client (or Groq) for other models
-            completion = client.chat.completions.create(
-                model=_current_model["id"],
-                messages=messages,
-                stream=True,
-                max_tokens=5000
-            )
-
-            content = ""
-            for chunk in completion:
-                if chunk.choices[0].delta.content:
-                    content += chunk.choices[0].delta.content
-                    clean_code = remove_code_block(content)
-                    if has_existing_html:
-                        # ... (rest of the modification logic) ...
-                        if clean_code.strip().startswith("<!DOCTYPE html>") or clean_code.strip().startswith("<html"):
-                            yield {
-                                code_output: gr.update(value=clean_code, language=get_gradio_language(language)),
-                                history_output: history_to_chatbot_messages(_history),
-                                sandbox: send_to_sandbox(clean_code) if language == "html" else "<div style='padding:1em;color:#888;text-align:center;'>Preview is only available for HTML. Please download your code using the download button above.</div>",
-                            }
-                        else:
-                            last_html = _history[-1][1] if _history else ""
-                            modified_html = apply_search_replace_changes(last_html, clean_code)
+                    yield {
+                        code_output: gr.update(value=clean_code, language=get_gradio_language(language)),
+                        history_output: history_to_chatbot_messages(_history),
+                        sandbox: send_to_sandbox(clean_code) if language == "html" else "<div style='padding:1em;color:#888;text-align:center;'>Preview is only available for HTML. Please download your code using the download button above.</div>",
+                    }
         # Handle response based on whether this is a modification or new generation
         if has_existing_html:
             # Fallback: If the model returns a full HTML file, use it directly
@@ -1272,12 +1142,10 @@ with gr.Blocks(
                     outputs=input
                 )
         if not tavily_client:
-            tavily_status = gr.Markdown("⚠️ Web search unavailable. Set TAVILY_API_KEY.", visible=True)
+            gr.Markdown("⚠️ Web search unavailable", visible=True)
         else:
-            tavily_status = gr.Markdown("✅ Web search available", visible=True)
+            gr.Markdown("✅ Web search available", visible=True)
         model_display = gr.Markdown(f"**Model:** {AVAILABLE_MODELS[0]['name']}", visible=True)  # Moonshot Kimi-K2
-
-
         def on_model_change(model_name):
             for m in AVAILABLE_MODELS:
                 if m['name'] == model_name: