app.py

from flask import Flask, request, jsonify
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import requests

app = Flask(__name__)

# Allow CORS for everything (so TurboWarp can connect)
from flask_cors import CORS
CORS(app)

# === Load Phi model ===
print("🚀 Loading Phi model... this may take a minute.")
model_name = "microsoft/phi-2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True
)
model.to("cuda" if torch.cuda.is_available() else "cpu")
print("✅ Model loaded!")

# === Main API ===
@app.route("/api/ask", methods=["POST"])
def ask():
    data = request.get_json()
    prompt = data.get("prompt", "")

    chat_prompt = f"### Instruction:\nYou are Acla, a helpful AI powered by phi-3 mini that can reason about math, code, and logic.\n\n### Input:\n{prompt}\n\n### Response:"
    inputs = tokenizer(chat_prompt, return_tensors="pt").to(model.device)

    outputs = model.generate(
        **inputs,
        max_new_tokens=300,
        temperature=0.7,
        top_p=0.9,
        do_sample=True
    )

    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    response = text.split("### Response:")[-1].strip()
    return jsonify({"reply": response})


# === Proxy endpoint ===
@app.route("/proxy", methods=["POST"])
def proxy():
    """Forward TurboWarp requests to /api/ask internally."""
    try:
        data = request.get_json()
        r = requests.post("http://localhost:7860/api/ask", json=data)
        return jsonify(r.json())
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@app.route("/")
def home():
    return "🧠 Phi-2 Chatbot + Proxy running! Send POST to /proxy or /api/ask"

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)