Update backend/ai_output.py

backend/ai_output.py

@@ -1,25 +1,52 @@
 # backend/ai_output.py
 
-import os
+import torch
+from fastapi import HTTPException
 
-BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+# 你原本在 api_server.py 裡的 model / tokenizer
+from api_server import model, tokenizer
 
-def run_ai_output(field_name: str):
+
+def run_ai_output(input_text: str):
     """
-    讀取 backend/skills/<field_name>_ai/<field_name>.md
-    目前先回傳內容（之後你可以在這裡接 LLM）
+    使用 LLM 生成 3 個補全結果
     """
-    folder = os.path.join(BASE_DIR, "skills", f"{field_name}_ai")
-    file_path = os.path.join(folder, f"{field_name}.md")
 
-    if not os.path.exists(file_path):
-        # 沒有對應 skill → fallback
-        return f"[AI] {field_name}"
+    if model is None or tokenizer is None:
+        raise HTTPException(status_code=503, detail="AI 模型尚未準備就緒")
+
+    if len(input_text) > 512:
+        raise HTTPException(status_code=400, detail="輸入過長，請限制在 512 字元內")
+
+    try:
+        input_ids = tokenizer.encode(input_text, return_tensors='pt', truncation=True)
+
+        output = model.generate(
+            input_ids,
+            max_length=len(input_text) + 150,
+            num_return_sequences=3,
+            no_repeat_ngram_size=3,
+            do_sample=True,
+            top_k=50,
+            top_p=0.95,
+            temperature=0.8,
+            pad_token_id=tokenizer.eos_token_id
+        )
+
+        all_completions = []
+        for sequence in output:
+            generated_text = tokenizer.decode(sequence, skip_special_tokens=True)
+
+            if generated_text.startswith(input_text):
+                all_completions.append(generated_text)
+
+        unique_completions = sorted(list(set(all_completions)), key=len, reverse=True)
+
+        if not unique_completions:
+            return [input_text]
 
-    with open(file_path, "r", encoding="utf-8") as f:
-        content = f.read().strip()
+        return unique_completions
 
-    # 之後你可以在這裡：
-    # response = call_llm(prompt=content)
-    # return response
-    return content
+    except Exception as e:
+        print(f"AI 推論錯誤: {e}")
+        raise HTTPException(status_code=500, detail=f"AI 推論失敗：{str(e)[:50]}...")