# 프로젝트 폴더 생성
mkdir company-rag-system
cd company-rag-system

# 가상환경 생성
python -m venv venv

# 가상환경 활성화
# Windows:
venv\Scripts\activate
# macOS/Linux:
source venv/bin/activate

# 필요한 패키지 설치
pip install langchain chromadb streamlit ollama pypdf python-docx openpyxl

company-rag-system/
├── app.py              # 메인 애플리케이션
├── document_loader.py  # 문서 로더
├── rag_engine.py      # RAG 엔진
├── documents/         # 회사 문서 폴더
└── chroma_db/        # 벡터 DB 저장소

import os
from typing import List
from langchain.document_loaders import (
    PyPDFLoader,
    Docx2txtLoader,
    UnstructuredExcelLoader,
    TextLoader
)
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.schema import Document

class DocumentLoader:
    """다양한 형식의 문서를 로드하는 클래스"""
    
    def __init__(self, docs_path: str = "./documents"):
        self.docs_path = docs_path
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200,
            separators=["\n\n", "\n", ".", "!", "?", ",", " ", ""],
            length_function=len,
        )
    
    def load_documents(self) -> List[Document]:
        """문서 폴더의 모든 문서를 로드"""
        all_documents = []
        
        for root, dirs, files in os.walk(self.docs_path):
            for file in files:
                file_path = os.path.join(root, file)
                print(f"📄 로딩 중: {file}")
                
                try:
                    # PDF 파일
                    if file.endswith('.pdf'):
                        loader = PyPDFLoader(file_path)
                    # Word 파일
                    elif file.endswith('.docx'):
                        loader = Docx2txtLoader(file_path)
                    # Excel 파일
                    elif file.endswith(('.xlsx', '.xls')):
                        loader = UnstructuredExcelLoader(file_path)
                    # 텍스트 파일
                    elif file.endswith('.txt'):
                        loader = TextLoader(file_path, encoding='utf-8')
                    else:
                        continue
                    
                    documents = loader.load()
                    # 문서를 작은 청크로 분할
                    splits = self.text_splitter.split_documents(documents)
                    
                    # 메타데이터에 파일명 추가
                    for split in splits:
                        split.metadata['source'] = file
                    
                    all_documents.extend(splits)
                    
                except Exception as e:
                    print(f"❌ 오류 발생 ({file}): {str(e)}")
        
        print(f"✅ 총 {len(all_documents)}개의 문서 청크 로드 완료!")
        return all_documents

from typing import List, Tuple
from langchain.embeddings import OllamaEmbeddings
from langchain.vectorstores import Chroma
from langchain.llms import Ollama
from langchain.chains import RetrievalQA
from langchain.schema import Document
import chromadb

