test_providers.py•4.18 kB
#!/usr/bin/env python3
"""
Provider 테스트 스크립트
OpenAI와 vLLM Provider를 모두 테스트합니다.
"""
import os
import sys
import json
import time
import traceback
from typing import Dict, Any
# 현재 디렉토리를 Python 경로에 추가
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from mcp_pii_tools import MCPPIIDetector
def test_provider(provider_type: str, test_text: str) -> Dict[str, Any]:
"""특정 Provider로 PII 탐지를 테스트합니다."""
print(f"\n{'='*60}")
print(f"Provider 테스트: {provider_type.upper()}")
print(f"{'='*60}")
# 환경변수 설정
original_provider = os.environ.get("PII_PROVIDER")
os.environ["PII_PROVIDER"] = provider_type
try:
# PII 탐지기 초기화
detector = MCPPIIDetector()
# PII 탐지 실행
start_time = time.time()
result = detector.detect_pii(test_text)
duration = time.time() - start_time
print(f"✅ {provider_type.upper()} 테스트 성공! (응답 시간: {duration:.2f}초)")
print(f"탐지된 PII 수: {result.get('count', 0)}")
if result.get('success') and result.get('pii_items'):
print("\n=== 탐지된 PII 항목 ===")
for i, item in enumerate(result['pii_items'], 1):
print(f"{i}. {item['type']}: {item['value']} (위치: {item['start_pos']}-{item['end_pos']})")
return {
"success": True,
"provider": provider_type,
"duration": duration,
"count": result.get('count', 0),
"items": result.get('pii_items', [])
}
except Exception as e:
print(f"❌ {provider_type.upper()} 테스트 실패: {e.__class__.__name__}")
print(f"오류 메시지: {str(e)}")
traceback.print_exc()
return {
"success": False,
"provider": provider_type,
"error": str(e)
}
finally:
# 환경변수 복원
if original_provider is not None:
os.environ["PII_PROVIDER"] = original_provider
elif "PII_PROVIDER" in os.environ:
del os.environ["PII_PROVIDER"]
def main():
"""메인 테스트 함수"""
print("MCP PII Tools - Provider 테스트")
print("=" * 60)
# 테스트 텍스트
test_text = "김철수는 010-1234-5678로 연락 가능하며, 이메일은 chulsoo.kim@example.com 입니다. 주소는 서울특별시 강남구 테헤란로 231 1203호 입니다."
print(f"테스트 텍스트: {test_text}")
# 결과 저장
results = []
# OpenAI 테스트
if os.getenv("OPENAI_API_KEY"):
result = test_provider("openai", test_text)
results.append(result)
else:
print("\n⚠️ OPENAI_API_KEY가 설정되지 않아 OpenAI 테스트를 건너뜁니다.")
# vLLM 테스트
if os.getenv("VLLM_API_KEY") or os.getenv("OPENAI_API_KEY"):
result = test_provider("vllm", test_text)
results.append(result)
else:
print("\n⚠️ VLLM_API_KEY가 설정되지 않아 vLLM 테스트를 건너뜁니다.")
# 결과 요약
print(f"\n{'='*60}")
print("테스트 결과 요약")
print(f"{'='*60}")
for result in results:
status = "✅ 성공" if result["success"] else "❌ 실패"
duration = f" ({result.get('duration', 0):.2f}초)" if result["success"] else ""
count = f" - {result.get('count', 0)}개 PII 탐지" if result["success"] else ""
print(f"{result['provider'].upper()}: {status}{duration}{count}")
if not result["success"] and "error" in result:
print(f" 오류: {result['error']}")
# 성공한 Provider가 있는지 확인
successful_providers = [r for r in results if r["success"]]
if successful_providers:
print(f"\n🎉 {len(successful_providers)}개의 Provider가 정상 작동합니다!")
else:
print("\n⚠️ 모든 Provider 테스트가 실패했습니다. 환경변수를 확인해주세요.")
if __name__ == "__main__":
main()