mirror of
https://github.com/snailyp/gemini-balance.git
synced 2026-06-09 17:49:42 +08:00
refactor: 项目结构优化与FastAPI生命周期更新
重构项目目录结构，提高代码组织性和可维护性 将schemas目录重命名为domain，更好地表达领域模型概念 将services目录细分为service/chat、service/image等子目录 将api目录重命名为router，更符合FastAPI惯例 创建utils目录存放通用工具函数 更新FastAPI应用程序生命周期管理 替换已弃用的on_event方法为推荐的lifespan事件处理器 添加应用程序关闭时的日志记录 代码质量改进 抽取常量到constants.py，减少硬编码值 添加helpers.py提供通用工具函数 优化配置管理，使用环境变量和默认值 完善文档字符串，提高代码可读性
This commit is contained in:
snaily
164
app/handler/message_converter.py
Normal file
164
app/handler/message_converter.py
Normal file
@@ -0,0 +1,164 @@
|
||||
# app/services/chat/message_converter.py
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
import re
|
||||
from typing import Any, Dict, List, Optional
|
||||
import requests
|
||||
import base64
|
||||
|
||||
from app.core.constants import DATA_URL_PATTERN, IMAGE_URL_PATTERN, SUPPORTED_ROLES
|
||||
|
||||
|
||||
class MessageConverter(ABC):
|
||||
"""消息转换器基类"""
|
||||
|
||||
@abstractmethod
|
||||
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 = DATA_URL_PATTERN
|
||||
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": mime_type,
|
||||
"data": encoded_data
|
||||
}
|
||||
}
|
||||
return {
|
||||
"image_url": {
|
||||
"url": image_url
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
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消息格式转换器"""
|
||||
|
||||
def convert(self, messages: List[Dict[str, Any]]) -> tuple[List[Dict[str, Any]], Optional[Dict[str, Any]]]:
|
||||
converted_messages = []
|
||||
system_instruction_parts = []
|
||||
|
||||
for idx, msg in enumerate(messages):
|
||||
role = msg.get("role", "")
|
||||
if role not in SUPPORTED_ROLES:
|
||||
if role == "tool":
|
||||
role = "user"
|
||||
else:
|
||||
# 如果是最后一条消息，则认为是用户消息
|
||||
if idx == len(messages) - 1:
|
||||
role = "user"
|
||||
else:
|
||||
role = "model"
|
||||
|
||||
parts = []
|
||||
# 特别处理最后一个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.extend(_process_text_with_image(msg["content"]))
|
||||
elif isinstance(msg["content"], list):
|
||||
for content in msg["content"]:
|
||||
if isinstance(content, str) and content:
|
||||
parts.append({"text": content})
|
||||
elif isinstance(content, dict):
|
||||
if content["type"] == "text" and content["text"]:
|
||||
parts.append({"text": content["text"]})
|
||||
elif content["type"] == "image_url":
|
||||
parts.append(_convert_image(content["image_url"]["url"]))
|
||||
|
||||
if parts:
|
||||
if role == "system":
|
||||
system_instruction_parts.extend(parts)
|
||||
else:
|
||||
converted_messages.append({"role": role, "parts": parts})
|
||||
|
||||
system_instruction = (
|
||||
None
|
||||
if not system_instruction_parts
|
||||
else {
|
||||
"role": "system",
|
||||
"parts": system_instruction_parts,
|
||||
}
|
||||
)
|
||||
return converted_messages, system_instruction
|
||||
306
app/handler/response_handler.py
Normal file
306
app/handler/response_handler.py
Normal file
@@ -0,0 +1,306 @@
|
||||
# app/services/chat/response_handler.py
|
||||
|
||||
import base64
|
||||
import json
|
||||
import random
|
||||
import string
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Dict, Any, List, Optional
|
||||
import time
|
||||
import uuid
|
||||
from app.config.config import settings
|
||||
from app.utils.uploader import ImageUploaderFactory
|
||||
|
||||
|
||||
class ResponseHandler(ABC):
|
||||
"""响应处理器基类"""
|
||||
|
||||
@abstractmethod
|
||||
def handle_response(self, response: Dict[str, Any], model: str, stream: bool = False) -> Dict[str, Any]:
|
||||
pass
|
||||
|
||||
|
||||
class GeminiResponseHandler(ResponseHandler):
|
||||
"""Gemini响应处理器"""
|
||||
|
||||
def __init__(self):
|
||||
self.thinking_first = True
|
||||
self.thinking_status = False
|
||||
|
||||
def handle_response(self, response: Dict[str, Any], model: str, stream: bool = False) -> Dict[str, Any]:
|
||||
if stream:
|
||||
return _handle_gemini_stream_response(response, model, stream)
|
||||
return _handle_gemini_normal_response(response, model, stream)
|
||||
|
||||
|
||||
def _handle_openai_stream_response(response: Dict[str, Any], model: str, finish_reason: str) -> Dict[str, Any]:
|
||||
text, tool_calls = _extract_result(response, model, stream=True, gemini_format=False)
|
||||
if not text and not tool_calls:
|
||||
delta = {}
|
||||
else:
|
||||
delta = {"content": text, "role": "assistant"}
|
||||
if tool_calls:
|
||||
delta["tool_calls"] = tool_calls
|
||||
|
||||
return {
|
||||
"id": f"chatcmpl-{uuid.uuid4()}",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": int(time.time()),
|
||||
"model": model,
|
||||
"choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}],
|
||||
}
|
||||
|
||||
|
||||
def _handle_openai_normal_response(response: Dict[str, Any], model: str, finish_reason: str) -> Dict[str, Any]:
|
||||
text, tool_calls = _extract_result(response, model, stream=False, gemini_format=False)
|
||||
return {
|
||||
"id": f"chatcmpl-{uuid.uuid4()}",
|
||||
"object": "chat.completion",
|
||||
"created": int(time.time()),
|
||||
"model": model,
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": text, "tool_calls": tool_calls},
|
||||
"finish_reason": finish_reason,
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
|
||||
}
|
||||
|
||||
|
||||
class OpenAIResponseHandler(ResponseHandler):
|
||||
"""OpenAI响应处理器"""
|
||||
|
||||
def __init__(self, config):
|
||||
self.config = config
|
||||
self.thinking_first = True
|
||||
self.thinking_status = False
|
||||
|
||||
def handle_response(
|
||||
self,
|
||||
response: Dict[str, Any],
|
||||
model: str,
|
||||
stream: bool = False,
|
||||
finish_reason: str = None
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
if stream:
|
||||
return _handle_openai_stream_response(response, model, finish_reason)
|
||||
return _handle_openai_normal_response(response, model, finish_reason)
|
||||
|
||||
def handle_image_chat_response(self, image_str: str, model: str, stream=False, finish_reason="stop"):
|
||||
if stream:
|
||||
return _handle_openai_stream_image_response(image_str,model,finish_reason)
|
||||
return _handle_openai_normal_image_response(image_str,model,finish_reason)
|
||||
|
||||
|
||||
def _handle_openai_stream_image_response(image_str: str,model: str,finish_reason: str) -> Dict[str, Any]:
|
||||
return {
|
||||
"id": f"chatcmpl-{uuid.uuid4()}",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": int(time.time()),
|
||||
"model": model,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {"content": image_str} if image_str else {},
|
||||
"finish_reason": finish_reason
|
||||
}]
|
||||
}
|
||||
|
||||
|
||||
def _handle_openai_normal_image_response(image_str: str,model: str,finish_reason: str) -> Dict[str, Any]:
|
||||
return {
|
||||
"id": f"chatcmpl-{uuid.uuid4()}",
|
||||
"object": "chat.completion",
|
||||
"created": int(time.time()),
|
||||
"model": model,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": image_str
|
||||
},
|
||||
"finish_reason": finish_reason
|
||||
}],
|
||||
"usage": {
|
||||
"prompt_tokens": 0,
|
||||
"completion_tokens": 0,
|
||||
"total_tokens": 0
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def _extract_result(response: Dict[str, Any], model: str, stream: bool = False, gemini_format: bool = False) -> tuple[str, List[Dict[str, Any]]]:
|
||||
text, tool_calls = "", []
|
||||
if stream:
|
||||
if response.get("candidates"):
|
||||
candidate = response["candidates"][0]
|
||||
content = candidate.get("content", {})
|
||||
parts = content.get("parts", [])
|
||||
if not parts:
|
||||
return "", []
|
||||
if "text" in parts[0]:
|
||||
text = parts[0].get("text")
|
||||
elif "executableCode" in parts[0]:
|
||||
text = _format_code_block(parts[0]["executableCode"])
|
||||
elif "codeExecution" in parts[0]:
|
||||
text = _format_code_block(parts[0]["codeExecution"])
|
||||
elif "executableCodeResult" in parts[0]:
|
||||
text = _format_execution_result(
|
||||
parts[0]["executableCodeResult"]
|
||||
)
|
||||
elif "codeExecutionResult" in parts[0]:
|
||||
text = _format_execution_result(
|
||||
parts[0]["codeExecutionResult"]
|
||||
)
|
||||
elif "inlineData" in parts[0]:
|
||||
text = _extract_image_data(parts[0])
|
||||
else:
|
||||
text = ""
|
||||
text = _add_search_link_text(model, candidate, text)
|
||||
tool_calls = _extract_tool_calls(parts, gemini_format)
|
||||
else:
|
||||
if response.get("candidates"):
|
||||
candidate = response["candidates"][0]
|
||||
if "thinking" in model:
|
||||
if settings.SHOW_THINKING_PROCESS:
|
||||
if len(candidate["content"]["parts"]) == 2:
|
||||
text = (
|
||||
"> thinking\n\n"
|
||||
+ candidate["content"]["parts"][0]["text"]
|
||||
+ "\n\n---\n> output\n\n"
|
||||
+ candidate["content"]["parts"][1]["text"]
|
||||
)
|
||||
else:
|
||||
text = candidate["content"]["parts"][0]["text"]
|
||||
else:
|
||||
if len(candidate["content"]["parts"]) == 2:
|
||||
text = candidate["content"]["parts"][1]["text"]
|
||||
else:
|
||||
text = candidate["content"]["parts"][0]["text"]
|
||||
else:
|
||||
text = ""
|
||||
if "parts" in candidate["content"]:
|
||||
for part in candidate["content"]["parts"]:
|
||||
if "text" in part:
|
||||
text += part["text"]
|
||||
elif "inlineData" in part:
|
||||
text += _extract_image_data(part)
|
||||
|
||||
|
||||
text = _add_search_link_text(model, candidate, text)
|
||||
tool_calls = _extract_tool_calls(candidate["content"]["parts"], gemini_format)
|
||||
else:
|
||||
text = "暂无返回"
|
||||
return text, tool_calls
|
||||
|
||||
def _extract_image_data(part: dict) -> str:
|
||||
image_uploader = None
|
||||
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)
|
||||
current_date = time.strftime("%Y/%m/%d")
|
||||
filename = f"{current_date}/{uuid.uuid4().hex[:8]}.png"
|
||||
base64_data = part["inlineData"]["data"]
|
||||
#将base64_data转成bytes数组
|
||||
bytes_data = base64.b64decode(base64_data)
|
||||
upload_response = image_uploader.upload(bytes_data,filename)
|
||||
if upload_response.success:
|
||||
text = f""
|
||||
else:
|
||||
text = ""
|
||||
return text
|
||||
|
||||
def _extract_tool_calls(parts: List[Dict[str, Any]], gemini_format: bool) -> List[Dict[str, Any]]:
|
||||
"""提取工具调用信息"""
|
||||
if not parts or not isinstance(parts, list):
|
||||
return []
|
||||
|
||||
letters = string.ascii_lowercase + string.digits
|
||||
|
||||
tool_calls = list()
|
||||
for i in range(len(parts)):
|
||||
part = parts[i]
|
||||
if not part or not isinstance(part, dict):
|
||||
continue
|
||||
|
||||
item = part.get("functionCall", {})
|
||||
if not item or not isinstance(item, dict):
|
||||
continue
|
||||
|
||||
if gemini_format:
|
||||
tool_calls.append(part)
|
||||
else:
|
||||
id = f"call_{''.join(random.sample(letters, 32))}"
|
||||
name = item.get("name", "")
|
||||
arguments = json.dumps(item.get("args", None) or {})
|
||||
|
||||
tool_calls.append(
|
||||
{
|
||||
"index": i,
|
||||
"id": id,
|
||||
"type": "function",
|
||||
"function": {"name": name, "arguments": arguments},
|
||||
}
|
||||
)
|
||||
|
||||
return tool_calls
|
||||
|
||||
|
||||
def _handle_gemini_stream_response(response: Dict[str, Any], model: str, stream: bool) -> Dict[str, Any]:
|
||||
text, tool_calls = _extract_result(response, model, stream=stream, gemini_format=True)
|
||||
if tool_calls:
|
||||
content = {"parts": tool_calls, "role": "model"}
|
||||
else:
|
||||
content = {"parts": [{"text": text}], "role": "model"}
|
||||
response["candidates"][0]["content"] = content
|
||||
return response
|
||||
|
||||
|
||||
def _handle_gemini_normal_response(response: Dict[str, Any], model: str, stream: bool) -> Dict[str, Any]:
|
||||
text, tool_calls = _extract_result(response, model, stream=stream, gemini_format=True)
|
||||
if tool_calls:
|
||||
content = {"parts": tool_calls, "role": "model"}
|
||||
else:
|
||||
content = {"parts": [{"text": text}], "role": "model"}
|
||||
response["candidates"][0]["content"] = content
|
||||
return response
|
||||
|
||||
|
||||
def _format_code_block(code_data: dict) -> str:
|
||||
"""格式化代码块输出"""
|
||||
language = code_data.get("language", "").lower()
|
||||
code = code_data.get("code", "").strip()
|
||||
return f"""\n\n---\n\n【代码执行】\n```{language}\n{code}\n```\n"""
|
||||
|
||||
|
||||
def _add_search_link_text(model: str, candidate: dict, text: str) -> str:
|
||||
if (
|
||||
settings.SHOW_SEARCH_LINK
|
||||
and model.endswith("-search")
|
||||
and "groundingMetadata" in candidate
|
||||
and "groundingChunks" in candidate["groundingMetadata"]
|
||||
):
|
||||
grounding_chunks = candidate["groundingMetadata"]["groundingChunks"]
|
||||
text += "\n\n---\n\n"
|
||||
text += "**【引用来源】**\n\n"
|
||||
for _, grounding_chunk in enumerate(grounding_chunks, 1):
|
||||
if "web" in grounding_chunk:
|
||||
text += _create_search_link(grounding_chunk["web"])
|
||||
return text
|
||||
else:
|
||||
return text
|
||||
|
||||
|
||||
def _create_search_link(grounding_chunk: dict) -> str:
|
||||
return f'\n- [{grounding_chunk["title"]}]({grounding_chunk["uri"]})'
|
||||
|
||||
|
||||
def _format_execution_result(result_data: dict) -> str:
|
||||
"""格式化执行结果输出"""
|
||||
outcome = result_data.get("outcome", "")
|
||||
output = result_data.get("output", "").strip()
|
||||
return f"""\n【执行结果】\n> outcome: {outcome}\n\n【输出结果】\n```plaintext\n{output}\n```\n\n---\n\n"""
|
||||
41
app/handler/retry_handler.py
Normal file
41
app/handler/retry_handler.py
Normal file
@@ -0,0 +1,41 @@
|
||||
# app/services/chat/retry_handler.py
|
||||
|
||||
from typing import TypeVar, Callable
|
||||
from functools import wraps
|
||||
from app.logger.logger import get_retry_logger
|
||||
|
||||
T = TypeVar('T')
|
||||
logger = get_retry_logger()
|
||||
|
||||
|
||||
class RetryHandler:
|
||||
"""重试处理装饰器"""
|
||||
|
||||
def __init__(self, max_retries: int = 3, key_arg: str = "api_key"):
|
||||
self.max_retries = max_retries
|
||||
self.key_arg = key_arg
|
||||
|
||||
def __call__(self, func: Callable[..., T]) -> Callable[..., T]:
|
||||
@wraps(func)
|
||||
async def wrapper(*args, **kwargs) -> T:
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(self.max_retries):
|
||||
try:
|
||||
return await func(*args, **kwargs)
|
||||
except Exception as e:
|
||||
last_exception = e
|
||||
logger.warning(f"API call failed with error: {str(e)}. Attempt {attempt + 1} of {self.max_retries}")
|
||||
|
||||
# 从函数参数中获取 key_manager
|
||||
key_manager = kwargs.get('key_manager')
|
||||
if key_manager:
|
||||
old_key = kwargs.get(self.key_arg)
|
||||
new_key = await key_manager.handle_api_failure(old_key)
|
||||
kwargs[self.key_arg] = new_key
|
||||
logger.info(f"Switched to new API key: {new_key}")
|
||||
|
||||
logger.error(f"All retry attempts failed, raising final exception: {str(last_exception)}")
|
||||
raise last_exception
|
||||
|
||||
return wrapper
|
||||
133
app/handler/stream_optimizer.py
Normal file
133
app/handler/stream_optimizer.py
Normal file
@@ -0,0 +1,133 @@
|
||||
# app/services/chat/stream_optimizer.py
|
||||
|
||||
import asyncio
|
||||
import math
|
||||
from typing import Any, List, AsyncGenerator, Callable
|
||||
from app.logger.logger import get_openai_logger, get_gemini_logger
|
||||
from app.config.config import settings
|
||||
from app.core.constants import DEFAULT_STREAM_CHUNK_SIZE, DEFAULT_STREAM_LONG_TEXT_THRESHOLD, DEFAULT_STREAM_MAX_DELAY, DEFAULT_STREAM_MIN_DELAY, DEFAULT_STREAM_SHORT_TEXT_THRESHOLD
|
||||
|
||||
logger_openai = get_openai_logger()
|
||||
logger_gemini = get_gemini_logger()
|
||||
|
||||
|
||||
class StreamOptimizer:
|
||||
"""流式输出优化器
|
||||
|
||||
提供流式输出优化功能，包括智能延迟调整和长文本分块输出。
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
logger=None,
|
||||
min_delay: float = DEFAULT_STREAM_MIN_DELAY,
|
||||
max_delay: float = DEFAULT_STREAM_MAX_DELAY,
|
||||
short_text_threshold: int = DEFAULT_STREAM_SHORT_TEXT_THRESHOLD,
|
||||
long_text_threshold: int = DEFAULT_STREAM_LONG_TEXT_THRESHOLD,
|
||||
chunk_size: int = DEFAULT_STREAM_CHUNK_SIZE):
|
||||
"""初始化流式输出优化器
|
||||
|
||||
参数:
|
||||
logger: 日志记录器
|
||||
min_delay: 最小延迟时间（秒）
|
||||
max_delay: 最大延迟时间（秒）
|
||||
short_text_threshold: 短文本阈值（字符数）
|
||||
long_text_threshold: 长文本阈值（字符数）
|
||||
chunk_size: 长文本分块大小（字符数）
|
||||
"""
|
||||
self.logger = logger
|
||||
self.min_delay = min_delay
|
||||
self.max_delay = max_delay
|
||||
self.short_text_threshold = short_text_threshold
|
||||
self.long_text_threshold = long_text_threshold
|
||||
self.chunk_size = chunk_size
|
||||
|
||||
def calculate_delay(self, text_length: int) -> float:
|
||||
"""根据文本长度计算延迟时间
|
||||
|
||||
参数:
|
||||
text_length: 文本长度
|
||||
|
||||
返回:
|
||||
延迟时间（秒）
|
||||
"""
|
||||
if text_length <= self.short_text_threshold:
|
||||
# 短文本使用较大延迟
|
||||
return self.max_delay
|
||||
elif text_length >= self.long_text_threshold:
|
||||
# 长文本使用较小延迟
|
||||
return self.min_delay
|
||||
else:
|
||||
# 中等长度文本使用线性插值计算延迟
|
||||
# 使用对数函数使延迟变化更平滑
|
||||
ratio = math.log(text_length / self.short_text_threshold) / math.log(self.long_text_threshold / self.short_text_threshold)
|
||||
return self.max_delay - ratio * (self.max_delay - self.min_delay)
|
||||
|
||||
def split_text_into_chunks(self, text: str) -> List[str]:
|
||||
"""将文本分割成小块
|
||||
|
||||
参数:
|
||||
text: 要分割的文本
|
||||
|
||||
返回:
|
||||
文本块列表
|
||||
"""
|
||||
return [text[i:i+self.chunk_size] for i in range(0, len(text), self.chunk_size)]
|
||||
|
||||
async def optimize_stream_output(self,
|
||||
text: str,
|
||||
create_response_chunk: Callable[[str], Any],
|
||||
format_chunk: Callable[[Any], str]) -> AsyncGenerator[str, None]:
|
||||
"""优化流式输出
|
||||
|
||||
参数:
|
||||
text: 要输出的文本
|
||||
create_response_chunk: 创建响应块的函数，接收文本，返回响应块
|
||||
format_chunk: 格式化响应块的函数，接收响应块，返回格式化后的字符串
|
||||
|
||||
返回:
|
||||
异步生成器，生成格式化后的响应块
|
||||
"""
|
||||
if not text:
|
||||
return
|
||||
|
||||
# 计算智能延迟时间
|
||||
delay = self.calculate_delay(len(text))
|
||||
if self.logger:
|
||||
self.logger.info(f"Text length: {len(text)}, delay: {delay:.4f}s")
|
||||
|
||||
# 根据文本长度决定输出方式
|
||||
if len(text) >= self.long_text_threshold:
|
||||
# 长文本：分块输出
|
||||
chunks = self.split_text_into_chunks(text)
|
||||
if self.logger:
|
||||
self.logger.info(f"Long text: splitting into {len(chunks)} chunks")
|
||||
for chunk_text in chunks:
|
||||
chunk_response = create_response_chunk(chunk_text)
|
||||
yield format_chunk(chunk_response)
|
||||
await asyncio.sleep(delay)
|
||||
else:
|
||||
# 短文本：逐字符输出
|
||||
for char in text:
|
||||
char_chunk = create_response_chunk(char)
|
||||
yield format_chunk(char_chunk)
|
||||
await asyncio.sleep(delay)
|
||||
|
||||
|
||||
# 创建默认的优化器实例，可以直接导入使用
|
||||
openai_optimizer = StreamOptimizer(
|
||||
logger=logger_openai,
|
||||
min_delay=settings.STREAM_MIN_DELAY,
|
||||
max_delay=settings.STREAM_MAX_DELAY,
|
||||
short_text_threshold=settings.STREAM_SHORT_TEXT_THRESHOLD,
|
||||
long_text_threshold=settings.STREAM_LONG_TEXT_THRESHOLD,
|
||||
chunk_size=settings.STREAM_CHUNK_SIZE
|
||||
)
|
||||
|
||||
gemini_optimizer = StreamOptimizer(
|
||||
logger=logger_gemini,
|
||||
min_delay=settings.STREAM_MIN_DELAY,
|
||||
max_delay=settings.STREAM_MAX_DELAY,
|
||||
short_text_threshold=settings.STREAM_SHORT_TEXT_THRESHOLD,
|
||||
long_text_threshold=settings.STREAM_LONG_TEXT_THRESHOLD,
|
||||
chunk_size=settings.STREAM_CHUNK_SIZE
|
||||
)
|
||||
Reference in New Issue
Block a user