From b124bd0d0ef7a84787ca5c6c52a0f14fe3646816 Mon Sep 17 00:00:00 2001 From: Soulter <905617992@qq.com> Date: Thu, 9 Jan 2025 14:05:48 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=94=AF=E6=8C=81=E9=80=9A=E8=BF=87=20?= =?UTF-8?q?Google=20GenAI=20=E8=AE=BF=E9=97=AE=20Gemini=20=E6=A8=A1?= =?UTF-8?q?=E5=9E=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- astrbot/core/config/default.py | 12 +- astrbot/core/provider/entites.py | 8 +- astrbot/core/provider/func_tool_manager.py | 17 ++ astrbot/core/provider/manager.py | 2 + .../core/provider/sources/gemini_source.py | 287 ++++++++++++++++++ 5 files changed, 321 insertions(+), 5 deletions(-) create mode 100644 astrbot/core/provider/sources/gemini_source.py diff --git a/astrbot/core/config/default.py b/astrbot/core/config/default.py index a1a0a2208..3eb3519fb 100644 --- a/astrbot/core/config/default.py +++ b/astrbot/core/config/default.py @@ -257,7 +257,7 @@ CONFIG_METADATA_2 = { "model": "llama3.1-8b", }, }, - "gemini": { + "gemini(OpenAI兼容)": { "id": "gemini_default", "type": "openai_chat_completion", "enable": True, @@ -267,6 +267,16 @@ CONFIG_METADATA_2 = { "model": "gemini-1.5-flash", }, }, + "gemini(googlegenai原生)": { + "id": "gemini_default", + "type": "googlegenai_chat_completion", + "enable": True, + "key": [], + "api_base": "https://generativelanguage.googleapis.com/", + "model_config": { + "model": "gemini-1.5-flash", + }, + }, "deepseek": { "id": "deepseek_default", "type": "openai_chat_completion", diff --git a/astrbot/core/provider/entites.py b/astrbot/core/provider/entites.py index 3ee3379f1..8dae2680d 100644 --- a/astrbot/core/provider/entites.py +++ b/astrbot/core/provider/entites.py @@ -1,4 +1,4 @@ -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import List, Dict from .func_tool_manager import FuncCall @@ -32,9 +32,9 @@ class ProviderRequest(): class LLMResponse: role: str '''角色''' - completion_text: str = None + completion_text: str = "" '''LLM 返回的文本''' - tools_call_args: List[Dict[str, any]] = None + tools_call_args: List[Dict[str, any]] = field(default_factory=list) '''工具调用参数''' - tools_call_name: List[str] = None + tools_call_name: List[str] = field(default_factory=list) '''工具调用名称''' \ No newline at end of file diff --git a/astrbot/core/provider/func_tool_manager.py b/astrbot/core/provider/func_tool_manager.py index f1d2f7b28..0c23e412c 100644 --- a/astrbot/core/provider/func_tool_manager.py +++ b/astrbot/core/provider/func_tool_manager.py @@ -101,6 +101,23 @@ class FuncCall: } ) return _l + + def get_func_desc_google_genai_style(self) -> Dict: + declarations = {} + tools = [] + for f in self.func_list: + if not f.active: + continue + tools.append( + { + "name": f.name, + "parameters": f.parameters, + "description": f.description, + } + ) + declarations["function_declarations"] = tools + return declarations + async def func_call(self, question: str, session_id: str, provider) -> tuple: _l = [] diff --git a/astrbot/core/provider/manager.py b/astrbot/core/provider/manager.py index e020148d6..19075338c 100644 --- a/astrbot/core/provider/manager.py +++ b/astrbot/core/provider/manager.py @@ -41,6 +41,8 @@ class ProviderManager(): from .sources.llmtuner_source import LLMTunerModelLoader # noqa: F401 case "dify": from .sources.dify_source import ProviderDify # noqa: F401 + case "googlegenai_chat_completion": + from .sources.gemini_source import ProviderGoogleGenAI # noqa: F401 async def initialize(self): diff --git a/astrbot/core/provider/sources/gemini_source.py b/astrbot/core/provider/sources/gemini_source.py new file mode 100644 index 000000000..ac37e989b --- /dev/null +++ b/astrbot/core/provider/sources/gemini_source.py @@ -0,0 +1,287 @@ +import traceback +import base64 +import json +import aiohttp +from astrbot.core.utils.io import download_image_by_url +from astrbot.core.db import BaseDatabase +from astrbot.api.provider import Provider +from astrbot import logger +from astrbot.core.provider.func_tool_manager import FuncCall +from typing import List +from ..register import register_provider_adapter +from astrbot.core.provider.entites import LLMResponse + +class SimpleGoogleGenAIClient(): + def __init__(self, api_key: str, api_base: str): + self.api_key = api_key + if api_base.endswith("/"): + self.api_base = api_base[:-1] + else: + self.api_base = api_base + self.client = aiohttp.ClientSession() + + async def models_list(self) -> List[str]: + request_url = f"{self.api_base}/v1beta/models?key={self.api_key}" + async with self.client.get(request_url, timeout=10) as resp: + response = await resp.json() + + models = [] + for model in response["models"]: + if 'generateContent' in model["supportedGenerationMethods"]: + models.append(model["name"].replace("models/", "")) + return models + + async def generate_content( + self, + contents: List[dict], + model: str="gemini-1.5-flash", + system_instruction: str="", + tools: dict=None + ): + payload = {} + if system_instruction: + payload["system_instruction"] = { + "parts": {"text": system_instruction} + } + if tools: + payload["tools"] = [tools] + payload["contents"] = contents + logger.debug(f"payload: {payload}") + request_url = f"{self.api_base}/v1beta/models/{model}:generateContent?key={self.api_key}" + async with self.client.post(request_url, json=payload, timeout=10) as resp: + response = await resp.json() + return response + + +@register_provider_adapter("googlegenai_chat_completion", "Google Gemini Chat Completion 提供商适配器") +class ProviderGoogleGenAI(Provider): + def __init__( + self, + provider_config: dict, + provider_settings: dict, + db_helper: BaseDatabase, + persistant_history = True + ) -> None: + super().__init__(provider_config, provider_settings, persistant_history, db_helper) + self.chosen_api_key = None + self.api_keys: List = provider_config.get("key", []) + self.chosen_api_key = self.api_keys[0] if len(self.api_keys) > 0 else None + + self.client = SimpleGoogleGenAIClient( + api_key=self.chosen_api_key, + api_base=provider_config.get("api_base", None) + ) + self.set_model(provider_config['model_config']['model']) + + async def get_human_readable_context(self, session_id, page, page_size): + if session_id not in self.session_memory: + raise Exception("会话 ID 不存在") + contexts = [] + temp_contexts = [] + for record in self.session_memory[session_id]: + if record['role'] == "user": + temp_contexts.append(f"User: {record['content']}") + elif record['role'] == "assistant": + temp_contexts.append(f"Assistant: {record['content']}") + contexts.insert(0, temp_contexts) + temp_contexts = [] + + # 展平 contexts 列表 + contexts = [item for sublist in contexts for item in sublist] + + # 计算分页 + paged_contexts = contexts[(page-1)*page_size:page*page_size] + total_pages = len(contexts) // page_size + if len(contexts) % page_size != 0: + total_pages += 1 + + return paged_contexts, total_pages + + async def get_models(self): + return await self.client.models_list() + + async def pop_record(self, session_id: str, pop_system_prompt: bool = False): + ''' + 弹出第一条记录 + ''' + if session_id not in self.session_memory: + raise Exception("会话 ID 不存在") + + if len(self.session_memory[session_id]) == 0: + return None + + for i in range(len(self.session_memory[session_id])): + # 检查是否是 system prompt + if not pop_system_prompt and self.session_memory[session_id][i]['user']['role'] == "system": + # 如果只有一个 system prompt，才不删掉 + f = False + for j in range(i+1, len(self.session_memory[session_id])): + if self.session_memory[session_id][j]['user']['role'] == "system": + f = True + break + if not f: + continue + record = self.session_memory[session_id].pop(i) + break + + return record + + async def _query(self, payloads: dict, tools: FuncCall) -> LLMResponse: + tool = None + if tools: + tool = tools.get_func_desc_google_genai_style() + + system_instruction = "" + for message in payloads["messages"]: + if message["role"] == "system": + system_instruction = message["content"] + break + + google_genai_conversation = [] + for message in payloads["messages"]: + if message["role"] == "user": + if isinstance(message["content"], str): + google_genai_conversation.append({ + "role": "user", + "parts": [{"text": message["content"]}] + }) + elif isinstance(message["content"], list): + # images + parts = [] + for part in message["content"]: + if part["type"] == "text": + parts.append({"text": part["text"]}) + elif part["type"] == "image_url": + parts.append({"inline_data": { + "mime_type": "image/jpeg", + "data": part["image_url"]["url"].replace("data:image/jpeg;base64,", "") # base64 + }}) + google_genai_conversation.append({ + "role": "user", + "parts": parts + }) + + elif message["role"] == "assistant": + google_genai_conversation.append({ + "role": "model", + "parts": [{"text": message["content"]}] + }) + + + logger.debug(f"google_genai_conversation: {google_genai_conversation}") + + result = await self.client.generate_content( + contents=google_genai_conversation, + model=self.get_model(), + system_instruction=system_instruction, + tools=tool + ) + logger.debug(f"result: {result}") + + candidates = result["candidates"][0]['content']['parts'] + llm_response = LLMResponse("assistant") + for candidate in candidates: + if 'text' in candidate: + llm_response.completion_text += candidate['text'] + elif 'functionCall' in candidate: + llm_response.role = "tool" + llm_response.tools_call_args.append(candidate['functionCall']['args']) + llm_response.tools_call_name.append(candidate['functionCall']['name']) + + return llm_response + + + async def text_chat( + self, + prompt: str, + session_id: str, + image_urls: List[str]=None, + func_tool: FuncCall=None, + contexts=None, + system_prompt=None, + **kwargs + ) -> LLMResponse: + new_record = await self.assemble_context(prompt, image_urls) + context_query = [] + if not contexts: + context_query = [*self.session_memory[session_id], new_record] + else: + context_query = [*contexts, new_record] + if system_prompt: + context_query.insert(0, {"role": "system", "content": system_prompt}) + + payloads = { + "messages": context_query, + **self.provider_config.get("model_config", {}) + } + + try: + llm_response = await self._query(payloads, func_tool) + except Exception as e: + if "maximum context length" in str(e): + logger.warning(f"请求失败：{e}。上下文长度超过限制。尝试弹出最早的记录然后重试。") + self.pop_record(session_id) + logger.warning(traceback.format_exc()) + + await self.save_history(contexts, new_record, session_id, llm_response) + + return llm_response + + async def save_history(self, contexts: List, new_record: dict, session_id: str, llm_response: LLMResponse): + if llm_response.role == "assistant" and session_id: + # 文本回复 + if not contexts: + # 添加用户 record + self.session_memory[session_id].append(new_record) + # 添加 assistant record + self.session_memory[session_id].append({ + "role": "assistant", + "content": llm_response.completion_text + }) + else: + self.session_memory[session_id] = [*contexts, new_record, { + "role": "assistant", + "content": llm_response.completion_text + }] + self.db_helper.update_llm_history(session_id, json.dumps(self.session_memory[session_id]), self.provider_config['type']) + + async def forget(self, session_id: str) -> bool: + self.session_memory[session_id] = [] + return True + + def get_current_key(self) -> str: + return self.client.api_key + + def get_keys(self) -> List[str]: + return self.api_keys + + def set_key(self, key): + self.client.api_key = key + + async def assemble_context(self, text: str, image_urls: List[str] = None): + ''' + 组装上下文。 + ''' + if image_urls: + user_content = {"role": "user","content": [{"type": "text", "text": text}]} + for image_url in image_urls: + if image_url.startswith("http"): + image_path = await download_image_by_url(image_url) + image_data = await self.encode_image_bs64(image_path) + else: + image_data = await self.encode_image_bs64(image_url) + user_content["content"].append({"type": "image_url", "image_url": {"url": image_data}}) + return user_content + else: + return {"role": "user","content": text} + + async def encode_image_bs64(self, image_url: str) -> str: + ''' + 将图片转换为 base64 + ''' + if image_url.startswith("base64://"): + return image_url.replace("base64://", "data:image/jpeg;base64,") + with open(image_url, "rb") as f: + image_bs64 = base64.b64encode(f.read()).decode('utf-8') + return "data:image/jpeg;base64," + image_bs64 + return '' \ No newline at end of file