Core API Reference¶

Complete reference for querygym's core classes and functions. All functions include comprehensive docstrings in the source code.

Data Structures¶

QueryItem¶

Represents a query with its ID and text.

from querygym import QueryItem

query = QueryItem(qid="q1", text="what causes diabetes")

Attributes: - qid (str): Unique query identifier - text (str): Query text content

Docstring in code: See querygym/core/base.py

ReformulationResult¶

Result of a query reformulation operation.

result = ReformulationResult(
    qid="q1",
    original="diabetes",
    reformulated="diabetes causes symptoms treatment",
    metadata={}
)

Attributes: - qid (str): Query identifier - original (str): Original query text - reformulated (str): Reformulated query text - metadata (Dict[str, Any]): Additional metadata about the reformulation

Docstring in code: See querygym/core/base.py

MethodConfig¶

Configuration for a reformulation method.

config = MethodConfig(
    name="genqr_ensemble",
    params={"repeat_query_weight": 3},
    llm={"model": "gpt-4", "temperature": 0.8},
    seed=42,
    retries=2
)

Attributes: - name (str): Method name (e.g., "genqr", "query2doc") - params (Dict[str, Any]): Method-specific parameters - llm (Dict[str, Any]): LLM configuration (model, temperature, max_tokens, etc.) - seed (int): Random seed for reproducibility (default: 42) - retries (int): Number of retries on LLM failure (default: 2)

Docstring in code: See querygym/core/base.py

BaseReformulator¶

Base class for all query reformulation methods.

All reformulation methods inherit from this class and implement the reformulate method. This class provides common functionality for concatenating results and batch processing.

Key Methods:

reformulate¶

result = reformulator.reformulate(query, contexts=None)

Reformulate a single query.

Parameters: - q (QueryItem): Query to reformulate - contexts (Optional[List[str]]): Optional retrieved contexts (required for some methods)

Returns: ReformulationResult

reformulate_batch¶

results = reformulator.reformulate_batch(queries, ctx_map=None)

Reformulate multiple queries in batch with progress bar.

Parameters: - queries (List[QueryItem]): List of queries to reformulate - ctx_map (Optional[Dict[str, List[str]]]): Optional mapping of query IDs to context lists

Returns: List[ReformulationResult]

Example:

queries = [QueryItem("q1", "diabetes"), QueryItem("q2", "cancer")]
results = reformulator.reformulate_batch(queries)
for r in results:
    print(f"{r.qid}: {r.reformulated}")

concatenate_result¶

reformulated = reformulator.concatenate_result(
    original_query="diabetes",
    generated_content="causes symptoms treatment",
    query_repeats=3
)

Concatenate the original query with generated content based on the method's strategy.

Parameters: - original_query (str): The original query text - generated_content (str | List[str]): Content generated by the LLM - query_repeats (Optional[int]): Number of times to repeat the original query - content_separator (str): Separator between multiple content pieces (default: " ")

Returns: str - Concatenated reformulated query

Docstrings in code: See querygym/core/base.py

High-Level API¶

create_reformulator¶

Create a reformulator instance with sensible defaults.

import querygym as qg

reformulator = qg.create_reformulator(
    method_name="genqr_ensemble",
    model="gpt-4",
    params={"repeat_query_weight": 3},
    llm_config={"temperature": 0.8, "max_tokens": 256}
)

Parameters: - method_name (str): Name of the method (e.g., "genqr", "genqr_ensemble", "query2doc") - model (str): LLM model name (default: "gpt-4") - params (dict): Method-specific parameters (optional) - llm_config (dict): Additional LLM configuration (temperature, max_tokens, etc.) (optional) - prompt_bank_path (str): Path to prompt bank YAML (optional, uses bundled by default) - **kwargs: Additional MethodConfig parameters (seed, retries)

Returns: BaseReformulator instance

Raises: ValueError if method_name is not recognized

Docstring in code: See querygym/__init__.py

load_queries¶

Load queries from a local file.

# Load from TSV
queries = qg.load_queries("queries.tsv", format="tsv")

