Owner/src/owner_voice_pet/llm.py

from __future__ import annotations

import json
import socket
import urllib.error
import urllib.request
from collections.abc import Callable, Iterable, Sequence
from typing import Any

from .config import AppConfig
from .models import ErrorCode, Message, ProviderError, ReplyDelta


class MockLlmProvider:
    def __init__(self, chunks: list[str] | None = None, fail: ProviderError | None = None) -> None:
        self.chunks = chunks or ["好的，我听到了。"]
        self.fail = fail
        self.calls: list[list[Message]] = []

    def stream_reply(self, messages: Sequence[Message]) -> Iterable[ReplyDelta]:
        self.calls.append(list(messages))
        if self.fail:
            raise self.fail
        for chunk in self.chunks:
            yield ReplyDelta(chunk, is_sentence_boundary=_ends_sentence(chunk))
        yield ReplyDelta("", finish_reason="stop")


class OpenAICompatibleLlmProvider:
    def __init__(
        self,
        config: AppConfig,
        timeout_s: float = 20.0,
        urlopen: Callable[..., Any] | None = None,
    ) -> None:
        self.config = config
        self.timeout_s = timeout_s
        self.urlopen = urlopen or urllib.request.urlopen

    def stream_reply(self, messages: Sequence[Message]) -> Iterable[ReplyDelta]:
        self.config.require_llm_credentials()
        if self.config.llm_api_style == "responses":
            yield from self._responses(messages)
        else:
            yield from self._chat_completions(messages)

    def _chat_completions(self, messages: Sequence[Message]) -> Iterable[ReplyDelta]:
        payload = {
            "model": self.config.llm_model,
            "messages": [{"role": m.role, "content": m.content} for m in messages],
            "stream": self.config.llm_stream,
        }
        yield from self._post_stream(
            self.config.api_url("/v1/chat/completions"),
            payload,
            parser=_parse_chat_completion_sse,
        )

    def _responses(self, messages: Sequence[Message]) -> Iterable[ReplyDelta]:
        payload = {
            "model": self.config.llm_model,
            "input": [{"role": m.role, "content": m.content} for m in messages],
            "stream": self.config.llm_stream,
        }
        yield from self._post_stream(
            self.config.api_url("/v1/responses"),
            payload,
            parser=_parse_responses_sse,
        )

    def _post_stream(
        self,
        url: str,
        payload: dict[str, Any],
        parser: Callable[[dict[str, Any]], ReplyDelta | None],
    ) -> Iterable[ReplyDelta]:
        body = json.dumps(payload).encode("utf-8")
        request = urllib.request.Request(
            url,
            data=body,
            headers={
                "Authorization": f"Bearer {self.config.llm_api_key}",
                "Content-Type": "application/json",
            },
            method="POST",
        )
        try:
            with self.urlopen(request, timeout=self.timeout_s) as response:
                raw_body = _read_response_body(response)
                emitted = False
                if raw_body.lstrip().startswith(b"data:"):
                    for raw_line in raw_body.splitlines():
                        line = raw_line.decode("utf-8", errors="replace").strip()
                        if not line or not line.startswith("data:"):
                            continue
                        data = line.removeprefix("data:").strip()
                        if data == "[DONE]":
                            yield ReplyDelta("", finish_reason="stop")
                            return
                        event = json.loads(data)
                        delta = parser(event)
                        if delta is not None:
                            emitted = True
                            yield delta
                else:
                    for delta in _parse_non_stream_json(json.loads(raw_body.decode("utf-8"))):
                        emitted = True
                        yield delta
                if not emitted:
                    raise ProviderError(
                        ErrorCode.LLM_EMPTY_REPLY,
                        "LLM stream ended without text",
                        True,
                        "openai-compatible",
                        "llm",
                    )
        except ProviderError:
            raise
        except urllib.error.HTTPError as exc:
            code = ErrorCode.LLM_RATE_LIMITED if exc.code == 429 else ErrorCode.LLM_NETWORK_ERROR
            raise ProviderError(code, f"LLM HTTP error {exc.code}", exc.code >= 500, "openai-compatible", "llm") from exc
        except (urllib.error.URLError, TimeoutError, socket.timeout) as exc:
            raise ProviderError(
                ErrorCode.LLM_REQUEST_TIMEOUT,
                f"LLM request failed or timed out: {exc}",
                True,
                "openai-compatible",
                "llm",
            ) from exc


def _parse_chat_completion_sse(event: dict[str, Any]) -> ReplyDelta | None:
    choices = event.get("choices") or []
    if not choices:
        return None
    choice = choices[0]
    text = str((choice.get("delta") or {}).get("content") or "")
    finish = choice.get("finish_reason")
    if not text and not finish:
        return None
    return ReplyDelta(text, is_sentence_boundary=_ends_sentence(text), finish_reason=finish)


def _parse_responses_sse(event: dict[str, Any]) -> ReplyDelta | None:
    event_type = event.get("type")
    if event_type == "response.output_text.delta":
        text = str(event.get("delta") or "")
        return ReplyDelta(text, is_sentence_boundary=_ends_sentence(text))
    if event_type in {"response.completed", "response.output_text.done"}:
        return ReplyDelta("", finish_reason="stop")
    return None


def _parse_non_stream_json(event: dict[str, Any]) -> Iterable[ReplyDelta]:
    if "choices" in event:
        choices = event.get("choices") or []
        if choices:
            message = choices[0].get("message") or {}
            text = str(message.get("content") or "")
            if text:
                yield ReplyDelta(text, is_sentence_boundary=_ends_sentence(text), finish_reason="stop")
        return
    if "output_text" in event:
        text = str(event.get("output_text") or "")
        if text:
            yield ReplyDelta(text, is_sentence_boundary=_ends_sentence(text), finish_reason="stop")
        return
    output = event.get("output") or []
    parts: list[str] = []
    for item in output:
        for content in item.get("content", []):
            if content.get("type") in {"output_text", "text"}:
                parts.append(str(content.get("text") or ""))
    text = "".join(parts)
    if text:
        yield ReplyDelta(text, is_sentence_boundary=_ends_sentence(text), finish_reason="stop")


def _read_response_body(response: Any) -> bytes:
    if hasattr(response, "read"):
        return response.read()
    return b"".join(response)


def _ends_sentence(text: str) -> bool:
    return text.endswith(("。", "！", "？", ".", "!", "?", "\n"))