mirror of
https://github.com/zhayujie/chatgpt-on-wechat.git
synced 2026-07-21 06:07:13 +08:00
feat: support skills
This commit is contained in:
saboteur7
@@ -6,6 +6,8 @@ Google gemini bot
|
||||
"""
|
||||
# encoding:utf-8
|
||||
|
||||
import json
|
||||
import time
|
||||
from bot.bot import Bot
|
||||
import google.generativeai as genai
|
||||
from bot.session_manager import SessionManager
|
||||
@@ -113,3 +115,224 @@ class GoogleGeminiBot(Bot):
|
||||
elif turn == "assistant":
|
||||
turn = "user"
|
||||
return res
|
||||
|
||||
def call_with_tools(self, messages, tools=None, stream=False, **kwargs):
|
||||
"""
|
||||
Call Gemini API with tool support for agent integration
|
||||
|
||||
Args:
|
||||
messages: List of messages
|
||||
tools: List of tool definitions (OpenAI format, will be converted to Gemini format)
|
||||
stream: Whether to use streaming
|
||||
**kwargs: Additional parameters
|
||||
|
||||
Returns:
|
||||
Formatted response compatible with OpenAI format or generator for streaming
|
||||
"""
|
||||
try:
|
||||
# Configure Gemini
|
||||
genai.configure(api_key=self.api_key)
|
||||
model_name = kwargs.get("model", self.model)
|
||||
|
||||
# Extract system prompt from messages
|
||||
system_prompt = kwargs.get("system", "")
|
||||
gemini_messages = []
|
||||
|
||||
for msg in messages:
|
||||
if msg.get("role") == "system":
|
||||
system_prompt = msg["content"]
|
||||
else:
|
||||
gemini_messages.append(msg)
|
||||
|
||||
# Convert messages to Gemini format
|
||||
gemini_messages = self._convert_to_gemini_messages(gemini_messages)
|
||||
|
||||
# Safety settings
|
||||
safety_settings = {
|
||||
HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
|
||||
HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
|
||||
HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
|
||||
HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
|
||||
}
|
||||
|
||||
# Convert tools from OpenAI format to Gemini format if provided
|
||||
gemini_tools = None
|
||||
if tools:
|
||||
gemini_tools = self._convert_tools_to_gemini_format(tools)
|
||||
|
||||
# Create model with system instruction if available
|
||||
model_kwargs = {"model_name": model_name}
|
||||
if system_prompt:
|
||||
model_kwargs["system_instruction"] = system_prompt
|
||||
|
||||
model = genai.GenerativeModel(**model_kwargs)
|
||||
|
||||
# Generate content
|
||||
generation_config = {}
|
||||
if kwargs.get("max_tokens"):
|
||||
generation_config["max_output_tokens"] = kwargs["max_tokens"]
|
||||
if kwargs.get("temperature") is not None:
|
||||
generation_config["temperature"] = kwargs["temperature"]
|
||||
|
||||
request_params = {
|
||||
"safety_settings": safety_settings
|
||||
}
|
||||
if generation_config:
|
||||
request_params["generation_config"] = generation_config
|
||||
if gemini_tools:
|
||||
request_params["tools"] = gemini_tools
|
||||
|
||||
if stream:
|
||||
return self._handle_gemini_stream_response(model, gemini_messages, request_params, model_name)
|
||||
else:
|
||||
return self._handle_gemini_sync_response(model, gemini_messages, request_params, model_name)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"[Gemini] call_with_tools error: {e}")
|
||||
if stream:
|
||||
def error_generator():
|
||||
yield {
|
||||
"error": True,
|
||||
"message": str(e),
|
||||
"status_code": 500
|
||||
}
|
||||
return error_generator()
|
||||
else:
|
||||
return {
|
||||
"error": True,
|
||||
"message": str(e),
|
||||
"status_code": 500
|
||||
}
|
||||
|
||||
def _convert_tools_to_gemini_format(self, openai_tools):
|
||||
"""Convert OpenAI tool format to Gemini function declarations"""
|
||||
import google.generativeai as genai
|
||||
|
||||
gemini_functions = []
|
||||
for tool in openai_tools:
|
||||
if tool.get("type") == "function":
|
||||
func = tool.get("function", {})
|
||||
gemini_functions.append(
|
||||
genai.protos.FunctionDeclaration(
|
||||
name=func.get("name"),
|
||||
description=func.get("description", ""),
|
||||
parameters=func.get("parameters", {})
|
||||
)
|
||||
)
|
||||
|
||||
if gemini_functions:
|
||||
return [genai.protos.Tool(function_declarations=gemini_functions)]
|
||||
return None
|
||||
|
||||
def _handle_gemini_sync_response(self, model, messages, request_params, model_name):
|
||||
"""Handle synchronous Gemini API response"""
|
||||
import json
|
||||
|
||||
response = model.generate_content(messages, **request_params)
|
||||
|
||||
# Extract text content and function calls
|
||||
text_content = ""
|
||||
tool_calls = []
|
||||
|
||||
if response.candidates and response.candidates[0].content:
|
||||
for part in response.candidates[0].content.parts:
|
||||
if hasattr(part, 'text') and part.text:
|
||||
text_content += part.text
|
||||
elif hasattr(part, 'function_call') and part.function_call:
|
||||
# Convert Gemini function call to OpenAI format
|
||||
func_call = part.function_call
|
||||
tool_calls.append({
|
||||
"id": f"call_{hash(func_call.name)}",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": func_call.name,
|
||||
"arguments": json.dumps(dict(func_call.args))
|
||||
}
|
||||
})
|
||||
|
||||
# Build message in OpenAI format
|
||||
message = {
|
||||
"role": "assistant",
|
||||
"content": text_content
|
||||
}
|
||||
if tool_calls:
|
||||
message["tool_calls"] = tool_calls
|
||||
|
||||
# Format response to match OpenAI structure
|
||||
formatted_response = {
|
||||
"id": f"gemini_{int(time.time())}",
|
||||
"object": "chat.completion",
|
||||
"created": int(time.time()),
|
||||
"model": model_name,
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": message,
|
||||
"finish_reason": "stop" if not tool_calls else "tool_calls"
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"prompt_tokens": 0, # Gemini doesn't provide token counts in the same way
|
||||
"completion_tokens": 0,
|
||||
"total_tokens": 0
|
||||
}
|
||||
}
|
||||
|
||||
logger.info(f"[Gemini] call_with_tools reply, model={model_name}")
|
||||
return formatted_response
|
||||
|
||||
def _handle_gemini_stream_response(self, model, messages, request_params, model_name):
|
||||
"""Handle streaming Gemini API response"""
|
||||
import json
|
||||
|
||||
try:
|
||||
response_stream = model.generate_content(messages, stream=True, **request_params)
|
||||
|
||||
for chunk in response_stream:
|
||||
if chunk.candidates and chunk.candidates[0].content:
|
||||
for part in chunk.candidates[0].content.parts:
|
||||
if hasattr(part, 'text') and part.text:
|
||||
# Text content
|
||||
yield {
|
||||
"id": f"gemini_{int(time.time())}",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": int(time.time()),
|
||||
"model": model_name,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {"content": part.text},
|
||||
"finish_reason": None
|
||||
}]
|
||||
}
|
||||
elif hasattr(part, 'function_call') and part.function_call:
|
||||
# Function call
|
||||
func_call = part.function_call
|
||||
yield {
|
||||
"id": f"gemini_{int(time.time())}",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": int(time.time()),
|
||||
"model": model_name,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"tool_calls": [{
|
||||
"index": 0,
|
||||
"id": f"call_{hash(func_call.name)}",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": func_call.name,
|
||||
"arguments": json.dumps(dict(func_call.args))
|
||||
}
|
||||
}]
|
||||
},
|
||||
"finish_reason": None
|
||||
}]
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"[Gemini] stream response error: {e}")
|
||||
yield {
|
||||
"error": True,
|
||||
"message": str(e),
|
||||
"status_code": 500
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user