# Load from JSONL
queries = qg.load_queries("queries.jsonl", format="jsonl")

Parameters: - path (str): Path to queries file - format (str): File format - "tsv" or "jsonl" (default: "tsv") - **kwargs: Additional parameters for DataLoader.load_queries()

Returns: List[QueryItem]

Docstring in code: See querygym/__init__.py

load_qrels¶

Load qrels (relevance judgments) from a local file.

qrels = qg.load_qrels("qrels.txt")

Parameters: - path (str): Path to qrels file - format (str): File format - "trec" (default: "trec") - **kwargs: Additional parameters for DataLoader.load_qrels()

Returns: Dict[str, Dict[str, int]] - Nested dict: qid → docid → relevance

Docstring in code: See querygym/__init__.py

load_contexts¶

Load contexts from a JSONL file.

contexts = qg.load_contexts("contexts.jsonl")

Parameters: - path (str): Path to contexts JSONL file - **kwargs: Additional parameters for DataLoader.load_contexts()

Returns: Dict[str, List[str]] - Dict: qid → list of context strings

Docstring in code: See querygym/__init__.py

LLM Client¶

OpenAICompatibleClient¶

Client for OpenAI-compatible chat completion APIs.

Supports any API that implements the OpenAI chat completions format, including OpenAI, Azure OpenAI, local models via vLLM, etc.

from querygym.core.llm import OpenAICompatibleClient

client = OpenAICompatibleClient(
    model="gpt-4",
    base_url="https://api.openai.com/v1",
    api_key="sk-..."
)

response = client.chat([
    {"role": "system", "content": "You are helpful."},
    {"role": "user", "content": "Hello!"}
])

Constructor Parameters: - model (str): Model identifier (e.g., "gpt-4", "gpt-3.5-turbo") - base_url (str | None): Optional base URL for the API. If None, uses OPENAI_BASE_URL environment variable or OpenAI's default. - api_key (str | None): Optional API key. If None, uses OPENAI_API_KEY environment variable.

Methods:

chat¶

response = client.chat(
    messages=[{"role": "user", "content": "Hello"}],
    temperature=0.8,
    max_tokens=256
)

Send a chat completion request.

Parameters: - messages (List[Dict[str, str]]): List of message dicts with 'role' and 'content' keys - **kw: Additional parameters (temperature, max_tokens, etc.)

Returns: str - Generated text response

Docstrings in code: See querygym/core/llm.py

Prompt Bank¶

PromptBank¶

Manages prompt templates from a YAML file.

Loads and provides access to prompt templates with metadata. Supports rendering templates with variable substitution.

from querygym.core.prompts import PromptBank
from pathlib import Path

pb = PromptBank(Path("querygym/prompt_bank.yaml"))
messages = pb.render("genqr_keywords", query="diabetes")

Constructor Parameters: - path (str | Path): Path to the prompt bank YAML file

Methods:

render¶

messages = pb.render("genqr_keywords", query="diabetes")
# Returns: [{"role": "system", "content": "..."}, {"role": "user", "content": "..."}]

Render a prompt template with variable substitution.

Parameters: - prompt_id (str): ID of the prompt to render - **vars: Variables to substitute in the template (e.g., query="text")

Returns: List[Dict[str, str]] - List of message dicts ready for LLM chat completion

Raises: KeyError if prompt_id doesn't exist or required template variables are missing

list¶

prompt_ids = pb.list()

List all available prompt IDs.

Returns: List[str]

get¶

prompt_spec = pb.get("genqr_keywords")

Get a prompt specification by ID.

Returns: PromptSpec object

Raises: KeyError if prompt_id doesn't exist

get_meta¶

metadata = pb.get_meta("genqr_keywords")
# Returns: {"authors": [...], "license": "...", ...}

Get metadata for a prompt.

Returns: Dict[str, Any] - Metadata dict (authors, license, etc.)

Raises: KeyError if prompt_id doesn't exist

Docstrings in code: See querygym/core/prompts.py