botrun-flow-lang 5.12.264__py3-none-any.whl → 6.2.61__py3-none-any.whl

botrun_flow_lang/langgraph_agents/agents/util/pdf_analyzer.py

@@ -11,11 +11,13 @@ import asyncio
 import base64
 import httpx
 import os
-from typing import List, Dict, Any
+from typing import List, Dict, Any, Tuple
 
 from dotenv import load_dotenv
 from google.oauth2 import service_account
 
+from botrun_flow_lang.langgraph_agents.agents.util.usage_metadata import UsageMetadata
+
 load_dotenv()
 
 # 檔案大小閾值（MB）
@@ -30,16 +32,17 @@ MAX_CONCURRENT_CHUNKS = 5
 
 def analyze_pdf_with_claude(
     pdf_data: str, user_input: str, model_name: str = "claude-sonnet-4-5-20250929"
-):
+) -> Tuple[str, UsageMetadata]:
     """
     Analyze a PDF file using Claude API
 
     Args:
         pdf_data: Base64-encoded PDF data
         user_input: User's query about the PDF content
+        model_name: Claude model name to use
 
     Returns:
-        str: Claude's analysis of the PDF content based on the query
+        Tuple[str, UsageMetadata]: Claude's analysis and usage metadata
     """
     # Initialize Anthropic client
     client = anthropic.Anthropic()
@@ -66,15 +69,25 @@ def analyze_pdf_with_claude(
         ],
     )
 
+    # Extract usage metadata
+    usage = UsageMetadata(
+        prompt_tokens=message.usage.input_tokens,
+        completion_tokens=message.usage.output_tokens,
+        total_tokens=message.usage.input_tokens + message.usage.output_tokens,
+        cache_creation_input_tokens=getattr(message.usage, 'cache_creation_input_tokens', 0) or 0,
+        cache_read_input_tokens=getattr(message.usage, 'cache_read_input_tokens', 0) or 0,
+        model=model_name,
+    )
+
     print(
         f"analyze_pdf_with_claude============> input_token: {message.usage.input_tokens} output_token: {message.usage.output_tokens}",
     )
-    return message.content[0].text
+    return message.content[0].text, usage
 
 
 def analyze_pdf_with_gemini(
     pdf_data: str, user_input: str, model_name: str = "gemini-2.5-flash", pdf_url: str = ""
-):
+) -> Tuple[str, UsageMetadata]:
     """
     Analyze a PDF file using Gemini API
 
@@ -82,9 +95,10 @@ def analyze_pdf_with_gemini(
         pdf_data: Base64-encoded PDF data
         user_input: User's query about the PDF content
         model_name: Gemini model name to use
+        pdf_url: Original PDF URL for logging
 
     Returns:
-        str: Gemini's analysis of the PDF content based on the query
+        Tuple[str, UsageMetadata]: Gemini's analysis and usage metadata
     """
     # 放到要用的時候才 import，不然loading 會花時間
     from google import genai
@@ -112,14 +126,25 @@ def analyze_pdf_with_gemini(
             ),
         ],
     )
-    # Log token usage if available
+
+    # Extract usage metadata
+    usage = UsageMetadata(model=model_name)
     if hasattr(response, "usage_metadata"):
+        usage_meta = response.usage_metadata
+        usage = UsageMetadata(
+            prompt_tokens=getattr(usage_meta, 'prompt_token_count', 0) or 0,
+            completion_tokens=getattr(usage_meta, 'candidates_token_count', 0) or 0,
+            total_tokens=getattr(usage_meta, 'total_token_count', 0) or 0,
+            cache_creation_input_tokens=0,
+            cache_read_input_tokens=getattr(usage_meta, 'cached_content_token_count', 0) or 0,
+            model=model_name,
+        )
         print(
-            f"analyze_pdf_with_gemini============> input_token: {response.usage_metadata.prompt_token_count} output_token: {response.usage_metadata.candidates_token_count}",
+            f"analyze_pdf_with_gemini============> input_token: {usage_meta.prompt_token_count} output_token: {usage_meta.candidates_token_count}",
         )
 
     print(f"{pdf_url} success")
