#!/usr/bin/env python3
"""Enhance a task description with context from referenced papers.

The enhance tool fetches and summarises papers referenced in the input
text, producing enriched text that gives AI agents much better context
to work with. This is useful as a preprocessing step before submitting
tasks for analysis.

Supported sources:
  - arXiv URLs
  - bioRxiv / medRxiv URLs
  - PubMed / PMC URLs
  - Generic PDF links (e.g. journal papers)
  - Loose bibliographic references (e.g. "Smith et al. 2020"),
    resolved to arXiv papers via Skepthical
"""

import os

import air

client = air.AIR(
    api_key=os.environ["AIR_API_KEY"],
    base_url=os.environ.get("AIR_BASE_URL", "http://localhost:8000"),
)

text = """
We want to implement the variational quantum eigensolver (VQE) described
in https://arxiv.org/abs/1304.3061 to estimate the ground state energy
of small molecules. We will compare with the quantum approximate
optimization algorithm (QAOA) from https://arxiv.org/abs/1411.4028
on a set of combinatorial benchmark problems.
"""

print(f"Input:\n{text.strip()}\n")
print("Enhancing (this may take a minute) ...")

# Options:
#   max_workers         — parallel download workers (default 2)
#   max_depth           — max summarisation depth (default 10)
#   resolve_references  — resolve loose refs like "Smith et al. 2020" (default True)
#   max_references      — max loose references to resolve (default 10)
enhanced = client.enhance(text)

print(f"\nEnhanced text:\n{enhanced}")

client.close()
