app.py

import streamlit as st
import os
import json
from datetime import datetime, timedelta
import base64
import pandas as pd
import pydeck as pdk
import requests
from paper import (
    literature_research_task, outline_task, draft_writing_task,
    citation_task, editing_task, chatbot_task,
    run_task
)

# Add web search functionality
def brave_search(query, count=10):
    """
    Perform a web search using Brave Search API
    """
    api_key = os.getenv("SEARCH_API")
    if not api_key:
        return "Error: Brave Search API key not found. Please set the SEARCH_API environment variable."
    
    url = "https://api.search.brave.com/res/v1/web/search"
    headers = {
        "Accept": "application/json",
        "Accept-Encoding": "gzip",
        "X-Subscription-Token": api_key
    }
    params = {
        "q": query,
        "count": count
    }
    
    try:
        response = requests.get(url, headers=headers, params=params)
        response.raise_for_status()
        
        results = response.json()
        formatted_results = []
        
        for result in results.get("web", {}).get("results", []):
            formatted_results.append({
                "title": result.get("title", ""),
                "url": result.get("url", ""),
                "description": result.get("description", "")
            })
        
        return formatted_results
    except Exception as e:
        return f"Search error: {str(e)}"

# st.set_page_config()는 다른 Streamlit 함수보다 가장 먼저 실행되어야 합니다.
st.set_page_config(
    page_title="Your AI Agent for Academic Research",
    page_icon="📚",
    layout="wide",
    initial_sidebar_state="expanded"
)

# ------------------------------------------
# 다국어 지원 (영어/한국어 예시)
# ------------------------------------------
translations = {
    "en": {
         "page_title": "Your AI Agent for Academic Research",
         "header": "Your AI Agent for Academic Research",
         "create_itinerary": "Generate Your Research Paper",
         "trip_details": "Research Details",
         "origin": "Research Topic",
         "destination": "Paper Title",
         "travel_dates": "Due Date",
         "duration": "Paper Length (pages)",
         "preferences": "Keywords/Focus",
         "special_requirements": "Additional Instructions",
         "submit": "🚀 Generate My Research Paper",
         "request_details": "Your Research Request",
         "from": "Topic",
         "when": "Due Date",
         "budget": "Paper Type",
         "travel_style": "Writing Style",
         "live_agent_outputs": "Live Agent Outputs",
         "full_itinerary": "Full Paper",
         "details": "Details",
         "download_share": "Download & Share",
         "save_itinerary": "Save Your Paper",
         "plan_another_trip": "🔄 Generate Another Paper",
         "about": "About",
         "how_it_works": "How it works",
         "travel_agents": "Research Agents",
         "share_itinerary": "Share Your Paper",
         "save_for_mobile": "Save for Mobile",
         "built_with": "Built with ❤️ for you",
         "itinerary_ready": "Your Research Paper is Ready! 🎉",
         "personalized_experience": "We've created a personalized academic paper based on your inputs. Explore your paper below.",
         "agent_activity": "Agent Activity",
         "error_origin_destination": "Please enter both the research topic and paper title.",
         "your_itinerary_file": "Your Paper File",
         "text_format": "Text format - Can be opened in any text editor",
         "web_search": "Web Search",
         "search_placeholder": "Search for academic papers, journals, or information...",
         "search_button": "Search",
         "search_results": "Search Results"
    },
    "ko": {
         "page_title": "당신의 학술 연구 AI 에이전트",
         "header": "당신의 학술 연구 AI 에이전트",
         "create_itinerary": "논문 생성",
         "trip_details": "연구 세부사항",
         "origin": "연구 주제",
         "destination": "논문 제목",
         "travel_dates": "제출 기한",
         "duration": "논문 분량 (페이지)",
         "preferences": "키워드/주요 초점",
         "special_requirements": "추가 지시사항",
         "submit": "🚀 나의 논문 생성",
         "request_details": "연구 요청 정보",
         "from": "주제",
         "when": "제출 기한",
         "budget": "논문 종류",
         "travel_style": "작성 스타일",
         "live_agent_outputs": "실시간 에이전트 결과",
         "full_itinerary": "전체 논문",
         "details": "세부사항",
         "download_share": "다운로드 및 공유",
         "save_itinerary": "논문 저장",
         "plan_another_trip": "🔄 다른 논문 생성",
         "about": "소개",
         "how_it_works": "작동 방식",
         "travel_agents": "연구 에이전트",
         "share_itinerary": "논문 공유",
         "save_for_mobile": "모바일 저장",
         "built_with": "당신을 위해 ❤️ 만들어졌습니다",
         "itinerary_ready": "논문이 준비되었습니다! 🎉",
         "personalized_experience": "입력하신 정보를 바탕으로 맞춤형 논문이 생성되었습니다. 아래에서 논문을 확인하세요.",
         "agent_activity": "에이전트 활동",
         "error_origin_destination": "연구 주제와 논문 제목을 모두 입력하세요.",
         "your_itinerary_file": "논문 파일",
         "text_format": "텍스트 형식 - 모든 텍스트 편집기에서 열 수 있습니다.",
         "web_search": "웹 검색",
         "search_placeholder": "학술 논문, 저널 또는 정보 검색...",
         "search_button": "검색",
         "search_results": "검색 결과"
    }
}

