Update main.py

main.py

@@ -1,78 +1,79 @@
-import json, os, glob, pathlib
+import json, os, glob, time
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi import FastAPI, Request, Header, BackgroundTasks, HTTPException, status
+from fastapi import FastAPI, Request, Header, BackgroundTasks, HTTPException
 from google import genai
 from linebot import LineBotApi, WebhookHandler
 from linebot.exceptions import InvalidSignatureError
-from linebot.models import MessageEvent, TextMessage, TextSendMessage, ImageSendMessage, AudioMessage
+from linebot.models import MessageEvent, TextMessage, TextSendMessage
 import PyPDF2
 import logging
 
-# 設定日誌
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# 檢查並設定 Google AI API 金鑰
 GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
 if not GOOGLE_API_KEY:
-    logger.error("錯誤：未設定 GOOGLE_API_KEY")
-    raise ValueError("請在環境變數中設定 GOOGLE_API_KEY")
+    raise ValueError("請設定 GOOGLE_API_KEY")
 
 client = genai.Client(api_key=GOOGLE_API_KEY)
 
-# 讀取 PDF 內容
+# 讀取 PDF
 files = glob.glob('docs/*.pdf')
 pdf_content = ''
 
-if not files:
-    logger.warning("警告：docs/ 資料夾中找不到 PDF 檔案")
-else:
-    logger.info(f"找到 {len(files)} 個 PDF 檔案")
+if files:
+    logger.info(f"找到 {len(files)} 個 PDF")
     for filename in files:
         try:
             with open(filename, 'rb') as pdf_file:
                 pdf_reader = PyPDF2.PdfReader(pdf_file)
-                for i in range(len(pdf_reader.pages)):
-                    page = pdf_reader.pages[i]
+                for page in pdf_reader.pages:
                     pdf_content += page.extract_text()
-            logger.info(f"成功讀取: {filename}")
         except Exception as e:
-            logger.error(f"讀取 {filename} 失敗: {str(e)}")
-
-# 檢查 PDF 內容長度
-logger.info(f"PDF 內容總長度: {len(pdf_content)} 字元")
-if len(pdf_content) > 100000:  # 如果內容過長
-    logger.warning("PDF 內容過長，可能影響 API 效能")
-
-# 設定生成文字的參數 + 角色扮演
-system_instruction = pdf_content + "\n" + "您是一位問答助手。請僅限使用以上提供的內容來回答問題，如果字數太多，請用列表方式呈現，如果您不知道答案，請說聯繫服務專員，不要捏造答案。"
-thinking_config = genai.types.ThinkingConfig(thinking_budget=0)
-generation_config = genai.types.GenerateContentConfig(
-    max_output_tokens=3000,
-    temperature=0.1,
-    top_p=0.2,
-    thinking_config=thinking_config,
-    system_instruction=system_instruction
-)
-
-# 設定 Line Bot 的 API 金鑰和秘密金鑰
-CHANNEL_ACCESS_TOKEN = os.getenv("CHANNEL_ACCESS_TOKEN")
-CHANNEL_SECRET = os.getenv("CHANNEL_SECRET")
-
-if not CHANNEL_ACCESS_TOKEN or not CHANNEL_SECRET:
-    logger.error("錯誤：未設定 LINE Bot 憑證")
-    raise ValueError("請設定 CHANNEL_ACCESS_TOKEN 和 CHANNEL_SECRET")
-
-line_bot_api = LineBotApi(CHANNEL_ACCESS_TOKEN)
-line_handler = WebhookHandler(CHANNEL_SECRET)
-
-# 設定是否正在與使用者交談
-working_status = os.getenv("DEFALUT_TALKING", default="true").lower() == "true"
+            logger.error(f"讀取失敗: {e}")
+
+logger.info(f"PDF 長度: {len(pdf_content)} 字元")
+
+# 限制長度
+if len(pdf_content) > 50000:
+    pdf_content = pdf_content[:50000]
+
+# 使用 Context Caching 儲存 PDF 內容
+cache_name = None
+try:
+    # 建立 cached content
+    cached_content = client.caches.create(
+        model="gemini-1.5-flash-001",
+        contents=[
+            genai.types.Content(
+                role="user",
+                parts=[genai.types.Part(text=f"""以下是參考資料，請記住這些內容：
+
+{pdf_content}
+
+你是問答助手，使用以上資料回答問題。規則：
+1. 只用提供的資料回答
+2. 不知道就說「請聯繫服務專員」
+3. 答案簡潔明瞭""")]
+            )
+        ],
+        ttl="300s",  # 快取 5 分鐘
+    )
+    cache_name = cached_content.name
+    logger.info(f"成功建立 Cache: {cache_name}")
+except Exception as e:
+    logger.error(f"無法建立 Cache: {e}")
+    logger.info("將使用一般模式（不使用 Cache）")
+
+# LINE Bot 設定
+line_bot_api = LineBotApi(os.getenv("CHANNEL_ACCESS_TOKEN"))
+line_handler = WebhookHandler(os.getenv("CHANNEL_SECRET"))
+working_status = True
+
+last_request_time = {}
+REQUEST_COOLDOWN = 3
 
