Update src/paper2DAG.py

src/paper2DAG.py

@@ -16,13 +16,15 @@ 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  ===
+    raw_url = config.get('api_base_url', '').strip().rstrip("/")
+    if raw_url.endswith("/v1"):
+        base_url = raw_url[:-3].rstrip("/") # 去掉最后的 /v1
+    else:
+        base_url = raw_url
+
     client = genai.Client(
-        api_key=config['api_keys']['gemini_api_key']
+        api_key=config['api_keys']['gemini_api_key'],
+        http_options={'base_url': base_url} if base_url else None
     )
 
     # === 读取 markdown 文件 ===
@@ -94,11 +96,6 @@ 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))
 
@@ -114,8 +111,15 @@ def split_paper(
         markdown_text = f.read()
 
     # === 2. 初始化 Gemini Client ===
+    raw_url = config.get('api_base_url', '').strip().rstrip("/")
+    if raw_url.endswith("/v1"):
+        base_url = raw_url[:-3].rstrip("/") # 去掉最后的 /v1
+    else:
+        base_url = raw_url
+
     client = genai.Client(
-        api_key=config['api_keys']['gemini_api_key']
+        api_key=config['api_keys']['gemini_api_key'],
+        http_options={'base_url': base_url} if base_url else None
     )
 
     # === 提取一级 section 信息供参考 (假设 SECTION_RE 已在外部定义) ===
@@ -208,14 +212,16 @@ 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 ---
+    raw_url = config.get('api_base_url', '').strip().rstrip("/")
+    if raw_url.endswith("/v1"):
+        base_url = raw_url[:-3].rstrip("/") # 去掉最后的 /v1
+    else:
+        base_url = raw_url
+
     client = genai.Client(
-        api_key=config['api_keys']['gemini_api_key']
+        api_key=config['api_keys']['gemini_api_key'],
+        http_options={'base_url': base_url} if base_url else None
     )
 
     # --- Gemini Call ---
@@ -314,14 +320,15 @@ 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}")
-
+    raw_url = config.get('api_base_url', '').strip().rstrip("/")
+    if raw_url.endswith("/v1"):
+        base_url = raw_url[:-3].rstrip("/") # 去掉最后的 /v1
+    else:
+        base_url = raw_url
     # 初始化 Client
     client = genai.Client(
-        api_key=config['api_keys']['gemini_api_key']
+        api_key=config['api_keys']['gemini_api_key'],
+        http_options={'base_url': base_url} if base_url else None
     )
 
     gpt_input = prompt_for_gpt + "\n\n" + \
@@ -544,15 +551,16 @@ def build_section_dags(
     FALLBACK_STRIP_BACKSLASH_ONLY_IN_CONTENT = True
     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}")
-        
+    # === Init Client (Gemini) ===      
+    raw_url = config.get('api_base_url', '').strip().rstrip("/")
+    if raw_url.endswith("/v1"):
+        base_url = raw_url[:-3].rstrip("/") # 去掉最后的 /v1
+    else:
+        base_url = raw_url
     # 使用 config 中的 key 
     client = genai.Client(
-        api_key=config['api_keys']['gemini_api_key']
+        api_key=config['api_keys']['gemini_api_key'],
+        http_options={'base_url': base_url} if base_url else None
     )
 
     def build_full_prompt(base_prompt: str, section_name: str, md_text: str) -> str: