How do I count brand mentions and citations in scraped ChatGPT, Perplexity, Gemini and AI Overview answers?

#!/usr/bin/env python3
"""Collect repeated answers from AI answer engines and save them as one JSON array."""
import argparse
import json
import os
import sys
import time
from urllib.error import HTTPError, URLError
from urllib.parse import quote
from urllib.request import Request, urlopen

BASE_URL = "https://api.querying.ai"
ENGINES = ("CHATGPT", "PERPLEXITY", "GEMINI", "GOOGLE_AIO")
DONE = ("COMPLETED", "FAILED")  # QUEUED and PROCESSING mean poll again


class RetryLater(Exception):
    def __init__(self, seconds):
        super().__init__(f"retry after {seconds}s")
        self.seconds = seconds


def call(method, path, api_key, payload=None):
    data = None if payload is None else json.dumps(payload).encode("utf-8")
    request = Request(BASE_URL + path, data=data, method=method, headers={
        "Authorization": "Bearer " + api_key, "Content-Type": "application/json"})
    try:
        with urlopen(request, timeout=30) as response:
            return json.load(response)
    except HTTPError as error:
        if error.code in (429, 503):  # plan concurrency, queue backpressure or a short outage
            wait = error.headers.get("Retry-After", "")
            raise RetryLater(int(wait) if wait.isdigit() else 5) from error
        raise RuntimeError(f"{method} {path} returned HTTP {error.code}") from error
    except (URLError, TimeoutError) as error:
        raise RuntimeError(f"{method} {path} failed; check the dashboard before resubmitting") from error


def answer_of(engine, response):
    """Return (text, sources), or None when the engine showed no answer."""
    if engine == "GOOGLE_AIO":
        response = response.get("aioverview")  # null means Google showed no AI Overview
        if response is None:
            return None
    return response["text"], response["sources"]


def collect(api_key, engines, prompts, runs, country, rows, max_in_flight=2, task_timeout=600):
    """Append one row per task to rows, so the caller keeps them if a call fails."""
    waiting = [(e, p, r) for e in engines for p in prompts for r in range(1, runs + 1)]
    running = {}
    while waiting or running:
        while waiting and len(running) < max_in_flight:
            engine, prompt, run = waiting[0]
            body = {"taskType": engine,
                    "payload": {"prompt": prompt, "query": prompt, "country": country}}
            try:
                task = call("POST", "/v1/async/task", api_key, body)["task"]
            except RetryLater as busy:
                time.sleep(busy.seconds)
                break
            waiting.pop(0)
            running[task["id"]] = (engine, prompt, run, time.monotonic())
        for task_id, (engine, prompt, run, started) in list(running.items()):
            try:
                result = call("GET", "/v1/async/task/" + quote(task_id, safe=""), api_key)
            except RetryLater as busy:
                time.sleep(busy.seconds)
                continue
            status = result["task"]["status"]
            if status not in DONE and time.monotonic() - started < task_timeout:
                continue
            del running[task_id]
            row = {"task_id": task_id, "engine": engine, "country": country,
                   "prompt": prompt, "run": run, "status": status}
            if status == "COMPLETED":
                answer = answer_of(engine, result["response"])
                if answer is None:
                    row["status"] = "NO_ANSWER"
                else:
                    row["text"], row["sources"] = answer
            elif status != "FAILED":
                row["status"] = "TIMED_OUT"  # still running; GET it later by task_id
            rows.append(row)
            print(f"{len(rows)} {engine} run {run}: {row['status']}", file=sys.stderr)
        if running:
            time.sleep(3)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("prompts", help="Text file with one prompt per line")
    parser.add_argument("--engine", action="append", choices=ENGINES, required=True,
                        help="Repeat for each engine")
    parser.add_argument("--runs", type=int, default=5, help="Runs per prompt and engine")
    parser.add_argument("--country", default="US")
    parser.add_argument("--max-in-flight", type=int, default=2,
                        help="Tasks running at once; match your plan's concurrency")
    args = parser.parse_args()
    api_key = os.environ.get("QUERYING_API_KEY")
    if not api_key:
        print("Set QUERYING_API_KEY first", file=sys.stderr)
        return 1
    with open(args.prompts, encoding="utf-8") as handle:
        prompts = [line.strip() for line in handle if line.strip()]
    rows, code = [], 0
    try:
        collect(api_key, args.engine, prompts, args.runs, args.country, rows, args.max_in_flight)
    except (RuntimeError, KeyboardInterrupt) as error:
        print(f"Stopped: {error}. Printing the {len(rows)} rows collected so far.", file=sys.stderr)
        code = 1
    print(json.dumps(rows, ensure_ascii=False, indent=2))
    return code


if __name__ == "__main__":
    sys.exit(main())

