from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline

app = FastAPI()

# Define request model for math operations
class CalculationRequest(BaseModel):
    a: float
    b: float
    operation: str

# Load a smaller Hugging Face model (example: distilgpt2)
model = pipeline('text-generation', model='distilgpt2')

@app.post("/calculate")
def calculate(request: CalculationRequest):
    a = request.a
    b = request.b
    operation = request.operation

    if operation == "add":
        result = a + b
    elif operation == "subtract":
        result = a - b
    elif operation == "multiply":
        result = a * b
    elif operation == "divide":
        result = a / b
    else:
        return {"error": "Invalid operation"}

    return {"result": result}

# Example endpoint using Hugging Face model
@app.post("/generate")
def generate_text(prompt: str):
    generated = model(prompt, max_length=50, clean_up_tokenization_spaces=True)
    return {"generated_text": generated[0]['generated_text']}

# New endpoint for testing math operations
@app.post("/test_math")
def test_math(request: CalculationRequest):
    a = request.a
    b = request.b
    operation = request.operation

    if operation == "add":
        result = a + b
    elif operation == "subtract":
        result = a - b
    elif operation == "multiply":
        result = a * b
    elif operation == "divide":
        result = a / b
    else:
        return {"error": "Invalid operation"}

    return {"result": result}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)