Update src/paper2DAG.py

src/paper2DAG.py

@@ -16,6 +16,10 @@ def clean_paper(markdown_path, clean_prompt, model, config):
     删除 Abstract / Related Work / Appendix / References 等部分，
     保留标题、作者、Introduction、Methods、Experiments、Conclusion。
     """
+    # === 可选：使用 config 中的 api_base_url 覆盖默认值 ===
+    if config['api_base_url'] is not None and config['api_base_url'].strip():
+        genai.config.api_base = config['api_base_url'].strip().rstrip("/") + "/v1"
+        print(f"🔧 Using custom API base: {genai.config.api_base}")
     # === 初始化 Client  ===
     client = genai.Client(
         api_key=config['api_keys']['gemini_api_key']
@@ -90,6 +94,11 @@ def split_paper(
     使用 Gemini 拆分论文，并将所有拆分后的 markdown 保存在：
         <parent_of_auto>/section_split_output/
     """
+    # === 可选：使用 config 中的 api_base_url 覆盖默认值 ===
+    if config['api_base_url'] is not None and config['api_base_url'].strip():
+        genai.config.api_base = config['api_base_url'].strip().rstrip("/") + "/v1"
+        print(f"🔧 Using custom API base: {genai.config.api_base}")
+
     # 1️⃣ 输入文件所在的 auto 文件夹
     auto_dir = os.path.dirname(os.path.abspath(cleaned_md_path))
 
@@ -199,6 +208,11 @@ def initialize_dag(markdown_path, initialize_dag_prompt, model, config=None):
     with open(markdown_path, "r", encoding="utf-8") as f:
         md_text = f.read()
 
+    # === 可选：使用 config 中的 api_base_url 覆盖默认值 ===
+    if config['api_base_url'] is not None and config['api_base_url'].strip():
+        genai.config.api_base = config['api_base_url'].strip().rstrip("/") + "/v1"
+        print(f"🔧 Using custom API base: {genai.config.api_base}")
+
     # --- Gemini Client Init ---
     client = genai.Client(
         api_key=config['api_keys']['gemini_api_key']
@@ -300,6 +314,11 @@ def extract_and_generate_visual_dag(
     normalized_refs = [f"![]({m})" for m in relative_imgs]
 
     # === 3. 发送给 Gemini ===
+    # === 可选：使用 config 中的 api_base_url 覆盖默认值 ===
+    if config['api_base_url'] is not None and config['api_base_url'].strip():
+        genai.config.api_base = config['api_base_url'].strip().rstrip("/") + "/v1"
+        print(f"🔧 Using custom API base: {genai.config.api_base}")
+
     # 初始化 Client
     client = genai.Client(
         api_key=config['api_keys']['gemini_api_key']
@@ -526,6 +545,11 @@ def build_section_dags(
     MAX_RETRIES_ON_FAIL = 2
 
     # === Init Client (Gemini) ===
+    # === 可选：使用 config 中的 api_base_url 覆盖默认值 ===
+    if config['api_base_url'] is not None and config['api_base_url'].strip():
+        genai.config.api_base = config['api_base_url'].strip().rstrip("/") + "/v1"
+        print(f"🔧 Using custom API base: {genai.config.api_base}")
+        
     # 使用 config 中的 key 
     client = genai.Client(
         api_key=config['api_keys']['gemini_api_key']