def t(key):
    lang = st.session_state.get("selected_language", "en")
    return translations[lang].get(key, key)

# ---------------------------
# 세션 초기화
# ---------------------------
if 'selected_language' not in st.session_state:
    st.session_state.selected_language = "en"

# ------------------------------------------
# 사이드바에 언어 선택 위젯 추가
# ------------------------------------------
with st.sidebar:
    language = st.selectbox(
        "Language / 언어",
        ["English", "한국어"]
    )
    lang_map = {
        "English": "en",
        "한국어": "ko"
    }
    st.session_state.selected_language = lang_map.get(language, "en")
    
    # Add web search component to sidebar
    st.markdown('<div class="modern-card">', unsafe_allow_html=True)
    st.markdown("### " + t("web_search"))
    search_query = st.text_input(t("search_placeholder"), key="search_query")
    search_button = st.button(t("search_button"))
    
    if search_button and search_query:
        with st.spinner("Searching..."):
            search_results = brave_search(search_query)
            if isinstance(search_results, str) and "error" in search_results.lower():
                st.error(search_results)
            else:
                st.success(f"Found {len(search_results)} results")
                st.session_state.search_results = search_results
    
    if 'search_results' in st.session_state and st.session_state.search_results:
        st.markdown(f"### {t('search_results')}")
        for i, result in enumerate(st.session_state.search_results[:5]):
            st.markdown(f"**{i+1}. [{result['title']}]({result['url']})**")
            st.markdown(f"{result['description']}")
            st.markdown("---")
    st.markdown('</div>', unsafe_allow_html=True)

# ------------------------------------------
# UI 시작
# ------------------------------------------
st.markdown("""
<style>
    :root {
        --primary: #3a86ff;
        --primary-light: #4895ef;
        --primary-dark: #2667ff;
        --background: #f8f9fa;
        --card-bg: #ffffff;
        --text: #212529;
        --border: #e9ecef;
    }
    .main-header {
        font-size: 2.5rem;
        color: var(--primary-dark);
        text-align: center;
        margin-bottom: 0.8rem;
        font-weight: 700;
    }
    .modern-card {
        background-color: var(--card-bg);
        border-radius: 10px;
        padding: 1.2rem;
        margin-bottom: 1.2rem;
        box-shadow: 0 2px 10px rgba(0,0,0,0.05);
        border: 1px solid var(--border);
    }
</style>
""", unsafe_allow_html=True)

