fix: streaming 状态更新时传递 provider 相关 ID 信息

在 update_usage_status 方法中增加 provider_id、provider_endpoint_id
和 provider_api_key_id 参数,确保流式请求进入 streaming 状态时
能正确记录这些字段。
This commit is contained in:
fawney19
2026-01-05 09:12:03 +08:00
parent 3064497636
commit b34cfb676d
2 changed files with 18 additions and 0 deletions

View File

@@ -1476,6 +1476,9 @@ class UsageService:
provider: Optional[str] = None,
target_model: Optional[str] = None,
first_byte_time_ms: Optional[int] = None,
provider_id: Optional[str] = None,
provider_endpoint_id: Optional[str] = None,
provider_api_key_id: Optional[str] = None,
) -> Optional[Usage]:
"""
快速更新使用记录状态
@@ -1488,6 +1491,9 @@ class UsageService:
provider: 提供商名称可选streaming 状态时更新)
target_model: 映射后的目标模型名(可选)
first_byte_time_ms: 首字时间/TTFB可选streaming 状态时更新)
provider_id: Provider ID可选streaming 状态时更新)
provider_endpoint_id: Endpoint ID可选streaming 状态时更新)
provider_api_key_id: Provider API Key ID可选streaming 状态时更新)
Returns:
更新后的 Usage 记录,如果未找到则返回 None
@@ -1513,6 +1519,12 @@ class UsageService:
usage.target_model = target_model
if first_byte_time_ms is not None:
usage.first_byte_time_ms = first_byte_time_ms
if provider_id is not None:
usage.provider_id = provider_id
if provider_endpoint_id is not None:
usage.provider_endpoint_id = provider_endpoint_id
if provider_api_key_id is not None:
usage.provider_api_key_id = provider_api_key_id
db.commit()

View File

@@ -484,6 +484,9 @@ class StreamUsageTracker:
status="streaming",
provider=self.provider,
first_byte_time_ms=first_byte_time_ms,
provider_id=self.provider_id,
provider_endpoint_id=self.provider_endpoint_id,
provider_api_key_id=self.provider_api_key_id,
)
except Exception as e:
logger.warning(f"更新使用记录状态为 streaming 失败: {e}")
@@ -942,6 +945,9 @@ class EnhancedStreamUsageTracker(StreamUsageTracker):
status="streaming",
provider=self.provider,
first_byte_time_ms=first_byte_time_ms,
provider_id=self.provider_id,
provider_endpoint_id=self.provider_endpoint_id,
provider_api_key_id=self.provider_api_key_id,
)
except Exception as e:
logger.warning(f"更新使用记录状态为 streaming 失败: {e}")