Compare commits

...

19 Commits

Author SHA1 Message Date
snaily
8ca62707ea feat: 添加搜索模型配置并改进Markdown链接处理
在Dockerfile中添加SEARCH_MODELS环境变量,支持gemini-2.0-flash-exp和gemini-2.0-pro-exp模型
改进message_converter中的图片链接正则表达式
2025-03-19 19:56:50 +08:00
Toddy
21444ed6c7 chore: 统一从model_service读取模型列表 2025-03-18 18:05:00 +00:00
Toddy
ba292dbedd chore: 规范变量名 2025-03-18 17:54:18 +00:00
snaily
6ba58ce9d1 fix: 重构图片MIME类型转换逻辑
将"image/jpg"到"image/jpeg"的MIME类型转换逻辑从_convert_image函数移至_get_mime_type_and_data函数,避免代码重复并提高一致性。这确保了MIME类型的标准化处理发生在数据提取的同一位置。
2025-03-18 21:50:27 +08:00
snaily
16f16a3ae9 Merge branch 'pr/yangtb2024/13' 2025-03-18 21:46:34 +08:00
snaily
26dcb64687 fix: 将image/jpg MIME类型转换为标准的image/jpeg
修复了图像转换过程中的MIME类型处理,确保当遇到非标准的"image/jpg"类型时,将其转换为标准的"image/jpeg"类型。这样可以提高与接收图像数据的API和系统的兼容性
2025-03-18 21:35:19 +08:00
yangtb2024
df88492113 将chat-bison-001、text-bison-001和embedding-gecko-001添加到FILTERED_MODELS列表 2025-03-18 15:21:29 +08:00
yangtb2024
851bb9c09b 将 filtered_models 从硬编码改为可配置参数
1. 在 config.py 中添加 FILTERED_MODELS 配置项
2. 在 .env.example 中添加 FILTERED_MODELS 示例
3. 修改 model_service.py 以使用配置的过滤模型列表
4. 优化模型过滤逻辑
2025-03-18 14:47:58 +08:00
yangtb2024
0cac178572 Merge branch 'snailyp:main' into model 2025-03-18 12:44:09 +08:00
snaily
67c85c994a Merge pull request #14 from cr-zhichen/main
fix: 更新Cloudflare ImgBed上传请求URL,新增uploadNameType参数,以保持正确的目录结构命名。
2025-03-17 15:24:39 +08:00
cr-zhichen
ee979dd568 Merge branch 'main' of https://github.com/cr-zhichen/gemini-balance 2025-03-17 07:12:43 +00:00
cr-zhichen
e79a1ba56c feat: 更新CloudFlare ImgBed上传请求URL,新增uploadNameType参数,以保持正确的日期命名目录结构。 2025-03-17 07:10:21 +00:00
yangtb2024
016e6e06ee Filter out vision-based Gemini models from model list 2025-03-17 13:56:01 +08:00
snaily
8779a5f0b3 feat: 添加对 image-generation 模型的支持
在 gemini_chat_service 和 openai_chat_service 中添加对 "-image-generation" 后缀模型的支持
确保 image-generation 模型与 image 模型有相同的处理逻辑
2025-03-16 23:53:53 +08:00
cr-zhichen
89f2825ac7 feat: 新增对CloudFlare ImgBed的支持,更新环境变量和文档 2025-03-16 04:39:40 +00:00
snaily
985a12554d fix:修改OpenAI消息转换器中assistant消息处理逻辑,将特殊处理的目标从最后一条消息调整为倒数第二条消息。 2025-03-15 21:18:20 +08:00
snaily
37a7a140fc feat:改进消息转换器中的图像处理和消息分割逻辑
添加 _get_mime_type_and_data 函数从 base64 字符串中提取 MIME 类型和数据
修改 _convert_image 函数使用动态检测的 MIME 类型,而非硬编码
将 _process_text_with_image 中的 MIME 类型从 image/jpeg 改为 image/png
简化异常处理逻辑
优化 OpenAIMessageConverter 中的消息分割逻辑,仅对最后一个 assistant 消息进行分割处理
2025-03-15 21:11:10 +08:00
zhanghaoyu
28e67cc3fa 1. modify IMAGE_URL_PATTERN
2. modify import
2025-03-15 12:37:56 +08:00
zhanghaoyu7
d99a0bde93 feat: 新增图文上下文同步 2025-03-14 16:29:03 +08:00
13 changed files with 315 additions and 40 deletions