def get_download_link(text_content, filename):
    b64 = base64.b64encode(text_content.encode()).decode()
    href = f'<a class="download-link" href="data:text/plain;base64,{b64}" download="{filename}"><i>📥</i> {t("save_itinerary")}</a>'
    return href

def display_modern_progress(current_step, total_steps=5):
    if 'progress_steps' not in st.session_state:
        st.session_state.progress_steps = {
            0: {'status': 'pending', 'name': t("trip_details")},
            1: {'status': 'pending', 'name': t("about")},
            2: {'status': 'pending', 'name': t("live_agent_outputs")},
            3: {'status': 'pending', 'name': t("download_share")},
            4: {'status': 'pending', 'name': t("full_itinerary")}
        }
    for i in range(total_steps):
        if i < current_step:
            st.session_state.progress_steps[i]['status'] = 'complete'
        elif i == current_step:
            st.session_state.progress_steps[i]['status'] = 'active'
        else:
            st.session_state.progress_steps[i]['status'] = 'pending'
    progress_percentage = (current_step / total_steps) * 100
    st.progress(progress_percentage / 100)
    st.markdown("<div>Progress: " + str(progress_percentage) + "% completed.</div>")
    return progress_percentage

def update_step_status(step_index, status):
    if 'progress_steps' in st.session_state and step_index in st.session_state.progress_steps:
        st.session_state.progress_steps[step_index]['status'] = status

def run_task_with_logs(task, input_text, log_container, output_container, results_key=None):
    log_message = f"🤖 Starting {task.agent.role}..."
    st.session_state.log_messages.append(log_message)
    with log_container:
        st.markdown("### " + t("agent_activity"))
        for msg in st.session_state.log_messages:
            st.markdown(msg)
    
    # Enhance with web search if needed
    if "include_search" in input_text.lower():
        search_query = task.agent.role + " about " + st.session_state.research_topic
        search_results = brave_search(search_query, count=5)
        
        if not isinstance(search_results, str):  # If not error message
            search_info = "\n\nWeb Search Results for reference:\n"
            for i, result in enumerate(search_results):
                search_info += f"{i+1}. {result['title']} - {result['url']}\n"
                search_info += f"   {result['description']}\n\n"
            
            input_text += search_info
            log_message = f"🔍 Added web search results for {search_query}"
            st.session_state.log_messages.append(log_message)
    
    result = run_task(task, input_text)
    if results_key:
        st.session_state.results[results_key] = result
    log_message = f"✅ {task.agent.role} completed!"
    st.session_state.log_messages.append(log_message)
    with log_container:
        st.markdown("### " + t("agent_activity"))
        for msg in st.session_state.log_messages:
            st.markdown(msg)
    with output_container:
        st.markdown(f"### {task.agent.role} Output")
        st.markdown("<div class='agent-output'>" + result + "</div>", unsafe_allow_html=True)
    return result

if 'generated_itinerary' not in st.session_state:
    st.session_state.generated_itinerary = None
if 'generation_complete' not in st.session_state:
    st.session_state.generation_complete = False
if 'current_step' not in st.session_state:
    st.session_state.current_step = 0
if 'results' not in st.session_state:
    st.session_state.results = {
        "literature_review": "",
        "outline": "",
        "draft": "",
        "citations": "",
        "edited": ""
    }
if 'log_messages' not in st.session_state:
    st.session_state.log_messages = []
if 'form_submitted' not in st.session_state:
    st.session_state.form_submitted = False

st.markdown(f"""
<div style="text-align: center;">
    <img src="https://img.icons8.com/fluency/96/book.png" width="90">
    <h1 class="main-header">{t("header")}</h1>
    <p>Generate your personalized research paper with AI-powered academic agents.</p>
</div>
""", unsafe_allow_html=True)
st.markdown('<hr>', unsafe_allow_html=True)

