jihuaqiang 2 주 전
부모
커밋
121c7fbd65
4개의 변경된 파일46개의 추가작업 그리고 21개의 파일을 삭제
  1. 4 2
      agent.py
  2. 15 4
      gemini.py
  3. 16 8
      tools/indentify/image_identifier.py
  4. 11 7
      tools/indentify/video_identifier.py

+ 4 - 2
agent.py

@@ -516,13 +516,15 @@ def create_langgraph_workflow():
             
             # 获取7个不同的GEMINI API密钥
             api_keys = []
+            default_key = os.getenv('GEMINI_API_KEY') or os.getenv('GEMINI_API_KEY_1')
+            
             for i in range(1, 8):  # GEMINI_API_KEY_1 到 GEMINI_API_KEY_7
                 api_key = os.getenv(f'GEMINI_API_KEY_{i}')
                 if api_key:
                     api_keys.append(api_key)
                 else:
-                    logger.warning(f"未找到 GEMINI_API_KEY_{i},使用默认 GEMINI_API_KEY")
-                    api_keys.append(os.getenv('GEMINI_API_KEY_1'))
+                    logger.warning(f"未找到 GEMINI_API_KEY_{i},使用默认密钥")
+                    api_keys.append(default_key)
             
             # 准备多进程参数,为每个任务分配API密钥
             process_args = []

+ 15 - 4
gemini.py

@@ -19,18 +19,29 @@ class GeminiProcessor:
         # 加载环境变量
         load_dotenv()
         
-        # 获取API密钥
+        # 获取API密钥,支持多密钥配置
         self.api_key = os.getenv('GEMINI_API_KEY')
         if not self.api_key:
-            raise ValueError("未找到GEMINI_API_KEY环境变量")
+            # 如果没有基础密钥,尝试使用第一个多密钥
+            self.api_key = os.getenv('GEMINI_API_KEY_1')
+            if not self.api_key:
+                raise ValueError("未找到GEMINI_API_KEY或GEMINI_API_KEY_1环境变量")
         
-        # 配置Gemini
-        genai.configure(api_key=self.api_key)
+        # 延迟配置Gemini,在真正使用时再设置
+        self._configured = False
 
+    def _ensure_configured(self):
+        """确保Gemini已配置"""
+        if not self._configured:
+            genai.configure(api_key=self.api_key)
+            self._configured = True
     
     def process(self, content: Any, system_prompt: str) -> Dict[str, Any]:
 
         try:
+            # 确保Gemini已配置
+            self._ensure_configured()
+            
             # 处理输入内容格式
             if isinstance(content, dict):
                 # 将字典转换为JSON字符串

+ 16 - 8
tools/indentify/image_identifier.py

@@ -40,15 +40,22 @@ prompt = """
 class ImageIdentifier:
     def __init__(self):
         # 加载环境变量
-        load_dotenv()
+        load_dotenv()        
         
-        # 初始化Gemini API
-        api_key = os.getenv('GEMINI_API_KEY')
-        if not api_key:
-            raise ValueError("请在环境变量中设置 GEMINI_API_KEY")
-        
-        genai.configure(api_key=api_key)
-        self.model = genai.GenerativeModel('gemini-2.5-flash')
+        # 延迟配置Gemini,在真正使用时再设置
+        self._configured = False
+        self.model = None
+    
+    def _ensure_configured(self):
+        """确保Gemini已配置"""
+        if not self._configured:
+            self.api_key = os.getenv('GEMINI_API_KEY')
+            print(f"配置Gemini: {self.api_key}")
+            if not self.api_key:
+                raise ValueError("请在环境变量中设置 GEMINI_API_KEY")
+            genai.configure(api_key=self.api_key)
+            self.model = genai.GenerativeModel('gemini-2.5-flash')
+            self._configured = True
     
     def download_image(self, image_url: str) -> Optional[Image.Image]:
         """下载图片并转换为PIL Image对象"""
@@ -93,6 +100,7 @@ class ImageIdentifier:
                         return {"idx": idx, "url": url, "content": "", "success": False, "error": "图片下载失败"}
 
                     # 使用 Gemini 直接分析图片
+                    self._ensure_configured()
                     response = self.model.generate_content([system_prompt, image])
                     
                     if response.text:

+ 11 - 7
tools/indentify/video_identifier.py

@@ -36,13 +36,8 @@ class VideoIdentifier:
         # 加载环境变量
         load_dotenv()
         
-        # 获取API密钥
-        self.api_key = os.getenv('GEMINI_API_KEY')
-        if not self.api_key:
-            raise ValueError("未找到GEMINI_API_KEY环境变量")
-        
-        # 配置Gemini
-        genai.configure(api_key=self.api_key)
+        # 延迟配置Gemini,在真正使用时再设置
+        self._configured = False
         
         # 初始化缓存清理时间
         self.last_cache_cleanup = time.time()
@@ -77,6 +72,14 @@ class VideoIdentifier:
     ]
 }"""
     
+    def _ensure_configured(self):
+        """确保Gemini已配置"""
+        if not self._configured:
+            self.api_key = os.getenv('GEMINI_API_KEY')
+            print(f"配置Gemini: {self.api_key}")
+            genai.configure(api_key=self.api_key)
+            self._configured = True
+    
     def download_video(self, video_url: str) -> Optional[str]:
         """下载视频到本地缓存"""
         file_path = os.path.join(CACHE_DIR, f'{str(uuid.uuid4())}.mp4')
@@ -391,6 +394,7 @@ class VideoIdentifier:
                         }
 
                     # 3. 使用 Gemini 直接分析视频文件
+                    self._ensure_configured()
                     model = genai.GenerativeModel(
                         model_name='gemini-2.5-flash',
                         generation_config=genai.GenerationConfig(