youlbot-webui/api_client.py

"""율봇 API 클라이언트 — APIClientProtocol 인터페이스 + HTTPAPIClient 구현."""
import json
import os
import urllib.parse
from typing import AsyncIterator, Protocol, runtime_checkable

import httpx
from tenacity import retry, retry_if_exception, stop_after_attempt, wait_exponential

from config import AppConfig


def _is_transient(exc: Exception) -> bool:
    if isinstance(exc, httpx.HTTPStatusError):
        return exc.response.status_code >= 500
    return isinstance(exc, httpx.TransportError)


_retry = retry(
    retry=retry_if_exception(_is_transient),
    stop=stop_after_attempt(3),
    wait=wait_exponential(multiplier=1, min=1, max=8),
    reraise=True,
)


@runtime_checkable
class APIClientProtocol(Protocol):
    def chat(
        self, message: str, user_id: str, show_thinking: bool
    ) -> AsyncIterator[tuple[str, str | None]]: ...

    async def reset(self, user_id: str) -> None: ...

    async def ingest(self, file_path: str) -> dict: ...

    async def list_documents(self) -> list[str]: ...

    async def delete_document(self, source: str) -> None: ...

    async def save_feedback(
        self, user_id: str, user_msg: str, asst_msg: str, rating: int, run_id: str | None
    ) -> None: ...


class HTTPAPIClient:
    def __init__(self, config: AppConfig):
        self._url = config.youlbot_api_url.rstrip("/")
        self._timeout = config.youlbot_api_timeout
        self._client = httpx.AsyncClient(
            headers={"Authorization": f"Bearer {config.youlbot_api_token}"}
            if config.youlbot_api_token else {},
        )

    async def chat(
        self,
        message: str,
        user_id: str = "default",
        show_thinking: bool = False,
        image_path: str | None = None,
    ) -> AsyncIterator[tuple[str, str | None]]:
        payload: dict = {"message": message, "user_id": user_id, "show_thinking": show_thinking}
        if image_path:
            import base64
            with open(image_path, "rb") as f:
                payload["image_base64"] = base64.b64encode(f.read()).decode()
        async with self._client.stream(
            "POST",
            f"{self._url}/chat",
            json=payload,
            timeout=self._timeout,
        ) as response:
            response.raise_for_status()
            async for line in response.aiter_lines():
                if not line.startswith("data: "):
                    continue
                raw = line[6:]
                try:
                    payload = json.loads(raw)
                except json.JSONDecodeError:
                    yield str(raw), None
                    continue
                if isinstance(payload, dict) and payload.get("__done"):
                    yield "", payload.get("run_id")
                    return
                yield payload, None

    @_retry
    async def reset(self, user_id: str = "default") -> None:
        r = await self._client.post(
            f"{self._url}/reset",
            params={"user_id": user_id},
            timeout=30,
        )
        r.raise_for_status()

    @_retry
    async def ingest(self, file_path: str) -> dict:
        with open(file_path, "rb") as f:
            filename = os.path.basename(file_path)
            r = await self._client.post(
                f"{self._url}/ingest",
                files={"file": (filename, f, "application/octet-stream")},
                timeout=300,
            )
        r.raise_for_status()
        return r.json()

    @_retry
    async def list_documents(self) -> list[str]:
        r = await self._client.get(f"{self._url}/documents", timeout=30)
        r.raise_for_status()
        return r.json().get("documents", [])

    @_retry
    async def delete_document(self, source: str) -> None:
        encoded = urllib.parse.quote(source, safe="")
        r = await self._client.delete(f"{self._url}/documents/{encoded}", timeout=30)
        r.raise_for_status()

    @_retry
    async def save_feedback(
        self,
        user_id: str,
        user_msg: str,
        asst_msg: str,
        rating: int,
        run_id: str | None = None,
    ) -> None:
        r = await self._client.post(
            f"{self._url}/feedback",
            json={
                "user_id": user_id,
                "user_msg": user_msg,
                "asst_msg": asst_msg,
                "rating": rating,
                "run_id": run_id,
            },
            timeout=30,
        )
        r.raise_for_status()

    async def aclose(self) -> None:
        await self._client.aclose()