with st.sidebar:
    st.markdown("""
    <div style="text-align: center;">
        <img src="https://img.icons8.com/fluency/96/book.png" width="80">
        <h3>Your AI Academic Research Assistant</h3>
        <p>AI-Powered Paper Generation</p>
    </div>
    """, unsafe_allow_html=True)
    st.markdown('<div class="modern-card">', unsafe_allow_html=True)
    st.markdown("### " + t("about"))
    st.info("This tool generates a personalized academic research paper based on your inputs. Fill in the form and let our specialized agents craft your paper!")
    st.markdown('</div>', unsafe_allow_html=True)
    st.markdown('<div class="modern-card">', unsafe_allow_html=True)
    st.markdown("### " + t("how_it_works"))
    st.markdown("""
    <ol>
        <li>Enter your research details</li>
        <li>AI conducts literature research</li>
        <li>Generate a paper outline</li>
        <li>Draft and edit your paper</li>
        <li>Download your final paper</li>
    </ol>
    """, unsafe_allow_html=True)
    st.markdown('</div>', unsafe_allow_html=True)

if not st.session_state.generation_complete:
    st.markdown('<div class="modern-card">', unsafe_allow_html=True)
    st.markdown("<h3>" + t("create_itinerary") + "</h3>", unsafe_allow_html=True)
    st.markdown("<p>Fill in the details below to generate your research paper.</p>", unsafe_allow_html=True)
    
    with st.form("research_form"):
        col1, col2 = st.columns(2)
        with col1:
            research_topic = st.text_input(t("origin"), placeholder="e.g., Deep Learning in Healthcare")
            paper_title = st.text_input(t("destination"), placeholder="e.g., Advances in Deep Learning for Medical Diagnosis")
            due_date = st.date_input(t("travel_dates"), min_value=datetime.now())
        with col2:
            paper_length = st.slider(t("duration"), min_value=5, max_value=50, value=10)
            paper_type_options = ["Journal", "Conference", "Thesis", "Review"]
            paper_type = st.selectbox(t("budget"), paper_type_options, help="Select the type of paper")
            writing_style = st.multiselect(t("travel_style"), options=["Formal", "Technical", "Creative"], default=["Formal"])
        additional_instructions = st.text_area(t("special_requirements"), placeholder="Any additional instructions or requirements...")
        keywords = st.text_area(t("preferences"), placeholder="Enter keywords or focus areas, separated by commas")
        
        # Add option to include web search
        include_search = st.checkbox("Include web search for latest references", value=True)
        
        submit_button = st.form_submit_button(t("submit"))
    st.markdown('</div>', unsafe_allow_html=True)
    
    if submit_button:
        if not research_topic or not paper_title:
            st.error(t("error_origin_destination"))
        else:
            st.session_state.form_submitted = True
            st.session_state.research_topic = research_topic
            user_input = {
                "research_topic": research_topic,
                "paper_title": paper_title,
                "due_date": due_date.strftime("%Y-%m-%d"),
                "paper_length": str(paper_length),
                "paper_type": paper_type,
                "writing_style": ", ".join(writing_style),
                "keywords": keywords,
                "additional_instructions": additional_instructions,
                "include_search": "include_search" if include_search else ""
            }
            st.session_state.user_input = user_input
            input_context = f"""Research Request Details:
Research Topic: {user_input['research_topic']}
Paper Title: {user_input['paper_title']}
Due Date: {user_input['due_date']}
Paper Length: {user_input['paper_length']} pages
Paper Type: {user_input['paper_type']}
Writing Style: {user_input['writing_style']}
Keywords/Focus: {user_input['keywords']}
Additional Instructions: {user_input['additional_instructions']}
{user_input['include_search']}
"""
            llm_language_instructions = {
                "en": "Please output the response in English.",
                "ko": "한국어로 출력해 주세요."
            }
            selected_lang = st.session_state.get("selected_language", "en")
            language_instruction = llm_language_instructions.get(selected_lang, "Please output the response in English.")
            modified_input_context = language_instruction + "\n" + input_context
            
            st.markdown("<div>Processing your request...</div>", unsafe_allow_html=True)
            st.session_state.current_step = 0
            update_step_status(0, 'active')
            progress_placeholder = st.empty()
            with progress_placeholder.container():
                display_modern_progress(st.session_state.current_step)
            log_container = st.container()
            st.session_state.log_messages = []
            output_container = st.container()
            st.session_state.results = {}
            
            # Step 1: Literature Research
            literature_review = run_task_with_logs(
                literature_research_task,
                modified_input_context.format(topic=user_input['research_topic'], keywords=user_input['keywords']),
                log_container,
                output_container,
                "literature_review"
            )
            update_step_status(0, 'complete')
            st.session_state.current_step = 1
            update_step_status(1, 'active')
            with progress_placeholder.container():
                display_modern_progress(st.session_state.current_step)
            
            # Step 2: Generate Outline
            outline = run_task_with_logs(
                outline_task,
                modified_input_context.format(topic=user_input['research_topic']),
                log_container,
                output_container,
                "outline"
            )
            update_step_status(1, 'complete')
            st.session_state.current_step = 2
            update_step_status(2, 'active')
            with progress_placeholder.container():
                display_modern_progress(st.session_state.current_step)
            
            # Step 3: Draft Writing
            draft = run_task_with_logs(
                draft_writing_task,
                modified_input_context.format(topic=user_input['research_topic']),
                log_container,
                output_container,
                "draft"
            )
            update_step_status(2, 'complete')
            st.session_state.current_step = 3
            update_step_status(3, 'active')
            with progress_placeholder.container():
                display_modern_progress(st.session_state.current_step)
            
            # Step 4: Citation Generation
            citations = run_task_with_logs(
                citation_task,
                modified_input_context.format(topic=user_input['research_topic']),
                log_container,
                output_container,
                "citations"
            )
            update_step_status(3, 'complete')
            st.session_state.current_step = 4
            update_step_status(4, 'active')
            with progress_placeholder.container():
                display_modern_progress(st.session_state.current_step)
            
            # Step 5: Editing and Polishing
            edited = run_task_with_logs(
                editing_task,
                modified_input_context.format(topic=user_input['research_topic']),
                log_container,
                output_container,
                "edited"
            )
            update_step_status(4, 'complete')
            st.session_state.current_step = 5
            with progress_placeholder.container():
                display_modern_progress(st.session_state.current_step)
            
            full_paper = f"""Research Paper:
{input_context}

Literature Review:
{literature_review}

Outline:
{outline}

Draft:
{draft}

Citations:
{citations}

Edited Version:
{edited}
"""
            st.session_state.generated_itinerary = full_paper
            st.session_state.generation_complete = True
            date_str = datetime.now().strftime("%Y-%m-%d")
            st.session_state.filename = f"{user_input['paper_title'].replace(' ', '_')}_{date_str}_paper.txt"

