pixelrag_mvp.py

import os, requests
from PIL import Image

OLLAMA = "http://localhost:11434/api"
MODEL = "gemma4:e2b"

def retrieve_wiki(query):
    res = requests.post("https://pixelrag.ai/api/search", json={"queries": [{"text": query}], "n_docs": 1}).json()
    hit = res["results"][0]["hits"][0]
    
    tile_url = f"https://pixelrag.ai/api/tile/{hit['article_id']}/{hit['tile_index']}/{hit['chunk_index']}"
    img_data = requests.get(tile_url).content
    os.makedirs("tiles", exist_ok=True)
    img_path = f"tiles/tile_{hit['article_id']}_{hit['tile_index']}.png"
    with open(img_path, "wb") as f:
        f.write(img_data)
        
    wiki_res = requests.get(
        f"https://en.wikipedia.org/w/api.php?action=query&prop=extracts&exintro&explaintext&format=json&titles={hit['url']}",
        headers={"User-Agent": "PixelRAGMVP/1.0"}
    ).json()
    pages = wiki_res["query"]["pages"]
    text = list(pages.values())[0].get("extract", "No extract found.")
    return {"path": img_path, "text": text, "url": hit["url"], "score": hit["score"]}

def ollama_generate(prompt):
    return requests.post(f"{OLLAMA}/generate", json={"model": MODEL, "prompt": prompt, "stream": False, "options": {"think": False}}).json()