app.py

# ── CELL 1: Imports & API Keys ──────────────────────────────────────────────
# Runs the LLM locally on Hugging Face Spaces' free ZeroGPU hardware.
# ZeroGPU only grants a GPU for the duration of a function decorated with
# @spaces.GPU — the model itself still lives in the Space's normal CPU RAM
# the rest of the time, so keep the model small enough to fit there too.

import os
import time
import uuid
from typing import TypedDict
from dotenv import load_dotenv

import torch
import spaces                                                  # HF ZeroGPU
from transformers import AutoModelForCausalLM, AutoTokenizer
from ddgs import DDGS                                          # DuckDuckGo search
from langchain_core.messages import SystemMessage, HumanMessage
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver
from langgraph.types import interrupt, Command

load_dotenv()


# ── Logging helper: BACKEND ONLY — this prints to the terminal/server console.
# It is intentionally NOT surfaced in the Streamlit UI; the UI only shows a
# generic "running" status. Watch your terminal to see step-by-step progress.
def log(msg: str):
    print(msg)


# ── CELL 2: State ────────────────────────────────────────────────────────────
class ResearchState(TypedDict):
    query:          str   # Original question from the user
    web_results:    str   # Formatted text of DuckDuckGo search results
    search_query:   str   # Rewritten / improved search query
    final_answer:   str   # LLM-synthesized answer
    needs_retry:    bool  # True if critic says answer is poor
    retry_count:    int   # Number of retry attempts so far
    human_approved: bool  # Whether the human approved the answer
    human_feedback: str   # Feedback from the human if they rejected


# ── CELL 3: LLM (local model, generated on ZeroGPU) ──────────────────────────
# Qwen2.5-3B-Instruct is the default: it follows formatting instructions well
# (Markdown, ## headings, [1]/[2] citations), is a reliable yes/no judge for
# the critic step, and its weights (~6GB in bf16) comfortably fit in the 16GB
# CPU RAM of a free Space alongside the rest of the app. Bump MODEL_ID up to
# Qwen/Qwen2.5-7B-Instruct via env var if you're on upgraded (more RAM) hardware.
MODEL_ID = os.environ.get("HF_MODEL_ID", "Qwen/Qwen2.5-3B-Instruct")

log_prefix = "🧩"
print(f"{log_prefix} Loading tokenizer & model weights for '{MODEL_ID}' into CPU RAM...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
)
print(f"{log_prefix} Model loaded. GPU will be attached on demand via ZeroGPU.")


class _LLMResponse:
    """Thin wrapper so call sites can keep using `response.content`."""
    def __init__(self, content: str):
        self.content = content


def _to_chat_messages(messages):
    chat = []
    for m in messages:
        role = "system" if isinstance(m, SystemMessage) else "user"
        chat.append({"role": role, "content": m.content})
    return chat


@spaces.GPU(duration=90)
def _generate(chat_messages, max_new_tokens: int = 1024) -> str:
    """Runs on a GPU that ZeroGPU attaches only for this call's duration."""
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device)
    prompt_text = tokenizer.apply_chat_template(
        chat_messages, tokenize=False, add_generation_prompt=True
    )
    inputs = tokenizer(prompt_text, return_tensors="pt").to(device)
    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=0.3,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
        )
    new_tokens = output_ids[0][inputs["input_ids"].shape[-1]:]
    return tokenizer.decode(new_tokens, skip_special_tokens=True).strip()


class _LocalLLM:
    def invoke(self, messages):
        chat_messages = _to_chat_messages(messages)
        text = _generate(chat_messages)
        return _LLMResponse(text)


llm = _LocalLLM()