View File

@@ -1,8 +1,10 @@
API_KEYS=["AIzaSyxxxxxxxxxxxxxxxxxxx","AIzaSyxxxxxxxxxxxxxxxxxxx"]
ALLOWED_TOKENS=["sk-123456"]
# AUTH_TOKEN=sk-123456
MODEL_SEARCH=["gemini-2.0-flash-exp","gemini-2.0-pro-exp"]
MODEL_IMAGE=["gemini-2.0-flash-exp"]
TEST_MODEL=gemini-1.5-flash
IMAGE_MODELS=["gemini-2.0-flash-exp"]
SEARCH_MODELS=["gemini-2.0-flash-exp","gemini-2.0-pro-exp"]
FILTERED_MODELS=["gemini-1.0-pro-vision-latest", "gemini-pro-vision", "chat-bison-001", "text-bison-001", "embedding-gecko-001"]
TOOLS_CODE_EXECUTION_ENABLED=false
SHOW_SEARCH_LINK=true
SHOW_THINKING_PROCESS=true
@@ -14,6 +16,8 @@ CREATE_IMAGE_MODEL=imagen-3.0-generate-002
UPLOAD_PROVIDER=smms
SMMS_SECRET_TOKEN=XXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
PICGO_API_KEY=xxxx
CLOUDFLARE_IMGBED_URL=https://xxxxxxx.pages.dev/upload
CLOUDFLARE_IMGBED_AUTH_CODE=xxxxxxxxx
##########################################################################
#########################stream_optimizer 相关配置########################
STREAM_MIN_DELAY=0.016

View File

@@ -11,7 +11,8 @@ ENV API_KEYS='["your_api_key_1"]'
ENV ALLOWED_TOKENS='["your_token_1"]'
ENV BASE_URL=https://generativelanguage.googleapis.com/v1beta
ENV TOOLS_CODE_EXECUTION_ENABLED=false
ENV MODEL_SEARCH='["gemini-2.0-flash-exp"]'
ENV IMAGE_MODELS='["gemini-2.0-flash-exp"]'
ENV SEARCH_MODELS='["gemini-2.0-flash-exp","gemini-2.0-pro-exp"]'
# Expose port
EXPOSE 8000

View File

