gemini_mcp_server.py•10.1 kB
import os
from typing import Optional
from fastmcp import FastMCP
from openai import OpenAI
import httpx
from dotenv import load_dotenv
# Load environment variables
load_dotenv()
# Initialize FastMCP server
mcp = FastMCP("Gemini MCP Server")
def setup_proxy():
"""Setup proxy configuration"""
http_proxy = os.getenv("http_proxy") or os.getenv("HTTP_PROXY")
https_proxy = os.getenv("https_proxy") or os.getenv("HTTPS_PROXY")
proxies = {}
if http_proxy:
proxies['http'] = http_proxy
if https_proxy:
proxies['https'] = https_proxy
return proxies
def get_config():
"""Get API configuration"""
api_key = os.getenv("GEMINI_API_KEY")
if not api_key:
raise ValueError("GEMINI_API_KEY environment variable is not set")
# Read model and URL configuration from environment variables with default values
model = os.getenv("GEMINI_MODEL", "gemini-2.5-flash")
base_url = os.getenv("GEMINI_BASE_URL", "https://generativelanguage.googleapis.com/v1beta/openai/")
return api_key, model, base_url
def get_gemini_client():
"""Get configured Gemini client"""
api_key, model, base_url = get_config()
proxies = setup_proxy()
http_client = None
if proxies:
# Use HTTPS proxy as API calls are usually HTTPS
proxy_url = proxies.get('https') or proxies.get('http')
http_client = httpx.Client(proxy=proxy_url)
client = OpenAI(
api_key=api_key,
base_url=base_url,
http_client=http_client
)
else:
client = OpenAI(
api_key=api_key,
base_url=base_url,
)
return client, model, http_client
@mcp.tool()
def call_gemini_stream_thinking(prompt: str, system_message: Optional[str] = "You are a helpful AI assistant.", reasoning_effort: Optional[str] = "high") -> dict:
"""Call Gemini API - Stream thinking mode"""
try:
client, model, http_client = get_gemini_client()
messages = [
{"role": "system", "content": system_message},
{"role": "user", "content": prompt}
]
# Stream thinking mode
response = client.chat.completions.create(
model=model,
messages=messages,
reasoning_effort=reasoning_effort,
stream=True
)
# Collect streaming response
full_response = ""
for chunk in response:
if chunk.choices[0].delta.content is not None:
full_response += chunk.choices[0].delta.content
if http_client:
http_client.close()
return {
"status": "success",
"message": "Stream thinking mode call successful",
"model": model,
"mode": "stream_thinking",
"reasoning_effort": reasoning_effort,
"response": full_response,
"prompt": prompt
}
except Exception as e:
if 'http_client' in locals() and http_client:
http_client.close()
return {
"status": "error",
"message": f"Stream thinking mode call failed: {str(e)}",
"mode": "stream_thinking"
}
@mcp.tool()
def call_gemini_stream_no_thinking(prompt: str, system_message: Optional[str] = "You are a helpful AI assistant.", reasoning_effort: Optional[str] = "high") -> dict:
"""Call Gemini API - Stream no thinking mode"""
try:
client, model, http_client = get_gemini_client()
messages = [
{"role": "system", "content": system_message},
{"role": "user", "content": prompt}
]
# Stream normal mode
response = client.chat.completions.create(
model=model,
messages=messages,
reasoning_effort=reasoning_effort,
stream=True
)
# Collect streaming response
full_response = ""
for chunk in response:
if chunk.choices[0].delta.content is not None:
full_response += chunk.choices[0].delta.content
if http_client:
http_client.close()
return {
"status": "success",
"message": "Stream no thinking mode call successful",
"model": model,
"mode": "stream_no_thinking",
"reasoning_effort": reasoning_effort,
"response": full_response,
"prompt": prompt
}
except Exception as e:
if 'http_client' in locals() and http_client:
http_client.close()
return {
"status": "error",
"message": f"Stream no thinking mode call failed: {str(e)}",
"mode": "stream_no_thinking"
}
@mcp.tool()
def call_gemini_non_stream_thinking(prompt: str, system_message: Optional[str] = "You are a helpful AI assistant.", reasoning_effort: Optional[str] = "high") -> dict:
"""Call Gemini API - Non-stream thinking mode"""
try:
client, model, http_client = get_gemini_client()
messages = [
{"role": "system", "content": system_message},
{"role": "user", "content": prompt}
]
# Non-stream thinking mode
response = client.chat.completions.create(
model=model,
messages=messages,
reasoning_effort=reasoning_effort
)
if http_client:
http_client.close()
return {
"status": "success",
"message": "Non-stream thinking mode call successful",
"model": model,
"mode": "non_stream_thinking",
"reasoning_effort": reasoning_effort,
"response": response.choices[0].message.content,
"prompt": prompt,
"usage": {
"prompt_tokens": getattr(response.usage, 'prompt_tokens', 0),
"completion_tokens": getattr(response.usage, 'completion_tokens', 0),
"total_tokens": getattr(response.usage, 'total_tokens', 0)
}
}
except Exception as e:
if 'http_client' in locals() and http_client:
http_client.close()
return {
"status": "error",
"message": f"Non-stream thinking mode call failed: {str(e)}",
"mode": "non_stream_thinking"
}
@mcp.tool()
def call_gemini_non_stream_no_thinking(prompt: str, system_message: Optional[str] = "You are a helpful AI assistant.", reasoning_effort: Optional[str] = "high") -> dict:
"""Call Gemini API - Non-stream no thinking mode"""
try:
client, model, http_client = get_gemini_client()
messages = [
{"role": "system", "content": system_message},
{"role": "user", "content": prompt}
]
# Non-stream normal mode
response = client.chat.completions.create(
model=model,
messages=messages,
reasoning_effort=reasoning_effort
)
if http_client:
http_client.close()
return {
"status": "success",
"message": "Non-stream no thinking mode call successful",
"model": model,
"mode": "non_stream_no_thinking",
"reasoning_effort": reasoning_effort,
"response": response.choices[0].message.content,
"prompt": prompt,
"usage": {
"prompt_tokens": getattr(response.usage, 'prompt_tokens', 0),
"completion_tokens": getattr(response.usage, 'completion_tokens', 0),
"total_tokens": getattr(response.usage, 'total_tokens', 0)
}
}
except Exception as e:
if 'http_client' in locals() and http_client:
http_client.close()
return {
"status": "error",
"message": f"Non-stream no thinking mode call failed: {str(e)}",
"mode": "non_stream_no_thinking"
}
@mcp.tool()
def get_gemini_config() -> dict:
"""Get current Gemini configuration information"""
try:
api_key, model, base_url = get_config()
proxies = setup_proxy()
return {
"status": "success",
"message": "Get configuration information successful",
"config": {
"model": model,
"base_url": base_url,
"api_key_set": bool(api_key),
"proxy_enabled": bool(proxies),
"proxy_config": proxies if proxies else "No proxy configured"
}
}
except Exception as e:
return {
"status": "error",
"message": f"Get gemini server configuration information failed: {str(e)}"
}
@mcp.tool()
def test_gemini_connection() -> dict:
"""Test Gemini API connection"""
try:
client, model, http_client = get_gemini_client()
# Send simple test request
test_messages = [
{"role": "system", "content": "You are an AI assistant."},
{"role": "user", "content": "Please simply answer: Hello"}
]
response = client.chat.completions.create(
model=model,
messages=test_messages,
reasoning_effort="low"
)
if http_client:
http_client.close()
return {
"status": "success",
"message": "Gemini API connection test successful",
"model": model,
"test_response": response.choices[0].message.content
}
except Exception as e:
if 'http_client' in locals() and http_client:
http_client.close()
return {
"status": "error",
"message": f"Gemini API connection test failed: {str(e)}"
}
if __name__ == "__main__":
mcp.run()