mirror of
https://github.com/snailyp/gemini-balance.git
synced 2026-05-12 19:40:03 +08:00
b14bb93d8f
重构项目目录结构，提高代码组织性和可维护性 将schemas目录重命名为domain，更好地表达领域模型概念 将services目录细分为service/chat、service/image等子目录 将api目录重命名为router，更符合FastAPI惯例 创建utils目录存放通用工具函数 更新FastAPI应用程序生命周期管理 替换已弃用的on_event方法为推荐的lifespan事件处理器 添加应用程序关闭时的日志记录 代码质量改进 抽取常量到constants.py，减少硬编码值 添加helpers.py提供通用工具函数 优化配置管理，使用环境变量和默认值 完善文档字符串，提高代码可读性
62 lines
2.4 KiB
Python
62 lines
2.4 KiB
Python
# app/services/chat/api_client.py
|
|
|
|
from typing import Dict, Any, AsyncGenerator
|
|
import httpx
|
|
from abc import ABC, abstractmethod
|
|
|
|
from app.core.constants import DEFAULT_TIMEOUT
|
|
|
|
|
|
class ApiClient(ABC):
|
|
"""API客户端基类"""
|
|
|
|
@abstractmethod
|
|
async def generate_content(self, payload: Dict[str, Any], model: str, api_key: str) -> Dict[str, Any]:
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def stream_generate_content(self, payload: Dict[str, Any], model: str, api_key: str) -> AsyncGenerator[str, None]:
|
|
pass
|
|
|
|
|
|
class GeminiApiClient(ApiClient):
|
|
"""Gemini API客户端"""
|
|
|
|
def __init__(self, base_url: str, timeout: int = DEFAULT_TIMEOUT):
|
|
self.base_url = base_url
|
|
self.timeout = timeout
|
|
|
|
def _get_real_model(self, model: str) -> str:
|
|
if model.endswith("-search"):
|
|
model = model[:-7]
|
|
if model.endswith("-image"):
|
|
model = model[:-6]
|
|
|
|
return model
|
|
|
|
async def generate_content(self, payload: Dict[str, Any], model: str, api_key: str) -> Dict[str, Any]:
|
|
timeout = httpx.Timeout(self.timeout, read=self.timeout)
|
|
model = self._get_real_model(model)
|
|
|
|
async with httpx.AsyncClient(timeout=timeout) as client:
|
|
url = f"{self.base_url}/models/{model}:generateContent?key={api_key}"
|
|
response = await client.post(url, json=payload)
|
|
if response.status_code != 200:
|
|
error_content = response.text
|
|
raise Exception(f"API call failed with status code {response.status_code}, {error_content}")
|
|
return response.json()
|
|
|
|
async def stream_generate_content(self, payload: Dict[str, Any], model: str, api_key: str) -> AsyncGenerator[str, None]:
|
|
timeout = httpx.Timeout(self.timeout, read=self.timeout)
|
|
model = self._get_real_model(model)
|
|
|
|
async with httpx.AsyncClient(timeout=timeout) as client:
|
|
url = f"{self.base_url}/models/{model}:streamGenerateContent?alt=sse&key={api_key}"
|
|
async with client.stream(method="POST", url=url, json=payload) as response:
|
|
if response.status_code != 200:
|
|
error_content = await response.aread()
|
|
error_msg = error_content.decode("utf-8")
|
|
raise Exception(f"API call failed with status code {response.status_code}, {error_msg}")
|
|
async for line in response.aiter_lines():
|
|
yield line
|