Update app.py

app.py

@@ -4,7 +4,7 @@ from transformers import pipeline
 
 app = FastAPI()
 
-# Define request model
+# Define request model for math operations
 class CalculationRequest(BaseModel):
     a: float
     b: float
@@ -36,4 +36,24 @@ def calculate(request: CalculationRequest):
 @app.post("/generate")
 def generate_text(prompt: str):
     generated = model(prompt, max_length=50)
-    return {"generated_text": generated[0]['generated_text']}
+    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}