feat:支持Gemini系列模型的接口

This commit is contained in:
lorge
2025-02-11 05:40:34 +08:00
parent 6d90d524b1
commit bc0fffaf9c
4 changed files with 41 additions and 3 deletions
+36
View File
@@ -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}")