Langchain으로 데이터 시각화하기(+Dash)

import os 
from langchain_openai import ChatOpenAI 
from langchain_core.prompts import ChatPromptTemplate  
from langchain.prompts.chat import MessagesPlaceholder 
from langchain.memory import ConversationSummaryBufferMemory 
from langchain.schema.runnable import RunnablePassthrough  

from dash import Dash, html, dcc, callback, Output, Input, State 
import dash_ag_grid as dag 
import pandas as pd  
import re  # 정규 표현식 사용

# 환경 변수 설정

# CSV 파일에서 데이터 로드 및 처리
df = pd.read_csv('StudentsPerformance.csv')  # CSV 파일 읽기
df_select = df.head(100)  # 상위 100개 행 선택
csv_string = df_select.to_string(index=False)  # 프롬프트에서 사용하기 위해 데이터프레임을 문자열로 변환

# ChatGPT 모델 설정
llm = ChatOpenAI(model="gpt-4o", temperature=0)  

# 대화 요약 메모리 설정
memory = ConversationSummaryBufferMemory(
   llm=llm, # 대화 요약에 사용할 모델 지정
   max_token_limit=200, # 메모리에 저장할 최대 토큰 수, 초과시 이전 것부터 삭제됨 
   memory_key="chat_history", # 대화 기록을 저장할 키
   return_messages=True # 메모리에서 전체 메시지를 반환하도록 설정
)

# 메모리 로드 함수 정의
def load_memory(input):
   return memory.load_memory_variables({})["chat_history"] #  메모리에서 대화 기록 변수를 불러와 반환

# 프롬프트 템플릿 정의
prompt = ChatPromptTemplate.from_messages([
   # 시스템 메시지, 데이터 시각화 전문가 역할
   (
       "system",
       "You're a data visualization expert using Plotly, a graphing library. "
       "The data to analyze is provided as a StudentsPerformance.csv file. Here are the selected data set: {csv_string}"
       "Follow the user's indications when creating the graph. Do not include fig.show() in your code. "
       "After providing the code for your graph, please write a structured and detailed description of your graph in Korean. "
       "Start your description with '### 그래프 설명' on a new line after the code block. "
       "Organize your explanation into the following categories, each starting with a subtitle in bold:\n"
       "**1. 주요 통계:**\n"
       "**2. 데이터 분포:**\n"
       "**3. 주목할 만한 점:**\n"
       "**4. 결론 및 인사이트:**\n"
       "Under each category, present information as bullet points, with each point starting with '• ' and followed by a line break. "
       "Limit your explanation to 2-3 bullet points per category, prioritizing the most important information. "
       "Include specific numeric data in your explanation, such as averages, maximums, minimums, or other relevant statistics. "
       "In the conclusion, directly address the user's specific question or request about the data. "
       "Provide insights that are relevant to the user's query and the visualized data. "
       "Ensure that your description accurately reflects the data shown in the generated graph and answers the user's question. "
   ),
   MessagesPlaceholder(variable_name="chat_history"),  # 대화 기록 플레이스홀더
   ("human", "{question}"),  # 사용자 질문 플레이스홀더
])

# RunnablePassthrough를 사용해 체인 설정, 대화 기록 불러오기
chain = RunnablePassthrough.assign(chat_history=load_memory) | prompt | llm

# 체인 실행 함수 정의
def invoke_chain(question):
   result = chain.invoke({"question": question, "csv_string": csv_string}) # question, csv_string를 체인에 전달, 결과 반환
   memory.save_context( # 메모리에 입력된 질문과 체인 결과 저장
       {"input": question},
       {"output": result.content},
   )
   return result.content

# Dash 초기화 및 레이아웃 설정
app = Dash()  
app.layout = html.Div([
   # 헤더 설정
   html.H2([
       html.Img(src=app.get_asset_url('chart.gif'), style={'height': '30px', 'marginRight': '10px', 'verticalAlign': 'middle'}),
       "AI 기반 데이터 시각화 도구"
   ], style={'marginBottom': '10px'}),    

   # 데이터 그리드 표시
   html.Div([
       dag.AgGrid(
           rowData=df_select.to_dict("records"),
           columnDefs=[{"field": i} for i in df_select.columns],
           defaultColDef={"filter": True, "sortable": True, "floatingFilter": True},
           style={'height': '30vh', 'width': '100%'}
       )
   ], style={'marginBottom': '20px'}),

  # 사용자 입력 영역
  html.Strong("데이터에서 무엇을 보고 싶으신가요?"),
  dcc.Textarea(id='user-request', style={'width': '100%', 'height': '30px', 'margin-top': '5px'}),
  html.Br(),
  html.Button('Submit', id='my-button'),
  
  # 결과 표시 영역
  dcc.Loading(
      [
           html.Div(id='my-figure', style={'marginTop': '20px', 'marginBottom': '20px'}),
           dcc.Markdown(id='content', children='', style={'whiteSpace': 'pre-wrap', 'wordBreak': 'break-word'})
      ],
      type='default'
  )  
], style={ # 여백을 위한 CSS 설정
   'margin-left': 'auto', 
   'margin-right': 'auto',
   'max-width': '1600px',
   'padding': '0 20px',  
   'paddingBottom': '30px'
})

# 코드에서 그래프 객체 추출 함수 정의
def get_fig_from_code(code):
   local_variables = {}
   exec(code, globals(), local_variables)
   return local_variables['fig']

# 콜백 함수 정의
@callback(
   Output('my-figure', 'children'), # 콜백 함수가 반환할 첫 번째 출력       
   Output('content', 'children'), # 두 번째 출력
   Input('my-button', 'n_clicks'), # submit 버튼을 누르면 콜백 함수가 반응
   State('user-request', 'value'), # 콜백 함수는 사용자가 입력한 값을 사용
   prevent_initial_call=True # 초기 페이지 로드 시 콜백 함수 실행 방지
)
def create_graph(_, user_input):
   result_output = invoke_chain(user_input) # 체인 실행 결과 반환
   # print(result_output)  # 디버깅용 출력

   # 코드 블록 추출 및 처리
   code_block_match = re.search(r'```(?:python)?(.*?)```', result_output, re.DOTALL | re.IGNORECASE)
   if code_block_match:
       code_block = code_block_match.group(1).strip() # 코드 블록을 변수에 저장
       cleaned_code = re.sub(r'(?m)^\s+fig\.show\(\)\s+$', '', code_block) # 코드 블록에서 불필요한 부분을 제거
       explanation = re.sub(r'```(?:python)?.*?```', '', result_output, flags=re.DOTALL | re.IGNORECASE).strip() # 설명 부분을 추출

       fig = get_fig_from_code(cleaned_code) # 추출한 코드 블록을 실행하여 그래프 객체 생성
       return_value = (dcc.Graph(figure=fig), explanation) # 그래프와 설명을 반환값으로 설정
       return return_value
   
   else:
       return "", "그래프를 생성할 수 없습니다."
   
# 메인 실행 부분
if __name__ == '__main__':
  app.run_server(debug=False, port=8887)  # 디버그 모드 비활성화, 포트 8887에서 서버 실행