""" 流式响应用量统计服务 处理流式响应的token计算和使用量记录 """ import json import re from typing import Any, AsyncIterator, Dict, Optional, Tuple from sqlalchemy.orm import Session from src.api.handlers.base.parsers import get_parser_for_format from src.api.handlers.base.response_parser import StreamStats from src.core.exceptions import EmptyStreamException from src.core.logger import logger from src.database.database import create_session from src.models.database import ApiKey, User from src.services.usage.service import UsageService class StreamUsageTracker: """流式响应用量跟踪器""" def __init__( self, db: Session, user: User, api_key: ApiKey, provider: str, model: str, request_headers: Optional[Dict[str, Any]] = None, provider_request_headers: Optional[Dict[str, Any]] = None, request_id: Optional[str] = None, start_time: Optional[float] = None, attempt_id: Optional[str] = None, # Provider 侧追踪信息（用于记录真实成本） provider_id: Optional[str] = None, provider_endpoint_id: Optional[str] = None, provider_api_key_id: Optional[str] = None, # API 格式（用于选择正确的响应解析器） api_format: Optional[str] = None, ): """ 初始化流式用量跟踪器 Args: db: 数据库会话 user: 用户对象 api_key: API密钥对象 provider: 提供商名称 model: 模型名称 request_headers: 实际的请求头 provider_request_headers: 向提供商发送的请求头 request_id: 请求ID（用于日志关联） start_time: 请求开始时间（用于计算总响应时间） attempt_id: RequestTrace 请求尝试ID provider_id: Provider ID（用于记录真实成本） provider_endpoint_id: Endpoint ID（用于记录真实成本） provider_api_key_id: API Key ID（用于记录真实成本） api_format: API 格式（CLAUDE, CLAUDE_CLI, OPENAI, OPENAI_CLI） """ self.db = db # 只存储ID，避免会话绑定问题 self.user_id = user.id if user else None self.api_key_id = api_key.id if api_key else None self.provider = provider self.model = model self.request_headers = request_headers or {} self.provider_request_headers = provider_request_headers or {} self.request_id = request_id self.request_start_time = start_time # Provider 侧追踪信息 self.provider_id = provider_id self.provider_endpoint_id = provider_endpoint_id self.provider_api_key_id = provider_api_key_id # API 格式和响应解析器 self.api_format = api_format or "CLAUDE" self.response_parser = get_parser_for_format(self.api_format) self.stream_stats = StreamStats() # 解析器统计信息 # Token计数 self.input_tokens = 0 self.output_tokens = 0 self.cache_creation_input_tokens = 0 self.cache_read_input_tokens = 0 self.accumulated_content = "" # 完整响应跟踪（仅用于内部统计，不记录到数据库） self.complete_response = { "id": None, "type": "message", "role": "assistant", "content": [], "model": model, "stop_reason": None, "stop_sequence": None, "usage": {}, } self.response_chunks = [] # 保存所有原始响应块 self.raw_chunks = [] # 保存所有原始字节流（用于错误诊断） # 时间跟踪 self.start_time = None self.end_time = None # 响应头 (将在track_stream中设置) self.response_headers = {} # SSE解析缓冲区 self.buffer = b"" # 用于处理不完整的字节流 self.current_line = "" # 用于累积SSE行 self.sse_event_buffer = { "event": None, "data": [], "id": None, "retry": None, } # SSE事件缓冲 # 错误状态跟踪 self.status_code = 200 # 默认成功状态码 self.error_message = None # 错误消息(如果有) self.attempt_id = attempt_id def set_error_status(self, status_code: int, error_message: str): """ 设置错误状态 Args: status_code: HTTP状态码 error_message: 错误消息 """ self.status_code = status_code self.error_message = error_message logger.debug(f"ID:{self.request_id} | 流式响应错误状态已设置 | 状态码:{status_code} | 错误:{error_message[:100]}") def _update_complete_response(self, chunk: Dict[str, Any]): """根据响应块更新完整响应结构""" try: # 更新响应ID if chunk.get("id"): self.complete_response["id"] = chunk["id"] # 更新模型 if chunk.get("model"): self.complete_response["model"] = chunk["model"] # 处理不同类型的事件 event_type = chunk.get("type") if event_type == "message_start": # 消息开始事件 message = chunk.get("message", {}) if message.get("id"): self.complete_response["id"] = message["id"] if message.get("model"): self.complete_response["model"] = message["model"] self.complete_response["usage"] = message.get("usage", {}) elif event_type == "content_block_start": # 内容块开始 content_block = chunk.get("content_block", {}) self.complete_response["content"].append(content_block) elif event_type == "content_block_delta": # 内容块增量更新 index = chunk.get("index", 0) delta = chunk.get("delta", {}) # 确保content列表有足够的元素 while len(self.complete_response["content"]) <= index: self.complete_response["content"].append({"type": "text", "text": ""}) current_block = self.complete_response["content"][index] if delta.get("type") == "text_delta": # 文本增量 if current_block.get("type") == "text": current_block["text"] = current_block.get("text", "") + delta.get( "text", "" ) elif delta.get("type") == "input_json_delta": # 工具调用输入增量 if current_block.get("type") == "tool_use": current_input = current_block.get("input", {}) if isinstance(current_input, str): current_input += delta.get("partial_json", "") current_block["input"] = current_input elif event_type == "content_block_stop": # 内容块结束 pass elif event_type == "message_delta": # 消息级别的增量更新 delta = chunk.get("delta", {}) if delta.get("stop_reason"): self.complete_response["stop_reason"] = delta["stop_reason"] if delta.get("stop_sequence"): self.complete_response["stop_sequence"] = delta["stop_sequence"] elif event_type == "message_stop": # 消息结束 pass # 更新usage信息 if chunk.get("usage"): self.complete_response["usage"].update(chunk["usage"]) except Exception as e: # 记录错误但不中断流处理 logger.warning(f"Failed to update complete response: {e}") def estimate_input_tokens(self, messages: list) -> int: """ 估算输入tokens Args: messages: 消息列表 Returns: 估算的token数 """ total_chars = 0 for msg in messages: if isinstance(msg, dict): content = msg.get("content", "") if isinstance(content, str): total_chars += len(content) elif isinstance(content, list): for block in content: if hasattr(block, "text"): total_chars += len(block.text) # 粗略估算：4个字符约等于1个token return max(1, total_chars // 4) def _process_sse_event(self) -> Tuple[Optional[str], Optional[Dict[str, Any]]]: """ 处理缓冲区中的完整SSE事件 Returns: (内容文本, 使用信息) """ content = None usage = None # 如果没有data，直接返回 if not self.sse_event_buffer["data"]: return None, None # 合并所有data行（根据SSE规范，每个data行之间加入换行符） data_str = "\n".join(self.sse_event_buffer["data"]) # 清空缓冲区 self.sse_event_buffer = {"event": None, "data": [], "id": None, "retry": None} if not data_str or data_str == "[DONE]": return None, None try: data = json.loads(data_str) if isinstance(data, dict): self.response_chunks.append(data) try: self._update_complete_response(data) except Exception as update_error: logger.warning(f"Failed to update complete response from chunk: {update_error}") # Claude格式 if "type" in data: if data["type"] == "content_block_delta": delta = data.get("delta", {}) content = delta.get("text", "") if content: logger.debug(f"Extracted content from delta: {len(content)} chars") elif data["type"] == "message_delta": usage_data = data.get("usage", {}) if usage_data: usage = usage_data logger.debug(f"Extracted usage from message_delta: {usage}") elif data["type"] == "message_stop": logger.debug("Received message_stop event") # OpenAI格式 elif "choices" in data: choices = data.get("choices", []) if choices: delta = choices[0].get("delta", {}) content = delta.get("content", "") if "usage" in data: usage = data["usage"] except json.JSONDecodeError as e: # 更详细的JSON解析错误日志 logger.warning(f"Failed to parse SSE JSON data: {str(e)}") return None, None except Exception as e: logger.error(f"Unexpected error processing SSE event: {type(e).__name__}: {str(e)}") return content, usage def parse_sse_line(self, line: str) -> Tuple[Optional[str], Optional[Dict[str, Any]]]: """ 解析单行SSE事件（使用统一响应解析器） Args: line: SSE格式的一行 Returns: (内容文本, 使用信息) - 当遇到空行时处理完整事件 """ # 使用统一的响应解析器 chunk = self.response_parser.parse_sse_line(line, self.stream_stats) if chunk is None: return None, None # 从 ParsedChunk 中提取内容和使用信息 content = chunk.text_delta # 构建 usage 字典（如果有 token 信息） usage = None if ( chunk.input_tokens or chunk.output_tokens or chunk.cache_creation_tokens or chunk.cache_read_tokens ): usage = { "input_tokens": chunk.input_tokens or self.stream_stats.input_tokens, "output_tokens": chunk.output_tokens or self.stream_stats.output_tokens, "cache_creation_input_tokens": chunk.cache_creation_tokens or self.stream_stats.cache_creation_tokens, "cache_read_input_tokens": chunk.cache_read_tokens or self.stream_stats.cache_read_tokens, } # 更新响应 ID if chunk.response_id and not self.complete_response.get("id"): self.complete_response["id"] = chunk.response_id # 更新完整响应（如果有数据） if chunk.data: self.response_chunks.append(chunk.data) try: self._update_complete_response(chunk.data) except Exception as update_error: logger.warning(f"Failed to update complete response from chunk: {update_error}") return content, usage def parse_stream_chunk(self, chunk: bytes) -> Tuple[Optional[str], Optional[Dict[str, Any]]]: """ 解析流式响应块（处理原始字节流） Args: chunk: 原始字节流 Returns: (累积的内容文本, 使用信息) """ total_content = "" final_usage = None # 将新chunk添加到缓冲区 # 确保 chunk 是字节类型 if isinstance(chunk, str): chunk = chunk.encode("utf-8") self.buffer += chunk # 尝试解码并处理完整的行 try: # 尝试解码整个缓冲区 text = self.buffer.decode("utf-8") self.buffer = b"" # 清空缓冲区 # 将文本添加到当前行 self.current_line += text # 按换行符分割，处理完整的行 lines = self.current_line.split("\n") # 最后一个可能是不完整的行，保留它 self.current_line = lines[-1] # 处理完整的行 for line in lines[:-1]: line = line.rstrip("\r") content, usage = self.parse_sse_line(line) if content: total_content += content if usage: final_usage = usage except UnicodeDecodeError: # 如果解码失败，说明缓冲区中有不完整的UTF-8序列 # 尝试找到最后一个完整的字符边界 for i in range(len(self.buffer) - 1, max(0, len(self.buffer) - 4), -1): try: text = self.buffer[:i].decode("utf-8") # 成功解码，处理这部分 remaining = self.buffer[i:] self.buffer = remaining # 保留未解码的部分 # 处理解码的文本 self.current_line += text lines = self.current_line.split("\n") self.current_line = lines[-1] for line in lines[:-1]: line = line.rstrip("\r") content, usage = self.parse_sse_line(line) if content: total_content += content if usage: final_usage = usage break except UnicodeDecodeError: continue return total_content if total_content else None, final_usage async def track_stream( self, stream: AsyncIterator[str], request_data: Dict[str, Any], response_headers: Optional[Dict[str, Any]] = None, ) -> AsyncIterator[str]: """ 跟踪流式响应并计算用量 Args: stream: 原始流式响应 request_data: 请求数据 response_headers: 实际的响应头 Yields: 流式响应块 """ import time self.start_time = time.time() self.request_data = request_data # 保存请求数据 # 保存响应头（如果没有提供，使用空字典而不是默认值） # 这样可以确保记录的是实际的响应头，而不是构造的默认值 self.response_headers = response_headers if response_headers is not None else {} # 估算输入tokens messages = request_data.get("messages", []) self.input_tokens = self.estimate_input_tokens(messages) logger.debug(f"ID:{self.request_id} | 开始跟踪流式响应 | 估算输入tokens:{self.input_tokens}") # 更新状态为 streaming，同时更新 provider if self.request_id: try: from src.services.usage.service import UsageService UsageService.update_usage_status( db=self.db, request_id=self.request_id, status="streaming", provider=self.provider, ) except Exception as e: logger.warning(f"更新使用记录状态为 streaming 失败: {e}") chunk_count = 0 try: async for chunk in stream: chunk_count += 1 # 保存原始字节流（用于错误诊断） self.raw_chunks.append(chunk) # 返回原始块给客户端 yield chunk # 解析块以提取内容和使用信息（chunk是原始字节） content, usage = self.parse_stream_chunk(chunk) if content: self.accumulated_content += content # 实时估算输出tokens self.output_tokens = max(1, len(self.accumulated_content) // 4) if usage: # 如果响应中包含准确的usage信息，使用它 self.input_tokens = usage.get("input_tokens", self.input_tokens) self.output_tokens = usage.get("output_tokens", self.output_tokens) self.cache_creation_input_tokens = usage.get( "cache_creation_input_tokens", self.cache_creation_input_tokens ) self.cache_read_input_tokens = usage.get( "cache_read_input_tokens", self.cache_read_input_tokens ) # 处理新的cache_creation格式 if "cache_creation" in usage: cache_creation_data = usage.get("cache_creation", {}) # 如果没有cache_creation_input_tokens，尝试从cache_creation中获取 if not self.cache_creation_input_tokens: self.cache_creation_input_tokens = cache_creation_data.get( "ephemeral_5m_input_tokens", 0 ) + cache_creation_data.get("ephemeral_1h_input_tokens", 0) finally: # 流结束后记录使用量 self.end_time = time.time() logger.debug(f"ID:{self.request_id} | 流式响应结束 | 共处理{chunk_count}个chunks | " f"累积内容长度:{len(self.accumulated_content)} | 输出tokens:{self.output_tokens}") # 检查是否收到了有效数据 # 情况1: 收到了原始数据但无法解析为有效的SSE JSON if chunk_count > 0 and not self.response_chunks: error_msg = f"流式响应完成但未解析到任何有效数据（收到 {chunk_count} 个原始chunk但无法解析）" logger.error(f"ID:{self.request_id} | {error_msg}") # 设置错误状态，避免被记录为成功 self.set_error_status(502, error_msg) # 抛出异常让 FallbackOrchestrator 捕获并触发故障转移 raise EmptyStreamException( provider_name=self.provider, chunk_count=chunk_count, request_metadata=None, ) # 情况2: 流式响应完成但没有收到完整的消息（没有 message_stop 事件） # 这种情况通常发生在服务器重启或连接中断时 if not self.stream_stats.has_completion and not self.response_chunks: error_msg = "流式响应中断：未收到任何有效数据（可能是连接中断或服务重启）" logger.warning(f"ID:{self.request_id} | {error_msg}") self.set_error_status(503, error_msg) # 确保日志一定会输出，即使记录usage失败 try: await self._record_usage() except Exception as e: # 如果记录失败，至少输出基本的汇总日志 logger.exception(f"Failed to record stream usage for request {self.request_id}: {e}") # 尝试输出基本的汇总日志，使用多层防护 try: # 计算响应时间，使用多层后备机制 try: if self.request_start_time and self.end_time: total_response_time = int( (self.end_time - self.request_start_time) * 1000 ) elif self.start_time and self.end_time: total_response_time = int((self.end_time - self.start_time) * 1000) else: total_response_time = 0 except Exception: total_response_time = 0 # 安全地输出汇总日志 logger.info(f"[请求完成] ID:{self.request_id or 'unknown'} | 200 | 耗时:{total_response_time}ms | " f"Token:输入{self.input_tokens}/输出{self.output_tokens} | 费用:未知(记录失败)") except Exception as log_error: # 最后的防线：输出最简单的完成标记 logger.error(f"Failed to output summary log: {log_error}") try: logger.info(f"[请求完成] ID:{self.request_id or 'unknown'} | 记录失败但流已完成") except Exception: # 如果连最简单的日志都失败了，放弃 pass async def _record_usage(self): """记录最终的使用量""" try: if self.request_start_time and self.end_time: response_time_ms = int((self.end_time - self.request_start_time) * 1000) elif self.start_time and self.end_time: response_time_ms = int((self.end_time - self.start_time) * 1000) else: response_time_ms = None # 如果没有准确的token计数，使用估算值 if self.output_tokens == 0 and self.accumulated_content: self.output_tokens = max(1, len(self.accumulated_content) // 4) # 使用完整的响应体（包含所有信息，包括工具调用） # 更新最终的usage信息 self.complete_response["usage"].update( { "input_tokens": self.input_tokens, "output_tokens": self.output_tokens, "cache_creation_input_tokens": self.cache_creation_input_tokens, "cache_read_input_tokens": self.cache_read_input_tokens, } ) # 记录响应数据 # 如果成功解析了SSE chunks，使用解析后的结构化数据 # 否则使用原始字节流（用于错误诊断，如403 HTML响应） if self.response_chunks: # 正常情况：成功解析的SSE JSON响应 response_body = { "chunks": self.response_chunks, "metadata": { "stream": True, "total_chunks": len(self.response_chunks), "content_length": len(self.accumulated_content), "response_time_ms": response_time_ms, }, } else: # 错误情况：无法解析为JSON（如HTML错误页面） # 尝试解码原始字节流为文本 raw_response_text = "" for chunk in self.raw_chunks: try: if isinstance(chunk, bytes): raw_response_text += chunk.decode("utf-8", errors="replace") else: raw_response_text += str(chunk) except Exception: pass response_body = { "chunks": [], "raw_response": raw_response_text[:10000], # 限制大小，避免过大 "metadata": { "stream": True, "total_chunks": 0, "raw_chunks_count": len(self.raw_chunks), "content_length": len(raw_response_text), "response_time_ms": response_time_ms, "parse_error": "Failed to parse response as SSE JSON format", }, } # 检查会话状态，如果会话处于不可用状态，需要回滚并创建新事务 from sqlalchemy.exc import InvalidRequestError user = None api_key = None db_for_usage = self.db created_temp_session = False def _load_user_and_key(db_session: Session) -> Tuple[Optional[User], Optional[ApiKey]]: local_user = ( db_session.query(User).filter(User.id == self.user_id).first() if self.user_id else None ) local_api_key = ( db_session.query(ApiKey).filter(ApiKey.id == self.api_key_id).first() if self.api_key_id else None ) return local_user, local_api_key try: # 检查会话是否可用 db_for_usage.info # 重新查询用户和API密钥对象，确保它们在会话中 user, api_key = _load_user_and_key(db_for_usage) except InvalidRequestError: # 会话处于不可用状态，需要回滚并重新开始 logger.warning(f"Session in invalid state for request {self.request_id}, rolling back and retrying") try: db_for_usage.rollback() except Exception: pass try: db_for_usage.close() except Exception: pass # 使用新的会话记录用量，避免 prepared 状态继续影响查询 try: db_for_usage = create_session() created_temp_session = True user, api_key = _load_user_and_key(db_for_usage) except Exception as session_error: logger.exception(f"Failed to recover from invalid session for request {self.request_id}: {session_error}") return # 根据状态码确定请求状态 final_status = "completed" if self.status_code == 200 else "failed" usage_record = await UsageService.record_usage_async( db=db_for_usage, user=user, api_key=api_key, provider=self.provider, model=self.model, input_tokens=self.input_tokens, output_tokens=self.output_tokens, cache_creation_input_tokens=self.cache_creation_input_tokens, cache_read_input_tokens=self.cache_read_input_tokens, request_type="chat", api_format=self.api_format, is_stream=True, response_time_ms=response_time_ms, status_code=self.status_code, # 使用实际的状态码 error_message=self.error_message, # 使用实际的错误消息 metadata={"stream": True, "content_length": len(self.accumulated_content)}, request_body=self.request_data if hasattr(self, "request_data") else None, request_headers=self.request_headers, provider_request_headers=self.provider_request_headers, response_headers=self.response_headers, response_body=response_body, request_id=self.request_id, # 传递 request_id # Provider 侧追踪信息（用于记录真实成本） provider_id=self.provider_id, provider_endpoint_id=self.provider_endpoint_id, provider_api_key_id=self.provider_api_key_id, # 请求状态 status=final_status, ) # 立即获取 total_cost_usd 的值，避免后续访问时对象已脱离会话 total_cost = 0.0 if usage_record: try: # 在 usage_record 仍在会话中时，立即获取所需属性 total_cost = usage_record.total_cost_usd or 0.0 except Exception as e: logger.warning(f"Failed to access total_cost_usd from usage_record: {e}") total_cost = 0.0 if db_for_usage and self.attempt_id: # RequestTrace 功能已移除，使用 RequestCandidate 表追踪 # 状态更新已在 RequestCandidateService 中完成 pass # 计算总响应时间（从请求开始到流结束） if self.request_start_time: total_response_time = int((self.end_time - self.request_start_time) * 1000) else: total_response_time = response_time_ms # 输出汇总日志（类似非流式请求的完成日志） # 根据状态码决定图标和日志级别 status_prefix = "[请求完成]" if self.status_code == 200 else "[请求失败]" # 根据费用大小选择合适的格式 if total_cost >= 0.01: cost_str = f"${total_cost:.4f}" elif total_cost > 0: cost_str = f"${total_cost:.6f}" else: cost_str = "$0" logger.info(f"{status_prefix} ID:{self.request_id} | {self.status_code} | 耗时:{total_response_time}ms | " f"Token:输入{self.input_tokens}/输出{self.output_tokens} | 费用:{cost_str}") # 记录提供商结果用于动态权重调整 # 记录提供商结果的健康监控已由 FallbackOrchestrator 自动处理 # 这里不再需要手动记录 except Exception as e: logger.exception(f"Failed to record stream usage: {e}") finally: if created_temp_session: try: db_for_usage.close() except Exception: pass class EnhancedStreamUsageTracker(StreamUsageTracker): """ 增强的流式用量跟踪器 支持更准确的token计算 """ def __init__( self, db: Session, user: User, api_key: ApiKey, provider: str, model: str, request_headers: Optional[Dict[str, Any]] = None, provider_request_headers: Optional[Dict[str, Any]] = None, request_id: Optional[str] = None, start_time: Optional[float] = None, attempt_id: Optional[str] = None, # Provider 侧追踪信息（用于记录真实成本） provider_id: Optional[str] = None, provider_endpoint_id: Optional[str] = None, provider_api_key_id: Optional[str] = None, # API 格式（用于选择正确的响应解析器） api_format: Optional[str] = None, ): super().__init__( db, user, api_key, provider, model, request_headers, provider_request_headers, request_id, start_time, attempt_id, provider_id, provider_endpoint_id, provider_api_key_id, api_format, ) # 用于更准确的token计算 self._init_tokenizer() # 继承父类的SSE解析缓冲区 # 这些已经在父类中初始化了 def _init_tokenizer(self): """初始化分词器（如果可用）""" try: # 尝试导入tiktoken用于更准确的token计算 import tiktoken # 根据模型选择合适的编码 if "gpt-4" in self.model: self.tokenizer = tiktoken.get_encoding("cl100k_base") elif "gpt-3.5" in self.model: self.tokenizer = tiktoken.get_encoding("cl100k_base") else: # Claude或其他模型，使用近似方法 self.tokenizer = None except ImportError: logger.debug("tiktoken not available, using estimation") self.tokenizer = None def count_tokens(self, text: str) -> int: """ 更准确地计算tokens Args: text: 文本内容 Returns: token数量 """ if self.tokenizer: try: return len(self.tokenizer.encode(text)) except Exception as e: logger.warning(f"Token encoding failed: {e}") # 回退到估算方法 # 中文字符通常是2个token，英文单词约1.3个token chinese_chars = len(re.findall(r"[\u4e00-\u9fff]", text)) english_words = len(re.findall(r"\b\w+\b", text)) estimated_tokens = chinese_chars * 2 + english_words * 1.3 return max(1, int(estimated_tokens)) def estimate_input_tokens(self, messages: list) -> int: """ 更准确地估算输入tokens Args: messages: 消息列表 Returns: 估算的token数 """ total_text = "" for msg in messages: if isinstance(msg, dict): content = msg.get("content", "") if isinstance(content, str): total_text += content + " " elif isinstance(content, list): for block in content: if hasattr(block, "text"): total_text += block.text + " " return self.count_tokens(total_text) async def track_stream( self, stream: AsyncIterator[str], request_data: Dict[str, Any], response_headers: Optional[Dict[str, Any]] = None, ) -> AsyncIterator[str]: """ 跟踪流式响应并更准确地计算用量 Args: stream: 原始流式响应 request_data: 请求数据 response_headers: 实际的响应头 Yields: 流式响应块 """ import time self.start_time = time.time() self.request_data = request_data # 保存请求数据 # 保存响应头（如果没有提供，使用空字典而不是默认值） # 这样可以确保记录的是实际的响应头，而不是构造的默认值 self.response_headers = response_headers if response_headers is not None else {} # 更准确地估算输入tokens messages = request_data.get("messages", []) self.input_tokens = self.estimate_input_tokens(messages) logger.debug(f"ID:{self.request_id} | 开始跟踪流式响应(Enhanced) | 估算输入tokens:{self.input_tokens}") chunk_count = 0 try: async for chunk in stream: chunk_count += 1 # 保存原始字节流（用于错误诊断） self.raw_chunks.append(chunk) # 返回原始块给客户端 yield chunk # 解析块以提取内容和使用信息（chunk是原始字节） content, usage = self.parse_stream_chunk(chunk) if content: self.accumulated_content += content # 使用更准确的方法计算输出tokens self.output_tokens = self.count_tokens(self.accumulated_content) if usage: # 如果响应中包含准确的usage信息，优先使用 if "input_tokens" in usage: self.input_tokens = usage["input_tokens"] if "output_tokens" in usage: self.output_tokens = usage["output_tokens"] if "cache_creation_input_tokens" in usage: self.cache_creation_input_tokens = usage["cache_creation_input_tokens"] if "cache_read_input_tokens" in usage: self.cache_read_input_tokens = usage["cache_read_input_tokens"] finally: # 流结束后记录使用量 self.end_time = time.time() logger.debug(f"ID:{self.request_id} | 流式响应结束 | 共处理{chunk_count}个chunks | " f"累积内容长度:{len(self.accumulated_content)} | 输出tokens:{self.output_tokens}") # 检查是否收到了有效数据 # 情况1: 收到了原始数据但无法解析为有效的SSE JSON if chunk_count > 0 and not self.response_chunks: error_msg = f"流式响应完成但未解析到任何有效数据（收到 {chunk_count} 个原始chunk但无法解析）" logger.error(f"ID:{self.request_id} | {error_msg}") # 设置错误状态，避免被记录为成功 self.set_error_status(502, error_msg) # 抛出异常让 FallbackOrchestrator 捕获并触发故障转移 raise EmptyStreamException( provider_name=self.provider, chunk_count=chunk_count, request_metadata=None, ) # 情况2: 流式响应完成但没有收到完整的消息（没有 message_stop 事件） # 这种情况通常发生在服务器重启或连接中断时 if not self.stream_stats.has_completion and not self.response_chunks: error_msg = "流式响应中断：未收到任何有效数据（可能是连接中断或服务重启）" logger.warning(f"ID:{self.request_id} | {error_msg}") self.set_error_status(503, error_msg) # 确保日志一定会输出，即使记录usage失败 try: await self._record_usage() except Exception as e: # 如果记录失败，至少输出基本的汇总日志 logger.exception(f"Failed to record stream usage for request {self.request_id}: {e}") # 尝试输出基本的汇总日志，使用多层防护 try: # 计算响应时间，使用多层后备机制 try: if self.request_start_time and self.end_time: total_response_time = int( (self.end_time - self.request_start_time) * 1000 ) elif self.start_time and self.end_time: total_response_time = int((self.end_time - self.start_time) * 1000) else: total_response_time = 0 except Exception: total_response_time = 0 # 安全地输出汇总日志 logger.info(f"[请求完成] ID:{self.request_id or 'unknown'} | 200 | 耗时:{total_response_time}ms | " f"Token:输入{self.input_tokens}/输出{self.output_tokens} | 费用:未知(记录失败)") except Exception as log_error: # 最后的防线：输出最简单的完成标记 logger.error(f"Failed to output summary log: {log_error}") try: logger.info(f"[请求完成] ID:{self.request_id or 'unknown'} | 记录失败但流已完成") except Exception: # 如果连最简单的日志都失败了，放弃 pass # 导出便捷函数 def create_stream_tracker( db: Session, user: User, api_key: ApiKey, provider: str, model: str, enhanced: bool = True, request_headers: Optional[Dict[str, Any]] = None, provider_request_headers: Optional[Dict[str, Any]] = None, request_id: Optional[str] = None, start_time: Optional[float] = None, attempt_id: Optional[str] = None, # Provider 侧追踪信息（用于记录真实成本） provider_id: Optional[str] = None, provider_endpoint_id: Optional[str] = None, provider_api_key_id: Optional[str] = None, # API 格式（用于选择正确的响应解析器） api_format: Optional[str] = None, ) -> StreamUsageTracker: """ 创建流式用量跟踪器 Args: db: 数据库会话 user: 用户对象 api_key: API密钥对象 provider: 提供商名称 model: 模型名称 enhanced: 是否使用增强版跟踪器 request_headers: 实际的请求头 provider_request_headers: 向提供商发送的请求头 request_id: 请求ID（用于日志关联） start_time: 请求开始时间（用于计算总响应时间） attempt_id: RequestTrace 请求尝试ID（可选） provider_id: Provider ID（用于记录真实成本） provider_endpoint_id: Endpoint ID（用于记录真实成本） provider_api_key_id: API Key ID（用于记录真实成本） api_format: API 格式（CLAUDE, CLAUDE_CLI, OPENAI, OPENAI_CLI） Returns: 流式用量跟踪器实例 """ if enhanced: return EnhancedStreamUsageTracker( db, user, api_key, provider, model, request_headers, provider_request_headers, request_id, start_time, attempt_id, provider_id, provider_endpoint_id, provider_api_key_id, api_format, ) else: return StreamUsageTracker( db, user, api_key, provider, model, request_headers, provider_request_headers, request_id, start_time, attempt_id, provider_id, provider_endpoint_id, provider_api_key_id, api_format, )