knowledge/gemini.py

import os
from google import genai
from google.genai import types
import config


class GeminiProcessor:
    def __init__(self, api_key: str = None):
        """
        初始化Gemini客户端
        :param api_key: Gemini API密钥，如果为None则使用config中的密钥
        """
        self.api_key = api_key if api_key else config.GEMINI_API_KEY
        self.client = genai.Client(api_key=self.api_key)

    def process(self, content: str, system_prompt: str, model_name: str = "gemini-2.5-flash"):
        """
        处理文本内容
        :param content: 要处理的文本内容
        :param system_prompt: 系统提示语
        :param model_name: 使用的模型名称，默认为gemini-2.5-flash
        :return: 处理后的文本结果
        """
        input_text = ({content})

        try:
            # 调用Gemini模型处理内容
            response = self.client.models.generate_content(
                model=model_name,
                config=types.GenerateContentConfig(
                    system_instruction=system_prompt
                ),
                contents=input_text
            )
            return response.text
        except Exception as e:
            print(f"处理内容时出错: {e}")
            return None


# 使用示例
if __name__ == "__main__":
    processor = GeminiProcessor()
    result = processor.process(
        content="你好，请介绍一下你自己",
        system_prompt="你是一个有帮助的AI助手"
    )
    if result:
        print(result)
    print("\n处理完成")