app.py

"""
RFx 문서 분석 AI 에이전트 (PDF Text Highlighting)
PDF 텍스트에 직접 하이라이트 표시
"""
import streamlit as st
import fitz  # PyMuPDF
import chromadb
from sentence_transformers import SentenceTransformer, util
import requests
import os
import re
import shutil
from collections import Counter
import numpy as np
from typing import List, Dict, Tuple
import base64

GROK_API_KEY = os.getenv("GROK_API_KEY")
GROK_API_BASE = "https://api.x.ai/v1"
CHROMA_DIR = "./chroma_db"
EMBEDDING_MODEL = 'jhgan/ko-sroberta-multitask'

st.set_page_config(
    page_title="RFx 문서 분석 AI",
    page_icon="📄",
    layout="wide",
    initial_sidebar_state="collapsed"
)

st.markdown("""
<style>
    .main-title {
        font-size: 1.8rem; 
        font-weight: bold; 
        color: #1E3A8A; 
        margin-bottom: 1rem;
        text-align: center;
    }
    .source-box {
        background: #F1F5F9; 
        padding: 1rem; 
        border-radius: 0.5rem; 
        margin: 0.5rem 0; 
        border-left: 3px solid #3B82F6;
    }
    .source-title {
        font-weight: bold; 
        color: #1E40AF; 
        margin-bottom: 0.5rem;
    }
    .keyword-badge {
        display: inline-block;
        background: #DBEAFE;
        color: #1E40AF;
        padding: 0.2rem 0.6rem;
        border-radius: 0.3rem;
        margin: 0.2rem;
        font-size: 0.85rem;
    }
    .pdf-container {
        border: 2px solid #E2E8F0;
        border-radius: 0.5rem;
        padding: 0.5rem;
        height: 800px;
        overflow-y: auto;
        background: white;
    }
    .page-indicator {
        background: #3B82F6;
        color: white;
        padding: 0.3rem 0.8rem;
        border-radius: 1rem;
        font-size: 0.85rem;
        display: inline-block;
        margin: 0.2rem;
    }
    .highlight-indicator {
        background: #FEF08A;
        color: #854D0E;
        padding: 0.5rem 1rem;
        border-radius: 0.5rem;
        margin: 0.5rem 0;
        font-weight: bold;
        border-left: 4px solid #EAB308;
    }
</style>
""", unsafe_allow_html=True)


def init_session():
    if 'processed' not in st.session_state:
        st.session_state.processed = False
    if 'vector_db' not in st.session_state:
        st.session_state.vector_db = None
    if 'embedder' not in st.session_state:
        st.session_state.embedder = None
    if 'chat_history' not in st.session_state:
        st.session_state.chat_history = []
    if 'doc_metadata' not in st.session_state:
        st.session_state.doc_metadata = {}
    if 'pdf_bytes' not in st.session_state:
        st.session_state.pdf_bytes = None
    if 'pdf_pages_text' not in st.session_state:
        st.session_state.pdf_pages_text = {}
    if 'current_highlights' not in st.session_state:
        st.session_state.current_highlights = []


