app.py

import gradio as gr
from huggingface_hub import InferenceClient
import os
import json
import base64
from PIL import Image
import io

# Load the default access token from environment variable at startup
# This will be used if no custom key is provided by the user.
ACCESS_TOKEN = os.getenv("HF_TOKEN")
print(f"Default HF_TOKEN from environment loaded: {'Present' if ACCESS_TOKEN else 'Not set'}")

# Function to encode image to base64
def encode_image(image_path):
    if not image_path:
        print("No image path provided")
        return None
    
    try:
        print(f"Encoding image from path: {image_path}")
        
        if isinstance(image_path, Image.Image):
            image = image_path
        else:
            image = Image.open(image_path)
        
        if image.mode == 'RGBA':
            image = image.convert('RGB')
        
        buffered = io.BytesIO()
        image.save(buffered, format="JPEG")
        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
        print("Image encoded successfully")
        return img_str
    except Exception as e:
        print(f"Error encoding image: {e}")
        return None

def respond(
    message,
    image_files,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
    frequency_penalty,
    seed,
    provider,
    custom_api_key, # This is the value from the BYOK textbox
    custom_model,    
    model_search_term,
    selected_model
):
    print(f"Received message: {message}")
    print(f"Received {len(image_files) if image_files else 0} images")
    # print(f"History: {history}") # Can be very verbose
    print(f"System message: {system_message}")
    print(f"Max tokens: {max_tokens}, Temperature: {temperature}, Top-P: {top_p}")
    print(f"Frequency Penalty: {frequency_penalty}, Seed: {seed}")
    print(f"Selected provider: {provider}")         
    print(f"Custom API Key input field value (raw): '{custom_api_key[:10]}...' (masked if long)")
    print(f"Selected model (custom_model input field): {custom_model}")  
    print(f"Model search term: {model_search_term}")
    print(f"Selected model from radio: {selected_model}")

    token_to_use = None
    original_hf_token_env_value = os.environ.get("HF_TOKEN")
    env_hf_token_temporarily_modified = False

    if custom_api_key and custom_api_key.strip():
        token_to_use = custom_api_key.strip()
        print(f"USING CUSTOM API KEY (BYOK): '{token_to_use[:5]}...' (masked for security).")
        # Aggressively ensure custom key is fundamental:
        # Temporarily remove HF_TOKEN from os.environ if it exists,
        # to prevent any possibility of InferenceClient picking it up.
        if "HF_TOKEN" in os.environ:
            print(f"Temporarily unsetting HF_TOKEN from environment (was: {'Present' if os.environ.get('HF_TOKEN') else 'Not set'}) to prioritize custom key.")
            del os.environ["HF_TOKEN"]
            env_hf_token_temporarily_modified = True
    elif ACCESS_TOKEN: # Use default token from environment if no custom key
        token_to_use = ACCESS_TOKEN
        print(f"USING DEFAULT API KEY (HF_TOKEN from environment variable at script start): '{token_to_use[:5]}...' (masked for security).")
        # Ensure HF_TOKEN is set in the current env if it was loaded at start
        # This handles cases where it might have been unset by a previous call with a custom key
        if original_hf_token_env_value is not None:
            os.environ["HF_TOKEN"] = original_hf_token_env_value
        elif "HF_TOKEN" in os.environ: # If ACCESS_TOKEN was loaded but original_hf_token_env_value was None (e.g. set by other means)
             pass # Let it be whatever it is
    else:
        print("No custom API key provided AND no default HF_TOKEN was found in environment at script start.")
        print("InferenceClient will be initialized without an explicit token. May fail or use public access.")
        # token_to_use remains None
        # If HF_TOKEN was in env and we want to ensure it's not used when token_to_use is None:
        if "HF_TOKEN" in os.environ:
            print(f"Temporarily unsetting HF_TOKEN from environment (was: {'Present' if os.environ.get('HF_TOKEN') else 'Not set'}) as no valid key is chosen.")
            del os.environ["HF_TOKEN"]
            env_hf_token_temporarily_modified = True # Mark for restoration

    print(f"Final token being passed to InferenceClient: '{str(token_to_use)[:5]}...' (masked)" if token_to_use else "None")

    try:
        client = InferenceClient(token=token_to_use, provider=provider)
        print(f"Hugging Face Inference Client initialized with {provider} provider.")

        if seed == -1:
            seed = None

        user_content = []
        if message and message.strip():
            user_content.append({"type": "text", "text": message})
        
        if image_files:
            for img_path in image_files:
                if img_path:
                    encoded_image = encode_image(img_path)
                    if encoded_image:
                        user_content.append({
                            "type": "image_url",
                            "image_url": {"url": f"data:image/jpeg;base64,{encoded_image}"}
                        })
        
        if not user_content: # If only images were sent and none encoded, or empty message
             if image_files: # If there were image files, it implies an image-only message
                 user_content = [{"type": "text", "text": ""}] # Send an empty text for context, or specific prompt
             else: # Truly empty input
                 yield "Error: Empty message content."
                 return


        messages = [{"role": "system", "content": system_message}]
        for val in history:
            user_part, assistant_part = val
            # Handle multimodal history if necessary (simplified for now)
            if isinstance(user_part, dict) and 'files' in user_part: # from MultimodalTextbox
                history_text = user_part.get("text", "")
                history_files = user_part.get("files", [])
                current_user_content_history = []
                if history_text:
                    current_user_content_history.append({"type": "text", "text": history_text})
                for h_img_path in history_files:
                    encoded_h_img = encode_image(h_img_path)
                    if encoded_h_img:
                        current_user_content_history.append({
                            "type": "image_url",
                            "image_url": {"url": f"data:image/jpeg;base64,{encoded_h_img}"}
                        })
                if current_user_content_history:
                     messages.append({"role": "user", "content": current_user_content_history})
            elif isinstance(user_part, str): # from simple text history
                 messages.append({"role": "user", "content": user_part})

            if assistant_part:
                messages.append({"role": "assistant", "content": assistant_part})

        messages.append({"role": "user", "content": user_content if len(user_content) > 1 or not isinstance(user_content[0], dict) or user_content[0].get("type") != "text" else user_content[0]["text"]})


        model_to_use = custom_model.strip() if custom_model.strip() else selected_model
        print(f"Model selected for inference: {model_to_use}")

        response_text = ""
        print(f"Sending request to {provider} with model {model_to_use}.")

        parameters = {
            "max_tokens": max_tokens,
            "temperature": temperature,
            "top_p": top_p,
            "frequency_penalty": frequency_penalty,
        }
        if seed is not None:
            parameters["seed"] = seed

        stream = client.chat_completion(
            model=model_to_use,
            messages=messages,
            stream=True,
            **parameters
        )
        
        print("Streaming response: ", end="", flush=True)
        for chunk in stream:
            if hasattr(chunk, 'choices') and chunk.choices:
                delta = chunk.choices[0].delta
                if hasattr(delta, 'content') and delta.content:
                    token_chunk = delta.content
                    print(token_chunk, end="", flush=True)
                    response_text += token_chunk
                    yield response_text
        print("\nStream finished.")
    
    except Exception as e:
        error_message = f"Error during inference: {e}"
        print(error_message)
        # If there was already some response, append error. Otherwise, yield error.
        if 'response_text' in locals() and response_text:
             response_text += f"\n{error_message}"
             yield response_text
        else:
             yield error_message
    finally:
        # Restore HF_TOKEN in os.environ if it was temporarily removed/modified
        if env_hf_token_temporarily_modified:
            if original_hf_token_env_value is not None:
                os.environ["HF_TOKEN"] = original_hf_token_env_value
                print("Restored HF_TOKEN in environment from its original value.")
            else:
                # If it was unset and originally not present, ensure it remains unset
                if "HF_TOKEN" in os.environ: # Should not happen if original was None and we deleted
                    del os.environ["HF_TOKEN"]
                print("HF_TOKEN was originally not set and was temporarily removed; ensuring it remains not set in env.")
        print("Response generation attempt complete.")


