Exa Search

#!/usr/bin/env python3
# /// script
# requires-python = ">=3.11"
# dependencies = ["exa-py>=1.14.0"]
# ///
"""Fetch and extract content from URLs using Exa's /contents endpoint.

Example:
    uv run exa_extract.py \\
        https://arxiv.org/abs/2401.04088 \\
        https://www.nature.com/articles/s41586-024-07566-y \\
        --text \\
        -o extracted.json
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from dataclasses import asdict, dataclass, field
from typing import Any

try:
    from exa_py import Exa
except ImportError:
    print(
        "exa_py not installed. Run: uv pip install exa-py  (or invoke with: uv run --with exa-py)",
        file=sys.stderr,
    )
    sys.exit(2)


EXA_INTEGRATION_HEADER = "k-dense-ai--scientific-agent-skills"


@dataclass
class ExtractedDocument:
    """Typed view of a single extracted document for JSON export."""

    url: str
    id: str | None
    title: str | None
    author: str | None
    published_date: str | None
    text: str | None = None
    highlights: list[str] = field(default_factory=list)


def _build_contents(text: bool, highlights: bool) -> dict[str, Any]:
    contents: dict[str, Any] = {}
    if text:
        contents["text"] = True
    if highlights:
        contents["highlights"] = True
    if not contents:
        # Default to full text when the caller doesn't pick anything.
        contents["text"] = True
    return contents


def _to_typed(item: Any) -> ExtractedDocument:
    return ExtractedDocument(
        url=getattr(item, "url", ""),
        id=getattr(item, "id", None),
        title=getattr(item, "title", None),
        author=getattr(item, "author", None),
        published_date=getattr(item, "published_date", None),
        text=getattr(item, "text", None),
        highlights=list(getattr(item, "highlights", None) or []),
    )


def run(args: argparse.Namespace) -> dict[str, Any]:
    api_key = os.environ.get("EXA_API_KEY")
    if not api_key:
        print("EXA_API_KEY environment variable is not set.", file=sys.stderr)
        sys.exit(2)

    client = Exa(api_key=api_key)
    client.headers["x-exa-integration"] = EXA_INTEGRATION_HEADER

    contents = _build_contents(args.text, args.highlights)
    response = client.get_contents(urls=args.urls, **contents)

    typed = [_to_typed(item) for item in getattr(response, "results", []) or []]
    return {
        "urls": list(args.urls),
        "num_results": len(typed),
        "results": [asdict(doc) for doc in typed],
    }


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Extract content from URLs with Exa.")
    parser.add_argument("urls", nargs="+", help="One or more URLs to extract.")
    parser.add_argument("--text", action="store_true", help="Return full-text content.")
    parser.add_argument("--highlights", action="store_true", help="Return extracted highlight snippets.")
    parser.add_argument("-o", "--output", default=None, help="Write JSON to this file (default: stdout).")
    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    payload = run(args)
    text = json.dumps(payload, indent=2, ensure_ascii=False)
    if args.output:
        with open(args.output, "w", encoding="utf-8") as fh:
            fh.write(text)
        print(f"Wrote {len(payload['results'])} documents to {args.output}")
    else:
        print(text)
    return 0


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

#!/usr/bin/env python3
# /// script
# requires-python = ">=3.11"
# dependencies = ["exa-py>=1.14.0"]
# ///
"""Run an Exa web search and write results to JSON.

Uses the Exa Python SDK. Auth via the EXA_API_KEY environment variable.

Example:
    uv run exa_search.py "transformer architectures" \\
        --category "research paper" \\
        --text --highlights \\
        -o results.json
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from dataclasses import asdict, dataclass, field
from typing import Any

try:
    from exa_py import Exa
except ImportError:
    print(
        "exa_py not installed. Run: uv pip install exa-py  (or invoke with: uv run --with exa-py)",
        file=sys.stderr,
    )
    sys.exit(2)


EXA_INTEGRATION_HEADER = "k-dense-ai--scientific-agent-skills"


