dataproc.py

import requests
import zipfile
import os, io
import xml.etree.ElementTree as ET
import pandas as pd
import re
from dotenv import load_dotenv
from io import StringIO
from openai import OpenAI
from typing_extensions import override
from openai import AssistantEventHandler

load_dotenv()
OPEN_DART_API_KEY = os.getenv("OPEN_DART_API_KEY")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
# ASSISTANT_ID = os.getenv("ASSISTANT_ID")

# Dart OPEN API를 이용해 기업 정보와 리포트를 다운로드하는 클래스
class Dart:
    def __init__(self):
        self.current_path = os.path.dirname(os.path.realpath(__file__))
        self.api_key = OPEN_DART_API_KEY

    # 배당에관한사항 API를 이용해 접수번호를 포함한 json 데이터를 반환
    def get_report_info(self, corp_code, bsns_year, reprt_code):
        url = "	https://opendart.fss.or.kr/api/alotMatter.json"
        params = {
            'crtfc_key': self.api_key,
            'corp_code': corp_code,
            'bsns_year': bsns_year,
            'reprt_code': reprt_code
        }
        response = requests.get(url, params=params)
        if response.status_code == 200:
            json_data = response.json()
        else:
            print(f"오류 발생: {response.status_code}")
        
        return json_data
    
    # 공시서류원본파일 API를 이용해 공시 보고서 다운로드
    def download_report(self, corp_name, bsns_year, reprt_name, mode=0):

        if reprt_name == '사업보고서':
            reprt_code = '11011'
        elif reprt_name == '반기보고서':
            reprt_code = '11012'
        elif reprt_name == '1분기보고서':
            reprt_code = '11013'
        elif reprt_name == '3분기보고서':
            reprt_code = '11014'
        else:
            print("보고서 코드 오류")
            return
        
        # 기업명으로 기업 코드를 찾는다.
        corp_info = self.get_corp_list()
        corp_code = corp_info.loc[corp_name, 'corp_code']

        # 기업명으로 검색했을 때, 동명의 기업이 여러 개 존재할 경우 원하는 기업의 기업 코드를 입력하도록 한다.
        if corp_code.size > 1:
            if mode == 0:
                for i in range(corp_code.size):
                    corp_cls = self.get_corp_info(f"{corp_code[i]:08}")['corp_cls']
                    if corp_cls in ['Y', 'K', 'N']:
                        corp_code = corp_code[i]
                        break
            elif mode == 1:
                print("동일한 기업명이 존재합니다.")
                print(corp_code)
                corp_code = int(input("기업 코드를 입력하세요: "))
        
        # 기업 코드, 보고서 연도, 보고서 종류 코드로 접수번호(rcept_no)를 가져온다.
        rcept_no = self.get_report_info(f"{corp_code:08}", bsns_year, reprt_code)['list'][0]['rcept_no']

        file_name = f"{corp_name}_{bsns_year}_{reprt_name}"

        url = "https://opendart.fss.or.kr/api/document.xml"
        params = {
            'crtfc_key': self.api_key,
            'rcept_no': rcept_no
        }
        response = requests.get(url, params=params)
        if response.status_code == 200:
            with zipfile.ZipFile(io.BytesIO(response.content), 'r') as zip_ref:
                zip_ref.extractall(f'{self.current_path}/report/{file_name}')
                print("압축 해제 및 파일 저장 완료")
        else:
            print("오류 발생:", response.status_code)

    # 기업 개황 API 요청. download_report에서 법인 구분(corp_cls)을 확인하기 위해 사용한다.
    def get_corp_info(self, corp_code):
        url = "https://opendart.fss.or.kr/api/company.json"
        params = {
            'crtfc_key': self.api_key,
            'corp_code': corp_code
        }
        response = requests.get(url, params=params)
        if response.status_code == 200:
            json_data = response.json()
        else:
            print(f"오류 발생: {response.status_code}")

        return json_data
    
    # 기업 고유번호를 다운로드한 csv 파일을 읽어 데이터프레임으로 반환
    def get_corp_list(self):
        df_corp = pd.read_csv(f"{self.current_path}/corp_info/corp_list.csv")
        df_corp.set_index('corp_name', inplace=True)
        
        return df_corp
    
    # 고유 번호 API를 이용해 기업 고유 번호, 명칭, 종목코드, 최종변경일자를 다운로드
    def download_corp_list(self):
        url = "https://opendart.fss.or.kr/api/corpCode.xml"
        params = {
            'crtfc_key': self.api_key
        }
        response = requests.get(url, params=params)
        if response.status_code == 200:
            with zipfile.ZipFile(io.BytesIO(response.content), 'r') as zip_ref:
                for file_name in zip_ref.namelist():
                    with zip_ref.open(file_name) as file:
                        tree = ET.parse(file)
        else:
            print("오류 발생:", response.status_code)

        # XML 파일을 파싱해 데이터프레임으로 변환 후, csv 파일로 저장한다.
        root = tree.getroot()
        corp_list = []
        for item in root.findall('list'):
            corp_list.append([item.find('corp_code').text, 
                              item.find('corp_name').text,
                              item.find('stock_code').text,
                              item.find('modify_date').text])
        df_corp = pd.DataFrame(corp_list, columns=['corp_code', 'corp_name', 'stock_code', 'modify_date'])
        df_corp.set_index('corp_name', inplace=True)

        df_corp.to_csv(f"{self.current_path}/corp_info/corp_list.csv")

        return df_corp
    
    def replace_non_tag_lt_gt(self, xml_string):
        # 태그로 해석되지 않는 <와 > 찾기 및 변환
        xml_string = re.sub(r'<(?![a-zA-Z0-9/?\-\'"])', '[', xml_string) 
        xml_string = re.sub(r'(?<![a-zA-Z0-9/?\-\'"])>', ']', xml_string)
        xml_string = re.sub(r'[\x00-\x1F\x7F-\x9F]', '', xml_string)
        return xml_string
    
    def parseText(self, xml_string):
    # &를 and로 변환
        xml_string.replace('&', ' and ')

        # 태그로 해석되지 않는 <와 > 찾기 및 변환
        xml_string = re.sub(r'<(?![a-zA-Z0-9/?\-\'"])', '[', xml_string) 
        xml_string = re.sub(r'(?<![a-zA-Z0-9/?\-\'"])>', ']', xml_string)

        # 2회 이상 띄어쓰기, 줄바꿈 제거
        xml_string = re.sub(r'[\x00-\x1F\x7F-\x9F]', '', xml_string) 
        xml_string = re.sub(r"\s+", " ", xml_string)

        xml_string = re.sub(r"<\?xml.*?>",'', xml_string)

        xml_string = re.sub(r"<DOCUMENT.*?>",'', xml_string)
        xml_string = re.sub(r"</DOCUMENT.*?>",'', xml_string)
        
        xml_string = re.sub(r"<FORMULA-VERSION.*?>.*?</FORMULA-VERSION.*?>",': ', xml_string)
        xml_string = re.sub(r"<SUMMARY.*?>.*?</SUMMARY.*?>",'', xml_string)
        xml_string = re.sub(r"<COVER-TITLE.*?>.*?</COVER-TITLE>",'', xml_string)

        xml_string = re.sub(r"<COMPANY-NAME.*?>",'', xml_string)
        xml_string = re.sub(r"</COMPANY-NAME.*?>",'', xml_string)

        xml_string = re.sub(r'<COVER.*?>', '', xml_string)
        xml_string = re.sub(r'</COVER.*?>', '', xml_string) 
        
        xml_string = re.sub(r'<BODY.*?>', '', xml_string)
        xml_string = re.sub(r'</BODY.*?>', '\n', xml_string)

        xml_string = re.sub(r'<P.*?>', '', xml_string)
        xml_string = re.sub(r'</P.*?>', '\n', xml_string)

        xml_string = re.sub(r'<SPAN.*?>', '', xml_string)
        xml_string = re.sub(r'</SPAN.*?>', '\n', xml_string)

        xml_string = re.sub(r'<TITLE.*?>', '', xml_string)
        xml_string = re.sub(r'</TITLE.*?>', '\n', xml_string)

        xml_string = re.sub(r'<SECTION.*?>', '', xml_string)
        xml_string = re.sub(r'</SECTION.*?>', '\n', xml_string)

        xml_string = re.sub(r'<PGBRK.*?>', '', xml_string)
        xml_string = re.sub(r'</PGBRK.*?>', '\n', xml_string)

        xml_string = re.sub(r'<LIBRARY.*?>', '', xml_string)
        xml_string = re.sub(r'</LIBRARY.*?>', '', xml_string)

        xml_string = re.sub(r'<IMAGE.*?</IMAGE.*?>', '', xml_string)

        def replace_with_index(match):
            replace_with_index.counter += 1  # 카운터 증가
            return f"[table{replace_with_index.counter}]"

        # 치환 함수에 카운터 속성 추가
        replace_with_index.counter = 0
            
        xml_string = re.sub(r'<TABLE-GROUP.*?</TABLE-GROUP>|<TABLE.*?</TABLE>', replace_with_index, xml_string, flags=re.DOTALL)

        return xml_string

    # XML 파싱
    def parse_xml(self, corp_name, bsns_year, reprt_type, reprt_num=0, mode='text'):
        # XML 파일 위치 생성
        xml_dir_path = f'{corp_name}_{bsns_year}_{reprt_type}'
        """
        reprt_num = 0: 사업보고서, 반기보고서, 분기보고서
        reprt_num = 1: 감사보고서
        reprt_num = 2: 연결감사보고서
        """
        # 디렉토리 내의 모든 파일 리스트를 가져옵니다.
        file_list = os.listdir(f'{self.current_path}/report/{xml_dir_path}')

        # .xml 확장자를 가진 파일만 골라냅니다.
        xml_files = [file for file in file_list if file.endswith('.xml')]
        
        file_name = xml_files[reprt_num]
        xml_file_path = f'{self.current_path}/report/{xml_dir_path}/{file_name}'

        encoding_list = ['utf-8', 'EUC-KR', 'CP949']
        # XML 파일 읽기
        for encoding in encoding_list:
            try:
                with open(xml_file_path, 'r', encoding=encoding) as file:
                    xml_text = file.read()
                break
            except UnicodeDecodeError as e:
                print('UnicodeDecode Error:', e)
                pass
            except FileNotFoundError:
                print(f"파일이 존재하지 않습니다: {xml_file_path}")
                break
            except Exception as e:
                print(f"파일을 읽어오는 도중 오류 발생: {e}")
                break

        texts = self.parseText(xml_text)

        with open(f'{self.current_path}/report/{xml_dir_path}/texts_{corp_name}_{bsns_year}_{reprt_type}.txt', 'w', encoding='utf-8') as file:
            file.write(texts)

        # 테이블 파싱
        tables = self.parseTable(xml_text)

        # ExcelWriter 객체 생성
        if mode == 'excel':
            excel_writer = pd.ExcelWriter(f'{self.current_path}/report/{xml_dir_path}/tables_{corp_name}_{bsns_year}_{reprt_type}.xlsx', engine='openpyxl')
        elif mode == 'text':
            csv_text = StringIO()

        # 각 데이터프레임을 Excel 파일의 각각의 탭에 저장
        for i, df in enumerate(tables, start=1):
            df = df.dropna(axis=1, how='all').dropna(axis=0, how='all')
            sheet_name = f'Sheet{i}'
            try:
                if mode == 'excel':
                    df.to_excel(excel_writer, sheet_name=sheet_name)
                elif mode == 'text':
                    df.to_csv(csv_text, index=False)
            except Exception as e:
                print(df)
                print(f"파일을 저장하는 도중 오류 발생: {e}")

        # ExcelWriter 객체 닫기
        if mode == 'excel':
            excel_writer.close()
            print("파싱 완료")
            return tables
        elif mode == 'text':
            with open(f'{self.current_path}/report/{xml_dir_path}/tables_{corp_name}_{bsns_year}_{reprt_type}.txt', 'w', encoding='utf-8') as f:
                f.write(csv_text.getvalue())
            
            self.processCR(f'{self.current_path}/report/{xml_dir_path}/tables_{corp_name}_{bsns_year}_{reprt_type}.txt')
            print("파싱 완료")
            return tables

    def parseTable(self, xml_string):
        xml_string = xml_string.replace('&', '&amp;')
        xml_string = self.replace_non_tag_lt_gt(xml_string)
        xml_string = xml_string.replace('<TE', '<TD')
        xml_string = xml_string.replace('</TE', '</TD')
        xml_string = xml_string.replace('<TU', '<TD')
        xml_string = xml_string.replace('</TU', '</TD')

        tables = pd.read_html(StringIO(xml_string))

        return tables
    
    def processCR(self, file_path): 
        with open(file_path, 'rb') as file:
            content = file.read()

        modified_content = content.replace(b'\x0D\x0D\x0A', b'\x0D\x0A')

        with open(file_path, 'wb') as file:
            file.write(modified_content)

