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]:
        """
        处理内容结构化
        
        Args:
            content: 要处理的内容
            system_prompt: 系统提示词
            
        Returns:
            处理结果
        """
        try:
            # 构建完整的提示词
            full_prompt = f"{system_prompt}\n\n内容：{json.dumps(content, ensure_ascii=False)}"
            
            # 调用 Gemini API
            response = self.model.generate_content(full_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}