-    return response.text
+    return response.text, usage
 
 
 def _analyze_single_chunk(
@@ -135,7 +160,7 @@ def _analyze_single_chunk(
         model_name: 使用的模型名稱
 
     Returns:
-        Dict: {"page_range": str, "answer": str, "relevant": bool, "error": str|None}
+        Dict: {"page_range": str, "answer": str, "relevant": bool, "error": str|None, "usage": UsageMetadata}
     """
     # 構建切片專用的 prompt
     chunk_prompt = f"""你正在閱讀一份大型 PDF 文件的其中一部分（{page_range}）。
@@ -149,15 +174,16 @@ def _analyze_single_chunk(
 
     try:
         if model_name.startswith("gemini-"):
-            answer = analyze_pdf_with_gemini(chunk_data, chunk_prompt, model_name)
+            answer, usage = analyze_pdf_with_gemini(chunk_data, chunk_prompt, model_name)
         elif model_name.startswith("claude-"):
-            answer = analyze_pdf_with_claude(chunk_data, chunk_prompt, model_name)
+            answer, usage = analyze_pdf_with_claude(chunk_data, chunk_prompt, model_name)
         else:
             return {
                 "page_range": page_range,
                 "answer": "",
                 "relevant": False,
                 "error": f"Unknown model type: {model_name}",
+                "usage": UsageMetadata(),
             }
 
         # 判斷是否相關
@@ -168,6 +194,7 @@ def _analyze_single_chunk(
             "answer": answer if is_relevant else "",
             "relevant": is_relevant,
             "error": None,
+            "usage": usage,
         }
 
     except Exception as e:
@@ -179,12 +206,13 @@ def _analyze_single_chunk(
             "answer": "",
             "relevant": False,
             "error": str(e),
+            "usage": UsageMetadata(model=model_name),
         }
 
 
 async def analyze_pdf_chunks_parallel(
     chunks: List[tuple], user_input: str, model_name: str, max_concurrent: int = 5
-) -> List[Dict[str, Any]]:
+) -> Tuple[List[Dict[str, Any]], List[UsageMetadata]]:
     """
     平行問答多個 PDF 切片
 
@@ -195,7 +223,7 @@ async def analyze_pdf_chunks_parallel(
         max_concurrent: 最大平行數量
 
     Returns:
-        List[Dict]: 每個切片的回答結果
+        Tuple[List[Dict], List[UsageMetadata]]: 每個切片的回答結果和每次呼叫的 usage list
     """
     semaphore = asyncio.Semaphore(max_concurrent)
 
@@ -224,8 +252,9 @@ async def analyze_pdf_chunks_parallel(
     # 平行執行
     results = await asyncio.gather(*tasks, return_exceptions=True)
 
-    # 處理例外
+    # 處理例外並收集 usage list
     processed_results = []
+    usage_list = []
     for i, result in enumerate(results):
         if isinstance(result, Exception):
             processed_results.append(
@@ -234,19 +263,24 @@ async def analyze_pdf_chunks_parallel(
                     "answer": "",
                     "relevant": False,
                     "error": str(result),
+                    "usage": UsageMetadata(model=model_name),
                 }
             )
+            usage_list.append(UsageMetadata(model=model_name))
         else:
             processed_results.append(result)
+            # 收集 usage
+            if "usage" in result and isinstance(result["usage"], UsageMetadata):
+                usage_list.append(result["usage"])
 
-    return processed_results
+    return processed_results, usage_list
 
 
 def merge_chunk_results(
     chunk_results: List[Dict[str, Any]],
     user_input: str,
     model_name: str = "gemini-2.5-flash",
-) -> str:
+) -> Tuple[str, UsageMetadata]:
     """
     使用 LLM 統整多個切片的回答
 
@@ -256,7 +290,7 @@ def merge_chunk_results(
         model_name: 統整使用的模型名稱
 
     Returns:
-        str: 統整後的回答
+        Tuple[str, UsageMetadata]: 統整後的回答和 usage metadata
     """
     # 過濾出相關的回答
     relevant_results = [r for r in chunk_results if r.get("relevant", False)]
@@ -266,12 +300,12 @@ def merge_chunk_results(
         error_results = [r for r in chunk_results if r.get("error")]
         if error_results:
             error_msgs = [f"{r['page_range']}: {r['error']}" for r in error_results]
-            return f"分析 PDF 時發生錯誤：\n" + "\n".join(error_msgs)
-        return "在 PDF 文件中未找到與您問題相關的內容。"
+            return f"分析 PDF 時發生錯誤：\n" + "\n".join(error_msgs), UsageMetadata(model=model_name)
+        return "在 PDF 文件中未找到與您問題相關的內容。", UsageMetadata(model=model_name)
 
-    # 只有一個相關結果，直接回傳
+    # 只有一個相關結果，直接回傳（不需要額外的 LLM 呼叫）
     if len(relevant_results) == 1:
-        return relevant_results[0]["answer"]
+        return relevant_results[0]["answer"], UsageMetadata(model=model_name)
 
     # 多個相關結果，需要統整
     combined_content = "\n\n".join(
@@ -310,22 +344,33 @@ def merge_chunk_results(
             contents=[merge_prompt],
         )
 
+        # Extract usage metadata
+        usage = UsageMetadata(model=model_name)
         if hasattr(response, "usage_metadata"):
+            usage_meta = response.usage_metadata
+            usage = UsageMetadata(
+                prompt_tokens=getattr(usage_meta, 'prompt_token_count', 0) or 0,
+                completion_tokens=getattr(usage_meta, 'candidates_token_count', 0) or 0,
+                total_tokens=getattr(usage_meta, 'total_token_count', 0) or 0,
+                cache_creation_input_tokens=0,
+                cache_read_input_tokens=getattr(usage_meta, 'cached_content_token_count', 0) or 0,
+                model=model_name,
+            )
             print(
-                f"merge_chunk_results============> input_token: {response.usage_metadata.prompt_token_count} output_token: {response.usage_metadata.candidates_token_count}",
+                f"merge_chunk_results============> input_token: {usage_meta.prompt_token_count} output_token: {usage_meta.candidates_token_count}",
             )
 
-        return response.text
+        return response.text, usage
 
     except Exception as e:
         import traceback
 
         traceback.print_exc()
         # 統整失敗，直接回傳合併的內容
-        return f"統整時發生錯誤，以下是各部分的回答：\n\n{combined_content}"
+        return f"統整時發生錯誤，以下是各部分的回答：\n\n{combined_content}", UsageMetadata(model=model_name)
 
 
-async def analyze_pdf_async(pdf_url: str, user_input: str) -> str:
+async def analyze_pdf_async(pdf_url: str, user_input: str) -> Dict[str, Any]:
     """
     非同步分析 PDF 檔案（智慧處理策略）
 
@@ -338,8 +383,13 @@ async def analyze_pdf_async(pdf_url: str, user_input: str) -> str:
         user_input: 使用者問題
 
     Returns:
-        str: 分析結果
+        Dict[str, Any]: {
+            "result": str,  # 分析結果
+            "usage_metadata": List[Dict]  # 每次 LLM 呼叫的 usage 資訊
+        }
     """
+    usage_list: List[UsageMetadata] = []
+
     try:
         # 1. 下載 PDF
         print(f"[analyze_pdf_async] 下載 PDF: {pdf_url}")
@@ -364,9 +414,19 @@ async def analyze_pdf_async(pdf_url: str, user_input: str) -> str:
             for model in models:
                 try:
                     if model.startswith("gemini-"):
-                        return analyze_pdf_with_gemini(pdf_data, user_input, model, pdf_url)
+                        result, usage = analyze_pdf_with_gemini(pdf_data, user_input, model, pdf_url)
+                        usage_list.append(usage)
+                        return {
+                            "result": result,
+                            "usage_metadata": [u.to_dict() for u in usage_list],
+                        }
                     elif model.startswith("claude-"):
-                        return analyze_pdf_with_claude(pdf_data, user_input, model)
+                        result, usage = analyze_pdf_with_claude(pdf_data, user_input, model)
+                        usage_list.append(usage)
+                        return {
+                            "result": result,
+                            "usage_metadata": [u.to_dict() for u in usage_list],
+                        }
                 except Exception as e:
                     import traceback
 
@@ -374,7 +434,10 @@ async def analyze_pdf_async(pdf_url: str, user_input: str) -> str:
                     last_error = str(e)
                     continue
 
-            return f"分析 PDF 時所有模型都失敗。最後錯誤: {last_error}"
+            return {
+                "result": f"分析 PDF 時所有模型都失敗。最後錯誤: {last_error}",
+                "usage_metadata": [u.to_dict() for u in usage_list],
+            }
 
         # 3. 大檔：壓縮 → 切割 → 平行問答 → 統整
         print(f"[analyze_pdf_async] 大檔模式 (>= {PDF_SIZE_THRESHOLD_MB}MB)")
@@ -427,9 +490,10 @@ async def analyze_pdf_async(pdf_url: str, user_input: str) -> str:
 
         # 3.3 平行問答
         print(f"[analyze_pdf_async] 開始平行問答 (最大並行: {MAX_CONCURRENT_CHUNKS})...")
-        chunk_results = await analyze_pdf_chunks_parallel(
+        chunk_results, chunk_usage_list = await analyze_pdf_chunks_parallel(
             chunks, user_input, primary_model, max_concurrent=MAX_CONCURRENT_CHUNKS
         )
+        usage_list.extend(chunk_usage_list)
 
         # 統計結果
         relevant_count = sum(1 for r in chunk_results if r.get("relevant", False))
@@ -441,19 +505,28 @@ async def analyze_pdf_async(pdf_url: str, user_input: str) -> str:
 
         # 3.4 統整結果
         print("[analyze_pdf_async] 統整結果...")
-        result = merge_chunk_results(chunk_results, user_input, primary_model)
+        result, merge_usage = merge_chunk_results(chunk_results, user_input, primary_model)
+        # 只有當 merge_usage 有實際 token 使用時才加入（避免加入空的 usage）
+        if merge_usage.prompt_tokens > 0 or merge_usage.completion_tokens > 0:
+            usage_list.append(merge_usage)
         print("[analyze_pdf_async] 完成")
 
-        return result
+        return {
+            "result": result,
+            "usage_metadata": [u.to_dict() for u in usage_list],
+        }
 
     except Exception as e:
         import traceback
 
         traceback.print_exc()
-        return f"分析 PDF {pdf_url} 時發生錯誤: {str(e)}"
+        return {
+            "result": f"分析 PDF {pdf_url} 時發生錯誤: {str(e)}",
+            "usage_metadata": [u.to_dict() for u in usage_list],
+        }
 
 
-def analyze_pdf(pdf_url: str, user_input: str) -> str:
+def analyze_pdf(pdf_url: str, user_input: str) -> Dict[str, Any]:
     """
     分析 PDF 檔案（同步包裝函數）
 
@@ -465,7 +538,10 @@ def analyze_pdf(pdf_url: str, user_input: str) -> str:
         user_input: 使用者問題
 
     Returns:
-        str: 分析結果
+        Dict[str, Any]: {
+            "result": str,  # 分析結果
+            "usage_metadata": List[Dict]  # 每次 LLM 呼叫的 usage 資訊
+        }
     """
     try:
         # 嘗試取得現有的事件迴圈

botrun_flow_lang/langgraph_agents/agents/util/usage_metadata.py

@@ -0,0 +1,34 @@
+"""
+Usage Metadata 模組
+
+提供 LLM 呼叫的 token 使用量追蹤功能。
+"""
+
+from dataclasses import dataclass, asdict
+from typing import Dict, Any
+
+
+@dataclass
+class UsageMetadata:
+    """Token usage metadata that matches the expected parsing format."""
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    total_tokens: int = 0
+    cache_creation_input_tokens: int = 0
+    cache_read_input_tokens: int = 0
+    model: str = ""
+
+    def __add__(self, other: "UsageMetadata") -> "UsageMetadata":
+        """Combine two UsageMetadata objects."""
+        return UsageMetadata(
+            prompt_tokens=self.prompt_tokens + other.prompt_tokens,
+            completion_tokens=self.completion_tokens + other.completion_tokens,
+            total_tokens=self.total_tokens + other.total_tokens,
+            cache_creation_input_tokens=self.cache_creation_input_tokens + other.cache_creation_input_tokens,
+            cache_read_input_tokens=self.cache_read_input_tokens + other.cache_read_input_tokens,
+            model=self.model or other.model,
+        )
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary."""
+        return asdict(self)