| """Module for streaming.""" |
|
|
| from __future__ import annotations |
|
|
| import logging |
| from collections.abc import AsyncIterator, Callable, Iterator |
| from typing import Any |
|
|
| from .backends import ILLMBackend, LLMStreamChunk |
| from .exceptions import ( |
| LLMAPIError, |
| LLMAuthenticationError, |
| LLMBackendUnavailableError, |
| LLMNetworkError, |
| LLMTimeoutError, |
| ) |
|
|
| logger = logging.getLogger("apps.core.llm.streaming") |
| _RETRIABLE_ERRORS = (LLMTimeoutError, LLMNetworkError, LLMAPIError) |
|
|
|
|
| def _resolve_backends( |
| get_backend: Callable[[str], ILLMBackend], |
| get_backends_by_priority: Callable[[], list[tuple[str, ILLMBackend]]], |
| backend: str | None, |
| fallback: bool, |
| ) -> list[tuple[str, ILLMBackend]]: |
| """解析要尝试的后端列表""" |
| if backend: |
| backends_to_try: list[tuple[str, ILLMBackend]] = [(backend, get_backend(backend))] |
| if fallback: |
| for name, b in get_backends_by_priority(): |
| if name != backend: |
| backends_to_try.append((name, b)) |
| return backends_to_try |
| return get_backends_by_priority() |
|
|
|
|
| def _handle_stream_error( |
| name: str, e: Exception, fallback: bool, errors: list[Any], log_prefix: str = "stream" |
| ) -> None: |
| """处理流式调用错误""" |
| errors.append((name, e)) |
| if isinstance(e, _RETRIABLE_ERRORS): |
| logger.warning( |
| f"后端 {log_prefix} 调用失败,尝试下一个", |
| extra={"backend": name, "error": str(e), "error_type": type(e).__name__}, |
| ) |
| if not fallback: |
| raise |
| else: |
| logger.warning( |
| f"后端 {log_prefix} 调用发生未知错误", |
| extra={"backend": name, "error": str(e), "error_type": type(e).__name__}, |
| ) |
| if not fallback: |
| raise LLMAPIError(message=f"调用后端 {name} 时发生错误: {e!s}", errors={"detail": str(e)}) from e |
|
|
|
|
| def _build_stream_kwargs( |
| messages: Any, |
| model: Any, |
| temperature: Any, |
| max_tokens: Any, |
| extra_kwargs: dict[str, Any] | None = None, |
| ) -> dict[str, Any]: |
| kwargs = {"messages": messages, "model": model, "temperature": temperature, "max_tokens": max_tokens} |
| if extra_kwargs: |
| kwargs.update(extra_kwargs) |
| return kwargs |
|
|
|
|
| def stream_with_fallback( |
| *, |
| get_backend: Callable[[str], ILLMBackend], |
| get_backends_by_priority: Callable[[], list[tuple[str, ILLMBackend]]], |
| backend: str | None, |
| fallback: bool, |
| messages: list[dict[str, str]], |
| model: str | None, |
| temperature: float, |
| max_tokens: int | None, |
| **kwargs: Any, |
| ) -> Iterator[LLMStreamChunk]: |
| from .fallback_policy import _diagnose_unavailable |
|
|
| kwargs = _build_stream_kwargs(messages, model, temperature, max_tokens, kwargs) |
| if backend and (not fallback): |
| yield from get_backend(backend).stream(**kwargs) |
| return |
| backends_to_try = _resolve_backends(get_backend, get_backends_by_priority, backend, fallback) |
| errors: list[tuple[str, Exception]] = [] |
| skipped: list[tuple[str, str]] = [] |
| for name, backend_instance in backends_to_try: |
| if not backend_instance.is_available(): |
| reason = _diagnose_unavailable(name, backend_instance) |
| logger.warning("后端不可用,跳过", extra={"backend": name, "reason": reason}) |
| skipped.append((name, reason)) |
| continue |
| it = None |
| try: |
| it = iter(backend_instance.stream(**kwargs)) |
| first = next(it) |
| except StopIteration: |
| return |
| except LLMAuthenticationError: |
| raise |
| except Exception as e: |
| logger.exception("操作失败") |
| _handle_stream_error(name, e, fallback, errors, "stream") |
| continue |
| yield first |
| if it is not None: |
| yield from it |
| return |
| attempts_detail = [(n, str(e)) for n, e in errors] |
| attempts_detail.extend([(n, reason) for n, reason in skipped]) |
| raise LLMBackendUnavailableError( |
| message="所有 LLM 后端均不可用", |
| errors={"attempts": attempts_detail, "skipped": skipped}, |
| ) |
|
|
|
|
| async def astream_with_fallback( |
| *, |
| get_backend: Callable[[str], ILLMBackend], |
| get_backends_by_priority: Callable[[], list[tuple[str, ILLMBackend]]], |
| backend: str | None, |
| fallback: bool, |
| messages: list[dict[str, str]], |
| model: str | None, |
| temperature: float, |
| max_tokens: int | None, |
| **kwargs: Any, |
| ) -> AsyncIterator[LLMStreamChunk]: |
| from .fallback_policy import _diagnose_unavailable |
|
|
| kwargs = _build_stream_kwargs(messages, model, temperature, max_tokens, kwargs) |
| if backend and (not fallback): |
| async for chunk in get_backend(backend).astream(**kwargs): |
| yield chunk |
| return |
| backends_to_try = _resolve_backends(get_backend, get_backends_by_priority, backend, fallback) |
| errors: list[tuple[str, Exception]] = [] |
| skipped: list[tuple[str, str]] = [] |
| for name, backend_instance in backends_to_try: |
| if not backend_instance.is_available(): |
| reason = _diagnose_unavailable(name, backend_instance) |
| logger.warning("后端不可用,跳过", extra={"backend": name, "reason": reason}) |
| skipped.append((name, reason)) |
| continue |
| try: |
| ait = backend_instance.astream(**kwargs).__aiter__() |
| first = await ait.__anext__() |
| except StopAsyncIteration: |
| return |
| except LLMAuthenticationError: |
| raise |
| except Exception as e: |
| logger.exception("操作失败") |
| _handle_stream_error(name, e, fallback, errors, "astream") |
| continue |
| yield first |
| async for chunk in ait: |
| yield chunk |
| return |
| attempts_detail = [(n, str(e)) for n, e in errors] |
| attempts_detail.extend([(n, reason) for n, reason in skipped]) |
| raise LLMBackendUnavailableError( |
| message="所有 LLM 后端均不可用", |
| errors={"attempts": attempts_detail, "skipped": skipped}, |
| ) |
|
|
