from flask import Flask, request, jsonify
from huggingface_hub import InferenceClient
import os

app = Flask(__name__)

# Initialize the InferenceClient
client = InferenceClient(
    token=os.getenv('HUGGING_FACE_API_KEY')  # Make sure to set this environment variable
)

@app.route('/', methods=['GET'])
def generate_text():
    try:
        # Verify API key is set
        if not os.getenv('HUGGING_FACE_API_KEY'):
            return jsonify({'error': 'HUGGING_FACE_API_KEY environment variable is not set'}), 500

        # Get the user's prompt from query parameters
        user_prompt = request.args.get('prompt', 'What is the capital of France?')

        # Prepare the messages
        messages = [
            {
                "role": "user",
                "content": user_prompt
            }
        ]

        # Generate completion
        completion = client.chat.completions.create(
            model="deepseek-ai/DeepSeek-V3",
            messages=messages,
            max_tokens=500
        )

        # Extract the response
        response_text = completion.choices[0].message.content
        
        return jsonify({'response': response_text})

    except Exception as e:
        print(f"Error occurred: {str(e)}")
        return jsonify({'error': f'An error occurred: {str(e)}'}), 500

if __name__ == '__main__':
    app.run(debug=True)