def extract_text_from_pdf(pdf_file) -> Tuple[List[str], List[Dict], bytes, Dict]:
    pdf_bytes = pdf_file.read()
    doc = fitz.open(stream=pdf_bytes, filetype="pdf")
    
    chunks = []
    metadata_list = []
    pages_text = {}
    
    CHUNK_SIZE = 300
    OVERLAP_SIZE = 60
    
    for page_num in range(len(doc)):
        page = doc[page_num]
        text = page.get_text("text")
        pages_text[page_num + 1] = text
        
        if not text.strip():
            continue
        
        lines = [line.strip() for line in text.split('\n') if line.strip()]
        cleaned_text = '\n'.join(lines)
        
        sentences = re.split(r'([.!?]\s+|\n{2,})', cleaned_text)
        sentences = [s for s in sentences if s.strip()]
        
        current_chunk = ""
        current_length = 0
        
        for sentence in sentences:
            sentence_length = len(sentence)
            
            if current_length + sentence_length > CHUNK_SIZE and current_chunk:
                chunks.append(current_chunk.strip())
                metadata_list.append({
                    "page": page_num + 1,
                    "source": pdf_file.name,
                    "chunk_type": "paragraph"
                })
                
                overlap_text = current_chunk[-OVERLAP_SIZE:] if len(current_chunk) > OVERLAP_SIZE else current_chunk
                current_chunk = overlap_text + sentence
                current_length = len(current_chunk)
            else:
                current_chunk += sentence
                current_length += sentence_length
        
        if current_chunk.strip():
            chunks.append(current_chunk.strip())
            metadata_list.append({
                "page": page_num + 1,
                "source": pdf_file.name,
                "chunk_type": "paragraph"
            })
    
    doc.close()
    return chunks, metadata_list, pdf_bytes, pages_text


@st.cache_resource
def load_embedding_model():
    return SentenceTransformer(EMBEDDING_MODEL)


def create_vector_db(chunks: List[str], metadata_list: List[Dict]):
    embedder = load_embedding_model()
    
    if os.path.exists(CHROMA_DIR):
        try:
            shutil.rmtree(CHROMA_DIR)
        except Exception:
            pass
    
    client = chromadb.PersistentClient(
        path=CHROMA_DIR,
        settings=chromadb.Settings(
            anonymized_telemetry=False,
            allow_reset=True,
            is_persistent=True
        )
    )
    
    try:
        collection = client.get_or_create_collection(
            name="rfx_docs",
            metadata={"hnsw:space": "cosine"}
        )
    except Exception:
        try:
            client.delete_collection("rfx_docs")
        except Exception:
            pass
        collection = client.create_collection(
            name="rfx_docs",
            metadata={"hnsw:space": "cosine"}
        )
    
    batch_size = 32
    all_embeddings = []
    
    for i in range(0, len(chunks), batch_size):
        batch = chunks[i:i + batch_size]
        embeddings = embedder.encode(batch, show_progress_bar=False, convert_to_numpy=True)
        all_embeddings.extend(embeddings)
    
    ids = [f"doc_{i}" for i in range(len(chunks))]
    collection.add(
        embeddings=[emb.tolist() for emb in all_embeddings],
        documents=chunks,
        metadatas=metadata_list,
        ids=ids
    )
    
    return collection, embedder


def extract_keywords(text: str, top_n: int = 5) -> List[str]:
    words_with_numbers = re.findall(r'[가-힣]*\d+[가-힣]*', text)
    words = re.findall(r'[가-힣]{2,}', text)
    
    stopwords = {
        '것', '등', '및', '그', '이', '저', '수', '때', '중', '내', '년', '월', '일',
        '경우', '대한', '통해', '위해', '관련', '있는', '하는', '되는', '이런', '저런',
        '어떤', '무슨', '어느', '누구', '언제', '어디', '무엇', '어떻게', '왜',
        '알려', '설명', '말해', '대해', '관하여', '있나요', '인가요', '무엇인가요',
        '얼마', '입니까', '합니까'
    }
    
    important_keywords = {
        '금액', '가격', '비용', '예산', '설계', '사업', '과업', '계약',
        '공사', '용역', '제안', '입찰', '낙찰', '견적', '단가'
    }
    
    filtered_words = [w for w in words if w not in stopwords and len(w) >= 2]
    word_freq = Counter(filtered_words)
    
    for word in word_freq:
        if word in important_keywords:
            word_freq[word] += 5
    
    result = []
    result.extend([w for w in words_with_numbers if w])
    
    for word, _ in word_freq.most_common(top_n * 2):
        if word not in result:
            result.append(word)
        if len(result) >= top_n:
            break
    
    return result[:top_n]


