LangBot/src/langbot/libs/deerflow_api/client.py

"""DeerFlow LangGraph HTTP API 客户端

参考 astrbot 的 deerflow_api_client 实现，使用 httpx 适配 LangBot 风格。
"""

from __future__ import annotations

import codecs
import json
import typing
from collections.abc import AsyncGenerator

import httpx

from langbot.pkg.utils import httpclient

from .errors import DeerFlowAPIError


SSE_MAX_BUFFER_CHARS = 1_048_576
SSE_MAX_TOTAL_BYTES = 16 * 1024 * 1024
ERROR_BODY_MAX_BYTES = 1024 * 1024


async def _read_error_body(response: httpx.Response) -> str:
    body = bytearray()
    async for chunk in response.aiter_bytes(8192):
        body.extend(chunk)
        if len(body) > ERROR_BODY_MAX_BYTES:
            raise DeerFlowAPIError(
                operation='read error response',
                body='response exceeds the runtime limit',
            )
    return body.decode('utf-8', errors='replace')


def _normalize_sse_newlines(text: str) -> str:
    """规范化 CRLF/CR 为 LF，确保 SSE 块分割稳定"""
    return text.replace('\r\n', '\n').replace('\r', '\n')


def _parse_sse_data_lines(data_lines: list[str]) -> typing.Any:
    raw_data = '\n'.join(data_lines)
    try:
        return json.loads(raw_data)
    except json.JSONDecodeError:
        # 某些 LangGraph 兼容服务端会在单个 SSE 事件中用多个 data 行
        # 发送多段 JSON 片段（例如 tuple payload）
        parsed_lines: list[typing.Any] = []
        can_parse_all = True
        for line in data_lines:
            line = line.strip()
            if not line:
                continue
            try:
                parsed_lines.append(json.loads(line))
            except json.JSONDecodeError:
                can_parse_all = False
                break
        if can_parse_all and parsed_lines:
            return parsed_lines[0] if len(parsed_lines) == 1 else parsed_lines
        return raw_data


def _parse_sse_block(block: str) -> dict[str, typing.Any] | None:
    if not block.strip():
        return None

    event_name = 'message'
    data_lines: list[str] = []
    for line in block.splitlines():
        if line.startswith('event:'):
            event_name = line[6:].strip()
        elif line.startswith('data:'):
            data_lines.append(line[5:].lstrip())

    if not data_lines:
        return None
    return {'event': event_name, 'data': _parse_sse_data_lines(data_lines)}


class AsyncDeerFlowClient:
    """DeerFlow LangGraph HTTP API 客户端"""

    api_base: str
    headers: dict[str, str]

    def __init__(
        self,
        api_base: str = 'http://127.0.0.1:2026',
        api_key: str = '',
        auth_header: str = '',
    ) -> None:
        self.api_base = api_base.rstrip('/')
        self.headers: dict[str, str] = {}
        if auth_header:
            self.headers['Authorization'] = auth_header
        elif api_key:
            self.headers['Authorization'] = f'Bearer {api_key}'

    async def create_thread(self, timeout: float = 20) -> dict[str, typing.Any]:
        """创建一个新的 LangGraph thread

        Returns:
            包含 thread_id 等信息的字典
        """
        url = f'{self.api_base}/api/langgraph/threads'
        payload = {'metadata': {}}

        async with httpx.AsyncClient(
            trust_env=True,
            timeout=timeout,
            event_hooks=httpclient.httpx_response_limit_hooks(),
        ) as client:
            response = await client.post(
                url,
                headers=self.headers,
                json=payload,
            )
            if response.status_code not in (200, 201):
                body = await httpclient.response_text(response)
                raise DeerFlowAPIError(
                    operation='create thread',
                    status=response.status_code,
                    body=body,
                    url=url,
                )
            return await httpclient.parse_json_response(response)

    async def delete_thread(self, thread_id: str, timeout: float = 20) -> None:
        """删除指定 thread"""
        url = f'{self.api_base}/api/threads/{thread_id}'

        async with httpx.AsyncClient(
            trust_env=True,
            timeout=timeout,
            event_hooks=httpclient.httpx_response_limit_hooks(),
        ) as client:
            response = await client.delete(url, headers=self.headers)
            if response.status_code not in (200, 202, 204, 404):
                body = await httpclient.response_text(response)
                raise DeerFlowAPIError(
                    operation='delete thread',
                    status=response.status_code,
                    body=body,
                    url=url,
                    thread_id=thread_id,
                )

    async def stream_run(
        self,
        thread_id: str,
        payload: dict[str, typing.Any],
        timeout: float = 120,
    ) -> AsyncGenerator[dict[str, typing.Any], None]:
        """运行一次 LangGraph stream 请求，逐事件 yield

        Yields:
            事件字典 {'event': event_name, 'data': parsed_data}
        """
        url = f'{self.api_base}/api/langgraph/threads/{thread_id}/runs/stream'

        # 流式请求使用单独的 read timeout 控制
        stream_timeout = httpx.Timeout(
            connect=min(timeout, 30),
            read=timeout,
            write=timeout,
            pool=timeout,
        )

        async with httpx.AsyncClient(
            trust_env=True,
            timeout=stream_timeout,
        ) as client:
            async with client.stream(
                'POST',
                url,
                headers={
                    **self.headers,
                    'Accept': 'text/event-stream',
                    'Content-Type': 'application/json',
                },
                json=payload,
            ) as resp:
                if resp.status_code != 200:
                    raise DeerFlowAPIError(
                        operation='runs/stream request',
                        status=resp.status_code,
                        body=await _read_error_body(resp),
                        url=url,
                        thread_id=thread_id,
                    )

                decoder = codecs.getincrementaldecoder('utf-8')('replace')
                buffer = ''
                total_bytes = 0

                async for chunk in resp.aiter_bytes(8192):
                    total_bytes += len(chunk)
                    if total_bytes > SSE_MAX_TOTAL_BYTES:
                        raise DeerFlowAPIError(
                            operation='runs/stream response',
                            body='response exceeds the runtime limit',
                            url=url,
                            thread_id=thread_id,
                        )
                    buffer += _normalize_sse_newlines(decoder.decode(chunk))

                    while '\n\n' in buffer:
                        block, buffer = buffer.split('\n\n', 1)
                        parsed = _parse_sse_block(block)
                        if parsed is not None:
                            yield parsed

                    if len(buffer) > SSE_MAX_BUFFER_CHARS:
                        # 缓冲区过大，强制 flush
                        parsed = _parse_sse_block(buffer)
                        if parsed is not None:
                            yield parsed
                        buffer = ''

                # flush 剩余内容
                buffer += _normalize_sse_newlines(decoder.decode(b'', final=True))
                while '\n\n' in buffer:
                    block, buffer = buffer.split('\n\n', 1)
                    parsed = _parse_sse_block(block)
                    if parsed is not None:
                        yield parsed
                if buffer.strip():
                    parsed = _parse_sse_block(buffer)
                    if parsed is not None:
                        yield parsed