How to Add Web Search to an AI Workflow or RAG Pipeline (Python, 2026)

"""Live web search for a RAG pipeline, using the Serpex Search API.

Needs:  pip install requests
Run:    SERPEX_API_KEY=sk_... python rag_example.py
"""
import math
import os
import re
import time
from collections import Counter
from datetime import date

import requests

SERPEX_URL = "https://api.serpex.dev/api/search"
API_KEY = os.environ["SERPEX_API_KEY"]
RETRYABLE = {429, 500, 503}  # safe to retry; errors are never charged
STOPWORDS = {"a", "an", "and", "are", "did", "does", "for", "how", "in", "is", "of",
             "on", "or", "the", "to", "what", "when", "which", "who", "why"}

def search_with_content(query: str, content_results: int = 5) -> dict:
    """One search that also returns the top results' pages as Markdown."""
    headers = {"Authorization": f"Bearer {API_KEY}"}
    payload = {"q": query, "include_content": True, "content_results": content_results}
    for attempt in range(4):  # the first try plus up to 3 retries
        resp = requests.post(SERPEX_URL, headers=headers, json=payload, timeout=60)
        if resp.status_code in RETRYABLE and attempt < 3:
            time.sleep(2 ** attempt)  # wait 1 s, 2 s, then 4 s
            continue
        resp.raise_for_status()
        return resp.json()

def chunk_markdown(markdown: str, max_chars: int = 1200) -> list:
    """Split Markdown on blank lines and pack the pieces into chunks of at most max_chars."""
    pieces = []
    for block in re.split(r"\n\s*\n", markdown):
        block = block.strip()
        # Hard-split any block longer than max_chars (big tables, long lines).
        pieces.extend(block[i:i + max_chars] for i in range(0, len(block), max_chars))
    chunks, current = [], ""
    for piece in pieces:
        if current and len(current) + 2 + len(piece) > max_chars:
            chunks.append(current)
            current = piece
        else:
            current = f"{current}\n\n{piece}" if current else piece
    if current:
        chunks.append(current)
    return chunks


def collect_chunks(results: list) -> list:
    """Page content becomes Markdown chunks; a result without content keeps its snippet."""
    chunks = []
    for r in results:
        if r.get("content"):
            texts, kind = chunk_markdown(r["content"]), "page"
        else:
            # content_error, or no content requested: keep the snippet as a short source.
            texts, kind = [r.get("snippet") or ""], "snippet"
        for text in texts:
            if text.strip():
                chunks.append({
                    "text": text,
                    "url": r["url"],
                    "title": r.get("title") or r["url"],
                    "position": r.get("position", 99),
                    "kind": kind,
                })
    return chunks

def terms(text: str) -> list:
    return [w for w in re.findall(r"[a-z0-9]+", text.lower()) if w not in STOPWORDS]


def score(text: str, query_terms: set) -> float:
    counts = Counter(terms(text))
    matched = [t for t in query_terms if counts[t]]
    # Each distinct query term found scores 1; repeats add a little (log-damped).
    return len(matched) + sum(math.log1p(counts[t]) for t in matched) / 10


def select_chunks(question: str, chunks: list, k: int = 6, per_source: int = 2) -> list:
    """Keep the k best chunks, at most per_source from any one URL."""
    query_terms = set(terms(question))
    ranked = sorted(chunks, key=lambda c: (-score(c["text"], query_terms), c["position"]))
    picked, used = [], Counter()
    for c in ranked:
        if used[c["url"]] < per_source:
            picked.append(c)
            used[c["url"]] += 1
        if len(picked) == k:
            break
    return picked

def build_prompt(question: str, picked: list) -> tuple:
    """Number each source once and list its chunks under that number."""
    sources = {}
    for c in picked:
        if c["url"] not in sources:
            sources[c["url"]] = {"n": len(sources) + 1, "title": c["title"],
                                 "url": c["url"], "kind": c["kind"], "texts": []}
        sources[c["url"]]["texts"].append(c["text"])
    blocks = []
    for s in sources.values():
        label = " (search snippet only)" if s["kind"] == "snippet" else ""
        body = "\n\n[...]\n\n".join(s["texts"])
        blocks.append(f"[{s['n']}] {s['title']}{label}\nURL: {s['url']}\n\n{body}")
    prompt = (
        f"Today's date is {date.today().isoformat()}.\n"
        "Answer the question using only the numbered sources below. "
        "Cite them inline as [1], [2] and so on. "
        "If the sources do not answer the question, say so.\n\n"
        f"Question: {question}\n\nSources:\n\n" + "\n\n---\n\n".join(blocks)
    )
    return prompt, list(sources.values())

def ask_llm(prompt: str) -> str:
    # Your model call goes here: send `prompt` to your LLM and return its reply.
    # This stub returns the first 500 characters so the script runs without one.
    return prompt[:500]


def main() -> None:
    question = "What changed in the EU AI Act in 2026?"
    data = search_with_content(question)
    meta = data.get("metadata", {})
    print(f"Pages with content: {meta.get('content_delivered')}/{meta.get('content_requested')}, "
          f"credits used: {meta.get('credits_used')}")
    results = data.get("results", [])
    if not results:
        print(data.get("message", "No results found."))
        return
    for r in results:
        if r.get("content_error"):
            print(f"  no content for {r['url']}: {r['content_error']}")
    picked = select_chunks(question, collect_chunks(results))
    prompt, sources = build_prompt(question, picked)
    print("\nSources:")
    for s in sources:
        print(f"[{s['n']}] {s['title']} - {s['url']}")
    print("\nAnswer:")
    print(ask_llm(prompt))


if __name__ == "__main__":
    main()