# ── Retry helper: handles transient errors (ZeroGPU quota momentarily busy, ──
# brief CUDA hiccup) with a short backoff. Not meant to retry hard failures
# like out-of-memory — those are raised immediately.
def llm_invoke_with_retry(messages, max_retries: int = 2, base_wait: float = 3.0):
    """
    Calls llm.invoke(messages). If a transient GPU-allocation or timeout error
    occurs, waits and retries up to max_retries times with exponential backoff.
    Raises the last exception if all retries are exhausted.
    """
    for attempt in range(max_retries + 1):
        try:
            if attempt == 0:
                log(f"    📡 Generating with local model '{MODEL_ID}' on ZeroGPU...")
            return llm.invoke(messages)
        except Exception as e:
            err_str = str(e)
            is_transient = any(s in err_str.lower() for s in
                                ["timeout", "timed out", "queue", "busy", "unavailable"])
            if is_transient and attempt < max_retries:
                wait_time = base_wait * (2 ** attempt)   # 3s, 6s
                log(f"    ⚠️  ZeroGPU not available yet. Waiting {wait_time:.0f}s "
                    f"before retry {attempt + 1}/{max_retries}...")
                time.sleep(wait_time)
                log(f"    📡 Retrying generation ({attempt + 1}/{max_retries})...")
            else:
                log(f"    ❌ Generation failed permanently: {err_str[:200]}")
                raise


# ── CELL 4: Node 1 — Query Rewrite ───────────────────────────────────────────
REWRITE_PROMPT = """\
You rewrite user questions into a short, high-signal web search query.
Return ONLY the search query text. No explanation, no quotes.
"""

def query_rewrite_node(state: ResearchState) -> dict:
    log("🔄 STEP: Query Rewrite — turning your question into a search query")
    feedback = state.get("human_feedback", "").strip()

    if feedback:
        user_content = (
            f"Original question: {state['query']}\n"
            f"Previous answer was rejected. Human feedback: {feedback}\n"
            f"Rewrite the search query to address this feedback."
        )
        log(f"    Using human feedback to improve query: '{feedback}'")
    else:
        user_content = state["query"]

    messages = [
        SystemMessage(content=REWRITE_PROMPT),
        HumanMessage(content=user_content)
    ]
    response = llm_invoke_with_retry(messages)
    rewritten = response.content.strip()

    log(f"    ✅ Search query: {rewritten}")
    return {"search_query": rewritten}


# ── CELL 5: Node 2 — Web Search (DuckDuckGo) ─────────────────────────────────
def web_search_node(state: ResearchState) -> dict:
    log(f"🌐 STEP: Web Search — querying DuckDuckGo for '{state['search_query']}'")
    formatted_results = []

    try:
        with DDGS() as ddgs:
            raw = list(ddgs.text(state["search_query"], max_results=5))
        log(f"    ✅ Got {len(raw)} results")
    except Exception as e:
        log(f"    ❌ DuckDuckGo search error: {e}")
        raw = []

    for i, result in enumerate(raw, start=1):
        title   = result.get("title",  "No title")
        url     = result.get("href",   "")          # DuckDuckGo uses 'href' not 'url'
        content = result.get("body",   "").strip()  # DuckDuckGo uses 'body' not 'content'
        formatted_results.append(
            f"[{i}] {title}\n"
            f"    URL: {url}\n"
            f"    {content}"
        )

    web_results_text = "\n\n".join(formatted_results)
    return {"web_results": web_results_text}


# ── CELL 6: Node 3 — Synthesizer (Ollama) ────────────────────────────────────
SYSTEM_PROMPT = """\
You are an expert research assistant.

You will be given:
  1. A user research query.
  2. Five web search results (title, URL, snippet).

Your job:
  - Synthesize the web results with your own knowledge.
  - Write a clear, well-structured, comprehensive answer using Markdown.
  - Use ## headings to organize sections (e.g. ## Overview, ## Key Findings).
  - Cite sources inline using [1], [2], etc. where relevant.
  - Highlight any conflicting information across sources.
  - End with a "## References" section listing all URLs as:
      [1] https://...
      [2] https://...
  - Be factual. Do not hallucinate beyond the provided data.
"""

