app.py

import gradio as gr
import logging
import os
from huggingface_hub import InferenceClient
from datetime import datetime
import uuid
import json

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("chatbot_logs.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger("CompanyChatbot")

# Environment variables (for production use)
HF_MODEL = os.environ.get("HF_MODEL", "HuggingFaceH4/zephyr-7b-beta")
HF_API_TOKEN = os.environ.get("HF_API_TOKEN", None)  # Set your API token as env variable
COMPANY_NAME = os.environ.get("COMPANY_NAME", "Your Company")
DEFAULT_SYSTEM_PROMPT = os.environ.get("DEFAULT_SYSTEM_PROMPT", 
    f"You are {COMPANY_NAME}'s professional AI assistant. Be helpful, accurate, and concise.")

# Initialize the client
try:
    client = InferenceClient(HF_MODEL, token=HF_API_TOKEN)
    logger.info(f"Successfully initialized InferenceClient with model: {HF_MODEL}")
except Exception as e:
    logger.error(f"Failed to initialize InferenceClient: {str(e)}")
    raise RuntimeError(f"Failed to initialize the model. Please check your configuration: {str(e)}")

# Conversation tracking
def save_conversation(user_id, conversation):
    filename = f"conversations/{user_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
    os.makedirs(os.path.dirname(filename), exist_ok=True)
    with open(filename, 'w') as f:
        json.dump(conversation, f)
    logger.info(f"Saved conversation for user {user_id}")

# Main chat function
def respond(
    message, 
    history: list[tuple[str, str]], 
    system_message, 
    max_tokens, 
    temperature, 
    top_p,
    user_id
):
    if not message.strip():
        return "I'm sorry, I didn't receive any input. How can I help you today?"
    
    # Log the incoming request
    logger.info(f"User {user_id} sent message - Length: {len(message)}")
    
    try:
        messages = [{"role": "system", "content": system_message}]
        
        # Build conversation history
        for user_msg, assistant_msg in history:
            if user_msg:
                messages.append({"role": "user", "content": user_msg})
            if assistant_msg:
                messages.append({"role": "assistant", "content": assistant_msg})
        
        messages.append({"role": "user", "content": message})
        
        # Generate response
        full_response = ""
        start_time = datetime.now()
        
        for message_chunk in client.chat_completion(
            messages, 
            max_tokens=max_tokens, 
            stream=True, 
            temperature=temperature, 
            top_p=top_p,
        ):
            token = message_chunk.choices[0].delta.content
            full_response += token if token else ""
            yield full_response
        
        # Log completion
        time_taken = (datetime.now() - start_time).total_seconds()
        logger.info(f"Response generated for user {user_id} in {time_taken:.2f}s - Length: {len(full_response)}")
        
        # Save conversation for audit/analytics
        conversation_data = {
            "timestamp": datetime.now().isoformat(),
            "user_id": user_id,
            "messages": messages,
            "response": full_response,
            "parameters": {
                "max_tokens": max_tokens,
                "temperature": temperature,
                "top_p": top_p
            },
            "time_taken": time_taken
        }
        save_conversation(user_id, conversation_data)
        
    except Exception as e:
        error_msg = f"An error occurred: {str(e)}"
        logger.error(f"Error generating response for user {user_id}: {str(e)}")
        return error_msg

# Authentication function (replace with your actual auth system)
def authenticate(username, password):
    # In production, this should check against your company's auth system
    valid_credentials = {"admin": "admin123", "user": "user123"}  # Example only
    
    if username in valid_credentials and valid_credentials[username] == password:
        return True, str(uuid.uuid4())  # Generate user session ID
    return False, None

# Login interface
def login(username, password):
    success, user_id = authenticate(username, password)
    if success:
        return gr.update(visible=False), gr.update(visible=True), user_id
    else:
        return gr.update(visible=True), gr.update(visible=False), None

# Main application
with gr.Blocks(css="styles.css", title=f"{COMPANY_NAME} AI Assistant") as demo:
    user_id = gr.State(None)
    
    with gr.Row():
        gr.Markdown(f"# {COMPANY_NAME} AI Assistant")
    
    with gr.Group(visible=True) as login_group:
        gr.Markdown("### Please log in to continue")
        username = gr.Textbox(label="Username")
        password = gr.Textbox(label="Password", type="password")
        login_button = gr.Button("Login")
    
    with gr.Group(visible=False) as chat_group:
        chatbot = gr.ChatInterface(
            respond,
            additional_inputs=[
                gr.Textbox(value=DEFAULT_SYSTEM_PROMPT, label="System Instructions"),
                gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max Response Length"),
                gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature (Creativity)"),
                gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (Variation)"),
                user_id
            ],
            analytics_enabled=True,
            title=None,
        )
    
    login_button.click(
        login,
        inputs=[username, password],
        outputs=[login_group, chat_group, user_id]
    )

# For CSS styling
css = """
body {
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
    background-color: #f9f9f9;
}
.gradio-container {
    max-width: 1200px !important;
    margin: auto;
}
.footer {
    text-align: center;
    margin-top: 20px;
    color: #666;
    font-size: 0.8em;
}
"""

with open("styles.css", "w") as f:
    f.write(css)

if __name__ == "__main__":
    # Check if we're running in production
    if os.environ.get("PRODUCTION", "false").lower() == "true":
        demo.launch(
            server_name="0.0.0.0",
            server_port=int(os.environ.get("PORT", 7860)),
            share=False,
            show_error=False,
            auth=None,  # We handle auth in the app
        )
    else:
        # Development mode
        demo.launch(share=True)