@dataclass
class SearchResult:
    """Typed view of a single Exa search result for JSON export."""

    title: str | None
    url: str
    id: str | None
    author: str | None
    published_date: str | None
    score: float | None
    text: str | None = None
    highlights: list[str] = field(default_factory=list)
    highlight_scores: list[float] = field(default_factory=list)


def _split_csv(value: str | None) -> list[str] | None:
    if not value:
        return None
    items = [item.strip() for item in value.split(",") if item.strip()]
    return items or None


def _build_contents(text: bool, highlights: bool) -> dict[str, Any] | None:
    contents: dict[str, Any] = {}
    if text:
        contents["text"] = True
    if highlights:
        contents["highlights"] = True
    return contents or None


def _result_to_typed(item: Any) -> SearchResult:
    highlights = list(getattr(item, "highlights", None) or [])
    scores = list(getattr(item, "highlight_scores", None) or [])
    return SearchResult(
        title=getattr(item, "title", None),
        url=getattr(item, "url", ""),
        id=getattr(item, "id", None),
        author=getattr(item, "author", None),
        published_date=getattr(item, "published_date", None),
        score=getattr(item, "score", None),
        text=getattr(item, "text", None),
        highlights=highlights,
        highlight_scores=scores,
    )


def run(args: argparse.Namespace) -> dict[str, Any]:
    api_key = os.environ.get("EXA_API_KEY")
    if not api_key:
        print("EXA_API_KEY environment variable is not set.", file=sys.stderr)
        sys.exit(2)

    client = Exa(api_key=api_key)
    # Attribute API usage to this skill for integration tracking.
    client.headers["x-exa-integration"] = EXA_INTEGRATION_HEADER

    contents = _build_contents(args.text, args.highlights)

    kwargs: dict[str, Any] = {
        "query": args.query,
        "num_results": args.num_results,
        "type": args.type,
    }
    if args.category:
        kwargs["category"] = args.category
    if include := _split_csv(args.include_domains):
        kwargs["include_domains"] = include
    if exclude := _split_csv(args.exclude_domains):
        kwargs["exclude_domains"] = exclude
    if args.start_published_date:
        kwargs["start_published_date"] = args.start_published_date
    if args.end_published_date:
        kwargs["end_published_date"] = args.end_published_date
    if args.user_location:
        kwargs["user_location"] = args.user_location

    if contents is not None:
        response = client.search_and_contents(**kwargs, **contents)
    else:
        response = client.search(**kwargs)

    typed = [_result_to_typed(item) for item in getattr(response, "results", []) or []]
    return {
        "query": args.query,
        "type": args.type,
        "num_results": len(typed),
        "autoprompt_string": getattr(response, "autoprompt_string", None),
        "results": [asdict(result) for result in typed],
    }


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Search the web with Exa.")
    parser.add_argument("query", help="Natural-language search query.")
    parser.add_argument(
        "--type",
        default="auto",
        choices=["auto", "fast", "deep"],
        help="Search type. 'auto' is Exa's general-purpose search; 'fast' is lowest latency; 'deep' is highest quality at higher latency.",
    )
    parser.add_argument("--num-results", type=int, default=10, help="Number of results (1-100).")
    parser.add_argument(
        "--category",
        default=None,
        choices=[
            "company",
            "research paper",
            "news",
            "github",
            "personal site",
            "financial report",
            "people",
        ],
        help="Bias results toward a content category.",
    )
    parser.add_argument("--include-domains", default=None, help="Comma-separated allowlist.")
    parser.add_argument("--exclude-domains", default=None, help="Comma-separated blocklist.")
    parser.add_argument("--start-published-date", default=None, help="ISO date, e.g. 2024-01-01.")
    parser.add_argument("--end-published-date", default=None, help="ISO date, e.g. 2024-12-31.")
    parser.add_argument("--user-location", default=None, help="Two-letter ISO country code.")
    parser.add_argument("--text", action="store_true", help="Return full-text content per result.")
    parser.add_argument("--highlights", action="store_true", help="Return extracted highlight snippets.")
    parser.add_argument("-o", "--output", default=None, help="Write JSON to this file (default: stdout).")
    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    payload = run(args)
    text = json.dumps(payload, indent=2, ensure_ascii=False)
    if args.output:
        with open(args.output, "w", encoding="utf-8") as fh:
            fh.write(text)
        print(f"Wrote {len(payload['results'])} results to {args.output}")
    else:
        print(text)
    return 0


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

