From bc0fffaf9c12d6c2ec148bb9db4d8e16b201868b Mon Sep 17 00:00:00 2001 From: lorge <1833935022@qq.com> Date: Tue, 11 Feb 2025 05:40:34 +0800 Subject: [PATCH] =?UTF-8?q?feat=EF=BC=9A=E6=94=AF=E6=8C=81Gemini=E7=B3=BB?= =?UTF-8?q?=E5=88=97=E6=A8=A1=E5=9E=8B=E7=9A=84=E6=8E=A5=E5=8F=A3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- llm_adapters.py | 36 ++++++++++++++++++++++++++++++++++++ requirements.txt | Bin 328 -> 366 bytes tooltips.py | 4 ++-- ui.py | 4 +++- 4 files changed, 41 insertions(+), 3 deletions(-) diff --git a/llm_adapters.py b/llm_adapters.py index f09d06b..525b117 100644 --- a/llm_adapters.py +++ b/llm_adapters.py @@ -3,6 +3,8 @@ import logging from typing import Optional from langchain_openai import ChatOpenAI, AzureChatOpenAI +from google import genai +from google.genai import types def ensure_openai_base_url_has_v1(url: str) -> str: import re @@ -77,6 +79,38 @@ class OpenAIAdapter(BaseLLMAdapter): return "" return response.content +class GeminiAdapter(BaseLLMAdapter): + """ + 适配 Google Gemini 接口 + """ + def __init__(self, api_key: str, model_name: str, max_tokens: int, temperature: float = 0.7, timeout: Optional[int] = 600): + self.api_key = api_key + self.model_name = model_name + self.max_tokens = max_tokens + self.temperature = temperature + self.timeout = timeout + + self._client = genai.Client(api_key=self.api_key) + + def invoke(self, prompt: str) -> str: + try: + response = self._client.models.generate_content( + model = self.model_name, + contents = prompt, + config = types.GenerateContentConfig( + max_output_tokens=self.max_tokens, + temperature=self.temperature, + ) + ) + if response and response.text: + return response.text + else: + logging.warning("No text response from Gemini API.") + return "" + except Exception as e: + logging.error(f"Gemini API 调用失败: {e}") + return "" + class AzureOpenAIAdapter(BaseLLMAdapter): """ 适配 Azure OpenAI 接口（使用 langchain.ChatOpenAI） @@ -190,5 +224,7 @@ def create_llm_adapter( return OllamaAdapter(api_key, base_url, model_name, max_tokens, temperature, timeout) elif interface_format.lower() == "ml studio": return MLStudioAdapter(api_key, base_url, model_name, max_tokens, temperature, timeout) + elif interface_format.lower() == "gemini": + return GeminiAdapter(api_key, model_name, max_tokens, temperature, timeout) else: raise ValueError(f"Unknown interface_format: {interface_format}") diff --git a/requirements.txt b/requirements.txt index 981d39b908cc7e1d861cc64db11f772fedd02975..3c699d82120bd45d56313eece6bd278ef18b1d92 100644 GIT binary patch delta 46 ucmX@X^p0tR2cud#Lp~6uGvqL&GUx*7RE9htE@DVzC}GHCC