修改获取网页内容代码

ArticleReplaceBatch/main_process.py

@@ -1,11 +1,9 @@
-import pandas as pd
-import getpass
-import sys  # 导入sys模块
+
 import threading
 import queue
 
 
-from ai_studio import call_dify_workflow
+from ai_studio import call_dify_workflow, call_coze_workflow
 from databases import *
 
 
@@ -15,9 +13,11 @@ from get_web_content import  *
 from config import *
 
 # ==============================主程序===========================
-def process_link(link):
+def process_link(link, ai_service="dify"):
     """
     处理单个链接
+    :param link: 要处理的链接
+    :param ai_service: AI服务提供商，可选值：dify, coze
     """
     try:
         if link.startswith("https://www.toutiao.com"):
@@ -63,7 +63,16 @@ def process_link(link):
             "old_article": article_text
         }
 
-        message_content = call_dify_workflow(input_data)
+        if ai_service == "dify":
+            input_data = {
+                "old_article": article_text
+            }
+            message_content = call_dify_workflow(input_data)
+        elif ai_service == "coze":
+            input_data = {
+                "old_article": article_text
+            }
+            message_content = call_coze_workflow(input_data)
         # 获取当前时间并格式化
         current_time = datetime.now().strftime("%H:%M:%S")
 
@@ -109,7 +118,7 @@ def process_link(link):
 
 
 
-def link_to_text(prompt1=None, prompt2=None, num_threads=None):
+def link_to_text(prompt1=None, prompt2=None, num_threads=None, ai_service="dify"):
     use_link_path = 'use_link_path.txt'
 
     # 读取链接
@@ -165,7 +174,7 @@ def worker():
             # 处理链接
             try:
                 logger.info(f"开始处理链接：{link}")
-                process_link(link)
+                process_link(link, ai_service)
                 result_queue.put((link, True, None))  # 成功
             except Exception as e:
                 result_queue.put((link, False, str(e)))  # 失败