app.py

import streamlit as st
import pandas as pd
import plotly.express as px
from pandasai import Agent
from langchain_community.embeddings.openai import OpenAIEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_openai import ChatOpenAI
from langchain.chains import RetrievalQA
from langchain.schema import Document
import os

# 제목 설정
st.title("PandasAI 데이터 분석기 with RAG")

# 사이드바에 API 키 입력 필드 추가
api_key = st.sidebar.text_input("OpenAI API Key", type="password")
pandasai_api_key = st.sidebar.text_input("PandasAI API Key", type="password")

# 파일 업로드
uploaded_file = st.file_uploader("엑셀 또는 CSV 파일을 업로드하세요", type=["xlsx", "csv"])

if uploaded_file is not None and api_key and pandasai_api_key:
    # API 키 설정
    os.environ["OPENAI_API_KEY"] = api_key
    os.environ["PANDASAI_API_KEY"] = pandasai_api_key

    # 데이터 로드
    if uploaded_file.name.endswith('.xlsx'):
        df = pd.read_excel(uploaded_file)
    else:
        df = pd.read_csv(uploaded_file)

    st.write("데이터 미리보기:")
    st.write(df.head())

    # PandasAI Agent 설정
    agent = Agent(df)

    # 데이터프레임을 문서로 변환
    documents = [
        Document(
            page_content=", ".join([f"{col}: {row[col]}" for col in df.columns]),
            metadata={"index": index}
        )
        for index, row in df.iterrows()
    ]

    # RAG 설정
    embeddings = OpenAIEmbeddings()
    vectorstore = FAISS.from_documents(documents, embeddings)
    retriever = vectorstore.as_retriever()
    qa_chain = RetrievalQA.from_chain_type(
        llm=ChatOpenAI(),
        chain_type="stuff",
        retriever=retriever
    )

    # 탭 생성
    tab1, tab2, tab3 = st.tabs(["PandasAI 분석", "RAG 질문답변", "데이터 시각화"])

    with tab1:
        st.header("PandasAI를 사용한 데이터 분석")
        pandas_question = st.text_input("데이터에 대해 질문하세요 (PandasAI):")
        if pandas_question:
            result = agent.chat(pandas_question)
            st.write("PandasAI 답변:", result)

    with tab2:
        st.header("RAG를 사용한 질문답변")
        rag_question = st.text_input("데이터에 대해 질문하세요 (RAG):")
        if rag_question:
            result = qa_chain.run(rag_question)
            st.write("RAG 답변:", result)

    with tab3:
        st.header("데이터 시각화")
        viz_question = st.text_input("어떤 그래프를 그리고 싶으신가요? (예: '연봉과 경력의 관계를 산점도로 보여줘')")
        
        if viz_question:
            try:
                result = agent.chat(viz_question)
                
                # PandasAI의 결과가 문자열이므로, 이를 실행 가능한 코드로 변환
                import re
                code_pattern = r'```python\n(.*?)\n```'
                code_match = re.search(code_pattern, result, re.DOTALL)
                
                if code_match:
                    viz_code = code_match.group(1)
                    # 'plt' 대신 'px'를 사용하도록 코드 수정
                    viz_code = viz_code.replace('plt.', 'px.')
                    viz_code = viz_code.replace('plt.show()', 'fig = px.scatter(df, x=x, y=y)')
                    
                    # 코드 실행 및 그래프 표시
                    exec(viz_code)
                    st.plotly_chart(fig)
                else:
                    st.write("그래프를 생성하지 못했습니다. 다른 질문을 해보세요.")
            except Exception as e:
                st.write(f"오류가 발생했습니다: {str(e)}")
                st.write("다른 방식으로 질문해보세요.")

elif not api_key:
    st.warning("OpenAI API 키를 입력해주세요.")
elif not pandasai_api_key:
    st.warning("PandasAI API 키를 입력해주세요.")