-# 建立 FastAPI 應用程式
 app = FastAPI()
-
-# 設定 CORS
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -83,79 +84,93 @@ app.add_middleware(
 
 @app.get("/")
 def root():
-    return {"title": "Line Bot", "status": "running"}
+    return {
+        "status": "running",
+        "cache_enabled": cache_name is not None,
+        "pdf_length": len(pdf_content)
+    }
 
 @app.post("/webhook")
-async def webhook(
-    request: Request,
-    background_tasks: BackgroundTasks,
-    x_line_signature=Header(None),
-):
+async def webhook(request: Request, background_tasks: BackgroundTasks, x_line_signature=Header(None)):
     body = await request.body()
     try:
-        background_tasks.add_task(
-            line_handler.handle, body.decode("utf-8"), x_line_signature
-        )
+        background_tasks.add_task(line_handler.handle, body.decode("utf-8"), x_line_signature)
     except InvalidSignatureError:
-        logger.error("無效的簽章")
         raise HTTPException(status_code=400, detail="Invalid signature")
     return "ok"
 
 @line_handler.add(MessageEvent, message=TextMessage)
 def handle_message(event):
-    global working_status
-    
     if event.type != "message" or event.message.type != "text":
-        line_bot_api.reply_message(
-            event.reply_token,
-            TextSendMessage(text="錯誤：訊息類型不支援")
-        )
         return
         
     if event.message.text == "再見":
-        line_bot_api.reply_message(
-            event.reply_token,
-            TextSendMessage(text="Bye!")
-        )
+        line_bot_api.reply_message(event.reply_token, TextSendMessage(text="Bye!"))
         return
+    
+    # 速率限制
+    user_id = event.source.user_id
+    current_time = time.time()
+    
+    if user_id in last_request_time:
+        elapsed = current_time - last_request_time[user_id]
+        if elapsed < REQUEST_COOLDOWN:
+            line_bot_api.reply_message(
+                event.reply_token,
+                TextSendMessage(text=f"請稍等 {REQUEST_COOLDOWN - int(elapsed)} 秒")
+            )
+            return
+    
+    last_request_time[user_id] = current_time
        
     if working_status:
         try: 
             prompt = event.message.text
-            logger.info(f"收到訊息: {prompt}")
+            logger.info(f"提問: {prompt[:50]}")
             
-            # 呼叫 Gemini API
-            response = client.models.generate_content(
-                model="gemini-2.0-flash",  # 使用最新的模型
-                contents=prompt,
-                config=generation_config
-            )
+            # 如果有 Cache，使用 Cache
+            if cache_name:
+                response = client.models.generate_content(
+                    model="gemini-1.5-flash-001",
+                    contents=prompt,
+                    config=genai.types.GenerateContentConfig(
+                        max_output_tokens=1000,
+                        temperature=0.1,
+                        cached_content=cache_name
+                    )
+                )
+            else:
+                # 沒有 Cache，使用一般模式
+                response = client.models.generate_content(
+                    model="gemini-1.5-flash",
+                    contents=prompt,
+                    config=genai.types.GenerateContentConfig(
+                        max_output_tokens=1000,
+                        temperature=0.1,
+                        system_instruction=f"""你是問答助手。參考資料：
+{pdf_content[:10000]}
+規則：只用資料回答，不知道就說聯繫專員。"""
+                    )
+                )
             
-            # 檢查回應
             if response and response.text:
                 out = response.text
-                logger.info("Gemini 回應成功")
+                logger.info("回應成功")
             else:
-                out = "Gemini沒答案！將由專員服務！"
-                logger.warning("Gemini 回應為空")
+                out = "無法回答，請聯繫服務專員。"
                 
         except Exception as e:
-            # 詳細記錄錯誤
             error_msg = str(e)
-            logger.error(f"Gemini API 錯誤: {error_msg}")
+            logger.error(f"錯誤: {error_msg}")
             
-            # 根據錯誤類型給予不同提示
-            if "quota" in error_msg.lower() or "limit" in error_msg.lower():
-                out = "API 使用額度已達上限，請稍後再試！"
-            elif "key" in error_msg.lower() or "auth" in error_msg.lower():
-                out = "API 金鑰驗證失敗，請聯繫服務專員！"
+            if "quota" in error_msg.lower() or "resource_exhausted" in error_msg.lower():
+                out = "系統忙碌，請稍後再試。"
+            elif "429" in error_msg:
+                out = "請求過於頻繁，請稍候。"
             else:
-                out = f"Gemini執行出錯！請聯繫服務專員。錯誤：{error_msg[:50]}"
+                out = "系統錯誤，請聯繫服務專員。"
   
-        line_bot_api.reply_message(
-            event.reply_token,
-            TextSendMessage(text=out)
-        )
+        line_bot_api.reply_message(event.reply_token, TextSendMessage(text=out))
 
 if __name__ == "__main__":
     import uvicorn