app.py

import streamlit as st
from arxiv_parser import (
    get_ar5iv_link,
    get_bibliography_from_html,
    get_html_page,
    get_paragraphs_from_html,
    get_title_from_html,
)
import arxiv_parser
from gist import answer_question, create_summary, get_next_page_break, get_client
from streamlit_helper import (
    compute_gist_metrics,
    delete_session_state,
    render_llm_metrics,
    render_new_page,
    render_processed_pages,
    reset_session_state,
    show_inference_stat_dist,
    unpack_summary,
    update_inference_client,
)
import os
import requests
import article_parser
from providers import PROVIDERS, check_provider_status
from db_manager import DocumentStore, ArxivDocument
import datetime
import re

def get_groq_models():
    """Fetch available Groq models via API"""
    api_key = os.environ.get("GROQ_API_KEY", "")
    url = "https://api.groq.com/openai/v1/models"
    
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    
    try:
        response = requests.get(url, headers=headers)
        if response.status_code == 200:
            models_data = response.json()
            # Create an ordered dictionary of models
            models_dict = {}
            # Sort models by ID for consistency
            for model in sorted(models_data["data"], key=lambda x: x["id"]):
                models_dict[model["id"]] = model["id"]
            return models_dict
        else:
            st.warning(f"Failed to fetch Groq models: {response.status_code}")
            return {}
    except Exception as e:
        st.warning(f"Error fetching Groq models: {str(e)}")
        return {}

@st.cache_data(ttl=3600)  # Cache for 1 hour
def get_groq_models_cached():
    return get_groq_models()

def process_url(input_url: str):
    """Process URL using appropriate parser."""
    if article_parser.is_arxiv_url(input_url):
        # Use existing arXiv parser
        ar5iv_url = arxiv_parser.get_ar5iv_link(input_url)
        page_html = arxiv_parser.get_html_page(ar5iv_url)
        title = arxiv_parser.get_title_from_html(page_html)
        paragraphs, paragraphs_html, is_rtl_list = arxiv_parser.get_paragraphs_from_html(page_html)
        bib = arxiv_parser.get_bibliography_from_html(page_html)
    else:
        # Use new general article parser
        page_html = article_parser.get_html_page(input_url)
        title = article_parser.get_title_from_html(page_html)
        paragraphs, paragraphs_html, is_rtl_list = article_parser.get_paragraphs_from_html(page_html)
        bib = article_parser.get_bibliography_from_html(page_html)
    
    return title, paragraphs, paragraphs_html, is_rtl_list, bib

def render_processed_pages(title, paragraphs, paragraphs_html, is_rtl_list):
    """Render the processed pages with appropriate text direction."""
    for i, (text, html, is_rtl) in enumerate(zip(paragraphs, paragraphs_html, is_rtl_list)):
        # Apply RTL styling if needed
        if is_rtl:
            st.markdown(
                f"""
                <div dir="rtl" style="text-align: right;">
                    {text}
                </div>
                """,
                unsafe_allow_html=True
            )
        else:
            st.markdown(text)

def get_openai_models():
    """Fetch available OpenAI models and their tiers"""
    models_dict = {
        "gpt-3.5-turbo": {
            "id": "gpt-3.5-turbo",
            "name": "GPT-3.5 Turbo",
            "tier": "paid",
            "context_length": "4K",
            "pricing": "$0.0015/1K tokens"
        },
        "gpt-3.5-turbo-16k": {
            "id": "gpt-3.5-turbo-16k",
            "name": "GPT-3.5 Turbo 16K",
            "tier": "paid",
            "context_length": "16K",
            "pricing": "$0.003/1K tokens"
        },
        "gpt-4": {
            "id": "gpt-4",
            "name": "GPT-4",
            "tier": "paid",
            "context_length": "8K",
            "pricing": "$0.03/1K tokens"
        },
        "gpt-4-turbo": {
            "id": "gpt-4-turbo-preview",
            "name": "GPT-4 Turbo",
            "tier": "paid",
            "context_length": "128K",
            "pricing": "$0.01/1K tokens"
        }
    }
    return models_dict