"""Unit tests for the exa-search skill scripts.

These tests mock the Exa SDK so they never hit the live API. Run from the
skill root:

    python -m unittest discover -s tests -v

Tests cover: CLI argument plumbing, the content-options builder (text /
highlights), CSV splitting for domain lists, the content fallback
cascade, and the x-exa-integration header wiring.
"""
from __future__ import annotations

import importlib.util
import json
import os
import sys
import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import MagicMock, patch

SKILL_ROOT = Path(__file__).resolve().parent.parent
SCRIPTS_DIR = SKILL_ROOT / "scripts"


def _load_script(name: str):
    """Load one of the scripts as a module regardless of cwd."""
    spec = importlib.util.spec_from_file_location(name, SCRIPTS_DIR / f"{name}.py")
    assert spec and spec.loader
    module = importlib.util.module_from_spec(spec)
    sys.modules[name] = module
    spec.loader.exec_module(module)
    return module


def _fake_result(**overrides):
    defaults = {
        "title": "Attention Is All You Need",
        "url": "https://arxiv.org/abs/1706.03762",
        "id": "abc123",
        "author": "Vaswani et al.",
        "published_date": "2017-06-12",
        "score": 0.91,
        "text": "Full paper text here",
        "highlights": ["The Transformer relies entirely on self-attention"],
        "highlight_scores": [0.87],
    }
    defaults.update(overrides)
    return SimpleNamespace(**defaults)


class BuildContentsTests(unittest.TestCase):
    """The content-options builder is shared logic across all three scripts."""

    def setUp(self):
        self.search_mod = _load_script("exa_search")
        self.extract_mod = _load_script("exa_extract")

    def test_search_no_flags_returns_none(self):
        self.assertIsNone(self.search_mod._build_contents(False, False))

    def test_search_text_only(self):
        self.assertEqual(self.search_mod._build_contents(True, False), {"text": True})

    def test_search_highlights_only(self):
        self.assertEqual(
            self.search_mod._build_contents(False, True),
            {"highlights": True},
        )

    def test_search_text_and_highlights_combine(self):
        self.assertEqual(
            self.search_mod._build_contents(True, True),
            {"text": True, "highlights": True},
        )

    def test_extract_defaults_to_text_when_nothing_specified(self):
        self.assertEqual(self.extract_mod._build_contents(False, False), {"text": True})


class SplitCsvTests(unittest.TestCase):
    def setUp(self):
        self.mod = _load_script("exa_search")

    def test_none_returns_none(self):
        self.assertIsNone(self.mod._split_csv(None))

    def test_empty_string_returns_none(self):
        self.assertIsNone(self.mod._split_csv(""))

    def test_splits_and_trims(self):
        self.assertEqual(
            self.mod._split_csv("arxiv.org, nature.com ,pubmed.gov"),
            ["arxiv.org", "nature.com", "pubmed.gov"],
        )

    def test_ignores_empty_segments(self):
        self.assertEqual(self.mod._split_csv("a,, b,"), ["a", "b"])


class ResultTypingTests(unittest.TestCase):
    """Verify the Any → typed-dataclass conversion handles missing fields."""

    def setUp(self):
        self.mod = _load_script("exa_search")

    def test_full_result(self):
        item = _fake_result()
        typed = self.mod._result_to_typed(item)
        self.assertEqual(typed.url, "https://arxiv.org/abs/1706.03762")
        self.assertEqual(typed.highlights, ["The Transformer relies entirely on self-attention"])
        self.assertEqual(typed.highlight_scores, [0.87])

    def test_missing_optional_fields_default_cleanly(self):
        item = SimpleNamespace(url="https://example.com")
        typed = self.mod._result_to_typed(item)
        self.assertEqual(typed.url, "https://example.com")
        self.assertIsNone(typed.title)
        self.assertIsNone(typed.text)
        self.assertEqual(typed.highlights, [])
        self.assertEqual(typed.highlight_scores, [])

    def test_highlights_only_response(self):
        """Content fallback: when only highlights are present, text must stay None."""
        item = _fake_result(text=None, highlights=["snippet A", "snippet B"])
        typed = self.mod._result_to_typed(item)
        self.assertIsNone(typed.text)
        self.assertEqual(typed.highlights, ["snippet A", "snippet B"])

    def test_text_only_response(self):
        item = _fake_result(highlights=[], highlight_scores=[])
        typed = self.mod._result_to_typed(item)
        self.assertEqual(typed.text, "Full paper text here")
        self.assertEqual(typed.highlights, [])


