Aether/_deprecated_py_src/api/handlers/gemini/stream_parser.py

"""
Gemini 流解析器（SSE + JSON-array 兼容）

Gemini streamGenerateContent 常见两种返回（与上游/代理实现有关）：
1) `?alt=sse`：SSE（`data: {GenerateContentResponse}`）
2) 默认：JSON-array / JSON-chunks（`[{...},{...},...]`，可能跨 chunk/跨行）

本解析器提供：
- parse_line(): 适用于 SSE data 行或逐行 JSON 对象
- parse_chunk(): 适用于 JSON-array/chunks（可跨 chunk 拼接）

参考:
- https://ai.google.dev/gemini-api/docs/text-generation?lang=python#generate-a-text-stream
- https://generativelanguage.googleapis.com/$discovery/rest?version=v1beta
"""

import json
from typing import Any


class GeminiStreamParser:
    """
    Gemini 流解析器

    解析 Gemini streamGenerateContent API 的响应流。

    Gemini 流式响应特点:
    - 每个事件块本质上都是一个 GenerateContentResponse JSON 对象（包含 candidates、usageMetadata 等）
    - 结束判定以 `candidates[].finishReason` 为准（存在且不为 FINISH_REASON_UNSPECIFIED）
    """

    # finishReason（官方枚举值很多，见 discovery；这里仅保留一个明确的“未结束”哨兵）
    FINISH_REASON_UNSPECIFIED = "FINISH_REASON_UNSPECIFIED"

    def __init__(self) -> None:
        self._buffer = ""
        self._in_array = False
        self._brace_depth = 0

    def reset(self) -> None:
        """重置解析器状态"""
        self._buffer = ""
        self._in_array = False
        self._brace_depth = 0

    def parse_chunk(self, chunk: bytes | str) -> list[dict[str, Any]]:
        """
        解析流式数据块

        Args:
            chunk: 原始数据（bytes 或 str）

        Returns:
            解析后的事件列表
        """
        if isinstance(chunk, bytes):
            text = chunk.decode("utf-8")
        else:
            text = chunk

        events: list[dict[str, Any]] = []

        for char in text:
            if char == "[" and not self._in_array:
                self._in_array = True
                continue

            if char == "]" and self._in_array and self._brace_depth == 0:
                # 数组结束
                self._in_array = False
                if self._buffer.strip():
                    try:
                        obj = json.loads(self._buffer.strip().rstrip(","))
                        events.append(obj)
                    except json.JSONDecodeError:
                        pass
                self._buffer = ""
                continue

            if self._in_array:
                if char == "{":
                    self._brace_depth += 1
                elif char == "}":
                    self._brace_depth -= 1

                self._buffer += char

                # 当 brace_depth 回到 0 时，说明一个完整的 JSON 对象结束
                if self._brace_depth == 0 and self._buffer.strip():
                    try:
                        obj = json.loads(self._buffer.strip().rstrip(","))
                        events.append(obj)
                        self._buffer = ""
                    except json.JSONDecodeError:
                        # 可能还不完整，继续累积
                        pass

        return events

    def parse_line(self, line: str) -> dict[str, Any] | None:
        """
        解析单行 JSON 数据

        Args:
            line: JSON 数据行

        Returns:
            解析后的事件字典，如果无法解析返回 None
        """
        if not line or line.strip() in ["[", "]", ","]:
            return None

        try:
            result = json.loads(line.strip().rstrip(","))
            if isinstance(result, dict):
                return result
            return None
        except json.JSONDecodeError:
            return None

    def is_done_event(self, event: dict[str, Any]) -> bool:
        """
        判断是否为结束事件

        Args:
            event: 事件字典

        Returns:
            True 如果是结束事件
        """
        candidates = event.get("candidates", [])
        if not candidates:
            return False

        for candidate in candidates:
            finish_reason = candidate.get("finishReason")
            if not finish_reason:
                continue
            # 只要出现非 UNSPECIFIED 的 finishReason，通常表示该 candidate 已结束。
            # 例如：STOP/MAX_TOKENS/SAFETY/RECITATION/MALFORMED_FUNCTION_CALL/...（枚举持续演进）
            if str(finish_reason) != self.FINISH_REASON_UNSPECIFIED:
                return True

        return False

    def is_error_event(self, event: dict[str, Any]) -> bool:
        """
        判断是否为错误事件

        检测多种 Gemini 错误格式:
        1. 顶层 error: {"error": {...}}
        2. chunks 内嵌套 error: {"chunks": [{"error": {...}}]}
        3. candidates 内的错误状态

        Args:
            event: 事件字典

        Returns:
            True 如果是错误事件
        """
        # 顶层 error
        if "error" in event:
            return True

        # chunks 内嵌套 error (某些 Gemini 响应格式)
        chunks = event.get("chunks", [])
        if chunks and isinstance(chunks, list):
            for chunk in chunks:
                if isinstance(chunk, dict) and "error" in chunk:
                    return True

        return False

    def extract_error_info(self, event: dict[str, Any]) -> dict[str, Any] | None:
        """
        从事件中提取错误信息

        Args:
            event: 事件字典

        Returns:
            错误信息字典 {"code": int, "message": str, "status": str}，无错误返回 None
        """
        # 顶层 error
        if "error" in event:
            error = event["error"]
            if isinstance(error, dict):
                return {
                    "code": error.get("code"),
                    "message": error.get("message", str(error)),
                    "status": error.get("status"),
                }
            return {"code": None, "message": str(error), "status": None}

        # chunks 内嵌套 error
        chunks = event.get("chunks", [])
        if chunks and isinstance(chunks, list):
            for chunk in chunks:
                if isinstance(chunk, dict) and "error" in chunk:
                    error = chunk["error"]
                    if isinstance(error, dict):
                        return {
                            "code": error.get("code"),
                            "message": error.get("message", str(error)),
                            "status": error.get("status"),
                        }
                    return {"code": None, "message": str(error), "status": None}

        return None

    def get_finish_reason(self, event: dict[str, Any]) -> str | None:
        """
        获取结束原因

        Args:
            event: 事件字典

        Returns:
            结束原因字符串
        """
        candidates = event.get("candidates", [])
        if candidates:
            reason = candidates[0].get("finishReason")
            return str(reason) if reason is not None else None
        return None

    def extract_text_delta(self, event: dict[str, Any]) -> str | None:
        """
        从响应中提取文本内容

        Args:
            event: 事件字典

        Returns:
            文本内容，如果没有文本返回 None
        """
        candidates = event.get("candidates", [])
        if not candidates:
            return None

        content = candidates[0].get("content", {})
        parts = content.get("parts", [])

        text_parts = []
        for part in parts:
            if "text" in part:
                text_parts.append(part["text"])

        return "".join(text_parts) if text_parts else None

    def extract_usage(self, event: dict[str, Any]) -> dict[str, int] | None:
        """
        从事件中提取 token 使用量

        这是 Gemini token 提取的单一实现源，其他地方都应该调用此方法。

        Args:
            event: 事件字典（包含 usageMetadata）

        Returns:
            使用量字典，如果没有完整的使用量信息返回 None

        注意:
            - 只有当 totalTokenCount 存在时才提取（确保是完整的 usage 数据）
            - 输出 token = thoughtsTokenCount + candidatesTokenCount
        """
        usage_metadata = event.get("usageMetadata", {})
        if not usage_metadata or "totalTokenCount" not in usage_metadata:
            return None

        # 输出 token = thoughtsTokenCount + candidatesTokenCount
        thoughts_tokens = usage_metadata.get("thoughtsTokenCount", 0)
        candidates_tokens = usage_metadata.get("candidatesTokenCount", 0)
        output_tokens = thoughts_tokens + candidates_tokens

        return {
            "input_tokens": usage_metadata.get("promptTokenCount", 0),
            "output_tokens": output_tokens,
            "total_tokens": usage_metadata.get("totalTokenCount", 0),
            "cached_tokens": usage_metadata.get("cachedContentTokenCount", 0),
        }

    def extract_model_version(self, event: dict[str, Any]) -> str | None:
        """
        从响应中提取模型版本

        Args:
            event: 事件字典

        Returns:
            模型版本，如果没有返回 None
        """
        version = event.get("modelVersion")
        return str(version) if version is not None else None

    def extract_safety_ratings(self, event: dict[str, Any]) -> list[dict[str, Any]] | None:
        """
        从响应中提取安全评级

        Args:
            event: 事件字典

        Returns:
            安全评级列表，如果没有返回 None
        """
        candidates = event.get("candidates", [])
        if not candidates:
            return None

        ratings = candidates[0].get("safetyRatings")
        if isinstance(ratings, list):
            return ratings
        return None


__all__ = ["GeminiStreamParser"]