if st.session_state.generation_complete:
    st.markdown(f"""
    <div class="modern-card">
      <div style="text-align: center;">
        <h2>{t("itinerary_ready")}</h2>
        <p>{t("personalized_experience")}</p>
      </div>
    </div>
    """, unsafe_allow_html=True)
    
    # 탭 생성 (전체 논문, 세부 정보, 다운로드/공유, 시각화, 챗봇)
    full_paper_tab, details_tab, download_tab, visualization_tab, chatbot_tab = st.tabs([
        "🗒️ " + t("full_itinerary"),
        "💼 " + t("details"),
        "💾 " + t("download_share"),
        "📊 Visualization",
        "🤖 챗봇 인터페이스"
    ])
    
    with full_paper_tab:
        st.text_area("Your Research Paper", st.session_state.generated_itinerary, height=600)
    
    with details_tab:
        agent_tabs = st.tabs(["📚 Literature Review", "📝 Outline", "✍️ Draft", "🔗 Citations", "🖋️ Edited Version"])
        with agent_tabs[0]:
            st.markdown("### Literature Review")
            st.markdown(st.session_state.results.get("literature_review", ""))
        with agent_tabs[1]:
            st.markdown("### Outline")
            st.markdown(st.session_state.results.get("outline", ""))
        with agent_tabs[2]:
            st.markdown("### Draft")
            st.markdown(st.session_state.results.get("draft", ""))
        with agent_tabs[3]:
            st.markdown("### Citations")
            st.markdown(st.session_state.results.get("citations", ""))
        with agent_tabs[4]:
            st.markdown("### Edited Version")
            st.markdown(st.session_state.results.get("edited", ""))
    
    with download_tab:
        col1, col2 = st.columns([2, 1])
        with col1:
            st.markdown("### " + t("save_itinerary"))
            st.markdown("Download your research paper to access it offline or share with your colleagues.")
            st.markdown(f"""
            <div style="background-color: #f8f9fa; padding: 15px; border-radius: 10px; margin-top: 20px;">
                <h4>{t("your_itinerary_file")}</h4>
                <p style="font-size: 0.9rem; color: #6c757d;">{t("text_format")}</p>
            </div>
            """, unsafe_allow_html=True)
            st.markdown("<div>" + get_download_link(st.session_state.generated_itinerary, st.session_state.filename) + "</div>", unsafe_allow_html=True)
            st.markdown("### " + t("share_itinerary"))
            st.markdown("*Coming soon: Email your paper or share via social media.*")
        with col2:
            st.markdown("### " + t("save_for_mobile"))
            st.markdown("*Coming soon: QR code for easy access on your phone*")
    
    with visualization_tab:
        st.markdown("### Visualization")
        st.markdown("A conceptual diagram or visualization related to your research paper can be displayed here. (Feature under development)")
    
    with chatbot_tab:
        st.markdown("### AI 챗봇 인터페이스")
        if "chat_history" not in st.session_state:
            st.session_state.chat_history = []
        
        user_message = st.text_input("메시지를 입력하세요:", key="chat_input")
        if st.button("전송", key="send_button"):
            if user_message:
                # Add option to search when chatting
                if "search" in user_message.lower() or "find" in user_message.lower() or "look up" in user_message.lower():
                    search_results = brave_search(user_message, count=3)
                    if not isinstance(search_results, str):
                        search_context = f"I found some information that might help:\n\n"
                        for i, result in enumerate(search_results):
                            search_context += f"{i+1}. {result['title']} - {result['url']}\n"
                            search_context += f"   {result['description']}\n\n"
                        
                        enhanced_prompt = f"{user_message}\n\nSearch context: {search_context}"
                        response = run_task(chatbot_task, enhanced_prompt)
                    else:
                        response = run_task(chatbot_task, user_message)
                else:
                    response = run_task(chatbot_task, user_message)
                
                st.session_state.chat_history.append({
                    "speaker": "사용자",
                    "message": user_message,
                    "time": datetime.now()
                })
                st.session_state.chat_history.append({
                    "speaker": "AI",
                    "message": response,
                    "time": datetime.now()
                })
        
        st.markdown("<div style='max-height:400px; overflow-y:auto; padding:10px; border:1px solid #eaeaea; border-radius:6px;'>", unsafe_allow_html=True)
        for chat in st.session_state.chat_history:
            time_str = chat["time"].strftime("%H:%M:%S")
            st.markdown(f"**{chat['speaker']}** ({time_str}): {chat['message']}")
        st.markdown("</div>", unsafe_allow_html=True)

st.markdown("""
<div style="text-align: center; padding: 20px; color: #6c757d; font-size: 0.8rem;">
    <p>""" + t("built_with") + """</p>
</div>
""", unsafe_allow_html=True)