app.py

import streamlit as st
import asyncio
import tokonomics
from utils import create_model_hierarchy

st.set_page_config(page_title="LLM Pricing App", layout="wide")

# --------------------------
# Async Data Loading Function
# --------------------------
async def load_data():
    """Simulate loading data asynchronously."""
    AVAILABLE_MODELS = await tokonomics.get_available_models()
    hierarchy = create_model_hierarchy(AVAILABLE_MODELS)
    FILTERED_MODELS = []
    MODEL_PRICING = {}
    PROVIDERS = list(hierarchy.keys())
    for provider in PROVIDERS:
        for model_family in hierarchy[provider]:
            for model_version in hierarchy[provider][model_family].keys():
                for region in hierarchy[provider][model_family][model_version]:
                    model_id = hierarchy[provider][model_family][model_version][region]
                    MODEL_PRICING[model_id] = await tokonomics.get_model_costs(model_id)
                    FILTERED_MODELS.append(model_id)
    return FILTERED_MODELS, MODEL_PRICING, PROVIDERS

# --------------------------
# Provider Change Function
# --------------------------
def provider_change(provider, selected_type, all_types=["text", "vision", "video", "image"]):
    """Filter models based on the selected provider and type."""
    all_models = st.session_state.get("models", [])
    new_models = []
    others = [a_type for a_type in all_types if selected_type != a_type]
    for model_name in all_models:
        if provider in model_name:
            if selected_type in model_name:
                new_models.append(model_name)
            elif any(other in model_name for other in others):
                continue
            else:
                new_models.append(model_name)
    return new_models if new_models else all_models

# --------------------------
# Estimate Cost Function (Updated)
# --------------------------
def estimate_cost(num_alerts, input_size, output_size, model_id):
    pricing = st.session_state.get("pricing", {})
    cost_token = pricing.get(model_id)
    if not cost_token:
        return "NA"
    input_tokens = round(input_size * 1.3)
    output_tokens = round(output_size * 1.3)
    price_day = cost_token.get("input_cost_per_token", 0) * input_tokens + \
                cost_token.get("output_cost_per_token", 0) * output_tokens
    price_total = price_day * num_alerts
    return f"""## Estimated Cost:

Day Price: {price_total:0.2f} USD  
Month Price: {price_total * 31:0.2f} USD  
Year Price: {price_total * 365:0.2f} USD  
"""

# --------------------------
# Load Data into Session State (only once)
# --------------------------
if "data_loaded" not in st.session_state:
    with st.spinner("Loading pricing data..."):
        models, pricing, providers = asyncio.run(load_data())
        st.session_state["models"] = models
        st.session_state["pricing"] = pricing
        st.session_state["providers"] = providers
        st.session_state["data_loaded"] = True

# --------------------------
# Sidebar
# --------------------------
with st.sidebar:
    st.image("https://cdn.prod.website-files.com/630f558f2a15ca1e88a2f774/631f1436ad7a0605fecc5e15_Logo.svg",
             use_container_width=True)
    st.markdown(
        """ Visit: [https://www.priam.ai](https://www.priam.ai)
        """
    )
    st.divider()
    st.sidebar.title("LLM Pricing Calculator")

# --------------------------
# Main Content Layout (Model Selection Tab)
# --------------------------
tab1, tab2 = st.tabs(["Model Selection", "About"])

with tab1:
    st.header("LLM Pricing App")

    # --- Row 1: Provider/Type and Model Selection ---
    col_left, col_right = st.columns(2)
    with col_left:
        selected_provider = st.selectbox(
            "Select a provider",
            st.session_state["providers"],
            index=st.session_state["providers"].index("azure") if "azure" in st.session_state["providers"] else 0
        )
        selected_type = st.radio("Select type", options=["text", "image"], index=0)

    with col_right:
        # Filter models based on the selected provider and type
        filtered_models = provider_change(selected_provider, selected_type)
        
        if filtered_models:
            # Force "gpt-4-turbo" as default if available; otherwise, default to the first model.
            default_model = "o1" if "o1" in filtered_models else filtered_models[0]
            selected_model = st.selectbox(
                "Select a model",
                options=filtered_models,
                index=filtered_models.index(default_model)
            )
        else:
            selected_model = None
            st.write("No models available")

    # --- Row 2: Alert Stats ---
    col1, col2, col3 = st.columns(3)
    with col1:
        num_alerts = st.number_input(
            "Security Alerts Per Day",
            value=100,
            min_value=1,
            step=1,
            help="Number of security alerts to analyze daily"
        )
    with col2:
        input_size = st.number_input(
            "Alert Content Size (characters)",
            value=1000,
            min_value=1,
            step=1,
            help="Include logs, metadata, and context per alert"
        )
    with col3:
        output_size = st.number_input(
            "Analysis Output Size (characters)",
            value=500,
            min_value=1,
            step=1,
            help="Expected length of security analysis and recommendations"
        )

    # --- Row 3: Buttons ---
    btn_col1, btn_col2 = st.columns(2)
    with btn_col1:
        if st.button("Estimate"):
            if selected_model:
                st.session_state["result"] = estimate_cost(num_alerts, input_size, output_size, selected_model)
            else:
                st.session_state["result"] = "No model selected."
    with btn_col2:
        if st.button("Refresh Pricing Data"):
            with st.spinner("Refreshing pricing data..."):
                models, pricing, providers = asyncio.run(load_data())
                st.session_state["models"] = models
                st.session_state["pricing"] = pricing
                st.session_state["providers"] = providers
                st.success("Pricing data refreshed!")

    st.divider()
    # --- Display Results ---
    st.markdown("### Results")
    if "result" in st.session_state:
        st.write(st.session_state["result"])
    else:
        st.write("Use the buttons above to estimate costs.")

    # --- Clear Button Below Results ---
    if st.button("Clear"):
        st.session_state.pop("result", None)
        st.rerun()

with tab2:
    st.markdown(
        """
        ## About This App

        This is based on the tokonomics package.

        - The app downloads the latest pricing from the LiteLLM repository.
        - Using simple maths to estimate the total tokens.
        - Version 0.1

        Website: [https://www.priam.ai](https://www.priam.ai)
        """
    )