HungLuong10/microservice

1from __future__ import annotations
2from fastapi import FastAPI, UploadFile, File, BackgroundTasks
3from pydantic import BaseModel
4
5import tempfile
6import zipfile
7import json
8import os
9import shutil
10from fastapi import HTTPException
11from fastapi.responses import FileResponse
12
13# import helper functions
14from llm import *
15
16from bert import * 
17
18app = FastAPI()
19
20def _cleanup_empty_dir(dir_path):
21    """Xóa folder nếu nó rỗng hoặc chỉ chứa folder rỗng."""
22    try:
23        if os.path. exists(dir_path) and os.path.isdir(dir_path):
24            # Nếu folder không rỗng, không xóa
25            if not os.listdir(dir_path):
26                os.rmdir(dir_path)
27    except Exception:
28        pass
29
30class AskRequest(BaseModel):
31    question: str
32
33
34class VectorizeRequest(BaseModel):
35    files: list[dict]
36
37# hỏi llm
38@app.post("/llm/ask")
39def ask(req: AskRequest):
40    result = run_pipeline_ask_llm(req.question)
41    return result
42
43@app.post("/llm/vectorize")
44async def vectorize(req: VectorizeRequest):
45    try:
46        files = req.files
47        run_pipeline_vectorize_files(files)
48        return {"status": "ok", "total": len(files)}
49    except Exception as e:
50        import traceback
51        traceback.print_exc()
52        raise HTTPException(status_code=500, detail=str(e))
53
54
55# vectorize files
56@app.post("/bert/process-docx")
57async def process_docx(
58    file: UploadFile = File(...),
59    background_tasks: BackgroundTasks = BackgroundTasks()
60):
61    # ===== 1. Lưu docx tạm =====
62    original_name = os.path. splitext(file.filename)[0]
63
64    with tempfile.NamedTemporaryFile(delete=False, suffix=".docx") as tmp_docx:
65        tmp_docx.write(await file.read())
66        docx_path = tmp_docx.name
67
68    # ===== 2. Tạo zip tạm =====
69    with tempfile.NamedTemporaryFile(delete=False, suffix=". zip") as tmp_zip:
70        zip_path = tmp_zip.name
71
72    try:
73        # ===== 3. Xử lý =====
74        data = extract(docx_path)
75        
76        if not data:
77            raise HTTPException(status_code=400, detail="Không trích xuất được nội dung từ file DOCX.")
78      
79        for x in data:
80            if x["text"]:
81                x["text"] = clean_one_paragraph_text(x["text"])
82            elif x["images"]:
83                print(
84                    f"para={x['para_index']} "
85                    f"images={[i['saved_path'] for i in x['images']]}"
86                )
87
88        result = classify_docx(data)
89        result = explode_mcq_in_grouped_result(result)
90        for i, qa in enumerate(result):
91            q = qa["question"]
92            if q.get("images"):
93                print(
94                    f"[Q{i}] images={[img['saved_path'] for img in q['images']]}"
95                )
96
97        # Media folder path
98        media_dir = os. path.abspath("./data/outputs/media")
99
100        with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as z:
101            # 👉 JSON cùng tên file docx
102            json_name = f"{original_name}.json"
103            z.writestr(
104                json_name,
105                json.dumps(result, ensure_ascii=False, indent=2)
106            )
107
108            # 👉 Thêm folder media/ vào zip
109            if os.path.exists(media_dir):
110                for fname in os.listdir(media_dir):
111                    file_path = os.path.join(media_dir, fname)
112                    if os.path.isfile(file_path):
113                        z.write(file_path, f"media/{fname}")
114
115        # ===== 4. Cleanup tasks =====
116        # Xóa docx tạm
117        background_tasks. add_task(os.remove, docx_path)
118
119        # Xóa toàn bộ media folder sau khi gửi zip
120        background_tasks.add_task(shutil.rmtree, media_dir, ignore_errors=True)
121
122        # Xóa outputs folder nếu nó rỗng
123        outputs_dir = os.path.abspath("./data/outputs")
124        background_tasks.add_task(_cleanup_empty_dir, outputs_dir)
125
126        # Xóa zip file sau cùng
127        background_tasks.add_task(os.remove, zip_path)
128
129        return FileResponse(
130            zip_path,
131            media_type="application/zip",
132            filename=f"{original_name}.zip",
133            background=background_tasks
134        )
135
136    except Exception as e:
137        # Nếu có lỗi, cleanup ngay
138        if os.path.exists(docx_path):
139            os.remove(docx_path)
140        if os.path.exists(zip_path):
141            os.remove(zip_path)
142
143        media_dir = os.path.abspath("./data/outputs/media")
144        if os.path. exists(media_dir):
145            shutil.rmtree(media_dir, ignore_errors=True)
146
147        raise
148