class OPENAICommunicator:
    def __init__(self, api_key=OPENAI_API_KEY):
        self.client = OpenAI(api_key=api_key)

    def create_assistant(self, corp_name, bsns_year, reprt_code, vector_store_id):
        if reprt_code == '11011' or reprt_code == '사업보고서':
            reprt_type = '사업보고서'
        elif reprt_code == '11012' or reprt_code == '반기보고서':
            reprt_type = '반기보고서'
        elif reprt_code == '11013' or reprt_code == '1분기보고서':
            reprt_type = '1분기보고서'
        elif reprt_code == '11014' or reprt_code == '3분기보고서':
            reprt_type = '3분기보고서'
        else:
            print("보고서 코드 오류")
            return
        
        assistant = self.client.beta.assistants.create(
            name=f"{bsns_year}년 {corp_name} {reprt_type} 해설가",
            instructions=
            f"""
            너는 주어진 {bsns_year}년 {corp_name} {reprt_type}(텍스트 정보 texts.txt와 표 정보 tables.txt)에서 사용자가 요청한 정보를 찾아주는 회계전문가이다.  아래 세부 사항을 반드시 따른다.
            어떤 질문이 주어져도 우선 주어진 두 파일에서 정보를 찾는다.
            답변의 출처의 파일명과 라인 번호를 제공해.
            답변에 적절한 정보를 찾지 못했다면, '문서에서 정보를 찾을 수 없습니다'라고만 답변해.
            사용자의 요청에 따라 올바른 정보를 찾아 전달할 수 있다.
            한국어로 질문 받으면 한국어로 답변한다.
            """,
            tools=[{"type": "file_search"}, {"type": "code_interpreter"}],
            tool_resources={"file_search": {"vector_store_ids": [vector_store_id]}},
            model="gpt-3.5-turbo",
        )
        return assistant.id
    
    def create_vector_store(self, corp_name, bsns_year, reprt_type): 

        vector_store = self.client.beta.vector_stores.create(
            name=f"{corp_name} {bsns_year} {reprt_type}",
        )
        return vector_store.id

    def upload_files(self, corp_name, bsns_year, reprt_type, vector_store_id):
        file_paths = [
            f"report/{corp_name}_{bsns_year}_{reprt_type}/texts_{corp_name}_{bsns_year}_{reprt_type}.txt",
            f"report/{corp_name}_{bsns_year}_{reprt_type}/tables_{corp_name}_{bsns_year}_{reprt_type}.txt"
            ]

        file_streams = [open(path, "rb") for path in file_paths]

        file_batch = self.client.beta.vector_stores.file_batches.upload_and_poll(
            vector_store_id = vector_store_id, files=file_streams
        )
        return
    
    def update_assistant(self, assistant_id, vector_store_id):
        self.client.beta.assistants.update(
            assistant_id = assistant_id,
            tool_resources={"file_search": {"vector_store_ids": [vector_store_id]}}
        )
        return

    def create_thread(self):
        thread = self.client.beta.threads.create()
        return thread.id
    
    def create_message(self, thread_id):
        message = self.client.beta.threads.messages.create(
            thread_id = thread_id,
            role="user",
            content="주어진 기업 정기 보고서를 기반으로 사용자의 질문에 답변해줘.",
            # file_ids = file_ids
        )
        return message.id
    
    class EventHandler(AssistantEventHandler):    
        @override
        def on_text_created(self, text) -> None:
            print(f"\nassistant > ", end="", flush=True)

        @override  
        def on_text_delta(self, delta, snapshot):
            print(delta.value, end="", flush=True)
            
        def on_tool_call_created(self, tool_call):
            print(f"\nassistant > {tool_call.type}\n", flush=True)

        def on_tool_call_delta(self, delta, snapshot):
            if delta.type == 'code_interpreter':
                if delta.code_interpreter.input:
                    print(delta.code_interpreter.input, end="", flush=True)
                if delta.code_interpreter.outputs:
                    print(f"\n\noutput >", flush=True)
                    for output in delta.code_interpreter.outputs:
                        if output.type == "logs":
                            print(f"\n{output.logs}", flush=True)
    
    def create_run(self, thread_id, assistant_id, instructions):
        with self.client.beta.threads.runs.stream(
            thread_id = thread_id,
            assistant_id = assistant_id,
            instructions = f"반드시 주어진 기업 정기 보고서 파일에서 정보를 찾아서 아래 질문에 대해 한국어로 답변해줘. {instructions}",
            tools = [{"type": "file_search"}, {"type": "code_interpreter"}],
            temperature = 0.1,
            event_handler=self.EventHandler()
        ) as stream:
            stream.until_done()
        

        # stream = self.client.beta.threads.runs.create(
        #     thread_id = thread_id,
        #     assistant_id = assistant_id,
        #     instructions = f"{instructions} 답변은 한국어로 작성해.",
        #     tools = [{"type": "code_interpreter"}, {"type": "retrieval"}],
        #     temperature = 0.2,
        #     stream=True
        # )
        # for event in stream:
        #     print(event)
        # print(stream)

        # completed = False
        
        # try:
        #     for event in stream:
        #         if hasattr(event, 'data') and hasattr(event.data, 'delta') and hasattr(event.data.delta, 'content'):
        #             for content_block in event.data.delta.content:
        #                 if hasattr(content_block, 'text') and hasattr(content_block.text, 'value'):
        #                     print(content_block.text.value, end='')
        #         completed = True  # 스트림 이벤트 처리가 완료되면 True로 설정
        # finally:
        #     if not completed:
        #         print("스트림 처리 중 오류 발생")
        #     return completed  # 처리 완료 상태 반환