def rewrite_query(query: str) -> Dict[str, any]:
    original = query.strip()
    cleaned = re.sub(r'[?!,.~]', '', original)
    keywords = extract_keywords(cleaned, top_n=7)
    
    variations = []
    variations.append(original)
    
    if keywords:
        if len(keywords) >= 2:
            variations.append(' '.join(keywords[:2]))
        if len(keywords) >= 3:
            variations.append(' '.join(keywords[:3]))
    
    for kw in keywords[:3]:
        if kw not in variations:
            variations.append(kw)
    
    synonym_map = {
        '금액': ['가격', '비용', '예산'],
        '설계': ['디자인', '계획'],
        '사업': ['프로젝트', '과업'],
    }
    
    for keyword in keywords[:2]:
        if keyword in synonym_map:
            for syn in synonym_map[keyword]:
                combined = original.replace(keyword, syn)
                if combined not in variations:
                    variations.append(combined)
                    break
    
    seen = set()
    unique_variations = []
    for v in variations:
        if v not in seen and v.strip():
            seen.add(v)
            unique_variations.append(v)
    
    return {
        'original': original,
        'cleaned': cleaned,
        'keywords': keywords,
        'variations': unique_variations[:7]
    }


def search_with_multiple_queries(queries: List[str], collection, embedder, top_k: int = 5) -> Dict:
    all_results = []
    seen_ids = set()
    
    for query in queries:
        query_embedding = embedder.encode([query], convert_to_numpy=True)[0]
        
        results = collection.query(
            query_embeddings=[query_embedding.tolist()],
            n_results=min(top_k * 5, 30),
            include=["documents", "metadatas", "distances"]
        )
        
        for i, doc_id in enumerate(results['ids'][0]):
            if doc_id not in seen_ids:
                seen_ids.add(doc_id)
                all_results.append({
                    'id': doc_id,
                    'document': results['documents'][0][i],
                    'metadata': results['metadatas'][0][i],
                    'distance': results['distances'][0][i],
                    'query': query
                })
    
    all_results.sort(key=lambda x: x['distance'])
    top_results = all_results[:top_k]
    
    return {
        'documents': [[r['document'] for r in top_results]],
        'metadatas': [[r['metadata'] for r in top_results]],
        'distances': [[r['distance'] for r in top_results]],
        'queries_used': queries,
        'total_found': len(all_results)
    }


def rerank_results(query: str, search_results: Dict, embedder, keywords: List[str]) -> Dict:
    docs = search_results['documents'][0]
    metas = search_results['metadatas'][0]
    distances = search_results['distances'][0]
    
    if not docs:
        return {
            'documents': [[]],
            'metadatas': [[]],
            'distances': [[]],
            'scores': []
        }
    
    query_embedding = embedder.encode([query], convert_to_numpy=True)[0]
    doc_embeddings = embedder.encode(docs, convert_to_numpy=True)
    
    similarities = util.cos_sim(query_embedding, doc_embeddings)[0].cpu().numpy()
    
    keyword_scores = []
    for doc in docs:
        doc_lower = doc.lower()
        score = sum(1 for kw in keywords if kw.lower() in doc_lower)
        keyword_scores.append(score)
    
    if max(keyword_scores) > 0:
        keyword_scores = [s / max(keyword_scores) for s in keyword_scores]
    
    numeric_query_terms = ['금액', '예산', '가격', '비용', '단가']
    is_numeric_query = any(term in query for term in numeric_query_terms)
    
    if is_numeric_query:
        money_patterns = [
            r'\d{1,3}(?:,\d{3})+원',
            r'\d+만원',
            r'\d+억원',
            r'\(일금\s*[^)]+\)'
        ]
        numeric_scores = []
        for doc in docs:
            score = 0
            for pattern in money_patterns:
                if re.search(pattern, doc):
                    score = 1
                    break
            numeric_scores.append(score)
        if max(numeric_scores) > 0:
            numeric_scores = [s / max(numeric_scores) for s in numeric_scores]
        else:
            numeric_scores = [0.0 for _ in numeric_scores]
        
        final_scores = [
            0.6 * sim + 0.25 * kw + 0.15 * num
            for sim, kw, num in zip(similarities, keyword_scores, numeric_scores)
        ]
    else:
        final_scores = [0.7 * sim + 0.3 * kw for sim, kw in zip(similarities, keyword_scores)]
    
    ranked_indices = np.argsort(final_scores)[::-1]
    
    return {
        'documents': [[docs[i] for i in ranked_indices]],
        'metadatas': [[metas[i] for i in ranked_indices]],
        'distances': [[distances[i] for i in ranked_indices]],
        'scores': [final_scores[i] for i in ranked_indices]
    }