def get_google_models():
    """Fetch available Google models and their tiers"""
    models_dict = {
        "gemini-pro": {
            "id": "gemini-pro",
            "name": "Gemini Pro",
            "tier": "quota_free",  # Free with quota
            "context_length": "32K",
            "pricing": "Free up to 60 requests/min"
        },
        "gemini-pro-vision": {
            "id": "gemini-pro-vision",
            "name": "Gemini Pro Vision",
            "tier": "quota_free",
            "context_length": "32K",
            "pricing": "Free up to 60 requests/min"
        },
        "palm-2": {
            "id": "palm-2",
            "name": "PaLM 2",
            "tier": "paid",
            "context_length": "8K",
            "pricing": "Contact sales"
        }
    }
    return models_dict

def get_available_models(inference_provider: str):
    """Get available models based on the selected provider"""
    model_name_to_id = {}
    
    if inference_provider == "Groq":
        model_name_to_id = {
            "Mixtral 8x7B": "mixtral-8x7b-32768",
            "LLaMA2 70B": "llama2-70b-4096"
        }
    elif inference_provider == "SambaNova":
        model_name_to_id = {
            "SambaNova 8B": "Meta-Llama-3.1-8B-Instruct"
        }
    elif inference_provider == "OpenAI":
        model_name_to_id = {
            "GPT-4": "gpt-4",
            "GPT-3.5 Turbo": "gpt-3.5-turbo"
        }
    elif inference_provider == "Cerebras":
        model_name_to_id = {
            "Cerebras 13B": "cerebras-13b"
        }
    elif inference_provider == "Fireworks":
        model_name_to_id = {
            "Llama 3.1 8B": "accounts/fireworks/models/llama-v3p1-8b-instruct"
        }
    elif inference_provider == "DeepSeek":
        model_name_to_id = {
            "DeepSeek 7B": "deepseek-7b"
        }
    
    return model_name_to_id

def render_model_selector(inference_provider: str):
    """Render the model selector with available models"""
    models_dict = get_available_models(inference_provider)
    
    if not models_dict:
        st.warning(f"No models available for {inference_provider}")
        return None
    
    # Create a radio group for model selection
    model_options = list(models_dict.keys())
    selected_model = st.radio("Select Model", model_options)
    
    return models_dict[selected_model]

# Initialize document store
doc_store = DocumentStore()

def get_arxiv_id(url: str) -> str:
    """Extract arxiv ID from URL"""
    match = re.search(r"(?:abs|pdf)/(\d+\.\d+)", url)
    return match.group(1) if match else None

if __name__ == "__main__":
    st.set_page_config(page_title="Arxiv GIST", layout="wide")

    with st.sidebar.container():
        # Get API key
        env_key = os.environ.get("OPENAI_API_KEY", "")
        api_key = st.text_input(
            "OpenAI API Key",
            type="password",
            value=env_key
        )

        if not api_key:
            st.warning("Please enter your OpenAI API key")
            st.stop()

        # Set model options
        model_options = {
            "GPT-4": "gpt-4",
            "GPT-3.5 Turbo": "gpt-3.5-turbo"
        }
        
        selected_model = st.selectbox(
            "Select Model",
            options=list(model_options.keys())
        )
        
        model_id = model_options[selected_model]

        try:
            if "client" not in st.session_state:
                st.session_state["client"] = get_client(
                    model_id=model_id,
                    api_key=api_key
                )
                st.success("Successfully connected to OpenAI")
        except Exception as e:
            st.error(f"Failed to connect to OpenAI: {str(e)}")
            st.stop()

    st.title("Q&A Documents")
