app.py

from flask import Flask, request, jsonify
from flask_cors import CORS
import os
from dotenv import load_dotenv
from dataproc import Dart, OPENAICommunicator

app = Flask(__name__)
CORS(app)  # CORS 설정 추가

# .env 파일에서 환경 변수를 로드
load_dotenv()

def extract_thread_info():
    thread_info_list = []
    for key, value in os.environ.items():
        if 'THREAD_ID' in key:
            parts = key.split('_')
            if len(parts) >= 6:
                corp_code = parts[0]
                bsns_year = parts[1]
                reprt_code = parts[2]
                thread_id = value
                common_identifier = f"{corp_code}_{bsns_year}_{reprt_code}"
                thread_info_list.append({
                    'corp_code': corp_code,
                    'bsns_year': bsns_year,
                    'reprt_code': reprt_code,
                    'thread_id': thread_id,
                    'common_identifier': common_identifier
                })
    return thread_info_list

def extract_assistant_info():
    assistant_info_list = []
    for key, value in os.environ.items():
        if 'ASSISTANT_ID' in key:
            parts = key.split('_')
            # Debug: Print the parts to verify the splitting
            print(f"Processing key: {key}, Parts: {parts}")
            if len(parts) >= 5:
                corp_code = parts[0]
                bsns_year = parts[1]
                reprt_code = parts[2]
                assistant_id = value
                common_identifier = f"{corp_code}_{bsns_year}_{reprt_code}"
                assistant_info_list.append({
                    'corp_code': corp_code,
                    'bsns_year': bsns_year,
                    'reprt_code': reprt_code,
                    'assistant_id': assistant_id,
                    'common_identifier': common_identifier
                })
    # Debug: Print the extracted assistant info list
    print(f"Assistant Info List: {assistant_info_list}")
    return assistant_info_list


@app.route('/threads', methods=['GET'])
def get_threads():
    thread_info_list = extract_thread_info()
    assistant_info_list = extract_assistant_info()

    # Convert the assistant_info_list to a dictionary for quick lookup
    assistant_info_dict = {item['common_identifier']: item['assistant_id'] for item in assistant_info_list}

    # Debug: Print the dictionaries to check identifiers
    print("Assistant Info Dict:", assistant_info_dict)

    # Add assistant_id to the thread_info_list based on the common identifier
    for thread_info in thread_info_list:
        common_identifier = thread_info['common_identifier']
        thread_info['assistant_id'] = assistant_info_dict.get(common_identifier, None)

    # Debug: Print the thread info list to check if assistant_id is added
    print("Thread Info List:", thread_info_list)

    # Remove the common_identifier key as it's no longer needed
    for thread_info in thread_info_list:
        del thread_info['common_identifier']

    return jsonify(thread_info_list)


@app.route('/process_report', methods=['POST'])
def process_report():
    data = request.json
    corp_name = data.get('corp_name')
    corp_code = data.get('corp_code')
    bsns_year = data.get('bsns_year')
    reprt_code = data.get('reprt_code')
    uid = data.get('uid', 'default_uid')

    d = Dart()
    communicator = OPENAICommunicator()

    if reprt_code == '11011':
        reprt_type = '사업보고서'
    elif reprt_code == '11012':
        reprt_type = '반기보고서'
    elif reprt_code == '11013':
        reprt_type = '1분기보고서'
    elif reprt_code == '11014':
        reprt_type = '3분기보고서'
    else:
        print("보고서 코드 오류")

    # 리포트 다운로드
    d.download_report(corp_name, bsns_year, reprt_type)
    d.parse_xml(corp_name, bsns_year, reprt_type)

    assistant_id_env_name = f"{corp_code}_{bsns_year}_{reprt_code}_assistant_id".upper()
    thread_id_env_name = f"{corp_code}_{bsns_year}_{reprt_code}_thread_id_{uid}".upper()
    vector_store_id_env_name = f"{corp_code}_{bsns_year}_{reprt_code}_vector_store_id".upper()

    # assistant_id 확인 및 생성
    if not os.getenv(assistant_id_env_name):
        vector_store_id = communicator.create_vector_store()
        file_ids = communicator.upload_files(corp_name, bsns_year, reprt_type, vector_store_id)  # 파일 업로드
        assistant_id = communicator.create_assistant(corp_name, bsns_year, reprt_code, vector_store_id)  # 어시스턴트 생성

        # 환경 변수에 assistant_id 저장
        with open(".env", "a") as env_file:
            env_file.write(f"{assistant_id_env_name}={assistant_id}\n")
            env_file.write(f"{vector_store_id_env_name}={vector_store_id}\n")
    else:
        assistant_id = os.getenv(assistant_id_env_name)
        vector_store_id = os.getenv(vector_store_id_env_name)

    # thread_id 확인 및 생성
    if not os.getenv(thread_id_env_name):
        thread_id = communicator.create_thread()
        
        # 환경 변수에 thread_id 저장
        with open(".env", "a") as env_file:
            env_file.write(f"{thread_id_env_name}={thread_id}\n")
    else:
        thread_id = os.getenv(thread_id_env_name)

    response = {
        ' ': "success in create the thread",
    }

    return jsonify(response)

if __name__ == '__main__':
    app.run(debug=True)