from flask import Flask, request, jsonify
from gradio_client import Client
import os

app = Flask(__name__)

# Initialize the Gradio client
gradio_client = Client("bhlewis/BERT-for-Patents_Semantic-Patent-Finder-v2")

@app.route('/')
def home():
    return "Server is running. Use /api/predict for predictions."

@app.route('/api/predict', methods=['POST'])
def predict():
    try:
        data = request.json
        query = data.get('query', '')
        top_k = data.get('top_k', 5)

        print(f"Received query: {query}")
        print(f"Top K: {top_k}")

        result = gradio_client.predict(
            query=query,
            top_k=top_k,
            api_name="/predict"
        )

        print(f"Result: {result}")

        return jsonify({"result": result})

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

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    app.run(host='0.0.0.0', port=port)