知微
fa45d8aa5f
- health_checklist.json: 192.168.1.122→node122
- ocr_client.py: docstring IP→node122
- docs/market-data-requirements.md: IP→node122
- 所有API调用通过ProxyHandler({})绕过系统代理
Privoxy对node122:18003返回500，直连正常
545 lines
17 KiB
Python
545 lines
17 KiB
Python
import base64
|
|
import datetime
|
|
import json
|
|
import math
|
|
from typing import Any, Dict, List, Optional, Sequence, Union
|
|
|
|
import httpx
|
|
|
|
import litellm
|
|
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
|
|
from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter
|
|
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
|
from litellm.secret_managers.main import get_secret_str
|
|
from litellm.types.llms.openai import AllMessageValues
|
|
from litellm.types.utils import TokenCountResponse
|
|
|
|
GEMINI_IMAGE_ASPECT_RATIOS: Dict[str, float] = {
|
|
"1:1": 1 / 1,
|
|
"1:4": 1 / 4,
|
|
"1:8": 1 / 8,
|
|
"2:3": 2 / 3,
|
|
"3:2": 3 / 2,
|
|
"3:4": 3 / 4,
|
|
"4:1": 4 / 1,
|
|
"4:3": 4 / 3,
|
|
"4:5": 4 / 5,
|
|
"5:4": 5 / 4,
|
|
"8:1": 8 / 1,
|
|
"9:16": 9 / 16,
|
|
"16:9": 16 / 9,
|
|
"21:9": 21 / 9,
|
|
}
|
|
|
|
# Supported aspect ratio dimensions from Google Gemini image generation docs:
|
|
# https://ai.google.dev/gemini-api/docs/image-generation#aspect_ratios_and_image_size
|
|
GEMINI_IMAGE_SIZE_TO_ASPECT_RATIO: Dict[tuple[int, int], str] = {
|
|
(512, 512): "1:1",
|
|
(1024, 1024): "1:1",
|
|
(2048, 2048): "1:1",
|
|
(4096, 4096): "1:1",
|
|
(256, 1024): "1:4",
|
|
(512, 2048): "1:4",
|
|
(1024, 4096): "1:4",
|
|
(2048, 8192): "1:4",
|
|
(192, 1536): "1:8",
|
|
(384, 3072): "1:8",
|
|
(768, 6144): "1:8",
|
|
(1536, 12288): "1:8",
|
|
(424, 632): "2:3",
|
|
(848, 1264): "2:3",
|
|
(1696, 2528): "2:3",
|
|
(3392, 5056): "2:3",
|
|
(632, 424): "3:2",
|
|
(1264, 848): "3:2",
|
|
(2528, 1696): "3:2",
|
|
(5056, 3392): "3:2",
|
|
(448, 600): "3:4",
|
|
(896, 1200): "3:4",
|
|
(1792, 2400): "3:4",
|
|
(3584, 4800): "3:4",
|
|
(1024, 256): "4:1",
|
|
(2048, 512): "4:1",
|
|
(4096, 1024): "4:1",
|
|
(8192, 2048): "4:1",
|
|
(600, 448): "4:3",
|
|
(1200, 896): "4:3",
|
|
(2400, 1792): "4:3",
|
|
(4800, 3584): "4:3",
|
|
(464, 576): "4:5",
|
|
(928, 1152): "4:5",
|
|
(1856, 2304): "4:5",
|
|
(3712, 4608): "4:5",
|
|
(576, 464): "5:4",
|
|
(1152, 928): "5:4",
|
|
(2304, 1856): "5:4",
|
|
(4608, 3712): "5:4",
|
|
(1536, 192): "8:1",
|
|
(3072, 384): "8:1",
|
|
(6144, 768): "8:1",
|
|
(12288, 1536): "8:1",
|
|
(384, 688): "9:16",
|
|
(768, 1376): "9:16",
|
|
(1536, 2752): "9:16",
|
|
(3072, 5504): "9:16",
|
|
(688, 384): "16:9",
|
|
(1376, 768): "16:9",
|
|
(2752, 1536): "16:9",
|
|
(5504, 3072): "16:9",
|
|
(792, 336): "21:9",
|
|
(1584, 672): "21:9",
|
|
(3168, 1344): "21:9",
|
|
(6336, 2688): "21:9",
|
|
(1280, 896): "4:3",
|
|
(896, 1280): "3:4",
|
|
}
|
|
|
|
|
|
def map_openai_size_to_gemini_image_config(
|
|
size: str, model: str
|
|
) -> Optional[Dict[str, str]]:
|
|
dimensions = _parse_openai_image_size(size)
|
|
if dimensions is None:
|
|
return None
|
|
|
|
width, height = dimensions
|
|
image_config = {
|
|
"aspectRatio": _map_dimensions_to_gemini_aspect_ratio(width, height)
|
|
}
|
|
image_size = _map_dimensions_to_gemini_image_size(width, height)
|
|
if is_gemini_image_model(model):
|
|
if supports_gemini_image_size(model):
|
|
image_config["imageSize"] = image_size
|
|
else:
|
|
image_config["imageSize"] = image_size
|
|
return image_config
|
|
|
|
|
|
def supports_gemini_image_size(model: str) -> bool:
|
|
try:
|
|
model_info = litellm.get_model_info(model=model)
|
|
value = model_info.get("supports_image_size")
|
|
if value is not None:
|
|
return bool(value)
|
|
except Exception:
|
|
pass
|
|
return "2.5-flash" not in model
|
|
|
|
|
|
def is_gemini_image_model(model: str) -> bool:
|
|
base_model = model.split("/", 1)[-1]
|
|
return "gemini" in base_model
|
|
|
|
|
|
def map_openai_image_params_to_gemini(
|
|
params: Dict[str, Any],
|
|
model: str,
|
|
supported_params: Sequence[str],
|
|
optional_params: Optional[Dict[str, Any]] = None,
|
|
parse_image_config_string: bool = False,
|
|
) -> Dict[str, Any]:
|
|
optional_params = optional_params or {}
|
|
filtered_params = {
|
|
key: value for key, value in params.items() if key in supported_params
|
|
}
|
|
|
|
mapped_params: Dict[str, Any] = {}
|
|
|
|
if "n" in filtered_params and "n" not in optional_params:
|
|
mapped_params["sampleCount"] = filtered_params["n"]
|
|
|
|
if "size" in filtered_params and "size" not in optional_params:
|
|
image_config = map_openai_size_to_gemini_image_config(
|
|
filtered_params["size"],
|
|
model,
|
|
)
|
|
if image_config is not None:
|
|
if is_gemini_image_model(model):
|
|
mapped_params["imageConfig"] = image_config
|
|
else:
|
|
mapped_params["aspectRatio"] = image_config["aspectRatio"]
|
|
if "imageSize" in image_config:
|
|
mapped_params["imageSize"] = image_config["imageSize"]
|
|
|
|
image_config_param = filtered_params.get("imageConfig")
|
|
if isinstance(image_config_param, str) and parse_image_config_string:
|
|
try:
|
|
image_config_param = json.loads(image_config_param)
|
|
except json.JSONDecodeError as exc:
|
|
raise litellm.UnsupportedParamsError(
|
|
model=model,
|
|
message="`imageConfig` must be valid JSON when provided as a string.",
|
|
) from exc
|
|
if isinstance(image_config_param, dict):
|
|
mapped_params["imageConfig"] = image_config_param
|
|
|
|
for key, value in filtered_params.items():
|
|
if (
|
|
key not in ("n", "size", "imageConfig", "tools", "web_search_options")
|
|
and key not in optional_params
|
|
):
|
|
mapped_params[key] = value
|
|
|
|
return mapped_params
|
|
|
|
|
|
def _dedupe_gemini_search_tools(tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
|
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
|
VertexGeminiConfig,
|
|
)
|
|
|
|
search_tool_keys = VertexGeminiConfig._search_tool_keys()
|
|
seen_search_keys: set[str] = set()
|
|
deduped_tools: List[Dict[str, Any]] = []
|
|
|
|
for tool in tools:
|
|
if not isinstance(tool, dict):
|
|
deduped_tools.append(tool)
|
|
continue
|
|
|
|
search_key = next((key for key in search_tool_keys if key in tool), None)
|
|
if search_key is None:
|
|
deduped_tools.append(tool)
|
|
continue
|
|
|
|
if search_key in seen_search_keys:
|
|
continue
|
|
|
|
seen_search_keys.add(search_key)
|
|
deduped_tools.append(tool)
|
|
|
|
return deduped_tools
|
|
|
|
|
|
def _has_gemini_search_tool(tools: List[Any]) -> bool:
|
|
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
|
VertexGeminiConfig,
|
|
)
|
|
|
|
search_tool_keys = VertexGeminiConfig._search_tool_keys()
|
|
return any(
|
|
isinstance(tool, dict) and any(key in tool for key in search_tool_keys)
|
|
for tool in tools
|
|
)
|
|
|
|
|
|
def map_gemini_image_tools_params(
|
|
non_default_params: Dict[str, Any],
|
|
mapped_params: Dict[str, Any],
|
|
) -> Dict[str, Any]:
|
|
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
|
VertexGeminiConfig,
|
|
)
|
|
|
|
gemini_config = VertexGeminiConfig()
|
|
result = dict(mapped_params)
|
|
result.pop("web_search_options", None)
|
|
|
|
tools_value = non_default_params.get("tools")
|
|
if isinstance(tools_value, list) and tools_value:
|
|
mapped_tools = gemini_config._map_function(
|
|
value=tools_value, optional_params=result
|
|
)
|
|
result = gemini_config._add_tools_to_optional_params(result, mapped_tools)
|
|
|
|
web_search_options = non_default_params.get("web_search_options")
|
|
existing_tools = result.get("tools")
|
|
if isinstance(web_search_options, dict) and not (
|
|
isinstance(existing_tools, list) and _has_gemini_search_tool(existing_tools)
|
|
):
|
|
search_tool = gemini_config._map_web_search_options(web_search_options)
|
|
result = gemini_config._add_tools_to_optional_params(result, [search_tool])
|
|
|
|
gemini_config._drop_search_tools_mixed_with_functions(result)
|
|
|
|
if isinstance(result.get("tools"), list):
|
|
result["tools"] = _dedupe_gemini_search_tools(result["tools"])
|
|
|
|
return result
|
|
|
|
|
|
def get_gemini_image_web_search_requests(
|
|
response_data: Dict[str, Any],
|
|
) -> Optional[int]:
|
|
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
|
VertexGeminiConfig,
|
|
)
|
|
|
|
grounding_metadata: List[Dict[str, Any]] = []
|
|
for candidate in response_data.get("candidates", []):
|
|
if not isinstance(candidate, dict):
|
|
continue
|
|
candidate_grounding = candidate.get("groundingMetadata")
|
|
if isinstance(candidate_grounding, list):
|
|
grounding_metadata.extend(candidate_grounding)
|
|
elif isinstance(candidate_grounding, dict):
|
|
grounding_metadata.append(candidate_grounding)
|
|
|
|
return VertexGeminiConfig._calculate_web_search_requests(grounding_metadata)
|
|
|
|
|
|
def get_gemini_image_generation_config(
|
|
model: str,
|
|
optional_params: Dict[str, Any],
|
|
) -> Dict[str, Any]:
|
|
generation_config: Dict[str, Any] = {"response_modalities": ["IMAGE", "TEXT"]}
|
|
|
|
image_config: Dict[str, Any] = {}
|
|
if isinstance(optional_params.get("imageConfig"), dict):
|
|
image_config.update(optional_params["imageConfig"])
|
|
|
|
if not supports_gemini_image_size(model):
|
|
image_config.pop("imageSize", None)
|
|
|
|
if image_config:
|
|
generation_config["imageConfig"] = image_config
|
|
|
|
candidate_count = next(
|
|
(
|
|
optional_params[key]
|
|
for key in ("candidateCount", "candidate_count", "sampleCount", "n")
|
|
if optional_params.get(key) is not None
|
|
),
|
|
None,
|
|
)
|
|
if candidate_count is not None:
|
|
generation_config["candidateCount"] = candidate_count
|
|
|
|
return generation_config
|
|
|
|
|
|
def _parse_openai_image_size(size: str) -> Optional[tuple[int, int]]:
|
|
if size == "auto":
|
|
return None
|
|
|
|
width_str, separator, height_str = size.lower().partition("x")
|
|
if not separator:
|
|
return None
|
|
|
|
try:
|
|
width = int(width_str)
|
|
height = int(height_str)
|
|
except ValueError:
|
|
return None
|
|
|
|
if width <= 0 or height <= 0:
|
|
return None
|
|
|
|
return width, height
|
|
|
|
|
|
def _map_dimensions_to_gemini_aspect_ratio(width: int, height: int) -> str:
|
|
if (width, height) in GEMINI_IMAGE_SIZE_TO_ASPECT_RATIO:
|
|
return GEMINI_IMAGE_SIZE_TO_ASPECT_RATIO[(width, height)]
|
|
|
|
requested_ratio = width / height
|
|
return min(
|
|
GEMINI_IMAGE_ASPECT_RATIOS,
|
|
key=lambda aspect_ratio: abs(
|
|
math.log(GEMINI_IMAGE_ASPECT_RATIOS[aspect_ratio] / requested_ratio)
|
|
),
|
|
)
|
|
|
|
|
|
def _map_dimensions_to_gemini_image_size(width: int, height: int) -> str:
|
|
effective_square_side = math.sqrt(width * height)
|
|
if effective_square_side < 768:
|
|
return "512"
|
|
if effective_square_side < 1536:
|
|
return "1K"
|
|
if effective_square_side < 3072:
|
|
return "2K"
|
|
return "4K"
|
|
|
|
|
|
class GeminiError(BaseLLMException):
|
|
pass
|
|
|
|
|
|
class GeminiModelInfo(BaseLLMModelInfo):
|
|
def validate_environment(
|
|
self,
|
|
headers: dict,
|
|
model: str,
|
|
messages: List[AllMessageValues],
|
|
optional_params: dict,
|
|
litellm_params: dict,
|
|
api_key: Optional[str] = None,
|
|
api_base: Optional[str] = None,
|
|
) -> dict:
|
|
"""Google AI Studio sends api key via x-goog-api-key header"""
|
|
return headers
|
|
|
|
@property
|
|
def api_version(self) -> str:
|
|
return "v1beta"
|
|
|
|
@staticmethod
|
|
def get_api_base(api_base: Optional[str] = None) -> Optional[str]:
|
|
return (
|
|
api_base
|
|
or get_secret_str("GEMINI_API_BASE")
|
|
or "https://generativelanguage.googleapis.com"
|
|
)
|
|
|
|
@staticmethod
|
|
def get_api_key(api_key: Optional[str] = None) -> Optional[str]:
|
|
return (
|
|
api_key
|
|
or (get_secret_str("GOOGLE_API_KEY"))
|
|
or (get_secret_str("GEMINI_API_KEY"))
|
|
)
|
|
|
|
@staticmethod
|
|
def get_base_model(model: str) -> Optional[str]:
|
|
return model.replace("gemini/", "")
|
|
|
|
def process_model_name(self, models: List[Dict[str, str]]) -> List[str]:
|
|
litellm_model_names = []
|
|
for model in models:
|
|
stripped_model_name = model["name"].replace("models/", "")
|
|
litellm_model_name = "gemini/" + stripped_model_name
|
|
litellm_model_names.append(litellm_model_name)
|
|
return litellm_model_names
|
|
|
|
def get_models(
|
|
self, api_key: Optional[str] = None, api_base: Optional[str] = None
|
|
) -> List[str]:
|
|
api_base = GeminiModelInfo.get_api_base(api_base)
|
|
api_key = GeminiModelInfo.get_api_key(api_key)
|
|
endpoint = f"/{self.api_version}/models"
|
|
if api_base is None or api_key is None:
|
|
raise ValueError(
|
|
"GEMINI_API_BASE or GEMINI_API_KEY/GOOGLE_API_KEY is not set. Please set the environment variable, to query Gemini's `/models` endpoint."
|
|
)
|
|
|
|
response = litellm.module_level_client.get(
|
|
url=f"{api_base}{endpoint}",
|
|
headers={"x-goog-api-key": api_key},
|
|
)
|
|
|
|
if response.status_code != 200:
|
|
raise ValueError(
|
|
f"Failed to fetch models from Gemini. Status code: {response.status_code}, Response: {response.json()}"
|
|
)
|
|
|
|
models = response.json()["models"]
|
|
|
|
litellm_model_names = self.process_model_name(models)
|
|
return litellm_model_names
|
|
|
|
def get_error_class(
|
|
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
|
|
) -> BaseLLMException:
|
|
return GeminiError(
|
|
status_code=status_code, message=error_message, headers=headers
|
|
)
|
|
|
|
def get_token_counter(self) -> Optional[BaseTokenCounter]:
|
|
"""
|
|
Factory method to create a token counter for this provider.
|
|
|
|
Returns:
|
|
Optional TokenCounterInterface implementation for this provider,
|
|
or None if token counting is not supported.
|
|
"""
|
|
return GoogleAIStudioTokenCounter()
|
|
|
|
|
|
def encode_unserializable_types(
|
|
data: Dict[str, object], depth: int = 0
|
|
) -> Dict[str, object]:
|
|
"""Converts unserializable types in dict to json.dumps() compatible types.
|
|
|
|
This function is called in models.py after calling convert_to_dict(). The
|
|
convert_to_dict() can convert pydantic object to dict. However, the input to
|
|
convert_to_dict() is dict mixed of pydantic object and nested dict(the output
|
|
of converters). So they may be bytes in the dict and they are out of
|
|
`ser_json_bytes` control in model_dump(mode='json') called in
|
|
`convert_to_dict`, as well as datetime deserialization in Pydantic json mode.
|
|
|
|
Returns:
|
|
A dictionary with json.dumps() incompatible type (e.g. bytes datetime)
|
|
to compatible type (e.g. base64 encoded string, isoformat date string).
|
|
"""
|
|
if depth > DEFAULT_MAX_RECURSE_DEPTH:
|
|
return data
|
|
processed_data: dict[str, object] = {}
|
|
if not isinstance(data, dict):
|
|
return data
|
|
for key, value in data.items():
|
|
if isinstance(value, bytes):
|
|
processed_data[key] = base64.urlsafe_b64encode(value).decode("ascii")
|
|
elif isinstance(value, datetime.datetime):
|
|
processed_data[key] = value.isoformat()
|
|
elif isinstance(value, dict):
|
|
processed_data[key] = encode_unserializable_types(value, depth + 1)
|
|
elif isinstance(value, list):
|
|
if all(isinstance(v, bytes) for v in value):
|
|
processed_data[key] = [
|
|
base64.urlsafe_b64encode(v).decode("ascii") for v in value
|
|
]
|
|
if all(isinstance(v, datetime.datetime) for v in value):
|
|
processed_data[key] = [v.isoformat() for v in value]
|
|
else:
|
|
processed_data[key] = [
|
|
encode_unserializable_types(v, depth + 1) for v in value
|
|
]
|
|
else:
|
|
processed_data[key] = value
|
|
return processed_data
|
|
|
|
|
|
def get_api_key_from_env() -> Optional[str]:
|
|
return get_secret_str("GOOGLE_API_KEY") or get_secret_str("GEMINI_API_KEY")
|
|
|
|
|
|
class GoogleAIStudioTokenCounter(BaseTokenCounter):
|
|
"""Token counter implementation for Google AI Studio provider."""
|
|
|
|
def should_use_token_counting_api(
|
|
self,
|
|
custom_llm_provider: Optional[str] = None,
|
|
) -> bool:
|
|
from litellm.types.utils import LlmProviders
|
|
|
|
return custom_llm_provider == LlmProviders.GEMINI.value
|
|
|
|
async def count_tokens(
|
|
self,
|
|
model_to_use: str,
|
|
messages: Optional[List[Dict[str, Any]]],
|
|
contents: Optional[List[Dict[str, Any]]],
|
|
deployment: Optional[Dict[str, Any]] = None,
|
|
request_model: str = "",
|
|
tools: Optional[List[Dict[str, Any]]] = None,
|
|
system: Optional[Any] = None,
|
|
) -> Optional[TokenCountResponse]:
|
|
import copy
|
|
|
|
from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter
|
|
|
|
deployment = deployment or {}
|
|
count_tokens_params_request = copy.deepcopy(
|
|
deployment.get("litellm_params", {})
|
|
)
|
|
count_tokens_params = {
|
|
"model": model_to_use,
|
|
"contents": contents,
|
|
}
|
|
count_tokens_params_request.update(count_tokens_params)
|
|
result = await GoogleAIStudioTokenCounter().acount_tokens(
|
|
**count_tokens_params_request,
|
|
)
|
|
|
|
if result is not None:
|
|
return TokenCountResponse(
|
|
total_tokens=result.get("totalTokens", 0),
|
|
request_model=request_model,
|
|
model_used=model_to_use,
|
|
tokenizer_type=result.get("tokenizer_used", ""),
|
|
original_response=result,
|
|
)
|
|
|
|
return None
|