knowledge-agent/gemini.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
简单的 Gemini 处理器
用于满足导入需求，实际功能可以根据需要扩展
"""

import os
import json
from typing import Any, Dict, Optional
from dotenv import load_dotenv
import google.generativeai as genai


class GeminiProcessor:
    """Gemini API 处理器"""
    
    def __init__(self):
        # 加载环境变量
        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)
        self.model = genai.GenerativeModel('gemini-2.5-flash')
    
    def process(self, content: Any, system_prompt: str) -> Dict[str, Any]:

        try:
            # 构建完整的提示词
            full_prompt = f"{system_prompt}\n\n内容：{json.dumps(content, ensure_ascii=False)}"
            
            # 调用 Gemini API
            response = self.model.generate_content(
                contents=content,
                config=types.GenerateContentConfig(
                    system_instruction=system_prompt
                )
            )
            
            # 尝试解析 JSON 响应
            try:
                result = json.loads(response.text)
                return result
            except json.JSONDecodeError:
                # 如果不是 JSON 格式，返回原始文本
                return {"result": response.text, "raw_response": response.text}
                
        except Exception as e:
            return {"error": str(e), "content": content} 
            
    def batch_process(self, contents: list, system_prompt: str) -> list:
        results = []
        for content in contents:
            result = self.process(content, system_prompt)
            results.append(result)
        return results