app.py

"""
Hunter Omega — Live NIAH Demo
One-click needle in a haystack test from 8K to 12M tokens.
Uses real Wikipedia content as filler. Shows needle BEFORE querying.
"""
import gradio as gr
import requests
import random
import string
import time
import os
import re
import json

API_URL = os.environ.get("HUNTER_API")
MODEL = "hunter-omega-12m-fast"

SIZE_MAP = {
    "8K": 8000, "32K": 32000, "64K": 64000, "128K": 128000,
    "256K": 256000, "512K": 512000, "1M": 1000000,
    "1.5M": 1500000, "3M": 3000000, "6M": 6000000, "12M": 12000000,
}

WIKI_CACHE = []

def load_wikipedia():
    global WIKI_CACHE
    try:
        with open("wiki_cache.json") as f:
            WIKI_CACHE = json.load(f)
        print(f"Loaded {len(WIKI_CACHE)} Wikipedia articles")
    except:
        # Fallback: download
        articles = [
            "Solar_System", "Roman_Empire", "Photosynthesis", "Industrial_Revolution",
            "Quantum_mechanics", "French_Revolution", "DNA", "Climate_change",
            "Ancient_Egypt", "Theory_of_relativity", "Renaissance", "American_Civil_War",
        ]
        for title in articles:
            try:
                r = requests.get(f"https://en.wikipedia.org/api/rest_v1/page/html/{title}",
                                headers={"User-Agent": "Demo/1.0"}, timeout=10)
                if r.status_code == 200:
                    text = re.sub(r'<[^>]+>', ' ', r.text)
                    text = re.sub(r'\s+', ' ', text).strip()
                    WIKI_CACHE.append(' '.join(text.split()[:5000]))
            except:
                pass
        print(f"Downloaded {len(WIKI_CACHE)} articles")

load_wikipedia()


def generate_needle():
    code = ''.join(random.choices(string.ascii_uppercase, k=4))
    nums = ''.join(random.choices(string.digits, k=4))
    word = random.choice(["ALPHA", "BRAVO", "DELTA", "ECHO", "FOXTROT", "GOLF", "HOTEL",
                          "INDIA", "JULIET", "KILO", "LIMA", "MIKE", "NOVEMBER", "OSCAR",
                          "PAPA", "QUEBEC", "ROMEO", "SIERRA", "TANGO", "UNIFORM", "VICTOR"])
    return f"{code}-{nums}-{word}"


def generate_haystack(target_tokens):
    words = []
    target_words = target_tokens * 3 // 4
    i = 0
    while len(words) < target_words:
        if WIKI_CACHE:
            words.extend(WIKI_CACHE[i % len(WIKI_CACHE)].split())
        else:
            words.extend("The quick brown fox jumps over the lazy dog.".split())
        i += 1
    return words[:target_words]


