import gradio as gr
from openai import OpenAI

# Initialize the OpenAI client
client = OpenAI(
    api_key="EMPTY",
    base_url='https://hub.societyai.com/models/llama-3-2-3b/openai/v1',
)

with gr.Blocks(css="footer {visibility: hidden}") as demo:
    chatbot = gr.Chatbot(type="messages")
    msg = gr.Textbox()
    clear = gr.Button("Clear")

    # System message to provide initial context to the conversation
    system_message = {
        "role": "system",
        "content": "You are a helpful AI assistant, capable of answering questions and providing information. Please assist the user in a friendly and informative manner. As an Expert SEO Content Creator, your task is to create a well-optimized, high-quality SEO article for the provided keyword. If the request is unclear, begin by asking for clarification on the keyword or intent. Follow these steps: clarify the keyword and search intent, propose a content structure with headings and estimated word count, explain the importance of external links, include the primary keyword in strategic locations, maintain natural language, use readability enhancements, optimize for featured snippets, suggest a meta description and SEO-friendly URL, propose internal links and image alt text, present the draft for review, and ensure all SEO elements are in place before final delivery. Prioritize clarity, accuracy, and user satisfaction while balancing SEO best practices with engaging content."
}

    def user(user_message, history: list):
        """Appends the user message to the conversation history."""
        if not history:
            # Initialize with system message if history is empty
            history = [system_message]
        return "", history + [{"role": "user", "content": user_message}]

    def bot(history: list):
        """Sends the conversation history to the vLLM API and streams the assistant's response."""
        # Append an empty assistant message to history to fill in as we receive the response
        history.append({"role": "assistant", "content": ""})

        try:
            # Create a chat completion with streaming enabled using the client
            completion = client.chat.completions.create(
                model="llama-3.2-3B-instruct",  # Adjust the model name if needed
                messages=history,
                stream=True
            )

            # Iterate over the streamed response
            for chunk in completion:
                # Access the delta content from the chunk
                delta = chunk.choices[0].delta
                content = getattr(delta, 'content', '')
                if content:
                    # Update the assistant's message with new content
                    history[-1]['content'] += content
                    yield history
        except Exception as e:
            # Handle exceptions and display an error message
            history[-1]['content'] += f"\n[Error]: {str(e)}"
            yield history

    # Set up the Gradio interface components
    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, chatbot, chatbot
    )
    clear.click(lambda: None, None, chatbot, queue=False)

if __name__ == "__main__":
    demo.launch()