def synthesizer_node(state: ResearchState) -> dict:
    log("🧠 STEP: Synthesizer — asking the local Ollama model to write the sourced answer")
    user_prompt = (
        f"Research Query: {state['query']}\n\n"
        f"Web Search Results:\n\n{state['web_results']}"
    )

    messages = [
        SystemMessage(content=SYSTEM_PROMPT),
        HumanMessage(content=user_prompt)
    ]

    response = llm_invoke_with_retry(messages)
    log("    ✅ Answer drafted")
    return {"final_answer": response.content}


# ── CELL 7: Node 4 — Critic ──────────────────────────────────────────────────
CRITIC_PROMPT = """\
You judge whether a research answer is well supported by the given sources.
Reply with exactly one word: "yes" if well supported,
"no" if it is thin, vague, or barely uses the sources.
"""

def critic_node(state: ResearchState) -> dict:
    log("🔍 STEP: Critic — checking whether the answer is well supported")
    user_prompt = (
        f"Sources:\n{state['web_results']}\n\n"
        f"Answer:\n{state['final_answer']}"
    )

    messages = [
        SystemMessage(content=CRITIC_PROMPT),
        HumanMessage(content=user_prompt)
    ]

    verdict = llm_invoke_with_retry(messages).content.strip().lower()

    if verdict.startswith("no") and state["retry_count"] < 3:
        needs_retry = True
        log(f"    🔁 Verdict: '{verdict}' — retrying search (attempt {state['retry_count'] + 1}/3)")
    else:
        needs_retry = False
        log(f"    ✅ Verdict: '{verdict}' — moving to human review")

    return {
        "needs_retry": needs_retry,
        "retry_count": state["retry_count"] + 1
    }


# ── CELL 8: Node 5 — Human Approval ──────────────────────────────────────────
def human_approval_node(state: ResearchState) -> dict:
    log("✅ STEP: Human Approval — pausing for your review")

    # Pause graph — resumes when user submits feedback via Streamlit
    user_input = interrupt("Type 'ok' to approve, or give feedback to improve the answer: ")

    approved = user_input.strip().lower() == "ok"

    if approved:
        log("✅ Approved!")
        return {
            "human_approved": True,
            "human_feedback": ""
        }
    else:
        feedback = user_input.strip()
        log(f"🔄 Rejected. Feedback: '{feedback}'. Retrying with improved query...")
        return {
            "human_approved": False,
            "human_feedback": feedback,
            "retry_count":    0
        }


# ── CELL 9: Build Graph ───────────────────────────────────────────────────────
#
#   START → query_rewrite → web_search → synthesizer → critic
#               ↑                                         |
#               └──── needs_retry=True ───────────────────┘
#                              (else)
#                                ↓
#                        human_approval   ← graph PAUSES here
#                        /           \
#                  approved         rejected (with feedback)
#                     ↓                  ↓
#                    END          query_rewrite  (smarter retry)

def route_after_critic(state: ResearchState):
    return "web_search" if state["needs_retry"] else "human_approval"

def route_after_human(state: ResearchState):
    return "end" if state["human_approved"] else "query_rewrite"


workflow = StateGraph(ResearchState)

workflow.add_node("query_rewrite",  query_rewrite_node)
workflow.add_node("web_search",     web_search_node)
workflow.add_node("synthesizer",    synthesizer_node)
workflow.add_node("critic",         critic_node)
workflow.add_node("human_approval", human_approval_node)

workflow.add_edge(START,           "query_rewrite")
workflow.add_edge("query_rewrite", "web_search")
workflow.add_edge("web_search",    "synthesizer")
workflow.add_edge("synthesizer",   "critic")

workflow.add_conditional_edges(
    "critic",
    route_after_critic,
    {"web_search": "web_search", "human_approval": "human_approval"}
)

workflow.add_conditional_edges(
    "human_approval",
    route_after_human,
    {"end": END, "query_rewrite": "query_rewrite"}
)

memory = MemorySaver()
agent  = workflow.compile(checkpointer=memory)


# ── Helper: create a fresh thread config ─────────────────────────────────────
def new_thread() -> dict:
    return {"configurable": {"thread_id": str(uuid.uuid4())}}