def run_test(context_size):
    tokens = SIZE_MAP.get(context_size, 128000)
    needle_code = generate_needle()
    needle = f"The classified authorization code for Operation Starfall is {needle_code}."
    query = "What is the classified authorization code for Operation Starfall? State ONLY the code."
    depth = random.randint(10, 90) if tokens <= 1000000 else random.randint(55, 95)

    words = generate_haystack(tokens)
    insert_pos = int(len(words) * depth / 100)
    words.insert(insert_pos, needle)
    actual_words = len(words)
    actual_tokens = actual_words * 4 // 3

    wiki_count = len(WIKI_CACHE)

    # Estimate time for user
    time_estimates = {
        8000: "2-10s", 32000: "2-10s", 64000: "3-10s", 128000: "5-15s",
        256000: "15-30s", 512000: "30-60s", 1000000: "1-2 min",
        1500000: "2-3 min", 3000000: "4-6 min", 6000000: "8-12 min", 12000000: "20-30 min",
    }
    est = time_estimates.get(tokens, "unknown")

    # STEP 1: Show needle before querying
    setup = f"""## Step 1: Test Generated

**Context size:** {actual_tokens:,} tokens (~{actual_words:,} words)
**Filler:** Real Wikipedia articles ({wiki_count} articles including Solar System, Roman Empire, Quantum Mechanics, DNA, Climate Change, etc.)
**Needle:** `{needle_code}`
**Full sentence:** "{needle}"
**Position:** {depth}% depth (word {insert_pos:,} of {actual_words:,})
**Estimated time:** {est}

The passphrase above was just generated randomly. It is now hidden in {actual_tokens:,} tokens of Wikipedia text.

**Do not refresh your browser.** The result will appear below when complete.

---
## Step 2: """

    if tokens <= 1000000:
        # Native attention
        yield setup + "Querying model (native attention)...", "", "QUERYING...", ""

        haystack = " ".join(words)
        t0 = time.time()
        try:
            r = requests.post(API_URL, json={
                "model": MODEL,
                "messages": [{"role": "user", "content": haystack + "\n\n" + query}],
                "max_tokens": 50,
                "temperature": 0,
            }, timeout=300)
            elapsed = time.time() - t0

            if r.status_code != 200:
                yield setup + f"Server returned {r.status_code}. May be busy, please retry.", "", "ERROR", ""
                return

            answer = r.json()["choices"][0]["message"]["content"]
            found = needle_code.lower() in answer.lower()

        except requests.exceptions.Timeout:
            yield setup + "Timed out. Server may be busy. Please retry.", "", "TIMEOUT", ""
            return
        except Exception:
            yield setup + "Connection error. Please retry.", "", "ERROR", ""
            return
    else:
        # Extended context
        doc = " ".join(words)
        user_id = f"demo_{random.randint(100000, 999999)}"
        api_base = API_URL.replace("/v1/chat/completions", "")

        yield setup + f"Indexing {actual_tokens:,} tokens (this will take a few minutes)...", "", "INDEXING...", ""

        t0 = time.time()
        try:
            r = requests.post(f"{api_base}/v1/documents", json={
                "user_id": user_id, "text": doc,
            }, timeout=1800)
            index_time = time.time() - t0
            info = r.json()
            chunks = info.get("overflow_chunks", 0)

            yield setup + f"Indexed {chunks} chunks in {index_time:.0f}s. Now querying...", "", "QUERYING...", f"Index: {index_time:.0f}s"

            r = requests.post(API_URL, json={
                "model": MODEL,
                "messages": [{"role": "user", "content": query}],
                "max_tokens": 50,
                "temperature": 0,
                "user": user_id,
            }, timeout=600)
            elapsed = time.time() - t0
            answer = r.json()["choices"][0]["message"]["content"]
            found = needle_code.lower() in answer.lower()

            requests.delete(f"{api_base}/v1/documents/{user_id}", timeout=10)

        except requests.exceptions.Timeout:
            try: requests.delete(f"{api_base}/v1/documents/{user_id}", timeout=10)
            except: pass
            yield setup + f"Timed out after {time.time()-t0:.0f}s. Large documents take time. Please retry.", "", "TIMEOUT", f"{time.time()-t0:.0f}s"
            return
        except Exception:
            try: requests.delete(f"{api_base}/v1/documents/{user_id}", timeout=10)
            except: pass
            yield setup + "Connection error.", "", "ERROR", ""
            return

    # STEP 3: Result
    result = f"""Complete!

---
## Step 3: Result {"✅ NEEDLE FOUND" if found else "❌ NOT FOUND"}

**Expected:** `{needle_code}`
**Model returned:** `{answer.strip()}`
**Match:** {"YES" if found else "NO"}
**Time:** {elapsed:.1f}s
**Method:** {"Native full attention" if tokens <= 1000000 else "Extended context engine"}

{"The model correctly found the randomly generated passphrase in " + f"{actual_tokens:,} tokens of real Wikipedia content." if found else "The model did not find the needle."}

*Click 'Find the Needle' again to run another test with a fresh random passphrase.*"""

    yield setup + result, "", "FOUND" if found else "MISSED", f"{elapsed:.1f}s"


with gr.Blocks(title="Hunter Omega — Live NIAH Demo", css="""
    .gradio-container { max-width: 900px !important; }
    footer { display: none !important; }
""") as demo:
    gr.Markdown("""
# Hunter Omega — Needle in a Haystack (Live Demo)

**How this works:**
1. Pick a context length (8K to 12M tokens)
2. A **random passphrase** is generated and hidden at a **random depth** in **real Wikipedia articles**
3. You see the passphrase and its position **before** the model runs
4. The model searches and retrieves it
5. You verify the match

Fresh random passphrase every time. Real Wikipedia content, not repetitive filler text.

**8K to 1M:** Native full attention &nbsp;|&nbsp; **1.5M to 12M:** Extended effective context

> **Do not refresh your browser during a test.** Results appear in place when complete. Large contexts (6M+) can take several minutes.
    """)

    with gr.Row():
        context_dropdown = gr.Dropdown(
            choices=["8K", "32K", "64K", "128K", "256K", "512K", "1M", "1.5M", "3M", "6M", "12M"],
            value="128K",
            label="Context Length",
        )
        run_btn = gr.Button("Find the Needle", variant="primary", size="lg")

    output = gr.Markdown(label="Test Progress", value="Select a context length above and click **Find the Needle** to begin.")

    with gr.Row():
        status_box = gr.Textbox(label="Status", interactive=False, scale=1)
        time_box = gr.Textbox(label="Time", interactive=False, scale=1)

    hidden = gr.Textbox(visible=False)

    gr.Markdown("""
---

<details>
<summary><b>Estimated times and details</b></summary>

| Size | Method | Time (no queue) |
|:----:|:------:|:---------------:|
| 8K - 128K | Native attention | 2-15s |
| 256K - 512K | Native attention | 15-60s |
| 1M | Native attention | 1-2 min |
| 1.5M - 3M | Extended context | 2-6 min |
| 6M | Extended context | 8-12 min |
| 12M | Extended context | 20-30 min |

If multiple users are testing at the same time, requests queue. The model processes large-context requests one at a time.

**Filler:** 16 real Wikipedia articles (Solar System, Roman Empire, Quantum Mechanics, French Revolution, DNA, Climate Change, Ancient Egypt, Industrial Revolution, and more).

</details>

[Full benchmarks and methodology](https://github.com/SovNodeAI/hunter-omega-benchmarks) &nbsp;|&nbsp; Running on a single GPU cluster
    """)

    run_btn.click(
        fn=run_test,
        inputs=[context_dropdown],
        outputs=[output, hidden, status_box, time_box],
    )


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)