class RAGEngine:
    """RAG 검색 엔진 클래스"""
    
    def __init__(self, model_name: str = "llama2", persist_directory: str = "./chroma_db"):
        self.model_name = model_name
        self.persist_directory = persist_directory
        
        # Ollama 임베딩 모델 초기화
        self.embeddings = OllamaEmbeddings(
            model=model_name,
            base_url="http://localhost:11434"
        )
        
        # Ollama LLM 초기화
        self.llm = Ollama(
            model=model_name,
            base_url="http://localhost:11434",
            temperature=0.3,  # 더 정확한 답변을 위해 낮은 온도 설정
        )
        
        self.vectorstore = None
        self.qa_chain = None
    
    def create_vectorstore(self, documents: List[Document]):
        """문서로부터 벡터 저장소 생성"""
        print("🔄 벡터 데이터베이스 생성 중...")
        
        # ChromaDB 클라이언트 설정
        client = chromadb.PersistentClient(path=self.persist_directory)
        
        # 벡터 저장소 생성
        self.vectorstore = Chroma.from_documents(
            documents=documents,
            embedding=self.embeddings,
            persist_directory=self.persist_directory,
            client=client,
            collection_name="company_docs"
        )
        
        print("✅ 벡터 데이터베이스 생성 완료!")
    
    def load_vectorstore(self):
        """기존 벡터 저장소 로드"""
        print("🔄 기존 벡터 데이터베이스 로드 중...")
        
        client = chromadb.PersistentClient(path=self.persist_directory)
        
        self.vectorstore = Chroma(
            persist_directory=self.persist_directory,
            embedding_function=self.embeddings,
            client=client,
            collection_name="company_docs"
        )
        
        print("✅ 벡터 데이터베이스 로드 완료!")
    
    def setup_qa_chain(self):
        """QA 체인 설정"""
        if not self.vectorstore:
            raise ValueError("벡터 저장소가 초기화되지 않았습니다!")
        
        # 프롬프트 템플릿 설정
        from langchain.prompts import PromptTemplate
        
        prompt_template = """주어진 문서를 참고하여 질문에 답변해주세요. 
        답변은 한국어로 작성하고, 문서에 없는 내용은 추측하지 마세요.
        
        문서 내용:
        {context}
        
        질문: {question}
        
        답변:"""
        
        PROMPT = PromptTemplate(
            template=prompt_template, 
            input_variables=["context", "question"]
        )
        
        # RetrievalQA 체인 생성
        self.qa_chain = RetrievalQA.from_chain_type(
            llm=self.llm,
            chain_type="stuff",
            retriever=self.vectorstore.as_retriever(
                search_kwargs={"k": 4}  # 상위 4개 문서 검색
            ),
            return_source_documents=True,
            chain_type_kwargs={"prompt": PROMPT}
        )
    
    def query(self, question: str) -> Tuple[str, List[str]]:
        """질문에 대한 답변과 출처 반환"""
        if not self.qa_chain:
            raise ValueError("QA 체인이 설정되지 않았습니다!")
        
        # 질문 처리
        result = self.qa_chain({"query": question})
        
        # 답변
        answer = result['result']
        
        # 출처 문서들
        sources = []
        for doc in result['source_documents']:
            source = doc.metadata.get('source', 'Unknown')
            if source not in sources:
                sources.append(source)
        
        return answer, sources

import streamlit as st
import os
from datetime import datetime
from document_loader import DocumentLoader
from rag_engine import RAGEngine

# 페이지 설정
st.set_page_config(
    page_title="회사 문서 AI 검색 시스템",
    page_icon="🔍",
    layout="wide"
)

# 세션 상태 초기화
if 'rag_engine' not in st.session_state:
    st.session_state.rag_engine = None
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []

def initialize_rag_engine():
    """RAG 엔진 초기화"""
    with st.spinner("🚀 AI 엔진을 초기화하는 중..."):
        try:
            rag_engine = RAGEngine()
            
            # 기존 DB가 있는지 확인
            if os.path.exists("./chroma_db"):
                rag_engine.load_vectorstore()
            else:
                st.warning("⚠️ 문서 데이터베이스가 없습니다. 문서를 업로드해주세요.")
                return None
            
            rag_engine.setup_qa_chain()
            return rag_engine
            
        except Exception as e:
            st.error(f"초기화 중 오류 발생: {str(e)}")
            return None

