Update app.py

app.py

@@ -1,14 +1,37 @@
 from fastapi import FastAPI
+from pydantic import BaseModel
 from transformers import pipeline
+import mysql.connector
+import os
 
 app = FastAPI()
 
-# Load model at startup (only once)
+# Database connection settings
+DB_HOST = "gateway01.ap-southeast-1.prod.aws.tidbcloud.com"
+DB_PORT = 4000
+DB_USER = "4V44XYoMA7okY9v.root"
+DB_PASS = "aW2CrSwcTgjFhNAb"
+DB_NAME = "final_project"
+
+# Create MySQL connection
+conn = mysql.connector.connect(
+    host=DB_HOST,
+    port=DB_PORT,
+    user=DB_USER,
+    password=DB_PASS,
+    database=DB_NAME,
+    ssl_verify_cert=True,
+    ssl_verify_identity=True
+)
+cursor = conn.cursor()
+
+# Load model
 classifier = pipeline(
     "zero-shot-classification",
     model="MoritzLaurer/deberta-v3-large-zeroshot-v1.1-all-33"
 )
 
+# Labels
 subject_labels = [
     "Physics", "Chemistry", "Biology", "Astronomy",
     "Earth Science", "Environmental Science",
@@ -25,12 +48,36 @@ subject_labels = [
     "Physical Education", "Health Science"
 ]
 
-@app.post("/predict/")
-async def predict_topic(text: str):
+# Request model
+class TextInput(BaseModel):
+    student_id: str
+    text: str
+
+@app.post("/predict")
+def predict_topic(data: TextInput):
+    # Predict subject
     result = classifier(
-        text,
+        data.text,
         candidate_labels=subject_labels,
         hypothesis_template="This text is about {}."
     )
-    predicted_topic = result["labels"][0]
-    return {"topic": predicted_topic}
+    predicted_subject = result["labels"][0]
+
+    # Get first 100 characters of the text
+    sample_text = data.text[:100]
+
+    # Save to DB
+    cursor.execute(
+        """
+        INSERT INTO log_table (student_id, input_sample, subject)
+        VALUES (%s, %s, %s)
+        """,
+        (data.student_id, sample_text, predicted_subject)
+    )
+    conn.commit()
+
+    return {
+        "student_id": data.student_id,
+        "predicted_subject": predicted_subject,
+        "sample_text": sample_text
+    }