def build_context(search_results: Dict, max_length: int = 3000) -> str:
    context_parts = []
    current_length = 0
    
    docs = search_results['documents'][0]
    metas = search_results['metadatas'][0]
    
    for i, (doc, meta) in enumerate(zip(docs, metas), 1):
        part = f"[문서 {i}] (페이지 {meta['page']})\n{doc}\n"
        part_length = len(part)
        
        if current_length + part_length > max_length:
            remaining = max_length - current_length
            if remaining > 200:
                part = f"[문서 {i}] (페이지 {meta['page']})\n{doc[:remaining-50]}...\n"
                context_parts.append(part)
            break
        
        context_parts.append(part)
        current_length += part_length
    
    return "\n".join(context_parts)


def generate_answer(query: str, search_results: Dict, api_key: str) -> str:
    context = build_context(search_results, max_length=4000)
    
    system_prompt = """당신은 RFx 문서 전문 분석가입니다.

**중요 원칙:**
1. 제공된 문서를 **매우 꼼꼼히** 읽고 정확한 정보를 찾으세요
2. 숫자, 금액, 날짜 등 구체적인 정보를 우선적으로 찾으세요
3. 문서에 정보가 있는데도 "없다"고 하지 마세요
4. 답변 시 반드시 [문서 N, 페이지 X] 형태로 출처 명시
5. 애매한 표현 대신 구체적인 수치를 제공하세요

**답변 형식:**
- 핵심 답변을 먼저 명확하게 제시
- 출처 명시 (페이지 번호 포함)
- 필요시 추가 관련 정보 제공"""
    
    user_prompt = f"""다음 문서들을 **매우 꼼꼼히** 읽고 질문에 답변하세요.

<문서>
{context}
</문서>

<질문>
{query}
</질문>

**중요**: 
- 문서를 처음부터 끝까지 주의 깊게 읽으세요
- 숫자, 금액 등 구체적인 정보를 찾으세요
- 찾은 정보는 정확히 인용하세요
- 정말로 문서에 없는 경우에만 "찾을 수 없습니다"라고 하세요"""
    
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }
    
    payload = {
        "model": "grok-3",
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ],
        "temperature": 0.1,
        "max_tokens": 2000,
        "stream": False
    }
    
    try:
        response = requests.post(
            f"{GROK_API_BASE}/chat/completions",
            headers=headers,
            json=payload,
            timeout=30
        )
        
        if response.status_code != 200:
            error_detail = ""
            try:
                error_data = response.json()
                error_detail = error_data.get('error', {}).get('message', '')
            except Exception:
                error_detail = response.text
            
            return f"❌ API 오류 (코드: {response.status_code})\n\n{error_detail}"
        
        result = response.json()
        return result["choices"][0]["message"]["content"]
        
    except Exception as e:
        return f"❌ 오류: {str(e)}"


def highlight_text_in_pdf(pdf_bytes: bytes, highlight_info: List[Dict]) -> bytes:
    doc = fitz.open(stream=pdf_bytes, filetype="pdf")
    
    for item in highlight_info:
        page_num = item['page'] - 1
        search_text = item['text']
        
        if page_num >= len(doc):
            continue
        
        page = doc[page_num]
        
        text_variations = [
            search_text,
            search_text.replace(' ', ''),
            search_text.replace(',', ''),
        ]
        
        for text_var in text_variations:
            text_instances = page.search_for(text_var)
            
            for inst in text_instances:
                highlight = page.add_highlight_annot(inst)
                highlight.set_colors(stroke=[1, 1, 0])
                highlight.update()
    
    output_bytes = doc.tobytes()
    doc.close()
    
    return output_bytes


