test_streaming.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
ChatBIA Streaming API 테스트 클라이언트
"""
import requests
import json
import sys
import io

if sys.platform == 'win32':
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')


def test_streaming_chat(base_url: str, message: str, mode: str = "bsl"):
    """
    스트리밍 채팅 테스트

    Args:
        base_url: API 서버 URL (예: http://localhost:8000)
        message: 사용자 메시지
        mode: "bsl" 또는 "general"
    """
    url = f"{base_url}/chat/stream"

    payload = {
        "message": message,
        "mode": mode,
        "max_tokens": 512,
        "temperature": 0.7
    }

    print(f"🔄 요청 전송: {message}")
    print(f"📡 모드: {mode}")
    print("=" * 60)

    try:
        # 스트리밍 요청
        response = requests.post(
            url,
            json=payload,
            stream=True,  # 중요: 스트리밍 활성화
            headers={
                "Content-Type": "application/json",
                "Accept": "text/event-stream"
            },
            timeout=120  # 2분 타임아웃 (충분히 긴 시간)
        )

        if response.status_code != 200:
            print(f"❌ 오류: {response.status_code}")
            print(response.text)
            return

        print("✅ 응답 수신 중...\n")

        full_text = ""
        token_count = 0

        # SSE 스트림 읽기
        for line in response.iter_lines():
            if line:
                line_str = line.decode('utf-8')

                # SSE 형식: "data: {json}"
                if line_str.startswith("data: "):
                    data_str = line_str[6:]  # "data: " 제거

                    try:
                        data = json.loads(data_str)

                        # 에러 체크
                        if "error" in data:
                            print(f"\n❌ 서버 오류: {data['error']}")
                            break

                        # 토큰 출력
                        if not data.get("done", False):
                            token = data.get("token", "")
                            print(token, end="", flush=True)
                            full_text += token
                            token_count = data.get("token_count", token_count)
                        else:
                            # 완료
                            print(f"\n\n✅ 완료!")
                            print(f"📊 토큰 수: {data.get('token_count', token_count)}")
                            print(f"🎯 모드: {data.get('mode', mode)}")
                            break

                    except json.JSONDecodeError as e:
                        print(f"\n⚠️ JSON 파싱 오류: {e}")
                        print(f"   원본: {data_str}")

        print("\n" + "=" * 60)
        print(f"전체 응답 길이: {len(full_text)} 글자")

    except requests.exceptions.Timeout:
        print("❌ 타임아웃 오류: 서버 응답이 너무 느립니다.")
    except requests.exceptions.ConnectionError:
        print("❌ 연결 오류: 서버에 연결할 수 없습니다.")
    except Exception as e:
        print(f"❌ 예상치 못한 오류: {e}")


def test_regular_chat(base_url: str, message: str, mode: str = "bsl"):
    """
    일반 (비스트리밍) 채팅 테스트
    """
    url = f"{base_url}/chat"

    payload = {
        "message": message,
        "mode": mode,
        "max_tokens": 512,
        "temperature": 0.7
    }

    print(f"🔄 일반 요청 전송: {message}")
    print(f"📡 모드: {mode}")
    print("=" * 60)

    try:
        response = requests.post(url, json=payload, timeout=60)

        if response.status_code == 200:
            result = response.json()
            print("✅ 응답:\n")
            print(result["response"])
            print(f"\n📊 토큰 수: {result['tokens']}")
        else:
            print(f"❌ 오류: {response.status_code}")
            print(response.text)

    except Exception as e:
        print(f"❌ 오류: {e}")


if __name__ == "__main__":
    # 로컬 테스트
    BASE_URL = "http://localhost:8000"

    # Hugging Face Spaces 테스트 (배포 후)
    # BASE_URL = "https://your-username-chatbia-server.hf.space"

    print("=" * 60)
    print("ChatBIA Streaming API 테스트")
    print("=" * 60)

    # 테스트 1: 스트리밍 채팅 (BSL 모드)
    print("\n[테스트 1] 스트리밍 채팅 - BSL 모드\n")
    test_streaming_chat(
        BASE_URL,
        "5천만원 설비의 감가상각 계산해줘. 내용연수는 10년이고 잔존가치는 10%야.",
        mode="bsl"
    )

    print("\n" + "=" * 60 + "\n")

    # 테스트 2: 스트리밍 채팅 (일반 모드)
    print("\n[테스트 2] 스트리밍 채팅 - 일반 모드\n")
    test_streaming_chat(
        BASE_URL,
        "회계에서 감가상각이 뭐야?",
        mode="general"
    )

    print("\n" + "=" * 60 + "\n")

    # 테스트 3: 일반 채팅 (비스트리밍) - 비교용
    print("\n[테스트 3] 일반 채팅 (비스트리밍) - 비교용\n")
    test_regular_chat(
        BASE_URL,
        "안녕하세요!",
        mode="bsl"
    )