app.py

"""
app.py
--------------
Agentic RAG — HuggingFace Space entry point (ZeroGPU version).
"""
import spaces # noqa: F401 - MUST stay first; see docstring above

import time
import gradio as gr

from config import CONFIDENCE_THRESHOLD, MAX_PROMPT_TOKENS
from retrieval import run_pipeline_phase1
from generation import generate_answer
from utils import extract_sources, clean_final_answer
from ui_helpers import format_metrics


# Chat handler

def chat(message: str, history: list):
  """
  Sync generator for gr.ChatInterface(fn=chat, additional_outputs=[...]).

  This function has no GPU-specific code. It calls run_pipeline_phase1() (FAISS +
  Tavily, both CPU-only) and generate_answer() (which internally
  calls the @spaces.GPU-decorated _generate_on_gpu() exactly once).

  Yields 2-tuples: (bot_message_string, metrics_markdown_string).
  """
  t_total_start = time.perf_counter()

  # Step 1: Immediate placeholder
  yield "🔍 Searching...", format_metrics()

  # Step 2: Retrieval (FAISS + optional Tavily, CPU-only)
  state = run_pipeline_phase1(message, history=history)

  t_retrieve = state.get("t_retrieve", 0.0)
  t_web = state.get("t_web", 0.0)
  confidence = state.get("confidence", 0.0)
  source_type = state.get("source_type", "none")
  sources_count = state.get("sources_count", 0)

  if state.get("needs_web") and sources_count > 0:
    status = (
        f"🌐 {sources_count} sources fetched ({t_web:.1f}s). "
        f"Generating on GPU…"
    )
  elif state.get("needs_web"):
      status = "🌐 Web search done. Generating on GPU…"
  else:
      status = f"📚 FAISS (sim={confidence:.3f}). Generating on GPU…"

  yield status, format_metrics(
      source_type = source_type,
      confidence = confidence,
      sources_count = sources_count,
      generating = True,
      elapsed = time.perf_counter() - t_total_start,
      t_retrieve = t_retrieve,
      t_web = t_web,
  )

  # Step 3: Word-by-word generation
  # generate_answer() internally calls the @spaces.GPU-decorated
  # _generate_on_gpu() once (blocking - but on ZeroGPU this should
  # take a few seconds), then yields the response word by word.
  answer = ""
  output_tokens = 0
  prompt_tokens = 0
  t_gen_start = time.perf_counter()
  last_metrics = format_metrics(
    source_type = source_type,
    confidence = confidence,
    sources_count = sources_count,
    generating = True,
    elapsed = 0.0,
    t_retrieve = t_retrieve,
    t_web = t_web,
  )

  try:
    for word_text, prompt_tok_count in generate_answer(state):
      answer += word_text
      output_tokens += 1
      prompt_tokens = prompt_tok_count

      if output_tokens % 5 == 0:
        last_metrics = format_metrics(
            source_type   = source_type,
            confidence    = confidence,
            sources_count = sources_count,
            generating    = True,
            elapsed       = time.perf_counter() - t_gen_start,
            t_retrieve    = t_retrieve,
            t_web         = t_web,
            token_so_far  = output_tokens,
        )

      yield answer, last_metrics

  except Exception as e:
      error_msg = f"{type(e).__name__}: {e}"
      yield f"❌ Generation failed: {error_msg}", format_metrics(error=error_msg)
      return

  # Step 4: Final answer + sources + metrics
  t_generate     = time.perf_counter() - t_gen_start
  t_total        = time.perf_counter() - t_total_start
  tokens_per_sec = output_tokens / t_generate if t_generate > 0 else 0.0

  sources_section = extract_sources(state)
  final_answer    = clean_final_answer(answer) + sources_section

  yield final_answer, format_metrics(
      source_type    = source_type,
      confidence     = confidence,
      sources_count  = sources_count,
      prompt_tokens  = prompt_tokens,
      t_retrieve     = t_retrieve,
      t_web          = t_web,
      t_generate     = t_generate,
      t_total        = t_total,
      tokens_per_sec = tokens_per_sec,
      output_tokens  = output_tokens,
  )


# UI Layout

with gr.Blocks(theme=gr.themes.Soft(), title="Agentic RAG (ZeroGPU)") as demo:

  gr.Markdown(
      "## 🤖 Agentic RAG  "
  )
  gr.Markdown(
      "💡**Tip:** Include full context in every question for best results. \n"
      "Follow-up questions are supported but specific questions always "
      "perform better.\n\n"
      "⚡**Note:** Running on a shared free GPU (ZeroGPU). Most responses "
      "complete in a few seconds, though you may occasionally wait in a "
      "short queue if the shared GPU is busy with other Spaces."
  )

  with gr.Row():

      with gr.Column(scale=3):
          metrics_panel = gr.Markdown(value=format_metrics(), render=False)

          gr.ChatInterface(
              fn=chat,
              additional_outputs=[metrics_panel],
          )

      with gr.Column(scale=1, min_width=220):
          gr.Markdown("**📊 Metrics**")
          metrics_panel.render()


# Launch
# No share=True — HF Spaces provides its own public URL automatically.

demo.queue(default_concurrency_limit=1)
demo.launch()