query_expansion.py

"""Query-side expansion for figure retrieval.

expand_query   generates alternative phrasings of the user's request, which are
               searched alongside the original and fused. This targets
               underspecified queries, where a single embedding lands in a
               crowded region of the space.

suggest_authors asks the model which researchers are associated with the topic.
               Names that match nobody in the corpus are ignored downstream, so
               a hallucinated name is inert rather than harmful, but a confident
               wrong name still spends rank budget. Use as a boost, not a filter.
"""

import json
import re

EXPAND_PROMPT = """A researcher is searching for a figure in the astrophysics literature. Their query:

"{query}"

Write {n} alternative phrasings of the same search. Vary them:
- one that spells out any symbols or acronyms in words
- one that writes the same concepts in standard symbolic or abbreviated notation
- one that is more specific about what is plotted (axes, plot type), guessing sensibly if the query does not say
- one that uses the vocabulary a paper caption would use
- the rest phrased differently but with the same meaning

Do not change what is being searched for. Do not add science that is not implied by the query.

JSON only: {{"queries": ["...", "..."]}}"""

AUTHORS_PROMPT = """A researcher is searching for a figure in the astrophysics literature. Their query:

"{query}"

Name up to {n} researchers who have published substantially on this specific topic.
Give full names as they would appear in an author list. If you are not confident that
someone works on this exact topic, leave them out; a short list is better than a
speculative one. If nothing comes to mind, return an empty list.

JSON only: {{"authors": ["Family, Given", "..."]}}"""


def _parse(text: str, key: str) -> list:
    text = re.sub(r"```(json)?", "", text)
    start, end = text.find("{"), text.rfind("}")
    if start == -1 or end == -1:
        return []
    data = json.loads(text[start:end + 1])
    out = data.get(key, [])
    return [x for x in out if isinstance(x, str) and x.strip()]


def expand_query(query: str, client, model: str, n: int = 5,
                 max_tokens: int = 500) -> list[str]:
    resp = client.messages.create(
        model=model, max_tokens=max_tokens,
        messages=[{"role": "user",
                   "content": EXPAND_PROMPT.format(query=query, n=n)}],
    )
    try:
        return _parse(resp.content[0].text, "queries")[:n]
    except Exception:
        return []


def suggest_authors(query: str, client, model: str, n: int = 8,
                    max_tokens: int = 300) -> list[str]:
    resp = client.messages.create(
        model=model, max_tokens=max_tokens,
        messages=[{"role": "user",
                   "content": AUTHORS_PROMPT.format(query=query, n=n)}],
    )
    try:
        return _parse(resp.content[0].text, "authors")[:n]
    except Exception:
        return []