Refactor API key handling and improve UI elements

app.py

@@ -1,17 +1,19 @@
 import sys
 import time
 from pathlib import Path
+
 sys.path.insert(0, str(Path(__file__).parent))
 
 import streamlit as st
-from anthropic import Anthropic
 
 from foto import (
-    MODEL_LABELS, get_model, CostTracker,
+    MODEL_LABELS, MODEL_REGISTRY, get_model, CostTracker,
     InputParser, PaperSearcher, PaperTriager,
     PDFStore, FigureExtractor, FigureScorer,
     build_zip, format_authors, get_confidence, confidence_badge_class,
 )
+from foto.models import PROVIDER_DISPLAY
+from foto.llm_client import LLMClient
 from foto.persistence import load_stats, log_search, log_rating
 
 st.set_page_config(
@@ -29,13 +31,14 @@ html, body, [class*="css"] { font-family: 'Inter', sans-serif; font-weight: 300;
 .stApp { background: #fafaf8; color: #1a1a1a; }
 
 .foto-header { padding: 3rem 0 1.5rem 0; border-bottom: 1px solid #e0e0d8; margin-bottom: 2.5rem; }
-.foto-title { font-family: 'DM Serif Display', serif; font-size: 2.8rem; font-weight: 400; letter-spacing: -0.02em; line-height: 1.1; color: #1a1a1a; margin: 0; }
-.foto-title em { font-style: italic; color: #4a5568; }
 .foto-subtitle { font-size: 0.85rem; color: #888; letter-spacing: 0.08em; text-transform: uppercase; margin-top: 0.4rem; font-family: 'DM Mono', monospace; }
 .foto-tagline { font-size: 1rem; color: #555; margin-top: 0.8rem; font-weight: 300; max-width: 560px; }
 
 .section-label { font-family: 'DM Mono', monospace; font-size: 0.7rem; letter-spacing: 0.12em; text-transform: uppercase; color: #999; margin-bottom: 0.5rem; }
 
+.api-help { font-size: 0.75rem; color: #777; margin-top: -0.5rem; margin-bottom: 0.8rem; }
+.api-help a { color: #4a5568; text-decoration: underline; }
+
 .result-card { background: white; border: 1px solid #e8e8e0; border-radius: 4px; padding: 1.2rem 1.4rem; margin-bottom: 1.2rem; }
 .result-title { font-family: 'DM Serif Display', serif; font-size: 1.05rem; color: #1a1a1a; margin-bottom: 0.2rem; line-height: 1.3; }
 .result-meta { font-size: 0.8rem; color: #777; font-style: italic; margin-bottom: 0.6rem; }
@@ -62,20 +65,27 @@ html, body, [class*="css"] { font-family: 'Inter', sans-serif; font-weight: 300;
 .tally-num { font-family: 'DM Serif Display', serif; font-size: 2rem; color: #fafaf8; }
 .tally-label { font-family: 'DM Mono', monospace; font-size: 0.65rem; color: #888; letter-spacing: 0.08em; margin-top: 0.2rem; }
 
+.pathfinder-cite { font-family: 'DM Mono', monospace; font-size: 0.72rem; color: #888; line-height: 1.5; }
+.pathfinder-cite a { color: #4a5568; text-decoration: underline; }
+
 .stTextArea textarea { font-family: 'Inter', sans-serif; font-size: 0.95rem; border: 1px solid #ddd; border-radius: 3px; }
 .stButton button { font-family: 'DM Mono', monospace; font-size: 0.8rem; letter-spacing: 0.06em; border-radius: 3px; }
-div[data-testid="stSelectbox"] label,
-div[data-testid="stTextInput"] label,
-div[data-testid="stCheckbox"] label { font-family: 'DM Mono', monospace; font-size: 0.75rem; letter-spacing: 0.08em; text-transform: uppercase; color: #888; }
 
-.pathfinder-row { display: flex; align-items: center; gap: 0.5rem; }
-.pathfinder-cite { font-family: 'DM Mono', monospace; font-size: 0.7rem; color: #888; }
-.pathfinder-cite a { color: #888; text-decoration: underline; }
+div[data-testid="stSelectbox"] label,
+div[data-testid="stTextInput"] label { font-family: 'DM Mono', monospace; font-size: 0.75rem; letter-spacing: 0.08em; text-transform: uppercase; color: #888; }
+
+div[data-testid="stCheckbox"] label p {
+    font-family: 'Inter', sans-serif !important;
+    font-size: 0.9rem !important;
+    letter-spacing: normal !important;
+    text-transform: none !important;
+    color: #1a1a1a !important;
+    font-weight: 400 !important;
+}
 </style>
 """, unsafe_allow_html=True)
 
 
-# Per-session state for caching, run status, log buffer, and feedback tally
 for key, default in {
     "pdf_cache": {},
     "results": None,
@@ -102,55 +112,104 @@ st.markdown("""
 
 col_left, col_right = st.columns([1, 1], gap="large")
 
-with col_left:
-    st.markdown('<p class="section-label">Model</p>', unsafe_allow_html=True)
-    model_label = st.selectbox("Model", options=MODEL_LABELS, label_visibility="collapsed")
 
-    st.markdown('<p class="section-label" style="margin-top:1.2rem;">Anthropic API Key</p>', unsafe_allow_html=True)
-    api_key = st.text_input("Anthropic API Key", type="password", label_visibility="collapsed", placeholder="sk-ant-...")
+def render_api_key_field(provider: str, model_cfg, key_state: str) -> str:
+    display_name = PROVIDER_DISPLAY.get(provider, "API Key")
+    st.markdown(f'<p class="section-label" style="margin-top:0.8rem;">{display_name}</p>', unsafe_allow_html=True)
+    key = st.text_input(
+        display_name, type="password", label_visibility="collapsed",
+        placeholder="...", key=key_state,
+    )
+    st.markdown(
+        f'<div class="api-help">{model_cfg.api_help_text} '
+        f'<a href="{model_cfg.api_help_url}" target="_blank">Get key →</a></div>',
+        unsafe_allow_html=True,
+    )
+    return key
 
-    st.markdown('<p class="section-label" style="margin-top:0.8rem;">Semantic Scholar Key (optional)</p>', unsafe_allow_html=True)
-    s2_key = st.text_input("S2 Key", type="password", label_visibility="collapsed", placeholder="(Recommended for keyword fallback)")
 
-    st.markdown('<p class="section-label" style="margin-top:1.5rem;">Describe the figure</p>', unsafe_allow_html=True)
-    user_text = st.text_area(
-        "Figure description", label_visibility="collapsed", height=120,
-        placeholder='e.g. "scatter plot of cosmological parameter constraints from wavelet scattering transform, Omega_m vs sigma_8"',
+with col_left:
+    st.markdown('<p class="section-label">Primary Model</p>', unsafe_allow_html=True)
+    primary_label = st.selectbox(
+        "Primary Model", options=MODEL_LABELS,
+        label_visibility="collapsed", key="primary_model",
     )
+    primary_cfg = get_model(primary_label)
+    primary_key = render_api_key_field(primary_cfg.provider, primary_cfg, "primary_api_key")
 
-    st.markdown('<p class="section-label" style="margin-top:0.8rem;">Upload a sketch (optional)</p>', unsafe_allow_html=True)
-    sketch_file = st.file_uploader("Sketch", type=["png", "jpg", "jpeg", "webp"], label_visibility="collapsed")
+    use_pathfinder = st.checkbox(
+        "Use Pathfinder (recommended)",
+        value=True,
+        key="use_pathfinder",
+    )
+    st.markdown(
+        '<div class="pathfinder-cite" style="margin-top:-0.4rem;margin-left:1.8rem;">'
+        'Based on <a href="https://arxiv.org/abs/2408.01556" target="_blank">arXiv:2408.01556</a> · '
+        'OpenAI key required for query embedding · '
+        '~$1 per 2M queries'
+        '</div>',
+        unsafe_allow_html=True,
+    )
 
-    # Pathfinder toggle + inline citation link
-    pf_col1, pf_col2 = st.columns([2, 3])
-    with pf_col1:
-        use_pathfinder = st.checkbox("Use Pathfinder", value=True)
-    with pf_col2:
+    openai_key = ""
+    if use_pathfinder:
+        st.markdown('<p class="section-label" style="margin-top:0.8rem;">OpenAI API Key</p>', unsafe_allow_html=True)
+        openai_key = st.text_input(
+            "OpenAI Key", type="password", label_visibility="collapsed",
+            placeholder="sk-...", key="openai_key",
+        )
         st.markdown(
-            '<div class="pathfinder-cite" style="padding-top:0.55rem;">'
-            'based on <a href="https://arxiv.org/abs/2408.01556" target="_blank">arXiv:2408.01556</a>'
-            '</div>',
+            '<div class="api-help">Used to embed queries with text-embedding-3-small. '
+            '<a href="https://platform.openai.com/api-keys" target="_blank">Get key →</a></div>',
             unsafe_allow_html=True,
         )
 
-    # OpenAI key only needed when Pathfinder is active (used to embed queries
-    # with text-embedding-3-small against the Pathfinder corpus)
-    openai_key = None
-    if use_pathfinder:
-        st.markdown('<p class="section-label" style="margin-top:0.6rem;">OpenAI API Key</p>', unsafe_allow_html=True)
-        openai_key = st.text_input(
-            "OpenAI Key", type="password", label_visibility="collapsed",
-            placeholder="sk-...",
+    s2_key = ""
+    if not use_pathfinder:
+        st.markdown('<p class="section-label" style="margin-top:0.8rem;">Semantic Scholar Key (optional)</p>', unsafe_allow_html=True)
+        s2_key = st.text_input(
+            "S2 Key", type="password", label_visibility="collapsed",
+            placeholder="(improves keyword search)", key="s2_key",
         )
         st.markdown(
-            '<p style="font-size:0.78rem;color:#888;margin-top:-0.4rem;">'
-            'Used to embed queries with text-embedding-3-small (~$0.40 per million queries).'
-            '</p>',
+            '<div class="api-help">Optional — speeds up the keyword-based paper search. '
+            '<a href="https://www.semanticscholar.org/product/api" target="_blank">Get key →</a></div>',
             unsafe_allow_html=True,
         )
 
-    run_verify = st.checkbox("Secondary verification — recommended, adds ~$0.05", value=True)
-    st.markdown('<p style="font-size:0.78rem;color:#888;margin-top:-0.8rem;margin-left:1.8rem;">Uses a smarter model to double-check top matches. Best results, small extra cost.</p>', unsafe_allow_html=True)
+    st.markdown('<p class="section-label" style="margin-top:1.2rem;">Describe the figure</p>', unsafe_allow_html=True)
+    user_text = st.text_area(
+        "Figure description", label_visibility="collapsed", height=120,
+        placeholder='e.g. "scatter plot of cosmological parameter constraints from wavelet scattering transform, Omega_m vs sigma_8"',
+    )
+
+    st.markdown('<p class="section-label" style="margin-top:0.8rem;">Upload a sketch (optional)</p>', unsafe_allow_html=True)
+    sketch_file = st.file_uploader("Sketch", type=["png", "jpg", "jpeg", "webp"], label_visibility="collapsed")
+
+    run_verify = st.checkbox("Secondary verification (recommended)", value=True, key="run_verify")
+    st.markdown(
+        '<p style="font-size:0.78rem;color:#888;margin-top:-0.6rem;margin-left:1.8rem;">'
+        'Double-checks top matches. Uses the primary model by default.</p>',
+        unsafe_allow_html=True,
+    )
+
+    verify_cfg = primary_cfg
+    verify_key = primary_key
+    if run_verify:
+        verify_options = ["Same as primary"] + MODEL_LABELS
+        verify_choice = st.selectbox(
+            "Verification model", options=verify_options,
+            label_visibility="collapsed", key="verify_model",
+        )
+        if verify_choice != "Same as primary":
+            verify_cfg = get_model(verify_choice)
+            if verify_cfg.provider != primary_cfg.provider:
+                verify_key = render_api_key_field(
+                    verify_cfg.provider, verify_cfg, "verify_api_key",
+                )
+            else:
+                verify_key = primary_key
+
     num_papers = st.slider("Papers to search", min_value=5, max_value=50, value=20, step=5)
 
     run_btn = st.button("🔭  Search", use_container_width=True, type="primary", disabled=st.session_state.running)
@@ -160,29 +219,30 @@ with col_right:
         st.markdown("""
 <div style="padding: 3rem 2rem; color: #aaa; text-align: center;">
   <div style="font-size: 3rem; margin-bottom: 1rem;">🔭</div>
-  <div style="font-family: 'DM Mono', monospace; font-size: 0.75rem; letter-spacing: 0.1em; text-transform: uppercase;">Results will appear here</div>
-  <div style="margin-top: 0.8rem; font-size: 0.85rem; max-width: 300px; margin-left: auto; margin-right: auto; line-height: 1.6;">
-    Enter a description, an optional sketch, and your API key then hit Search (On average per search cost 20-30¢).
-  </div>
+  <div style="font-family: 'DM Mono', monospace; font-size: 0.75rem; letter-spacing: 0.1em; text-transform: uppercase;">Your search progress will appear here</div>
 </div>
 """, unsafe_allow_html=True)
 
 
-# Full pipeline runs on button press
 if run_btn:
-    if not api_key:
-        st.error("Please enter your Anthropic API key.")
+    if not primary_key:
+        st.error(f"Please enter your {PROVIDER_DISPLAY[primary_cfg.provider]}.")
     elif use_pathfinder and not openai_key:
-        st.error("Pathfinder is checked — please enter your OpenAI API key, or uncheck Pathfinder to use keyword search.")
+        st.error("Pathfinder is checked — please enter your OpenAI API key, or uncheck Pathfinder.")
+    elif run_verify and verify_cfg.provider != primary_cfg.provider and not verify_key:
+        st.error(f"Verification model needs its own key — please enter your {PROVIDER_DISPLAY[verify_cfg.provider]}.")
     elif not user_text and not sketch_file:
         st.error("Please enter a description or upload a sketch (or both).")
     else:
         st.session_state.running = True
         st.session_state.results = None
 
-        model_cfg = get_model(model_label)
-        client = Anthropic(api_key=api_key)
-        tracker = CostTracker(model_cfg.prices)
+        tracker = CostTracker()
+        primary_client = LLMClient(provider=primary_cfg.provider, api_key=primary_key)
+        verify_client = (
+            primary_client if verify_cfg.provider == primary_cfg.provider
+            else LLMClient(provider=verify_cfg.provider, api_key=verify_key)
+        )
         sketch_bytes = sketch_file.read() if sketch_file else None
 
         with col_right:
@@ -199,33 +259,38 @@ if run_btn:
             st.session_state.log = []
 
             try:
-                # Parse text + optional sketch into a structured search spec
                 log("⟳ Parsing your description...")
-                parser = InputParser(client, model_cfg.smart, tracker)
+                parser = InputParser(
+                    primary_client, primary_cfg.model_id, primary_cfg.prices, tracker,
+                    max_tokens=primary_cfg.triage_max_tokens * 4,
+                )
                 spec = parser.parse(text=user_text or None, sketch_bytes=sketch_bytes)
                 query = spec["science_query"] or user_text or "(no query)"
                 log(f"✓ Query: <em>{query}</em>")
                 if spec.get("plot_type"):
                     log(f"  Plot type: {spec['plot_type']}")
 
-                # Pathfinder semantic retrieval, or legacy keyword expansion as fallback
                 searcher = PaperSearcher(s2_key=s2_key or None)
                 if use_pathfinder:
                     all_papers = searcher.expanded_search_pathfinder(query, openai_key, log=log)
                 else:
                     all_papers = searcher.expanded_search(
-                        query, client, model_cfg.smart, tracker, log=log)
+                        query, primary_client, primary_cfg.model_id, primary_cfg.prices, tracker,
+                        max_tokens=primary_cfg.triage_max_tokens * 4, log=log,
+                    )
                 log(f"✓ {len(all_papers)} unique papers found")
 
-                # Abstract-level relevance filter to cut downstream cost
-                log("⟳ Triaging with Claude...")
-                triager = PaperTriager(client, model_cfg.cheap, tracker)
+                log(f"⟳ Triaging papers (batches of {primary_cfg.batch_size})...")
+                triager = PaperTriager(
+                    primary_client, primary_cfg.model_id, primary_cfg.prices, tracker,
+                    max_tokens=primary_cfg.triage_max_tokens,
+                    batch_size=primary_cfg.batch_size,
+                )
                 triaged = triager.triage(all_papers, spec)
                 top = triaged[:num_papers]
                 log(f"✓ {len(top)} papers passed triage")
                 paper_lookup = {p["paperId"]: p for p in top}
 
-                # PDF fetch with arxiv-first URL preference, polite spacing
                 log("⟳ Fetching PDFs...")
                 downloaded = []
                 for i, paper in enumerate(top):
@@ -241,7 +306,6 @@ if run_btn:
                 progress_placeholder.empty()
                 log(f"✓ {len(downloaded)} PDFs ready")
 
-                # Pull raster figures + captions from each PDF, then caption pre-filter
                 log("⟳ Extracting figures...")
                 extractor = FigureExtractor()
                 all_figures = []
@@ -255,23 +319,27 @@ if run_btn:
                 filtered = extractor.caption_filter(all_figures, query)
                 log(f"  {len(filtered)} figures after caption filter (from {len(all_figures)} total)")
 
-                # Cheap vision pass: score every surviving figure against the spec
-                log(f"⟳ Scoring {len(filtered)} figures...")
-                scorer = FigureScorer(client, model_cfg.cheap, tracker)
-                primary_matches = []
-                for i, fig in enumerate(filtered):
-                    progress_placeholder.progress((i + 1) / len(filtered), text=f"Scoring {i+1}/{len(filtered)}")
-                    result = scorer.score(fig, spec)
-                    if result.get("confidence", 0) >= 0.5:
-                        primary_matches.append(fig)
-                progress_placeholder.empty()
+                log(f"⟳ Scoring {len(filtered)} figures (batches of {primary_cfg.batch_size})...")
+                scorer = FigureScorer(
+                    primary_client, primary_cfg.model_id, primary_cfg.prices, tracker,
+                    score_max_tokens=primary_cfg.score_max_tokens,
+                    verify_max_tokens=primary_cfg.verify_max_tokens,
+                    batch_size=primary_cfg.batch_size,
+                )
+                results = scorer.score_batch(filtered, spec)
+                primary_matches = [fig for fig, result in zip(filtered, results)
+                                   if result.get("confidence", 0) >= 0.5]
                 log(f"✓ {len(primary_matches)} primary matches")
 
-                # Optional smart-model verification on figures that passed primary scoring
                 verified = primary_matches
                 if run_verify and primary_matches:
                     log(f"⟳ Verifying {len(primary_matches)} matches...")
-                    verifier = FigureScorer(client, model_cfg.smart, tracker)
+                    verifier = FigureScorer(
+                        verify_client, verify_cfg.model_id, verify_cfg.prices, tracker,
+                        score_max_tokens=verify_cfg.score_max_tokens,
+                        verify_max_tokens=verify_cfg.verify_max_tokens,
+                        batch_size=1,
+                    )
                     verified = []
                     for i, fig in enumerate(primary_matches):
                         progress_placeholder.progress((i + 1) / len(primary_matches), text=f"Verifying {i+1}/{len(primary_matches)}")
@@ -298,6 +366,8 @@ if run_btn:
                 st.session_state.global_stats["searches"] += 1
 
             except Exception as e:
+                import traceback
+                traceback.print_exc()
                 log(f"✗ Error: {e}")
                 st.error(f"Pipeline error: {e}")
             finally:
@@ -306,7 +376,6 @@ if run_btn:
         st.rerun()
 
 
-# Render results: stats row, downloadable zip, then per-figure cards with metadata
 if st.session_state.results:
     res = st.session_state.results
     matches = res["matches"]
@@ -367,7 +436,6 @@ if st.session_state.results:
                 st.markdown("---")
 
 
-# Post-search feedback slider; submission logs to persistence layer
 if st.session_state.results and st.session_state.results.get("matches"):
     st.markdown("""
 <div class="feedback-box">
@@ -387,7 +455,6 @@ if st.session_state.results and st.session_state.results.get("matches"):
         st.success("Thanks!")
 
 
-# Aggregate stats across all sessions, loaded from persistence
 stats = st.session_state.global_stats
 n_ratings = len(stats["ratings"])
 avg = sum(stats["ratings"]) / n_ratings if n_ratings else 0
@@ -401,4 +468,4 @@ st.markdown(f"""
     <div class="stat-item"><div class="tally-num">{"—" if not n_ratings else f"{avg:.1f}"}</div><div class="tally-label">Avg score</div></div>
   </div>
 </div>
-""", unsafe_allow_html=True)
+""", unsafe_allow_html=True)