Files

62 lines
1.9 KiB
Python

"""율봇 API 클라이언트 — youlbot REST API를 httpx로 호출."""
import json
import os
from typing import AsyncIterator
import httpx
from dotenv import load_dotenv
load_dotenv()
_API_URL = os.getenv("YOULBOT_API_URL", "http://localhost:8000").rstrip("/")
_API_TOKEN = os.getenv("YOULBOT_API_TOKEN", "")
def _headers() -> dict:
if _API_TOKEN:
return {"Authorization": f"Bearer {_API_TOKEN}"}
return {}
async def chat(
message: str,
user_id: str = "default",
show_thinking: bool = False,
) -> AsyncIterator[tuple[str, str | None]]:
"""SSE 스트림을 읽어 (token, run_id) 튜플을 yield.
- 일반 토큰: (token_str, None)
- 스트림 종료: ("", run_id_or_None) ← __done 이벤트
"""
async with httpx.AsyncClient(timeout=180) as client:
async with client.stream(
"POST",
f"{_API_URL}/chat",
json={"message": message, "user_id": user_id, "show_thinking": show_thinking},
headers=_headers(),
) 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 raw, None
continue
if isinstance(payload, dict) and payload.get("__done"):
yield "", payload.get("run_id")
return
yield payload, None
async def reset(user_id: str = "default") -> None:
async with httpx.AsyncClient(timeout=30) as client:
r = await client.post(
f"{_API_URL}/reset",
params={"user_id": user_id},
headers=_headers(),
)
r.raise_for_status()