Update app.py

app.py

@@ -46,6 +46,7 @@ print("Model loaded successfully!")
 # ============================================================
 class SummarizeRequest(BaseModel):
     text: str
+    length_level: int = 1  # 0: Ngắn (2-3 ý), 1: Trung bình (4-5 ý), 2: Chi tiết (6+ ý)
 
 
 # ============================================================
@@ -129,10 +130,11 @@ def fix_truncated_text(text: str) -> str:
     return text + '.'
 
 
-def format_as_bullet_points(summaries: list[str]) -> str:
+def format_as_bullet_points(summaries: list[str], max_points: int = None) -> str:
     """
     Chuyển đổi các đoạn tóm tắt thành Bullet Points.
     Mỗi ý một dòng, bắt đầu bằng '•'.
+    max_points: Số lượng bullet points tối đa (None = không giới hạn)
     """
     bullet_points = []
     
@@ -147,6 +149,10 @@ def format_as_bullet_points(summaries: list[str]) -> str:
                 sentence = fix_truncated_text(sentence)
                 bullet_points.append(f"• {sentence}")
     
+    # Giới hạn số lượng bullet points nếu có
+    if max_points is not None and len(bullet_points) > max_points:
+        bullet_points = bullet_points[:max_points]
+    
     return '\n'.join(bullet_points)
 
 
@@ -277,8 +283,18 @@ async def summarize_text(request: SummarizeRequest):
     """
     Tóm tắt văn bản tiếng Việt.
     Trả về danh sách Bullet Points.
+    length_level: 0 = Ngắn (2-3 ý), 1 = Trung bình (4-5 ý), 2 = Chi tiết (6+ ý)
     """
     text = request.text
+    length_level = request.length_level
+    
+    # Map length_level to max_points
+    max_points_map = {
+        0: 3,   # Ngắn: 2-3 ý
+        1: 5,   # Trung bình: 4-5 ý
+        2: None # Chi tiết: không giới hạn
+    }
+    max_points = max_points_map.get(length_level, 5)
     
     # Validate text
     if not text or len(text.strip()) < 50:
@@ -296,25 +312,35 @@ async def summarize_text(request: SummarizeRequest):
             detail="Không thể tạo tóm tắt."
         )
     
-    # Format as bullet points
-    bullet_points = format_as_bullet_points(summaries)
+    # Format as bullet points với giới hạn số lượng
+    bullet_points = format_as_bullet_points(summaries, max_points=max_points)
     
     return {
         "success": True,
         "original_length": len(text),
         "word_count": len(text.split()),
         "num_chunks": len(summaries),
+        "length_level": length_level,
         "summary": bullet_points
     }
 
 
 @app.post("/upload-pdf")
-async def upload_pdf(file: UploadFile = File(...)):
+async def upload_pdf(file: UploadFile = File(...), length_level: int = 1):
     """
     Upload và tóm tắt file PDF.
     Đọc qua byte stream, KHÔNG lưu file ra đĩa.
     Trả về danh sách Bullet Points.
+    length_level: 0 = Ngắn (2-3 ý), 1 = Trung bình (4-5 ý), 2 = Chi tiết (6+ ý)
     """
+    # Map length_level to max_points
+    max_points_map = {
+        0: 3,   # Ngắn: 2-3 ý
+        1: 5,   # Trung bình: 4-5 ý
+        2: None # Chi tiết: không giới hạn
+    }
+    max_points = max_points_map.get(length_level, 5)
+    
     # Validate file type
     if not file.filename.lower().endswith('.pdf'):
         raise HTTPException(
@@ -358,8 +384,8 @@ async def upload_pdf(file: UploadFile = File(...)):
             detail="Không thể tạo tóm tắt."
         )
     
-    # Format as bullet points
-    bullet_points = format_as_bullet_points(summaries)
+    # Format as bullet points với giới hạn số lượng
+    bullet_points = format_as_bullet_points(summaries, max_points=max_points)
     
     return {
         "success": True,
@@ -367,6 +393,7 @@ async def upload_pdf(file: UploadFile = File(...)):
         "original_length": len(text),
         "word_count": len(text.split()),
         "num_chunks": len(summaries),
+        "length_level": length_level,
         "summary": bullet_points
     }