src/stream/processor.py

"""
流式响应处理器（简化版）

实现"假流式"：收集所有上游数据后，解析并聚合为一个包含完整上下文（思考过程+最终内容）的响应，一次性发送给下游。
"""

import json
import time
from typing import Any, cast
import collections.abc

from src.core.errors import (
    VcoreError,
    EmptyResponseError,
    InternalError,
    NotFoundError,
    InvalidArgumentError,
    RateLimitError, 
)
from src.utils.logger import get_logger
from src.utils.token_counter import calculate_usage_metadata
from .parser import parse_upstream_data


# 初始化日志
logger = get_logger(__name__)


class StreamProcessor:
    """
    简化的流式响应处理器 (v3 - 假流式 + 思考过程)
    
    职责:
    1. 收集上游所有流式数据
    2. 使用 parser 解析聚合后的数据
    3. 构建 Gemini 格式的 SSE 响应
    """
    
    def __init__(self):
        logger.debug("初始化流处理器")
        
        # 状态追踪
        self._actual_content_sent = False
        self._request_context: dict[str, Any] = {}
    
    def has_actual_content_sent(self) -> bool:
        """检查是否已发送实际文本内容"""
        return self._actual_content_sent
    
    def set_request_context(self, downstream_payload: dict[str, Any], upstream_payload: dict[str, Any]):
        """设置请求上下文"""
        logger.debug("设置流处理器请求上下文")
        self._request_context = {
            'downstream_payload': downstream_payload,
            'upstream_payload': upstream_payload
        }
    
    def _create_gemini_chunk(
        self,
        parts: list[dict[str, Any]],
        finish_reason: str | None,
        safety_ratings: list[dict[str, Any]],
        citation_metadata: dict[str, Any],
        grounding_metadata: dict[str, Any],
        candidate_index: int,
        prompt_feedback: dict[str, Any],
        usage_metadata: dict[str, Any],
        finish_message: str | None = None,
        token_count: int | None = None,
        avg_logprobs: float | None = None,
        logprobs_result: dict[str, Any] | None = None,
        create_time: str | None = None,
        model_version: str | None = None,
        response_id: str | None = None,
        model_status: dict[str, Any] | None = None,
    ) -> str:
        """根据聚合后的内容，创建一个包含完整上下文的Gemini格式SSE事件。"""
        candidate: dict[str, Any] = {
            "index": candidate_index
        }
        if finish_reason and isinstance(finish_reason, str):
            candidate["finishReason"] = finish_reason.upper()
        if finish_message:
            candidate["finishMessage"] = finish_message
        
        if parts:
            candidate["content"] = {
                "parts": parts,
                "role": "model"
            }
        
        # 只添加实际存在的字段
        if safety_ratings:
            candidate["safetyRatings"] = safety_ratings
        if citation_metadata:
            candidate["citationMetadata"] = citation_metadata
        if grounding_metadata:
            candidate["groundingMetadata"] = grounding_metadata
        if token_count is not None:
            candidate["tokenCount"] = token_count
        if avg_logprobs is not None:
            candidate["avgLogprobs"] = avg_logprobs
        if logprobs_result:
            candidate["logprobsResult"] = logprobs_result
            
        # 根据 doc.md，确保 candidate 其他字段正确
        
        chunk: dict[str, Any] = {"candidates": [candidate]}
        
        # 按照 doc.md 定义
        if prompt_feedback:
            chunk["promptFeedback"] = prompt_feedback
        if usage_metadata:
            chunk["usageMetadata"] = usage_metadata
        if create_time:
            chunk["createTime"] = create_time
        if model_version:
            chunk["modelVersion"] = model_version
        if response_id:
            chunk["responseId"] = response_id
        if model_status:
            chunk["modelStatus"] = model_status
             
        return "data: " + json.dumps(chunk, ensure_ascii=False, separators=(',', ':')) + "\n\n"

    async def process_stream(
        self,
        response_iterator: collections.abc.AsyncIterator[str],
        model: str = "vcore-ai-proxy"
    ) -> collections.abc.AsyncGenerator[str, None]:
        """
        处理流式响应（假流式 v3）。
        """
        # 移除重复日志，只在开始时记录一次
        # logger.info(f"开始处理流式响应: 模型={model}")
        start_time = time.time()
        raw_chunks: list[str] = []
        
        try:
            chunk_count = 0
            logger.debug("开始收集上游数据块")
            async for chunk in response_iterator:
                chunk_count += 1
                raw_chunks.append(chunk)
            
            logger.debug(f"收集完成，共 {chunk_count} 个数据块")
            raw_data = '\n'.join(raw_chunks)
            
            # 记录完整的原始上游响应
            try:
                parsed_data = json.loads(raw_data)
                logger.debug_json("完整原始上游响应", parsed_data)
            except json.JSONDecodeError:
                logger.debug_large("完整原始上游响应", raw_data)
            
            if not raw_data:
                logger.error("上游返回空数据")
                # 触发错误备份
                from src.utils.error_logger import save_error_snapshot
                save_error_snapshot(
                    downstream_payload=self._request_context.get('downstream_payload', {}),
                    upstream_payload=self._request_context.get('upstream_payload', {}),
                    upstream_response="[EMPTY RESPONSE]",
                    error_type="empty_response"
                )
                raise EmptyResponseError("Upstream returned no data")

            # 使用独立解析函数
            result = parse_upstream_data(raw_data)
            
            # 关键修复：在这里处理解析出的上游错误
            if result["has_error"] and not result["parts"]:
                error_msg = result["error_message"]
                error_obj = result.get("error_obj")
                
                # 排除特定的、不需要备份的错误
                is_auth_error = "Failed to verify action" in error_msg or "The caller does not have permission" in error_msg
                is_rate_limit = isinstance(error_obj, RateLimitError) or "resource has been exhausted" in error_msg.lower() or "quota" in error_msg.lower()
                
                if not is_auth_error and not is_rate_limit:
                    logger.error(f"API 错误且无内容: {error_msg}")
                    
                    # 确定错误类型用于备份目录名
                    error_type = "api_error"
                    if error_obj:
                        # 如果有解析出的错误对象，使用其类名或 code
                        error_type = f"upstream_{error_obj.code}_{type(error_obj).__name__}"
                    
                    # 触发错误备份
                    from src.utils.error_logger import save_error_snapshot
                    save_error_snapshot(
                        downstream_payload=self._request_context.get('downstream_payload', {}),
                        upstream_payload=self._request_context.get('upstream_payload', {}),
                        upstream_response=raw_data,
                        error_type=error_type
                    )
                
                # 如果 parser 已经解析出了错误对象，直接抛出
                if error_obj:
                    raise error_obj

                # 降级处理：基于上游错误消息抛出错误
                error_msg_lower = error_msg.lower()
                if "not found" in error_msg_lower:
                    raise NotFoundError(message=error_msg)
                elif is_rate_limit:
                    raise RateLimitError(message=error_msg)
                elif is_auth_error:
                    from src.core.errors import AuthenticationError
                    raise AuthenticationError(
                        message=f"Authentication/Recaptcha failed: {error_msg}",
                        details={"upstream_response": error_msg},
                        upstream_response=error_msg
                    )
                else:
                    raise InvalidArgumentError(message=error_msg)
            
            finish_reason = result.get("finish_reason") or "STOP"
            
            if not result["parts"] and finish_reason == "STOP" and not result["has_error"]:
                 if not result.get("prompt_feedback"):
                    logger.error("上游返回空响应 (无 parts 且 finish_reason=STOP)")
                    # 触发错误备份
                    from src.utils.error_logger import save_error_snapshot
                    save_error_snapshot(
                        downstream_payload=self._request_context.get('downstream_payload', {}),
                        upstream_payload=self._request_context.get('upstream_payload', {}),
                        upstream_response=raw_data,
                        error_type="stop_no_content"
                    )
                    # 快照将由 VcoreAIClient 统一保存
                    raise EmptyResponseError("Upstream returned empty response (STOP with no content/metadata)")
            
            # 计算 usage metadata
            usage_metadata = result.get("usage_metadata", {})
            if not usage_metadata and self._request_context:
                try:
                    downstream_payload = self._request_context.get('downstream_payload', {})
                    
                    # 从请求上下文中提取输入内容
                    prompt_contents: list[dict[str, Any]] = []
                    if 'gemini_payload' in downstream_payload:
                        gemini_payload = downstream_payload['gemini_payload']
                        if isinstance(gemini_payload, dict) and 'contents' in gemini_payload:
                            prompt_contents = cast(list[dict[str, Any]], gemini_payload['contents'])
                    
                    # 计算 token 使用情况
                    usage_metadata = await calculate_usage_metadata(
                        prompt_contents=prompt_contents,
                        response_parts=result["parts"],
                        request_context=self._request_context
                    )
                except Exception as e:
                    logger.warning(f"计算 usage metadata 失败: {e}")
                    usage_metadata = {}
            
            final_chunk = self._create_gemini_chunk(
                parts=result["parts"],
                finish_reason=result.get("finish_reason"),
                safety_ratings=result.get("safety_ratings", []),
                citation_metadata=result.get("citation_metadata", {}),
                grounding_metadata=result.get("grounding_metadata", {}),
                candidate_index=result.get("candidate_index", 0),
                prompt_feedback=result.get("prompt_feedback", {}),
                usage_metadata=usage_metadata,
                finish_message=result.get("finish_message"),
                token_count=result.get("token_count"),
                avg_logprobs=result.get("avg_logprobs"),
                logprobs_result=result.get("logprobs_result"),
                create_time=result.get("create_time"),
                model_version=result.get("model_version"),
                response_id=result.get("response_id"),
                model_status=result.get("model_status")
            )
            
            process_time = time.time() - start_time
            logger.success(f"流式响应处理完成: 耗时={process_time:.3f}s, 完成原因={result.get('finish_reason', 'UNKNOWN')}")
            
            yield final_chunk
            self._actual_content_sent = True
            
        except VcoreError as e:
            if "Failed to verify action" in e.message or "The caller does not have permission" in e.message:
                from src.core.errors import AuthenticationError
                if not isinstance(e, AuthenticationError):
                    raise AuthenticationError(
                        message=f"Stream contained Authentication error: {e.message}",
                        details={"upstream_response": e.upstream_response or e.message},
                        upstream_response=e.upstream_response or e.message
                    )
            else:
                logger.error(f"流处理 Vcore 错误: {e.message}")
            raise
        except Exception as e:
            logger.error(f"流处理未知错误: {e}")
            # 触发内部错误备份
            from src.utils.error_logger import save_error_snapshot
            save_error_snapshot(
                downstream_payload=self._request_context.get('downstream_payload', {}),
                upstream_payload=self._request_context.get('upstream_payload', {}),
                upstream_response=str(e),
                error_type="internal_exception"
            )
            raise InternalError(message=f"Unknown stream processing error: {e}")


def get_stream_processor() -> StreamProcessor:
    """创建流处理器实例"""
    return StreamProcessor()