@@ -64,18 +64,24 @@
AUTH_TOKEN="" # 超级管理员token具有所有权限默认使用 ALLOWED_TOKENS 的第一个
# 模型功能配置
MODEL_SEARCH=["gemini-2.0-flash-exp"] # 支持搜索功能的模型列表
TEST_MODEL="gemini-1.5-flash" # 用于测试密钥是否可用的模型
SEARCH_MODELS=["gemini-2.0-flash-exp"] # 支持搜索功能的模型列表
IMAGE_MODELS=["gemini-2.0-flash-exp"] # 支持绘图功能的模型列表
TOOLS_CODE_EXECUTION_ENABLED=false # 是否启用代码执行工具默认false
SHOW_SEARCH_LINK=true # 是否在响应中显示搜索结果链接默认true
SHOW_THINKING_PROCESS=true # 是否显示模型思考过程默认true
FILTERED_MODELS=["gemini-1.0-pro-vision-latest", "gemini-pro-vision", "chat-bison-001", "text-bison-001", "embedding-gecko-001"] # 被禁用的模型列表
# 图片生成配置
PAID_KEY="your-paid-api-key" # 付费版API Key用于图片生成等高级功能
CREATE_IMAGE_MODEL="imagen-3.0-generate-002" # 图片生成模型默认使用imagen-3.0
# 图片上传配置
UPLOAD_PROVIDER="smms" # 图片上传提供商目前支持smms
UPLOAD_PROVIDER="smms" # 图片上传提供商目前支持smms、picgo、cloudflare_imgbed
SMMS_SECRET_TOKEN="your-smms-token" # SM.MS图床的API Token
PICGO_API_KEY="your-picogo-apikey" # PicoGo图床的API Key 可在 `https://www.picgo.net/settings/api` 获取
CLOUDFLARE_IMGBED_URL="https://xxxxxxx.pages.dev/upload" # CloudFlare 图床上传地址,可自行搭建:`https://github.com/MarSeventh/CloudFlare-ImgBed`
CLOUDFLARE_IMGBED_AUTH_CODE="your-cloudflare-imgber-auth-code" # CloudFlare图床的鉴权key可在项目后台设置若无鉴权则可直接置空。
# stream_optimizer 相关配置
STREAM_MIN_DELAY=0.016
@@ -112,9 +118,17 @@
#### 模型功能配置
- `MODEL_SEARCH`: 搜索功能支持的模型
- `TEST_MODEL`: 用于测试密钥可用性的模型
- 默认值: `"gemini-1.5-flash"`
- `SEARCH_MODELS`: 搜索功能支持的模型
- 默认值: `["gemini-2.0-flash-exp"]`
- 说明: 仅列表中的模型可使用搜索功能
- `IMAGE_MODELS`: 绘图功能支持的模型
- 默认值: `["gemini-2.0-flash-exp"]`
- 说明: 仅列表中的模型可使用绘图功能
- `FILTERED_MODELS`: 被禁用的模型列表
- 默认值: `["gemini-1.0-pro-vision-latest", "gemini-pro-vision", "chat-bison-001", "text-bison-001", "embedding-gecko-001"]`
- 说明: 列表中的模型将被禁用
- `TOOLS_CODE_EXECUTION_ENABLED`: 代码执行功能
- 默认值: `false`
- 安全提示: 生产环境建议禁用
@@ -138,10 +152,26 @@
- `UPLOAD_PROVIDER`: 图片上传服务提供商
- 默认值: `smms`
- 说明: 目前支持 SM.MS 图床
- 可选值: `smms`, `picgo`, `cloudflare_imgbed`
- 说明: 用于选择图片上传的服务提供商。目前支持 SM.MS 图床, PicGo 图床, 以及 Cloudflare ImgBed。
- `SMMS_SECRET_TOKEN`: SM.MS API Token
- 用途: 用于图片上传到 SM.MS 图床
- 获取方式: 需要在 SM.MS 官网注册并获取
- 用途: 用于图片上传到 SM.MS 图床的身份验证。
- 获取方式: 需要在 [SM.MS 官网](https://sm.ms/) 注册并获取
- `PICGO_API_KEY`: PicGo API Key
- 用途: 用于图片上传到 PicGo 图床的身份验证。
- 获取方式: 可在 [PicGo 官网](https://www.picgo.net/settings/api) 的设置页面 API 选项中获取。
- `CLOUDFLARE_IMGBED_URL`: Cloudflare ImgBed 上传地址
- 用途: 指定 Cloudflare ImgBed 图床的上传 API 地址。
- 获取方式: 如果您自行搭建了 Cloudflare ImgBed 服务,请填写您的服务部署地址。参考 [Cloudflare-ImgBed 项目](https://github.com/MarSeventh/CloudFlare-ImgBed) 自行搭建。
- 注意: URL 必须以 `https://` 开头,并指向 `/upload` 路径 ,例如 `https://cloudflare-imgbed-7b0.pages.dev/upload`。
- `CLOUDFLARE_IMGBED_AUTH_CODE`: Cloudflare ImgBed 鉴权 Key
- 用途: 用于 Cloudflare ImgBed 图床的身份验证。
- 说明: 如果您的 Cloudflare ImgBed 服务启用了鉴权,请填写鉴权 Key。若未启用鉴权则留空即可。
- 获取方式: 在 Cloudflare ImgBed 项目的后台设置中获取,或在搭建时自行设置。
#### 流式输出优化配置

View File

@@ -23,7 +23,7 @@ async def get_key_manager():
async def get_next_working_key_wrapper(key_manager: KeyManager = Depends(get_key_manager)):
return await key_manager.get_next_working_key()
model_service = ModelService(settings.MODEL_SEARCH,settings.MODEL_IMAGE)
model_service = ModelService(settings.SEARCH_MODELS,settings.IMAGE_MODELS)
@router.get("/models")
@@ -41,8 +41,8 @@ async def list_models(_=Depends(security_service.verify_key),
model_mapping = {x.get("name", "").split("/", maxsplit=1)[1]: x for x in models_json["models"]}
# 添加搜索模型
if settings.MODEL_SEARCH:
for name in settings.MODEL_SEARCH:
if model_service.search_models:
for name in model_service.search_models:
model = model_mapping.get(name, None)
if not model:
continue
@@ -56,8 +56,8 @@ async def list_models(_=Depends(security_service.verify_key),
models_json["models"].append(item)
# 添加图像生成模型
if settings.MODEL_IMAGE:
for name in settings.MODEL_IMAGE:
if model_service.image_models:
for name in model_service.image_models:
model = model_mapping.get(name, None)
if not model:
continue

View File

@@ -17,7 +17,7 @@ logger = get_openai_logger()
# 初始化服务
security_service = SecurityService(settings.ALLOWED_TOKENS, settings.AUTH_TOKEN)
model_service = ModelService(settings.MODEL_SEARCH,settings.MODEL_IMAGE)
model_service = ModelService(settings.SEARCH_MODELS,settings.IMAGE_MODELS)
embedding_service = EmbeddingService(settings.BASE_URL)
image_create_service = ImageCreateService()

View File

@@ -6,8 +6,9 @@ class Settings(BaseSettings):
API_KEYS: List[str]
ALLOWED_TOKENS: List[str]
BASE_URL: str = "https://generativelanguage.googleapis.com/v1beta"
MODEL_SEARCH: List[str] = ["gemini-2.0-flash-exp"]
MODEL_IMAGE: List[str] = ["gemini-2.0-flash-exp"]
SEARCH_MODELS: List[str] = ["gemini-2.0-flash-exp"]
IMAGE_MODELS: List[str] = ["gemini-2.0-flash-exp"]
FILTERED_MODELS: List[str] = ["gemini-1.0-pro-vision-latest", "gemini-pro-vision", "chat-bison-001", "text-bison-001", "embedding-gecko-001"]
TOOLS_CODE_EXECUTION_ENABLED: bool = False
SHOW_SEARCH_LINK: bool = True
SHOW_THINKING_PROCESS: bool = True
@@ -18,6 +19,8 @@ class Settings(BaseSettings):
UPLOAD_PROVIDER: str = "smms"
SMMS_SECRET_TOKEN: str = ""
PICGO_API_KEY: str = ""
CLOUDFLARE_IMGBED_URL: str = ""
CLOUDFLARE_IMGBED_AUTH_CODE: str = ""
TEST_MODEL: str = "gemini-1.5-flash"
# 流式输出优化器配置

View File

@@ -258,6 +258,119 @@ class PicGoUploader(ImageUploader):
original_error=e
)
class CloudFlareImgBedUploader(ImageUploader):
"""CloudFlare图床上传器"""
def __init__(self, auth_code: str, api_url: str):
"""
初始化CloudFlare图床上传器
Args:
auth_code: 认证码
api_url: 上传API地址
"""
self.auth_code = auth_code
self.api_url = api_url
def upload(self, file: bytes, filename: str) -> UploadResponse:
"""
上传图片到CloudFlare图床
Args:
file: 图片文件二进制数据
filename: 文件名
Returns:
UploadResponse: 上传响应对象
Raises:
UploadError: 上传失败时抛出异常
"""
try:
# 准备请求URL添加认证码参数如果存在
if self.auth_code:
request_url = f"{self.api_url}?authCode={self.auth_code}&uploadNameType=origin"
else:
request_url = f"{self.api_url}?uploadNameType=origin"
# 准备文件数据
files = {
"file": (filename, file)
}
# 发送请求
response = requests.post(
request_url,
files=files
)
# 检查响应状态
response.raise_for_status()
# 解析响应
result = response.json()
# 验证响应格式
if not result or not isinstance(result, list) or len(result) == 0:
raise UploadError(
message="Invalid response format",
error_type=UploadErrorType.PARSE_ERROR
)
# 获取文件URL
file_path = result[0].get("src")
if not file_path:
raise UploadError(
message="Missing file URL in response",
error_type=UploadErrorType.PARSE_ERROR
)
# 构建完整URL如果返回的是相对路径
base_url = self.api_url.split("/upload")[0]
full_url = file_path if file_path.startswith(("http://", "https://")) else f"{base_url}{file_path}"
# 构建图片元数据注意CloudFlare-ImgBed不返回所有元数据所以部分字段为默认值
image_metadata = ImageMetadata(
width=0, # CloudFlare-ImgBed不返回宽度
height=0, # CloudFlare-ImgBed不返回高度
filename=filename,
size=0, # CloudFlare-ImgBed不返回大小
url=full_url,
delete_url=None # CloudFlare-ImgBed不返回删除URL
)
return UploadResponse(
success=True,
code="success",
message="Upload success",
data=image_metadata
)
except requests.RequestException as e:
# 处理网络请求相关错误
raise UploadError(
message=f"Upload request failed: {str(e)}",
error_type=UploadErrorType.NETWORK_ERROR,
original_error=e
)
except (KeyError, ValueError, TypeError, IndexError) as e:
# 处理响应解析错误
raise UploadError(
message=f"Invalid response format: {str(e)}",
error_type=UploadErrorType.PARSE_ERROR,
original_error=e
)
except UploadError:
# 重新抛出已经是 UploadError 类型的异常
raise
except Exception as e:
# 处理其他未预期的错误
raise UploadError(
message=f"Upload failed: {str(e)}",
error_type=UploadErrorType.UNKNOWN,
original_error=e
)
class ImageUploaderFactory:
@staticmethod
@@ -272,4 +385,9 @@ class ImageUploaderFactory:
elif provider == "picgo":
api_url = credentials.get("api_url", "https://www.picgo.net/api/1/upload")
return PicGoUploader(credentials["api_key"], api_url)
elif provider == "cloudflare_imgbed":
return CloudFlareImgBedUploader(
credentials["auth_code"],
credentials["base_url"]
)
raise ValueError(f"Unknown provider: {provider}")

View File

@@ -1,9 +1,13 @@
# app/services/chat/message_converter.py
from abc import ABC, abstractmethod
import re
from typing import Any, Dict, List, Optional
import requests
import base64
SUPPORTED_ROLES = ["user", "model", "system"]
IMAGE_URL_PATTERN = r'\[(.*?)\]\((.*?)\)'
class MessageConverter(ABC):
@@ -13,13 +17,36 @@ class MessageConverter(ABC):
def convert(self, messages: List[Dict[str, Any]]) -> tuple[List[Dict[str, Any]], Optional[Dict[str, Any]]]:
pass
def _get_mime_type_and_data(base64_string):
"""
从 base64 字符串中提取 MIME 类型和数据。
参数:
base64_string (str): 可能包含 MIME 类型信息的 base64 字符串
返回:
tuple: (mime_type, encoded_data)
"""
# 检查字符串是否以 "data:" 格式开始
if base64_string.startswith('data:'):
# 提取 MIME 类型和数据
pattern = r'data:([^;]+);base64,(.+)'
match = re.match(pattern, base64_string)
if match:
mime_type = "image/jpeg" if match.group(1) == "image/jpg" else match.group(1)
encoded_data = match.group(2)
return mime_type, encoded_data
# 如果不是预期格式,假定它只是数据部分
return None, base64_string
def _convert_image(image_url: str) -> Dict[str, Any]:
if image_url.startswith("data:image"):
mime_type, encoded_data = _get_mime_type_and_data(image_url)
return {
"inline_data": {
"mime_type": "image/jpeg",
"data": image_url.split(",")[1]
"mime_type": mime_type,
"data": encoded_data
}
}
return {
@@ -29,6 +56,56 @@ def _convert_image(image_url: str) -> Dict[str, Any]:
}
def _convert_image_to_base64(url: str) -> str:
"""
将图片URL转换为base64编码
Args:
url: 图片URL
Returns:
str: base64编码的图片数据
"""
response = requests.get(url)
if response.status_code == 200:
# 将图片内容转换为base64
img_data = base64.b64encode(response.content).decode('utf-8')
return img_data
else:
raise Exception(f"Failed to fetch image: {response.status_code}")
def _process_text_with_image(text: str) -> List[Dict[str, Any]]:
"""
处理可能包含图片URL的文本提取图片并转换为base64
Args:
text: 可能包含图片URL的文本
Returns:
List[Dict[str, Any]]: 包含文本和图片的部分列表
"""
parts = []
img_url_match = re.search(IMAGE_URL_PATTERN, text)
if img_url_match:
# 提取URL
img_url = img_url_match.group(2)
# 将URL对应的图片转换为base64
try:
base64_data = _convert_image_to_base64(img_url)
parts.append({
"inlineData": {
"mimeType": "image/png",
"data": base64_data
}
})
except Exception:
# 如果转换失败,回退到文本模式
parts.append({"text": text})
else:
# 没有图片URL作为纯文本处理
parts.append({"text": text})
return parts
class OpenAIMessageConverter(MessageConverter):
"""OpenAI消息格式转换器"""
@@ -49,9 +126,18 @@ class OpenAIMessageConverter(MessageConverter):
role = "model"
parts = []
if isinstance(msg["content"], str) and msg["content"]:
# 特别处理最后一个assistant的消息按\n\n分割
if role == "assistant" and idx == len(messages) - 2 and isinstance(msg["content"], str) and msg["content"]:
# 按\n\n分割消息
content_parts = msg["content"].split("\n\n")
for part in content_parts:
if not part.strip(): # 跳过空内容
continue
# 处理可能包含图片的文本
parts.extend(_process_text_with_image(part))
elif isinstance(msg["content"], str) and msg["content"]:
# 请求 gemini 接口时如果包含 content 字段但内容为空时会返回 400 错误,所以需要判断是否为空并移除
parts.append({"text": msg["content"]})
parts.extend(_process_text_with_image(msg["content"]))
elif isinstance(msg["content"], list):
for content in msg["content"]:
if isinstance(content, str) and content:
@@ -76,4 +162,4 @@ class OpenAIMessageConverter(MessageConverter):
"parts": system_instruction_parts,
}
)
return converted_messages, system_instruction
return converted_messages, system_instruction

View File

@@ -200,14 +200,16 @@ def _extract_image_data(part: dict) -> str:
image_uploader = ImageUploaderFactory.create(provider=settings.UPLOAD_PROVIDER,api_key=settings.SMMS_SECRET_TOKEN)
elif settings.UPLOAD_PROVIDER == "picgo":
image_uploader = ImageUploaderFactory.create(provider=settings.UPLOAD_PROVIDER,api_key=settings.PICGO_API_KEY)
elif settings.UPLOAD_PROVIDER == "cloudflare_imgbed":
image_uploader = ImageUploaderFactory.create(provider=settings.UPLOAD_PROVIDER,base_url=settings.CLOUDFLARE_IMGBED_URL,auth_code=settings.CLOUDFLARE_IMGBED_AUTH_CODE)
current_date = time.strftime("%Y/%m/%d")
filename = f"{current_date}/{uuid.uuid4().hex[:8]}.png"
base64_data = part["inlineData"]["data"]
#将base64_data转成bytes数组
#将base64_data转成bytes数组
bytes_data = base64.b64decode(base64_data)
upload_response = image_uploader.upload(bytes_data,filename)
if upload_response.success:
text = f"\n![image]({upload_response.data.url})\n"
text = f"![image]({upload_response.data.url})"
else:
text = ""
return text

View File

@@ -71,7 +71,7 @@ def _build_payload(model: str, request: GeminiRequest) -> Dict[str, Any]:
"systemInstruction": request_dict.get("systemInstruction", "")
}
if model.endswith("-image"):
if model.endswith("-image") or model.endswith("-image-generation"):
payload.pop("systemInstruction")
payload["generationConfig"]["responseModalities"] = ["Text","Image"]
return payload

View File

@@ -96,11 +96,6 @@ class ImageCreateService:
for index, generated_image in enumerate(response.generated_images):
image_data = generated_image.image.image_bytes
image_uploader = None
if settings.UPLOAD_PROVIDER == "smms":
image_uploader = ImageUploaderFactory.create(provider=settings.UPLOAD_PROVIDER,api_key=settings.SMMS_SECRET_TOKEN)
current_date = time.strftime("%Y/%m/%d")
filename = f"{current_date}/{uuid.uuid4().hex[:8]}.png"
upload_response = image_uploader.upload(image_data,filename)
if request.response_format == "b64_json":
base64_image = base64.b64encode(image_data).decode('utf-8')
@@ -109,6 +104,30 @@ class ImageCreateService:
"revised_prompt": request.prompt
})
else:
current_date = time.strftime("%Y/%m/%d")
filename = f"{current_date}/{uuid.uuid4().hex[:8]}.png"
if settings.UPLOAD_PROVIDER == "smms":
image_uploader = ImageUploaderFactory.create(
provider=settings.UPLOAD_PROVIDER,
api_key=settings.SMMS_SECRET_TOKEN
)
elif settings.UPLOAD_PROVIDER == "picgo":
image_uploader = ImageUploaderFactory.create(
provider=settings.UPLOAD_PROVIDER,
api_key=settings.PICGO_API_KEY
)
elif settings.UPLOAD_PROVIDER == "cloudflare_imgbed":
image_uploader = ImageUploaderFactory.create(
provider=settings.UPLOAD_PROVIDER,
base_url=settings.CLOUDFLARE_IMGBED_URL,
auth_code=settings.CLOUDFLARE_IMGBED_AUTH_CODE
)
else:
raise ValueError(f"Unsupported upload provider: {settings.UPLOAD_PROVIDER}")
upload_response = image_uploader.upload(image_data, filename)
images_data.append({
"url": f"{upload_response.data.url}",
"revised_prompt": request.prompt

View File

@@ -7,10 +7,11 @@ from app.core.config import settings
logger = get_model_logger()
class ModelService:
def __init__(self, model_search: list, model_image: list):
self.model_search = model_search
self.model_image = model_image
def __init__(self, search_models: list, image_models: list):
self.search_models = search_models
self.image_models = image_models
self.base_url = "https://generativelanguage.googleapis.com/v1beta"
self.filtered_models = settings.FILTERED_MODELS
def get_gemini_models(self, api_key: str) -> Optional[Dict[str, Any]]:
url = f"{self.base_url}/models?key={api_key}"
@@ -19,6 +20,16 @@ class ModelService:
response = requests.get(url)
if response.status_code == 200:
gemini_models = response.json()
filtered_models_list = []
for model in gemini_models.get("models", []):
model_id = model["name"].split("/")[-1]
if model_id not in self.filtered_models:
filtered_models_list.append(model)
else:
logger.info(f"Filtered out model: {model_id}")
gemini_models["models"] = filtered_models_list
return gemini_models
else:
logger.error(f"Error: {response.status_code}")
@@ -54,11 +65,11 @@ class ModelService:
}
openai_format["data"].append(openai_model)
if model_id in self.model_search:
if model_id in self.search_models:
search_model = openai_model.copy()
search_model["id"] = f"{model_id}-search"
openai_format["data"].append(search_model)
if model_id in self.model_image:
if model_id in self.image_models:
image_model = openai_model.copy()
image_model["id"] = f"{model_id}-image"
openai_format["data"].append(image_model)
@@ -76,9 +87,9 @@ class ModelService:
model = model.strip()
if model.endswith("-search"):
model = model[:-7]
return model in settings.MODEL_SEARCH
return model in self.search_models
if model.endswith("-image"):
model = model[:-6]
return model in settings.MODEL_IMAGE
return model in self.image_models
return True
return model not in self.filtered_models

View File

@@ -35,7 +35,7 @@ def _build_tools(
if (
settings.TOOLS_CODE_EXECUTION_ENABLED
and not (model.endswith("-search") or "-thinking" in model or model.endswith("-image"))
and not (model.endswith("-search") or "-thinking" in model or model.endswith("-image") or model.endswith("-image-generation"))
and not _has_image_parts(messages)
):
tools.append({"code_execution": {}})
@@ -110,7 +110,7 @@ def _build_payload(
"tools": _build_tools(request, messages),
"safetySettings": _get_safety_settings(request.model),
}
if request.model.endswith("-image"):
if request.model.endswith("-image") or request.model.endswith("-image-generation"):
payload["generationConfig"]["responseModalities"] = ["Text","Image"]
if (
@@ -119,6 +119,7 @@ def _build_payload(
and instruction.get("role") == "system"
and instruction.get("parts")
and not request.model.endswith("-image")
and not request.model.endswith("-image-generation")
):
payload["systemInstruction"] = instruction