def extract_highlight_texts(documents: List[str], keywords: List[str]) -> List[str]:
    highlight_texts = []
    
    for doc in documents:
        money_patterns = [
            r'\d{1,3}(?:,\d{3})+원',
            r'\d+만원',
            r'\d+억원',
            r'\(일금\s*[^)]+\)',
        ]
        
        for pattern in money_patterns:
            matches = re.findall(pattern, doc)
            highlight_texts.extend(matches)
        
        date_patterns = [
            r'\d{4}[년.]\d{1,2}[월.]\d{1,2}일?',
            r'\d{2}\.\d{2}\.\d{2}',
        ]
        
        for pattern in date_patterns:
            matches = re.findall(pattern, doc)
            highlight_texts.extend(matches)
        
        for keyword in keywords:
            if keyword in doc:
                sentences = re.split(r'[.!?]\s+', doc)
                for sent in sentences:
                    if keyword in sent and len(sent) < 100:
                        highlight_texts.append(sent.strip())
    
    unique_texts = list(set(highlight_texts))
    unique_texts.sort(key=len)
    
    return unique_texts[:10]


def render_pdf_with_highlights(pdf_bytes: bytes, highlight_info: List[Dict]):
    highlighted_pdf = highlight_text_in_pdf(pdf_bytes, highlight_info)
    
    doc = fitz.open(stream=highlighted_pdf, filetype="pdf")
    
    highlighted_pages = set(h['page'] for h in highlight_info)
    
    pdf_html = '<div class="pdf-container">'
    
    for page_num in range(len(doc)):
        page = doc[page_num]
        
        pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))
        img_data = pix.tobytes("png")
        img_base64 = base64.b64encode(img_data).decode()
        
        pdf_html += '<div style="margin-bottom: 2rem; position: relative;">'
        pdf_html += f'<div style="background: #3B82F6; color: white; padding: 0.5rem; margin-bottom: 0.5rem; border-radius: 0.3rem; font-weight: bold;">📄 페이지 {page_num + 1}</div>'
        
        if (page_num + 1) in highlighted_pages:
            page_highlights = [h for h in highlight_info if h['page'] == page_num + 1]
            highlight_texts = ', '.join([f'"{h["text"][:30]}..."' for h in page_highlights[:3]])
            pdf_html += f'<div class="highlight-indicator">⭐ 하이라이트: {highlight_texts}</div>'
        
        pdf_html += f'<img src="data:image/png;base64,{img_base64}" style="width: 100%; border: 1px solid #E2E8F0; border-radius: 0.3rem; box-shadow: 0 1px 3px rgba(0,0,0,0.1);" />'
        pdf_html += '</div>'
    
    pdf_html += '</div>'
    doc.close()
    
    return pdf_html


