import os
import chainlit as cl
from huggingface_hub import InferenceClient

# Make sure you set HF_API_KEY in your environment or Spaces Secrets.
# E.g., os.environ["HF_API_KEY"] = "hf_xxxxxxxxxxxxxx"

@cl.on_message
async def main(message: str):
    hf_api_key = os.getenv("HF_API_KEY", "")
    if not hf_api_key:
        await cl.Message(
            content="Error: No HF API key found. Please set HF_API_KEY."
        ).send()
        return

    client = InferenceClient(provider="together", api_key=hf_api_key)

    messages = [
        {
            "role": "user",
            "content": message
        }
    ]

    try:
        completion = client.chat.completions.create(
            model="deepseek-ai/DeepSeek-R1",
            messages=messages,
            max_tokens=500
        )
        response = completion.choices[0].message.content
        await cl.Message(content=response).send()
    except Exception as e:
        await cl.Message(content=f"Error: {str(e)}").send()