ZNJJ-api-server/src/backends/fastgpt.py

"""FastGPT implementation of the backend-neutral chat contract."""

import json
from collections.abc import AsyncIterator
from typing import Any

from fastgpt_client import AsyncChatClient, aiter_stream_events
from fastgpt_client.exceptions import APIError, AuthenticationError, RateLimitError

from .chat import (
    ChatBackendAPIError,
    ChatBackendAuthenticationError,
    ChatBackendRateLimitError,
    ChatInput,
    ChatResult,
    FormUpdate,
    TextDelta,
)


FORM_EXTRACT_MODULE_NAME = "文本内容提取事故信息"


def _parse_json_value(value: Any) -> Any:
    parsed = value
    for _ in range(3):
        if not isinstance(parsed, str):
            return parsed
        parsed = parsed.strip()
        if not parsed:
            return {}
        try:
            parsed = json.loads(parsed)
        except json.JSONDecodeError:
            return parsed
    return parsed


def _extract_form_update(nodes: Any) -> Any:
    if not isinstance(nodes, list):
        return {}

    for node in nodes:
        if not isinstance(node, dict):
            continue
        if node.get("moduleName") != FORM_EXTRACT_MODULE_NAME:
            continue

        extract_result = node.get("extractResult", {})
        if not isinstance(extract_result, dict):
            return {}

        form_update = extract_result.get("formUpdate", "")
        return _parse_json_value(form_update) if form_update else {}

    return {}


def _extract_content(data: dict[str, Any]) -> str:
    try:
        content = data["choices"][0]["message"]["content"]
    except (KeyError, IndexError, TypeError) as exc:
        raise ChatBackendAPIError("FastGPT response is missing message content") from exc

    if isinstance(content, list):
        try:
            content = content[0]["text"]["content"]
        except (KeyError, IndexError, TypeError) as exc:
            raise ChatBackendAPIError(
                "FastGPT response contains invalid list content"
            ) from exc

    if not isinstance(content, str):
        raise ChatBackendAPIError("FastGPT message content is not text")
    return content


class FastGPTBackend:
    """Translate FastGPT SDK calls and events into the neutral chat contract."""

    def __init__(self, client: AsyncChatClient):
        self._client = client

    async def stream(self, chat_input: ChatInput) -> AsyncIterator[TextDelta | FormUpdate]:
        try:
            response = await self._client.create_chat_completion(
                messages=[{"role": "user", "content": chat_input.text}],
                chatId=chat_input.session_id,
                stream=True,
                detail=True,
                variables={"needFormUpdate": chat_input.need_form_update},
            )

            async for event in aiter_stream_events(response):
                if event.kind == "flowResponses":
                    form_update = _extract_form_update(event.data)
                    if form_update:
                        yield FormUpdate(form_update)
                    continue

                if event.kind not in {"answer", "fastAnswer", "data"}:
                    continue
                if not isinstance(event.data, dict):
                    continue

                try:
                    content = event.data["choices"][0]["delta"].get("content", "")
                except (KeyError, IndexError, TypeError, AttributeError):
                    content = ""
                if content:
                    yield TextDelta(content)
        except AuthenticationError as exc:
            raise ChatBackendAuthenticationError(str(exc)) from exc
        except RateLimitError as exc:
            raise ChatBackendRateLimitError(str(exc)) from exc
        except APIError as exc:
            raise ChatBackendAPIError(str(exc)) from exc

    async def complete(self, chat_input: ChatInput) -> ChatResult:
        try:
            response = await self._client.create_chat_completion(
                messages=[{"role": "user", "content": chat_input.text}],
                chatId=chat_input.session_id,
                stream=False,
                detail=True,
                variables={"needFormUpdate": chat_input.need_form_update},
            )
            response.raise_for_status()
            data = response.json()
        except AuthenticationError as exc:
            raise ChatBackendAuthenticationError(str(exc)) from exc
        except RateLimitError as exc:
            raise ChatBackendRateLimitError(str(exc)) from exc
        except APIError as exc:
            raise ChatBackendAPIError(str(exc)) from exc

        if not isinstance(data, dict):
            raise ChatBackendAPIError("FastGPT response body is not an object")

        status_code = data.get("newVariables", {}).get("status_code")
        if status_code is not None:
            status_code = str(status_code)

        return ChatResult(
            content=_extract_content(data),
            status_code=status_code,
            form_update=_extract_form_update(data.get("responseData", [])),
        )