def main():
    init_session()
    
    st.markdown('<div class="main-title">📄 RFx 문서 분석 AI 에이전트</div>', unsafe_allow_html=True)
    
    with st.sidebar:
        st.header("⚙️ 설정")
        grok_key = st.text_input("Grok API Key", value=GROK_API_KEY or "", type="password")
        
        if grok_key:
            os.environ["GROK_API_KEY"] = grok_key
            st.session_state.grok_key = grok_key
        
        st.divider()
        
        if st.button("🔄 데이터베이스 초기화", help="ChromaDB 오류 발생 시 클릭"):
            if os.path.exists(CHROMA_DIR):
                try:
                    shutil.rmtree(CHROMA_DIR)
                    st.success("✅ 데이터베이스 초기화 완료!")
                    st.session_state.processed = False
                    st.session_state.vector_db = None
                    st.rerun()
                except Exception as e:
                    st.error(f"초기화 실패: {str(e)}")
        
        st.divider()
        
        st.subheader("📤 문서 업로드")
        uploaded_file = st.file_uploader("PDF 파일 선택", type=['pdf'])
        
        if uploaded_file:
            if st.button("📄 문서 처리", type="primary", disabled=st.session_state.get('processing', False)):
                if not grok_key:
                    st.error("⚠️ Grok API 키를 입력하세요!")
                    return
                
                st.session_state.processing = True
                
                with st.spinner("📄 문서 처리 중..."):
                    try:
                        chunks, metadata_list, pdf_bytes, pages_text = extract_text_from_pdf(uploaded_file)
                        
                        st.info(f"📑 {len(chunks)}개 청크 추출 완료")
                        
                        with st.expander("📝 추출된 텍스트 샘플", expanded=False):
                            if chunks:
                                st.text(f"첫 번째 청크 (총 {len(chunks[0])}자):")
                                st.code(chunks[0][:500] + "..." if len(chunks[0]) > 500 else chunks[0])
                        
                        with st.spinner("🔧 벡터 데이터베이스 생성 중..."):
                            collection, embedder = create_vector_db(chunks, metadata_list)
                        
                        st.session_state.vector_db = collection
                        st.session_state.embedder = embedder
                        st.session_state.pdf_bytes = pdf_bytes
                        st.session_state.pdf_pages_text = pages_text
                        st.session_state.processed = True
                        st.session_state.doc_metadata = {
                            "filename": uploaded_file.name,
                            "chunks": len(chunks),
                            "pages": len(set(m['page'] for m in metadata_list))
                        }
                        
                        st.success("✅ 문서 처리 완료!")
                        
                    except Exception as e:
                        st.error(f"오류: {str(e)}")
                    finally:
                        st.session_state.processing = False
        
        st.divider()
        
        if st.session_state.processed:
            st.subheader("📊 문서 정보")
            meta = st.session_state.doc_metadata
            st.write(f"**파일명:** {meta['filename']}")
            st.write(f"**페이지:** {meta['pages']}페이지")
            st.write(f"**청크:** {meta['chunks']}개")
        
        if st.button("🗑️ 채팅 초기화"):
            st.session_state.chat_history = []
            st.session_state.current_highlights = []
            st.rerun()
    
    if not st.session_state.processed:
        st.info("👈 왼쪽 사이드바에서 PDF 문서를 업로드하세요")
        
        col1, col2, col3 = st.columns(3)
        with col1:
            st.markdown("### 📄 PDF 뷰어\n원본 문서 확인")
        with col2:
            st.markdown("### 🎨 하이라이트\n핵심 내용 강조")
        with col3:
            st.markdown("### 💬 AI 챗봇\n정확한 답변")
    
    else:
        col1, col2 = st.columns([1, 1])
        
        with col1:
            st.markdown("### 📄 문서 뷰어")
            
            if st.session_state.pdf_bytes:
                pdf_html = render_pdf_with_highlights(
                    st.session_state.pdf_bytes,
                    st.session_state.current_highlights
                )
                st.markdown(pdf_html, unsafe_allow_html=True)
        
        with col2:
            st.markdown("### 💬 AI 챗봇")
            
            chat_container = st.container()
            with chat_container:
                for msg in st.session_state.chat_history:
                    with st.chat_message(msg["role"]):
                        st.markdown(msg["content"])
                        
                        if msg["role"] == "assistant" and "sources" in msg:
                            with st.expander("📚 참조 문서"):
                                for i, (doc, meta) in enumerate(zip(
                                    msg["sources"]["docs"],
                                    msg["sources"]["metas"]
                                ), 1):
                                    score = msg["sources"]["scores"][i-1] if "scores" in msg["sources"] else None
                                    score_text = f" (관련도: {score:.2%})" if score else ""
                                    
                                    st.markdown(f"""
                                    <div class="source-box">
                                        <div class="source-title">
                                            <span class="page-indicator">페이지 {meta['page']}</span>
                                            {score_text}
                                        </div>
                                        <div style="font-size: 0.9rem; color: #475569;">
                                            {doc[:300]}{'...' if len(doc) > 300 else ''}
                                        </div>
                                    </div>
                                    """, unsafe_allow_html=True)
            
            if prompt := st.chat_input("질문을 입력하세요...", disabled=st.session_state.get('processing', False)):
                
                if not st.session_state.get('grok_key'):
                    st.error("⚠️ Grok API 키를 입력해주세요!")
                    return
                
                with st.chat_message("user"):
                    st.markdown(prompt)
                st.session_state.chat_history.append({"role": "user", "content": prompt})
                
                with st.chat_message("assistant"):
                    with st.spinner("🔍 검색 및 분석 중..."):
                        try:
                            query_info = rewrite_query(prompt)
                            
                            with st.expander("🔍 검색 디버그 정보", expanded=False):
                                st.write("**추출된 키워드:**", query_info['keywords'])
                                st.write("**검색 쿼리 변형:**", query_info['variations'])
                                
                            search_results = search_with_multiple_queries(
                                query_info['variations'],
                                st.session_state.vector_db,
                                st.session_state.embedder,
                                top_k=7
                            )
                            
                            with st.expander("📄 검색된 문서 내용", expanded=False):
                                st.write(f"**총 {search_results.get('total_found', 0)}개 문서 발견**")
                                for i, doc in enumerate(search_results['documents'][0][:3], 1):
                                    st.write(f"**문서 {i}:**")
                                    st.text(doc[:300] + "..." if len(doc) > 300 else doc)
                                    st.divider()
                            
                            if 'total_found' in search_results:
                                st.success(f"✅ {search_results['total_found']}개 문서에서 상위 7개 선택")
                            
                            reranked_results = rerank_results(
                                query_info['original'],
                                search_results,
                                st.session_state.embedder,
                                query_info['keywords']
                            )
                            
                            answer = generate_answer(
                                query_info['original'],
                                reranked_results,
                                st.session_state.grok_key
                            )
                            
                            st.markdown(answer)
                            
                            highlight_texts = extract_highlight_texts(
                                reranked_results['documents'][0],
                                query_info['keywords']
                            )
                            
                            highlights = []
                            for doc, meta in zip(reranked_results['documents'][0], 
                                                 reranked_results['metadatas'][0]):
                                for text in highlight_texts:
                                    if text in doc:
                                        highlights.append({
                                            'page': meta['page'],
                                            'text': text
                                        })
                            
                            st.session_state.current_highlights = highlights
                            
                            st.session_state.chat_history.append({
                                "role": "assistant",
                                "content": answer,
                                "sources": {
                                    "docs": reranked_results['documents'][0],
                                    "metas": reranked_results['metadatas'][0],
                                    "scores": reranked_results.get('scores', []),
                                    "keywords": query_info['keywords']
                                }
                            })
                            
                            with st.expander("📚 참조 문서", expanded=True):
                                for i, (doc, meta) in enumerate(zip(
                                    reranked_results['documents'][0],
                                    reranked_results['metadatas'][0]
                                ), 1):
                                    score = reranked_results.get('scores', [None])[i-1]
                                    score_text = f" (관련도: {score:.2%})" if score else ""
                                    
                                    st.markdown(f"""
                                    <div class="source-box">
                                        <div class="source-title">
                                            <span class="page-indicator">페이지 {meta['page']}</span>
                                            {score_text}
                                        </div>
                                        <div style="font-size: 0.9rem; color: #475569;">
                                            {doc[:300]}{'...' if len(doc) > 300 else ''}
                                        </div>
                                    </div>
                                    """, unsafe_allow_html=True)
                            
                            st.rerun()
                        
                        except Exception as e:
                            st.error(f"오류: {str(e)}")
                            import traceback
                            st.code(traceback.format_exc())


if __name__ == "__main__":
    main()