def main():
    st.title("🏢 회사 문서 AI 검색 시스템")
    st.markdown("---")
    
    # 사이드바: 문서 관리
    with st.sidebar:
        st.header("📁 문서 관리")
        
        # 문서 업로드
        uploaded_files = st.file_uploader(
            "문서 업로드",
            type=['pdf', 'docx', 'txt', 'xlsx', 'xls'],
            accept_multiple_files=True
        )
        
        if uploaded_files:
            if st.button("📤 업로드한 문서 처리하기"):
                # documents 폴더 생성
                os.makedirs("documents", exist_ok=True)
                
                # 파일 저장
                for uploaded_file in uploaded_files:
                    file_path = os.path.join("documents", uploaded_file.name)
                    with open(file_path, "wb") as f:
                        f.write(uploaded_file.getbuffer())
                
                # 문서 로드 및 벡터화
                with st.spinner("문서를 처리하는 중..."):
                    loader = DocumentLoader()
                    documents = loader.load_documents()
                    
                    if documents:
                        rag_engine = RAGEngine()
                        rag_engine.create_vectorstore(documents)
                        rag_engine.setup_qa_chain()
                        st.session_state.rag_engine = rag_engine
                        st.success("✅ 문서 처리 완료!")
                        st.rerun()
        
        st.markdown("---")
        
        # 현재 등록된 문서 표시
        if os.path.exists("documents"):
            st.subheader("📚 등록된 문서")
            files = os.listdir("documents")
            for file in files:
                st.text(f"• {file}")
        
        # 채팅 기록 초기화
        if st.button("🗑️ 대화 기록 삭제"):
            st.session_state.chat_history = []
            st.rerun()
    
    # 메인 영역: 채팅 인터페이스
    col1, col2 = st.columns([2, 1])
    
    with col1:
        st.header("💬 AI와 대화하기")
        
        # RAG 엔진 초기화
        if st.session_state.rag_engine is None:
            st.session_state.rag_engine = initialize_rag_engine()
        
        # 채팅 인터페이스
        if st.session_state.rag_engine:
            # 질문 입력
            user_question = st.text_input(
                "질문을 입력하세요:",
                placeholder="예: 우리 회사의 연차 사용 규정이 어떻게 되나요?"
            )
            
            if st.button("🔍 검색") and user_question:
                with st.spinner("답변을 생성하는 중..."):
                    try:
                        # 답변 생성
                        answer, sources = st.session_state.rag_engine.query(user_question)
                        
                        # 채팅 기록에 추가
                        st.session_state.chat_history.append({
                            'timestamp': datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                            'question': user_question,
                            'answer': answer,
                            'sources': sources
                        })
                        
                    except Exception as e:
                        st.error(f"오류 발생: {str(e)}")
            
            # 채팅 기록 표시
            st.markdown("---")
            st.subheader("📝 대화 기록")
            
            for chat in reversed(st.session_state.chat_history):
                with st.container():
                    st.markdown(f"**🕐 {chat['timestamp']}**")
                    st.markdown(f"**👤 질문:** {chat['question']}")
                    st.markdown(f"**🤖 답변:** {chat['answer']}")
                    if chat['sources']:
                        st.markdown(f"**📌 출처:** {', '.join(chat['sources'])}")
                    st.markdown("---")
        else:
            st.info("👈 왼쪽 사이드바에서 문서를 업로드해주세요.")
    
    with col2:
        st.header("📊 시스템 상태")
        
        # 시스템 상태 표시
        if st.session_state.rag_engine:
            st.success("✅ AI 엔진 활성화")
            
            # 벡터 저장소 정보
            if st.session_state.rag_engine.vectorstore:
                collection = st.session_state.rag_engine.vectorstore._collection
                st.metric("저장된 문서 청크", collection.count())
        else:
            st.warning("⚠️ AI 엔진 비활성화")
        
        # 사용 가이드
        st.markdown("---")
        st.subheader("💡 사용 방법")
        st.markdown("""
        1. **문서 업로드**: 왼쪽 사이드바에서 회사 문서를 업로드
        2. **질문하기**: 자연스러운 한국어로 질문 입력
        3. **답변 확인**: AI가 문서에서 찾은 정보로 답변 생성
        
        **지원 파일 형식:**
        - PDF (.pdf)
        - Word (.docx)
        - Excel (.xlsx, .xls)
        - Text (.txt)
        """)

if __name__ == "__main__":
    main()

# Ollama가 실행 중인지 확인
ollama list

# 만약 llama2가 없다면 다운로드
ollama pull llama2

# 프로젝트 폴더에서 실행
streamlit run app.py

ConnectionError: Failed to connect to Ollama

# Ollama 서비스 재시작
# Windows
ollama serve

# macOS
brew services restart ollama

# TextLoader에 인코딩 명시
loader = TextLoader(file_path, encoding='utf-8-sig')

# 검색 결과 개수 늘리기
retriever = vectorstore.as_retriever(search_kwargs={"k": 6})

# 유사도 임계값 설정
retriever = vectorstore.as_retriever(
    search_type="similarity_score_threshold",
    search_kwargs={"score_threshold": 0.7}
)

# 질문 언어 감지 및 번역
from langdetect import detect
if detect(question) != 'ko':
    # 번역 API 활용
    pass

# 업로드한 문서의 요약 생성
from langchain.chains.summarize import load_summarize_chain
summary_chain = load_summarize_chain(llm, chain_type="map_reduce")

# 문서 기반 자동 질문 생성
suggested_questions = [
    "이 문서의 주요 내용은 무엇인가요?",
    "핵심 정책은 무엇인가요?",
    # ...
]