class IntegrationHeaderAndFlowTests(unittest.TestCase):
    """End-to-end: run() sets the integration header and calls the right SDK method."""

    def setUp(self):
        self.mod = _load_script("exa_search")

    def _run_with_mock(self, argv):
        fake_client = MagicMock()
        fake_client.headers = {}
        fake_client.search_and_contents.return_value = SimpleNamespace(
            results=[_fake_result()],
            autoprompt_string=None,
        )
        fake_client.search.return_value = SimpleNamespace(
            results=[_fake_result()],
            autoprompt_string=None,
        )
        # Mock at the source module so the import inside the script gets the fake.
        # NOTE: the script does `from exa_py import Exa` at module top, so the name
        # to patch is the already-imported binding on the script module itself.
        with patch.object(self.mod, "Exa", return_value=fake_client) as ctor, \
             patch.dict(os.environ, {"EXA_API_KEY": "test-key"}, clear=False):
            args = self.mod.build_parser().parse_args(argv)
            payload = self.mod.run(args)
        return payload, fake_client, ctor

    def test_integration_header_is_set(self):
        _, client, _ = self._run_with_mock(["what is attention", "--highlights"])
        self.assertEqual(
            client.headers.get("x-exa-integration"),
            "k-dense-ai--scientific-agent-skills",
        )

    def test_calls_search_and_contents_when_contents_requested(self):
        _, client, _ = self._run_with_mock(["query", "--text"])
        client.search_and_contents.assert_called_once()
        client.search.assert_not_called()

    def test_calls_search_when_no_contents_flags(self):
        _, client, _ = self._run_with_mock(["query"])
        client.search.assert_called_once()
        client.search_and_contents.assert_not_called()

    def test_domain_filters_are_split_and_passed(self):
        _, client, _ = self._run_with_mock(
            ["q", "--include-domains", "arxiv.org, nature.com", "--exclude-domains", "spam.com"]
        )
        kwargs = client.search.call_args.kwargs
        self.assertEqual(kwargs["include_domains"], ["arxiv.org", "nature.com"])
        self.assertEqual(kwargs["exclude_domains"], ["spam.com"])

    def test_missing_api_key_exits_with_code_2(self):
        # Remove EXA_API_KEY for this call only.
        with patch.dict(os.environ, {}, clear=True):
            with self.assertRaises(SystemExit) as ctx:
                self.mod.run(self.mod.build_parser().parse_args(["q"]))
            self.assertEqual(ctx.exception.code, 2)

    def test_output_file_written(self):
        with tempfile.TemporaryDirectory() as tmp:
            out = Path(tmp) / "results.json"
            fake_client = MagicMock()
            fake_client.headers = {}
            fake_client.search.return_value = SimpleNamespace(
                results=[_fake_result()],
                autoprompt_string=None,
            )
            with patch.object(self.mod, "Exa", return_value=fake_client), \
                 patch.dict(os.environ, {"EXA_API_KEY": "k"}, clear=False):
                rc = self.mod.main(["q", "-o", str(out)])
            self.assertEqual(rc, 0)
            self.assertTrue(out.exists())
            payload = json.loads(out.read_text())
            self.assertEqual(payload["num_results"], 1)
            self.assertEqual(payload["results"][0]["title"], "Attention Is All You Need")


if __name__ == "__main__":
    unittest.main()