#     st.write(
#         """Current Large Language Models (LLMs) are not only limited to some maximum context length, but also
# are not able to robustly consume long inputs. GIST addresses this by proposing a human-inspired reading agent.
# 1. The model "reads" a document and breaks it into "pages".
# 2. Each page is then summarized.
# 3. During Q&A, the model is given the summaries of each page and asked if it wants to expand ("reread") any of the pages.
# 4. The summaries of all pages + whichever pages the model chose to expand are used to answer the question.

# This approach allows the model to retain the high level flow of the original document while mixing granularities that allow it to capture the finer details that are needed to answer a specific question.
# Read more about GIST here: https://arxiv.org/pdf/1706.03762
# """
#     )

    input_url = st.text_input(
        "example Arxiv link:", placeholder="https://arxiv.org/pdf/1706.03762"
    )

    if input_url:
        if st.session_state.get("client") is None:
            st.stop()

        arxiv_id = get_arxiv_id(input_url)
        if arxiv_id:
            # Check if document exists in DB
            doc = doc_store.get_document(arxiv_id)
            if doc:
                st.success(f"Loading cached version of {doc.title}")
                title = doc.title
                paragraphs = doc.paragraphs
                paragraphs_html = doc.html_content
                bib = doc.bibliography
            else:
                # Process new document
                title, paragraphs, paragraphs_html, is_rtl_list, bib = process_url(input_url)
                # Save to DB
                doc = ArxivDocument(
                    arxiv_id=arxiv_id,
                    title=title,
                    paragraphs=paragraphs,
                    summaries=[],
                    html_content=paragraphs_html,
                    bibliography=bib,
                    processed_date=datetime.datetime.now().isoformat()
                )
                doc_store.save_document(doc)

        if (
            "pause_point" not in st.session_state
            or "url" not in st.session_state
            or st.session_state["url"] != input_url
        ):
            reset_session_state(input_url)

        render_llm_metrics(navbar_placeholder)
        render_processed_pages(title, paragraphs, paragraphs_html, is_rtl_list)

        # Phase 1: Preprocessing the arxiv paper.
        # The LLM will iteratively group paragraphs together based on
        # narration by selecting "pause points". Paragraphs contained from the
        # old pause point to the new pause point are referred to as a "page". The
        # LLM then summarizes this new page into a summarized page.
        while st.session_state["pause_point"] < len(paragraphs):
            old_pause_point = st.session_state["pause_point"]
            st.session_state["pages"], new_pause_point = get_next_page_break(
                st.session_state["client"],
                title,
                paragraphs,
                st.session_state["pages"],
                old_pause_point,
                llm_metrics=st.session_state["llm_metrics"],
                verbose=False,
            )

            page_html = paragraphs_html[old_pause_point:new_pause_point]
            st.session_state["pages_html"].append(page_html)
            st.session_state["pause_point"] = new_pause_point
            added_page_idx = len(st.session_state["pages"]) - 1

            render_llm_metrics(navbar_placeholder)
            cols = render_new_page()

            with cols[1]:
                summary_stream = create_summary(
                    st.session_state["client"],
                    title,
                    st.session_state["pages"][added_page_idx],
                    llm_metrics=st.session_state["llm_metrics"],
                    verbose=False,
                    stream=True,
                )
                st.write_stream(unpack_summary(summary_stream))
                render_llm_metrics(navbar_placeholder)

        with st.expander("Bibliography"):
            st.markdown(bib, unsafe_allow_html=True)

        compute_gist_metrics()
        render_llm_metrics(navbar_placeholder)
        show_inference_stat_dist(inference_stat_dist_placeholder)

        question = st.text_input("Question:")
        if question:
            intermediate, answer_stream = answer_question(
                st.session_state["client"],
                title,
                st.session_state["pages"],
                st.session_state["shortened_pages"],
                question,
                llm_metrics=st.session_state["llm_metrics"],
                verbose=False,
                stream=True,
            )
            if not isinstance(intermediate, str):
                st.error(intermediate)
                st.stop()

            with st.expander("Show reader's thoughts"):
                st.write(intermediate)
            st.write_stream(answer_stream)
            render_llm_metrics(navbar_placeholder)