178 lines
7.1 KiB
Python
178 lines
7.1 KiB
Python
import requests
|
|
import time
|
|
import json
|
|
import os
|
|
import argparse
|
|
import sys
|
|
from urllib.parse import urljoin
|
|
import logging
|
|
from dotenv import load_dotenv
|
|
|
|
# --- 로거 설정 ---
|
|
# 전역 로거 객체 생성
|
|
logger = logging.getLogger(__name__)
|
|
|
|
def setup_logger():
|
|
"""로거를 설정하여 콘솔과 파일에 모두 출력하도록 합니다."""
|
|
logger.setLevel(logging.INFO) # 로거의 최소 레벨 설정
|
|
|
|
# 로그 포맷 지정
|
|
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
|
|
|
|
# 콘솔 핸들러 설정
|
|
console_handler = logging.StreamHandler()
|
|
console_handler.setFormatter(formatter)
|
|
logger.addHandler(console_handler)
|
|
|
|
# 파일 핸들러 설정 (예: 'script_run.log' 파일에 저장)
|
|
file_handler = logging.FileHandler('script_run.log', encoding='utf-8')
|
|
file_handler.setFormatter(formatter)
|
|
logger.addHandler(file_handler)
|
|
|
|
# --- API 요청 함수 ---
|
|
|
|
def start_extraction(post_url, file_path, filename, headers, model_name=None):
|
|
"""POST /extract/inner: 문서 추출 시작"""
|
|
try:
|
|
with open(file_path, 'rb') as input_f:
|
|
files_to_upload = {'input_file': (filename, input_f)}
|
|
data_payload = {}
|
|
if model_name:
|
|
data_payload['model'] = model_name
|
|
|
|
response = requests.post(post_url, files=files_to_upload, data=data_payload, headers=headers)
|
|
response.raise_for_status()
|
|
|
|
return response.json()
|
|
except Exception:
|
|
# logger.exception은 오류의 상세 정보(traceback)까지 기록해줍니다.
|
|
logger.exception(f"[{filename}] POST 요청 중 오류 발생")
|
|
return None
|
|
|
|
def check_progress(base_url, progress_path, filename, headers):
|
|
"""GET /extract/progress/{request_id}: 진행 상태 확인 (로깅 적용)"""
|
|
get_url = urljoin(base_url + '/', progress_path.lstrip('/'))
|
|
|
|
RETRY_COUNT_ON_404 = 3
|
|
RETRY_DELAY_ON_404 = 5
|
|
retries_left = RETRY_COUNT_ON_404
|
|
last_status = ""
|
|
|
|
while True:
|
|
try:
|
|
response = requests.get(get_url, headers=headers, timeout=30)
|
|
|
|
if response.status_code == 404:
|
|
if retries_left > 0:
|
|
logger.warning(f"[{filename}] 작업을 찾을 수 없어(404) {RETRY_DELAY_ON_404}초 후 재시도합니다... ({retries_left}회 남음)")
|
|
retries_left -= 1
|
|
time.sleep(RETRY_DELAY_ON_404)
|
|
continue
|
|
else:
|
|
logger.error(f"[{filename}] 재시도 횟수 초과 후에도 작업을 찾을 수 없습니다 (404).")
|
|
return None
|
|
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
|
|
if "final_result" in data and data.get("final_result") is not None:
|
|
logger.info(f"[{filename}] 처리 완료.")
|
|
return data["final_result"]
|
|
|
|
if "progress_logs" in data and data["progress_logs"]:
|
|
status_message = data["progress_logs"][-1].get("status", "상태 확인 중...")
|
|
if status_message != last_status:
|
|
last_status = status_message
|
|
logger.info(f"[{filename}] 진행 상태: {last_status}")
|
|
|
|
time.sleep(2)
|
|
except requests.exceptions.ReadTimeout:
|
|
logger.warning(f"[{filename}] 상태 확인 타임아웃. 재시도...")
|
|
time.sleep(2)
|
|
except Exception:
|
|
logger.exception(f"[{filename}] 상태 확인 중 예측하지 못한 오류 발생")
|
|
return None
|
|
|
|
# --- 메인 실행 로직 ---
|
|
|
|
def main():
|
|
# .env 파일에서 환경 변수 로드 (workspace 디렉터리 기준)
|
|
dotenv_path = os.path.join(os.path.dirname(__file__), '.env')
|
|
load_dotenv(dotenv_path=dotenv_path)
|
|
|
|
# 로거를 가장 먼저 설정합니다.
|
|
setup_logger()
|
|
|
|
# 환경 변수에서 API 정보 가져오기
|
|
BASE_URL = os.getenv("BASE_URL")
|
|
API_KEY = os.getenv("API_KEY")
|
|
|
|
if not BASE_URL or not API_KEY:
|
|
logger.error("환경 변수(BASE_URL, API_KEY)가 설정되지 않았습니다. workspace/.env 파일을 확인하세요.")
|
|
return
|
|
|
|
parser = argparse.ArgumentParser(description="문서 정보 추출 자동화 스크립트")
|
|
parser.add_argument("input_dir", help="입력 디렉터리 경로")
|
|
parser.add_argument("-o", "--output_dir", default="results", help="출력 디렉터리 경로")
|
|
parser.add_argument("--endpoint", choices=['i18n', 'd6c'], default='i18n', help="추출 API 엔드포인트 선택 (i18n 또는 d6c)")
|
|
parser.add_argument("--model", dest="model_name", help="사용할 LLM 모델 이름")
|
|
args = parser.parse_args()
|
|
|
|
if not os.path.isdir(args.input_dir):
|
|
logger.error(f"입력 디렉터리를 찾을 수 없습니다 - {args.input_dir}")
|
|
return
|
|
|
|
os.makedirs(args.output_dir, exist_ok=True)
|
|
headers = {'X-API-KEY': API_KEY}
|
|
|
|
post_url = f"{BASE_URL}/extract/inner/{args.endpoint}"
|
|
|
|
logger.info("="*20 + " 스크립트 시작 " + "="*20)
|
|
logger.info(f"API 서버: {BASE_URL}")
|
|
logger.info(f"요청 API: {post_url}")
|
|
logger.info(f"입력 디렉터리: {args.input_dir}")
|
|
logger.info(f"출력 디렉터리: {args.output_dir}")
|
|
|
|
for filename in sorted(os.listdir(args.input_dir)):
|
|
file_path = os.path.join(args.input_dir, filename)
|
|
if not os.path.isfile(file_path):
|
|
continue
|
|
|
|
logger.info(f"--- 처리 시작: {filename} ---")
|
|
|
|
initial_response = start_extraction(post_url, file_path, filename, headers, args.model_name)
|
|
if not initial_response:
|
|
logger.error(f"[{filename}] 파일 처리 실패 (추출 시작 단계)")
|
|
continue
|
|
|
|
request_id = initial_response.get("request_id")
|
|
status_check_url = initial_response.get("status_check_url")
|
|
|
|
if not request_id or not status_check_url:
|
|
logger.error(f"[{filename}] 초기 응답이 잘못되었습니다: {initial_response}")
|
|
continue
|
|
|
|
logger.info(f"[{filename}] 작업 요청 성공. Request ID: {request_id}")
|
|
|
|
final_result = check_progress(BASE_URL, status_check_url, filename, headers)
|
|
|
|
if final_result:
|
|
output_path = os.path.join(args.output_dir, f"{os.path.splitext(filename)[0]}.json")
|
|
try:
|
|
with open(output_path, 'w', encoding='utf-8') as f:
|
|
json.dump(final_result, f, indent=2, ensure_ascii=False)
|
|
logger.info(f"[{filename}] 결과 저장 완료: {output_path}")
|
|
except IOError:
|
|
logger.exception(f"[{filename}] 파일 저장 중 오류 발생")
|
|
else:
|
|
logger.error(f"[{filename}] 파일 처리 실패 (결과 확인 단계)")
|
|
|
|
logger.info("="*20 + " 모든 작업 완료 " + "="*20)
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
main()
|
|
except KeyboardInterrupt:
|
|
# KeyboardInterrupt는 main 밖에서 처리해야 할 수 있으므로 로거를 직접 호출
|
|
logging.getLogger(__name__).warning("사용자에 의해 작업이 중단되었습니다.")
|