def validate_provider(api_key, provider_choice):
    # This validation might need adjustment based on providers.
    # For now, it assumes any custom key might work with other providers.
    # If HF_TOKEN is the only one available (no custom key), restrict to hf-inference.
    if not api_key.strip() and provider_choice != "hf-inference" and ACCESS_TOKEN:
        gr.Warning("Default HF_TOKEN can only be used with 'hf-inference' provider. Switching to 'hf-inference'.")
        return gr.update(value="hf-inference")
    return gr.update(value=provider_choice)

with gr.Blocks(theme="Nymbo/Nymbo_Theme") as demo:
    chatbot = gr.Chatbot(
        height=600, 
        show_copy_button=True, 
        placeholder="Select a model and begin chatting. Supports multimodal inputs.",
        layout="panel",
        avatar_images=(None, "https://hf.co/front/assets/huggingface_logo.svg") # Bot avatar
    )
    
    msg = gr.MultimodalTextbox(
        placeholder="Type a message or upload images...",
        show_label=False,
        container=False,
        scale=12,
        file_types=["image"],
        file_count="multiple",
        sources=["upload"]
    )
    
    with gr.Accordion("Settings", open=False):
        system_message_box = gr.Textbox(
            value="You are a helpful AI assistant that can understand images and text.", 
            placeholder="You are a helpful assistant.",
            label="System Prompt"
        )
        
        with gr.Row():
            with gr.Column():
                max_tokens_slider = gr.Slider(minimum=1, maximum=4096, value=1024, step=1, label="Max tokens")
                temperature_slider = gr.Slider(minimum=0.0, maximum=2.0, value=0.7, step=0.01, label="Temperature") # Allow 0 for deterministic
                top_p_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.95, step=0.01, label="Top-P") # Allow 0
            with gr.Column():
                frequency_penalty_slider = gr.Slider(minimum=-2.0, maximum=2.0, value=0.0, step=0.1, label="Frequency Penalty")
                seed_slider = gr.Slider(minimum=-1, maximum=65535, value=-1, step=1, label="Seed (-1 for random)")
        
        providers_list = ["hf-inference", "cerebras", "together", "sambanova", "novita", "cohere", "fireworks-ai", "hyperbolic", "nebius"]
        provider_radio = gr.Radio(choices=providers_list, value="hf-inference", label="Inference Provider")
        
        byok_textbox = gr.Textbox(
            value="", label="BYOK (Bring Your Own Key)", 
            info="Enter your Hugging Face API key (or provider-specific key). Overrides default. If empty, uses Space's HF_TOKEN (if set) for 'hf-inference'.",
            placeholder="hf_... or provider_specific_key", type="password"
        )
        
        custom_model_box = gr.Textbox(
            value="", label="Custom Model ID", 
            info="(Optional) Provide a model ID (e.g., 'meta-llama/Llama-3-8B-Instruct'). Overrides featured model selection.",
            placeholder="org/model-name"
        )
        
        model_search_box = gr.Textbox(label="Filter Featured Models", placeholder="Search...", lines=1)
        
        models_list = [
            "meta-llama/Llama-3.2-11B-Vision-Instruct", "meta-llama/Llama-3.1-70B-Instruct", 
            "mistralai/Mistral-Nemo-Instruct-2407", "Qwen/Qwen2.5-72B-Instruct",
            "microsoft/Phi-3.5-mini-instruct", "NousResearch/Hermes-3-Llama-3.1-8B",
            # Add more or fetch dynamically if possible
        ]
        featured_model_radio = gr.Radio(
            label="Select a Featured Model", choices=models_list, 
            value="meta-llama/Llama-3.2-11B-Vision-Instruct", interactive=True
        )
        
        gr.Markdown("[All Text Gen Models](https://huggingface.co/models?pipeline_tag=text-generation&sort=trending) | [All Multimodal Models](https://huggingface.co/models?pipeline_tag=image-text-to-text&sort=trending)")

    # Chat history state (using chatbot component directly for history)

    def handle_user_message_submission(user_input_mmtb, chat_history_list):
        # user_input_mmtb is a dict: {"text": "...", "files": ["path1", "path2"]}
        text_content = user_input_mmtb.get("text", "")
        files = user_input_mmtb.get("files", [])
        
        # Construct the display for the user message in the chat
        # For Gradio Chatbot, user message can be a string or a tuple (text, filepath) or (None, filepath)
        # If multiple files, they need to be sent as separate messages or handled in display
        
        if not text_content and not files:
            return chat_history_list # Or raise an error/warning

        # Append user message to history.
        # The actual content for the API will be constructed in respond()
        # For display, we can show text and a placeholder for images, or actual images if supported well.
        # Let's pass the raw MultimodalTextbox output to history for now.
        chat_history_list.append([user_input_mmtb, None])
        return chat_history_list

    def handle_bot_response_generation(
        chat_history_list, system_msg, max_tokens, temp, top_p, freq_pen, seed_val, 
        prov, api_key_val, cust_model_val, search_term_val, feat_model_val
    ):
        if not chat_history_list or chat_history_list[-1][0] is None:
            yield chat_history_list # Or an error message
            return

        # The last user message is chat_history_list[-1][0]
        # It's the dict from MultimodalTextbox: {"text": "...", "files": ["path1", ...]}
        last_user_input_mmtb = chat_history_list[-1][0]
        
        current_message_text = last_user_input_mmtb.get("text", "")
        current_image_files = last_user_input_mmtb.get("files", [])

        # Prepare history for the `respond` function (excluding the current turn's user message)
        api_history = []
        for user_msg_item, bot_msg_item in chat_history_list[:-1]:
            # Convert past user messages (which are MMTB dicts) to API format or simple strings
            past_user_text = user_msg_item.get("text", "")
            # For simplicity, not including past images in API history here, but could be added
            api_history.append((past_user_text, bot_msg_item))


        # Stream the response
        full_response = ""
        for_stream_chunk in respond(
            message=current_message_text,
            image_files=current_image_files,
            history=api_history, # Pass the processed history
            system_message=system_msg,
            max_tokens=max_tokens,
            temperature=temp,
            top_p=top_p,
            frequency_penalty=freq_pen,
            seed=seed_val,
            provider=prov,
            custom_api_key=api_key_val,
            custom_model=cust_model_val,
            model_search_term=search_term_val, # Note: search_term is for UI filtering, not API
            selected_model=feat_model_val
        ):
            full_response = for_stream_chunk
            chat_history_list[-1][1] = full_response
            yield chat_history_list
            
    msg.submit(
        handle_user_message_submission,
        [msg, chatbot],
        [chatbot],
        queue=False
    ).then(
        handle_bot_response_generation,
        [chatbot, system_message_box, max_tokens_slider, temperature_slider, top_p_slider, 
         frequency_penalty_slider, seed_slider, provider_radio, byok_textbox, custom_model_box, 
         model_search_box, featured_model_radio],
        [chatbot]
    ).then(
        lambda: gr.update(value=None),  # Clears MultimodalTextbox: {"text": None, "files": None}
        [], # No inputs needed for this
        [msg]
    )
    
    def filter_models_ui(search_term):
        filtered = [m for m in models_list if search_term.lower() in m.lower()] if search_term else models_list
        return gr.update(choices=filtered, value=filtered[0] if filtered else None)

    model_search_box.change(fn=filter_models_ui, inputs=model_search_box, outputs=featured_model_radio)
    
    # No need for set_custom_model_from_radio if custom_model_box overrides featured_model_radio directly in respond()

    byok_textbox.change(fn=validate_provider, inputs=[byok_textbox, provider_radio], outputs=provider_radio)
    provider_radio.change(fn=validate_provider, inputs=[byok_textbox, provider_radio], outputs=provider_radio)

print("Gradio interface initialized.")

if __name__ == "__main__":
    print("Launching the demo application.")
    # ForSpaces, share=True is often implied or handled by Spaces platform
    # For local, share=True makes it public via Gradio link
    demo.queue().launch(show_api=False) # .queue() is good for handling multiple users / long tasks