export QUERYING_API_KEY='<your-key>'
python3 collect_answers.py prompts.txt --engine CHATGPT --engine PERPLEXITY --engine GEMINI --engine GOOGLE_AIO --runs 5 --country US > answers.json

#!/usr/bin/env python3
"""Count brand mentions and owned-site citations per engine in collected answers."""
import argparse
import json
import re
import sys
from pathlib import Path
from urllib.parse import urlsplit


def count_engine(rows, patterns, fold, domain):
    counts = {'completed_answers': 0, 'mentioned_answers': 0,
              'owned_cited_answers': 0, 'both_answers': 0,
              'owned_cited_without_mention': 0, 'answers_with_sources': 0,
              'answers_without_sources': 0, 'not_counted': {}}
    mentioned_ids = []
    for row in rows:
        status = row.get('status', 'COMPLETED')
        if status != 'COMPLETED':
            counts['not_counted'][status] = counts['not_counted'].get(status, 0) + 1
            continue
        task_id, text, sources = row['task_id'], row.get('text'), row.get('sources')
        if not isinstance(text, str) or not isinstance(sources, list):
            raise ValueError(f'{task_id}: a completed answer needs text and sources; missing data is not zero')
        mention = any(p.search(fold(text)) for p in patterns)
        owned = False
        for source in sources:
            if not isinstance(source, dict) or not isinstance(source.get('url'), str):
                raise ValueError(f'{task_id}: every source needs a URL string')
            try:
                raw = source['url'].strip()
                url = urlsplit(raw)
                if not url.scheme and not url.netloc:
                    # Some answers list a bare host such as "schema.org"; read it as that host.
                    url = urlsplit('//' + raw)
                host = (url.hostname or '').casefold().rstrip('.')
            except ValueError as exc:
                raise ValueError(f'{task_id}: invalid source URL') from exc
            if url.scheme not in ('', 'http', 'https') or not host:
                raise ValueError(f'{task_id}: source URLs must be HTTP(S) URLs or bare hosts')
            owned |= host == domain or host.endswith('.' + domain)
        counts['completed_answers'] += 1
        if mention:
            mentioned_ids.append(task_id)
        counts['answers_with_sources'] += int(bool(sources))
        counts['answers_without_sources'] += int(not sources)
        counts['mentioned_answers'] += int(mention)
        counts['owned_cited_answers'] += int(owned)
        counts['both_answers'] += int(mention and owned)
        counts['owned_cited_without_mention'] += int(owned and not mention)
    denominator = counts['completed_answers']
    for name, part in (('mention_rate', 'mentioned_answers'),
                       ('owned_citation_rate', 'owned_cited_answers'),
                       ('source_coverage_rate', 'answers_with_sources')):
        counts[name] = counts[part] / denominator if denominator else None
    counts['mentioned_task_ids'] = mentioned_ids
    return counts


def summarize(rows, aliases, domain, case_sensitive=False):
    if not isinstance(rows, list) or not aliases or any(not a.strip() for a in aliases):
        raise ValueError('Provide a JSON array and at least one non-empty brand alias')
    domain = domain.casefold().rstrip('.')
    if not domain or '/' in domain or ':' in domain or ' ' in domain:
        raise ValueError('Domain must be a hostname, for example example.com')
    fold = (lambda value: value) if case_sensitive else str.casefold
    patterns = [re.compile(r'(?<![A-Za-z0-9])' + re.escape(fold(a)) + r'(?![A-Za-z0-9])')
                if a.isascii() else re.compile(re.escape(fold(a))) for a in aliases]
    seen, by_engine = set(), {}
    for row in rows:
        if not isinstance(row, dict) or not isinstance(row.get('task_id'), str) or not row['task_id']:
            raise ValueError('Every row must be an answer object with a task_id')
        if row['task_id'] in seen:
            raise ValueError('Duplicate task_id; deduplicate the export before counting')
        seen.add(row['task_id'])
        by_engine.setdefault(row.get('engine', 'unknown'), []).append(row)
    # Engines answer differently, so each engine keeps its own denominator.
    return {engine: count_engine(group, patterns, fold, domain)
            for engine, group in by_engine.items()}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('export', type=Path, help='JSON array written by collect_answers.py')
    parser.add_argument('--brand', action='append', required=True, help='Repeat for each literal brand alias')
    parser.add_argument('--domain', required=True, help='Owned hostname; subdomains count too')
    parser.add_argument('--case-sensitive', action='store_true',
                        help='Match aliases with their exact capitalization')
    args = parser.parse_args()
    try:
        result = summarize(json.loads(args.export.read_text()), args.brand, args.domain, args.case_sensitive)
    except (OSError, ValueError) as exc:
        print(str(exc), file=sys.stderr)
        return 1
    print(json.dumps(result, ensure_ascii=False, indent=2))
    return 0


if __name__ == '__main__':
    sys.exit(main())

python3 count_by_engine.py answers.json --brand ExampleBrand --domain example.com